trie.go

Documentation: golang.org/x/tools/go/callgraph/vta/internal/trie

     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) (interface{}, 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 interface{}) 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, interface{}) 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]interface{} {
   104  	dest := make(map[uint64]interface{}, m.Size())
   105  	m.Range(func(k uint64, v interface{}) 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, interface{}) 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 interface{}
   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, interface{}) bool) bool {
   219  	return true
   220  }
   221  func (l *leaf) visit(cb func(uint64, interface{}) bool) bool {
   222  	return cb(uint64(l.k), l.v)
   223  }
   224  func (br *branch) visit(cb func(uint64, interface{}) bool) bool {
   225  	if !br.left.visit(cb) {
   226  		return false
   227  	}
   228  	return br.right.visit(cb)
   229  }
   230
View as plain text