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::HeaderRecord;
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::VariableRecord,
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::DocumentRecord;
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 trait TextRecord
368 where
369     Self: Sized,
370 {
371     const NAME: &'static str;
372     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error>;
373 }
374
375 pub struct VariableSet {
376     pub name: String,
377     pub vars: Vec<String>,
378 }
379
380 impl VariableSet {
381     fn parse(input: &str) -> Result<Self, Error> {
382         let (name, input) = input.split_once('=').ok_or(Error::TBD)?;
383         let vars = input.split_ascii_whitespace().map(String::from).collect();
384         Ok(VariableSet {
385             name: name.into(),
386             vars,
387         })
388     }
389 }
390
391 #[derive(Clone)]
392 pub enum Value {
393     Number(Option<f64>),
394     String(String),
395 }
396
397 impl Value {
398     pub fn decode(raw: raw::Value, decoder: &Decoder) -> Self {
399         match raw {
400             raw::Value::Number(x) => Value::Number(x),
401             raw::Value::String(s) => Value::String(decoder.decode_exact_length(&s.0).into()),
402         }
403     }
404 }
405
406 pub struct ValueLabelRecord {
407     pub var_type: VarType,
408     pub labels: Vec<(Value, String)>,
409     pub variables: Vec<Identifier>,
410 }
411
412 trait WarnOnError<T> {
413     fn warn_on_error<F: Fn(Error)>(self, warn: &F) -> Option<T>;
414 }
415 impl<T> WarnOnError<T> for Result<T, Error> {
416     fn warn_on_error<F: Fn(Error)>(self, warn: &F) -> Option<T> {
417         match self {
418             Ok(result) => Some(result),
419             Err(error) => {
420                 warn(error);
421                 None
422             }
423         }
424     }
425 }
426
427 impl ValueLabelRecord {
428     pub fn decode(
429         decoder: &mut Decoder,
430         raw_value_label: &crate::raw::ValueLabelRecord,
431         dict_indexes: &crate::raw::VarIndexRecord,
432         warn: impl Fn(Error),
433     ) -> Result<Option<ValueLabelRecord>, Error> {
434         let variables: Vec<&Variable> = dict_indexes
435             .dict_indexes
436             .iter()
437             .filter_map(|&dict_index| {
438                 decoder
439                     .get_var_by_index(dict_index as usize)
440                     .warn_on_error(&warn)
441             })
442             .filter(|&variable| match variable.width {
443                 VarWidth::String(width) if width > 8 => {
444                     warn(Error::InvalidLongStringValueLabel(
445                         variable.short_name.clone(),
446                     ));
447                     false
448                 }
449                 _ => true,
450             })
451             .collect();
452         let mut i = variables.iter();
453         let Some(&first_var) = i.next() else {
454             return Ok(None);
455         };
456         let var_type: VarType = first_var.width.into();
457         for &variable in i {
458             let this_type: VarType = variable.width.into();
459             if var_type != this_type {
460                 let (numeric_var, string_var) = match var_type {
461                     VarType::Numeric => (first_var, variable),
462                     VarType::String => (variable, first_var),
463                 };
464                 warn(Error::ValueLabelsDifferentTypes {
465                     numeric_var: numeric_var.short_name.clone(),
466                     string_var: string_var.short_name.clone(),
467                 });
468                 return Ok(None);
469             }
470         }
471         let labels = raw_value_label
472             .labels
473             .iter()
474             .map(|(value, label)| {
475                 let label = decoder.decode_string(&label.0, &warn);
476                 let value = Value::decode(raw::Value::from_raw(*value, var_type, decoder.endian), &decoder);
477                 (value, label.into())
478             })
479             .collect();
480         let variables = variables
481             .iter()
482             .map(|&variable| variable.short_name.clone())
483             .collect();
484         Ok(Some(ValueLabelRecord {
485             var_type,
486             labels,
487             variables,
488         }))
489     }
490 }
491
492 pub struct VariableSetRecord(Vec<VariableSet>);
493
494 impl TextRecord for VariableSetRecord {
495     const NAME: &'static str = "variable set";
496     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
497         let mut sets = Vec::new();
498         for line in input.lines() {
499             if let Some(set) = VariableSet::parse(line).warn_on_error(&warn) {
500                 sets.push(set)
501             }
502         }
503         Ok(VariableSetRecord(sets))
504     }
505 }
506
507 pub struct ProductInfo(pub String);
508
509 impl TextRecord for ProductInfo {
510     const NAME: &'static str = "extra product info";
511     fn parse(input: &str, _warn: impl Fn(Error)) -> Result<Self, Error> {
512         Ok(ProductInfo(input.into()))
513     }
514 }
515
516 pub struct LongVariableName {
517     pub short_name: String,
518     pub long_name: String,
519 }
520
521 pub struct LongVariableNameRecord(Vec<LongVariableName>);
522
523 impl TextRecord for LongVariableNameRecord {
524     const NAME: &'static str = "long variable names";
525     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
526         let mut names = Vec::new();
527         for pair in input.split('\t').filter(|s| !s.is_empty()) {
528             if let Some((short_name, long_name)) = pair.split_once('=') {
529                 let name = LongVariableName {
530                     short_name: short_name.into(),
531                     long_name: long_name.into(),
532                 };
533                 names.push(name);
534             } else {
535                 warn(Error::TBD)
536             }
537         }
538         Ok(LongVariableNameRecord(names))
539     }
540 }
541
542 pub struct VeryLongString {
543     pub short_name: String,
544     pub length: usize,
545 }
546
547 impl VeryLongString {
548     fn parse(input: &str) -> Result<VeryLongString, Error> {
549         let Some((short_name, length)) = input.split_once('=') else {
550             return Err(Error::TBD);
551         };
552         let length: usize = length.parse().map_err(|_| Error::TBD)?;
553         Ok(VeryLongString {
554             short_name: short_name.into(),
555             length,
556         })
557     }
558 }
559
560 pub struct VeryLongStringRecord(Vec<VeryLongString>);
561
562 impl TextRecord for VeryLongStringRecord {
563     const NAME: &'static str = "very long strings";
564     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
565         let mut very_long_strings = Vec::new();
566         for tuple in input
567             .split('\0')
568             .map(|s| s.trim_end_matches('\t'))
569             .filter(|s| !s.is_empty())
570         {
571             if let Some(vls) = VeryLongString::parse(tuple).warn_on_error(&warn) {
572                 very_long_strings.push(vls)
573             }
574         }
575         Ok(VeryLongStringRecord(very_long_strings))
576     }
577 }
578
579 pub struct Attribute {
580     pub name: String,
581     pub values: Vec<String>,
582 }
583
584 impl Attribute {
585     fn parse<'a>(input: &'a str, warn: &impl Fn(Error)) -> Result<(Attribute, &'a str), Error> {
586         let Some((name, mut input)) = input.split_once('(') else {
587             return Err(Error::TBD);
588         };
589         let mut values = Vec::new();
590         loop {
591             let Some((value, rest)) = input.split_once('\n') else {
592                 return Err(Error::TBD);
593             };
594             if let Some(stripped) = value
595                 .strip_prefix('\'')
596                 .and_then(|value| value.strip_suffix('\''))
597             {
598                 values.push(stripped.into());
599             } else {
600                 warn(Error::TBD);
601                 values.push(value.into());
602             }
603             if let Some(rest) = rest.strip_prefix(')') {
604                 return Ok((
605                     Attribute {
606                         name: name.into(),
607                         values,
608                     },
609                     rest,
610                 ));
611             }
612             input = rest;
613         }
614     }
615 }
616
617 pub struct AttributeSet(pub Vec<Attribute>);
618
619 impl AttributeSet {
620     fn parse<'a>(
621         mut input: &'a str,
622         sentinel: Option<char>,
623         warn: &impl Fn(Error),
624     ) -> Result<(AttributeSet, &'a str), Error> {
625         let mut attributes = Vec::new();
626         let rest = loop {
627             match input.chars().next() {
628                 None => break input,
629                 c if c == sentinel => break &input[1..],
630                 _ => {
631                     let (attribute, rest) = Attribute::parse(input, &warn)?;
632                     attributes.push(attribute);
633                     input = rest;
634                 }
635             }
636         };
637         Ok((AttributeSet(attributes), rest))
638     }
639 }
640
641 pub struct FileAttributeRecord(AttributeSet);
642
643 impl TextRecord for FileAttributeRecord {
644     const NAME: &'static str = "data file attributes";
645     fn parse(input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
646         let (set, rest) = AttributeSet::parse(input, None, &warn)?;
647         if !rest.is_empty() {
648             warn(Error::TBD);
649         }
650         Ok(FileAttributeRecord(set))
651     }
652 }
653
654 pub struct VarAttributeSet {
655     pub long_var_name: String,
656     pub attributes: AttributeSet,
657 }
658
659 impl VarAttributeSet {
660     fn parse<'a>(
661         input: &'a str,
662         warn: &impl Fn(Error),
663     ) -> Result<(VarAttributeSet, &'a str), Error> {
664         let Some((long_var_name, rest)) = input.split_once(':') else {
665             return Err(Error::TBD);
666         };
667         let (attributes, rest) = AttributeSet::parse(rest, Some('/'), warn)?;
668         Ok((
669             VarAttributeSet {
670                 long_var_name: long_var_name.into(),
671                 attributes,
672             },
673             rest,
674         ))
675     }
676 }
677
678 pub struct VariableAttributeRecord(Vec<VarAttributeSet>);
679
680 impl TextRecord for VariableAttributeRecord {
681     const NAME: &'static str = "variable attributes";
682     fn parse(mut input: &str, warn: impl Fn(Error)) -> Result<Self, Error> {
683         let mut var_attribute_sets = Vec::new();
684         while !input.is_empty() {
685             let Some((var_attribute, rest)) =
686                 VarAttributeSet::parse(input, &warn).warn_on_error(&warn)
687             else {
688                 break;
689             };
690             var_attribute_sets.push(var_attribute);
691             input = rest;
692         }
693         Ok(VariableAttributeRecord(var_attribute_sets))
694     }
695 }
696
697 pub enum Measure {
698     Nominal,
699     Ordinal,
700     Scale,
701 }
702
703 pub enum Alignment {
704     Left,
705     Right,
706     Center,
707 }
708
709 pub struct VarDisplay {
710     pub measure: Option<Measure>,
711     pub width: u32,
712     pub align: Option<Alignment>,
713 }
714
715 pub struct VarDisplayRecord(pub Vec<VarDisplay>);
716
717 #[cfg(test)]
718 mod test {
719     use encoding_rs::WINDOWS_1252;
720
721     #[test]
722     fn test() {
723         let mut s = String::new();
724         s.push(char::REPLACEMENT_CHARACTER);
725         let encoded = WINDOWS_1252.encode(&s).0;
726         let decoded = WINDOWS_1252.decode(&encoded[..]).0;
727         println!("{:?}", decoded);
728     }
729
730     #[test]
731     fn test2() {
732         let charset: Vec<u8> = (0..=255).collect();
733         println!("{}", charset.len());
734         let decoded = WINDOWS_1252.decode(&charset[..]).0;
735         println!("{}", decoded.len());
736         let encoded = WINDOWS_1252.encode(&decoded[..]).0;
737         println!("{}", encoded.len());
738         assert_eq!(&charset[..], &encoded[..]);
739     }
740 }