1// Copyright 2016 The Go Authors. All rights reserved. 2// Use of this source code is governed by a BSD-style 3// license that can be found in the LICENSE file. 4 5package ssa 6 7import "fmt" 8 9const ( 10 rankLeaf rbrank = 1 11 rankZero rbrank = 0 12) 13 14type rbrank int8 15 16// RBTint32 is a red-black tree with data stored at internal nodes, 17// following Tarjan, Data Structures and Network Algorithms, 18// pp 48-52, using explicit rank instead of red and black. 19// Deletion is not yet implemented because it is not yet needed. 20// Extra operations glb, lub, glbEq, lubEq are provided for 21// use in sparse lookup algorithms. 22type RBTint32 struct { 23 root *node32 24 // An extra-clever implementation will have special cases 25 // for small sets, but we are not extra-clever today. 26} 27 28func (t *RBTint32) String() string { 29 if t.root == nil { 30 return "[]" 31 } 32 return "[" + t.root.String() + "]" 33} 34 35func (t *node32) String() string { 36 s := "" 37 if t.left != nil { 38 s = t.left.String() + " " 39 } 40 s = s + fmt.Sprintf("k=%d,d=%v", t.key, t.data) 41 if t.right != nil { 42 s = s + " " + t.right.String() 43 } 44 return s 45} 46 47type node32 struct { 48 // Standard conventions hold for left = smaller, right = larger 49 left, right, parent *node32 50 data interface{} 51 key int32 52 rank rbrank // From Tarjan pp 48-49: 53 // If x is a node with a parent, then x.rank <= x.parent.rank <= x.rank+1. 54 // If x is a node with a grandparent, then x.rank < x.parent.parent.rank. 55 // If x is an "external [null] node", then x.rank = 0 && x.parent.rank = 1. 56 // Any node with one or more null children should have rank = 1. 57} 58 59// makeNode returns a new leaf node with the given key and nil data. 60func (t *RBTint32) makeNode(key int32) *node32 { 61 return &node32{key: key, rank: rankLeaf} 62} 63 64// IsEmpty reports whether t is empty. 65func (t *RBTint32) IsEmpty() bool { 66 return t.root == nil 67} 68 69// IsSingle reports whether t is a singleton (leaf). 70func (t *RBTint32) IsSingle() bool { 71 return t.root != nil && t.root.isLeaf() 72} 73 74// VisitInOrder applies f to the key and data pairs in t, 75// with keys ordered from smallest to largest. 76func (t *RBTint32) VisitInOrder(f func(int32, interface{})) { 77 if t.root == nil { 78 return 79 } 80 t.root.visitInOrder(f) 81} 82 83func (n *node32) Data() interface{} { 84 if n == nil { 85 return nil 86 } 87 return n.data 88} 89 90func (n *node32) keyAndData() (k int32, d interface{}) { 91 if n == nil { 92 k = 0 93 d = nil 94 } else { 95 k = n.key 96 d = n.data 97 } 98 return 99} 100 101func (n *node32) Rank() rbrank { 102 if n == nil { 103 return 0 104 } 105 return n.rank 106} 107 108// Find returns the data associated with key in the tree, or 109// nil if key is not in the tree. 110func (t *RBTint32) Find(key int32) interface{} { 111 return t.root.find(key).Data() 112} 113 114// Insert adds key to the tree and associates key with data. 115// If key was already in the tree, it updates the associated data. 116// Insert returns the previous data associated with key, 117// or nil if key was not present. 118// Insert panics if data is nil. 119func (t *RBTint32) Insert(key int32, data interface{}) interface{} { 120 if data == nil { 121 panic("Cannot insert nil data into tree") 122 } 123 n := t.root 124 var newroot *node32 125 if n == nil { 126 n = t.makeNode(key) 127 newroot = n 128 } else { 129 newroot, n = n.insert(key, t) 130 } 131 r := n.data 132 n.data = data 133 t.root = newroot 134 return r 135} 136 137// Min returns the minimum element of t and its associated data. 138// If t is empty, then (0, nil) is returned. 139func (t *RBTint32) Min() (k int32, d interface{}) { 140 return t.root.min().keyAndData() 141} 142 143// Max returns the maximum element of t and its associated data. 144// If t is empty, then (0, nil) is returned. 145func (t *RBTint32) Max() (k int32, d interface{}) { 146 return t.root.max().keyAndData() 147} 148 149// Glb returns the greatest-lower-bound-exclusive of x and its associated 150// data. If x has no glb in the tree, then (0, nil) is returned. 151func (t *RBTint32) Glb(x int32) (k int32, d interface{}) { 152 return t.root.glb(x, false).keyAndData() 153} 154 155// GlbEq returns the greatest-lower-bound-inclusive of x and its associated 156// data. If x has no glbEQ in the tree, then (0, nil) is returned. 157func (t *RBTint32) GlbEq(x int32) (k int32, d interface{}) { 158 return t.root.glb(x, true).keyAndData() 159} 160 161// Lub returns the least-upper-bound-exclusive of x and its associated 162// data. If x has no lub in the tree, then (0, nil) is returned. 163func (t *RBTint32) Lub(x int32) (k int32, d interface{}) { 164 return t.root.lub(x, false).keyAndData() 165} 166 167// LubEq returns the least-upper-bound-inclusive of x and its associated 168// data. If x has no lubEq in the tree, then (0, nil) is returned. 169func (t *RBTint32) LubEq(x int32) (k int32, d interface{}) { 170 return t.root.lub(x, true).keyAndData() 171} 172 173func (t *node32) isLeaf() bool { 174 return t.left == nil && t.right == nil 175} 176 177func (t *node32) visitInOrder(f func(int32, interface{})) { 178 if t.left != nil { 179 t.left.visitInOrder(f) 180 } 181 f(t.key, t.data) 182 if t.right != nil { 183 t.right.visitInOrder(f) 184 } 185} 186 187func (t *node32) maxChildRank() rbrank { 188 if t.left == nil { 189 if t.right == nil { 190 return rankZero 191 } 192 return t.right.rank 193 } 194 if t.right == nil { 195 return t.left.rank 196 } 197 if t.right.rank > t.left.rank { 198 return t.right.rank 199 } 200 return t.left.rank 201} 202 203func (t *node32) minChildRank() rbrank { 204 if t.left == nil || t.right == nil { 205 return rankZero 206 } 207 if t.right.rank < t.left.rank { 208 return t.right.rank 209 } 210 return t.left.rank 211} 212 213func (t *node32) find(key int32) *node32 { 214 for t != nil { 215 if key < t.key { 216 t = t.left 217 } else if key > t.key { 218 t = t.right 219 } else { 220 return t 221 } 222 } 223 return nil 224} 225 226func (t *node32) min() *node32 { 227 if t == nil { 228 return t 229 } 230 for t.left != nil { 231 t = t.left 232 } 233 return t 234} 235 236func (t *node32) max() *node32 { 237 if t == nil { 238 return t 239 } 240 for t.right != nil { 241 t = t.right 242 } 243 return t 244} 245 246func (t *node32) glb(key int32, allow_eq bool) *node32 { 247 var best *node32 248 for t != nil { 249 if key <= t.key { 250 if key == t.key && allow_eq { 251 return t 252 } 253 // t is too big, glb is to left. 254 t = t.left 255 } else { 256 // t is a lower bound, record it and seek a better one. 257 best = t 258 t = t.right 259 } 260 } 261 return best 262} 263 264func (t *node32) lub(key int32, allow_eq bool) *node32 { 265 var best *node32 266 for t != nil { 267 if key >= t.key { 268 if key == t.key && allow_eq { 269 return t 270 } 271 // t is too small, lub is to right. 272 t = t.right 273 } else { 274 // t is a upper bound, record it and seek a better one. 275 best = t 276 t = t.left 277 } 278 } 279 return best 280} 281 282func (t *node32) insert(x int32, w *RBTint32) (newroot, newnode *node32) { 283 // defaults 284 newroot = t 285 newnode = t 286 if x == t.key { 287 return 288 } 289 if x < t.key { 290 if t.left == nil { 291 n := w.makeNode(x) 292 n.parent = t 293 t.left = n 294 newnode = n 295 return 296 } 297 var new_l *node32 298 new_l, newnode = t.left.insert(x, w) 299 t.left = new_l 300 new_l.parent = t 301 newrank := 1 + new_l.maxChildRank() 302 if newrank > t.rank { 303 if newrank > 1+t.right.Rank() { // rotations required 304 if new_l.left.Rank() < new_l.right.Rank() { 305 // double rotation 306 t.left = new_l.rightToRoot() 307 } 308 newroot = t.leftToRoot() 309 return 310 } else { 311 t.rank = newrank 312 } 313 } 314 } else { // x > t.key 315 if t.right == nil { 316 n := w.makeNode(x) 317 n.parent = t 318 t.right = n 319 newnode = n 320 return 321 } 322 var new_r *node32 323 new_r, newnode = t.right.insert(x, w) 324 t.right = new_r 325 new_r.parent = t 326 newrank := 1 + new_r.maxChildRank() 327 if newrank > t.rank { 328 if newrank > 1+t.left.Rank() { // rotations required 329 if new_r.right.Rank() < new_r.left.Rank() { 330 // double rotation 331 t.right = new_r.leftToRoot() 332 } 333 newroot = t.rightToRoot() 334 return 335 } else { 336 t.rank = newrank 337 } 338 } 339 } 340 return 341} 342 343func (t *node32) rightToRoot() *node32 { 344 // this 345 // left right 346 // rl rr 347 // 348 // becomes 349 // 350 // right 351 // this rr 352 // left rl 353 // 354 right := t.right 355 rl := right.left 356 right.parent = t.parent 357 right.left = t 358 t.parent = right 359 // parent's child ptr fixed in caller 360 t.right = rl 361 if rl != nil { 362 rl.parent = t 363 } 364 return right 365} 366 367func (t *node32) leftToRoot() *node32 { 368 // this 369 // left right 370 // ll lr 371 // 372 // becomes 373 // 374 // left 375 // ll this 376 // lr right 377 // 378 left := t.left 379 lr := left.right 380 left.parent = t.parent 381 left.right = t 382 t.parent = left 383 // parent's child ptr fixed in caller 384 t.left = lr 385 if lr != nil { 386 lr.parent = t 387 } 388 return left 389} 390 391// next returns the successor of t in a left-to-right 392// walk of the tree in which t is embedded. 393func (t *node32) next() *node32 { 394 // If there is a right child, it is to the right 395 r := t.right 396 if r != nil { 397 return r.min() 398 } 399 // if t is p.left, then p, else repeat. 400 p := t.parent 401 for p != nil { 402 if p.left == t { 403 return p 404 } 405 t = p 406 p = t.parent 407 } 408 return nil 409} 410 411// prev returns the predecessor of t in a left-to-right 412// walk of the tree in which t is embedded. 413func (t *node32) prev() *node32 { 414 // If there is a left child, it is to the left 415 l := t.left 416 if l != nil { 417 return l.max() 418 } 419 // if t is p.right, then p, else repeat. 420 p := t.parent 421 for p != nil { 422 if p.right == t { 423 return p 424 } 425 t = p 426 p = t.parent 427 } 428 return nil 429} 430