work
[pspp] / rust / src / cooked.rs
1 use std::{borrow::Cow, cmp::Ordering, collections::HashMap, iter::repeat};
2
3 use crate::{
4     format::{Error as FormatError, Spec, UncheckedSpec},
5     identifier::{Error as IdError, Identifier},
6     raw::{self, MissingValues, VarType},
7     {endian::Endian, Compression},
8 };
9 use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
10 use encoding_rs::{DecoderResult, Encoding};
11 use num::integer::div_ceil;
12 use thiserror::Error as ThisError;
13
14 #[derive(ThisError, Debug)]
15 pub enum Error {
16     #[error("Variable record at offset {offset:#x} specifies width {width} not in valid range [-1,255).")]
17     InvalidVariableWidth { offset: u64, width: i32 },
18
19     #[error("This file has corrupted metadata written by a buggy version of PSPP.  To ensure that other software can read it correctly, save a new copy of the file.")]
20     InvalidLongMissingValueFormat,
21
22     #[error("File creation date {creation_date} is not in the expected format \"DD MMM YY\" format.  Using 01 Jan 1970.")]
23     InvalidCreationDate { creation_date: String },
24
25     #[error("File creation time {creation_time} is not in the expected format \"HH:MM:SS\" format.  Using midnight.")]
26     InvalidCreationTime { creation_time: String },
27
28     #[error("{id_error}  Renaming variable to {new_name}.")]
29     InvalidVariableName {
30         id_error: IdError,
31         new_name: Identifier,
32     },
33
34     #[error(
35         "Substituting {new_spec} for invalid print format on variable {variable}.  {format_error}"
36     )]
37     InvalidPrintFormat {
38         new_spec: Spec,
39         variable: Identifier,
40         format_error: FormatError,
41     },
42
43     #[error(
44         "Substituting {new_spec} for invalid write format on variable {variable}.  {format_error}"
45     )]
46     InvalidWriteFormat {
47         new_spec: Spec,
48         variable: Identifier,
49         format_error: FormatError,
50     },
51
52     #[error("Renaming variable with duplicate name {duplicate_name} to {new_name}.")]
53     DuplicateVariableName {
54         duplicate_name: Identifier,
55         new_name: Identifier,
56     },
57
58     #[error("Dictionary index {dict_index} is outside valid range [1,{max_index}].")]
59     InvalidDictIndex { dict_index: usize, max_index: usize },
60
61     #[error("Dictionary index {0} refers to a long string continuation.")]
62     DictIndexIsContinuation(usize),
63
64     #[error("Variables associated with value label are not all of identical type.  Variable {numeric_var} is numeric, but variable {string_var} is string.")]
65     ValueLabelsDifferentTypes {
66         numeric_var: Identifier,
67         string_var: Identifier,
68     },
69
70     #[error(
71         "Value labels may not be added to long string variable {0} using record types 3 or 4."
72     )]
73     InvalidLongStringValueLabel(Identifier),
74
75     #[error("Details TBD")]
76     TBD,
77 }
78
79 type DictIndex = usize;
80
81 pub struct Variable {
82     pub dict_index: DictIndex,
83     pub short_name: Identifier,
84     pub long_name: Option<Identifier>,
85     pub width: VarWidth,
86 }
87
88 pub struct Decoder {
89     pub compression: Option<Compression>,
90     pub endian: Endian,
91     pub encoding: &'static Encoding,
92     pub variables: HashMap<DictIndex, Variable>,
93     pub var_names: HashMap<Identifier, DictIndex>,
94     n_dict_indexes: usize,
95     n_generated_names: usize,
96 }
97
98 impl Decoder {
99     fn generate_name(&mut self) -> Identifier {
100         loop {
101             self.n_generated_names += 1;
102             let name = Identifier::new(&format!("VAR{:03}", self.n_generated_names), self.encoding)
103                 .unwrap();
104             if !self.var_names.contains_key(&name) {
105                 return name;
106             }
107             assert!(self.n_generated_names < usize::MAX);
108         }
109     }
110     fn decode_string<'a>(&self, input: &'a [u8], warn: &impl Fn(Error)) -> Cow<'a, str> {
111         let (output, malformed) = self.encoding.decode_without_bom_handling(input);
112         if malformed {
113             warn(Error::TBD);
114         }
115         output
116     }
117     fn get_var_by_index(&self, dict_index: usize) -> Result<&Variable, Error> {
118         let max_index = self.n_dict_indexes - 1;
119         if dict_index == 0 || dict_index as usize > max_index {
120             return Err(Error::InvalidDictIndex {
121                 dict_index,
122                 max_index,
123             });
124         }
125         let Some(variable) = self.variables.get(&dict_index) else {
126             return Err(Error::DictIndexIsContinuation(dict_index));
127         };
128         Ok(variable)
129     }
130
131     /// Returns `input` decoded from `self.encoding` into UTF-8 such that
132     /// re-encoding the result back into `self.encoding` will have exactly the
133     /// same length in bytes.
134     ///
135     /// XXX warn about errors?
136     fn decode_exact_length<'a>(&self, input: &'a [u8]) -> Cow<'a, str> {
137         if let (s, false) = self.encoding.decode_without_bom_handling(input) {
138             // This is the common case.  Usually there will be no errors.
139             s.into()
140         } else {
141             // Unusual case.  Don't bother to optimize it much.
142             let mut decoder = self.encoding.new_decoder_without_bom_handling();
143             let mut output = String::with_capacity(
144                 decoder
145                     .max_utf8_buffer_length_without_replacement(input.len())
146                     .unwrap(),
147             );
148             let mut rest = input;
149             while !rest.is_empty() {
150                 match decoder.decode_to_string_without_replacement(rest, &mut output, true) {
151                     (DecoderResult::InputEmpty, _) => break,
152                     (DecoderResult::OutputFull, _) => unreachable!(),
153                     (DecoderResult::Malformed(a, b), consumed) => {
154                         let skipped = a as usize + b as usize;
155                         output.extend(repeat('?').take(skipped));
156                         rest = &rest[consumed..];
157                     }
158                 }
159             }
160             assert_eq!(self.encoding.encode(&output).0.len(), input.len());
161             output.into()
162         }
163     }
164 }
165
166 pub trait Decode: Sized {
167     type Input;
168     fn decode(decoder: &Decoder, input: &Self::Input, warn: impl Fn(Error)) -> Result<Self, Error>;
169 }
170
171 #[derive(Clone)]
172 pub struct Header {
173     pub eye_catcher: String,
174     pub weight_index: Option<usize>,
175     pub n_cases: Option<u64>,
176     pub creation: NaiveDateTime,
177     pub file_label: String,
178 }
179
180 impl Decode for Header {
181     type Input = crate::raw::Header;
182
183     fn decode(decoder: &Decoder, input: &Self::Input, warn: impl Fn(Error)) -> Result<Self, Error> {
184         let eye_catcher = decoder.decode_string(&input.eye_catcher.0, &warn);
185         let file_label = decoder.decode_string(&input.file_label.0, &warn);
186         let creation_date = decoder.decode_string(&input.creation_date.0, &warn);
187         let creation_date = NaiveDate::parse_from_str(&creation_date, "%v").unwrap_or_else(|_| {
188             warn(Error::InvalidCreationDate {
189                 creation_date: creation_date.into(),
190             });
191             Default::default()
192         });
193         let creation_time = decoder.decode_string(&input.creation_time.0, &warn);
194         let creation_time =
195             NaiveTime::parse_from_str(&creation_time, "%H:%M:%S").unwrap_or_else(|_| {
196                 warn(Error::InvalidCreationTime {
197                     creation_time: creation_time.into(),
198                 });
199                 Default::default()
200             });
201         Ok(Header {
202             eye_catcher: eye_catcher.into(),
203             weight_index: input.weight_index.map(|n| n as usize),
204             n_cases: input.n_cases.map(|n| n as u64),
205             creation: NaiveDateTime::new(creation_date, creation_time),
206             file_label: file_label.into(),
207         })
208     }
209 }
210
211 #[derive(Copy, Clone, PartialEq, Eq)]
212 pub enum VarWidth {
213     Numeric,
214     String(u16),
215 }
216
217 impl PartialOrd for VarWidth {
218     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
219         match (self, other) {
220             (VarWidth::Numeric, VarWidth::Numeric) => Some(Ordering::Equal),
221             (VarWidth::String(a), VarWidth::String(b)) => Some(a.cmp(b)),
222             _ => None,
223         }
224     }
225 }
226
227 impl VarWidth {
228     fn n_dict_indexes(self) -> usize {
229         match self {
230             VarWidth::Numeric => 1,
231             VarWidth::String(w) => div_ceil(w as usize, 8),
232         }
233     }
234 }
235
236 impl From<VarWidth> for VarType {
237     fn from(source: VarWidth) -> Self {
238         match source {
239             VarWidth::Numeric => VarType::Numeric,
240             VarWidth::String(_) => VarType::String,
241         }
242     }
243 }
244
245 pub struct VariableRecord {
246     pub width: VarWidth,
247     pub name: Identifier,
248     pub print_format: Spec,
249     pub write_format: Spec,
250     pub missing_values: MissingValues,
251     pub label: Option<String>,
252 }
253
254 fn decode_format(raw: raw::Spec, width: VarWidth, warn: impl Fn(Spec, FormatError)) -> Spec {
255     UncheckedSpec::try_from(raw)
256         .and_then(Spec::try_from)
257         .and_then(|x| x.check_width_compatibility(width))
258         .unwrap_or_else(|error| {
259             let new_format = Spec::default_for_width(width);
260             warn(new_format, error);
261             new_format
262         })
263 }
264
265 impl VariableRecord {
266     pub fn decode(
267         decoder: &mut Decoder,
268         input: &crate::raw::Variable,
269         warn: impl Fn(Error),
270     ) -> Result<Option<VariableRecord>, Error> {
271         let width = match input.width {
272             0 => VarWidth::Numeric,
273             w @ 1..=255 => VarWidth::String(w as u16),
274             -1 => return Ok(None),
275             _ => {
276                 return Err(Error::InvalidVariableWidth {
277                     offset: input.offset,
278                     width: input.width,
279                 })
280             }
281         };
282         let name = decoder.decode_string(&input.name.0, &warn);
283         let name = match Identifier::new(&name, decoder.encoding) {
284             Ok(name) => {
285                 if !decoder.var_names.contains_key(&name) {
286                     name
287                 } else {
288                     let new_name = decoder.generate_name();
289                     warn(Error::DuplicateVariableName {
290                         duplicate_name: name.clone(),
291                         new_name: new_name.clone(),
292                     });
293                     new_name
294                 }
295             }
296             Err(id_error) => {
297                 let new_name = decoder.generate_name();
298                 warn(Error::InvalidVariableName {
299                     id_error,
300                     new_name: new_name.clone(),
301                 });
302                 new_name
303             }
304         };
305         let variable = Variable {
306             dict_index: decoder.n_dict_indexes,
307             short_name: name.clone(),
308             long_name: None,
309             width,
310         };
311         decoder.n_dict_indexes += width.n_dict_indexes();
312         assert!(decoder
313             .var_names
314             .insert(name.clone(), variable.dict_index)
315             .is_none());
316         assert!(decoder
317             .variables
318             .insert(variable.dict_index, variable)
319             .is_none());
320
321         let print_format = decode_format(input.print_format, width, |new_spec, format_error| {
322             warn(Error::InvalidPrintFormat {
323                 new_spec,
324                 variable: name.clone(),
325                 format_error,
326             })
327         });
328         let write_format = decode_format(input.write_format, width, |new_spec, format_error| {
329             warn(Error::InvalidWriteFormat {
330                 new_spec,
331                 variable: name.clone(),
332                 format_error,
333             })
334         });
335         let label = input
336             .label
337             .as_ref()
338             .map(|label| decoder.decode_string(&label.0, &warn).into());
339         Ok(Some(VariableRecord {
340             width,
341             name,
342             print_format,
343             write_format,
344             missing_values: input.missing_values.clone(),
345             label,
346         }))
347     }
348 }
349
350 #[derive(Clone)]
351 pub struct Document(Vec<String>);
352
353 impl Decode for Document {
354     type Input = crate::raw::Document;
355
356     fn decode(decoder: &Decoder, input: &Self::Input, warn: impl Fn(Error)) -> Result<Self, Error> {
357         Ok(Document(
358             input
359                 .lines
360                 .iter()
361                 .map(|s| decoder.decode_string(&s.0, &warn).into())
362                 .collect(),
363         ))
364     }
365 }
366
367 pub use crate::raw::FloatInfo;
368 pub use crate::raw::IntegerInfo;
369
370 trait TextRecord
371 where
372     Self: Sized,
373 {
374     const NAME: &'static str;
375     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error>;
376 }
377
378 pub struct VariableSet {
379     pub name: String,
380     pub vars: Vec<String>,
381 }
382
383 impl VariableSet {
384     fn parse(input: &str) -> Result<Self, Error> {
385         let (name, input) = input.split_once('=').ok_or(Error::TBD)?;
386         let vars = input.split_ascii_whitespace().map(String::from).collect();
387         Ok(VariableSet {
388             name: name.into(),
389             vars,
390         })
391     }
392 }
393
394 #[derive(Clone)]
395 pub enum Value {
396     Number(Option<f64>),
397     String(String),
398 }
399
400 impl Value {
401     pub fn decode(raw: raw::Value, decoder: &Decoder) -> Self {
402         match raw {
403             raw::Value::Number(x) => Value::Number(x),
404             raw::Value::String(s) => Value::String(decoder.decode_exact_length(&s.0).into()),
405         }
406     }
407 }
408
409 pub struct ValueLabelRecord {
410     pub var_type: VarType,
411     pub labels: Vec<(Value, String)>,
412     pub variables: Vec<Identifier>,
413 }
414
415 trait WarnOnError<T> {
416     fn warn_on_error<F: Fn(Error)>(self, warn: &F) -> Option<T>;
417 }
418 impl<T> WarnOnError<T> for Result<T, Error> {
419     fn warn_on_error<F: Fn(Error)>(self, warn: &F) -> Option<T> {
420         match self {
421             Ok(result) => Some(result),
422             Err(error) => {
423                 warn(error);
424                 None
425             }
426         }
427     }
428 }
429
430 impl ValueLabelRecord {
431     pub fn decode(
432         decoder: &mut Decoder,
433         raw_value_label: &crate::raw::ValueLabel,
434         dict_indexes: &crate::raw::VarIndexes,
435         warn: impl Fn(Error),
436     ) -> Result<Option<ValueLabelRecord>, Error> {
437         let variables: Vec<&Variable> = dict_indexes
438             .dict_indexes
439             .iter()
440             .filter_map(|&dict_index| {
441                 decoder
442                     .get_var_by_index(dict_index as usize)
443                     .warn_on_error(&warn)
444             })
445             .filter(|&variable| match variable.width {
446                 VarWidth::String(width) if width > 8 => {
447                     warn(Error::InvalidLongStringValueLabel(
448                         variable.short_name.clone(),
449                     ));
450                     false
451                 }
452                 _ => true,
453             })
454             .collect();
455         let mut i = variables.iter();
456         let Some(&first_var) = i.next() else {
457             return Ok(None);
458         };
459         let var_type: VarType = first_var.width.into();
460         for &variable in i {
461             let this_type: VarType = variable.width.into();
462             if var_type != this_type {
463                 let (numeric_var, string_var) = match var_type {
464                     VarType::Numeric => (first_var, variable),
465                     VarType::String => (variable, first_var),
466                 };
467                 warn(Error::ValueLabelsDifferentTypes {
468                     numeric_var: numeric_var.short_name.clone(),
469                     string_var: string_var.short_name.clone(),
470                 });
471                 return Ok(None);
472             }
473         }
474         let labels = raw_value_label
475             .labels
476             .iter()
477             .map(|(value, label)| {
478                 let label = decoder.decode_string(&label.0, &warn);
479                 let value = Value::decode(raw::Value::from_raw(*value, var_type, decoder.endian), &decoder);
480                 (value, label.into())
481             })
482             .collect();
483         let variables = variables
484             .iter()
485             .map(|&variable| variable.short_name.clone())
486             .collect();
487         Ok(Some(ValueLabelRecord {
488             var_type,
489             labels,
490             variables,
491         }))
492     }
493 }
494
495 pub struct VariableSetRecord(Vec<VariableSet>);
496
497 impl TextRecord for VariableSetRecord {
498     const NAME: &'static str = "variable set";
499     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
500         let mut sets = Vec::new();
501         for line in input.lines() {
502             if let Some(set) = VariableSet::parse(line).warn_on_error(&warn) {
503                 sets.push(set)
504             }
505         }
506         Ok(VariableSetRecord(sets))
507     }
508 }
509
510 pub struct ProductInfo(pub String);
511
512 impl TextRecord for ProductInfo {
513     const NAME: &'static str = "extra product info";
514     fn parse(input: &str, _warn: impl Fn(Error)) -> Result<Self, Error> {
515         Ok(ProductInfo(input.into()))
516     }
517 }
518
519 pub struct LongVariableName {
520     pub short_name: String,
521     pub long_name: String,
522 }
523
524 pub struct LongVariableNameRecord(Vec<LongVariableName>);
525
526 impl TextRecord for LongVariableNameRecord {
527     const NAME: &'static str = "long variable names";
528     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
529         let mut names = Vec::new();
530         for pair in input.split('\t').filter(|s| !s.is_empty()) {
531             if let Some((short_name, long_name)) = pair.split_once('=') {
532                 let name = LongVariableName {
533                     short_name: short_name.into(),
534                     long_name: long_name.into(),
535                 };
536                 names.push(name);
537             } else {
538                 warn(Error::TBD)
539             }
540         }
541         Ok(LongVariableNameRecord(names))
542     }
543 }
544
545 pub struct VeryLongString {
546     pub short_name: String,
547     pub length: usize,
548 }
549
550 impl VeryLongString {
551     fn parse(input: &str) -> Result<VeryLongString, Error> {
552         let Some((short_name, length)) = input.split_once('=') else {
553             return Err(Error::TBD);
554         };
555         let length: usize = length.parse().map_err(|_| Error::TBD)?;
556         Ok(VeryLongString {
557             short_name: short_name.into(),
558             length,
559         })
560     }
561 }
562
563 pub struct VeryLongStringRecord(Vec<VeryLongString>);
564
565 impl TextRecord for VeryLongStringRecord {
566     const NAME: &'static str = "very long strings";
567     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
568         let mut very_long_strings = Vec::new();
569         for tuple in input
570             .split('\0')
571             .map(|s| s.trim_end_matches('\t'))
572             .filter(|s| !s.is_empty())
573         {
574             if let Some(vls) = VeryLongString::parse(tuple).warn_on_error(&warn) {
575                 very_long_strings.push(vls)
576             }
577         }
578         Ok(VeryLongStringRecord(very_long_strings))
579     }
580 }
581
582 pub struct Attribute {
583     pub name: String,
584     pub values: Vec<String>,
585 }
586
587 impl Attribute {
588     fn parse<'a>(input: &'a str, warn: &impl Fn(Error)) -> Result<(Attribute, &'a str), Error> {
589         let Some((name, mut input)) = input.split_once('(') else {
590             return Err(Error::TBD);
591         };
592         let mut values = Vec::new();
593         loop {
594             let Some((value, rest)) = input.split_once('\n') else {
595                 return Err(Error::TBD);
596             };
597             if let Some(stripped) = value
598                 .strip_prefix('\'')
599                 .and_then(|value| value.strip_suffix('\''))
600             {
601                 values.push(stripped.into());
602             } else {
603                 warn(Error::TBD);
604                 values.push(value.into());
605             }
606             if let Some(rest) = rest.strip_prefix(')') {
607                 return Ok((
608                     Attribute {
609                         name: name.into(),
610                         values,
611                     },
612                     rest,
613                 ));
614             }
615             input = rest;
616         }
617     }
618 }
619
620 pub struct AttributeSet(pub Vec<Attribute>);
621
622 impl AttributeSet {
623     fn parse<'a>(
624         mut input: &'a str,
625         sentinel: Option<char>,
626         warn: &impl Fn(Error),
627     ) -> Result<(AttributeSet, &'a str), Error> {
628         let mut attributes = Vec::new();
629         let rest = loop {
630             match input.chars().next() {
631                 None => break input,
632                 c if c == sentinel => break &input[1..],
633                 _ => {
634                     let (attribute, rest) = Attribute::parse(input, &warn)?;
635                     attributes.push(attribute);
636                     input = rest;
637                 }
638             }
639         };
640         Ok((AttributeSet(attributes), rest))
641     }
642 }
643
644 pub struct FileAttributeRecord(AttributeSet);
645
646 impl TextRecord for FileAttributeRecord {
647     const NAME: &'static str = "data file attributes";
648     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
649         let (set, rest) = AttributeSet::parse(input, None, &warn)?;
650         if !rest.is_empty() {
651             warn(Error::TBD);
652         }
653         Ok(FileAttributeRecord(set))
654     }
655 }
656
657 pub struct VarAttributeSet {
658     pub long_var_name: String,
659     pub attributes: AttributeSet,
660 }
661
662 impl VarAttributeSet {
663     fn parse<'a>(
664         input: &'a str,
665         warn: &impl Fn(Error),
666     ) -> Result<(VarAttributeSet, &'a str), Error> {
667         let Some((long_var_name, rest)) = input.split_once(':') else {
668             return Err(Error::TBD);
669         };
670         let (attributes, rest) = AttributeSet::parse(rest, Some('/'), warn)?;
671         Ok((
672             VarAttributeSet {
673                 long_var_name: long_var_name.into(),
674                 attributes,
675             },
676             rest,
677         ))
678     }
679 }
680
681 pub struct VariableAttributeRecord(Vec<VarAttributeSet>);
682
683 impl TextRecord for VariableAttributeRecord {
684     const NAME: &'static str = "variable attributes";
685     fn parse(mut input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
686         let mut var_attribute_sets = Vec::new();
687         while !input.is_empty() {
688             let Some((var_attribute, rest)) =
689                 VarAttributeSet::parse(input, &warn).warn_on_error(&warn)
690             else {
691                 break;
692             };
693             var_attribute_sets.push(var_attribute);
694             input = rest;
695         }
696         Ok(VariableAttributeRecord(var_attribute_sets))
697     }
698 }
699
700 pub enum Measure {
701     Nominal,
702     Ordinal,
703     Scale,
704 }
705
706 pub enum Alignment {
707     Left,
708     Right,
709     Center,
710 }
711
712 pub struct VarDisplay {
713     pub measure: Option<Measure>,
714     pub width: u32,
715     pub align: Option<Alignment>,
716 }
717
718 pub struct VarDisplayRecord(pub Vec<VarDisplay>);
719
720 #[cfg(test)]
721 mod test {
722     use encoding_rs::WINDOWS_1252;
723
724     #[test]
725     fn test() {
726         let mut s = String::new();
727         s.push(char::REPLACEMENT_CHARACTER);
728         let encoded = WINDOWS_1252.encode(&s).0;
729         let decoded = WINDOWS_1252.decode(&encoded[..]).0;
730         println!("{:?}", decoded);
731     }
732
733     #[test]
734     fn test2() {
735         let charset: Vec<u8> = (0..=255).collect();
736         println!("{}", charset.len());
737         let decoded = WINDOWS_1252.decode(&charset[..]).0;
738         println!("{}", decoded.len());
739         let encoded = WINDOWS_1252.encode(&decoded[..]).0;
740         println!("{}", encoded.len());
741         assert_eq!(&charset[..], &encoded[..]);
742     }
743 }