From 1938e162f46ce7c879ccb3aa5e3eb6bbf7e63969 Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Wed, 30 Oct 2024 16:24:02 -0500
Subject: [PATCH] feat: update root testing functions

---
 librustzcash-wrapper/src/lib.rs | 18 +++++++++++++-----
 src/C/Zcash.chs                 |  3 ++-
 src/ZcashHaskell/Orchard.hs     |  6 +++---
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/librustzcash-wrapper/src/lib.rs b/librustzcash-wrapper/src/lib.rs
index 3b060fd..1b22434 100644
--- a/librustzcash-wrapper/src/lib.rs
+++ b/librustzcash-wrapper/src/lib.rs
@@ -1626,13 +1626,14 @@ pub extern "C" fn rust_wrapper_get_orchard_root(
     out_len: &mut usize
     ){
     let tree: CommitmentTree<MerkleHashOrchard, 32> = CommitmentTree::empty();
-    let root = tree.root_at_depth(level, PathFiller::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,
@@ -1643,10 +1644,17 @@ pub extern "C" fn rust_wrapper_orchard_add_node(
     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());
-        tree.append(n);
-        let root = tree.root();
-        let h = Hhex { bytes: root.to_bytes().to_vec() };
-        marshall_to_haskell_var(&h, out, out_len, RW);
+        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);
diff --git a/src/C/Zcash.chs b/src/C/Zcash.chs
index 915058f..fda259a 100644
--- a/src/C/Zcash.chs
+++ b/src/C/Zcash.chs
@@ -312,7 +312,8 @@ import ZcashHaskell.Types
 #}
 
 {# fun unsafe rust_wrapper_orchard_add_node as rustWrapperOrchardAddNodeTest
-     { toBorshVar* `BS.ByteString'&
+     { `Int8'
+     , toBorshVar* `BS.ByteString'&
      , getVarBuffer `Buffer HexString'&
      }
   -> `()'
diff --git a/src/ZcashHaskell/Orchard.hs b/src/ZcashHaskell/Orchard.hs
index 7e30158..3fdb25a 100644
--- a/src/ZcashHaskell/Orchard.hs
+++ b/src/ZcashHaskell/Orchard.hs
@@ -224,9 +224,9 @@ getOrchardRootTest :: Int -> HexString
 getOrchardRootTest level =
   withPureBorshVarBuffer $ rustWrapperGetOrchardRootTest $ fromIntegral level
 
-addOrchardNodeGetRoot :: BS.ByteString -> HexString
-addOrchardNodeGetRoot n =
-  withPureBorshVarBuffer $ rustWrapperOrchardAddNodeTest n
+addOrchardNodeGetRoot :: Int -> BS.ByteString -> HexString
+addOrchardNodeGetRoot l n =
+  withPureBorshVarBuffer $ rustWrapperOrchardAddNodeTest (fromIntegral l) n
 
 -- | Update a Orchard commitment tree
 updateOrchardCommitmentTree ::