risc0_zeroio/
serialize.rs

1// Copyright 2023 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use alloc::vec::Vec;
16use core::mem::MaybeUninit;
17
18use impl_trait_for_tuples::impl_for_tuples;
19
20use super::{align_bytes_to_words, as_words_padded, pad_words, Result, ZeroioError};
21
22pub struct Alloc<'a> {
23    buf_left: MaybeUninit<&'a mut [u32]>,
24    base: usize,
25}
26
27pub struct AllocBuf<'a> {
28    buf: &'a mut [u32],
29    base: usize,
30}
31
32impl<'a> Alloc<'a> {
33    pub fn alloc(&mut self, words: usize) -> Result<AllocBuf<'a>> {
34        // Juggle around the slice contianing remaining allocation buffer, since
35        // the borrow checker doesn't understand that we're returning part of the
36        // buffer back to self.buf_left when we're done.
37        let mut tmp: MaybeUninit<&'a mut [u32]> = MaybeUninit::uninit();
38        core::mem::swap(&mut self.buf_left, &mut tmp);
39        let buf_left = unsafe { tmp.assume_init() };
40
41        if buf_left.len() < words {
42            return Err(ZeroioError::AllocationSizeMismatch);
43        }
44        let (new_buf, rest) = buf_left.split_at_mut(words);
45        self.buf_left.write(rest);
46        let old_base = self.base;
47        self.base += words;
48        Ok(AllocBuf {
49            buf: new_buf,
50            base: old_base,
51        })
52    }
53}
54
55impl<'a> AllocBuf<'a> {
56    pub fn descend(&mut self, offset: usize, len: usize) -> Result<AllocBuf> {
57        if offset + len > self.buf.len() {
58            return Err(ZeroioError::FillOverrun);
59        }
60        Ok(AllocBuf {
61            buf: &mut self.buf[offset..offset + len],
62            base: self.base + offset,
63        })
64    }
65
66    pub fn fill_from<const N: usize>(&mut self, val: [u32; N]) -> Result<()> {
67        if self.buf.len() != N {
68            return Err(ZeroioError::FillOverrun);
69        }
70
71        self.buf.clone_from_slice(&val[..]);
72        Ok(())
73    }
74
75    pub fn buf(&mut self, len: usize) -> Result<&mut [u32]> {
76        if self.buf.len() != len {
77            return Err(ZeroioError::FillOverrun);
78        }
79        Ok(&mut self.buf)
80    }
81
82    pub fn rel_ptr_from(&self, other: &AllocBuf) -> u32 {
83        (self.base - other.base) as u32
84    }
85}
86
87pub fn serialize<T>(val: &T) -> Result<Vec<u32>>
88where
89    T: Serialize,
90{
91    let tot_len = val.tot_len();
92    let padded = pad_words(tot_len);
93    let mut buf: Vec<u32> = Vec::new();
94    buf.resize(padded, 0);
95
96    let (mut alloc_buf, _) = buf.as_mut_slice().split_at_mut(tot_len);
97
98    log::trace!(
99        "Starting to serialize type {:?} with {} words and {} fixed words",
100        core::any::type_name::<T>(),
101        tot_len,
102        T::FIXED_WORDS
103    );
104
105    let mut alloc = Alloc {
106        buf_left: MaybeUninit::new(&mut alloc_buf),
107        base: 0,
108    };
109
110    let mut fixed = alloc.alloc(T::FIXED_WORDS)?;
111    val.fill(&mut fixed, &mut alloc)?;
112
113    Ok(buf)
114}
115
116pub trait Serialize {
117    const FIXED_WORDS: usize;
118
119    fn tot_len(&self) -> usize;
120
121    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()>;
122}
123
124impl Serialize for u32 {
125    const FIXED_WORDS: usize = 1;
126
127    fn tot_len(&self) -> usize {
128        1
129    }
130
131    fn fill(&self, buf: &mut AllocBuf, _a: &mut Alloc) -> Result<()> {
132        buf.fill_from([*self])?;
133        Ok(())
134    }
135}
136
137impl Serialize for alloc::string::String {
138    const FIXED_WORDS: usize = 2;
139
140    fn tot_len(&self) -> usize {
141        Self::FIXED_WORDS + align_bytes_to_words(self.len())
142    }
143
144    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
145        let words = align_bytes_to_words(self.len());
146        let str_data = a.alloc(words)?;
147
148        for (w, val) in core::iter::zip(
149            str_data.buf.iter_mut(),
150            as_words_padded(self.as_bytes().into_iter().cloned()),
151        ) {
152            *w = val;
153        }
154
155        buf.fill_from([self.len() as u32, str_data.rel_ptr_from(buf)])?;
156        Ok(())
157    }
158}
159
160impl Serialize for Vec<u8> {
161    const FIXED_WORDS: usize = 2;
162
163    fn tot_len(&self) -> usize {
164        Self::FIXED_WORDS + align_bytes_to_words(self.len())
165    }
166
167    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
168        let words = align_bytes_to_words(self.len());
169        let str_data = a.alloc(words)?;
170
171        for (w, val) in core::iter::zip(
172            str_data.buf.iter_mut(),
173            as_words_padded(self.iter().cloned()),
174        ) {
175            *w = val;
176        }
177
178        buf.fill_from([self.len() as u32, str_data.rel_ptr_from(buf)])?;
179        Ok(())
180    }
181}
182
183impl<T: Serialize> Serialize for core::option::Option<T> {
184    const FIXED_WORDS: usize = 1;
185
186    fn tot_len(&self) -> usize {
187        match self {
188            None => 1,
189            Some(val) => 1 + val.tot_len(),
190        }
191    }
192
193    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
194        match self {
195            None => {
196                buf.fill_from([0])?;
197                Ok(())
198            }
199            Some(val) => {
200                let mut sub_buf = a.alloc(T::FIXED_WORDS)?;
201                val.fill(&mut sub_buf, a)?;
202                buf.fill_from([sub_buf.rel_ptr_from(buf)])?;
203                Ok(())
204            }
205        }
206    }
207}
208
209impl<T: Serialize> Serialize for alloc::boxed::Box<T> {
210    const FIXED_WORDS: usize = T::FIXED_WORDS;
211
212    fn tot_len(&self) -> usize {
213        self.as_ref().tot_len()
214    }
215
216    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
217        self.as_ref().fill(buf, a)
218    }
219}
220
221impl<T: Serialize> Serialize for alloc::vec::Vec<T> {
222    const FIXED_WORDS: usize = 2;
223
224    fn tot_len(&self) -> usize {
225        self.iter().map(|x| x.tot_len()).sum::<usize>() + Self::FIXED_WORDS
226    }
227
228    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
229        let mut sub_buf = a.alloc(T::FIXED_WORDS * self.len())?;
230        let mut pos = 0;
231        for val in self {
232            val.fill(&mut sub_buf.descend(pos, T::FIXED_WORDS)?, a)?;
233            pos += T::FIXED_WORDS;
234        }
235        buf.fill_from([self.len() as u32, sub_buf.rel_ptr_from(buf)])?;
236        Ok(())
237    }
238}
239
240impl<T: Serialize, const N: usize> Serialize for [T; N] {
241    const FIXED_WORDS: usize = T::FIXED_WORDS * N;
242
243    fn tot_len(&self) -> usize {
244        self.iter().map(|x| x.tot_len()).sum()
245    }
246
247    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
248        let mut pos = 0;
249        for val in self {
250            val.fill(&mut buf.descend(pos, T::FIXED_WORDS)?, a)?;
251            pos += T::FIXED_WORDS;
252        }
253        Ok(())
254    }
255}
256
257impl<const N: usize> Serialize for [u8; N] {
258    const FIXED_WORDS: usize = align_bytes_to_words(N);
259
260    fn tot_len(&self) -> usize {
261        Self::FIXED_WORDS
262    }
263
264    fn fill(&self, buf: &mut AllocBuf, _a: &mut Alloc) -> Result<()> {
265        for (w, val) in core::iter::zip(
266            buf.buf(Self::FIXED_WORDS)?.iter_mut(),
267            as_words_padded(self.into_iter().cloned()),
268        ) {
269            *w = val;
270        }
271        Ok(())
272    }
273}
274
275#[impl_for_tuples(1, 5)]
276impl Serialize for Tuple {
277    for_tuples!(const FIXED_WORDS: usize = #(Tuple::FIXED_WORDS)+*; );
278
279    fn tot_len(&self) -> usize {
280        for_tuples!(#(Tuple.tot_len())+*)
281    }
282
283    fn fill(&self, buf: &mut AllocBuf, a: &mut Alloc) -> Result<()> {
284        let mut pos = 0;
285        for_tuples!(
286            #(
287                let fixed = Tuple::FIXED_WORDS;
288                Tuple.fill(&mut buf.descend(pos, fixed)?, a)?;
289                pos += fixed;
290            )*);
291        Ok(())
292    }
293}
294
295impl Serialize for () {
296    const FIXED_WORDS: usize = 0;
297
298    fn tot_len(&self) -> usize {
299        0
300    }
301
302    fn fill(&self, _buf: &mut AllocBuf, _a: &mut Alloc) -> Result<()> {
303        Ok(())
304    }
305}
risc0_zeroio/serialize.rs

risc0_zeroio/
serialize.rs
