nockchain-grpc/wallet/ztree.go

package wallet

import (
	"math/big"
	"slices"

	"github.com/btcsuite/btcd/btcutil/base58"
	"github.com/phamminh0811/private-grpc/crypto"
)

type ZNode struct {
	Key   interface{}
	Value interface{}

	Left  *ZNode
	Right *ZNode
}

type ZTree struct {
	Root          *ZNode
	HashKeyFunc   func(interface{}) [5]uint64
	HashValueFunc func(interface{}) ([5]uint64, error)
}

type ZPair struct {
	Key   interface{}
	Value interface{}
}

func NewZTree(hashKeyFunc func(interface{}) [5]uint64, hashValueFunc func(interface{}) ([5]uint64, error)) *ZTree {
	return &ZTree{
		Root:          nil,
		HashKeyFunc:   hashKeyFunc,
		HashValueFunc: hashValueFunc,
	}
}

func (z *ZTree) Insert(key, value interface{}) {
	z.Root = z.Root.InsertNode(z.HashKeyFunc, key, value)
}

func (z *ZTree) Hash() ([5]uint64, error) {
	return z.Root.HashNode(z.HashValueFunc)
}

func (t *ZTree) Tap() []ZPair {
	tap := []ZPair{}
	TapNode(t.Root, &tap)
	slices.Reverse(tap)
	return tap
}

func (node *ZNode) InsertNode(hashFunc func(interface{}) [5]uint64, key, value interface{}) *ZNode {
	if node == nil {
		node = &ZNode{
			Key:   key,
			Value: value,
			Left:  nil,
			Right: nil,
		}
		return node
	}
	keyHash := hashFunc(key)
	keyDoubleHash := crypto.Tip5RehashTenCell(keyHash, keyHash)
	nodeKeyHash := hashFunc(node.Key)
	nodeKeyDoubleHash := crypto.Tip5RehashTenCell(nodeKeyHash, nodeKeyHash)

	keyHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(keyHash)))
	nodeKeyHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(nodeKeyHash)))
	keyDoubleHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(keyDoubleHash)))
	nodeKeyDoubleHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(nodeKeyDoubleHash)))
	if keyHashBigInt.Cmp(nodeKeyHashBigInt) == 1 {
		// key < node key
		if keyDoubleHashBigInt.Cmp(nodeKeyDoubleHashBigInt) == 1 {
			// reinsert in right
			node.Right = node.Right.InsertNode(hashFunc, key, value)
		} else {
			//      new key
			// 	    /    \
			//  old key   ~
			//   /  \
			// ...   ...
			nodeKey := node.Key
			nodeValue := node.Value
			leftNode := node.Left
			rightNode := node.Right
			if rightNode == nil {
				node = &ZNode{
					Key:   key,
					Value: value,
					Right: nil,
					Left: &ZNode{
						Key:   nodeKey,
						Value: nodeValue,
						Left:  leftNode,
						Right: nil,
					},
				}
			} else {
				rightNodeKeyHash := hashFunc(rightNode.Key)
				rightNodeKeyHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(rightNodeKeyHash)))
				if rightNodeKeyHashBigInt.Cmp(keyHashBigInt) == -1 {
					node = &ZNode{
						Key:   key,
						Value: value,
						Right: nil,
						Left: &ZNode{
							Key:   nodeKey,
							Value: nodeValue,
							Left:  leftNode,
							Right: rightNode,
						},
					}
				} else {
					node = &ZNode{
						Key:   key,
						Value: value,
						Right: rightNode,
						Left: &ZNode{
							Key:   nodeKey,
							Value: nodeValue,
							Left:  leftNode,
							Right: nil,
						},
					}
				}
			}

		}
	} else {
		// key > node key
		if keyDoubleHashBigInt.Cmp(nodeKeyDoubleHashBigInt) == 1 {
			// reinsert in left
			node.Left = node.Left.InsertNode(hashFunc, key, value)
		} else {
			//    new key
			// 	   /    \
			//    ~   old key
			// 	        /  \
			// 	     ...   ...
			nodeKey := node.Key
			nodeValue := node.Value
			leftNode := node.Left
			rightNode := node.Right
			if leftNode == nil {
				node = &ZNode{
					Key:   key,
					Value: value,
					Right: &ZNode{
						Key:   nodeKey,
						Value: nodeValue,
						Left:  nil,
						Right: rightNode,
					},
					Left: nil,
				}
			} else {
				leftNodeKeyHash := hashFunc(leftNode.Key)
				leftNodeKeyHashBigInt := new(big.Int).SetBytes(base58.Decode(crypto.Tip5HashToBase58(leftNodeKeyHash)))
				if leftNodeKeyHashBigInt.Cmp(keyHashBigInt) == -1 {
					node = &ZNode{
						Key:   key,
						Value: value,
						Right: &ZNode{
							Key:   nodeKey,
							Value: nodeValue,
							Left:  nil,
							Right: rightNode,
						},
						Left: leftNode,
					}
				} else {
					node = &ZNode{
						Key:   key,
						Value: value,
						Right: &ZNode{
							Key:   nodeKey,
							Value: nodeValue,
							Left:  leftNode,
							Right: rightNode,
						},
						Left: nil,
					}
				}
			}
		}
	}
	return node
}

func (node *ZNode) HashNode(hashFunc func(interface{}) ([5]uint64, error)) ([5]uint64, error) {
	if node == nil {
		return crypto.Tip5Zero, nil
	}

	leftHash, err := node.Left.HashNode(hashFunc)
	if err != nil {
		return [5]uint64{}, err
	}
	rightHash, err := node.Right.HashNode(hashFunc)
	if err != nil {
		return [5]uint64{}, err
	}

	hashLeftRight := crypto.Tip5RehashTenCell(leftHash, rightHash)
	valHash, err := hashFunc(node.Value)
	if err != nil {
		return [5]uint64{}, err
	}
	return crypto.Tip5RehashTenCell(valHash, hashLeftRight), nil
}

func TapNode(node *ZNode, acc *[]ZPair) {
	if node == nil {
		return
	}

	stored := false
	if node.Right != nil {
		TapNode(node.Right, acc)
	} else {
		*acc = append(*acc, ZPair{Key: node.Key, Value: node.Value})
		stored = true
	}

	if node.Left != nil {
		TapNode(node.Left, acc)
	}

	if !stored {
		*acc = append(*acc, ZPair{Key: node.Key, Value: node.Value})
	}
}