blob: 08f14c6793df868ed452ef95dbea5f65bf864f27 [file] [log] [blame]
// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package trie
// Collision functions combine a left and right hand side (lhs and rhs) values
// the two values are associated with the same key and produces the value that
// will be stored for the key.
//
// Collision functions must be idempotent:
//
// collision(x, x) == x for all x.
//
// Collisions functions may be applied whenever a value is inserted
// or two maps are merged, or intersected.
type Collision func(lhs interface{}, rhs interface{}) interface{}
// TakeLhs always returns the left value in a collision.
func TakeLhs(lhs, rhs interface{}) interface{} { return lhs }
// TakeRhs always returns the right hand side in a collision.
func TakeRhs(lhs, rhs interface{}) interface{} { return rhs }
// Builder creates new Map. Each Builder has a unique Scope.
//
// IMPORTANT: Nodes are hash-consed internally to reduce memory consumption. To
// support hash-consing Builders keep an internal Map of all of the Maps that they
// have created. To GC any of the Maps created by the Builder, all references to
// the Builder must be dropped. This includes MutMaps.
type Builder struct {
scope Scope
// hash-consing maps for each node type.
empty *empty
leaves map[leaf]*leaf
branches map[branch]*branch
// It may be possible to support more types of patricia tries
// (e.g. non-hash-consed) by making Builder an interface and abstracting
// the mkLeaf and mkBranch functions.
}
// NewBuilder creates a new Builder with a unique Scope.
func NewBuilder() *Builder {
s := newScope()
return &Builder{
scope: s,
empty: &empty{s},
leaves: make(map[leaf]*leaf),
branches: make(map[branch]*branch),
}
}
func (b *Builder) Scope() Scope { return b.scope }
// Rescope changes the builder's scope to a new unique Scope.
//
// Any Maps created using the previous scope need to be Cloned
// before any operation.
//
// This makes the old internals of the Builder eligible to be GC'ed.
func (b *Builder) Rescope() {
s := newScope()
b.scope = s
b.empty = &empty{s}
b.leaves = make(map[leaf]*leaf)
b.branches = make(map[branch]*branch)
}
// Empty is the empty map.
func (b *Builder) Empty() Map { return Map{b.Scope(), b.empty} }
// InsertWith inserts a new association from k to v into the Map m to create a new map
// in the current scope and handle collisions using the collision function c.
//
// This is roughly corresponds to updating a map[uint64]interface{} by:
//
// if _, ok := m[k]; ok { m[k] = c(m[k], v} else { m[k] = v}
//
// An insertion or update happened whenever Insert(m, ...) != m .
func (b *Builder) InsertWith(c Collision, m Map, k uint64, v interface{}) Map {
m = b.Clone(m)
return Map{b.Scope(), b.insert(c, m.n, b.mkLeaf(key(k), v), false)}
}
// Inserts a new association from key to value into the Map m to create
// a new map in the current scope.
//
// If there was a previous value mapped by key, keep the previously mapped value.
// This is roughly corresponds to updating a map[uint64]interface{} by:
//
// if _, ok := m[k]; ok { m[k] = val }
//
// This is equivalent to b.Merge(m, b.Create({k: v})).
func (b *Builder) Insert(m Map, k uint64, v interface{}) Map {
return b.InsertWith(TakeLhs, m, k, v)
}
// Updates a (key, value) in the map. This is roughly corresponds to
// updating a map[uint64]interface{} by:
//
// m[key] = val
func (b *Builder) Update(m Map, key uint64, val interface{}) Map {
return b.InsertWith(TakeRhs, m, key, val)
}
// Merge two maps lhs and rhs to create a new map in the current scope.
//
// Whenever there is a key in both maps (a collision), the resulting value mapped by
// the key will be `c(lhs[key], rhs[key])`.
func (b *Builder) MergeWith(c Collision, lhs, rhs Map) Map {
lhs, rhs = b.Clone(lhs), b.Clone(rhs)
return Map{b.Scope(), b.merge(c, lhs.n, rhs.n)}
}
// Merge two maps lhs and rhs to create a new map in the current scope.
//
// Whenever there is a key in both maps (a collision), the resulting value mapped by
// the key will be the value in lhs `b.Collision(lhs[key], rhs[key])`.
func (b *Builder) Merge(lhs, rhs Map) Map {
return b.MergeWith(TakeLhs, lhs, rhs)
}
// Clone returns a Map that contains the same (key, value) elements
// within b.Scope(), i.e. return m if m.Scope() == b.Scope() or return
// a deep copy of m within b.Scope() otherwise.
func (b *Builder) Clone(m Map) Map {
if m.Scope() == b.Scope() {
return m
} else if m.n == nil {
return Map{b.Scope(), b.empty}
}
return Map{b.Scope(), b.clone(m.n)}
}
func (b *Builder) clone(n node) node {
switch n := n.(type) {
case *empty:
return b.empty
case *leaf:
return b.mkLeaf(n.k, n.v)
case *branch:
return b.mkBranch(n.prefix, n.branching, b.clone(n.left), b.clone(n.right))
default:
panic("unreachable")
}
}
// Remove a key from a Map m and return the resulting Map.
func (b *Builder) Remove(m Map, k uint64) Map {
m = b.Clone(m)
return Map{b.Scope(), b.remove(m.n, key(k))}
}
// Intersect Maps lhs and rhs and returns a map with all of the keys in
// both lhs and rhs and the value comes from lhs, i.e.
//
// {(k, lhs[k]) | k in lhs, k in rhs}.
func (b *Builder) Intersect(lhs, rhs Map) Map {
return b.IntersectWith(TakeLhs, lhs, rhs)
}
// IntersectWith take lhs and rhs and returns the intersection
// with the value coming from the collision function, i.e.
//
// {(k, c(lhs[k], rhs[k]) ) | k in lhs, k in rhs}.
//
// The elements of the resulting map are always { <k, c(lhs[k], rhs[k]) > }
// for each key k that a key in both lhs and rhs.
func (b *Builder) IntersectWith(c Collision, lhs, rhs Map) Map {
l, r := b.Clone(lhs), b.Clone(rhs)
return Map{b.Scope(), b.intersect(c, l.n, r.n)}
}
// MutMap is a convenient wrapper for a Map and a *Builder that will be used to create
// new Maps from it.
type MutMap struct {
B *Builder
M Map
}
// MutEmpty is an empty MutMap for a builder.
func (b *Builder) MutEmpty() MutMap {
return MutMap{b, b.Empty()}
}
// Insert an element into the map using the collision function for the builder.
// Returns true if the element was inserted.
func (mm *MutMap) Insert(k uint64, v interface{}) bool {
old := mm.M
mm.M = mm.B.Insert(old, k, v)
return old != mm.M
}
// Updates an element in the map. Returns true if the map was updated.
func (mm *MutMap) Update(k uint64, v interface{}) bool {
old := mm.M
mm.M = mm.B.Update(old, k, v)
return old != mm.M
}
// Removes a key from the map. Returns true if the element was removed.
func (mm *MutMap) Remove(k uint64) bool {
old := mm.M
mm.M = mm.B.Remove(old, k)
return old != mm.M
}
// Merge another map into the current one using the collision function
// for the builder. Returns true if the map changed.
func (mm *MutMap) Merge(other Map) bool {
old := mm.M
mm.M = mm.B.Merge(old, other)
return old != mm.M
}
// Intersect another map into the current one using the collision function
// for the builder. Returns true if the map changed.
func (mm *MutMap) Intersect(other Map) bool {
old := mm.M
mm.M = mm.B.Intersect(old, other)
return old != mm.M
}
func (b *Builder) Create(m map[uint64]interface{}) Map {
var leaves []*leaf
for k, v := range m {
leaves = append(leaves, b.mkLeaf(key(k), v))
}
return Map{b.Scope(), b.create(leaves)}
}
// Merge another map into the current one using the collision function
// for the builder. Returns true if the map changed.
func (mm *MutMap) MergeWith(c Collision, other Map) bool {
old := mm.M
mm.M = mm.B.MergeWith(c, old, other)
return old != mm.M
}
// creates a map for a collection of leaf nodes.
func (b *Builder) create(leaves []*leaf) node {
n := len(leaves)
if n == 0 {
return b.empty
} else if n == 1 {
return leaves[0]
}
// Note: we can do a more sophisicated algorithm by:
// - sorting the leaves ahead of time,
// - taking the prefix and branching bit of the min and max key,
// - binary searching for the branching bit,
// - splitting exactly where the branch will be, and
// - making the branch node for this prefix + branching bit.
// Skipping until this is a performance bottleneck.
m := n / 2 // (n >= 2) ==> 1 <= m < n
l, r := leaves[:m], leaves[m:]
return b.merge(nil, b.create(l), b.create(r))
}
// mkLeaf returns the hash-consed representative of (k, v) in the current scope.
func (b *Builder) mkLeaf(k key, v interface{}) *leaf {
l := &leaf{k: k, v: v}
if rep, ok := b.leaves[*l]; ok {
return rep
}
b.leaves[*l] = l
return l
}
// mkBranch returns the hash-consed representative of the tuple
//
// (prefix, branch, left, right)
//
// in the current scope.
func (b *Builder) mkBranch(p prefix, bp bitpos, left node, right node) *branch {
br := &branch{
sz: left.size() + right.size(),
prefix: p,
branching: bp,
left: left,
right: right,
}
if rep, ok := b.branches[*br]; ok {
return rep
}
b.branches[*br] = br
return br
}
// join two maps with prefixes p0 and p1 that are *known* to disagree.
func (b *Builder) join(p0 prefix, t0 node, p1 prefix, t1 node) *branch {
m := branchingBit(p0, p1)
var left, right node
if zeroBit(p0, m) {
left, right = t0, t1
} else {
left, right = t1, t0
}
prefix := mask(p0, m)
return b.mkBranch(prefix, m, left, right)
}
// collide two leaves with the same key to create a leaf
// with the collided value.
func (b *Builder) collide(c Collision, left, right *leaf) *leaf {
if left == right {
return left // c is idempotent: c(x, x) == x
}
val := left.v // keep the left value by default if c is nil
if c != nil {
val = c(left.v, right.v)
}
switch val {
case left.v:
return left
case right.v:
return right
default:
return b.mkLeaf(left.k, val)
}
}
// inserts a leaf l into a map m and returns the resulting map.
// When lhs is true, l is the left hand side in a collision.
// Both l and m are in the current scope.
func (b *Builder) insert(c Collision, m node, l *leaf, lhs bool) node {
switch m := m.(type) {
case *empty:
return l
case *leaf:
if m.k == l.k {
left, right := l, m
if !lhs {
left, right = right, left
}
return b.collide(c, left, right)
}
return b.join(prefix(l.k), l, prefix(m.k), m)
case *branch:
// fallthrough
}
// m is a branch
br := m.(*branch)
if !matchPrefix(prefix(l.k), br.prefix, br.branching) {
return b.join(prefix(l.k), l, br.prefix, br)
}
var left, right node
if zeroBit(prefix(l.k), br.branching) {
left, right = b.insert(c, br.left, l, lhs), br.right
} else {
left, right = br.left, b.insert(c, br.right, l, lhs)
}
if left == br.left && right == br.right {
return m
}
return b.mkBranch(br.prefix, br.branching, left, right)
}
// merge two maps in the current scope.
func (b *Builder) merge(c Collision, lhs, rhs node) node {
if lhs == rhs {
return lhs
}
switch lhs := lhs.(type) {
case *empty:
return rhs
case *leaf:
return b.insert(c, rhs, lhs, true)
case *branch:
switch rhs := rhs.(type) {
case *empty:
return lhs
case *leaf:
return b.insert(c, lhs, rhs, false)
case *branch:
// fallthrough
}
}
// Last remaining case is branch merging.
// For brevity, we adopt the Okasaki and Gill naming conventions
// for branching and prefixes.
s, t := lhs.(*branch), rhs.(*branch)
p, m := s.prefix, s.branching
q, n := t.prefix, t.branching
if m == n && p == q { // prefixes are identical.
left, right := b.merge(c, s.left, t.left), b.merge(c, s.right, t.right)
return b.mkBranch(p, m, left, right)
}
if !prefixesOverlap(p, m, q, n) {
return b.join(p, s, q, t) // prefixes are disjoint.
}
// prefixesOverlap(p, m, q, n) && !(m ==n && p == q)
// By prefixesOverlap(...), either:
// higher(m, n) && matchPrefix(q, p, m), or
// higher(n, m) && matchPrefix(p, q, n)
// So either s or t may can be merged with one branch or the other.
switch {
case ord(m, n) && zeroBit(q, m):
return b.mkBranch(p, m, b.merge(c, s.left, t), s.right)
case ord(m, n) && !zeroBit(q, m):
return b.mkBranch(p, m, s.left, b.merge(c, s.right, t))
case ord(n, m) && zeroBit(p, n):
return b.mkBranch(q, n, b.merge(c, s, t.left), t.right)
default:
return b.mkBranch(q, n, t.left, b.merge(c, s, t.right))
}
}
func (b *Builder) remove(m node, k key) node {
switch m := m.(type) {
case *empty:
return m
case *leaf:
if m.k == k {
return b.empty
}
return m
case *branch:
// fallthrough
}
br := m.(*branch)
kp := prefix(k)
if !matchPrefix(kp, br.prefix, br.branching) {
// The prefix does not match. kp is not in br.
return br
}
// the prefix matches. try to remove from the left or right branch.
left, right := br.left, br.right
if zeroBit(kp, br.branching) {
left = b.remove(left, k) // k may be in the left branch.
} else {
right = b.remove(right, k) // k may be in the right branch.
}
if left == br.left && right == br.right {
return br // no update
} else if _, ok := left.(*empty); ok {
return right // left updated and is empty.
} else if _, ok := right.(*empty); ok {
return left // right updated and is empty.
}
// Either left or right updated. Both left and right are not empty.
// The left and right branches still share the same prefix and disagree
// on the same branching bit. It is safe to directly create the branch.
return b.mkBranch(br.prefix, br.branching, left, right)
}
func (b *Builder) intersect(c Collision, l, r node) node {
if l == r {
return l
}
switch l := l.(type) {
case *empty:
return b.empty
case *leaf:
if rleaf := r.find(l.k); rleaf != nil {
return b.collide(c, l, rleaf)
}
return b.empty
case *branch:
switch r := r.(type) {
case *empty:
return b.empty
case *leaf:
if lleaf := l.find(r.k); lleaf != nil {
return b.collide(c, lleaf, r)
}
return b.empty
case *branch:
// fallthrough
}
}
// Last remaining case is branch intersection.
s, t := l.(*branch), r.(*branch)
p, m := s.prefix, s.branching
q, n := t.prefix, t.branching
if m == n && p == q {
// prefixes are identical.
left, right := b.intersect(c, s.left, t.left), b.intersect(c, s.right, t.right)
if _, ok := left.(*empty); ok {
return right
} else if _, ok := right.(*empty); ok {
return left
}
// The left and right branches are both non-empty.
// They still share the same prefix and disagree on the same branching bit.
// It is safe to directly create the branch.
return b.mkBranch(p, m, left, right)
}
if !prefixesOverlap(p, m, q, n) {
return b.empty // The prefixes share no keys.
}
// prefixesOverlap(p, m, q, n) && !(m ==n && p == q)
// By prefixesOverlap(...), either:
// ord(m, n) && matchPrefix(q, p, m), or
// ord(n, m) && matchPrefix(p, q, n)
// So either s or t may be a strict subtree of the other.
var lhs, rhs node
switch {
case ord(m, n) && zeroBit(q, m):
lhs, rhs = s.left, t
case ord(m, n) && !zeroBit(q, m):
lhs, rhs = s.right, t
case ord(n, m) && zeroBit(p, n):
lhs, rhs = s, t.left
default:
lhs, rhs = s, t.right
}
return b.intersect(c, lhs, rhs)
}