Skip to main content

powdr_number/
serialize.rs

1use std::{
2    fs::File,
3    io::{self, BufWriter, Read, Write},
4    path::Path,
5};
6
7use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Compress, Validate};
8use csv::{Reader, Writer};
9use serde::{de::DeserializeOwned, Serialize};
10use serde_with::{DeserializeAs, SerializeAs};
11
12use crate::FieldElement;
13
14#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
15pub enum CsvRenderMode {
16    SignedBase10,
17    UnsignedBase10,
18    #[default]
19    Hex,
20}
21
22const ROW_NAME: &str = "Row";
23
24pub fn write_polys_csv_file<T: FieldElement>(
25    file: impl Write,
26    render_mode: CsvRenderMode,
27    polys: &[(&String, &[T])],
28) {
29    let mut writer = Writer::from_writer(file);
30
31    // Write headers, adding a "Row" column
32    let mut headers = vec![ROW_NAME];
33    headers.extend(polys.iter().map(|(name, _)| {
34        assert!(*name != ROW_NAME);
35        name.as_str()
36    }));
37    writer.write_record(&headers).unwrap();
38
39    let max_len = polys.iter().map(|p| p.1.len()).max().unwrap();
40    for row_index in 0..max_len {
41        let mut row = Vec::new();
42        row.push(format!("{row_index}"));
43        for (_, values) in polys {
44            let value = values
45                .get(row_index)
46                .map(|v| match render_mode {
47                    CsvRenderMode::SignedBase10 => format!("{v}"),
48                    CsvRenderMode::UnsignedBase10 => format!("{}", v.to_integer()),
49                    CsvRenderMode::Hex => format!("0x{:x}", v.to_integer()),
50                })
51                .unwrap_or_default();
52            row.push(value);
53        }
54        writer.write_record(&row).unwrap();
55    }
56
57    writer.flush().unwrap();
58}
59
60pub fn read_polys_csv_file<T: FieldElement>(file: impl Read) -> Vec<(String, Vec<T>)> {
61    let mut reader = Reader::from_reader(file);
62    let headers = reader.headers().unwrap();
63
64    let mut polys = headers
65        .iter()
66        .map(|name| (name.to_string(), Vec::new()))
67        .collect::<Vec<_>>();
68
69    for result in reader.records() {
70        let record = result.unwrap();
71        for (idx, value) in record.iter().enumerate() {
72            // shorter polys/columns end in empty cells
73            if value.trim().is_empty() {
74                continue;
75            }
76            let value = if let Some(value) = value.strip_prefix("0x") {
77                T::from_str_radix(value, 16).unwrap()
78            } else if let Some(value) = value.strip_prefix('-') {
79                -T::from_str(value).unwrap()
80            } else {
81                T::from_str(value).unwrap()
82            };
83            polys[idx].1.push(value);
84        }
85    }
86
87    // Remove "Row" column, which was added by write_polys_csv_file()
88    polys
89        .into_iter()
90        .filter(|(name, _)| name != ROW_NAME)
91        .collect()
92}
93
94pub fn buffered_write_file<R>(
95    path: &Path,
96    do_write: impl FnOnce(&mut BufWriter<File>) -> R,
97) -> Result<R, io::Error> {
98    let mut writer = BufWriter::new(File::create(path)?);
99    let result = do_write(&mut writer);
100    writer.flush()?;
101
102    Ok(result)
103}
104
105pub trait ReadWrite {
106    fn read(file: &mut impl Read) -> Self;
107    fn write(&self, path: &Path) -> Result<(), serde_cbor::Error>;
108}
109
110impl<T: DeserializeOwned + Serialize> ReadWrite for T {
111    fn read(file: &mut impl Read) -> Self {
112        serde_cbor::from_reader(file).unwrap()
113    }
114    fn write(&self, path: &Path) -> Result<(), serde_cbor::Error> {
115        buffered_write_file(path, |writer| serde_cbor::to_writer(writer, &self))??;
116        Ok(())
117    }
118}
119
120// Serde wrappers for serialize/deserialize
121
122pub fn ark_se<S, A: CanonicalSerialize>(a: &A, s: S) -> Result<S::Ok, S::Error>
123where
124    S: serde::Serializer,
125{
126    let mut bytes = vec![];
127    a.serialize_with_mode(&mut bytes, Compress::Yes)
128        .map_err(serde::ser::Error::custom)?;
129    serde_with::Bytes::serialize_as(&bytes, s)
130}
131
132pub fn ark_de<'de, D, A: CanonicalDeserialize>(data: D) -> Result<A, D::Error>
133where
134    D: serde::de::Deserializer<'de>,
135{
136    let s: Vec<u8> = serde_with::Bytes::deserialize_as(data)?;
137    let a = A::deserialize_with_mode(s.as_slice(), Compress::Yes, Validate::Yes);
138    a.map_err(serde::de::Error::custom)
139}
140
141#[cfg(test)]
142mod tests {
143    use crate::Bn254Field;
144    use std::io::Cursor;
145
146    use super::*;
147    use test_log::test;
148
149    fn test_polys() -> Vec<(String, Vec<Bn254Field>)> {
150        vec![
151            ("a".to_string(), (0..16).map(Bn254Field::from).collect()),
152            ("b".to_string(), (-16..0).map(Bn254Field::from).collect()),
153        ]
154    }
155
156    #[test]
157    fn write_read() {
158        let mut buf: Vec<u8> = vec![];
159
160        let polys = test_polys();
161
162        serde_cbor::to_writer(&mut buf, &polys).unwrap();
163        let read_polys: Vec<(String, Vec<Bn254Field>)> = ReadWrite::read(&mut Cursor::new(buf));
164
165        assert_eq!(read_polys, polys);
166    }
167
168    #[test]
169    fn write_read_csv() {
170        let polys = test_polys()
171            .into_iter()
172            .map(|(name, values)| (name.to_string(), values))
173            .collect::<Vec<_>>();
174        let polys_ref = polys
175            .iter()
176            .map(|(name, values)| (name, values.as_slice()))
177            .collect::<Vec<_>>();
178
179        for render_mode in &[
180            CsvRenderMode::SignedBase10,
181            CsvRenderMode::UnsignedBase10,
182            CsvRenderMode::Hex,
183        ] {
184            let mut buf: Vec<u8> = vec![];
185            write_polys_csv_file(&mut buf, *render_mode, &polys_ref);
186            let read_polys = read_polys_csv_file::<Bn254Field>(&mut Cursor::new(buf));
187
188            assert_eq!(read_polys, polys);
189        }
190    }
191}