String slices as keys

pathmap_test.go

// Ref: https://old.reddit.com/r/golang/comments/1bcfe95/slice_as_map_key_in_go/
package main

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"slices"
	"strings"
	"testing"
)

type Path []string

func (p Path) Hash() uint64 {
	h := uint64(-len(p))
	for _, part := range p {
		for _, b := range []byte(part) {
			h ^= uint64(b)
			h *= 0x100000001b3
		}
		h ^= h >> 32
	}
	return h
}

func (p Path) EncodeText() string {
	var w strings.Builder
	for _, part := range p {
		fmt.Fprintf(&w, "%d%s", len(part), part)
	}
	return w.String()
}

func (p Path) EncodeBinary() string {
	var w strings.Builder

	// Optmization: pre-allocate, ~3x faster overall
	n := 8 * len(p)
	for _, part := range p {
		n += len(part) // TODO: mind overflow
	}
	w.Grow(n)

	for _, part := range p {
		var prefix [8]byte
		binary.LittleEndian.PutUint64(prefix[:], uint64(len(part)))
		w.Write(prefix[:])
		w.WriteString(part)
	}
	return w.String()
}

// PathMap is a 4-ary hash-trie operating like map[Path]int.
type PathMap struct {
	child [4]*PathMap
	path  Path
	value int
}

// Upsert returns a pointer to the map's value for the path key.
func Upsert(m **PathMap, path Path) *int {
	for h := path.Hash(); *m != nil; h <<= 2 {
		if slices.Equal(path, (*m).path) {
			return &(*m).value
		}
		m = &(*m).child[h>>62]
	}
	*m = new(PathMap)
	(*m).path = slices.Clone(path) // NOTE: may not be necessary
	return &(*m).value
}

func BenchmarkMap(b *testing.B) {
	parts := Path{
		"foo", "world", "", "hello", "this is a long test",
		"a", "b", "c", "d",
	}
	rng := rand.New(rand.NewSource(int64(b.N)))
	m := make(map[string]int)
	for i := 0; i < b.N; i++ {
		rng.Shuffle(len(parts), func(i, j int) {
			parts[i], parts[j] = parts[j], parts[i]
		})
		m[parts.EncodeBinary()] = i
	}
}

func BenchmarkPathMap(b *testing.B) {
	parts := []string{
		"foo", "world", "", "hello", "this is a long test",
		"a", "b", "c", "d",
	}
	rng := rand.New(rand.NewSource(int64(b.N)))
	var m *PathMap
	for i := 0; i < b.N; i++ {
		rng.Shuffle(len(parts), func(i, j int) {
			parts[i], parts[j] = parts[j], parts[i]
		})
		*Upsert(&m, parts) = i
	}
}

func BenchmarkText(b *testing.B) {
	parts := Path{"foo", "world", "", "hello", "this is a long test"}
	for i := 0; i < b.N; i++ {
		parts.EncodeText()
	}
}

func BenchmarkBinary(b *testing.B) {
	parts := Path{"foo", "world", "", "hello", "this is a long test"}
	for i := 0; i < b.N; i++ {
		parts.EncodeBinary()
	}
}