From b0ca5b7d4b88f63071d026fd4c38a491c606e7eb Mon Sep 17 00:00:00 2001
From: Rene Vergara <rene@vergara.network>
Date: Fri, 25 Oct 2024 09:09:21 -0500
Subject: [PATCH] feat: structure for a balanced monoidal tree

---
 src/Zenith/Tree.hs | 92 +++++++++++++++++++++++++++++++++++++---------
 test/Spec.hs       | 27 ++++++++++++++
 2 files changed, 101 insertions(+), 18 deletions(-)

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