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