Merge f377726 into 0592879

ergo-lib/src/ast.rs

@@ -84,6 +84,12 @@ impl fmt::Display for Expr {
     }
 }
 
+impl From<Constant> for Expr {
+    fn from(c: Constant) -> Self {
+        Self::Const(c)
+    }
+}
+
 #[derive(PartialEq, Eq, Debug, Clone)]
 /// Methods for Collection type instance
 pub enum CollMethods {
@@ -151,3 +157,20 @@ pub enum PredefFunc {
         input: Box<Expr>,
     },
 }
+
+#[cfg(test)]
+mod tests {
+    #![allow(unused_imports)]
+    use super::*;
+    use crate::sigma_protocol::sigma_boolean::SigmaProp;
+    use proptest::prelude::*;
+
+    impl Arbitrary for Expr {
+        type Parameters = ();
+        type Strategy = BoxedStrategy<Self>;
+
+        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+            prop_oneof![any::<Constant>().prop_map(Expr::Const)].boxed()
+        }
+    }
+}

ergo-lib/src/ast/constant.rs

@@ -1,4 +1,7 @@
 use crate::chain::{Base16DecodedBytes, Base16EncodedBytes};
+use crate::sigma_protocol::sigma_boolean::ProveDlog;
+use crate::sigma_protocol::sigma_boolean::SigmaBoolean;
+use crate::sigma_protocol::sigma_boolean::SigmaProofOfKnowledgeTree;
 use crate::{
     chain::ergo_box::ErgoBox,
     serialization::{SerializationError, SigmaSerializable},
@@ -277,7 +280,7 @@ impl From<Vec<i8>> for Constant {
 
 /// Underlying type is different from requested value type
 #[derive(PartialEq, Eq, Debug, Clone)]
-pub struct TryExtractFromError(String);
+pub struct TryExtractFromError(pub String);
 
 /// Extract underlying value if type matches
 pub trait TryExtractFrom<T>: Sized {
@@ -357,6 +360,27 @@ impl TryExtractFrom<ConstantVal> for SigmaProp {
     }
 }
 
+impl TryFrom<ConstantVal> for ProveDlog {
+    type Error = TryExtractFromError;
+    fn try_from(cv: ConstantVal) -> Result<Self, Self::Error> {
+        match cv {
+            ConstantVal::SigmaProp(sp) => match sp.value() {
+                SigmaBoolean::ProofOfKnowledge(SigmaProofOfKnowledgeTree::ProveDlog(
+                    prove_dlog,
+                )) => Ok(prove_dlog.clone()),
+                _ => Err(TryExtractFromError(format!(
+                    "expected ProveDlog, found {:?}",
+                    sp
+                ))),
+            },
+            _ => Err(TryExtractFromError(format!(
+                "expected SigmaProp, found {:?}",
+                cv
+            ))),
+        }
+    }
+}
+
 impl<T: TryExtractFrom<ConstantVal>> TryExtractFrom<Constant> for T {
     fn try_extract_from(cv: Constant) -> Result<Self, TryExtractFromError> {
         T::try_extract_from(cv.v)

ergo-lib/src/chain/address.rs

@@ -1,6 +1,7 @@
 //! Address types
 
 use super::digest32;
+use crate::types::SType;
 use crate::{
     ast::{Constant, ConstantVal, Expr},
     ergo_tree::{ErgoTree, ErgoTreeParsingError},
@@ -99,6 +100,26 @@ impl Address {
             .map(Address::P2PK)
     }
 
+    /// Re-create the address from ErgoTree that was built from the address
+    ///
+    /// At some point in the past a user entered an address from which the ErgoTree was built.
+    /// Re-create the address from this ErgoTree.
+    /// `tree` - ErgoTree that was created from an Address
+    pub fn recreate_from_ergo_tree(tree: &ErgoTree) -> Result<Address, AddressError> {
+        match tree.proposition() {
+            Ok(expr) => Ok(match &*expr {
+                Expr::Const(Constant {
+                    tpe: SType::SSigmaProp,
+                    v,
+                }) => ProveDlog::try_from(v.clone())
+                    .map(Address::P2PK)
+                    .unwrap_or_else(|_| Address::P2S(tree.sigma_serialize_bytes())),
+                _ => Address::P2S(tree.sigma_serialize_bytes()),
+            }),
+            Err(_) => Ok(Address::P2S(tree.sigma_serialize_bytes())),
+        }
+    }
+
     /// address type prefix (for encoding)
     pub fn address_type_prefix(&self) -> AddressTypePrefix {
         match self {
@@ -320,6 +341,7 @@ impl AddressEncoder {
 
 #[cfg(test)]
 mod tests {
+    use crate::chain::Base16DecodedBytes;
     use crate::types::SType;
 
     use super::*;
@@ -330,9 +352,17 @@ mod tests {
         type Strategy = BoxedStrategy<Self>;
 
         fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+            let non_parseable_tree = "100204a00b08cd021dde34603426402615658f1d970cfa7c7bd92ac81a8b16eeebff264d59ce4604ea02d192a39a8cc7a70173007301";
             prop_oneof![
-                any::<ProveDlog>().prop_map(Address::P2PK),
-                any::<ErgoTree>().prop_map(|t| Address::P2S(t.sigma_serialize_bytes())),
+                any::<ErgoTree>().prop_map(|t| match ProveDlog::try_from(t.clone()) {
+                    Ok(dlog) => Address::P2PK(dlog),
+                    Err(_) => Address::P2S(t.sigma_serialize_bytes()),
+                }),
+                Just(Address::P2S(
+                    Base16DecodedBytes::try_from(non_parseable_tree)
+                        .unwrap()
+                        .into()
+                ))
             ]
             .boxed()
         }
@@ -368,6 +398,13 @@ mod tests {
             prop_assert_eq![decoded_addr, v];
         }
 
+        #[test]
+        fn recreate_roundtrip(v in any::<Address>()) {
+            let tree = v.script().unwrap();
+            let recreated = Address::recreate_from_ergo_tree(&tree).unwrap();
+            prop_assert_eq![recreated, v];
+        }
+
         #[test]
         fn doesnt_crash_on_invalid_input(s in "\\w+") {
             let encoder = AddressEncoder::new(NetworkPrefix::Testnet);

ergo-lib/src/chain/base16_bytes.rs

@@ -35,3 +35,16 @@ impl TryFrom<String> for Base16DecodedBytes {
         Ok(Base16DecodedBytes(base16::decode(&str)?))
     }
 }
+
+impl TryFrom<&str> for Base16DecodedBytes {
+    type Error = base16::DecodeError;
+    fn try_from(v: &str) -> Result<Self, Self::Error> {
+        Base16DecodedBytes::try_from(v.to_string())
+    }
+}
+
+impl From<Base16DecodedBytes> for Vec<u8> {
+    fn from(b: Base16DecodedBytes) -> Self {
+        b.0
+    }
+}

ergo-lib/src/ergo_tree.rs

@@ -1,9 +1,11 @@
 //! ErgoTree
+use crate::ast::TryExtractFromError;
 use crate::serialization::{
     sigma_byte_reader::{SigmaByteRead, SigmaByteReader},
     sigma_byte_writer::{SigmaByteWrite, SigmaByteWriter},
     SerializationError, SigmaSerializable,
 };
+use crate::sigma_protocol::sigma_boolean::ProveDlog;
 use crate::{
     ast::{Constant, Expr},
     types::SType,
@@ -12,6 +14,7 @@ use io::{Cursor, Read};
 
 use crate::serialization::constant_store::ConstantStore;
 use sigma_ser::{peekable_reader::PeekableReader, vlq_encode};
+use std::convert::TryFrom;
 use std::io;
 use std::rc::Rc;
 use thiserror::Error;
@@ -175,15 +178,18 @@ impl SigmaSerializable for ErgoTree {
 
     fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
         let header = ErgoTreeHeader::sigma_parse(r)?;
-        if header.is_constant_segregation() {
+        let constants = if header.is_constant_segregation() {
             let constants_len = r.get_u32()?;
-            if constants_len != 0 {
-                return Err(SerializationError::NotImplementedYet(
-                    "separate constants serialization is not yet supported".to_string(),
-                ));
+            let mut constants = Vec::with_capacity(constants_len as usize);
+            for _ in 0..constants_len {
+                let c = Constant::sigma_parse(r)?;
+                constants.push(c);
             }
-        }
-        let constants = Vec::new();
+            constants
+        } else {
+            vec![]
+        };
+        r.set_constant_store(ConstantStore::new(constants.clone()));
         let root = Expr::sigma_parse(r)?;
         Ok(ErgoTree {
             header,
@@ -251,26 +257,47 @@ impl SigmaSerializable for ErgoTree {
     }
 }
 
+impl TryFrom<ErgoTree> for ProveDlog {
+    type Error = TryExtractFromError;
+
+    fn try_from(tree: ErgoTree) -> Result<Self, Self::Error> {
+        let expr = &*tree
+            .proposition()
+            .map_err(|_| TryExtractFromError("cannot read root expr".to_string()))?;
+        match expr {
+            Expr::Const(Constant {
+                tpe: SType::SSigmaProp,
+                v,
+            }) => ProveDlog::try_from(v.clone()),
+            _ => Err(TryExtractFromError(
+                "expected ProveDlog in the root".to_string(),
+            )),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
+    #![allow(unused_imports)]
     use super::*;
+    use crate::chain::Base16DecodedBytes;
     use crate::serialization::sigma_serialize_roundtrip;
-    use crate::{ast::ConstantVal, chain, sigma_protocol::sigma_boolean::SigmaProp, types::SType};
+    use crate::sigma_protocol::sigma_boolean::SigmaProp;
+    use crate::{ast::ConstantVal, chain, types::SType};
     use proptest::prelude::*;
 
     impl Arbitrary for ErgoTree {
         type Parameters = ();
         type Strategy = BoxedStrategy<Self>;
 
         fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-            (any::<SigmaProp>())
-                .prop_map(|p| {
-                    ErgoTree::from(Rc::new(Expr::Const(Constant {
-                        tpe: SType::SSigmaProp,
-                        v: ConstantVal::SigmaProp(Box::new(p)),
-                    })))
-                })
-                .boxed()
+            prop_oneof![
+                // make sure that P2PK tree is included
+                any::<ProveDlog>().prop_map(|p| ErgoTree::from(Rc::new(Expr::from(
+                    Constant::from(SigmaProp::from(p))
+                )))),
+            ]
+            .boxed()
         }
     }
 

ergo-lib/src/serialization/sigma_byte_reader.rs

@@ -41,6 +41,9 @@ pub trait SigmaByteRead: ReadSigmaVlqExt {
 
     /// Option to substitute ConstantPlaceholder with Constant from the store
     fn substitute_placeholders(&self) -> bool;
+
+    /// Set new constant store
+    fn set_constant_store(&mut self, constant_store: ConstantStore);
 }
 
 impl<R: Peekable> Read for SigmaByteReader<R> {
@@ -63,4 +66,8 @@ impl<R: ReadSigmaVlqExt> SigmaByteRead for SigmaByteReader<R> {
     fn substitute_placeholders(&self) -> bool {
         self.substitute_placeholders
     }
+
+    fn set_constant_store(&mut self, constant_store: ConstantStore) {
+        self.constant_store = constant_store;
+    }
 }