work
[pspp] / rust / src / lib.rs
1 #![allow(unused_variables)]
2 use endian::{Endian, Parse};
3 use num::Integer;
4 use num_derive::FromPrimitive;
5 use std::io::{BufReader, Error as IoError, ErrorKind, Read, Seek};
6 use thiserror::Error;
7
8 pub mod endian;
9
10 #[derive(Error, Debug)]
11 pub enum Error {
12     #[error("Not an SPSS system file")]
13     NotASystemFile,
14
15     #[error("Invalid magic number {0:?}")]
16     BadMagic([u8; 4]),
17
18     #[error("I/O error ({0})")]
19     Io(
20         #[from]
21         IoError,
22     ),
23
24     #[error("Invalid SAV compression code {0}")]
25     InvalidSavCompression(u32),
26
27     #[error("Invalid ZSAV compression code {0}")]
28     InvalidZsavCompression(u32),
29
30     #[error("Variable record at offset {offset:#x} specifies width {width} not in valid range [-1,255).")]
31     BadVariableWidth { offset: u64, width: i32 },
32
33     #[error("Misplaced type 4 record near offset {0:#x}.")]
34     MisplacedType4Record(u64),
35
36     #[error("Document record at offset {offset:#x} has document line count ({n}) greater than the maximum number {max}.")]
37     BadDocumentLength { offset: u64, n: u32, max: u32 },
38
39     #[error("At offset {offset:#x}, Unrecognized record type {rec_type}.")]
40     BadRecordType { offset: u64, rec_type: u32 },
41
42     #[error("At offset {offset:#x}, variable label code ({code}) is not 0 or 1.")]
43     BadVariableLabelCode { offset: u64, code: u32 },
44
45     #[error(
46         "At offset {offset:#x}, numeric missing value code ({code}) is not -3, -2, 0, 1, 2, or 3."
47     )]
48     BadNumericMissingValueCode { offset: u64, code: i32 },
49
50     #[error("At offset {offset:#x}, string missing value code ({code}) is not 0, 1, 2, or 3.")]
51     BadStringMissingValueCode { offset: u64, code: i32 },
52
53     #[error("At offset {offset:#x}, number of value labels ({n}) is greater than the maximum number {max}.")]
54     BadNumberOfValueLabels { offset: u64, n: u32, max: u32 },
55
56     #[error("At offset {offset:#x}, variable index record (type 4) does not immediately follow value label record (type 3) as it should.")]
57     MissingVariableIndexRecord { offset: u64 },
58
59     #[error("At offset {offset:#x}, number of variables indexes ({n}) is greater than the maximum number ({max}).")]
60     BadNumberOfVarIndexes { offset: u64, n: u32, max: u32 },
61
62     #[error("At offset {offset:#x}, record type 7 subtype {subtype} is too large with element size {size} and {count} elements.")]
63     ExtensionRecordTooLarge {
64         offset: u64,
65         subtype: u32,
66         size: u32,
67         count: u32,
68     },
69
70     #[error("Wrong ZLIB data header offset {zheader_offset:#x} (expected {offset:#x}).")]
71     BadZlibHeaderOffset { offset: u64, zheader_offset: u64 },
72
73     #[error("At offset {offset:#x}, impossible ZLIB trailer offset {ztrailer_offset:#x}.")]
74     BadZlibTrailerOffset { offset: u64, ztrailer_offset: u64 },
75
76     #[error("At offset {offset:#x}, impossible ZLIB trailer length {ztrailer_len}.")]
77     BadZlibTrailerLen { offset: u64, ztrailer_len: u64 },
78 }
79
80 #[derive(Error, Debug)]
81 pub enum Warning {
82     #[error("Unexpected floating-point bias {0} or unrecognized floating-point format.")]
83     UnexpectedBias(f64),
84
85     #[error("Duplicate type 6 (document) record.")]
86     DuplicateDocumentRecord,
87 }
88
89 #[derive(Copy, Clone, Debug)]
90 pub enum Compression {
91     Simple,
92     ZLib,
93 }
94
95 pub enum Record {
96     Header(Header),
97     Document(Document),
98     Variable(Variable),
99     ValueLabel(ValueLabel),
100     VarIndexes(VarIndexes),
101     Extension(Extension),
102     EndOfHeaders,
103     Case(Vec<Value>),
104 }
105
106 pub struct Header {
107     /// Magic number.
108     pub magic: Magic,
109
110     /// Eye-catcher string, product name, in the file's encoding.  Padded
111     /// on the right with spaces.
112     pub eye_catcher: [u8; 60],
113
114     /// Layout code, normally either 2 or 3.
115     pub layout_code: u32,
116
117     /// Number of variable positions, or `None` if the value in the file is
118     /// questionably trustworthy.
119     pub nominal_case_size: Option<u32>,
120
121     /// Compression type, if any,
122     pub compression: Option<Compression>,
123
124     /// 0-based variable index of the weight variable, or `None` if the file is
125     /// unweighted.
126     pub weight_index: Option<u32>,
127
128     /// Claimed number of cases, if known.
129     pub n_cases: Option<u32>,
130
131     /// Compression bias, usually 100.0.
132     pub bias: f64,
133
134     /// `dd mmm yy` in the file's encoding.
135     pub creation_date: [u8; 9],
136
137     /// `HH:MM:SS` in the file's encoding.
138     pub creation_time: [u8; 8],
139
140     /// File label, in the file's encoding.  Padded on the right with spaces.
141     pub file_label: [u8; 64],
142
143     /// Endianness of the data in the file header.
144     pub endianness: Endian,
145 }
146
147 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
148 pub struct Magic([u8; 4]);
149
150 impl Magic {
151     /// Magic number for a regular system file.
152     pub const SAV: Magic = Magic(*b"$FL2");
153
154     /// Magic number for a system file that contains zlib-compressed data.
155     pub const ZSAV: Magic = Magic(*b"$FL3");
156
157     /// Magic number for an EBDIC-encoded system file.  This is `$FL2` encoded
158     /// in EBCDIC.
159     pub const EBCDIC: Magic = Magic([0x5b, 0xc6, 0xd3, 0xf2]);
160 }
161
162 impl TryFrom<[u8; 4]> for Magic {
163     type Error = Error;
164
165     fn try_from(value: [u8; 4]) -> Result<Self, Self::Error> {
166         let magic = Magic(value);
167         match magic {
168             Magic::SAV | Magic::ZSAV | Magic::EBCDIC => Ok(magic),
169             _ => Err(Error::BadMagic(value)),
170         }
171     }
172 }
173
174 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
175 pub enum VarType {
176     Number,
177     String,
178 }
179
180 impl VarType {
181     fn from_width(width: i32) -> VarType {
182         match width {
183             0 => VarType::Number,
184             _ => VarType::String,
185         }
186     }
187 }
188
189 pub struct Reader<R: Read> {
190     r: BufReader<R>,
191     var_types: Vec<VarType>,
192     state: ReaderState,
193 }
194
195 enum ReaderState {
196     Start,
197     Headers(Endian, Option<Compression>),
198     Data(Endian),
199     End,
200 }
201
202 #[derive(Copy, Clone)]
203 pub enum Value {
204     Number(Option<f64>),
205     String([u8; 8]),
206 }
207
208 impl Value {
209     pub fn from_raw(var_type: VarType, raw: [u8; 8], endian: Endian) -> Value {
210         match var_type {
211             VarType::String => Value::String(raw),
212             VarType::Number => {
213                 let number: f64 = endian.parse(raw);
214                 Value::Number((number != -f64::MAX).then_some(number))
215             }
216         }
217     }
218 }
219
220 impl<R: Read + Seek> Reader<R> {
221     pub fn new(r: R) -> Result<Reader<R>, Error> {
222         Ok(Reader {
223             r: BufReader::new(r),
224             var_types: Vec::new(),
225             state: ReaderState::Start,
226         })
227     }
228     fn _next(&mut self) -> Result<Option<(Record, ReaderState)>, Error> {
229         match self.state {
230             ReaderState::Start => {
231                 let header = read_header(&mut self.r)?;
232                 let next_state = ReaderState::Headers(header.endianness, header.compression);
233                 Ok(Some((Record::Header(header), next_state)))
234             }
235             ReaderState::Headers(endian, compression) => {
236                 let rec_type: u32 = endian.parse(read_bytes(&mut self.r)?);
237                 let record = match rec_type {
238                     2 => {
239                         let variable = read_variable_record(&mut self.r, endian)?;
240                         self.var_types.push(VarType::from_width(variable.width));
241                         Record::Variable(variable)
242                     }
243                     3 => Record::ValueLabel(read_value_label_record(&mut self.r, endian)?),
244                     4 => Record::VarIndexes(read_var_indexes_record(&mut self.r, endian)?),
245                     6 => Record::Document(read_document_record(&mut self.r, endian)?),
246                     7 => Record::Extension(read_extension_record(&mut self.r, endian)?),
247                     999 => {
248                         let _: [u8; 4] = read_bytes(&mut self.r)?;
249                         let next_state = match compression {
250                             None => ReaderState::Data(endian),
251                             _ => ReaderState::End,
252                         };
253                         return Ok(Some((Record::EndOfHeaders, next_state)));
254                     }
255                     _ => {
256                         return Err(Error::BadRecordType {
257                             offset: self.r.stream_position()?,
258                             rec_type,
259                         })
260                     }
261                 };
262                 Ok(Some((record, ReaderState::Headers(endian, compression))))
263             }
264             ReaderState::Data(endian) => {
265                 let mut values = Vec::with_capacity(self.var_types.len());
266                 for (i, &var_type) in self.var_types.iter().enumerate() {
267                     let raw = match read_bytes(&mut self.r) {
268                         Ok(raw) => raw,
269                         Err(err) => {
270                             if i == 0 && err.kind() == ErrorKind::UnexpectedEof {
271                                 return Ok(None);
272                             } else {
273                                 return Err(Error::Io(err));
274                             }
275                         }
276                     };
277                     values.push(Value::from_raw(var_type, raw, endian));
278                 }
279                 Ok(Some((Record::Case(values), ReaderState::Data(endian))))
280             }
281             ReaderState::End => Ok(None),
282         }
283     }
284 }
285
286 impl<R: Read + Seek> Iterator for Reader<R> {
287     type Item = Result<Record, Error>;
288
289     fn next(&mut self) -> Option<Self::Item> {
290         let retval = self._next();
291         match retval {
292             Ok(None) => {
293                 self.state = ReaderState::End;
294                 None
295             }
296             Ok(Some((record, next_state))) => {
297                 self.state = next_state;
298                 Some(Ok(record))
299             }
300             Err(error) => {
301                 self.state = ReaderState::End;
302                 Some(Err(error))
303             }
304         }
305     }
306 }
307
308 fn read_header<R: Read>(r: &mut R) -> Result<Header, Error> {
309     let magic: [u8; 4] = read_bytes(r)?;
310     let magic: Magic = magic.try_into().map_err(|_| Error::NotASystemFile)?;
311
312     let eye_catcher: [u8; 60] = read_bytes(r)?;
313     let layout_code: [u8; 4] = read_bytes(r)?;
314     let endianness = Endian::identify_u32(2, layout_code)
315         .or_else(|| Endian::identify_u32(2, layout_code))
316         .ok_or_else(|| Error::NotASystemFile)?;
317     let layout_code = endianness.parse(layout_code);
318
319     let nominal_case_size: u32 = endianness.parse(read_bytes(r)?);
320     let nominal_case_size =
321         (nominal_case_size <= i32::MAX as u32 / 16).then_some(nominal_case_size);
322
323     let compression_code: u32 = endianness.parse(read_bytes(r)?);
324     let compression = match (magic, compression_code) {
325         (Magic::ZSAV, 2) => Some(Compression::ZLib),
326         (Magic::ZSAV, code) => return Err(Error::InvalidZsavCompression(code)),
327         (_, 0) => None,
328         (_, 1) => Some(Compression::Simple),
329         (_, code) => return Err(Error::InvalidSavCompression(code)),
330     };
331
332     let weight_index: u32 = endianness.parse(read_bytes(r)?);
333     let weight_index = (weight_index > 0).then_some(weight_index - 1);
334
335     let n_cases: u32 = endianness.parse(read_bytes(r)?);
336     let n_cases = (n_cases < i32::MAX as u32 / 2).then_some(n_cases);
337
338     let bias: f64 = endianness.parse(read_bytes(r)?);
339
340     let creation_date: [u8; 9] = read_bytes(r)?;
341     let creation_time: [u8; 8] = read_bytes(r)?;
342     let file_label: [u8; 64] = read_bytes(r)?;
343     let _: [u8; 3] = read_bytes(r)?;
344
345     Ok(Header {
346         magic,
347         layout_code,
348         nominal_case_size,
349         compression,
350         weight_index,
351         n_cases,
352         bias,
353         creation_date,
354         creation_time,
355         eye_catcher,
356         file_label,
357         endianness,
358     })
359 }
360
361 pub struct Variable {
362     /// Offset from the start of the file to the start of the record.
363     pub offset: u64,
364
365     /// Variable width, in the range -1..=255.
366     pub width: i32,
367
368     /// Variable name, padded on the right with spaces.
369     pub name: [u8; 8],
370
371     /// Print format.
372     pub print_format: u32,
373
374     /// Write format.
375     pub write_format: u32,
376
377     /// Missing value code, one of -3, -2, 0, 1, 2, or 3.
378     pub missing_value_code: i32,
379
380     /// Raw missing values, up to 3 of them.
381     pub missing: Vec<[u8; 8]>,
382
383     /// Optional variable label.
384     pub label: Option<Vec<u8>>,
385 }
386
387 fn read_variable_record<R: Read + Seek>(
388     r: &mut BufReader<R>,
389     e: Endian,
390 ) -> Result<Variable, Error> {
391     let offset = r.stream_position()?;
392     let width: i32 = e.parse(read_bytes(r)?);
393     let has_variable_label: u32 = e.parse(read_bytes(r)?);
394     let missing_value_code: i32 = e.parse(read_bytes(r)?);
395     let print_format: u32 = e.parse(read_bytes(r)?);
396     let write_format: u32 = e.parse(read_bytes(r)?);
397     let name: [u8; 8] = read_bytes(r)?;
398
399     let label = match has_variable_label {
400         0 => None,
401         1 => {
402             let len: u32 = e.parse(read_bytes(r)?);
403             let read_len = len.min(65535) as usize;
404             let label = Some(read_vec(r, read_len)?);
405
406             let padding_bytes = Integer::next_multiple_of(&len, &4) - len;
407             let _ = read_vec(r, padding_bytes as usize)?;
408
409             label
410         }
411         _ => {
412             return Err(Error::BadVariableLabelCode {
413                 offset,
414                 code: has_variable_label,
415             })
416         }
417     };
418
419     let mut missing = Vec::new();
420     if missing_value_code != 0 {
421         match (width, missing_value_code) {
422             (0, -3 | -2 | 1 | 2 | 3) => (),
423             (0, _) => {
424                 return Err(Error::BadNumericMissingValueCode {
425                     offset,
426                     code: missing_value_code,
427                 })
428             }
429             (_, 0..=3) => (),
430             (_, _) => {
431                 return Err(Error::BadStringMissingValueCode {
432                     offset,
433                     code: missing_value_code,
434                 })
435             }
436         }
437
438         for _ in 0..missing_value_code.abs() {
439             missing.push(read_bytes(r)?);
440         }
441     }
442
443     Ok(Variable {
444         offset,
445         width,
446         name,
447         print_format,
448         write_format,
449         missing_value_code,
450         missing,
451         label,
452     })
453 }
454
455 pub struct ValueLabel {
456     /// Offset from the start of the file to the start of the record.
457     pub offset: u64,
458
459     /// The labels.
460     pub labels: Vec<([u8; 8], Vec<u8>)>,
461 }
462
463 impl ValueLabel {
464     /// Maximum number of value labels in a record.
465     pub const MAX: u32 = u32::MAX / 8;
466 }
467
468 fn read_value_label_record<R: Read + Seek>(
469     r: &mut BufReader<R>,
470     e: Endian,
471 ) -> Result<ValueLabel, Error> {
472     let offset = r.stream_position()?;
473     let n: u32 = e.parse(read_bytes(r)?);
474     if n > ValueLabel::MAX {
475         return Err(Error::BadNumberOfValueLabels {
476             offset,
477             n,
478             max: ValueLabel::MAX,
479         });
480     }
481
482     let mut labels = Vec::new();
483     for _ in 0..n {
484         let value: [u8; 8] = read_bytes(r)?;
485         let label_len: u8 = e.parse(read_bytes(r)?);
486         let label_len = label_len as usize;
487         let padded_len = Integer::next_multiple_of(&(label_len + 1), &8);
488
489         let mut label = read_vec(r, padded_len)?;
490         label.truncate(label_len);
491         labels.push((value, label));
492     }
493     Ok(ValueLabel { offset, labels })
494 }
495
496 pub struct VarIndexes {
497     /// Offset from the start of the file to the start of the record.
498     pub offset: u64,
499
500     /// The 0-based indexes of the variable indexes.
501     pub var_indexes: Vec<u32>,
502 }
503
504 impl VarIndexes {
505     /// Maximum number of variable indexes in a record.
506     pub const MAX: u32 = u32::MAX / 8;
507 }
508
509 fn read_var_indexes_record<R: Read + Seek>(
510     r: &mut BufReader<R>,
511     e: Endian,
512 ) -> Result<VarIndexes, Error> {
513     let offset = r.stream_position()?;
514     let n: u32 = e.parse(read_bytes(r)?);
515     if n > VarIndexes::MAX {
516         return Err(Error::BadNumberOfVarIndexes {
517             offset,
518             n,
519             max: VarIndexes::MAX,
520         });
521     }
522     let mut var_indexes = Vec::with_capacity(n as usize);
523     for _ in 0..n {
524         var_indexes.push(e.parse(read_bytes(r)?));
525     }
526
527     Ok(VarIndexes {
528         offset,
529         var_indexes,
530     })
531 }
532
533 pub const DOC_LINE_LEN: u32 = 80;
534 pub const DOC_MAX_LINES: u32 = i32::MAX as u32 / DOC_LINE_LEN;
535
536 pub struct Document {
537     /// Offset from the start of the file to the start of the record.
538     pub pos: u64,
539
540     /// The document, as an array of 80-byte lines.
541     pub lines: Vec<[u8; DOC_LINE_LEN as usize]>,
542 }
543
544 fn read_document_record<R: Read + Seek>(
545     r: &mut BufReader<R>,
546     e: Endian,
547 ) -> Result<Document, Error> {
548     let offset = r.stream_position()?;
549     let n: u32 = e.parse(read_bytes(r)?);
550     match n {
551         0..=DOC_MAX_LINES => {
552             let pos = r.stream_position()?;
553             let mut lines = Vec::with_capacity(n as usize);
554             for _ in 0..n {
555                 let line: [u8; 80] = read_bytes(r)?;
556                 lines.push(line);
557             }
558             Ok(Document { pos, lines })
559         }
560         _ => Err(Error::BadDocumentLength {
561             offset,
562             n,
563             max: DOC_MAX_LINES,
564         }),
565     }
566 }
567
568 #[derive(FromPrimitive)]
569 enum ExtensionType {
570     /// Machine integer info.
571     Integer = 3,
572     /// Machine floating-point info.
573     Float = 4,
574     /// Variable sets.
575     VarSets = 5,
576     /// DATE.
577     Date = 6,
578     /// Multiple response sets.
579     Mrsets = 7,
580     /// SPSS Data Entry.
581     DataEntry = 8,
582     /// Extra product info text.
583     ProductInfo = 10,
584     /// Variable display parameters.
585     Display = 11,
586     /// Long variable names.
587     LongNames = 13,
588     /// Long strings.
589     LongStrings = 14,
590     /// Extended number of cases.
591     Ncases = 16,
592     /// Data file attributes.
593     FileAttrs = 17,
594     /// Variable attributes.
595     VarAttrs = 18,
596     /// Multiple response sets (extended).
597     Mrsets2 = 19,
598     /// Character encoding.
599     Encoding = 20,
600     /// Value labels for long strings.
601     LongLabels = 21,
602     /// Missing values for long strings.
603     LongMissing = 22,
604     /// "Format properties in dataview table".
605     Dataview = 24,
606 }
607
608 pub struct Extension {
609     /// Offset from the start of the file to the start of the record.
610     pub offset: u64,
611
612     /// Record subtype.
613     pub subtype: u32,
614
615     /// Size of each data element.
616     pub size: u32,
617
618     /// Number of data elements.
619     pub count: u32,
620
621     /// `size * count` bytes of data.
622     pub data: Vec<u8>,
623 }
624
625 fn extension_record_size_requirements(extension: ExtensionType) -> (u32, u32) {
626     match extension {
627         /* Implemented record types. */
628         ExtensionType::Integer => (4, 8),
629         ExtensionType::Float => (8, 3),
630         ExtensionType::VarSets => (1, 0),
631         ExtensionType::Mrsets => (1, 0),
632         ExtensionType::ProductInfo => (1, 0),
633         ExtensionType::Display => (4, 0),
634         ExtensionType::LongNames => (1, 0),
635         ExtensionType::LongStrings => (1, 0),
636         ExtensionType::Ncases => (8, 2),
637         ExtensionType::FileAttrs => (1, 0),
638         ExtensionType::VarAttrs => (1, 0),
639         ExtensionType::Mrsets2 => (1, 0),
640         ExtensionType::Encoding => (1, 0),
641         ExtensionType::LongLabels => (1, 0),
642         ExtensionType::LongMissing => (1, 0),
643
644         /* Ignored record types. */
645         ExtensionType::Date => (0, 0),
646         ExtensionType::DataEntry => (0, 0),
647         ExtensionType::Dataview => (0, 0),
648     }
649 }
650
651 fn read_extension_record<R: Read + Seek>(
652     r: &mut BufReader<R>,
653     e: Endian,
654 ) -> Result<Extension, Error> {
655     let subtype = e.parse(read_bytes(r)?);
656     let offset = r.stream_position()?;
657     let size: u32 = e.parse(read_bytes(r)?);
658     let count = e.parse(read_bytes(r)?);
659     let Some(product) = size.checked_mul(count) else {
660         return Err(Error::ExtensionRecordTooLarge {
661             offset,
662             subtype,
663             size,
664             count,
665         });
666     };
667     let offset = r.stream_position()?;
668     let data = read_vec(r, product as usize)?;
669     Ok(Extension {
670         offset,
671         subtype,
672         size,
673         count,
674         data,
675     })
676 }
677
678 struct ZHeader {
679     /// File offset to the start of the record.
680     offset: u64,
681
682     /// File offset to the ZLIB data header.
683     zheader_offset: u64,
684
685     /// File offset to the ZLIB trailer.
686     ztrailer_offset: u64,
687
688     /// Length of the ZLIB trailer in bytes.
689     ztrailer_len: u64,
690 }
691
692 fn read_zheader<R: Read + Seek>(r: &mut BufReader<R>, e: Endian) -> Result<ZHeader, Error> {
693     let offset = r.stream_position()?;
694     let zheader_offset: u64 = e.parse(read_bytes(r)?);
695     let ztrailer_offset: u64 = e.parse(read_bytes(r)?);
696     let ztrailer_len: u64 = e.parse(read_bytes(r)?);
697
698     Ok(ZHeader {
699         offset,
700         zheader_offset,
701         ztrailer_offset,
702         ztrailer_len,
703     })
704 }
705
706 fn read_bytes<const N: usize, R: Read>(r: &mut R) -> Result<[u8; N], IoError> {
707     let mut buf = [0; N];
708     r.read_exact(&mut buf)?;
709     Ok(buf)
710 }
711
712 fn read_vec<R: Read>(r: &mut BufReader<R>, n: usize) -> Result<Vec<u8>, IoError> {
713     let mut vec = vec![0; n];
714     r.read_exact(&mut vec)?;
715     Ok(vec)
716 }
717
718 /*
719 fn trim_end(mut s: Vec<u8>, c: u8) -> Vec<u8> {
720     while s.last() == Some(&c) {
721         s.pop();
722     }
723     s
724 }
725
726 fn skip_bytes<R: Read>(r: &mut R, mut n: u64) -> Result<(), IoError> {
727     let mut buf = [0; 1024];
728     while n > 0 {
729         let chunk = u64::min(n, buf.len() as u64);
730         r.read_exact(&mut buf[0..chunk as usize])?;
731         n -= chunk;
732     }
733     Ok(())
734 }
735
736 */