wasmbin/builtins/
lazy.rs

1// Copyright 2020 Google Inc. All Rights Reserved.
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 crate::builtins::WasmbinCountable;
16use crate::io::{Decode, DecodeError, DecodeErrorKind, Encode};
17use crate::visit::{Visit, VisitError};
18use custom_debug::Debug as CustomDebug;
19use once_cell::sync::OnceCell;
20use std::hash::Hash;
21
22/// A storage for unparsed bytes.
23///
24/// Unlike `Vec<u8>`, these raw bytes are not length-prefixed when encoded.
25#[derive(Default, CustomDebug, Clone, PartialEq, Eq, Hash, Visit)]
26pub struct UnparsedBytes {
27    #[allow(missing_docs)]
28    #[debug(with = "custom_debug::hexbuf_str")]
29    pub bytes: Vec<u8>,
30}
31
32impl Encode for UnparsedBytes {
33    fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
34        w.write_all(&self.bytes)
35    }
36}
37
38impl Decode for UnparsedBytes {
39    fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
40        let mut res = Self::default();
41        r.read_to_end(&mut res.bytes)?;
42        Ok(res)
43    }
44}
45
46impl AsRef<[u8]> for UnparsedBytes {
47    fn as_ref(&self) -> &[u8] {
48        &self.bytes
49    }
50}
51
52impl std::ops::Deref for UnparsedBytes {
53    type Target = [u8];
54
55    fn deref(&self) -> &Self::Target {
56        &self.bytes
57    }
58}
59
60impl std::ops::DerefMut for UnparsedBytes {
61    fn deref_mut(&mut self) -> &mut Self::Target {
62        &mut self.bytes
63    }
64}
65
66impl From<Vec<u8>> for UnparsedBytes {
67    fn from(bytes: Vec<u8>) -> Self {
68        Self { bytes }
69    }
70}
71
72#[derive(CustomDebug, Clone)]
73enum LazyStatus<T> {
74    FromInput {
75        raw: UnparsedBytes,
76        parsed: OnceCell<T>,
77    },
78    Output {
79        value: T,
80    },
81}
82
83/// A wrapper around a type that allows it to be lazily decoded.
84///
85/// This is useful for types that are expensive to decode, but allowed
86/// to be skipped over by the spec (e.g. as part of a length-prefixed
87/// [`Blob`](crate::builtins::Blob)). During decoding, this type will
88/// store the raw bytes of the value, and only decode them when
89/// explicitly requested.
90///
91/// When re-encoding, it will check if the value has ever been accessed mutably,
92/// and if so, re-encode it. Otherwise it will do a cheap copy of the original
93/// raw bytes.
94#[derive(Clone)]
95pub struct Lazy<T> {
96    status: LazyStatus<T>,
97}
98
99impl<T> Lazy<T> {
100    /// Create a new undecoded `Lazy` from a raw byte vector.
101    pub fn from_raw(raw: UnparsedBytes) -> Self {
102        Lazy {
103            status: LazyStatus::FromInput {
104                raw,
105                parsed: OnceCell::new(),
106            },
107        }
108    }
109
110    /// Retrieve the raw bytes if the value has not been modified yet.
111    pub fn try_as_raw(&self) -> Result<&UnparsedBytes, &T> {
112        match &self.status {
113            LazyStatus::FromInput { raw, .. } => Ok(raw),
114            LazyStatus::Output { value } => Err(value),
115        }
116    }
117}
118
119impl<T> From<T> for Lazy<T> {
120    fn from(value: T) -> Self {
121        Lazy {
122            status: LazyStatus::Output { value },
123        }
124    }
125}
126
127impl<T: Default> Default for Lazy<T> {
128    fn default() -> Self {
129        T::default().into()
130    }
131}
132
133impl<T: Encode> Encode for Lazy<T> {
134    fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
135        match &self.status {
136            LazyStatus::FromInput { raw, .. } => raw.encode(w),
137            LazyStatus::Output { value } => value.encode(w),
138        }
139    }
140}
141
142impl<T: Decode> Decode for Lazy<T> {
143    fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
144        UnparsedBytes::decode(r).map(Self::from_raw)
145    }
146}
147
148impl<T: std::fmt::Debug> std::fmt::Debug for Lazy<T> {
149    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150        self.status.fmt(f)
151    }
152}
153
154fn decode_raw<T: Decode>(raw: &UnparsedBytes) -> Result<T, DecodeError> {
155    let mut raw: &[u8] = raw;
156    let value = T::decode(&mut raw)?;
157    if !raw.is_empty() {
158        return Err(DecodeErrorKind::UnrecognizedData.into());
159    }
160    Ok(value)
161}
162
163impl<T: Decode> Lazy<T> {
164    /// Retrieve a reference to the inner value, decoding it if it wasn't already.
165    pub fn try_contents(&self) -> Result<&T, DecodeError> {
166        match &self.status {
167            LazyStatus::FromInput { raw, parsed } => parsed.get_or_try_init(|| decode_raw(raw)),
168            LazyStatus::Output { value } => Ok(value),
169        }
170    }
171
172    /// Retrieve a mutable reference to the inner value, decoding it if it wasn't already.
173    ///
174    /// This will invalidate the original raw bytes.
175    pub fn try_contents_mut(&mut self) -> Result<&mut T, DecodeError> {
176        if let LazyStatus::FromInput { raw, parsed } = &mut self.status {
177            // We can't trust input and output to match once we obtained a mutable reference,
178            // so get the value and change the status to just Output.
179            let parsed = std::mem::take(parsed);
180            self.status = LazyStatus::Output {
181                value: match parsed.into_inner() {
182                    Some(value) => value,
183                    None => decode_raw(raw)?,
184                },
185            };
186        }
187        if let LazyStatus::Output { value } = &mut self.status {
188            return Ok(value);
189        }
190        unreachable!()
191    }
192
193    /// Unwrap the inner value, decoding it if it wasn't already.
194    pub fn try_into_contents(self) -> Result<T, DecodeError> {
195        match self.status {
196            LazyStatus::FromInput { raw, parsed } => match parsed.into_inner() {
197                Some(value) => Ok(value),
198                None => decode_raw(&raw),
199            },
200            LazyStatus::Output { value } => Ok(value),
201        }
202    }
203}
204
205impl<T: Decode + PartialEq> PartialEq for Lazy<T> {
206    fn eq(&self, other: &Self) -> bool {
207        if let (LazyStatus::FromInput { raw: raw1, .. }, LazyStatus::FromInput { raw: raw2, .. }) =
208            (&self.status, &other.status)
209        {
210            if raw1 == raw2 {
211                return true;
212            }
213        }
214        if let (Ok(value1), Ok(value2)) = (self.try_contents(), other.try_contents()) {
215            return value1 == value2;
216        }
217        false
218    }
219}
220
221impl<T: Decode + Eq> Eq for Lazy<T> {}
222
223impl<T: Decode + Hash> Hash for Lazy<T> {
224    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
225        self.try_contents().ok().hash(state);
226    }
227}
228
229impl<T: WasmbinCountable> WasmbinCountable for Lazy<T> {}
230
231impl<T: Decode + Visit> Visit for Lazy<T> {
232    fn visit_children<'a, VisitT: 'static, E, F: FnMut(&'a VisitT) -> Result<(), E>>(
233        &'a self,
234        f: &mut F,
235    ) -> Result<(), VisitError<E>> {
236        match self.try_contents() {
237            Ok(contents) => contents.visit_child(f),
238            Err(err) => Err(VisitError::LazyDecode(err)),
239        }
240    }
241
242    fn visit_children_mut<VisitT: 'static, E, F: FnMut(&mut VisitT) -> Result<(), E>>(
243        &mut self,
244        f: &mut F,
245    ) -> Result<(), VisitError<E>> {
246        match self.try_contents_mut() {
247            Ok(contents) => contents.visit_child_mut(f),
248            Err(err) => Err(VisitError::LazyDecode(err)),
249        }
250    }
251}
wasmbin/builtins/lazy.rs

wasmbin/builtins/
lazy.rs
