trie.go (6274B)
1 // Copyright 2021 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 5 // trie implements persistent Patricia trie maps. 6 // 7 // Each Map is effectively a map from uint64 to interface{}. Patricia tries are 8 // a form of radix tree that are particularly appropriate when many maps will be 9 // created, merged together and large amounts of sharing are expected (e.g. 10 // environment abstract domains in program analysis). 11 // 12 // This implementation closely follows the paper: 13 // 14 // C. Okasaki and A. Gill, “Fast mergeable integer maps,” in ACM SIGPLAN 15 // Workshop on ML, September 1998, pp. 77–86. 16 // 17 // Each Map is immutable and can be read from concurrently. The map does not 18 // guarantee that the value pointed to by the interface{} value is not updated 19 // concurrently. 20 // 21 // These Maps are optimized for situations where there will be many maps created at 22 // with a high degree of sharing and combining of maps together. If you do not expect, 23 // significant amount of sharing, the builtin map[T]U is much better choice! 24 // 25 // Each Map is created by a Builder. Each Builder has a unique Scope and each node is 26 // created within this scope. Maps x and y are == if they contains the same 27 // (key,value) mappings and have equal scopes. 28 // 29 // Internally these are big endian Patricia trie nodes, and the keys are sorted. 30 package trie 31 32 import ( 33 "fmt" 34 "strings" 35 ) 36 37 // Map is effectively a finite mapping from uint64 keys to interface{} values. 38 // Maps are immutable and can be read from concurrently. 39 // 40 // Notes on concurrency: 41 // - A Map value itself is an interface and assignments to a Map value can race. 42 // - Map does not guarantee that the value pointed to by the interface{} value 43 // is not updated concurrently. 44 type Map struct { 45 s Scope 46 n node 47 } 48 49 func (m Map) Scope() Scope { 50 return m.s 51 } 52 func (m Map) Size() int { 53 if m.n == nil { 54 return 0 55 } 56 return m.n.size() 57 } 58 func (m Map) Lookup(k uint64) (any, bool) { 59 if m.n != nil { 60 if leaf := m.n.find(key(k)); leaf != nil { 61 return leaf.v, true 62 } 63 } 64 return nil, false 65 } 66 67 // Converts the map into a {<key>: <value>[, ...]} string. This uses the default 68 // %s string conversion for <value>. 69 func (m Map) String() string { 70 var kvs []string 71 m.Range(func(u uint64, i any) bool { 72 kvs = append(kvs, fmt.Sprintf("%d: %s", u, i)) 73 return true 74 }) 75 return fmt.Sprintf("{%s}", strings.Join(kvs, ", ")) 76 } 77 78 // Range over the leaf (key, value) pairs in the map in order and 79 // applies cb(key, value) to each. Stops early if cb returns false. 80 // Returns true if all elements were visited without stopping early. 81 func (m Map) Range(cb func(uint64, any) bool) bool { 82 if m.n != nil { 83 return m.n.visit(cb) 84 } 85 return true 86 } 87 88 // DeepEqual returns true if m and other contain the same (k, v) mappings 89 // [regardless of Scope]. 90 // 91 // Equivalently m.DeepEqual(other) <=> reflect.DeepEqual(Elems(m), Elems(other)) 92 func (m Map) DeepEqual(other Map) bool { 93 if m.Scope() == other.Scope() { 94 return m.n == other.n 95 } 96 if (m.n == nil) || (other.n == nil) { 97 return m.Size() == 0 && other.Size() == 0 98 } 99 return m.n.deepEqual(other.n) 100 } 101 102 // Elems are the (k,v) elements in the Map as a map[uint64]interface{} 103 func Elems(m Map) map[uint64]any { 104 dest := make(map[uint64]any, m.Size()) 105 m.Range(func(k uint64, v any) bool { 106 dest[k] = v 107 return true 108 }) 109 return dest 110 } 111 112 // node is an internal node within a trie map. 113 // A node is either empty, a leaf or a branch. 114 type node interface { 115 size() int 116 117 // visit the leaves (key, value) pairs in the map in order and 118 // applies cb(key, value) to each. Stops early if cb returns false. 119 // Returns true if all elements were visited without stopping early. 120 visit(cb func(uint64, any) bool) bool 121 122 // Two nodes contain the same elements regardless of scope. 123 deepEqual(node) bool 124 125 // find the leaf for the given key value or nil if it is not present. 126 find(k key) *leaf 127 128 // implementations must implement this. 129 nodeImpl() 130 } 131 132 // empty represents the empty map within a scope. 133 // 134 // The current builder ensure 135 type empty struct { 136 s Scope 137 } 138 139 // leaf represents a single <key, value> pair. 140 type leaf struct { 141 k key 142 v any 143 } 144 145 // branch represents a tree node within the Patricia trie. 146 // 147 // All keys within the branch match a `prefix` of the key 148 // up to a `branching` bit, and the left and right nodes 149 // contain keys that disagree on the bit at the `branching` bit. 150 type branch struct { 151 sz int // size. cached for O(1) lookup 152 prefix prefix // == mask(p0, branching) for some p0 153 branching bitpos 154 155 // Invariants: 156 // - neither is nil. 157 // - neither is *empty. 158 // - all keys in left are <= p. 159 // - all keys in right are > p. 160 left, right node 161 } 162 163 // all of these types are Maps. 164 var _ node = &empty{} 165 var _ node = &leaf{} 166 var _ node = &branch{} 167 168 func (*empty) nodeImpl() {} 169 func (*leaf) nodeImpl() {} 170 func (*branch) nodeImpl() {} 171 172 func (*empty) find(k key) *leaf { return nil } 173 func (l *leaf) find(k key) *leaf { 174 if k == l.k { 175 return l 176 } 177 return nil 178 } 179 func (br *branch) find(k key) *leaf { 180 kp := prefix(k) 181 if !matchPrefix(kp, br.prefix, br.branching) { 182 return nil 183 } 184 if zeroBit(kp, br.branching) { 185 return br.left.find(k) 186 } 187 return br.right.find(k) 188 } 189 190 func (*empty) size() int { return 0 } 191 func (*leaf) size() int { return 1 } 192 func (br *branch) size() int { return br.sz } 193 194 func (*empty) deepEqual(m node) bool { 195 _, ok := m.(*empty) 196 return ok 197 } 198 func (l *leaf) deepEqual(m node) bool { 199 if m, ok := m.(*leaf); ok { 200 return m == l || (l.k == m.k && l.v == m.v) 201 } 202 return false 203 } 204 205 func (br *branch) deepEqual(m node) bool { 206 if m, ok := m.(*branch); ok { 207 if br == m { 208 return true 209 } 210 return br.sz == m.sz && br.branching == m.branching && br.prefix == m.prefix && 211 br.left.deepEqual(m.left) && br.right.deepEqual(m.right) 212 } 213 // if m is not a branch, m contains 0 or 1 elem. 214 // br contains at least 2 keys that disagree on a prefix. 215 return false 216 } 217 218 func (*empty) visit(cb func(uint64, any) bool) bool { 219 return true 220 } 221 func (l *leaf) visit(cb func(uint64, any) bool) bool { 222 return cb(uint64(l.k), l.v) 223 } 224 func (br *branch) visit(cb func(uint64, any) bool) bool { 225 if !br.left.visit(cb) { 226 return false 227 } 228 return br.right.visit(cb) 229 }