diff --git a/ergotree-ir/src/ergo_tree.rs b/ergotree-ir/src/ergo_tree.rs index d7fcaa086..5f8187f70 100644 --- a/ergotree-ir/src/ergo_tree.rs +++ b/ergotree-ir/src/ergo_tree.rs @@ -211,23 +211,11 @@ impl ErgoTree { /// get Expr out of ErgoTree pub fn proposition(&self) -> Result { let tree = self.parsed_tree()?.clone(); + let root = tree.root; // This tree has ConstantPlaceholder nodes instead of Constant nodes. // We need to substitute placeholders with constant values. - // So far the easiest way to do it is during deserialization (after the serialization) - let root = tree.root; if tree.header.is_constant_segregation() { - let mut data = Vec::new(); - let constants = { - let cs = ConstantStore::new(tree.constants.clone()); - let mut w = SigmaByteWriter::new(&mut data, Some(cs)); - root.sigma_serialize(&mut w)?; - #[allow(clippy::unwrap_used)] // constant store is specified in SigmaByteWriter::new - w.constant_store.unwrap() - }; - let cursor = Cursor::new(&mut data[..]); - let mut sr = SigmaByteReader::new_with_substitute_placeholders(cursor, constants); - let parsed_expr = Expr::sigma_parse(&mut sr)?; - Ok(parsed_expr) + Ok(root.substitute_constants(&tree.constants)?) } else { Ok(root) } diff --git a/ergotree-ir/src/mir/expr.rs b/ergotree-ir/src/mir/expr.rs index 868cb2dde..8ee339211 100644 --- a/ergotree-ir/src/mir/expr.rs +++ b/ergotree-ir/src/mir/expr.rs @@ -1,10 +1,12 @@ //! IR expression +use std::convert::Infallible; use std::convert::TryFrom; use std::convert::TryInto; use crate::pretty_printer::PosTrackingWriter; use crate::pretty_printer::Print; +use crate::serialization::SigmaParsingError; use crate::source_span::Spanned; use crate::traversable::Traversable; use crate::types::stype::LiftIntoSType; @@ -326,6 +328,57 @@ impl Expr { } } + /// Rewrite tree, matching nodes that satisfy [`rule`] and applying [`subst`] to them + pub fn rewrite_bu(&self, rule: impl Fn(&Expr) -> bool, subst: impl Fn(&mut Expr)) -> Self { + let mut new_root = self.clone(); + Self::rewrite_bu_inner(&mut new_root, &rule, &|node| { + Ok::<_, Infallible>(subst(node)) + }) + .unwrap(); + new_root + } + + /// Attempt to rewrite tree, returning early if [`subst`] returns `Err(E)` + pub fn try_rewrite_bu( + &self, + rule: impl Fn(&Expr) -> bool, + subst: impl Fn(&mut Expr) -> Result<(), E>, + ) -> Result { + let mut new_root = self.clone(); + Self::rewrite_bu_inner(&mut new_root, &rule, &subst).map(|_| new_root) + } + + fn rewrite_bu_inner( + root: &mut Expr, + rule: &impl Fn(&Expr) -> bool, + subst: &impl Fn(&mut Expr) -> Result<(), E>, + ) -> Result<(), E> { + root.children_mut() + .try_for_each(|expr| Self::rewrite_bu_inner(expr, rule, subst))?; + if rule(root) { + subst(root)?; + } + Ok(()) + } + + /// Substitute [`ConstantPlaceholder`] nodes in `self` with [`Constant`] + /// Errors if a constant can not be found in `constants` + pub fn substitute_constants(&self, constants: &[Constant]) -> Result { + self.try_rewrite_bu::( + |expr| matches!(expr, Expr::ConstPlaceholder(_)), + |expr| { + if let Expr::ConstPlaceholder(ConstantPlaceholder { id, tpe: _ }) = expr { + *expr = constants + .get(*id as usize) + .cloned() + .map(Expr::from) + .ok_or(SigmaParsingError::ConstantForPlaceholderNotFound(*id))?; + } + Ok(()) + }, + ) + } + /// Prints the tree with newlines pub fn debug_tree(&self) -> String { let tree = format!("{:#?}", self);