From 91b9c97b1a260d5191b444ae054fc4dbaedf7cfa Mon Sep 17 00:00:00 2001
From: Kamal Ahmad <kamalahmad22@pm.me>
Date: Fri, 25 Oct 2024 12:52:59 +0500
Subject: [PATCH] Lazy has_deserialize for ErgoTree

Use serializer flags to avoid traversing tree to check for existence of deserialize nodes
---
 ergotree-interpreter/src/eval.rs |  6 ++--
 ergotree-ir/src/ergo_tree.rs     | 58 ++++++++++++++++++++++++++++++++
 ergotree-ir/src/mir/expr.rs      |  3 +-
 3 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/ergotree-interpreter/src/eval.rs b/ergotree-interpreter/src/eval.rs
index 748b73fae..d043ebdee 100644
--- a/ergotree-interpreter/src/eval.rs
+++ b/ergotree-interpreter/src/eval.rs
@@ -123,7 +123,7 @@ pub struct ReductionResult {
 }
 
 /// Evaluate the given expression by reducing it to SigmaBoolean value.
-pub fn reduce_to_crypto(expr: &ErgoTree, ctx: &Context) -> Result<ReductionResult, EvalError> {
+pub fn reduce_to_crypto(tree: &ErgoTree, ctx: &Context) -> Result<ReductionResult, EvalError> {
     fn inner<'ctx>(expr: &'ctx Expr, ctx: &Context<'ctx>) -> Result<ReductionResult, EvalError> {
         let mut env_mut = Env::empty();
         expr.eval(&mut env_mut, ctx)
@@ -150,8 +150,8 @@ pub fn reduce_to_crypto(expr: &ErgoTree, ctx: &Context) -> Result<ReductionResul
             })
     }
 
-    let expr = expr.proposition()?;
-    let expr = if expr.has_deserialize() {
+    let expr = tree.proposition()?;
+    let expr = if tree.has_deserialize() {
         expr.substitute_deserialize(ctx)?
     } else {
         expr
diff --git a/ergotree-ir/src/ergo_tree.rs b/ergotree-ir/src/ergo_tree.rs
index 5f8187f70..7cf55a8d1 100644
--- a/ergotree-ir/src/ergo_tree.rs
+++ b/ergotree-ir/src/ergo_tree.rs
@@ -18,6 +18,7 @@ use crate::serialization::constant_store::ConstantStore;
 use derive_more::From;
 use std::convert::TryFrom;
 use std::io;
+use std::sync::OnceLock;
 use thiserror::Error;
 
 mod tree_header;
@@ -29,6 +30,7 @@ pub struct ParsedErgoTree {
     header: ErgoTreeHeader,
     constants: Vec<Constant>,
     root: Expr,
+    has_deserialize: OnceLock<bool>,
 }
 
 impl ParsedErgoTree {
@@ -146,7 +148,11 @@ impl ErgoTree {
             vec![]
         };
         r.set_constant_store(ConstantStore::new(constants.clone()));
+        let was_deserialize = r.was_deserialize();
+        r.set_deserialize(false);
         let root = Expr::sigma_parse(r)?;
+        let has_deserialize = r.was_deserialize();
+        r.set_deserialize(was_deserialize);
         if root.tpe() != SType::SSigmaProp {
             return Err(ErgoTreeError::RootTpeError(root.tpe()));
         }
@@ -154,6 +160,7 @@ impl ErgoTree {
             header,
             constants,
             root,
+            has_deserialize: has_deserialize.into(),
         })
     }
 
@@ -195,12 +202,14 @@ impl ErgoTree {
                 header,
                 constants,
                 root: parsed_expr,
+                has_deserialize: OnceLock::new(),
             })
         } else {
             ErgoTree::Parsed(ParsedErgoTree {
                 header,
                 constants: Vec::new(),
                 root: expr.clone(),
+                has_deserialize: OnceLock::new(),
             })
         })
     }
@@ -221,6 +230,19 @@ impl ErgoTree {
         }
     }
 
+    /// Check if ErgoTree root has [`crate::mir::deserialize_context::DeserializeContext`] or [`crate::mir::deserialize_register::DeserializeRegister`] nodes
+    pub fn has_deserialize(&self) -> bool {
+        match self {
+            ErgoTree::Unparsed { .. } => false,
+            ErgoTree::Parsed(ParsedErgoTree {
+                header: _,
+                constants: _,
+                root,
+                has_deserialize,
+            }) => *has_deserialize.get_or_init(|| root.has_deserialize()),
+        }
+    }
+
     /// Prints with newlines
     pub fn debug_tree(&self) -> String {
         let tree = format!("{:#?}", self);
@@ -372,6 +394,7 @@ impl SigmaSerializable for ErgoTree {
                 header,
                 constants,
                 root,
+                has_deserialize: OnceLock::new(),
             }))
         }
     }
@@ -447,6 +470,7 @@ pub(crate) mod arbitrary {
 
 #[cfg(test)]
 #[cfg(feature = "arbitrary")]
+#[allow(clippy::unreachable)]
 #[allow(clippy::unwrap_used)]
 #[allow(clippy::panic)]
 #[allow(clippy::expect_used)]
@@ -454,7 +478,9 @@ mod tests {
     use super::*;
     use crate::chain::address::AddressEncoder;
     use crate::chain::address::NetworkPrefix;
+    use crate::mir::bool_to_sigma::BoolToSigmaProp;
     use crate::mir::constant::Literal;
+    use crate::mir::deserialize_context::DeserializeContext;
     use crate::sigma_protocol::sigma_boolean::SigmaProp;
     use proptest::prelude::*;
 
@@ -703,4 +729,36 @@ mod tests {
         let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
         tree.proposition().unwrap();
     }
+
+    fn has_deserialize(tree: ErgoTree) -> bool {
+        let ErgoTree::Parsed(ParsedErgoTree {
+            has_deserialize, ..
+        }) = ErgoTree::sigma_parse_bytes(&tree.sigma_serialize_bytes().unwrap()).unwrap()
+        else {
+            unreachable!();
+        };
+        *has_deserialize.get().unwrap()
+    }
+
+    #[test]
+    fn test_lazy_has_deserialize() {
+        let has_deserialize_expr: Expr = BoolToSigmaProp {
+            input: Box::new(
+                DeserializeContext {
+                    tpe: SType::SBoolean,
+                    id: 0,
+                }
+                .into(),
+            ),
+        }
+        .into();
+        let tree = ErgoTree::new(ErgoTreeHeader::v1(false), &has_deserialize_expr).unwrap();
+        assert!(has_deserialize(tree));
+        let no_deserialize_expr: Expr = BoolToSigmaProp {
+            input: Box::new(true.into()),
+        }
+        .into();
+        let tree = ErgoTree::new(ErgoTreeHeader::v1(false), &no_deserialize_expr).unwrap();
+        assert!(!has_deserialize(tree));
+    }
 }
diff --git a/ergotree-ir/src/mir/expr.rs b/ergotree-ir/src/mir/expr.rs
index c98fce05a..6e8a5f62b 100644
--- a/ergotree-ir/src/mir/expr.rs
+++ b/ergotree-ir/src/mir/expr.rs
@@ -387,8 +387,7 @@ impl Expr {
     }
 
     /// Returns true if the [`Expr`] has deserialize nodes, see: [`DeserializeContext`] and [`DeserializeRegister`]
-    pub fn has_deserialize(&self) -> bool {
-        // TODO: if expr is deserialized, use has_deserialize flag from reader to skip traversing tree
+    pub(crate) fn has_deserialize(&self) -> bool {
         self.iter().any(|c| {
             matches!(
                 c,