nockchain-grpc/wallet/ztree.go

package wallet

import (
	"slices"

	"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)
}

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 (z *ZTree) KeyLeftMost() interface{} {
	node := z.Root
	for node.Left != nil {
		node = node.Left
	}
	return node.Key
}
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)

	if slices.Compare(keyHash[:], nodeKeyHash[:]) == -1 {
		// key < node key
		if slices.Compare(keyDoubleHash[:], nodeKeyDoubleHash[:]) == -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
			node = &ZNode{
				Key:   key,
				Value: value,
				Right: nil,
				Left: &ZNode{
					Key:   nodeKey,
					Value: nodeValue,
					Left:  leftNode,
					Right: rightNode,
				},
			}
		}
	} else {
		// key > node key
		if slices.Compare(keyDoubleHash[:], nodeKeyDoubleHash[:]) == -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
			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
}