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