tree.go

Documentation: github.com/transparency-dev/merkle/testonly

     1  // Copyright 2022 Google LLC. All Rights Reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package testonly
    16  
    17  import (
    18  	"github.com/transparency-dev/merkle"
    19  	"github.com/transparency-dev/merkle/compact"
    20  	"github.com/transparency-dev/merkle/proof"
    21  )
    22  
    23  // Tree implements an append-only Merkle tree. For testing.
    24  type Tree struct {
    25  	hasher merkle.LogHasher
    26  	size   uint64
    27  	hashes [][][]byte // Node hashes, indexed by node (level, index).
    28  }
    29  
    30  // New returns a new empty Merkle tree.
    31  func New(hasher merkle.LogHasher) *Tree {
    32  	return &Tree{hasher: hasher}
    33  }
    34  
    35  // AppendData adds the leaf hashes of the given entries to the end of the tree.
    36  func (t *Tree) AppendData(entries ...[]byte) {
    37  	for _, data := range entries {
    38  		t.appendImpl(t.hasher.HashLeaf(data))
    39  	}
    40  }
    41  
    42  // Append adds the given leaf hashes to the end of the tree.
    43  func (t *Tree) Append(hashes ...[]byte) {
    44  	for _, hash := range hashes {
    45  		t.appendImpl(hash)
    46  	}
    47  }
    48  
    49  func (t *Tree) appendImpl(hash []byte) {
    50  	level := 0
    51  	for ; (t.size>>level)&1 == 1; level++ {
    52  		row := append(t.hashes[level], hash)
    53  		hash = t.hasher.HashChildren(row[len(row)-2], hash)
    54  		t.hashes[level] = row
    55  	}
    56  	if level > len(t.hashes) {
    57  		panic("gap in tree appends")
    58  	} else if level == len(t.hashes) {
    59  		t.hashes = append(t.hashes, nil)
    60  	}
    61  
    62  	t.hashes[level] = append(t.hashes[level], hash)
    63  	t.size++
    64  }
    65  
    66  // Size returns the current number of leaves in the tree.
    67  func (t *Tree) Size() uint64 {
    68  	return t.size
    69  }
    70  
    71  // LeafHash returns the leaf hash at the given index.
    72  // Requires 0 <= index < Size(), otherwise panics.
    73  func (t *Tree) LeafHash(index uint64) []byte {
    74  	return t.hashes[0][index]
    75  }
    76  
    77  // Hash returns the current root hash of the tree.
    78  func (t *Tree) Hash() []byte {
    79  	return t.HashAt(t.size)
    80  }
    81  
    82  // HashAt returns the root hash at the given size.
    83  // Requires 0 <= size <= Size(), otherwise panics.
    84  func (t *Tree) HashAt(size uint64) []byte {
    85  	if size == 0 {
    86  		return t.hasher.EmptyRoot()
    87  	}
    88  	hashes := t.getNodes(compact.RangeNodes(0, size, nil))
    89  
    90  	hash := hashes[len(hashes)-1]
    91  	for i := len(hashes) - 2; i >= 0; i-- {
    92  		hash = t.hasher.HashChildren(hashes[i], hash)
    93  	}
    94  	return hash
    95  }
    96  
    97  // InclusionProof returns the inclusion proof for the given leaf index in the
    98  // tree of the given size. Requires 0 <= index < size <= Size(), otherwise may
    99  // panic.
   100  func (t *Tree) InclusionProof(index, size uint64) ([][]byte, error) {
   101  	nodes, err := proof.Inclusion(index, size)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  	return nodes.Rehash(t.getNodes(nodes.IDs), t.hasher.HashChildren)
   106  }
   107  
   108  // ConsistencyProof returns the consistency proof between the two given tree
   109  // sizes. Requires 0 <= size1 <= size2 <= Size(), otherwise may panic.
   110  func (t *Tree) ConsistencyProof(size1, size2 uint64) ([][]byte, error) {
   111  	nodes, err := proof.Consistency(size1, size2)
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  	return nodes.Rehash(t.getNodes(nodes.IDs), t.hasher.HashChildren)
   116  }
   117  
   118  func (t *Tree) getNodes(ids []compact.NodeID) [][]byte {
   119  	hashes := make([][]byte, len(ids))
   120  	for i, id := range ids {
   121  		hashes[i] = t.hashes[id.Level][id.Index]
   122  	}
   123  	return hashes
   124  }
   125
View as plain text