From a13eab17d7b5765d820f96f81ec177893e79b263 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Wed, 6 Nov 2024 13:17:58 -0600
Subject: [PATCH] feat: add sapling commitment trees

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

diff --git a/librustzcash-wrapper/src/lib.rs b/librustzcash-wrapper/src/lib.rs
index 7a03a49..903c256 100644
--- a/librustzcash-wrapper/src/lib.rs
+++ b/librustzcash-wrapper/src/lib.rs
@@ -1614,6 +1614,52 @@ pub extern "C" fn rust_wrapper_get_sapling_root(
     marshall_to_haskell_var(&h, out, out_len, RW);
 }
 
+#[no_mangle]
+pub extern "C" fn rust_wrapper_read_sapling_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<Node, 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_frontier(
     tree: *const u8,
diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index 8e8947b..0bcdf6d 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -269,6 +269,13 @@ import ZcashHaskell.Types
   -> `()'
 #}
 
+{# fun unsafe rust_wrapper_read_sapling_commitment_tree_parts as rustWrapperReadSaplingTreeParts
+     { toBorshVar* `BS.ByteString'&
+     , getVarBuffer `Buffer SaplingRawTree'&
+     }
+  -> `()'
+#}
+
 {# fun unsafe rust_wrapper_read_orchard_node as rustWrapperReadOrchardNode
      { toBorshVar* `BS.ByteString'&
      , getVarBuffer `Buffer HexString'&
diff --git a/src/ZcashHaskell/Sapling.hs b/src/ZcashHaskell/Sapling.hs
index 362a8e9..f0b743c 100644
--- a/src/ZcashHaskell/Sapling.hs
+++ b/src/ZcashHaskell/Sapling.hs
@@ -26,6 +26,7 @@ import C.Zcash
   , rustWrapperReadSaplingFrontier
   , rustWrapperReadSaplingNode
   , rustWrapperReadSaplingPosition
+  , rustWrapperReadSaplingTreeParts
   , rustWrapperReadSaplingWitness
   , rustWrapperSaplingCheck
   , rustWrapperSaplingChgPaymentAddress
@@ -211,6 +212,25 @@ getSaplingRootTest :: Int8 -> HexString
 getSaplingRootTest level =
   withPureBorshVarBuffer $ rustWrapperGetSaplingRootTest level
 
+getSaplingTreeParts :: SaplingCommitmentTree -> Maybe SaplingTree
+getSaplingTreeParts h =
+  if isBlank (srt_left tree) && isBlank (srt_right tree)
+    then Nothing
+    else Just $
+         SaplingTree
+           (parseHex $ srt_left tree)
+           (parseHex $ srt_right tree)
+           (map parseHex (srt_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 $
+      rustWrapperReadSaplingTreeParts $ toBytes $ sapTree h
+
 getSaplingFrontier :: SaplingCommitmentTree -> Maybe SaplingFrontier
 getSaplingFrontier tree =
   if sf_pos updatedTree > 1
diff --git a/src/ZcashHaskell/Types.hs b/src/ZcashHaskell/Types.hs
index 8b8a90b..0226ea1 100644
--- a/src/ZcashHaskell/Types.hs
+++ b/src/ZcashHaskell/Types.hs
@@ -611,6 +611,21 @@ newtype SaplingCommitmentTree = SaplingCommitmentTree
   { sapTree :: HexString
   } deriving (Eq, Prelude.Show, Read)
 
+data SaplingRawTree = SaplingRawTree
+  { srt_left :: !HexString
+  , srt_right :: !HexString
+  , srt_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 SaplingRawTree
+
+data SaplingTree = SaplingTree
+  { ot_left :: !(Maybe HexString)
+  , ot_right :: !(Maybe HexString)
+  , ot_parents :: ![Maybe HexString]
+  } deriving (Eq, Prelude.Show, Read)
+
 data SaplingFrontier = SaplingFrontier
   { sf_pos :: !Int64
   , sf_leaf :: !HexString