1 package mls
2 3 // This package uses an array-based representation of complete balanced binary
4 // trees, as described in appendix C. For example, a tree with 8 leaves:
5 //
6 // X
7 // |
8 // .---------+---------.
9 // / \
10 // X X
11 // | |
12 // .---+---. .---+---.
13 // / \ / \
14 // X X X X
15 // / \ / \ / \ / \
16 // / \ / \ / \ / \
17 // X X X X X X X X
18 //
19 // Node: 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14
20 //
21 // Leaf: 0 1 2 3 4 5 6 7
22 23 // numLeaves exposes operations on a tree with a given number of leaves.
24 type numLeaves uint32
25 26 func numLeavesFromWidth(w uint32) numLeaves {
27 if w == 0 {
28 return 0
29 }
30 return numLeaves((w-1)/2 + 1)
31 }
32 33 // width computes the minimum length of the array, ie. the number of nodes.
34 func (n numLeaves) width() uint32 {
35 if n == 0 {
36 return 0
37 }
38 return 2*(uint32(n)-1) + 1
39 }
40 41 // root returns the index of the root node.
42 func (n numLeaves) root() nodeIndex {
43 return nodeIndex((1 << log2(n.width())) - 1)
44 }
45 46 // parent returns the index of the parent node for a non-root node index.
47 func (n numLeaves) parent(x nodeIndex) (nodeIndex, bool) {
48 if x == n.root() {
49 return 0, false
50 }
51 lvl := nodeIndex(x.level())
52 b := (x >> (lvl + 1)) & 1
53 p := (x | (1 << lvl)) ^ (b << (lvl + 1))
54 return p, true
55 }
56 57 // sibling returns the index of the other child of the node's parent.
58 func (n numLeaves) sibling(x nodeIndex) (nodeIndex, bool) {
59 p, ok := n.parent(x)
60 if !ok {
61 return 0, false
62 }
63 if x < p {
64 return p.right()
65 } else {
66 return p.left()
67 }
68 }
69 70 // directPath computes the direct path of a node, ordered from leaf to root.
71 func (n numLeaves) directPath(x nodeIndex) []nodeIndex {
72 var path []nodeIndex
73 for {
74 p, ok := n.parent(x)
75 if !ok {
76 break
77 }
78 path = append(path, p)
79 x = p
80 }
81 return path
82 }
83 84 // copath computes the copath of a node, ordered from leaf to root.
85 func (n numLeaves) copath(x nodeIndex) []nodeIndex {
86 path := n.directPath(x)
87 if len(path) == 0 {
88 return nil
89 }
90 path = append([]nodeIndex{x}, path...)
91 path = path[:len(path)-1]
92 93 var copath []nodeIndex
94 for _, y := range path {
95 s, ok := n.sibling(y)
96 if !ok {
97 panic("unreachable")
98 }
99 copath = append(copath, s)
100 }
101 102 return copath
103 }
104 105 // nodeIndex is the index of a node in a tree.
106 type nodeIndex uint32
107 108 // isLeaf returns true if this is a leaf node, false if this is an intermediate
109 // node.
110 func (x nodeIndex) isLeaf() bool {
111 return x%2 == 0
112 }
113 114 // leafIndex returns the index of the leaf from a node index.
115 func (x nodeIndex) leafIndex() (leafIndex, bool) {
116 if !x.isLeaf() {
117 return 0, false
118 }
119 return leafIndex(x) >> 1, true
120 }
121 122 // left returns the index of the left child for an intermediate node index.
123 func (x nodeIndex) left() (nodeIndex, bool) {
124 lvl := x.level()
125 if lvl == 0 {
126 return 0, false
127 }
128 l := x ^ (1 << (nodeIndex(lvl) - 1))
129 return l, true
130 }
131 132 // right returns the index of the right child for an intermediate node index.
133 func (x nodeIndex) right() (nodeIndex, bool) {
134 lvl := x.level()
135 if lvl == 0 {
136 return 0, false
137 }
138 r := x ^ (3 << (nodeIndex(lvl) - 1))
139 return r, true
140 }
141 142 // children returns the indices of the left and right children for an
143 // intermediate node index.
144 func (x nodeIndex) children() (left, right nodeIndex, ok bool) {
145 l, ok := x.left()
146 if !ok {
147 return 0, 0, false
148 }
149 r, _ := x.right()
150 return l, r, true
151 }
152 153 // level returns the level of a node in the tree. Leaves are at level 0, their
154 // parents are at level 1, etc.
155 func (x nodeIndex) level() uint32 {
156 if x&1 == 0 {
157 return 0
158 }
159 lvl := uint32(0)
160 for (x>>lvl)&1 == 1 {
161 lvl++
162 }
163 return lvl
164 }
165 166 // commonAncestor returns the the lowest node that is in the direct paths of
167 // both leaves.
168 func commonAncestor(x, y nodeIndex) nodeIndex {
169 // Handle cases where one is an ancestor of the other
170 lx, ly := x.level()+1, y.level()+1
171 if lx <= ly && x>>ly == y>>ly {
172 return y
173 } else if ly <= lx && x>>lx == y>>lx {
174 return x
175 }
176 177 // Handle other cases
178 xn, yn := x, y
179 k := 0
180 for xn != yn {
181 xn, yn = xn>>1, yn>>1
182 k++
183 }
184 return (xn << k) + (1 << (k - 1)) - 1
185 }
186 187 type leafIndex uint32
188 189 // nodeIndex returns the index of the node from a leaf index.
190 func (li leafIndex) nodeIndex() nodeIndex {
191 return nodeIndex(2 * li)
192 }
193 194 // log2 computes the exponent of the largest power of 2 less than x.
195 func log2(x uint32) uint32 {
196 if x == 0 {
197 return 0
198 }
199 200 k := uint32(0)
201 for x>>k > 0 {
202 k++
203 }
204 return k - 1
205 }
206 207 func isPowerOf2(x uint32) bool {
208 return x != 0 && x&(x-1) == 0
209 }
210