From 169b963a4efe53f2d16f91d907310dd74c2e7af6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Cauderlier?=
 <raphael.cauderlier@nomadic-labs.com>
Date: Mon, 20 Nov 2023 15:21:36 +0100
Subject: [PATCH] tezos-encoding: add lazy decoding

---
 tezos-encoding/src/nom.rs | 84 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/tezos-encoding/src/nom.rs b/tezos-encoding/src/nom.rs
index fe1d260694..3e50d9e3b1 100644
--- a/tezos-encoding/src/nom.rs
+++ b/tezos-encoding/src/nom.rs
@@ -1,5 +1,6 @@
 // Copyright (c) SimpleStaking, Viable Systems, Nomadic Labs and Tezedge Contributors
 // SPDX-CopyrightText: 2022-2023 TriliTech <contact@trili.tech>
+// SPDX-CopyrightText: 2023 Nomadic Labs <contact@nomadic-labs.com>
 // SPDX-License-Identifier: MIT
 
 use bitvec::slice::BitSlice;
@@ -572,6 +573,56 @@ where
     }
 }
 
+#[derive(Debug, PartialEq)]
+/// Lazy, memoized deserialization of a dynamic block.
+pub struct LazyCell<'a, T>
+where
+    T: NomReader<'a>,
+{
+    // Invariant: when the `deserialized` field is some, it's the
+    // result of deserialization of the `serialized` field using
+    // dynamic(T::nom_read)
+    serialized: &'a [u8],
+    deserialized: Option<T>,
+}
+
+impl<'a, T> NomReader<'a> for LazyCell<'a, T>
+where
+    T: NomReader<'a>,
+{
+    fn nom_read(input: &'a [u8]) -> NomResult<'a, Self> {
+        map(dynamic(rest), |serialized| LazyCell {
+            serialized,
+            deserialized: None,
+        })(input)
+    }
+}
+
+impl<'a, T> LazyCell<'a, T>
+where
+    T: NomReader<'a>,
+{
+    pub fn to_bytes(self) -> &'a [u8] {
+        self.serialized
+    }
+
+    pub fn has_been_forced(&self) -> bool {
+        self.deserialized.is_some()
+    }
+
+    pub fn force<'b>(&'b mut self) -> Result<&'b T, nom::Err<NomError<'a>>> {
+        match self.deserialized.as_mut() {
+            None => {
+                let (_remaining, deserialized) = T::nom_read(self.serialized)?;
+                self.deserialized = Some(deserialized);
+                Ok(())
+            }
+            Some(_) => Ok(()),
+        }?;
+        Ok(self.deserialized.as_ref().unwrap())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use num_bigint::BigInt;
@@ -721,6 +772,39 @@ mod test {
         assert_eq!(err, limit_error(input, BoundedEncodingKind::Dynamic));
     }
 
+    impl<'a> NomReader<'a> for u8 {
+        fn nom_read(input: &'a [u8]) -> NomResult<'a, Self> {
+            u8(input)
+        }
+    }
+
+    impl<'a, T> NomReader<'a> for Vec<T>
+    where
+        T: NomReader<'a>,
+    {
+        fn nom_read(input: &'a [u8]) -> NomResult<'a, Self> {
+            (list(T::nom_read))(input)
+        }
+    }
+
+    #[test]
+    fn test_lazy_dynamic() {
+        let input = &[0, 0, 0, 3, 0x78, 0x78, 0x78, 0xff];
+
+        let (remaining, mut res): (NomInput, LazyCell<Vec<u8>>) =
+            LazyCell::nom_read(input).unwrap();
+        assert_eq!(remaining, &[0xffu8][..]);
+        assert_eq!(
+            res,
+            (LazyCell {
+                serialized: &[0x78; 3],
+                deserialized: None
+            })
+        );
+        let deserialized: Vec<u8> = res.force().unwrap().to_owned();
+        assert_eq!(deserialized, vec![0x78; 3])
+    }
+
     #[test]
     fn test_bounded() {
         let input = &[1, 2, 3, 4, 5];