trie.go

     1  // Copyright 2026 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  package types2
     6  
     7  import "sort"
     8  
     9  // A trie[V] maps keys to values V, similar to a map.
    10  // A key is a list of integers; two keys collide if one of them is a prefix of the other.
    11  // For instance, [1, 2] and [1, 2, 3] collide, but [1, 2, 3] and [1, 2, 4] do not.
    12  // If all keys have length 1, a trie degenerates into an ordinary map[int]V.
    13  type trie[V any] map[int]any // map value is either trie[V] (non-leaf node), or V (leaf node)
    14  
    15  // insert inserts a value with a key into the trie; key must not be empty.
    16  // If key doesn't collide with any other key in the trie, insert succeeds
    17  // and returns (val, 0). Otherwise, insert fails and returns (alt, n) where
    18  // alt is an unspecified but deterministically chosen value with a colliding
    19  // key and n is the length > 0 of the common key prefix. The trie is not
    20  // changed in this case.
    21  func (tr trie[V]) insert(key []int, val V) (V, int) {
    22  	// A key serves as the path from the trie root to its corresponding value.
    23  	for l, index := range key {
    24  		if v, exists := tr[index]; exists {
    25  			// An entry already exists for this index; check its type to determine collision.
    26  			switch v := v.(type) {
    27  			case trie[V]:
    28  				// Path continues.
    29  				// Must check this case first in case V is any which would act as catch-all.
    30  				tr = v
    31  			case V:
    32  				// Collision: An existing key ends here.
    33  				// This means the existing key is a prefix of (or exactly equal to) key.
    34  				return v, l + 1
    35  			case nil:
    36  				// Handle esoteric case where V is any and val is nil.
    37  				var zero V
    38  				return zero, l + 1
    39  			default:
    40  				panic("trie.insert: invalid entry")
    41  			}
    42  		} else {
    43  			// Path doesn't exist yet, we need to build it.
    44  			if l == len(key)-1 {
    45  				// No prefix collision detected; insert val as a new leaf node.
    46  				tr[index] = val
    47  				return val, 0
    48  			}
    49  			node := make(trie[V])
    50  			tr[index] = node
    51  			tr = node
    52  		}
    53  	}
    54  
    55  	if len(key) == 0 {
    56  		panic("trie.insert: key must not be empty")
    57  	}
    58  
    59  	// Collision: path ends here, but the trie continues.
    60  	// This means key is a prefix of an existing key.
    61  	// Return a value from the subtrie.
    62  	for len(tr) != 0 {
    63  		// see switch above
    64  		switch v := tr.pickValue().(type) {
    65  		case trie[V]:
    66  			tr = v
    67  		case V:
    68  			return v, len(key)
    69  		case nil:
    70  			var zero V
    71  			return zero, len(key)
    72  		default:
    73  			panic("trie.insert: invalid entry")
    74  		}
    75  	}
    76  
    77  	panic("trie.insert: unreachable")
    78  }
    79  
    80  // pickValue deterministically picks a value of the trie and returns it.
    81  // Only called in case of a collision.
    82  // The trie must not be empty.
    83  func (tr trie[V]) pickValue() any {
    84  	var keys []int
    85  	for k := range tr {
    86  		keys = append(keys, k)
    87  	}
    88  	sort.Ints(keys) // guarantee deterministic element pick
    89  	return tr[keys[0]]
    90  }
    91
View as plain text