132 lines
3.2 KiB
Go
132 lines
3.2 KiB
Go
// Copyright (C) 2022 Opsmate, Inc.
|
|
//
|
|
// This Source Code Form is subject to the terms of the Mozilla
|
|
// Public License, v. 2.0. If a copy of the MPL was not distributed
|
|
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
//
|
|
// This software is distributed WITHOUT A WARRANTY OF ANY KIND.
|
|
// See the Mozilla Public License for details.
|
|
|
|
package merkletree
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"math/bits"
|
|
"slices"
|
|
)
|
|
|
|
type CollapsedTree struct {
|
|
nodes []Hash
|
|
size uint64
|
|
}
|
|
|
|
func calculateNumNodes(size uint64) int {
|
|
return bits.OnesCount64(size)
|
|
}
|
|
|
|
// TODO: phase out this function
|
|
func EmptyCollapsedTree() *CollapsedTree {
|
|
return &CollapsedTree{nodes: []Hash{}, size: 0}
|
|
}
|
|
|
|
// TODO: phase out this function
|
|
func NewCollapsedTree(nodes []Hash, size uint64) (*CollapsedTree, error) {
|
|
tree := new(CollapsedTree)
|
|
if err := tree.Init(nodes, size); err != nil {
|
|
return nil, err
|
|
}
|
|
return tree, nil
|
|
}
|
|
|
|
func (tree CollapsedTree) Equal(other CollapsedTree) bool {
|
|
return tree.size == other.size && slices.Equal(tree.nodes, other.nodes)
|
|
}
|
|
|
|
func (tree CollapsedTree) Clone() CollapsedTree {
|
|
return CollapsedTree{
|
|
nodes: slices.Clone(tree.nodes),
|
|
size: tree.size,
|
|
}
|
|
}
|
|
|
|
func (tree *CollapsedTree) Add(hash Hash) {
|
|
tree.nodes = append(tree.nodes, hash)
|
|
tree.size++
|
|
tree.collapse()
|
|
}
|
|
|
|
func (tree *CollapsedTree) Append(other CollapsedTree) error {
|
|
if tree.size > 0 {
|
|
maxSize := uint64(1) << bits.TrailingZeros64(tree.size)
|
|
if other.size > maxSize {
|
|
return fmt.Errorf("tree of size %d is too large to append to a tree of size %d (maximum size is %d)", other.size, tree.size, maxSize)
|
|
}
|
|
}
|
|
|
|
tree.nodes = append(tree.nodes, other.nodes...)
|
|
tree.size += other.size
|
|
tree.collapse()
|
|
return nil
|
|
}
|
|
|
|
func (tree *CollapsedTree) collapse() {
|
|
numNodes := calculateNumNodes(tree.size)
|
|
for len(tree.nodes) > numNodes {
|
|
left, right := tree.nodes[len(tree.nodes)-2], tree.nodes[len(tree.nodes)-1]
|
|
tree.nodes = tree.nodes[:len(tree.nodes)-2]
|
|
tree.nodes = append(tree.nodes, HashChildren(left, right))
|
|
}
|
|
}
|
|
|
|
func (tree *CollapsedTree) CalculateRoot() Hash {
|
|
if len(tree.nodes) == 0 {
|
|
return HashNothing()
|
|
}
|
|
i := len(tree.nodes) - 1
|
|
hash := tree.nodes[i]
|
|
for i > 0 {
|
|
i -= 1
|
|
hash = HashChildren(tree.nodes[i], hash)
|
|
}
|
|
return hash
|
|
}
|
|
|
|
func (tree *CollapsedTree) Nodes() []Hash {
|
|
return tree.nodes
|
|
}
|
|
|
|
func (tree *CollapsedTree) Size() uint64 {
|
|
return tree.size
|
|
}
|
|
|
|
func (tree CollapsedTree) MarshalJSON() ([]byte, error) {
|
|
return json.Marshal(map[string]interface{}{
|
|
"nodes": tree.nodes,
|
|
"size": tree.size,
|
|
})
|
|
}
|
|
|
|
func (tree *CollapsedTree) UnmarshalJSON(b []byte) error {
|
|
var rawTree struct {
|
|
Nodes []Hash `json:"nodes"`
|
|
Size uint64 `json:"size"`
|
|
}
|
|
if err := json.Unmarshal(b, &rawTree); err != nil {
|
|
return fmt.Errorf("error unmarshalling Collapsed Merkle Tree: %w", err)
|
|
}
|
|
if err := tree.Init(rawTree.Nodes, rawTree.Size); err != nil {
|
|
return fmt.Errorf("error unmarshalling Collapsed Merkle Tree: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (tree *CollapsedTree) Init(nodes []Hash, size uint64) error {
|
|
if len(nodes) != calculateNumNodes(size) {
|
|
return fmt.Errorf("nodes has wrong length (should be %d, not %d)", calculateNumNodes(size), len(nodes))
|
|
}
|
|
tree.size = size
|
|
tree.nodes = nodes
|
|
return nil
|
|
}
|