From 8ede441191cf083ec847faf243964683ddb7a830 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Sat, 17 Apr 2021 09:33:00 +0200
Subject: [PATCH] Add Graphviz generation methods

---
 tree.go      | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 tree_test.go |  5 +++-
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/tree.go b/tree.go
index 8194a6b..14a31b0 100644
--- a/tree.go
+++ b/tree.go
@@ -598,6 +598,12 @@ func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool,
 }
 
 func (t *Tree) dbGet(k []byte) ([]byte, error) {
+	// if key is empty, return empty as value
+	empty := make([]byte, t.hashFunction.Len())
+	if bytes.Equal(k, empty) {
+		return empty, nil
+	}
+
 	v, err := t.db.Get(k)
 	if err == nil {
 		return v, nil
@@ -696,3 +702,76 @@ func (t *Tree) ImportDump(b []byte) error {
 	}
 	return nil
 }
+
+// Graphviz iterates across the full tree to generate a string Graphviz
+// representation of the tree and writes it to w
+func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
+	fmt.Fprintf(w, `digraph hierarchy {
+node [fontname=Monospace,fontsize=10,shape=box]
+`)
+	nChars := 4
+	nEmpties := 0
+	empty := make([]byte, t.hashFunction.Len())
+	err := t.Iterate(func(k, v []byte) {
+		switch v[0] {
+		case PrefixValueEmpty:
+		case PrefixValueLeaf:
+			fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
+			// key & value from the leaf
+			kB, vB := readLeafValue(v)
+			fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
+				hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
+				hex.EncodeToString(vB[:nChars]))
+			fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
+				hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
+		case PrefixValueIntermediate:
+			l, r := readIntermediateChilds(v)
+			lStr := hex.EncodeToString(l[:nChars])
+			rStr := hex.EncodeToString(r[:nChars])
+			eStr := ""
+			if bytes.Equal(l, empty) {
+				lStr = fmt.Sprintf("empty%v", nEmpties)
+				eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
+					lStr)
+				nEmpties++
+			}
+			if bytes.Equal(r, empty) {
+				rStr = fmt.Sprintf("empty%v", nEmpties)
+				eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
+					rStr)
+				nEmpties++
+			}
+			fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
+				lStr, rStr)
+			fmt.Fprint(w, eStr)
+		default:
+		}
+	})
+	fmt.Fprintf(w, "}\n")
+	return err
+}
+
+// PrintGraphviz prints the output of Tree.Graphviz
+func (t *Tree) PrintGraphviz(rootKey []byte) error {
+	if rootKey == nil {
+		rootKey = t.Root()
+	}
+	w := bytes.NewBufferString("")
+	fmt.Fprintf(w,
+		"--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
+	err := t.Graphviz(w, nil)
+	if err != nil {
+		fmt.Println(w)
+		return err
+	}
+	fmt.Fprintf(w,
+		"End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
+
+	fmt.Println(w)
+	return nil
+}
+
+// Purge WIP: unimplemented
+func (t *Tree) Purge(keys [][]byte) error {
+	return nil
+}
diff --git a/tree_test.go b/tree_test.go
index 6412e72..53d1c15 100644
--- a/tree_test.go
+++ b/tree_test.go
@@ -195,7 +195,7 @@ func TestUpdate(t *testing.T) {
 	c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(11)))
 }
 
-func TestAux(t *testing.T) { // TMP
+func TestAux(t *testing.T) { // TODO split in proper tests
 	c := qt.New(t)
 	tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
 	c.Assert(err, qt.IsNil)
@@ -219,6 +219,9 @@ func TestAux(t *testing.T) { // TMP
 	k = BigIntToBytes(big.NewInt(int64(770)))
 	err = tree.Add(k, v)
 	c.Assert(err, qt.IsNil)
+	//
+	// err = tree.PrintGraphviz(nil)
+	// c.Assert(err, qt.IsNil)
 }
 
 func TestGet(t *testing.T) {