builder.go

package prefixtree

const (
	head = byte(0)
	tail = byte(1)
)

type Leaf struct {
	char   byte
	leaves []Leaf
}

func newLeaf(char byte) *Leaf {
	return &Leaf{
		char:   char,
		leaves: make([]Leaf, 0, 1),
	}
}

func newTailLeaf() *Leaf {
	return &Leaf{
		char:   tail,
		leaves: nil,
	}
}

func (l Leaf) IsTail() bool {
	return getLeaf(l.leaves, tail) != nil
}

func Build(branches [][]byte) Leaf {
	longestBranchIndex := longestArray(branches)
	masterBranch := buildMasterBranch(branches[longestBranchIndex])
	otherBranches := append(branches[:longestBranchIndex], branches[longestBranchIndex+1:]...)

	for _, branch := range otherBranches {
		masterBranch = addBranch(masterBranch, branch)
	}

	return masterBranch
}

/*
buildMasterBranch builds a root leaf from zero
*/
func buildMasterBranch(branch []byte) (leaf Leaf) {
	if len(branch) == 0 {
		return leaf
	}

	rootLeaf := &Leaf{
		char:   head,
		leaves: make([]Leaf, 0, 1),
	}

	currentLeaf := rootLeaf

	for _, char := range branch {
		leaf = *newLeaf(char)
		currentLeaf.leaves = append(currentLeaf.leaves, leaf)
		currentLeaf = &currentLeaf.leaves[len(currentLeaf.leaves)-1]
	}

	currentLeaf.leaves = append(currentLeaf.leaves, *newTailLeaf())

	return *rootLeaf
}

func addBranch(base Leaf, branch []byte) Leaf {
	commonPrefixOffset := -1
	currentLeaf := &base

	for i, char := range branch {
		leaf := getLeaf(currentLeaf.leaves, char)

		if leaf == nil {
			commonPrefixOffset = i
			break
		}

		currentLeaf = leaf
	}

	if commonPrefixOffset == -1 {
		if getLeaf(currentLeaf.leaves, tail) == nil {
			currentLeaf.leaves = append(currentLeaf.leaves, *newTailLeaf())
		}

		return base
	}

	for _, char := range branch[commonPrefixOffset:] {
		leaf := *newLeaf(char)
		currentLeaf.leaves = append(currentLeaf.leaves, leaf)
		currentLeaf = &currentLeaf.leaves[len(currentLeaf.leaves)-1]
	}

	currentLeaf.leaves = append(currentLeaf.leaves, *newTailLeaf())

	return base
}

func getLeaf(leaves []Leaf, char byte) *Leaf {
	for i := 0; i < len(leaves); i++ {
		if leaves[i].char != char {
			continue
		}

		return &leaves[i]
	}

	return nil
}

func longestArray(arrays [][]byte) (index int) {
	var max int

	for i, arr := range arrays {
		if len(arr) > max {
			max = len(arr)
			index = i
		}
	}

	return index
}

builder_benchmark_test.go

package prefixtree

import (
	"strings"
	"testing"
)

var (
	shortTreeBranch = []byte(strings.Repeat("def", 6))
	longTreeBranch  = []byte(strings.Repeat("def", 100))
)

var (
	shortTree = Build(multipleBranchesShort)
	longTree  = Build(multipleBranchesLong)
)

func BenchmarkBuiltInWalker(b *testing.B) {
	b.Run("shortTree", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			walker := NewWalker(shortTree)
			walker.Walk(shortTreeBranch)
		}
	})

	b.Run("longTree", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			walker := NewWalker(longTree)
			walker.Walk(longTreeBranch)
		}
	})
}

func BenchmarkManualWalker(b *testing.B) {
	b.Run("shortTree", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			currentLeaf := shortTree

			for _, char := range shortTreeBranch {
				for _, leaf := range currentLeaf.leaves {
					if leaf.char == char {
						currentLeaf = leaf
						break
					}
				}
			}
		}
	})

	b.Run("longTree", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			currentLeaf := longTree

			for _, char := range longTreeBranch {
				for _, leaf := range currentLeaf.leaves {
					if leaf.char == char {
						currentLeaf = leaf
						break
					}
				}
			}
		}
	})

	b.Run("hashmap", func(b *testing.B) {
		key := strings.Repeat("def", 100)
		hashmap := map[string][]byte{
			key: []byte("ok"),
		}

		for i := 0; i < b.N; i++ {
			val := hashmap[key]
			noop(val)
		}
	})
}

func noop(_ []byte) {}

lazywalker.go

package prefixtree

type Walker struct {
	currentLeaf Leaf
}

func NewWalker(root Leaf) Walker {
	return Walker{currentLeaf: root}
}

func (w *Walker) Walk(data []byte) (leaf Leaf, err bool) {
	for _, char := range data {
		childLeaf := getLeaf(w.currentLeaf.leaves, char)

		if childLeaf == nil {
			return leaf, true
		}

		w.currentLeaf = *childLeaf
	}

	return w.currentLeaf, false
}