github.com/cilium/statedb@v0.3.2/part/registry.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package part
     5  
     6  import (
     7  	"encoding/binary"
     8  	"fmt"
     9  	"math"
    10  	"reflect"
    11  	"sync"
    12  	"unicode/utf8"
    13  )
    14  
    15  // keyTypeRegistry is a registry of functions to convert to/from keys (of type K).
    16  // This mechanism enables use of zero value and JSON marshalling and unmarshalling
    17  // with Map and Set.
    18  var keyTypeRegistry sync.Map // map[reflect.Type]func(K) []byte
    19  
    20  // RegisterKeyType registers a new key type to be used with the Map and Set types.
    21  // Intended to be called from init() functions.
    22  // For Set-only usage only the [bytesFromKey] function is needed.
    23  func RegisterKeyType[K any](bytesFromKey func(K) []byte) {
    24  	keyType := reflect.TypeFor[K]()
    25  	keyTypeRegistry.Store(
    26  		keyType,
    27  		bytesFromKey,
    28  	)
    29  }
    30  
    31  func lookupKeyType[K any]() func(K) []byte {
    32  	keyType := reflect.TypeFor[K]()
    33  	funcAny, ok := keyTypeRegistry.Load(keyType)
    34  	if !ok {
    35  		panic(fmt.Sprintf("Key type %q not registered with part.RegisterMapKeyType()", keyType))
    36  	}
    37  	return funcAny.(func(K) []byte)
    38  }
    39  
    40  func init() {
    41  	// Register common key types.
    42  	RegisterKeyType[string](func(s string) []byte { return []byte(s) })
    43  	RegisterKeyType[[]byte](func(b []byte) []byte { return b })
    44  	RegisterKeyType[byte](func(b byte) []byte { return []byte{b} })
    45  	RegisterKeyType[rune](func(r rune) []byte { return utf8.AppendRune(nil, r) })
    46  	RegisterKeyType[complex128](func(c complex128) []byte {
    47  		buf := make([]byte, 0, 16)
    48  		buf = binary.BigEndian.AppendUint64(buf, math.Float64bits(real(c)))
    49  		buf = binary.BigEndian.AppendUint64(buf, math.Float64bits(imag(c)))
    50  		return buf
    51  	})
    52  	RegisterKeyType[float64](func(x float64) []byte { return binary.BigEndian.AppendUint64(nil, math.Float64bits(x)) })
    53  	RegisterKeyType[float32](func(x float32) []byte { return binary.BigEndian.AppendUint32(nil, math.Float32bits(x)) })
    54  	RegisterKeyType[uint64](func(x uint64) []byte { return binary.BigEndian.AppendUint64(nil, x) })
    55  	RegisterKeyType[uint32](func(x uint32) []byte { return binary.BigEndian.AppendUint32(nil, x) })
    56  	RegisterKeyType[uint16](func(x uint16) []byte { return binary.BigEndian.AppendUint16(nil, x) })
    57  	RegisterKeyType[int64](func(x int64) []byte { return binary.BigEndian.AppendUint64(nil, uint64(x)) })
    58  	RegisterKeyType[int32](func(x int32) []byte { return binary.BigEndian.AppendUint32(nil, uint32(x)) })
    59  	RegisterKeyType[int16](func(x int16) []byte { return binary.BigEndian.AppendUint16(nil, uint16(x)) })
    60  	RegisterKeyType[int](func(x int) []byte { return binary.BigEndian.AppendUint64(nil, uint64(x)) })
    61  
    62  	var (
    63  		trueBytes  = []byte{'T'}
    64  		falseBytes = []byte{'F'}
    65  	)
    66  	RegisterKeyType[bool](func(b bool) []byte {
    67  		if b {
    68  			return trueBytes
    69  		} else {
    70  			return falseBytes
    71  		}
    72  	})
    73  
    74  }