Add merkletree.FragmentedCollapsedTree

merkletree/fragment.go

@@ -0,0 +1,120 @@
+// 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
+}