matlab_mat/
lib.rs

1#[macro_use]
2extern crate nom;
3#[macro_use]
4extern crate enum_primitive_derive;
5extern crate log;
6mod mat_error;
7#[cfg(feature = "ndarray")]
8pub mod ndarray;
9mod parse;
10mod writer;
11
12use std::io::Write;
13
14use bytes::{BufMut, BytesMut};
15use nom::number::Endianness;
16use parse::Header;
17
18use crate::mat_error::MatError;
19
20#[derive(Clone, Debug)]
21pub struct Array {
22    array_flags: parse::ArrayFlags,
23    name: String,
24    size: Vec<usize>,
25    data: NumericData,
26}
27impl Array {
28    pub fn name(&self) -> &str {
29        &self.name
30    }
31    pub fn size(&self) -> &Vec<usize> {
32        &self.size
33    }
34    pub fn ndims(&self) -> usize {
35        self.size.len()
36    }
37    pub fn data(&self) -> &NumericData {
38        &self.data
39    }
40    pub fn get_num_elements(&self) -> u32 {
41        let mut count = 1;
42        for val in self.size.iter() {
43            count *= val;
44        }
45        count as u32
46    }
47    pub fn check_index_bound(&self, index: usize, dim: usize) -> usize {
48        if index >= 0 && index < self.size[dim] {
49            index
50        } else {
51            dim
52        }
53    }
54    pub fn set_int8(&mut self, row: usize, col: usize, value: i8) {
55        let ix0 = self.check_index_bound(row, 0);
56        let ix1 = self.check_index_bound(col, 1);
57        let index = row + col * self.size[0];
58        // println!("index={} x={} y={} value={}", index, ix0, ix1, value);
59        match &mut self.data {
60            NumericData::Int8 { real, .. } => {
61                real[index] = value;
62            }
63            _ => {}
64        }
65    }
66
67    // pub fn set_imaginary_int8(&mut self, value: Vec<i8>) {}
68}
69
70#[derive(Clone, Debug)]
71pub enum NumericData {
72    Int8 {
73        real: Vec<i8>,
74        imag: Option<Vec<i8>>,
75    },
76    UInt8 {
77        real: Vec<u8>,
78        imag: Option<Vec<u8>>,
79    },
80    Int16 {
81        real: Vec<i16>,
82        imag: Option<Vec<i16>>,
83    },
84    UInt16 {
85        real: Vec<u16>,
86        imag: Option<Vec<u16>>,
87    },
88    Int32 {
89        real: Vec<i32>,
90        imag: Option<Vec<i32>>,
91    },
92    UInt32 {
93        real: Vec<u32>,
94        imag: Option<Vec<u32>>,
95    },
96    Int64 {
97        real: Vec<i64>,
98        imag: Option<Vec<i64>>,
99    },
100    UInt64 {
101        real: Vec<u64>,
102        imag: Option<Vec<u64>>,
103    },
104    Single {
105        real: Vec<f32>,
106        imag: Option<Vec<f32>>,
107    },
108    Double {
109        real: Vec<f64>,
110        imag: Option<Vec<f64>>,
111    },
112}
113impl NumericData {
114    fn to_numberic_bytes(&self, endianness: nom::number::Endianness) -> (BytesMut, BytesMut) {
115        let mut real_bytes = BytesMut::new();
116        let mut imag_bytes = BytesMut::new();
117        match self {
118            NumericData::Int8 { real, imag } => {
119                for v in real {
120                    real_bytes.put_i8(v.to_owned());
121                }
122            }
123            NumericData::UInt8 { real, imag } => {
124                for v in real {
125                    real_bytes.put_u8(v.to_owned());
126                }
127            }
128            NumericData::Int16 { real, imag } => {
129                for v in real {
130                    if endianness == nom::number::Endianness::Big {
131                        real_bytes.put_i16(v.to_owned());
132                    } else {
133                        real_bytes.put_i16_le(v.to_owned());
134                    }
135                }
136            }
137            NumericData::UInt16 { real, imag } => {
138                for v in real {
139                    if endianness == nom::number::Endianness::Big {
140                        real_bytes.put_u16(v.to_owned());
141                    } else {
142                        real_bytes.put_u16_le(v.to_owned());
143                    }
144                }
145            }
146            NumericData::Int32 { real, imag } => {
147                for v in real {
148                    if endianness == nom::number::Endianness::Big {
149                        real_bytes.put_i32(v.to_owned());
150                    } else {
151                        real_bytes.put_i32_le(v.to_owned());
152                    }
153                }
154            }
155            NumericData::UInt32 { real, imag } => {
156                for v in real {
157                    if endianness == nom::number::Endianness::Big {
158                        real_bytes.put_u32(v.to_owned());
159                    } else {
160                        real_bytes.put_u32_le(v.to_owned());
161                    }
162                }
163            }
164            NumericData::Int64 { real, imag } => {
165                for v in real {
166                    if endianness == nom::number::Endianness::Big {
167                        real_bytes.put_i64(v.to_owned());
168                    } else {
169                        real_bytes.put_i64_le(v.to_owned());
170                    }
171                }
172            }
173            NumericData::UInt64 { real, imag } => {
174                for v in real {
175                    if endianness == nom::number::Endianness::Big {
176                        real_bytes.put_u64(v.to_owned());
177                    } else {
178                        real_bytes.put_u64_le(v.to_owned());
179                    }
180                }
181            }
182            NumericData::Single { real, imag } => {
183                for v in real {
184                    if endianness == nom::number::Endianness::Big {
185                        real_bytes.put_f32(v.to_owned());
186                    } else {
187                        real_bytes.put_f32_le(v.to_owned());
188                    }
189                }
190            }
191            NumericData::Double { real, imag } => {
192                for v in real {
193                    if endianness == nom::number::Endianness::Big {
194                        real_bytes.put_f64(v.to_owned());
195                    } else {
196                        real_bytes.put_f64_le(v.to_owned());
197                    }
198                }
199            }
200            _ => {}
201        };
202        (real_bytes, imag_bytes)
203    }
204    fn to_numberic_size(&self) -> (usize, usize) {
205        match self {
206            NumericData::Int8 { real, imag } => {
207                (real.len(), if let Some(v) = imag { v.len() } else { 0 })
208            }
209            NumericData::UInt8 { real, imag } => {
210                (real.len(), if let Some(v) = imag { v.len() } else { 0 })
211            }
212            NumericData::Int16 { real, imag } => (
213                2 * real.len(),
214                if let Some(v) = imag { v.len() * 2 } else { 0 },
215            ),
216            NumericData::UInt16 { real, imag } => (
217                2 * real.len(),
218                if let Some(v) = imag { v.len() * 2 } else { 0 },
219            ),
220            NumericData::Int32 { real, imag } => (
221                4 * real.len(),
222                if let Some(v) = imag { v.len() * 4 } else { 0 },
223            ),
224            NumericData::UInt32 { real, imag } => (
225                4 * real.len(),
226                if let Some(v) = imag { v.len() * 4 } else { 0 },
227            ),
228            NumericData::Int64 { real, imag } => (
229                8 * real.len(),
230                if let Some(v) = imag { v.len() * 8 } else { 0 },
231            ),
232            NumericData::UInt64 { real, imag } => (
233                8 * real.len(),
234                if let Some(v) = imag { v.len() * 8 } else { 0 },
235            ),
236            NumericData::Single { real, imag } => (
237                4 * real.len(),
238                if let Some(v) = imag { v.len() * 4 } else { 0 },
239            ),
240            NumericData::Double { real, imag } => (
241                8 * real.len(),
242                if let Some(v) = imag { v.len() * 8 } else { 0 },
243            ),
244            _ => (0usize, 0usize),
245        }
246    }
247    fn to_real_size(&self) -> usize {
248        match self {
249            NumericData::Int8 { real, imag } => real.len(),
250            NumericData::UInt8 { real, imag } => real.len(),
251            NumericData::Int16 { real, imag } => 2 * real.len(),
252            NumericData::UInt16 { real, imag } => 2 * real.len(),
253            NumericData::Int32 { real, imag } => 4 * real.len(),
254            NumericData::UInt32 { real, imag } => 4 * real.len(),
255            NumericData::Int64 { real, imag } => 8 * real.len(),
256            NumericData::UInt64 { real, imag } => 8 * real.len(),
257            NumericData::Single { real, imag } => 4 * real.len(),
258            NumericData::Double { real, imag } => 8 * real.len(),
259            _ => 0,
260        }
261    }
262}
263
264fn try_convert_number_format(
265    target_type: parse::MatlabType,
266    data: parse::NumericData,
267) -> Result<parse::NumericData, MatError> {
268    match target_type {
269        parse::MatlabType::Double => match data {
270            parse::NumericData::UInt8(data) => Ok(parse::NumericData::Double(
271                data.into_iter().map(|x| x as f64).collect(),
272            )),
273            parse::NumericData::Int16(data) => Ok(parse::NumericData::Double(
274                data.into_iter().map(|x| x as f64).collect(),
275            )),
276            parse::NumericData::UInt16(data) => Ok(parse::NumericData::Double(
277                data.into_iter().map(|x| x as f64).collect(),
278            )),
279            parse::NumericData::Int32(data) => Ok(parse::NumericData::Double(
280                data.into_iter().map(|x| x as f64).collect(),
281            )),
282            parse::NumericData::Double(data) => Ok(parse::NumericData::Double(data)),
283            _ => Err(MatError::ConversionError),
284        },
285        parse::MatlabType::Single => match data {
286            parse::NumericData::UInt8(data) => Ok(parse::NumericData::Single(
287                data.into_iter().map(|x| x as f32).collect(),
288            )),
289            parse::NumericData::Int16(data) => Ok(parse::NumericData::Single(
290                data.into_iter().map(|x| x as f32).collect(),
291            )),
292            parse::NumericData::UInt16(data) => Ok(parse::NumericData::Single(
293                data.into_iter().map(|x| x as f32).collect(),
294            )),
295            parse::NumericData::Int32(data) => Ok(parse::NumericData::Single(
296                data.into_iter().map(|x| x as f32).collect(),
297            )),
298            parse::NumericData::Single(data) => Ok(parse::NumericData::Single(data)),
299            _ => Err(MatError::ConversionError),
300        },
301        parse::MatlabType::UInt64 => match data {
302            parse::NumericData::UInt8(data) => Ok(parse::NumericData::UInt64(
303                data.into_iter().map(|x| x as u64).collect(),
304            )),
305            parse::NumericData::Int16(data) => Ok(parse::NumericData::UInt64(
306                data.into_iter().map(|x| x as u64).collect(),
307            )),
308            parse::NumericData::UInt16(data) => Ok(parse::NumericData::UInt64(
309                data.into_iter().map(|x| x as u64).collect(),
310            )),
311            parse::NumericData::Int32(data) => Ok(parse::NumericData::UInt64(
312                data.into_iter().map(|x| x as u64).collect(),
313            )),
314            parse::NumericData::UInt64(data) => Ok(parse::NumericData::UInt64(data)),
315            _ => Err(MatError::ConversionError),
316        },
317        parse::MatlabType::Int64 => match data {
318            parse::NumericData::UInt8(data) => Ok(parse::NumericData::Int64(
319                data.into_iter().map(|x| x as i64).collect(),
320            )),
321            parse::NumericData::Int16(data) => Ok(parse::NumericData::Int64(
322                data.into_iter().map(|x| x as i64).collect(),
323            )),
324            parse::NumericData::UInt16(data) => Ok(parse::NumericData::Int64(
325                data.into_iter().map(|x| x as i64).collect(),
326            )),
327            parse::NumericData::Int32(data) => Ok(parse::NumericData::Int64(
328                data.into_iter().map(|x| x as i64).collect(),
329            )),
330            parse::NumericData::Int64(data) => Ok(parse::NumericData::Int64(data)),
331            _ => Err(MatError::ConversionError),
332        },
333        parse::MatlabType::UInt32 => match data {
334            parse::NumericData::UInt8(data) => Ok(parse::NumericData::UInt32(
335                data.into_iter().map(|x| x as u32).collect(),
336            )),
337            parse::NumericData::Int16(data) => Ok(parse::NumericData::UInt32(
338                data.into_iter().map(|x| x as u32).collect(),
339            )),
340            parse::NumericData::UInt16(data) => Ok(parse::NumericData::UInt32(
341                data.into_iter().map(|x| x as u32).collect(),
342            )),
343            parse::NumericData::UInt32(data) => Ok(parse::NumericData::UInt32(data)),
344            _ => Err(MatError::ConversionError),
345        },
346        parse::MatlabType::Int32 => match data {
347            parse::NumericData::UInt8(data) => Ok(parse::NumericData::Int32(
348                data.into_iter().map(|x| x as i32).collect(),
349            )),
350            parse::NumericData::Int16(data) => Ok(parse::NumericData::Int32(
351                data.into_iter().map(|x| x as i32).collect(),
352            )),
353            parse::NumericData::UInt16(data) => Ok(parse::NumericData::Int32(
354                data.into_iter().map(|x| x as i32).collect(),
355            )),
356            parse::NumericData::Int32(data) => Ok(parse::NumericData::Int32(data)),
357            _ => Err(MatError::ConversionError),
358        },
359        parse::MatlabType::UInt16 => match data {
360            parse::NumericData::UInt8(data) => Ok(parse::NumericData::UInt16(
361                data.into_iter().map(|x| x as u16).collect(),
362            )),
363            parse::NumericData::UInt16(data) => Ok(parse::NumericData::UInt16(data)),
364            _ => Err(MatError::ConversionError),
365        },
366        parse::MatlabType::Int16 => match data {
367            parse::NumericData::UInt8(data) => Ok(parse::NumericData::Int16(
368                data.into_iter().map(|x| x as i16).collect(),
369            )),
370            parse::NumericData::Int16(data) => Ok(parse::NumericData::Int16(data)),
371            _ => Err(MatError::ConversionError),
372        },
373        parse::MatlabType::UInt8 => match data {
374            parse::NumericData::UInt8(data) => Ok(parse::NumericData::UInt8(data)),
375            _ => Err(MatError::ConversionError),
376        },
377        parse::MatlabType::Int8 => match data {
378            parse::NumericData::Int8(data) => Ok(parse::NumericData::Int8(data)),
379            _ => Err(MatError::ConversionError),
380        },
381        _ => Err(MatError::ConversionError),
382    }
383}
384
385impl NumericData {
386    fn try_from(
387        target_type: parse::MatlabType,
388        real: parse::NumericData,
389        imag: Option<parse::NumericData>,
390    ) -> Result<Self, MatError> {
391        let real = try_convert_number_format(target_type, real)?;
392        let imag = match imag {
393            Some(imag) => Some(try_convert_number_format(target_type, imag)?),
394            None => None,
395        };
396        match (real, imag) {
397            (parse::NumericData::Double(real), None) => Ok(NumericData::Double {
398                real: real,
399                imag: None,
400            }),
401            (parse::NumericData::Double(real), Some(parse::NumericData::Double(imag))) => {
402                Ok(NumericData::Double {
403                    real: real,
404                    imag: Some(imag),
405                })
406            }
407            (parse::NumericData::Single(real), None) => Ok(NumericData::Single {
408                real: real,
409                imag: None,
410            }),
411            (parse::NumericData::Single(real), Some(parse::NumericData::Single(imag))) => {
412                Ok(NumericData::Single {
413                    real: real,
414                    imag: Some(imag),
415                })
416            }
417            (parse::NumericData::UInt64(real), None) => Ok(NumericData::UInt64 {
418                real: real,
419                imag: None,
420            }),
421            (parse::NumericData::UInt64(real), Some(parse::NumericData::UInt64(imag))) => {
422                Ok(NumericData::UInt64 {
423                    real: real,
424                    imag: Some(imag),
425                })
426            }
427            (parse::NumericData::Int64(real), None) => Ok(NumericData::Int64 {
428                real: real,
429                imag: None,
430            }),
431            (parse::NumericData::Int64(real), Some(parse::NumericData::Int64(imag))) => {
432                Ok(NumericData::Int64 {
433                    real: real,
434                    imag: Some(imag),
435                })
436            }
437            (parse::NumericData::UInt32(real), None) => Ok(NumericData::UInt32 {
438                real: real,
439                imag: None,
440            }),
441            (parse::NumericData::UInt32(real), Some(parse::NumericData::UInt32(imag))) => {
442                Ok(NumericData::UInt32 {
443                    real: real,
444                    imag: Some(imag),
445                })
446            }
447            (parse::NumericData::Int32(real), None) => Ok(NumericData::Int32 {
448                real: real,
449                imag: None,
450            }),
451            (parse::NumericData::Int32(real), Some(parse::NumericData::Int32(imag))) => {
452                Ok(NumericData::Int32 {
453                    real: real,
454                    imag: Some(imag),
455                })
456            }
457            (parse::NumericData::UInt16(real), None) => Ok(NumericData::UInt16 {
458                real: real,
459                imag: None,
460            }),
461            (parse::NumericData::UInt16(real), Some(parse::NumericData::UInt16(imag))) => {
462                Ok(NumericData::UInt16 {
463                    real: real,
464                    imag: Some(imag),
465                })
466            }
467            (parse::NumericData::Int16(real), None) => Ok(NumericData::Int16 {
468                real: real,
469                imag: None,
470            }),
471            (parse::NumericData::Int16(real), Some(parse::NumericData::Int16(imag))) => {
472                Ok(NumericData::Int16 {
473                    real: real,
474                    imag: Some(imag),
475                })
476            }
477            (parse::NumericData::UInt8(real), None) => Ok(NumericData::UInt8 {
478                real: real,
479                imag: None,
480            }),
481            (parse::NumericData::UInt8(real), Some(parse::NumericData::UInt8(imag))) => {
482                Ok(NumericData::UInt8 {
483                    real: real,
484                    imag: Some(imag),
485                })
486            }
487            (parse::NumericData::Int8(real), None) => Ok(NumericData::Int8 {
488                real: real,
489                imag: None,
490            }),
491            (parse::NumericData::Int8(real), Some(parse::NumericData::Int8(imag))) => {
492                Ok(NumericData::Int8 {
493                    real: real,
494                    imag: Some(imag),
495                })
496            }
497            _ => return Err(MatError::InternalError),
498        }
499    }
500}
501
502#[derive(Clone, Debug)]
503pub struct MatFile {
504    header: Header,
505    arrays: Vec<Array>,
506}
507impl MatFile {
508    pub fn add_array(&mut self, array: Array) -> &mut MatFile {
509        self.arrays.push(array);
510        self
511    }
512    pub fn find_by_name<'a>(&'a self, name: &'_ str) -> Option<&'a Array> {
513        for array in &self.arrays {
514            if array.name == name {
515                return Some(array);
516            }
517        }
518        None
519    }
520
521    pub fn new_mat_file() -> MatFile {
522        // 检查CPU大端和小端模式
523        let v: u16 = 0x00ff;
524        let first_octet: u8 = unsafe {
525            let ptr = &v as *const u16;
526            let ptr = ptr as *const u8;
527            *ptr
528        };
529        MatFile {
530            arrays: vec![],
531            header: Header {
532                version: 1,
533                mat_identifier: "MATLAB 5.0 MAT-file".to_string(),
534                description: "".to_string(),
535                byte_order: if first_octet == 0xff {
536                    Endianness::Little
537                } else {
538                    Endianness::Big
539                },
540                subsys_offset: 0,
541                deflate_level: 1,
542            },
543        }
544    }
545    pub fn new_matrix(
546        name: &str,
547        rows: usize,
548        cols: usize,
549        complex: bool,
550        mat_type: parse::MatlabType,
551    ) -> Result<Array, MatError> {
552        let array_flags = parse::ArrayFlags {
553            complex: complex,
554            global: false,
555            logical: false,
556            class: mat_type,
557            nzmax: 0,
558        };
559        let real = if let Some(data) = parse::NumericData::try_from(mat_type, rows, cols) {
560            data
561        } else {
562            return Err(MatError::ParamsError("暂时不支持该数据类型".to_string()));
563        };
564        let imag = if complex {
565            parse::NumericData::try_from(mat_type, rows, cols)
566        } else {
567            None
568        };
569        let data = NumericData::try_from(mat_type, real, imag)?;
570        let array = Array {
571            array_flags: array_flags,
572            name: name.to_string(),
573            size: vec![rows, cols],
574            data: data,
575        };
576        Ok(array)
577    }
578    pub fn parse<R: std::io::Read>(mut read: R) -> Result<Self, MatError> {
579        let mut buf = Vec::new();
580        read.read_to_end(&mut buf)
581            .map_err(|err| MatError::IOError(err))?;
582        let (_remaining, parse_result) = parse::parse_all(&buf)
583            .map_err(|err| MatError::ParseError(parse::replace_err_slice(err, &[])))?;
584        let arrays: Result<Vec<Array>, MatError> = parse_result
585            .data_elements
586            .into_iter()
587            .filter_map(|data_element| match data_element {
588                parse::DataElement::NumericMatrix(flags, dims, name, real, imag) => {
589                    let size = dims.into_iter().map(|d| d as usize).collect();
590                    let numeric_data = match NumericData::try_from(flags.class, real, imag) {
591                        Ok(numeric_data) => numeric_data,
592                        Err(err) => return Some(Err(err)),
593                    };
594                    Some(Ok(Array {
595                        array_flags: flags,
596                        size: size,
597                        name: name,
598                        data: numeric_data,
599                    }))
600                }
601                _ => None,
602            })
603            .collect();
604        let arrays = arrays?;
605        Ok(MatFile {
606            arrays: arrays,
607            header: Header {
608                version: 1,
609                mat_identifier: "".to_string(),
610                description: "".to_string(),
611                byte_order: Endianness::Little,
612                subsys_offset: 0,
613                deflate_level: 1,
614            },
615        })
616    }
617    pub fn save_matfile<T: AsRef<str>>(&self, path: T) -> Result<(), MatError> {
618        let mut file = std::fs::File::create(path.as_ref())?;
619        let header = writer::write_header(self)?;
620        let _r = file.write_all(header.as_ref());
621        let body = writer::write_body(self)?;
622        let _r = file.write_all(body.as_ref());
623        Ok(())
624    }
625}
626
627mod tests {
628
629    #[test]
630    fn write_matfile() -> std::result::Result<(), crate::mat_error::MatError> {
631        let mut new_matfile = super::MatFile::new_mat_file();
632        let rows = 4;
633        let cols = 5;
634        let mut matrix = super::MatFile::new_matrix(
635            "matrixIdentity",
636            rows,
637            cols,
638            false,
639            crate::parse::MatlabType::Int8,
640        )?;
641        let mut count = 1;
642        for i in 0..rows {
643            for j in 0..cols {
644                matrix.set_int8(i, j, count);
645                count += 1;
646            }
647        }
648        println!("matrix==>{:?}", matrix);
649        new_matfile.add_array(matrix);
650        let _r = new_matfile.save_matfile("d:/newmyfile.mat")?;
651        println!("加载生成文件");
652        let file = std::fs::File::open("d:/newmyfile.mat")?;
653        let matfile = super::MatFile::parse(file)?;
654        let array = matfile.find_by_name("matrixIdentity");
655        println!("matrixIdentity={:?}", array);
656        Ok(())
657    }
658    #[test]
659    fn read_matfile() -> std::result::Result<(), crate::mat_error::MatError> {
660        let file = std::fs::File::open("d:/myfile.mat")?;
661        let matfile = super::MatFile::parse(file)?;
662        let array = matfile.find_by_name("matrixIdentity");
663        println!("matrixIdentity={:?}", array);
664        let array = matfile.find_by_name("tout");
665        println!("tout={:?}", array);
666        let array = matfile.find_by_name("x");
667        println!("x={:?}", array);
668        let array = matfile.find_by_name("y");
669        println!("y={:?}", array);
670        Ok(())
671    }
672}
matlab_mat/lib.rs

matlab_mat/
lib.rs
