Repository: ergoplatform/sigma-rust (language: rust)

Issue:
Expose access to the constants in ErgoTree
- get a constant by index;
- set new constant for given index;


Discussion / hints:

Fix patch:
diff --git a/bindings/ergo-lib-wasm/src/ergo_tree.rs b/bindings/ergo-lib-wasm/src/ergo_tree.rs
index 9c79737a6..6787452a4 100644
--- a/bindings/ergo-lib-wasm/src/ergo_tree.rs
+++ b/bindings/ergo-lib-wasm/src/ergo_tree.rs
@@ -9,6 +9,8 @@ use wasm_bindgen::prelude::*;
 extern crate derive_more;
 use derive_more::{From, Into};
 
+use crate::ast::Constant;
+
 /// The root of ErgoScript IR. Serialized instances of this class are self sufficient and can be passed around.
 #[wasm_bindgen]
 #[derive(PartialEq, Eq, Debug, Clone, From, Into)]
@@ -38,4 +40,36 @@ impl ErgoTree {
     pub fn to_base16_bytes(&self) -> String {
         self.0.to_base16_bytes()
     }
+
+    /// Returns constants number as stored in serialized ErgoTree or error if the parsing of
+    /// constants is failed
+    pub fn constants_len(&self) -> Result<usize, JsValue> {
+        self.0
+            .constants_len()
+            .map_err(|e| JsValue::from_str(&format!("{:?}", e)))
+    }
+
+    /// Returns constant with given index (as stored in serialized ErgoTree)
+    /// or None if index is out of bounds
+    /// or error if constants parsing were failed
+    pub fn get_constant(&self, index: usize) -> Result<Option<Constant>, JsValue> {
+        self.0
+            .get_constant(index)
+            .map(|opt| opt.map(|c| c.into()))
+            .map_err(|e| JsValue::from_str(&format!("{:?}", e)))
+    }
+
+    /// Sets new constant value for a given index in constants list (as stored in serialized ErgoTree),
+    /// and returns previous constant or None if index is out of bounds
+    /// or error if constants parsing were failed
+    pub fn set_constant(
+        &mut self,
+        index: usize,
+        constant: &Constant,
+    ) -> Result<Option<Constant>, JsValue> {
+        self.0
+            .set_constant(index, constant.clone().into())
+            .map(|opt| opt.map(|c| c.into()))
+            .map_err(|e| JsValue::from_str(&format!("{:?}", e)))
+    }
 }
diff --git a/ergo-lib/CHANGELOG.md b/ergo-lib/CHANGELOG.md
index 50862b763..2da68cfb4 100644
--- a/ergo-lib/CHANGELOG.md
+++ b/ergo-lib/CHANGELOG.md
@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 <!-- next-header -->
 ## [Unreleased] - ReleaseDate
+
+### Added 
+- `ErgoTree::constants_len`, `get_constant`, `set_constant` API for accessing the constants list in `ErgoTree` [#261](https://github.com/ergoplatform/sigma-rust/issues/261);
+
 ## [0.10.0] - 2021-04-22
 
 ### Added 
diff --git a/ergotree-ir/src/ergo_tree.rs b/ergotree-ir/src/ergo_tree.rs
index 9d4910a49..7981290e8 100644
--- a/ergotree-ir/src/ergo_tree.rs
+++ b/ergotree-ir/src/ergo_tree.rs
@@ -25,6 +25,19 @@ struct ParsedTree {
     root: Result<Rc<Expr>, ErgoTreeRootParsingError>,
 }
 
+impl ParsedTree {
+    /// Sets new constant value for a given index in constants list (as stored in serialized ErgoTree),
+    /// and returns either previous constant or None if index is out of bounds
+    pub fn set_constant(&mut self, index: usize, constant: Constant) -> Option<Constant> {
+        if index >= self.constants.len() {
+            None
+        } else {
+            let replaced = std::mem::replace(&mut self.constants[index], constant);
+            Some(replaced)
+        }
+    }
+}
+
 /** The root of ErgoScript IR. Serialized instances of this class are self sufficient and can be passed around.
  */
 #[derive(PartialEq, Eq, Debug, Clone)]
@@ -154,6 +167,42 @@ impl ErgoTree {
         let bytes = self.sigma_serialize_bytes();
         base16::encode_lower(&bytes)
     }
+
+    /// Returns constants number as stored in serialized ErgoTree or error if the parsing of
+    /// constants is failed
+    pub fn constants_len(&self) -> Result<usize, ErgoTreeConstantsParsingError> {
+        self.tree
+            .as_ref()
+            .map(|tree| tree.constants.len())
+            .map_err(|e| e.clone())
+    }
+
+    /// Returns constant with given index (as stored in serialized ErgoTree)
+    /// or None if index is out of bounds
+    /// or error if constants parsing were failed
+    pub fn get_constant(
+        &self,
+        index: usize,
+    ) -> Result<Option<Constant>, ErgoTreeConstantsParsingError> {
+        self.tree
+            .as_ref()
+            .map(|tree| tree.constants.get(index).cloned())
+            .map_err(|e| e.clone())
+    }
+
+    /// Sets new constant value for a given index in constants list (as stored in serialized ErgoTree),
+    /// and returns previous constant or None if index is out of bounds
+    /// or error if constants parsing were failed
+    pub fn set_constant(
+        &mut self,
+        index: usize,
+        constant: Constant,
+    ) -> Result<Option<Constant>, ErgoTreeConstantsParsingError> {
+        self.tree
+            .as_mut()
+            .map(|tree| tree.set_constant(index, constant))
+            .map_err(|e| e.clone())
+    }
 }
 
 impl From<Expr> for ErgoTree {
@@ -390,4 +439,39 @@ mod tests {
             .unwrap();
         assert_eq!(*parsed_expr, expr)
     }
+
+    #[test]
+    fn test_constant_len() {
+        let expr = Expr::Const(Constant {
+            tpe: SType::SBoolean,
+            v: Value::Boolean(false),
+        });
+        let ergo_tree = ErgoTree::with_segregation(&expr);
+        assert_eq!(ergo_tree.constants_len().unwrap(), 1);
+    }
+
+    #[test]
+    fn test_get_constant() {
+        let expr = Expr::Const(Constant {
+            tpe: SType::SBoolean,
+            v: Value::Boolean(false),
+        });
+        let ergo_tree = ErgoTree::with_segregation(&expr);
+        assert_eq!(ergo_tree.constants_len().unwrap(), 1);
+        assert_eq!(ergo_tree.get_constant(0).unwrap().unwrap(), false.into());
+    }
+
+    #[test]
+    fn test_set_constant() {
+        let expr = Expr::Const(Constant {
+            tpe: SType::SBoolean,
+            v: Value::Boolean(false),
+        });
+        let mut ergo_tree = ErgoTree::with_segregation(&expr);
+        assert_eq!(
+            ergo_tree.set_constant(0, true.into()).unwrap().unwrap(),
+            false.into()
+        );
+        assert_eq!(ergo_tree.get_constant(0).unwrap().unwrap(), true.into());
+    }
 }


Test patch:
diff --git a/bindings/ergo-lib-wasm/tests/test_ergo_tree.js b/bindings/ergo-lib-wasm/tests/test_ergo_tree.js
new file mode 100644
index 000000000..5b9aee3f0
--- /dev/null
+++ b/bindings/ergo-lib-wasm/tests/test_ergo_tree.js
@@ -0,0 +1,47 @@
+import { expect, assert } from 'chai';
+
+import {
+  ErgoTree,
+  Constant
+} from '../pkg/ergo_lib_wasm';
+
+it('constants_len', async () => {
+  let tree_bytes_base16_str = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
+  let tree = ErgoTree.from_base16_bytes(tree_bytes_base16_str);
+  assert(tree != null);
+  assert(tree.constants_len() == 2);
+});
+
+it('get_constant', async () => {
+  let tree_bytes_base16_str = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
+  let tree = ErgoTree.from_base16_bytes(tree_bytes_base16_str);
+  assert(tree != null);
+  assert(tree.constants_len() == 2);
+  assert(tree.get_constant(0) != null);
+  assert(tree.get_constant(1) != null);
+});
+
+it('get_constant out of bounds', async () => {
+  let tree_bytes_base16_str = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
+  let tree = ErgoTree.from_base16_bytes(tree_bytes_base16_str);
+  assert(tree != null);
+  assert(tree.constants_len() == 2);
+  assert(tree.get_constant(3) == null);
+});
+
+it('set_constant', async () => {
+  let tree_bytes_base16_str = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
+  let tree = ErgoTree.from_base16_bytes(tree_bytes_base16_str);
+  assert(tree.constants_len() == 2);
+  let constant = Constant.from_i32(99);
+  assert(tree.set_constant(0, constant) != null);
+  assert(tree.get_constant(0).to_i32() == 99);
+});
+
+it('set_constant out of bounds', async () => {
+  let tree_bytes_base16_str = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
+  let tree = ErgoTree.from_base16_bytes(tree_bytes_base16_str);
+  assert(tree.constants_len() == 2);
+  let constant = Constant.from_i32(99);
+  assert(tree.set_constant(3, constant) == null);
+});