121 lines
3.4 KiB
Go
121 lines
3.4 KiB
Go
// Copyright (C) 2024 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"
|
|
"slices"
|
|
)
|
|
|
|
// FragmentedCollapsedTree represents a sequence of non-overlapping subtrees
|
|
type FragmentedCollapsedTree struct {
|
|
subtrees []CollapsedTree // sorted by offset
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) AddHash(position uint64, hash Hash) error {
|
|
return tree.Add(CollapsedTree{
|
|
offset: position,
|
|
nodes: []Hash{hash},
|
|
size: 1,
|
|
})
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Add(subtree CollapsedTree) error {
|
|
if subtree.size == 0 {
|
|
return nil
|
|
}
|
|
i := len(tree.subtrees)
|
|
for i > 0 && tree.subtrees[i-1].offset > subtree.offset {
|
|
i--
|
|
}
|
|
if i > 0 && tree.subtrees[i-1].offset+tree.subtrees[i-1].size > subtree.offset {
|
|
return fmt.Errorf("overlaps with subtree ending at %d", tree.subtrees[i-1].offset+tree.subtrees[i-1].size)
|
|
}
|
|
if i < len(tree.subtrees) && subtree.offset+subtree.size > tree.subtrees[i].offset {
|
|
return fmt.Errorf("overlaps with subtree starting at %d", tree.subtrees[i].offset)
|
|
}
|
|
if i == 0 || tree.subtrees[i-1].Append(subtree) != nil {
|
|
tree.subtrees = slices.Insert(tree.subtrees, i, subtree)
|
|
i++
|
|
}
|
|
for i < len(tree.subtrees) && tree.subtrees[i-1].Append(tree.subtrees[i]) == nil {
|
|
tree.subtrees = slices.Delete(tree.subtrees, i, i+1)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Merge(other FragmentedCollapsedTree) error {
|
|
for _, subtree := range other.subtrees {
|
|
if err := tree.Add(subtree); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Gaps(yield func(uint64, uint64) bool) {
|
|
var prevEnd uint64
|
|
for i := range tree.subtrees {
|
|
if prevEnd != tree.subtrees[i].offset {
|
|
if !yield(prevEnd, tree.subtrees[i].offset) {
|
|
return
|
|
}
|
|
}
|
|
prevEnd = tree.subtrees[i].offset + tree.subtrees[i].size
|
|
}
|
|
yield(prevEnd, 0)
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) NumSubtrees() int {
|
|
return len(tree.subtrees)
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Subtree(i int) CollapsedTree {
|
|
return tree.subtrees[i]
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Subtrees() []CollapsedTree {
|
|
if tree.subtrees == nil {
|
|
return []CollapsedTree{}
|
|
} else {
|
|
return tree.subtrees
|
|
}
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) IsComplete(size uint64) bool {
|
|
return len(tree.subtrees) == 1 && tree.subtrees[0].offset == 0 && tree.subtrees[0].size == size
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) Init(subtrees []CollapsedTree) error {
|
|
for i := 1; i < len(subtrees); i++ {
|
|
if subtrees[i-1].offset+subtrees[i-1].size > subtrees[i].offset {
|
|
return fmt.Errorf("subtrees %d and %d overlap", i-1, i)
|
|
}
|
|
}
|
|
tree.subtrees = subtrees
|
|
return nil
|
|
}
|
|
|
|
func (tree FragmentedCollapsedTree) MarshalJSON() ([]byte, error) {
|
|
return json.Marshal(tree.Subtrees())
|
|
}
|
|
|
|
func (tree *FragmentedCollapsedTree) UnmarshalJSON(b []byte) error {
|
|
var subtrees []CollapsedTree
|
|
if err := json.Unmarshal(b, &subtrees); err != nil {
|
|
return fmt.Errorf("error unmarshaling Fragmented Collapsed Merkle Tree: %w", err)
|
|
}
|
|
if err := tree.Init(subtrees); err != nil {
|
|
return fmt.Errorf("error unmarshaling Fragmented Collapsed Merkle Tree: %w", err)
|
|
}
|
|
return nil
|
|
}
|