| 1 | // Copyright 2022 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 | // The persistent package defines various persistent data structures; |
| 6 | // that is, data structures that can be efficiently copied and modified |
| 7 | // in sublinear time. |
| 8 | package persistent |
| 9 | |
| 10 | import ( |
| 11 | "fmt" |
| 12 | "math/rand" |
| 13 | "strings" |
| 14 | "sync/atomic" |
| 15 | ) |
| 16 | |
| 17 | // Implementation details: |
| 18 | // * Each value is reference counted by nodes which hold it. |
| 19 | // * Each node is reference counted by its parent nodes. |
| 20 | // * Each map is considered a top-level parent node from reference counting perspective. |
| 21 | // * Each change does always effectivelly produce a new top level node. |
| 22 | // |
| 23 | // Functions which operate directly with nodes do have a notation in form of |
| 24 | // `foo(arg1:+n1, arg2:+n2) (ret1:+n3)`. |
| 25 | // Each argument is followed by a delta change to its reference counter. |
| 26 | // In case if no change is expected, the delta will be `-0`. |
| 27 | |
| 28 | // Map is an associative mapping from keys to values, both represented as |
| 29 | // interface{}. Key comparison and iteration order is defined by a |
| 30 | // client-provided function that implements a strict weak order. |
| 31 | // |
| 32 | // Maps can be Cloned in constant time. |
| 33 | // Get, Store, and Delete operations are done on average in logarithmic time. |
| 34 | // Maps can be Updated in O(m log(n/m)) time for maps of size n and m, where m < n. |
| 35 | // |
| 36 | // Values are reference counted, and a client-supplied release function |
| 37 | // is called when a value is no longer referenced by a map or any clone. |
| 38 | // |
| 39 | // Internally the implementation is based on a randomized persistent treap: |
| 40 | // https://en.wikipedia.org/wiki/Treap. |
| 41 | type Map struct { |
| 42 | less func(a, b interface{}) bool |
| 43 | root *mapNode |
| 44 | } |
| 45 | |
| 46 | func (m *Map) String() string { |
| 47 | var buf strings.Builder |
| 48 | buf.WriteByte('{') |
| 49 | var sep string |
| 50 | m.Range(func(k, v interface{}) { |
| 51 | fmt.Fprintf(&buf, "%s%v: %v", sep, k, v) |
| 52 | sep = ", " |
| 53 | }) |
| 54 | buf.WriteByte('}') |
| 55 | return buf.String() |
| 56 | } |
| 57 | |
| 58 | type mapNode struct { |
| 59 | key interface{} |
| 60 | value *refValue |
| 61 | weight uint64 |
| 62 | refCount int32 |
| 63 | left, right *mapNode |
| 64 | } |
| 65 | |
| 66 | type refValue struct { |
| 67 | refCount int32 |
| 68 | value interface{} |
| 69 | release func(key, value interface{}) |
| 70 | } |
| 71 | |
| 72 | func newNodeWithRef(key, value interface{}, release func(key, value interface{})) *mapNode { |
| 73 | return &mapNode{ |
| 74 | key: key, |
| 75 | value: &refValue{ |
| 76 | value: value, |
| 77 | release: release, |
| 78 | refCount: 1, |
| 79 | }, |
| 80 | refCount: 1, |
| 81 | weight: rand.Uint64(), |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | func (node *mapNode) shallowCloneWithRef() *mapNode { |
| 86 | atomic.AddInt32(&node.value.refCount, 1) |
| 87 | return &mapNode{ |
| 88 | key: node.key, |
| 89 | value: node.value, |
| 90 | weight: node.weight, |
| 91 | refCount: 1, |
| 92 | } |
| 93 | } |
| 94 | |
| 95 | func (node *mapNode) incref() *mapNode { |
| 96 | if node != nil { |
| 97 | atomic.AddInt32(&node.refCount, 1) |
| 98 | } |
| 99 | return node |
| 100 | } |
| 101 | |
| 102 | func (node *mapNode) decref() { |
| 103 | if node == nil { |
| 104 | return |
| 105 | } |
| 106 | if atomic.AddInt32(&node.refCount, -1) == 0 { |
| 107 | if atomic.AddInt32(&node.value.refCount, -1) == 0 { |
| 108 | if node.value.release != nil { |
| 109 | node.value.release(node.key, node.value.value) |
| 110 | } |
| 111 | node.value.value = nil |
| 112 | node.value.release = nil |
| 113 | } |
| 114 | node.left.decref() |
| 115 | node.right.decref() |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | // NewMap returns a new map whose keys are ordered by the given comparison |
| 120 | // function (a strict weak order). It is the responsibility of the caller to |
| 121 | // Destroy it at later time. |
| 122 | func NewMap(less func(a, b interface{}) bool) *Map { |
| 123 | return &Map{ |
| 124 | less: less, |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | // Clone returns a copy of the given map. It is a responsibility of the caller |
| 129 | // to Destroy it at later time. |
| 130 | func (pm *Map) Clone() *Map { |
| 131 | return &Map{ |
| 132 | less: pm.less, |
| 133 | root: pm.root.incref(), |
| 134 | } |
| 135 | } |
| 136 | |
| 137 | // Destroy destroys the map. |
| 138 | // |
| 139 | // After Destroy, the Map should not be used again. |
| 140 | func (pm *Map) Destroy() { |
| 141 | // The implementation of these two functions is the same, |
| 142 | // but their intent is different. |
| 143 | pm.Clear() |
| 144 | } |
| 145 | |
| 146 | // Clear removes all entries from the map. |
| 147 | func (pm *Map) Clear() { |
| 148 | pm.root.decref() |
| 149 | pm.root = nil |
| 150 | } |
| 151 | |
| 152 | // Range calls f sequentially in ascending key order for all entries in the map. |
| 153 | func (pm *Map) Range(f func(key, value interface{})) { |
| 154 | pm.root.forEach(f) |
| 155 | } |
| 156 | |
| 157 | func (node *mapNode) forEach(f func(key, value interface{})) { |
| 158 | if node == nil { |
| 159 | return |
| 160 | } |
| 161 | node.left.forEach(f) |
| 162 | f(node.key, node.value.value) |
| 163 | node.right.forEach(f) |
| 164 | } |
| 165 | |
| 166 | // Get returns the map value associated with the specified key, or nil if no entry |
| 167 | // is present. The ok result indicates whether an entry was found in the map. |
| 168 | func (pm *Map) Get(key interface{}) (interface{}, bool) { |
| 169 | node := pm.root |
| 170 | for node != nil { |
| 171 | if pm.less(key, node.key) { |
| 172 | node = node.left |
| 173 | } else if pm.less(node.key, key) { |
| 174 | node = node.right |
| 175 | } else { |
| 176 | return node.value.value, true |
| 177 | } |
| 178 | } |
| 179 | return nil, false |
| 180 | } |
| 181 | |
| 182 | // SetAll updates the map with key/value pairs from the other map, overwriting existing keys. |
| 183 | // It is equivalent to calling Set for each entry in the other map but is more efficient. |
| 184 | // Both maps must have the same comparison function, otherwise behavior is undefined. |
| 185 | func (pm *Map) SetAll(other *Map) { |
| 186 | root := pm.root |
| 187 | pm.root = union(root, other.root, pm.less, true) |
| 188 | root.decref() |
| 189 | } |
| 190 | |
| 191 | // Set updates the value associated with the specified key. |
| 192 | // If release is non-nil, it will be called with entry's key and value once the |
| 193 | // key is no longer contained in the map or any clone. |
| 194 | func (pm *Map) Set(key, value interface{}, release func(key, value interface{})) { |
| 195 | first := pm.root |
| 196 | second := newNodeWithRef(key, value, release) |
| 197 | pm.root = union(first, second, pm.less, true) |
| 198 | first.decref() |
| 199 | second.decref() |
| 200 | } |
| 201 | |
| 202 | // union returns a new tree which is a union of first and second one. |
| 203 | // If overwrite is set to true, second one would override a value for any duplicate keys. |
| 204 | // |
| 205 | // union(first:-0, second:-0) (result:+1) |
| 206 | // Union borrows both subtrees without affecting their refcount and returns a |
| 207 | // new reference that the caller is expected to call decref. |
| 208 | func union(first, second *mapNode, less func(a, b interface{}) bool, overwrite bool) *mapNode { |
| 209 | if first == nil { |
| 210 | return second.incref() |
| 211 | } |
| 212 | if second == nil { |
| 213 | return first.incref() |
| 214 | } |
| 215 | |
| 216 | if first.weight < second.weight { |
| 217 | second, first, overwrite = first, second, !overwrite |
| 218 | } |
| 219 | |
| 220 | left, mid, right := split(second, first.key, less, false) |
| 221 | var result *mapNode |
| 222 | if overwrite && mid != nil { |
| 223 | result = mid.shallowCloneWithRef() |
| 224 | } else { |
| 225 | result = first.shallowCloneWithRef() |
| 226 | } |
| 227 | result.weight = first.weight |
| 228 | result.left = union(first.left, left, less, overwrite) |
| 229 | result.right = union(first.right, right, less, overwrite) |
| 230 | left.decref() |
| 231 | mid.decref() |
| 232 | right.decref() |
| 233 | return result |
| 234 | } |
| 235 | |
| 236 | // split the tree midway by the key into three different ones. |
| 237 | // Return three new trees: left with all nodes with smaller than key, mid with |
| 238 | // the node matching the key, right with all nodes larger than key. |
| 239 | // If there are no nodes in one of trees, return nil instead of it. |
| 240 | // If requireMid is set (such as during deletion), then all return arguments |
| 241 | // are nil if mid is not found. |
| 242 | // |
| 243 | // split(n:-0) (left:+1, mid:+1, right:+1) |
| 244 | // Split borrows n without affecting its refcount, and returns three |
| 245 | // new references that that caller is expected to call decref. |
| 246 | func split(n *mapNode, key interface{}, less func(a, b interface{}) bool, requireMid bool) (left, mid, right *mapNode) { |
| 247 | if n == nil { |
| 248 | return nil, nil, nil |
| 249 | } |
| 250 | |
| 251 | if less(n.key, key) { |
| 252 | left, mid, right := split(n.right, key, less, requireMid) |
| 253 | if requireMid && mid == nil { |
| 254 | return nil, nil, nil |
| 255 | } |
| 256 | newN := n.shallowCloneWithRef() |
| 257 | newN.left = n.left.incref() |
| 258 | newN.right = left |
| 259 | return newN, mid, right |
| 260 | } else if less(key, n.key) { |
| 261 | left, mid, right := split(n.left, key, less, requireMid) |
| 262 | if requireMid && mid == nil { |
| 263 | return nil, nil, nil |
| 264 | } |
| 265 | newN := n.shallowCloneWithRef() |
| 266 | newN.left = right |
| 267 | newN.right = n.right.incref() |
| 268 | return left, mid, newN |
| 269 | } |
| 270 | mid = n.shallowCloneWithRef() |
| 271 | return n.left.incref(), mid, n.right.incref() |
| 272 | } |
| 273 | |
| 274 | // Delete deletes the value for a key. |
| 275 | func (pm *Map) Delete(key interface{}) { |
| 276 | root := pm.root |
| 277 | left, mid, right := split(root, key, pm.less, true) |
| 278 | if mid == nil { |
| 279 | return |
| 280 | } |
| 281 | pm.root = merge(left, right) |
| 282 | left.decref() |
| 283 | mid.decref() |
| 284 | right.decref() |
| 285 | root.decref() |
| 286 | } |
| 287 | |
| 288 | // merge two trees while preserving the weight invariant. |
| 289 | // All nodes in left must have smaller keys than any node in right. |
| 290 | // |
| 291 | // merge(left:-0, right:-0) (result:+1) |
| 292 | // Merge borrows its arguments without affecting their refcount |
| 293 | // and returns a new reference that the caller is expected to call decref. |
| 294 | func merge(left, right *mapNode) *mapNode { |
| 295 | switch { |
| 296 | case left == nil: |
| 297 | return right.incref() |
| 298 | case right == nil: |
| 299 | return left.incref() |
| 300 | case left.weight > right.weight: |
| 301 | root := left.shallowCloneWithRef() |
| 302 | root.left = left.left.incref() |
| 303 | root.right = merge(left.right, right) |
| 304 | return root |
| 305 | default: |
| 306 | root := right.shallowCloneWithRef() |
| 307 | root.left = merge(left, right.left) |
| 308 | root.right = right.right.incref() |
| 309 | return root |
| 310 | } |
| 311 | } |
| 312 |
Members