use crate::ast::constant::Constant;
use crate::ast::constant::TryExtractFromError;
use crate::ast::expr::Expr;
use crate::serialization::{
sigma_byte_reader::{SigmaByteRead, SigmaByteReader},
sigma_byte_writer::{SigmaByteWrite, SigmaByteWriter},
SerializationError, SigmaSerializable,
};
use crate::sigma_protocol::sigma_boolean::ProveDlog;
use crate::types::stype::SType;
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;
use vlq_encode::ReadSigmaVlqExt;
#[derive(PartialEq, Eq, Debug, Clone)]
struct ParsedTree {
constants: Vec<Constant>,
root: Result<Rc<Expr>, ErgoTreeRootParsingError>,
}
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct ErgoTree {
header: ErgoTreeHeader,
tree: Result<ParsedTree, ErgoTreeConstantsParsingError>,
}
#[derive(PartialEq, Eq, Debug, Clone)]
struct ErgoTreeHeader(u8);
impl ErgoTreeHeader {
const CONSTANT_SEGREGATION_FLAG: u8 = 0x10;
pub fn is_constant_segregation(&self) -> bool {
self.0 & ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG != 0
}
}
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct ErgoTreeConstantsParsingError {
pub bytes: Vec<u8>,
pub error: SerializationError,
}
#[derive(PartialEq, Eq, Debug, Clone)]
pub struct ErgoTreeRootParsingError {
pub bytes: Vec<u8>,
pub error: SerializationError,
}
#[derive(Error, PartialEq, Eq, Debug, Clone)]
pub enum ErgoTreeParsingError {
#[error("Whole ErgoTree parsing (deserialization) error")]
TreeParsingError(ErgoTreeConstantsParsingError),
#[error("ErgoTree root expr parsing (deserialization) error")]
RootParsingError(ErgoTreeRootParsingError),
}
impl ErgoTree {
const DEFAULT_HEADER: ErgoTreeHeader = ErgoTreeHeader(0);
pub const MAX_CONSTANTS_COUNT: usize = 4096;
pub fn proposition(&self) -> Result<Rc<Expr>, ErgoTreeParsingError> {
let root = self
.tree
.clone()
.map_err(ErgoTreeParsingError::TreeParsingError)
.and_then(|t| t.root.map_err(ErgoTreeParsingError::RootParsingError))?;
if self.header.is_constant_segregation() {
let mut data = Vec::new();
let mut cs = ConstantStore::empty();
let mut w = SigmaByteWriter::new(&mut data, Some(&mut cs));
root.sigma_serialize(&mut w).unwrap();
let cursor = Cursor::new(&mut data[..]);
let pr = PeekableReader::new(cursor);
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();
Ok(Rc::new(parsed_expr))
} else {
Ok(root)
}
}
pub fn without_segregation(expr: Rc<Expr>) -> ErgoTree {
ErgoTree {
header: ErgoTree::DEFAULT_HEADER,
tree: Ok(ParsedTree {
constants: Vec::new(),
root: Ok(expr),
}),
}
}
pub fn with_segregation(expr: Rc<Expr>) -> ErgoTree {
let mut data = Vec::new();
let mut cs = ConstantStore::empty();
let mut w = SigmaByteWriter::new(&mut data, Some(&mut cs));
expr.sigma_serialize(&mut w).unwrap();
let cursor = Cursor::new(&mut data[..]);
let pr = PeekableReader::new(cursor);
let constants = cs.get_all();
let new_cs = ConstantStore::new(constants.clone());
let mut sr = SigmaByteReader::new(pr, new_cs);
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)),
}),
}
}
}
impl From<Rc<Expr>> for ErgoTree {
fn from(expr: Rc<Expr>) -> Self {
match expr.as_ref() {
Expr::Const(c) => match &**c {
Constant { tpe, .. } if *tpe == SType::SSigmaProp => {
ErgoTree::without_segregation(expr)
}
_ => ErgoTree::with_segregation(expr),
},
_ => ErgoTree::with_segregation(expr),
}
}
}
impl SigmaSerializable for ErgoTreeHeader {
fn sigma_serialize<W: SigmaByteWrite>(&self, w: &mut W) -> Result<(), io::Error> {
w.put_u8(self.0)?;
Ok(())
}
fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
let header = r.get_u8()?;
Ok(ErgoTreeHeader(header))
}
}
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[..])?,
}
}
Err(ErgoTreeConstantsParsingError { bytes, .. }) => w.write_all(&bytes[..])?,
}
Ok(())
}
fn sigma_parse<R: SigmaByteRead>(r: &mut R) -> Result<Self, SerializationError> {
let header = ErgoTreeHeader::sigma_parse(r)?;
let constants = if header.is_constant_segregation() {
let constants_len = r.get_u32()?;
if constants_len as usize > ErgoTree::MAX_CONSTANTS_COUNT {
return Err(SerializationError::ValueOutOfBounds(
"too many constants".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);
}
constants
} 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(mut bytes: Vec<u8>) -> Result<Self, SerializationError> {
let cursor = Cursor::new(&mut bytes[..]);
let mut r = SigmaByteReader::new(PeekableReader::new(cursor), ConstantStore::empty());
let header = ErgoTreeHeader::sigma_parse(&mut r)?;
let constants = if header.is_constant_segregation() {
let constants_len = r.get_u32()?;
if constants_len as usize > ErgoTree::MAX_CONSTANTS_COUNT {
return Err(SerializationError::ValueOutOfBounds(
"too many constants".to_string(),
));
}
let mut constants = Vec::with_capacity(constants_len as usize);
for _ in 0..constants_len {
match Constant::sigma_parse(&mut r) {
Ok(c) => constants.push(c),
Err(_) => {
return Ok(ErgoTree {
header,
tree: Err(ErgoTreeConstantsParsingError {
bytes: bytes[1..].to_vec(),
error: SerializationError::NotImplementedYet(
"not all constant types serialization is supported".to_string(),
),
}),
})
}
}
}
constants
} else {
vec![]
};
let mut rest_of_the_bytes = Vec::new();
let _ = r.read_to_end(&mut rest_of_the_bytes);
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,
}),
}),
}),
}
}
}
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(c) => match &**c {
Constant {
tpe: SType::SSigmaProp,
v,
} => ProveDlog::try_from(v.clone()),
_ => Err(TryExtractFromError(
"expected ProveDlog in the root".to_string(),
)),
},
_ => Err(TryExtractFromError(
"expected ProveDlog in the root".to_string(),
)),
}
}
}
#[cfg(test)]
mod tests {
#![allow(unused_imports)]
use super::*;
use crate::ast::value::Value;
use crate::chain;
use crate::chain::Base16DecodedBytes;
use crate::serialization::sigma_serialize_roundtrip;
use crate::sigma_protocol::sigma_boolean::SigmaBoolean;
use crate::sigma_protocol::sigma_boolean::SigmaProp;
use proptest::prelude::*;
impl Arbitrary for ErgoTree {
type Parameters = ();
type Strategy = BoxedStrategy<Self>;
fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
prop_oneof![
any::<ProveDlog>().prop_map(|p| ErgoTree::from(Rc::new(Expr::Const(Box::new(
Constant::from(SigmaProp::from(SigmaBoolean::from(p)))
))))),
]
.boxed()
}
}
proptest! {
#[test]
fn ser_roundtrip(v in any::<ErgoTree>()) {
prop_assert_eq![sigma_serialize_roundtrip(&(v)), v];
}
}
#[test]
fn deserialization_non_parseable_tree_ok() {
assert!(ErgoTree::sigma_parse_bytes(vec![
ErgoTreeHeader::CONSTANT_SEGREGATION_FLAG,
1,
0,
99,
99
])
.is_ok());
}
#[test]
fn deserialization_non_parseable_root_ok() {
assert!(ErgoTree::sigma_parse_bytes(vec![0, 0, 1]).is_ok());
}
#[test]
fn test_constant_segregation_header_flag_support() {
let encoder = chain::address::AddressEncoder::new(chain::address::NetworkPrefix::Mainnet);
let address = encoder
.parse_address_from_str("9hzP24a2q8KLPVCUk7gdMDXYc7vinmGuxmLp5KU7k9UwptgYBYV")
.unwrap();
let contract = chain::contract::Contract::pay_to_address(&address).unwrap();
let bytes = &contract.ergo_tree().sigma_serialize_bytes();
assert_eq!(&bytes[..2], vec![0u8, 8u8].as_slice());
}
#[test]
fn test_constant_segregation() {
let expr = Expr::Const(
Constant {
tpe: SType::SBoolean,
v: Value::Boolean(true),
}
.into(),
);
let ergo_tree = ErgoTree::with_segregation(Rc::new(expr.clone()));
let bytes = ergo_tree.sigma_serialize_bytes();
let parsed_expr = ErgoTree::sigma_parse_bytes(bytes)
.unwrap()
.proposition()
.unwrap();
assert_eq!(*parsed_expr, expr)
}
}