Functions to create and manage Orchard commitment trees (#99)

This PR contains the functions to create, update and validate Orchard commitment trees.

Reviewed-on: https://git.vergara.tech///Vergara_Tech/zcash-haskell/pulls/99
Co-authored-by: Rene Vergara <rene@vergara.network>
Co-committed-by: Rene Vergara <rene@vergara.network>

librustzcash-wrapper/src/lib.rs

@@ -36,7 +36,8 @@ use incrementalmerkletree::{
     frontier::{
         CommitmentTree,
         Frontier,
-        NonEmptyFrontier
+        NonEmptyFrontier,
+        PathFiller
     },
     witness::IncrementalWitness
 };
@@ -271,6 +272,12 @@ impl<RW> ToHaskell<RW> for Hhex {
     }
 }
 
+impl<RW> FromHaskell<RW> for Hhex {
+    fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
+        let x = Hhex::deserialize(buf)?;
+        Ok(x)
+    }
+}
 
 #[derive(Debug, BorshSerialize, BorshDeserialize)]
 pub struct Haction {
@@ -716,6 +723,47 @@ impl<RW> FromHaskell<RW> for Hfrontier {
     }
 }
 
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct Htree {
+    left: Hhex,
+    right: Hhex,
+    parents: Vec<Hhex>
+}
+
+impl<RW> ToHaskell<RW> for Htree {
+    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
+        self.serialize(writer)?;
+        Ok(())
+    }
+}
+
+impl<RW> FromHaskell<RW> for Htree {
+    fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
+        let x = Htree::deserialize(buf)?;
+        Ok(x)
+    }
+}
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct Hpath {
+    position: u32,
+    path: Vec<Hhex>
+}
+
+impl<RW> ToHaskell<RW> for Hpath {
+    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
+        self.serialize(writer)?;
+        Ok(())
+    }
+}
+
+impl<RW> FromHaskell<RW> for Hpath {
+    fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
+        let x = Hpath::deserialize(buf)?;
+        Ok(x)
+    }
+}
+
 fn to_array<T, const N: usize>(v: Vec<T>) -> [T; N] {
     v.try_into().unwrap_or_else(|v: Vec<T>| panic!("Expected a Vec of length {} but it was {}", N, v.len()))
 }
@@ -1513,7 +1561,13 @@ pub extern "C" fn rust_wrapper_read_orchard_frontier(
                 Some(f2) => {
                     let (pos, leaf, omm) = f2.clone().into_parts();
                     let f = Hfrontier { position: <u64>::from(pos), leaf: Hhex { bytes: leaf.to_bytes().to_vec()}, ommers: omm.iter().map(|&x| Hhex { bytes: x.to_bytes().to_vec()}).collect()};
+                    let comm_tree2: Frontier<MerkleHashOrchard,32> = Frontier::from_parts(Position::from(f.position), MerkleHashOrchard::from_bytes(&to_array(f.leaf.bytes.clone())).unwrap(), f.ommers.iter().map(|x| MerkleHashOrchard::from_bytes(&to_array(x.bytes.clone())).unwrap() ).collect()).unwrap();
+                    if f1.root() == comm_tree2.root() {
                         marshall_to_haskell_var(&f, out, out_len, RW);
+                    } else {
+                        let f0 = Hfrontier { position: 0, leaf: Hhex { bytes: vec![0]}, ommers: vec![Hhex { bytes: vec![0]}]};
+                        marshall_to_haskell_var(&f0, out, out_len, RW);
+                    }
                 },
                 None => {
                     let f0 = Hfrontier { position: 0, leaf: Hhex { bytes: vec![0]}, ommers: vec![Hhex { bytes: vec![0]}]};
@@ -1535,12 +1589,66 @@ pub extern "C" fn rust_wrapper_read_orchard_tree_anchor(
     out: *mut u8,
     out_len: &mut usize
     ){
-    let tree_in: Hfrontier = marshall_from_haskell_var(tree, tree_len, RW);
+    let tree_in: Vec<u8> = marshall_from_haskell_var(tree, tree_len, RW);
-    let leaf = MerkleHashOrchard::from_bytes(&to_array(tree_in.leaf.bytes)).unwrap();
+    let tree_reader = Cursor::new(tree_in);
-    let comm_tree: NonEmptyFrontier<MerkleHashOrchard> = NonEmptyFrontier::from_parts(Position::from(tree_in.position), leaf, tree_in.ommers.iter().map(|x| MerkleHashOrchard::from_bytes(&to_array(x.bytes.clone())).unwrap() ).collect()).unwrap();
+    let comm_tree = read_commitment_tree(tree_reader);
-    let root = comm_tree.root(None);
+    match comm_tree {
+        Ok::<CommitmentTree<MerkleHashOrchard, 32>, _>(c1) => {
+            let root = c1.root();
             let h = Hhex { bytes: root.to_bytes().to_vec() };
             marshall_to_haskell_var(&h, out, out_len, RW);
+        },
+        Err(_) => {
+            let h0 = Hhex { bytes: vec![0] };
+            marshall_to_haskell_var(&h0, out, out_len, RW);
+        }
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn rust_wrapper_read_orchard_commitment_tree_parts(
+    tree: *const u8,
+    tree_len: usize,
+    out: *mut u8,
+    out_len: &mut usize
+    ){
+    let tree_in: Vec<u8> = marshall_from_haskell_var(tree, tree_len, RW);
+    let tree_reader = Cursor::new(tree_in);
+    let comm_tree = read_commitment_tree(tree_reader);
+    match comm_tree {
+        Ok::<CommitmentTree<MerkleHashOrchard, 32>, _>(c1) => {
+            let left = match c1.left() {
+                Some(x) => {
+                    Hhex { bytes: x.to_bytes().to_vec() }
+                },
+                None => {
+                    Hhex { bytes: vec![0] }
+                }
+            };
+            let right = match c1.right() {
+                Some(x) => {
+                    Hhex { bytes: x.to_bytes().to_vec() }
+                },
+                None => {
+                    Hhex { bytes: vec![0] }
+                }
+            };
+            let parents = c1.parents().iter().map(|x| match x {
+                Some(y) => {
+                    Hhex { bytes: y.to_bytes().to_vec() }
+                },
+                None => {
+                    Hhex { bytes: vec![0] }
+                }
+            }).collect();
+            let ht = Htree { left, right, parents};
+            marshall_to_haskell_var(&ht, out, out_len, RW);
+        },
+        Err(_e) => {
+            let ht0 = Htree { left: Hhex { bytes: vec![0] } , right: Hhex { bytes: vec![0] }, parents: vec![Hhex { bytes: vec![0] }]};
+            marshall_to_haskell_var(&ht0, out, out_len, RW);
+        }
+    }
 }
 
 #[no_mangle]
@@ -1577,6 +1685,29 @@ pub extern "C" fn rust_wrapper_read_orchard_node(
     }
 }
 
+#[no_mangle]
+pub extern "C" fn rust_wrapper_read_orchard_path_anchor(
+    path: *const u8,
+    path_len: usize,
+    cmx: *const u8,
+    cmx_len: usize,
+    out: *mut u8,
+    out_len: &mut usize
+    ){
+    let path_in: Hpath = marshall_from_haskell_var(path, path_len, RW);
+    let cmx_in: Vec<u8> = marshall_from_haskell_var(cmx, cmx_len, RW);
+    let mk_path = orchard::tree::MerklePath::from_parts(path_in.position, to_array(path_in.path.iter().map(|x| MerkleHashOrchard::from_bytes(&to_array(x.bytes.clone())).unwrap()).collect()));
+    let nc = ExtractedNoteCommitment::from_bytes(&to_array(cmx_in));
+    if nc.is_some().into() {
+        let anchor = mk_path.root(nc.unwrap());
+        let h = Hhex { bytes: anchor.to_bytes().to_vec() };
+        marshall_to_haskell_var(&h, out, out_len, RW);
+    } else {
+        let h0 = Hhex { bytes: vec![0] };
+        marshall_to_haskell_var(&h0, out, out_len, RW);
+    }
+}
+
 #[no_mangle]
 pub extern "C" fn rust_wrapper_combine_orchard_nodes(
     level: u8,
@@ -1618,6 +1749,47 @@ pub extern "C" fn rust_wrapper_combine_orchard_nodes(
     }
 }
 
+#[no_mangle]
+pub extern "C" fn rust_wrapper_get_orchard_root(
+    level: u8,
+    out: *mut u8,
+    out_len: &mut usize
+    ){
+    let tree: CommitmentTree<MerkleHashOrchard, 32> = CommitmentTree::empty();
+    let root = tree.root();
+    let h = Hhex { bytes: root.to_bytes().to_vec() };
+    marshall_to_haskell_var(&h, out, out_len, RW);
+}
+
+#[no_mangle]
+pub extern "C" fn rust_wrapper_orchard_add_node(
+    level: u8,
+    node: *const u8,
+    node_len: usize,
+    out: *mut u8,
+    out_len: &mut usize
+    ){
+    let mut tree: CommitmentTree<MerkleHashOrchard, 32> = CommitmentTree::empty();
+    let node_in: Vec<u8> = marshall_from_haskell_var(node, node_len, RW);
+    let orchard_note_comm = ExtractedNoteCommitment::from_bytes(&to_array(node_in));
+    if orchard_note_comm.is_some().into() {
+        let n = MerkleHashOrchard::from_cmx(&orchard_note_comm.unwrap());
+        match tree.append(n) {
+            Ok(()) => {
+                let root = tree.root_at_depth(level, PathFiller::empty());
+                let h = Hhex { bytes: root.to_bytes().to_vec() };
+                marshall_to_haskell_var(&h, out, out_len, RW);
+            },
+            Err(_) => {
+                let h0 = Hhex { bytes: vec![0] };
+                marshall_to_haskell_var(&h0, out, out_len, RW);
+            }
+        }
+    } else {
+        let h0 = Hhex { bytes: vec![0] };
+        marshall_to_haskell_var(&h0, out, out_len, RW);
+    }
+}
 
 #[no_mangle]
 pub extern "C" fn rust_wrapper_read_orchard_commitment_tree(

src/C/Zcash.chs

@@ -263,7 +263,7 @@ import ZcashHaskell.Types
 #}
 
 {# fun unsafe rust_wrapper_read_orchard_tree_anchor as rustWrapperReadOrchardTreeAnchor
-     { toBorshVar* `OrchardFrontier'&
+     { toBorshVar* `BS.ByteString'&
      , getVarBuffer `Buffer HexString'&
      }
   -> `()'
@@ -276,6 +276,21 @@ import ZcashHaskell.Types
   -> `()'
 #} 
 
+{# fun unsafe rust_wrapper_read_orchard_path_anchor as rustWrapperReadOrchardPathAnchor
+     { toBorshVar* `MerklePath'&
+     , toBorshVar* `BS.ByteString'&
+     , getVarBuffer `Buffer HexString'&
+     }
+  -> `()'
+#}
+
+{# fun unsafe rust_wrapper_get_orchard_root as rustWrapperGetOrchardRootTest
+     { `Int8'
+     , getVarBuffer `Buffer HexString'&
+     }
+  -> `()'
+#}
+
 {# fun unsafe rust_wrapper_read_orchard_commitment_tree as rustWrapperReadOrchardCommitmentTree
      { toBorshVar* `OrchardFrontier'&
      , toBorshVar* `BS.ByteString'&
@@ -284,6 +299,13 @@ import ZcashHaskell.Types
   -> `()'
 #}
 
+{# fun unsafe rust_wrapper_read_orchard_commitment_tree_parts as rustWrapperReadOrchardTreeParts
+     { toBorshVar* `BS.ByteString'&
+     , getVarBuffer `Buffer OrchardRawTree'&
+     }
+  -> `()'
+#}
+
 {# fun unsafe rust_wrapper_read_orchard_frontier as rustWrapperReadOrchardFrontier
      { toBorshVar* `BS.ByteString'&
      , getVarBuffer `Buffer OrchardFrontier'&
@@ -304,6 +326,15 @@ import ZcashHaskell.Types
   -> `Word64'
 #}
 
+{# fun unsafe rust_wrapper_orchard_add_node as rustWrapperOrchardAddNodeTest
+     { `Int8'
+     , toBorshVar* `BS.ByteString'&
+     , getVarBuffer `Buffer HexString'&
+     }
+  -> `()'
+
+#}
+
 {# fun unsafe rust_wrapper_update_sapling_witness as rustWrapperUpdateSaplingWitness
      { toBorshVar* `BS.ByteString'&
      , toBorshVar* `[BS.ByteString]'&

src/ZcashHaskell/Orchard.hs

@@ -21,14 +21,18 @@ import C.Zcash
   ( rustWrapperCombineOrchardNodes
   , rustWrapperGenOrchardReceiver
   , rustWrapperGenOrchardSpendKey
+  , rustWrapperGetOrchardRootTest
+  , rustWrapperOrchardAddNodeTest
   , rustWrapperOrchardCheck
   , rustWrapperOrchardNoteDecode
   , rustWrapperOrchardNoteDecodeSK
   , rustWrapperReadOrchardCommitmentTree
   , rustWrapperReadOrchardFrontier
   , rustWrapperReadOrchardNode
+  , rustWrapperReadOrchardPathAnchor
   , rustWrapperReadOrchardPosition
   , rustWrapperReadOrchardTreeAnchor
+  , rustWrapperReadOrchardTreeParts
   , rustWrapperReadOrchardWitness
   , rustWrapperReadOrchardWitnessAnchor
   , rustWrapperUADecode
@@ -209,15 +213,47 @@ getOrchardFrontier tree =
       withPureBorshVarBuffer $
       rustWrapperReadOrchardFrontier $ toBytes $ orchTree tree
 
-getOrchardTreeAnchor :: OrchardFrontier -> HexString
+getOrchardTreeAnchor :: OrchardCommitmentTree -> HexString
 getOrchardTreeAnchor tree =
-  withPureBorshVarBuffer $ rustWrapperReadOrchardTreeAnchor tree
+  withPureBorshVarBuffer $
+  rustWrapperReadOrchardTreeAnchor $ toBytes $ orchTree tree
 
 getOrchardWitnessAnchor :: OrchardWitness -> HexString
 getOrchardWitnessAnchor wit =
   withPureBorshVarBuffer $
   rustWrapperReadOrchardWitnessAnchor $ toBytes $ orchWit wit
 
+getOrchardRootTest :: Int -> HexString
+getOrchardRootTest level =
+  withPureBorshVarBuffer $ rustWrapperGetOrchardRootTest $ fromIntegral level
+
+addOrchardNodeGetRoot :: Int -> BS.ByteString -> HexString
+addOrchardNodeGetRoot l n =
+  withPureBorshVarBuffer $ rustWrapperOrchardAddNodeTest (fromIntegral l) n
+
+getOrchardTreeParts :: OrchardCommitmentTree -> Maybe OrchardTree
+getOrchardTreeParts h =
+  if isBlank (ort_left tree) && isBlank (ort_right tree)
+    then Nothing
+    else Just $
+         OrchardTree
+           (parseHex $ ort_left tree)
+           (parseHex $ ort_right tree)
+           (map parseHex (ort_parents tree))
+  where
+    isBlank h = (BS.length $ hexBytes $ h) == 1
+    parseHex h =
+      if (BS.length $ hexBytes $ h) > 1
+        then Just h
+        else Nothing
+    tree =
+      withPureBorshVarBuffer $
+      rustWrapperReadOrchardTreeParts $ toBytes $ orchTree h
+
+getOrchardPathAnchor :: HexString -> MerklePath -> HexString
+getOrchardPathAnchor hex p =
+  withPureBorshVarBuffer $ rustWrapperReadOrchardPathAnchor p (hexBytes hex)
+
 -- | Update a Orchard commitment tree
 updateOrchardCommitmentTree ::
      OrchardFrontier -- ^ the base tree

src/ZcashHaskell/Types.hs

@@ -722,11 +722,34 @@ instance FromJSON OrchardAction where
       a <- obj .: "spendAuthSig"
       pure $ OrchardAction n r c ephKey encText outText cval a
 
+data MerklePath = MerklePath
+  { mp_position :: !Int32
+  , mp_path :: ![HexString]
+  } deriving stock (Eq, Prelude.Show, GHC.Generic)
+    deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo)
+    deriving anyclass (Data.Structured.Show)
+    deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct MerklePath
+
 -- | Type for a Orchard note commitment tree
 newtype OrchardCommitmentTree = OrchardCommitmentTree
   { orchTree :: HexString
   } deriving (Eq, Prelude.Show, Read)
 
+data OrchardRawTree = OrchardRawTree
+  { ort_left :: !HexString
+  , ort_right :: !HexString
+  , ort_parents :: ![HexString]
+  } deriving stock (Eq, Prelude.Show, GHC.Generic)
+    deriving anyclass (SOP.Generic, SOP.HasDatatypeInfo)
+    deriving anyclass (Data.Structured.Show)
+    deriving (BorshSize, ToBorsh, FromBorsh) via AsStruct OrchardRawTree
+
+data OrchardTree = OrchardTree
+  { ot_left :: !(Maybe HexString)
+  , ot_right :: !(Maybe HexString)
+  , ot_parents :: ![Maybe HexString]
+  } deriving (Eq, Prelude.Show, Read)
+
 data OrchardFrontier = OrchardFrontier
   { of_pos :: !Int64
   , of_leaf :: !HexString