From 8f03811fcad2e18aa814ac11e457b887ce4e7e53 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Thu, 31 Oct 2024 13:36:28 -0500
Subject: [PATCH] feat: add orchard comm tree parts functionality

---
 librustzcash-wrapper/src/lib.rs | 67 +++++++++++++++++++++++++++++++++
 src/C/Zcash.chs                 |  7 ++++
 src/ZcashHaskell/Orchard.hs     | 20 ++++++++++
 src/ZcashHaskell/Types.hs       | 15 ++++++++
 4 files changed, 109 insertions(+)

diff --git a/librustzcash-wrapper/src/lib.rs b/librustzcash-wrapper/src/lib.rs
index 93b96ee..926aa42 100644
--- a/librustzcash-wrapper/src/lib.rs
+++ b/librustzcash-wrapper/src/lib.rs
@@ -717,6 +717,27 @@ 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)
+    }
+}
+
 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()))
 }
@@ -1550,6 +1571,52 @@ pub extern "C" fn rust_wrapper_read_orchard_tree_anchor(
     marshall_to_haskell_var(&h, 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]
 pub extern "C" fn rust_wrapper_read_orchard_witness_anchor(
     wit: *const u8,
diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index fda259a..6e721de 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -291,6 +291,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'&
diff --git a/src/ZcashHaskell/Orchard.hs b/src/ZcashHaskell/Orchard.hs
index 3fdb25a..6a21eb8 100644
--- a/src/ZcashHaskell/Orchard.hs
+++ b/src/ZcashHaskell/Orchard.hs
@@ -31,6 +31,7 @@ import C.Zcash
   , rustWrapperReadOrchardNode
   , rustWrapperReadOrchardPosition
   , rustWrapperReadOrchardTreeAnchor
+  , rustWrapperReadOrchardTreeParts
   , rustWrapperReadOrchardWitness
   , rustWrapperReadOrchardWitnessAnchor
   , rustWrapperUADecode
@@ -228,6 +229,25 @@ 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
+
 -- | Update a Orchard commitment tree
 updateOrchardCommitmentTree ::
      OrchardFrontier -- ^ the base tree
diff --git a/src/ZcashHaskell/Types.hs b/src/ZcashHaskell/Types.hs
index 324216f..afa6a81 100644
--- a/src/ZcashHaskell/Types.hs
+++ b/src/ZcashHaskell/Types.hs
@@ -727,6 +727,21 @@ 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