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

Issue:
ErgoTree v1 headers support



Discussion / hints:

Fix patch:
diff --git a/ergo-lib/CHANGELOG.md b/ergo-lib/CHANGELOG.md
index 486cfedb3..57940cb2c 100644
--- a/ergo-lib/CHANGELOG.md
+++ b/ergo-lib/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added 
 - AND, OR conjecture support in sigma protocol [#226](https://github.com/ergoplatform/sigma-rust/pull/226);
+- Support for ErgoTree header v1 with tree size [#281](https://github.com/ergoplatform/sigma-rust/pull/281);
 
 ## [0.13.1] - 2021-05-28
 
diff --git a/ergotree-interpreter/src/sigma_protocol/fiat_shamir.rs b/ergotree-interpreter/src/sigma_protocol/fiat_shamir.rs
index 15dc8e96b..24b610c53 100644
--- a/ergotree-interpreter/src/sigma_protocol/fiat_shamir.rs
+++ b/ergotree-interpreter/src/sigma_protocol/fiat_shamir.rs
@@ -5,7 +5,7 @@ use super::proof_tree::ProofTreeKind;
 use crate::sigma_protocol::ProverMessage;
 use blake2::digest::{Update, VariableOutput};
 use blake2::VarBlake2b;
-use ergotree_ir::ergo_tree::ErgoTree;
+use ergotree_ir::ergo_tree::{ErgoTree, ErgoTreeHeader};
 use ergotree_ir::mir::expr::Expr;
 use ergotree_ir::serialization::sigma_byte_writer::SigmaByteWrite;
 use ergotree_ir::serialization::sigma_byte_writer::SigmaByteWriter;
@@ -108,8 +108,10 @@ fn fiat_shamir_write_bytes<W: SigmaByteWrite>(
 
     Ok(match tree.as_tree_kind() {
         ProofTreeKind::Leaf(leaf) => {
-            let prop_tree =
-                ErgoTree::with_segregation(&Expr::Const(SigmaProp::new(leaf.proposition()).into()));
+            let prop_tree = ErgoTree::new(
+                ErgoTreeHeader::v0(true),
+                &Expr::Const(SigmaProp::new(leaf.proposition()).into()),
+            );
             let prop_bytes = prop_tree.sigma_serialize_bytes();
             let commitment_bytes = leaf
                 .commitment_opt()
diff --git a/ergotree-ir/src/ergo_tree.rs b/ergotree-ir/src/ergo_tree.rs
index 142e2e652..ddce0fed7 100644
--- a/ergotree-ir/src/ergo_tree.rs
+++ b/ergotree-ir/src/ergo_tree.rs
@@ -10,11 +10,17 @@ use crate::serialization::{
 use crate::sigma_protocol::sigma_boolean::ProveDlog;
 use crate::types::stype::SType;
 use io::Cursor;
+use sigma_ser::vlq_encode::ReadSigmaVlqExt;
+use sigma_ser::vlq_encode::WriteSigmaVlqExt;
 
 use crate::serialization::constant_store::ConstantStore;
+use derive_more::From;
+use derive_more::Into;
 use sigma_ser::peekable_reader::PeekableReader;
 use std::convert::TryFrom;
 use std::io;
+use std::io::Read;
+use std::io::Write;
 use std::rc::Rc;
 use thiserror::Error;
 
@@ -39,28 +45,118 @@ impl ParsedTree {
     fn template_bytes(&self) -> Vec<u8> {
         match self.root.clone().map(|root| root.sigma_serialize_bytes()) {
             Ok(bytes) => bytes,
-            Err(e) => e.bytes, // if tree was failed to parse we already have it's bytes
+            Err(e) => e.root_expr_bytes, // if tree was failed to parse we already have it's bytes
         }
     }
-}
 
-/** The root of ErgoScript IR. Serialized instances of this class are self sufficient and can be passed around.
- */
-#[derive(PartialEq, Eq, Debug, Clone)]
-pub struct ErgoTree {
-    header: ErgoTreeHeader,
-    tree: Result<ParsedTree, ErgoTreeConstantsParsingError>,
+    #[allow(clippy::unwrap_used)] // writer can fail only from OOM, so unwrap is pretty safe here
+    fn sigma_serialize_without_size(&self, header: &ErgoTreeHeader) -> Vec<u8> {
+        let mut data = Vec::new();
+        let mut w = SigmaByteWriter::new(&mut data, None);
+        header.sigma_serialize(&mut w).unwrap();
+        if header.is_constant_segregation() {
+            w.put_usize_as_u32(self.constants.len()).unwrap();
+            self.constants
+                .iter()
+                .try_for_each(|c| c.sigma_serialize(&mut w))
+                .unwrap();
+        };
+        match self.clone().root {
+            Ok(expr) => expr.sigma_serialize(&mut w).unwrap(),
+            Err(ErgoTreeRootParsingError {
+                root_expr_bytes: bytes,
+                ..
+            }) => w.write_all(&bytes).unwrap(),
+        }
+        data
+    }
 }
 
-#[derive(PartialEq, Eq, Debug, Clone)]
-struct ErgoTreeHeader(u8);
+/// Currently we define meaning for only first byte, which may be extended in future versions.
+///  7  6  5  4  3  2  1  0
+///  -------------------------
+///  |  |  |  |  |  |  |  |  |
+///  -------------------------
+///  Bit 7 == 1 if the header contains more than 1 byte (default == 0)
+///  Bit 6 - reserved for GZIP compression (should be 0)
+///  Bit 5 == 1 - reserved for context dependent costing (should be = 0)
+///  Bit 4 == 1 if constant segregation is used for this ErgoTree (default = 0)
+///  (see https://github.com/ScorexFoundation/sigmastate-interpreter/issues/264)
+///  Bit 3 == 1 if size of the whole tree is serialized after the header byte (default = 0)
+///  Bits 2-0 - language version (current version == 0)
+///
+///  Currently we don't specify interpretation for the second and other bytes of the header.
+///  We reserve the possibility to extend header by using Bit 7 == 1 and chain additional bytes as in VLQ.
+///  Once the new bytes are required, a new version of the language should be created and implemented.
+///  That new language will give an interpretation for the new bytes.
+#[derive(PartialEq, Eq, Debug, Clone, From, Into)]
+pub struct ErgoTreeHeader(u8);
+
+/// ErgoTree version 0..=7, should fit in 3 bits
+#[derive(PartialEq, Eq, Debug, Clone, Into)]
+pub struct ErgoTreeVersion(u8);
+
+impl ErgoTreeVersion {
+    /// Header mask to extract version bits.
+    pub const VERSION_MASK: u8 = 0x07;
+    /// Version 0
+    pub const V0: Self = ErgoTreeVersion(0);
+    /// Version 1 (size flag is mandatory)
+    pub const V1: Self = ErgoTreeVersion(1);
+
+    /// Returns a value of the version bits from the given header byte.
+    pub fn parse_version(header: &ErgoTreeHeader) -> ErgoTreeVersion {
+        let header_byte: u8 = header.clone().into();
+        ErgoTreeVersion(header_byte & ErgoTreeVersion::VERSION_MASK)
+    }
+}
 
 impl ErgoTreeHeader {
     const CONSTANT_SEGREGATION_FLAG: u8 = 0x10;
+    const HAS_SIZE_FLAG: u8 = 0x08;
+
+    /// Return a header with version set to 0 and constant segregation flag set to the given value
+    pub fn v0(constant_segregation: bool) -> Self {
+        if constant_segregation {
+            Self::CONSTANT_SEGREGATION_FLAG.into()
+        } else {
+            ErgoTreeHeader::default()
+        }
+    }
+
+    /// Return a header with version set to 1 (with size flag set) and constant segregation flag set to the given value
+    pub fn v1(constant_segregation: bool) -> Self {
+        let version: u8 = ErgoTreeVersion::V1.into();
+        // size flag should be set for version > 0
+        let mut header_byte: u8 = version | Self::HAS_SIZE_FLAG;
+        header_byte = if constant_segregation {
+            header_byte | Self::CONSTANT_SEGREGATION_FLAG
+        } else {
+            header_byte
+        };
+        header_byte.into()
+    }
 
+    /// Returns true if constant segregation flag is set
     pub fn is_constant_segregation(&self) -> bool {
         self.0 & ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG != 0
     }
+
+    /// Returns true if size flag is set
+    pub fn has_size(&self) -> bool {
+        self.0 & ErgoTreeHeader::HAS_SIZE_FLAG != 0
+    }
+
+    /// Returns ErgoTree version
+    pub fn version(&self) -> ErgoTreeVersion {
+        ErgoTreeVersion::parse_version(self)
+    }
+}
+
+impl Default for ErgoTreeHeader {
+    fn default() -> Self {
+        ErgoTreeHeader(ErgoTreeVersion::V0.into())
+    }
 }
 
 /// Whole ErgoTree parsing (deserialization) error
@@ -76,7 +172,7 @@ pub struct ErgoTreeConstantsParsingError {
 #[derive(PartialEq, Eq, Debug, Clone)]
 pub struct ErgoTreeRootParsingError {
     /// Ergo tree root expr bytes (failed to deserialize)
-    pub bytes: Vec<u8>,
+    pub root_expr_bytes: Vec<u8>,
     /// Deserialization error
     pub error: SerializationError,
 }
@@ -92,15 +188,79 @@ pub enum ErgoTreeParsingError {
     RootParsingError(ErgoTreeRootParsingError),
 }
 
-impl ErgoTree {
-    const DEFAULT_HEADER: ErgoTreeHeader = ErgoTreeHeader(0);
+/** The root of ErgoScript IR. Serialized instances of this class are self sufficient and can be passed around.
+ */
+#[derive(PartialEq, Eq, Debug, Clone)]
+pub struct ErgoTree {
+    header: ErgoTreeHeader,
+    tree: Result<ParsedTree, ErgoTreeConstantsParsingError>,
+}
 
-    fn sigma_parse_tree_bytes<R: SigmaByteRead>(
+impl ErgoTree {
+    fn sigma_parse_sized<R: SigmaByteRead>(
         r: &mut R,
+        header: ErgoTreeHeader,
+        size: u32,
+    ) -> Result<Self, SerializationError> {
+        let mut buf = vec![0u8; size as usize];
+        r.read_exact(buf.as_mut_slice())?;
+        if let Ok((constants, mut tree_bytes)) =
+            ErgoTree::sigma_parse_tree_bytes(buf.as_mut_slice(), header.is_constant_segregation())
+        {
+            let tree_bytes_copy = tree_bytes.clone();
+            let mut tree_reader = SigmaByteReader::new(
+                PeekableReader::new(Cursor::new(&mut tree_bytes[..])),
+                ConstantStore::new(constants.clone()),
+            );
+            match Expr::sigma_parse(&mut tree_reader) {
+                Ok(parsed) => Ok(ErgoTree {
+                    header,
+                    tree: Ok(ParsedTree {
+                        constants,
+                        root: Ok(Rc::new(parsed)),
+                    }),
+                }),
+                Err(err) => Ok(ErgoTree {
+                    header,
+                    tree: Ok(ParsedTree {
+                        constants,
+                        root: Err(ErgoTreeRootParsingError {
+                            root_expr_bytes: tree_bytes_copy,
+                            error: err,
+                        }),
+                    }),
+                }),
+            }
+        } else {
+            let mut whole_tree_bytes = Vec::new();
+            let mut w = SigmaByteWriter::new(&mut whole_tree_bytes, None);
+            header.sigma_serialize(&mut w)?;
+            if header.has_size() {
+                w.put_u32(size)?;
+            }
+            w.write_all(&buf)?;
+            Ok(ErgoTree {
+                header,
+                tree: Err(ErgoTreeConstantsParsingError {
+                    bytes: whole_tree_bytes,
+                    error: SerializationError::NotImplementedYet(
+                        "not all constant types serialization is supported".to_string(),
+                    ),
+                }),
+            })
+        }
+    }
+
+    fn sigma_parse_tree_bytes(
+        bytes: &mut [u8],
         is_constant_segregation: bool,
     ) -> Result<(Vec<Constant>, Vec<u8>), SerializationError> {
+        let mut r = SigmaByteReader::new(
+            PeekableReader::new(Cursor::new(&bytes)),
+            ConstantStore::empty(),
+        );
         let constants = if is_constant_segregation {
-            ErgoTree::sigma_parse_constants(r)?
+            ErgoTree::sigma_parse_constants(&mut r)?
         } else {
             vec![]
         };
@@ -126,6 +286,41 @@ impl ErgoTree {
         Ok(constants)
     }
 
+    /// Creates a tree using provided header and root expression
+    pub fn new(header: ErgoTreeHeader, expr: &Expr) -> Self {
+        if header.is_constant_segregation() {
+            let mut data = Vec::new();
+            let cs = ConstantStore::empty();
+            let mut w = SigmaByteWriter::new(&mut data, Some(cs));
+            #[allow(clippy::unwrap_used)]
+            expr.sigma_serialize(&mut w).unwrap();
+            #[allow(clippy::unwrap_used)]
+            let constants = w.constant_store_mut_ref().unwrap().get_all();
+            let cursor = Cursor::new(&mut data[..]);
+            let pr = PeekableReader::new(cursor);
+            let new_cs = ConstantStore::new(constants.clone());
+            let mut sr = SigmaByteReader::new(pr, new_cs);
+            #[allow(clippy::unwrap_used)]
+            // if it was serialized, then we should deserialize it without error
+            let parsed_expr = Expr::sigma_parse(&mut sr).unwrap();
+            ErgoTree {
+                header: ErgoTreeHeader(ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG | header.0),
+                tree: Ok(ParsedTree {
+                    constants,
+                    root: Ok(Rc::new(parsed_expr)),
+                }),
+            }
+        } else {
+            ErgoTree {
+                header,
+                tree: Ok(ParsedTree {
+                    constants: Vec::new(),
+                    root: Ok(Rc::new(expr.clone())),
+                }),
+            }
+        }
+    }
+
     /// Reasonable limit for the number of constants allowed in the ErgoTree
     pub const MAX_CONSTANTS_COUNT: usize = 4096;
 
@@ -160,42 +355,6 @@ impl ErgoTree {
         }
     }
 
-    /// Build ErgoTree using expr as is, without constants segregated
-    pub fn without_segregation(expr: Expr) -> ErgoTree {
-        ErgoTree {
-            header: ErgoTree::DEFAULT_HEADER,
-            tree: Ok(ParsedTree {
-                constants: Vec::new(),
-                root: Ok(Rc::new(expr)),
-            }),
-        }
-    }
-
-    /// Build ErgoTree with constants segregated from expr
-    pub fn with_segregation(expr: &Expr) -> ErgoTree {
-        let mut data = Vec::new();
-        let cs = ConstantStore::empty();
-        let mut w = SigmaByteWriter::new(&mut data, Some(cs));
-        #[allow(clippy::unwrap_used)]
-        expr.sigma_serialize(&mut w).unwrap();
-        #[allow(clippy::unwrap_used)]
-        let constants = w.constant_store_mut_ref().unwrap().get_all();
-        let cursor = Cursor::new(&mut data[..]);
-        let pr = PeekableReader::new(cursor);
-        let new_cs = ConstantStore::new(constants.clone());
-        let mut sr = SigmaByteReader::new(pr, new_cs);
-        #[allow(clippy::unwrap_used)]
-        // if it was serialized, then we should deserialize it without error
-        let parsed_expr = Expr::sigma_parse(&mut sr).unwrap();
-        ErgoTree {
-            header: ErgoTreeHeader(ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG),
-            tree: Ok(ParsedTree {
-                constants,
-                root: Ok(Rc::new(parsed_expr)),
-            }),
-        }
-    }
-
     /// Prints with newlines
     pub fn debug_tree(&self) -> String {
         let tree = format!("{:#?}", self);
@@ -256,11 +415,11 @@ impl From<Expr> for ErgoTree {
         match &expr {
             Expr::Const(c) => match c {
                 Constant { tpe, .. } if *tpe == SType::SSigmaProp => {
-                    ErgoTree::without_segregation(expr)
+                    ErgoTree::new(ErgoTreeHeader::v0(false), &expr)
                 }
-                _ => ErgoTree::with_segregation(&expr),
+                _ => ErgoTree::new(ErgoTreeHeader::v0(true), &expr),
             },
-            _ => ErgoTree::with_segregation(&expr),
+            _ => ErgoTree::new(ErgoTreeHeader::v0(true), &expr),
         }
     }
 }
@@ -277,16 +436,15 @@ impl SigmaSerializable for ErgoTreeHeader {
 
 impl SigmaSerializable for ErgoTree {
     fn sigma_serialize<W: SigmaByteWrite>(&self, w: &mut W) -> Result<(), io::Error> {
-        self.header.sigma_serialize(w)?;
         match &self.tree {
-            Ok(ParsedTree { constants, root }) => {
-                if self.header.is_constant_segregation() {
-                    w.put_usize_as_u32(constants.len())?;
-                    constants.iter().try_for_each(|c| c.sigma_serialize(w))?;
-                }
-                match root {
-                    Ok(expr) => expr.sigma_serialize(w)?,
-                    Err(ErgoTreeRootParsingError { bytes, .. }) => w.write_all(&bytes[..])?,
+            Ok(parsed_tree) => {
+                let bytes = parsed_tree.sigma_serialize_without_size(&self.header);
+                if self.header.has_size() {
+                    self.header.sigma_serialize(w)?;
+                    w.put_usize_as_u32(bytes.len() - 1)?; // skip the header byte
+                    w.write_all(&bytes[1..])?; // skip the header byte
+                } else {
+                    w.write_all(&bytes)?;
                 }
             }
             Err(ErgoTreeConstantsParsingError { bytes, .. }) => w.write_all(&bytes[..])?,
@@ -296,64 +454,37 @@ impl SigmaSerializable for ErgoTree {
 
     fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
         let header = ErgoTreeHeader::sigma_parse(r)?;
-        let constants = if header.is_constant_segregation() {
-            ErgoTree::sigma_parse_constants(r)?
+        if header.has_size() {
+            let tree_size_bytes = r.get_u32()?;
+            ErgoTree::sigma_parse_sized(r, header, tree_size_bytes)
         } else {
-            vec![]
-        };
-        r.set_constant_store(ConstantStore::new(constants.clone()));
-        let root = Expr::sigma_parse(r)?;
-        Ok(ErgoTree {
-            header,
-            tree: Ok(ParsedTree {
-                constants,
-                root: Ok(Rc::new(root)),
-            }),
-        })
+            let constants = if header.is_constant_segregation() {
+                ErgoTree::sigma_parse_constants(r)?
+            } else {
+                vec![]
+            };
+            r.set_constant_store(ConstantStore::new(constants.clone()));
+            let root = Expr::sigma_parse(r)?;
+            Ok(ErgoTree {
+                header,
+                tree: Ok(ParsedTree {
+                    constants,
+                    root: Ok(Rc::new(root)),
+                }),
+            })
+        }
     }
 
     fn sigma_parse_bytes(bytes: &[u8]) -> Result<Self, SerializationError> {
         let cursor = Cursor::new(bytes);
         let mut r = SigmaByteReader::new(PeekableReader::new(cursor), ConstantStore::empty());
         let header = ErgoTreeHeader::sigma_parse(&mut r)?;
-        if let Ok((constants, mut rest_of_the_bytes)) =
-            ErgoTree::sigma_parse_tree_bytes(&mut r, header.is_constant_segregation())
-        {
-            let rest_of_the_bytes_copy = rest_of_the_bytes.clone();
-            let mut new_r = SigmaByteReader::new(
-                PeekableReader::new(Cursor::new(&mut rest_of_the_bytes[..])),
-                ConstantStore::new(constants.clone()),
-            );
-            match Expr::sigma_parse(&mut new_r) {
-                Ok(parsed) => Ok(ErgoTree {
-                    header,
-                    tree: Ok(ParsedTree {
-                        constants,
-                        root: Ok(Rc::new(parsed)),
-                    }),
-                }),
-                Err(err) => Ok(ErgoTree {
-                    header,
-                    tree: Ok(ParsedTree {
-                        constants,
-                        root: Err(ErgoTreeRootParsingError {
-                            bytes: rest_of_the_bytes_copy,
-                            error: err,
-                        }),
-                    }),
-                }),
-            }
+        let rest_of_the_bytes_len = if header.has_size() {
+            r.get_u32()?
         } else {
-            Ok(ErgoTree {
-                header,
-                tree: Err(ErgoTreeConstantsParsingError {
-                    bytes: bytes[1..].to_vec(),
-                    error: SerializationError::NotImplementedYet(
-                        "not all constant types serialization is supported".to_string(),
-                    ),
-                }),
-            })
-        }
+            bytes.len() as u32 - 1 // skip the header byte
+        };
+        ErgoTree::sigma_parse_sized(&mut r, header, rest_of_the_bytes_len)
     }
 }
 
@@ -379,6 +510,8 @@ impl TryFrom<ErgoTree> for ProveDlog {
 #[cfg(feature = "arbitrary")]
 pub(crate) mod arbitrary {
 
+    use crate::mir::expr::arbitrary::ArbExprParams;
+
     use super::*;
     use proptest::prelude::*;
 
@@ -387,9 +520,23 @@ pub(crate) mod arbitrary {
         type Strategy = BoxedStrategy<Self>;
 
         fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+            // make sure that P2PK tree is included
             prop_oneof![
-                // make sure that P2PK tree is included
-                any::<ProveDlog>().prop_map(|p| ErgoTree::from(Expr::Const(p.into()))),
+                any::<ProveDlog>()
+                    .prop_map(|p| ErgoTree::new(ErgoTreeHeader::v0(false), &Expr::Const(p.into()))),
+                any::<ProveDlog>()
+                    .prop_map(|p| ErgoTree::new(ErgoTreeHeader::v1(false), &Expr::Const(p.into()))),
+                // SigmaProp with constant segregation using both v0 and v1 versions
+                any_with::<Expr>(ArbExprParams {
+                    tpe: SType::SSigmaProp,
+                    depth: 1
+                })
+                .prop_map(|e| ErgoTree::new(ErgoTreeHeader::v1(true), &e)),
+                any_with::<Expr>(ArbExprParams {
+                    tpe: SType::SSigmaProp,
+                    depth: 1
+                })
+                .prop_map(|e| ErgoTree::new(ErgoTreeHeader::v0(true), &e)),
             ]
             .boxed()
         }
@@ -411,56 +558,116 @@ mod tests {
 
         #[test]
         fn ser_roundtrip(v in any::<ErgoTree>()) {
+            dbg!(&v);
             let mut data = Vec::new();
             let mut w = SigmaByteWriter::new(&mut data, None);
             v.sigma_serialize(&mut w).expect("serialization failed");
+            // sigma_parse
             let cursor = Cursor::new(&mut data[..]);
             let pr = PeekableReader::new(cursor);
             let mut sr = SigmaByteReader::new(pr, ConstantStore::empty());
             let res = ErgoTree::sigma_parse(&mut sr).expect("parse failed");
             prop_assert_eq!(&res.template_bytes().unwrap(), &v.template_bytes().unwrap());
+            prop_assert_eq![&res, &v];
+            // sigma_parse_bytes
+            let res = ErgoTree::sigma_parse_bytes(&data).expect("parse failed");
+            prop_assert_eq!(&res.template_bytes().unwrap(), &v.template_bytes().unwrap());
             prop_assert_eq![res, v];
         }
     }
 
     #[test]
-    fn deserialization_non_parseable_tree_ok() {
+    fn deserialization_non_parseable_tree_v0() {
         // constants length is set, invalid constant
-        assert!(ErgoTree::sigma_parse_bytes(&[
-            ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG,
-            1,
-            0,
+        let bytes = [
+            ErgoTreeHeader::v0(true).into(),
+            1, // constants quantity
+            0, // invalid constant type
             99,
-            99
-        ])
-        .is_ok());
+            99,
+        ];
+        let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
+        assert!(tree.tree.is_err(), "parsing constants should fail");
+        assert_eq!(
+            tree.sigma_serialize_bytes(),
+            bytes,
+            "serialization should return original bytes"
+        );
+        assert!(
+            tree.template_bytes().is_err(),
+            "template bytes should not be parsed"
+        );
     }
 
     #[test]
-    fn serialization_non_parseable_tree_ok() {
-        // constants length is set, invalid constant
-        let original = &[ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG, 1, 0, 99, 99];
-        let tree = ErgoTree::sigma_parse_bytes(original).unwrap();
-        let bytes = tree.sigma_serialize_bytes();
-        assert_eq!(bytes, original);
-        assert!(tree.template_bytes().is_err());
+    fn deserialization_non_parseable_tree_v1() {
+        // v1(size is set), constants length is set, invalid constant
+        let bytes = [
+            ErgoTreeHeader::v1(true).into(),
+            4, // tree size
+            1, // constants quantity
+            0, // invalid constant type
+            99,
+            99,
+        ];
+        let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
+        assert!(tree.tree.is_err(), "parsing constants should fail");
+        assert_eq!(
+            tree.sigma_serialize_bytes(),
+            bytes,
+            "serialization should return original bytes"
+        );
+        assert!(
+            tree.template_bytes().is_err(),
+            "template bytes should not be parsed"
+        );
     }
 
     #[test]
-    fn deserialization_non_parseable_root_ok() {
+    fn deserialization_non_parseable_root_v0() {
         // no constant segregation, Expr is invalid
-        assert!(ErgoTree::sigma_parse_bytes(&[0, 0, 1]).is_ok());
+        let bytes = [ErgoTreeHeader::v0(false).into(), 0, 1];
+        let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
+        assert!(
+            tree.tree.clone().unwrap().root.is_err(),
+            "parsing root should fail"
+        );
+        assert_eq!(
+            tree.sigma_serialize_bytes(),
+            bytes,
+            "serialization should return original bytes"
+        );
+        assert_eq!(
+            tree.template_bytes().unwrap(),
+            bytes[1..],
+            "template bytes should be parsed"
+        );
     }
 
     #[test]
-    fn serialization_non_parseable_root_ok() {
+    fn deserialization_non_parseable_root_v1() {
         // no constant segregation, Expr is invalid
-        let original = &[0, 0, 1];
-        let tree = ErgoTree::sigma_parse_bytes(original).unwrap();
-        // serialization should return bytes that were failed to parse
-        let bytes = tree.sigma_serialize_bytes();
-        assert_eq!(bytes, original);
-        assert_eq!(tree.template_bytes().unwrap(), [0, 1]); // header byte is skipped
+        let bytes = [
+            ErgoTreeHeader::v1(false).into(),
+            2, // tree size
+            0,
+            1,
+        ];
+        let tree = ErgoTree::sigma_parse_bytes(&bytes).unwrap();
+        assert!(
+            tree.tree.clone().unwrap().root.is_err(),
+            "parsing root should fail"
+        );
+        assert_eq!(
+            tree.sigma_serialize_bytes(),
+            bytes,
+            "serialization should return original bytes"
+        );
+        assert_eq!(
+            tree.template_bytes().unwrap(),
+            bytes[2..],
+            "template bytes should be parsed"
+        );
     }
 
     #[test]
@@ -479,7 +686,7 @@ mod tests {
             tpe: SType::SBoolean,
             v: Value::Boolean(true),
         });
-        let ergo_tree = ErgoTree::with_segregation(&expr);
+        let ergo_tree = ErgoTree::new(ErgoTreeHeader::default(), &expr);
         let bytes = ergo_tree.sigma_serialize_bytes();
         let parsed_expr = ErgoTree::sigma_parse_bytes(&bytes)
             .unwrap()
@@ -494,7 +701,7 @@ mod tests {
             tpe: SType::SBoolean,
             v: Value::Boolean(false),
         });
-        let ergo_tree = ErgoTree::with_segregation(&expr);
+        let ergo_tree = ErgoTree::new(ErgoTreeHeader::v0(true), &expr);
         assert_eq!(ergo_tree.constants_len().unwrap(), 1);
     }
 
@@ -504,7 +711,7 @@ mod tests {
             tpe: SType::SBoolean,
             v: Value::Boolean(false),
         });
-        let ergo_tree = ErgoTree::with_segregation(&expr);
+        let ergo_tree = ErgoTree::new(ErgoTreeHeader::v0(true), &expr);
         assert_eq!(ergo_tree.constants_len().unwrap(), 1);
         assert_eq!(ergo_tree.get_constant(0).unwrap().unwrap(), false.into());
     }
@@ -515,11 +722,27 @@ mod tests {
             tpe: SType::SBoolean,
             v: Value::Boolean(false),
         });
-        let mut ergo_tree = ErgoTree::with_segregation(&expr);
+        let mut ergo_tree = ErgoTree::new(ErgoTreeHeader::v0(true), &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]
+    fn dex_t2tpool_parse() {
+        let base16_str = "19a3030f0400040204020404040404060406058080a0f6f4acdbe01b058080a0f6f4acdbe01b050004d00f0400040005000500d81ad601b2a5730000d602e4c6a70405d603db63087201d604db6308a7d605b27203730100d606b27204730200d607b27203730300d608b27204730400d609b27203730500d60ab27204730600d60b9973078c720602d60c999973088c720502720bd60d8c720802d60e998c720702720dd60f91720e7309d6108c720a02d6117e721006d6127e720e06d613998c7209027210d6147e720d06d615730ad6167e721306d6177e720c06d6187e720b06d6199c72127218d61a9c72167218d1edededededed93c27201c2a793e4c672010405720292c17201c1a793b27203730b00b27204730c00938c7205018c720601ed938c7207018c720801938c7209018c720a019593720c730d95720f929c9c721172127e7202069c7ef07213069a9c72147e7215067e9c720e720206929c9c721472167e7202069c7ef0720e069a9c72117e7215067e9c721372020695ed720f917213730e907217a19d721972149d721a7211ed9272199c7217721492721a9c72177211";
+        let tree_bytes = base16::decode(base16_str.as_bytes()).unwrap();
+        let mut tree = ErgoTree::sigma_parse_bytes(&tree_bytes).unwrap();
+        dbg!(&tree);
+        assert!(tree.header.has_size());
+        assert!(tree.header.is_constant_segregation());
+        assert_eq!(tree.header.version(), ErgoTreeVersion::V1);
+        tree.set_constant(7, 1i64.into()).unwrap();
+        assert_eq!(tree.get_constant(7).unwrap().unwrap(), 1i64.into());
+        tree.set_constant(8, 2i64.into()).unwrap();
+        assert_eq!(tree.get_constant(8).unwrap().unwrap(), 2i64.into());
+        assert!(tree.sigma_serialize_bytes().len() > 1);
+    }
 }
diff --git a/ergotree-ir/src/mir/expr.rs b/ergotree-ir/src/mir/expr.rs
index d4665f4c6..092fd7407 100644
--- a/ergotree-ir/src/mir/expr.rs
+++ b/ergotree-ir/src/mir/expr.rs
@@ -375,7 +375,7 @@ impl<T: TryFrom<Expr>> TryExtractFrom<Expr> for T {
 pub(crate) mod arbitrary {
     use super::*;
     use crate::mir::func_value::FuncArg;
-    use crate::sigma_protocol::dlog_group::EcPoint;
+    use crate::sigma_protocol::sigma_boolean::ProveDlog;
     use crate::types::sfunc::SFunc;
     use proptest::collection::*;
     use proptest::prelude::*;
@@ -429,11 +429,13 @@ pub(crate) mod arbitrary {
             .boxed()
     }
 
-    fn sigma_prop_nester_expr(_depth: usize) -> BoxedStrategy<Expr> {
-        // FIXME: Here we only generate leaf with proof for single key. No connectives yet
-        any::<EcPoint>()
-            .prop_map(|pk| Expr::Const(pk.into()))
-            .boxed()
+    fn sigma_prop_nested_expr(_depth: usize) -> BoxedStrategy<Expr> {
+        prop_oneof![
+            any::<ProveDlog>().prop_map(|pk| Expr::Const(pk.into())),
+            any::<SigmaAnd>().prop_map_into(),
+            any::<SigmaOr>().prop_map_into(),
+        ]
+        .boxed()
     }
 
     fn coll_nested_expr(depth: usize, elem_tpe: &SType) -> BoxedStrategy<Expr> {
@@ -459,7 +461,7 @@ pub(crate) mod arbitrary {
             ]
             .prop_map_into()
             .boxed(),
-            SType::SSigmaProp => sigma_prop_nester_expr(depth),
+            SType::SSigmaProp => sigma_prop_nested_expr(depth),
             _ => panic!("Nested expression not implemented for {:?}", &elem_tpe),
         }
     }
@@ -486,6 +488,7 @@ pub(crate) mod arbitrary {
             SType::SLong => numeric_nested_expr(depth, &tpe),
             SType::SBigInt => numeric_nested_expr(depth, &tpe),
             SType::SColl(elem_type) => coll_nested_expr(depth, elem_type.as_ref()),
+            SType::SSigmaProp => sigma_prop_nested_expr(depth),
             _ => todo!("nested expr is not implemented for type: {:?}", tpe),
         }
         .boxed()
diff --git a/ergotree-ir/src/mir/sigma_and.rs b/ergotree-ir/src/mir/sigma_and.rs
index b463644f7..9503d7fd0 100644
--- a/ergotree-ir/src/mir/sigma_and.rs
+++ b/ergotree-ir/src/mir/sigma_and.rs
@@ -72,7 +72,7 @@ mod arbitrary {
         type Parameters = ();
 
         fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-            vec(any_with::<Constant>(SType::SSigmaProp.into()), 0..5)
+            vec(any_with::<Constant>(SType::SSigmaProp.into()), 2..5)
                 .prop_map(|constants| Self {
                     items: constants.into_iter().map(|c| c.into()).collect(),
                 })
diff --git a/ergotree-ir/src/sigma_protocol/sigma_boolean.rs b/ergotree-ir/src/sigma_protocol/sigma_boolean.rs
index 4714064c9..d4e58c2af 100644
--- a/ergotree-ir/src/sigma_protocol/sigma_boolean.rs
+++ b/ergotree-ir/src/sigma_protocol/sigma_boolean.rs
@@ -7,6 +7,7 @@ use super::dlog_group::EcPoint;
 use crate::ergo_tree::ErgoTree;
 use crate::has_opcode::HasOpCode;
 use crate::mir::constant::Constant;
+use crate::mir::expr::Expr;
 use crate::serialization::op_code::OpCode;
 use crate::serialization::SigmaSerializable;
 use std::convert::TryFrom;
@@ -138,7 +139,8 @@ impl SigmaProp {
         // in order to have comparisons like  `box.propositionBytes == pk.propBytes` we need to make sure
         // the same serialization method is used in both cases
         let c: Constant = self.clone().into();
-        let ergo_tree = ErgoTree::without_segregation(c.into());
+        let e: Expr = c.into();
+        let ergo_tree: ErgoTree = e.into();
         ergo_tree.sigma_serialize_bytes()
     }
 }


Test patch: