feat: add function to fill tree to depth

src/Zenith/Tree.hs

@@ -5,42 +5,68 @@
 module Zenith.Tree where
 
 import Data.HexString
+import Data.Int (Int64)
+import ZcashHaskell.Orchard (combineOrchardNodes, getOrchardNodeValue)
+import ZcashHaskell.Types (OrchardFrontier(..))
+
+type Level = Integer
+
+maxLevel :: Level
+maxLevel = 32
+
+type Position = Int64
 
 class Monoid v =>
       Measured a v
   where
-  measure :: a -> Integer -> v
+  measure :: a -> Position -> v
 
 class Node v where
   getLevel :: v -> Level
-  getTag :: v -> String
+  getTag :: v -> HexString
-  getPosition :: v -> Integer
+  getPosition :: v -> Position
   isFull :: v -> Bool
+  mkNode :: Level -> Position -> HexString -> v
 
-type Level = Int
+type OrchardCommitment = HexString
 
-type OrchardCommitment = String
+instance Measured OrchardCommitment OrchardNode where
-
+  measure oc p =
-instance Measured [Char] OrchardNode where
+    case getOrchardNodeValue (hexBytes oc) of
-  measure oc p = OrchardNode p oc 0 True
+      Nothing -> OrchardNode 0 (hexString "00") 0 True
+      Just val -> OrchardNode p val 0 True
 
 data Tree v
-  = Leaf !v
+  = EmptyLeaf
+  | Leaf !v
+  | PrunedBranch !v
   | Branch !v !(Tree v) !(Tree v)
-  | EmptyLeaf
   deriving (Eq)
 
 instance (Node v, Show v) => Show (Tree v) where
   show EmptyLeaf = "()"
   show (Leaf v) = "(" ++ show v ++ ")"
-  show (Branch s x y) = "<" ++ getTag s ++ ">\n" ++ show x ++ "\n" ++ show y
+  show (PrunedBranch v) = "{" ++ show v ++ "}"
+  show (Branch s x y) =
+    "<" ++ show (getTag s) ++ ">\n" ++ show x ++ "\n" ++ show y
 
 instance (Monoid v, Node v) => Semigroup (Tree v) where
+  (<>) EmptyLeaf EmptyLeaf = PrunedBranch $ value $ branch EmptyLeaf EmptyLeaf
   (<>) EmptyLeaf x = x
   (<>) (Leaf x) EmptyLeaf = branch (Leaf x) EmptyLeaf
   (<>) (Leaf x) (Leaf y) = branch (Leaf x) (Leaf y)
   (<>) (Leaf x) Branch {} = Leaf x
-  (<>) (Branch s x y) EmptyLeaf = branch (Branch s x y) EmptyLeaf
+  (<>) (Leaf x) (PrunedBranch _) = Leaf x
+  (<>) (PrunedBranch x) EmptyLeaf = PrunedBranch $ x <> x
+  (<>) (PrunedBranch x) (Leaf _) = PrunedBranch x
+  (<>) (PrunedBranch x) (Branch s t u) =
+    if getLevel x == getLevel s
+      then branch (PrunedBranch x) (Branch s t u)
+      else EmptyLeaf
+  (<>) (PrunedBranch x) (PrunedBranch y) = PrunedBranch $ x <> y
+  (<>) (Branch s x y) EmptyLeaf =
+    branch (Branch s x y) $ getEmptyRoot (getLevel s)
+  (<>) (Branch s x y) (PrunedBranch _) = Branch s x y
   (<>) (Branch s x y) (Leaf w)
     | isFull s = Branch s x y <> mkSubTree (getLevel s) (Leaf w)
     | isFull (value x) = branch x (y <> Leaf w)
@@ -52,31 +78,47 @@ instance (Monoid v, Node v) => Semigroup (Tree v) where
 value :: Monoid v => Tree v -> v
 value EmptyLeaf = mempty
 value (Leaf v) = v
+value (PrunedBranch v) = v
 value (Branch v _ _) = v
 
 branch :: Monoid v => Tree v -> Tree v -> Tree v
 branch x y = Branch (value x <> value y) x y
 
-leaf :: Measured a v => a -> Integer -> Tree v
+leaf :: Measured a v => a -> Int64 -> Tree v
 leaf a p = Leaf (measure a p)
 
+prunedBranch :: Monoid v => Node v => Level -> Position -> HexString -> Tree v
+prunedBranch level pos val = PrunedBranch $ mkNode level pos val
+
+root :: Monoid v => Node v => Tree v -> Tree v
+root tree =
+  if getLevel (value tree) == maxLevel - 1
+    then tree
+    else mkSubTree (maxLevel - 1) tree
+
+getEmptyRoot :: Monoid v => Node v => Level -> Tree v
+getEmptyRoot level = iterate (\x -> x <> x) EmptyLeaf !! fromIntegral level
+
 data OrchardNode = OrchardNode
-  { on_position :: !Integer
+  { on_position :: !Position
-  , on_value :: !String
+  , on_value :: !HexString
   , on_level :: !Level
   , on_full :: !Bool
   } deriving (Eq)
 
 instance Semigroup OrchardNode where
   (<>) x y =
+    case combineOrchardNodes (on_level x) (on_value x) (on_value y) of
+      Nothing -> x
+      Just newHash ->
         OrchardNode
           (max (on_position x) (on_position y))
-      (on_value x <> on_value y)
+          newHash
           (1 + on_level x)
           (on_full x && on_full y)
 
 instance Monoid OrchardNode where
-  mempty = OrchardNode 0 "" 0 False
+  mempty = OrchardNode 0 (hexString "00") 0 False
   mappend = (<>)
 
 instance Node OrchardNode where
@@ -84,9 +126,10 @@ instance Node OrchardNode where
   getTag = on_value
   getPosition = on_position
   isFull = on_full
+  mkNode l p v = OrchardNode p v l True
 
 instance Show OrchardNode where
-  show = on_value
+  show = show . on_value
 
 instance Measured OrchardNode OrchardNode where
   measure o p = OrchardNode p (on_value o) (on_level o) (on_full o)
@@ -98,3 +141,21 @@ mkSubTree level t =
     else mkSubTree level subtree
   where
     subtree = t <> EmptyLeaf
+
+mkOrchardTree :: OrchardFrontier -> Tree OrchardNode
+mkOrchardTree (OrchardFrontier p l o) =
+  if odd p
+    then addOrchardOmmers (tail o) $
+         Leaf (OrchardNode (p - 1) (head o) 0 True) <>
+         Leaf (OrchardNode p l 0 True)
+    else addOrchardOmmers o $ Leaf (OrchardNode p l 0 True) <> EmptyLeaf
+
+addOrchardOmmers :: [HexString] -> Tree OrchardNode -> Tree OrchardNode
+addOrchardOmmers xs t =
+  foldl
+    (\s x -> PrunedBranch (mkNode (getLevel $ value s) (p (value s)) x) <> s)
+    t
+    xs
+  where
+    p :: OrchardNode -> Position
+    p (OrchardNode pos _ l _) = pos - (2 ^ l)

test/Spec.hs

@@ -10,7 +10,12 @@ import Database.Persist.Sqlite
 import System.Directory
 import Test.HUnit
 import Test.Hspec
-import ZcashHaskell.Orchard (isValidUnifiedAddress, parseAddress)
+import ZcashHaskell.Orchard
+  ( getOrchardFrontier
+  , getOrchardTreeAnchor
+  , isValidUnifiedAddress
+  , parseAddress
+  )
 import ZcashHaskell.Sapling
   ( decodeSaplingOutputEsk
   , encodeSaplingAddress
@@ -25,6 +30,8 @@ import ZcashHaskell.Transparent
   )
 import ZcashHaskell.Types
   ( DecodedNote(..)
+  , OrchardCommitmentTree(..)
+  , OrchardFrontier(..)
   , OrchardSpendingKey(..)
   , Phrase(..)
   , SaplingCommitmentTree(..)
@@ -583,28 +590,49 @@ main = do
                   Left e -> assertFailure $ show e
                   Right h -> h `shouldNotBe` (hexString "deadbeef")
     describe "Tree tests" $ do
+      let cmx1 =
+            hexString
+              "1712ead46028d4349e234abf59e94e0640fe7a0829e2e2e17e1a931631810400"
+      let t0 = EmptyLeaf <> EmptyLeaf :: Tree OrchardNode
+      let t1 = t0 <> EmptyLeaf :: Tree OrchardNode
+      let t1a = t0 <> t0
       it "Create leaf" $ do
-        let a = "a" :: OrchardCommitment
+        let n = leaf cmx1 1 :: Tree OrchardNode
-        let n = append EmptyLeaf a :: Tree OrchardNode
         getLevel (value n) `shouldBe` 0
       it "Create minimal tree" $ do
-        let a = "a" :: OrchardCommitment
+        let t = (leaf cmx1 1) <> EmptyLeaf :: Tree OrchardNode
-        let b = "b" :: OrchardCommitment
-        let n = append EmptyLeaf a :: Tree OrchardNode
-        let t = append n b :: Tree OrchardNode
         getLevel (value t) `shouldBe` 1
-      it "Create bigger tree" $ do
+      it "Create minimal empty tree" $ do
-        let a = "a" :: OrchardCommitment
+        getTag (value t0) `shouldNotBe` hexString "00"
-        let b = "b" :: OrchardCommitment
+      it "Expand empty tree" $ do t1 `shouldBe` t1a
-        let c = "c" :: OrchardCommitment
+      it "Create empty tree" $ mkSubTree 2 EmptyLeaf `shouldBe` t1
-        let d = "d" :: OrchardCommitment
+      it "Create empty tree non-rec" $ getEmptyRoot 2 `shouldBe` t1
-        let n = append EmptyLeaf a :: Tree OrchardNode
+      it "Create a tree from Frontier" $ do
-        let t1 = append n b :: Tree OrchardNode
+        let tree =
-        let t2 = append t1 c :: Tree OrchardNode
+              OrchardCommitmentTree $
-        {-
+              hexString
-         -let t3 = append t2 d :: Tree OrchardNode
+                "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000"
-         -}
+        case getOrchardFrontier tree of
-        {-
+          Nothing -> assertFailure "Failed to get frontier"
-         -getLevel (value t2) `shouldBe` 2
+          Just t1 -> do
-         -}
+            of_ommers t1 `shouldBe` []
-        t2 `shouldBe` EmptyLeaf
+      it "Validate a tree's depth from Frontier" $ do
+        let tree =
+              OrchardCommitmentTree $
+              hexString
+                "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000"
+        case getOrchardFrontier tree of
+          Nothing -> assertFailure "Failed to get frontier"
+          Just t1 -> do
+            let t = root $ mkOrchardTree t1
+            getLevel (value t) `shouldBe` 31
+      it "Validate a tree from Frontier" $ do
+        let tree =
+              OrchardCommitmentTree $
+              hexString
+                "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000"
+        case getOrchardFrontier tree of
+          Nothing -> assertFailure "Failed to get frontier"
+          Just t1 -> do
+            let t = root $ mkOrchardTree t1
+            getTag (value t) `shouldBe` getOrchardTreeAnchor t1

zcash-haskell

@@ -1 +1 @@
-Subproject commit 662a0d1148d3f52e2683157a4c9280bb8e81b0cb
+Subproject commit b6d490d05300a9db9cdf9929baa9b984bee9f3f6