github.com/MetalBlockchain/metalgo@v1.11.9/snow/engine/snowman/ancestor/tree.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package ancestor
     5  
     6  import (
     7  	"github.com/MetalBlockchain/metalgo/ids"
     8  	"github.com/MetalBlockchain/metalgo/utils/set"
     9  )
    10  
    11  var _ Tree = (*tree)(nil)
    12  
    13  // Tree manages a (potentially partial) view of a tree.
    14  //
    15  // For example, assume this is the full tree:
    16  //
    17  //		A
    18  //	  /	  \
    19  //	B		D
    20  //	|		|
    21  //	C		E
    22  //
    23  // A partial view of this tree may be:
    24  //
    25  //		A
    26  //	  /
    27  //	B		D
    28  //	|		|
    29  //	C		E
    30  //
    31  // Or:
    32  //
    33  //	B		D
    34  //	|		|
    35  //	C		E
    36  //
    37  // This structure is designed to update and traverse these partial views.
    38  type Tree interface {
    39  	// Add a mapping from blkID to parentID.
    40  	//
    41  	// Invariant: blkID must not be equal to parentID
    42  	// Invariant: a given blkID must only ever have one parentID
    43  	Add(blkID ids.ID, parentID ids.ID)
    44  
    45  	// Has returns if blkID's parentID is known by the tree.
    46  	Has(blkID ids.ID) bool
    47  
    48  	// GetAncestor returns the oldest known ancestor of blkID. If there is no
    49  	// known parentID of blkID, blkID will be returned.
    50  	GetAncestor(blkID ids.ID) ids.ID
    51  
    52  	// Remove the mapping from blkID to its parentID from the tree.
    53  	Remove(blkID ids.ID)
    54  
    55  	// RemoveDescendants removes blkID from the tree along with all of its known
    56  	// descendants.
    57  	RemoveDescendants(blkID ids.ID)
    58  
    59  	// Len returns the total number of blkID to parentID mappings that are
    60  	// currently tracked by the tree.
    61  	Len() int
    62  }
    63  
    64  type tree struct {
    65  	childToParent    map[ids.ID]ids.ID
    66  	parentToChildren map[ids.ID]set.Set[ids.ID]
    67  }
    68  
    69  func NewTree() Tree {
    70  	return &tree{
    71  		childToParent:    make(map[ids.ID]ids.ID),
    72  		parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
    73  	}
    74  }
    75  
    76  func (t *tree) Add(blkID ids.ID, parentID ids.ID) {
    77  	t.childToParent[blkID] = parentID
    78  
    79  	children := t.parentToChildren[parentID]
    80  	children.Add(blkID)
    81  	t.parentToChildren[parentID] = children
    82  }
    83  
    84  func (t *tree) Has(blkID ids.ID) bool {
    85  	_, ok := t.childToParent[blkID]
    86  	return ok
    87  }
    88  
    89  func (t *tree) GetAncestor(blkID ids.ID) ids.ID {
    90  	for {
    91  		parentID, ok := t.childToParent[blkID]
    92  		// this is the furthest parent available, break loop and return blkID
    93  		if !ok {
    94  			return blkID
    95  		}
    96  		// continue to loop with parentID
    97  		blkID = parentID
    98  	}
    99  }
   100  
   101  func (t *tree) Remove(blkID ids.ID) {
   102  	parent, ok := t.childToParent[blkID]
   103  	if !ok {
   104  		return
   105  	}
   106  	delete(t.childToParent, blkID)
   107  	// remove blkID from children
   108  	children := t.parentToChildren[parent]
   109  	children.Remove(blkID)
   110  	// this parent has no more children, remove it from map
   111  	if children.Len() == 0 {
   112  		delete(t.parentToChildren, parent)
   113  	}
   114  }
   115  
   116  func (t *tree) RemoveDescendants(blkID ids.ID) {
   117  	childrenList := []ids.ID{blkID}
   118  	for len(childrenList) > 0 {
   119  		newChildrenSize := len(childrenList) - 1
   120  		childID := childrenList[newChildrenSize]
   121  		childrenList = childrenList[:newChildrenSize]
   122  		t.Remove(childID)
   123  		// get children of child
   124  		for grandChildID := range t.parentToChildren[childID] {
   125  			childrenList = append(childrenList, grandChildID)
   126  		}
   127  	}
   128  }
   129  
   130  func (t *tree) Len() int {
   131  	return len(t.childToParent)
   132  }