add an option to substitute constant placeholders on ErgoTree parsing;

ergoplatform · Aug 11, 2020 · 4fe0450 · 4fe0450
1 parent f6e3ec1
commit 4fe0450
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/sigma-tree/src/ergo_tree.rs b/sigma-tree/src/ergo_tree.rs
@@ -85,9 +85,11 @@ impl ErgoTree {
             root.sigma_serialize(&mut w).unwrap();
             let cursor = Cursor::new(&mut data[..]);
             let pr = PeekableReader::new(cursor);
-            // TODO make reader substitute constants
-            let mut sr =
-                SigmaByteReader::new(pr, ConstantStore::new(self.tree.clone().unwrap().constants));
+            // TODO: make reader substitute constants
+            let mut sr = SigmaByteReader::new_with_substitute_placeholders(
+                pr,
+                ConstantStore::new(self.tree.clone().unwrap().constants),
+            );
             let parsed_expr = Expr::sigma_parse(&mut sr).unwrap();
             // todo!("substitute placeholders: {:?}", self.tree);
             Ok(Rc::new(parsed_expr))

diff --git a/sigma-tree/src/serialization/expr.rs b/sigma-tree/src/serialization/expr.rs
@@ -49,7 +49,14 @@ impl SigmaSerializable for Expr {
                 FoldSerializer::OP_CODE => FoldSerializer::sigma_parse(r),
                 ConstantPlaceholder::OP_CODE => {
                     let cp = ConstantPlaceholder::sigma_parse(r)?;
-                    Ok(Expr::ConstPlaceholder(cp))
+                    if r.substitute_placeholders() {
+                        // ConstantPlaceholder itself can be created only if a corresponding
+                        // constant is in the constant_store, thus unwrap() is safe here
+                        let c = r.constant_store().get(cp.id).unwrap();
+                        Ok(Expr::Const(c.clone()))
+                    } else {
+                        Ok(Expr::ConstPlaceholder(cp))
+                    }
                 }
                 o => Err(SerializationError::NotImplementedOpCode(o.value())),
             }

diff --git a/sigma-tree/src/serialization/sigma_byte_reader.rs b/sigma-tree/src/serialization/sigma_byte_reader.rs
@@ -5,6 +5,7 @@ use std::io::Read;
 pub struct SigmaByteReader<R> {
     inner: R,
     constant_store: ConstantStore,
+    substitute_placeholders: bool,
 }
 
 impl<R: Peekable> SigmaByteReader<R> {
@@ -13,12 +14,26 @@ impl<R: Peekable> SigmaByteReader<R> {
         SigmaByteReader {
             inner: pr,
             constant_store,
+            substitute_placeholders: false,
+        }
+    }
+
+    pub fn new_with_substitute_placeholders(
+        pr: R,
+        constant_store: ConstantStore,
+    ) -> SigmaByteReader<R> {
+        SigmaByteReader {
+            inner: pr,
+            constant_store,
+            substitute_placeholders: true,
         }
     }
 }
 
 pub trait SigmaByteRead: ReadSigmaVlqExt {
     fn constant_store(&mut self) -> &mut ConstantStore;
+
+    fn substitute_placeholders(&self) -> bool;
 }
 
 impl<R: Peekable> Read for SigmaByteReader<R> {
@@ -37,4 +52,8 @@ impl<R: ReadSigmaVlqExt> SigmaByteRead for SigmaByteReader<R> {
     fn constant_store(&mut self) -> &mut ConstantStore {
         &mut self.constant_store
     }
+
+    fn substitute_placeholders(&self) -> bool {
+        self.substitute_placeholders
+    }
 }