gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/trie/trie.go (about)

     1  // Copyright 2022 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package trie provides a character-based prefix trie data structure for storing arbitrary payloads
    16  // in an efficiently retrievable manner.
    17  package trie
    18  
    19  // Visitor accepts a prefix string and an associated value, and returns true iff searching should
    20  // continue deeper into the Trie. It is used by FindMatching().
    21  type Visitor func(prefix string, value any) bool
    22  
    23  // Trie stores data at given strings in tree structure, for linear-time retrieval.
    24  // Call New() to obtain a valid Trie.
    25  type Trie struct {
    26  	root *node
    27  	size int
    28  }
    29  
    30  // New creates a new instance of the Trie interface.
    31  func New() *Trie {
    32  	return &Trie{root: &node{children: make(map[rune]*node)}, size: 0}
    33  }
    34  
    35  type node struct {
    36  	value    any
    37  	children map[rune]*node
    38  }
    39  
    40  // FindPrefixes invokes the Visitor with all key-value pairs where the key is a prefix of `key`,
    41  // including exact matches. It does this in increasing order of key length, and terminates early if
    42  // Visitor returns false.
    43  func (t *Trie) FindPrefixes(key string, f Visitor) {
    44  	cur := t.root
    45  	if cur.value != nil && !f("", cur.value) {
    46  		return
    47  	}
    48  
    49  	for i, r := range key {
    50  		next, ok := cur.children[r]
    51  		if !ok {
    52  			return
    53  		}
    54  
    55  		if next.value != nil && !f(key[:(i+1)], next.value) {
    56  			return
    57  		}
    58  		cur = next
    59  	}
    60  }
    61  
    62  func (t *Trie) updateNode(n *node, newValue any) {
    63  	if n.value != nil {
    64  		t.size--
    65  	}
    66  	if newValue != nil {
    67  		t.size++
    68  	}
    69  	n.value = newValue
    70  }
    71  
    72  // SetValue associates the specified key with the given value, replacing any existing value.
    73  func (t *Trie) SetValue(key string, value any) {
    74  	cur := t.root
    75  	for _, r := range key {
    76  		next, ok := cur.children[r]
    77  		if !ok {
    78  			next = &node{children: make(map[rune]*node)}
    79  			cur.children[r] = next
    80  		}
    81  		cur = next
    82  	}
    83  
    84  	if cur.value != nil {
    85  		t.size--
    86  	}
    87  	if value != nil {
    88  		t.size++
    89  	}
    90  	cur.value = value
    91  }
    92  
    93  type queueEntry struct {
    94  	key   string
    95  	value *node
    96  }
    97  
    98  // FindSuffixes invokes the Visitor with all key-value pairs where the key is prefixed by `key`,
    99  // including exact matches. It does this in an unspecified order, and terminates early if the
   100  // Visitor returns false.
   101  //
   102  // Invoking FindSuffixes with the empty string as a key will iterate over all values.
   103  func (t *Trie) FindSuffixes(key string, f Visitor) {
   104  	cur := t.root
   105  	for _, r := range key {
   106  		next, ok := cur.children[r]
   107  		if !ok {
   108  			return
   109  		}
   110  		cur = next
   111  	}
   112  
   113  	queue := make([]queueEntry, 0)
   114  	queue = append(queue, queueEntry{key: key, value: cur})
   115  
   116  	for len(queue) > 0 {
   117  		cur := queue[0]
   118  		queue = queue[1:]
   119  
   120  		if cur.value.value != nil && !f(cur.key, cur.value.value) {
   121  			return
   122  		}
   123  
   124  		for r, v := range cur.value.children {
   125  			queue = append(queue, queueEntry{key: cur.key + string(r), value: v})
   126  		}
   127  	}
   128  }
   129  
   130  // Size returns the total number of values in the Trie.
   131  func (t *Trie) Size() int {
   132  	return t.size
   133  }