diff --git a/src/Zenith/Tree.hs b/src/Zenith/Tree.hs index a2e1ec9..09dd323 100644 --- a/src/Zenith/Tree.hs +++ b/src/Zenith/Tree.hs @@ -1,5 +1,6 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} module Zenith.Tree where @@ -8,37 +9,92 @@ import Data.HexString class Monoid v => Measured a v where - measure :: a -> v + measure :: a -> Integer -> v -data Tree v a - = Leaf !v !a - | Branch !v !(Tree v a) !(Tree v a) +class Node v where + getLevel :: v -> Level + getTag :: v -> String + getPosition :: v -> Integer + isFull :: v -> Bool -value :: Tree v a -> v -value (Leaf v _) = v +type Level = Int + +type OrchardCommitment = String + +instance Measured [Char] OrchardNode where + measure oc p = OrchardNode p oc 0 True + +data Tree v + = Leaf !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 + +instance (Monoid v, Node v) => Semigroup (Tree v) where + (<>) 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 + (<>) (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) + | otherwise = branch (x <> Leaf w) y + (<>) (Branch s x y) (Branch s1 x1 y1) + | getLevel s == getLevel s1 = branch (Branch s x y) (Branch s1 x1 y1) + | otherwise = Branch s x y + +value :: Monoid v => Tree v -> v +value EmptyLeaf = mempty +value (Leaf v) = v value (Branch v _ _) = v -branch :: Monoid v => Tree v a -> Tree v a -> Tree v a +branch :: Monoid v => Tree v -> Tree v -> Tree v branch x y = Branch (value x <> value y) x y -leaf :: Measured a v => a -> Tree v a -leaf a = Leaf (measure a) a +leaf :: Measured a v => a -> Integer -> Tree v +leaf a p = Leaf (measure a p) data OrchardNode = OrchardNode { on_position :: !Integer , on_value :: !String - , on_level :: !Int - } deriving (Show, Eq) + , on_level :: !Level + , on_full :: !Bool + } deriving (Eq) instance Semigroup OrchardNode where (<>) x y = - if on_level x == on_level y - then OrchardNode - (max (on_position x) (on_position y)) - (on_value x <> on_value y) - (on_level x) - else x + OrchardNode + (max (on_position x) (on_position y)) + (on_value x <> on_value y) + (1 + on_level x) + (on_full x && on_full y) instance Monoid OrchardNode where - mempty = OrchardNode 0 "" (-1) + mempty = OrchardNode 0 "" 0 False mappend = (<>) + +instance Node OrchardNode where + getLevel = on_level + getTag = on_value + getPosition = on_position + isFull = on_full + +instance Show OrchardNode where + show = on_value + +instance Measured OrchardNode OrchardNode where + measure o p = OrchardNode p (on_value o) (on_level o) (on_full o) + +mkSubTree :: Node v => Monoid v => Level -> Tree v -> Tree v +mkSubTree level t = + if getLevel (value subtree) == level + then subtree + else mkSubTree level subtree + where + subtree = t <> EmptyLeaf diff --git a/test/Spec.hs b/test/Spec.hs index 5d548b4..e5f0351 100644 --- a/test/Spec.hs +++ b/test/Spec.hs @@ -39,6 +39,7 @@ import ZcashHaskell.Types import ZcashHaskell.Utils (readZebraTransaction) import Zenith.Core import Zenith.DB +import Zenith.Tree import Zenith.Types main :: IO () @@ -581,3 +582,29 @@ main = do case tx of Left e -> assertFailure $ show e Right h -> h `shouldNotBe` (hexString "deadbeef") + describe "Tree tests" $ do + it "Create leaf" $ do + let a = "a" :: OrchardCommitment + let n = append EmptyLeaf a :: Tree OrchardNode + getLevel (value n) `shouldBe` 0 + it "Create minimal tree" $ do + let a = "a" :: OrchardCommitment + 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 + let a = "a" :: OrchardCommitment + let b = "b" :: OrchardCommitment + let c = "c" :: OrchardCommitment + let d = "d" :: OrchardCommitment + let n = append EmptyLeaf a :: Tree OrchardNode + let t1 = append n b :: Tree OrchardNode + let t2 = append t1 c :: Tree OrchardNode + {- + -let t3 = append t2 d :: Tree OrchardNode + -} + {- + -getLevel (value t2) `shouldBe` 2 + -} + t2 `shouldBe` EmptyLeaf