Design search autocomplete

package main
 
import (
	"fmt"
	"sort"
)
 
type node struct {
	children map[rune]*node
	word     string
	count    int
	top      []string // best completions in this subtree, computed once
}
 
func build(counts map[string]int, k int) *node {
	root := &node{children: map[rune]*node{}}
	for w, c := range counts {
		n := root
		for _, r := range w {
			if n.children[r] == nil {
				n.children[r] = &node{children: map[rune]*node{}}
			}
			n = n.children[r]
		}
		n.word, n.count = w, c
	}
	var fill func(n *node) []string
	fill = func(n *node) []string {
		var cands []string
		if n.word != "" {
			cands = append(cands, n.word)
		}
		for _, c := range n.children {
			cands = append(cands, fill(c)...)
		}
		sort.Slice(cands, func(i, j int) bool {
			if counts[cands[i]] != counts[cands[j]] {
				return counts[cands[i]] > counts[cands[j]]
			}
			return cands[i] < cands[j]
		})
		n.top = cands[:min(k, len(cands))]
		return n.top
	}
	fill(root)
	return root
}
 
func main() {
	counts := map[string]int{"tree": 50, "trie": 80, "true": 120, "try": 30, "trick": 45, "tea": 60}
	root := build(counts, 3)
	for _, prefix := range []string{"t", "tr", "tri", "x"} {
		n := root
		for _, r := range prefix {
			if n = n.children[r]; n == nil {
				break
			}
		}
		var top []string
		if n != nil {
			top = n.top
		}
		fmt.Printf("%-4s -> %v\n", prefix, top)
	}
}

t    -> [true trie tea]
tr   -> [true trie tree]
tri  -> [trie trick]
x    -> []

typeahead/
├── go.mod
├── main.go                         HTTP API, seed data, periodic rebuild
├── main_test.go
├── internal/suggest/trie.go        immutable trie with top-k per node
├── internal/suggest/service.go     counting, decay, blocklist, atomic swap
├── internal/suggest/suggest_test.go
└── web/index.html                  the search box

mkdir typeahead && cd typeahead && go mod init typeahead
mkdir -p internal/suggest web

module typeahead
 
go 1.24

// Package suggest serves prefix suggestions from a trie rebuilt periodically from search counts.
package suggest
 
import (
	"cmp"
	"slices"
)
 
// Entry is one suggestion and its score.
type Entry struct {
	Term  string `json:"term"`
	Count int    `json:"count"`
}
 
type node struct {
	children map[rune]*node
	entry    *Entry // set when a query ends here
	top      []Entry
}
 
// Trie is immutable after Build.
type Trie struct {
	root  *node
	terms int
}
 
// better orders entries by count descending, then term ascending, so results are deterministic.
func better(a, b Entry) int {
	if c := cmp.Compare(b.Count, a.Count); c != 0 {
		return c
	}
	return cmp.Compare(a.Term, b.Term)
}
 
// Build creates a trie whose every node stores its k best completions.
func Build(counts map[string]int, k int) *Trie {
	t := &Trie{root: &node{children: map[rune]*node{}}}
	for term, c := range counts {
		if term == "" || c <= 0 {
			continue
		}
		n := t.root
		for _, r := range term {
			child := n.children[r]
			if child == nil {
				child = &node{children: map[rune]*node{}}
				n.children[r] = child
			}
			n = child
		}
		n.entry = &Entry{term, c}
		t.terms++
	}
	fillTop(t.root, k)
	return t
}
 
func fillTop(n *node, k int) []Entry {
	var cands []Entry
	if n.entry != nil {
		cands = append(cands, *n.entry)
	}
	for _, c := range n.children {
		cands = append(cands, fillTop(c, k)...) // each child returns at most k
	}
	slices.SortFunc(cands, better)
	n.top = cands[:min(k, len(cands))]
	return n.top
}
 
// Suggest returns up to n completions of prefix, best first.
func (t *Trie) Suggest(prefix string, n int) []Entry {
	node := t.root
	for _, r := range prefix {
		if node = node.children[r]; node == nil {
			return []Entry{}
		}
	}
	return slices.Clone(node.top[:min(n, len(node.top))])
}
 
// Len is the number of distinct terms.
func (t *Trie) Len() int { return t.terms }

package suggest
 
import (
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"unicode/utf8"
)
 
// MaxQuery is the longest query recorded, in characters.
const MaxQuery = 100
 
// Normalize lowercases a query and collapses runs of whitespace into one space.
func Normalize(q string) string {
	return strings.Join(strings.Fields(strings.ToLower(q)), " ")
}
 
// normalizePrefix is Normalize but keeps one trailing space, which is meaningful while typing:
// "new " should suggest "new york", not "newsletter".
func normalizePrefix(q string) string {
	n := Normalize(q)
	if n != "" && strings.HasSuffix(q, " ") {
		n += " "
	}
	return n
}
 
// Service holds the current trie and the counts that feed the next one.
type Service struct {
	k       int
	decay   float64
	blocked map[string]bool
	trie    atomic.Pointer[Trie]
 
	mu     sync.Mutex
	window map[string]int     // searches since the last rebuild
	scores map[string]float64 // decayed history
}
 
// NewService keeps k suggestions per prefix. Each rebuild multiplies old scores by decay.
func NewService(k int, decay float64, blocked []string) *Service {
	s := &Service{k: k, decay: decay, blocked: map[string]bool{}, window: map[string]int{}, scores: map[string]float64{}}
	for _, b := range blocked {
		s.blocked[Normalize(b)] = true
	}
	s.trie.Store(Build(nil, k))
	return s
}
 
// Seed adds historical counts, for example from yesterday's logs. Like Record, they
// take effect at the next Rebuild.
func (s *Service) Seed(counts map[string]int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for q, c := range counts {
		s.window[Normalize(q)] += c
	}
}
 
// Record counts one completed search. It is not visible until the next Rebuild.
func (s *Service) Record(q string) bool {
	q = Normalize(q)
	if q == "" || utf8.RuneCountInString(q) > MaxQuery {
		return false
	}
	s.mu.Lock()
	s.window[q]++
	s.mu.Unlock()
	return true
}
 
// Rebuild folds the window into the decayed scores, builds a new trie and swaps it in.
// It returns the number of terms in the new trie.
func (s *Service) Rebuild() int {
	s.mu.Lock()
	for q, v := range s.scores {
		s.scores[q] = v * s.decay
	}
	for q, c := range s.window {
		s.scores[q] += float64(c)
	}
	s.window = map[string]int{}
	counts := make(map[string]int, len(s.scores))
	for q, v := range s.scores {
		if v < 0.5 {
			delete(s.scores, q) // faded away: forget it
			continue
		}
		if !s.blocked[q] {
			counts[q] = int(math.Round(v))
		}
	}
	s.mu.Unlock()
 
	t := Build(counts, s.k) // the slow part runs without holding the lock
	s.trie.Store(t)
	return t.Len()
}
 
// Suggest answers from the current trie without locking.
func (s *Service) Suggest(prefix string, n int) []Entry {
	p := normalizePrefix(prefix)
	if p == "" {
		return []Entry{}
	}
	return s.trie.Load().Suggest(p, min(n, s.k))
}

package suggest
 
import (
	"slices"
	"sync"
	"testing"
)
 
func terms(es []Entry) []string {
	var out []string
	for _, e := range es {
		out = append(out, e.Term)
	}
	return out
}
 
func TestTrieTopK(t *testing.T) {
	tr := Build(map[string]int{"go": 90, "golang": 100, "gopher": 40, "google": 100, "gin": 10, "café": 7, "cafe": 3}, 3)
	cases := []struct {
		prefix string
		want   []string
	}{
		{"g", []string{"golang", "google", "go"}}, // tie on 100: alphabetical
		{"gop", []string{"gopher"}},
		{"caf", []string{"café", "cafe"}},
		{"café", []string{"café"}},
		{"x", nil},
	}
	for _, c := range cases {
		if got := terms(tr.Suggest(c.prefix, 10)); !slices.Equal(got, c.want) {
			t.Errorf("Suggest(%q) = %v, want %v", c.prefix, got, c.want)
		}
	}
	if got := tr.Suggest("g", 1); len(got) != 1 {
		t.Errorf("n=1 returned %d", len(got))
	}
	if tr.Len() != 7 {
		t.Errorf("Len = %d", tr.Len())
	}
}
 
func TestRecordIsVisibleAfterRebuild(t *testing.T) {
	s := NewService(5, 0.5, nil)
	s.Record("Go  Channels")
	s.Record("go channels")
	if got := s.Suggest("go", 5); len(got) != 0 {
		t.Fatalf("before rebuild: %v", got)
	}
	s.Rebuild()
	got := s.Suggest("GO c", 5)
	if len(got) != 1 || got[0] != (Entry{"go channels", 2}) {
		t.Errorf("after rebuild: %+v", got)
	}
}
 
func TestDecayLetsTrendsWin(t *testing.T) {
	s := NewService(5, 0.5, nil)
	for range 4 {
		s.Record("go")
	}
	s.Rebuild() // go = 4
	s.Rebuild() // go = 2
	for range 3 {
		s.Record("gopher")
	}
	s.Rebuild() // go = 1, gopher = 3
	if got := terms(s.Suggest("go", 5)); !slices.Equal(got, []string{"gopher", "go"}) {
		t.Errorf("after decay = %v", got)
	}
	s.Rebuild() // go = 0.5, still kept
	// go = 0.25: forgotten
	if n := s.Rebuild(); n != 1 {
		t.Errorf("terms after fading = %d, want 1", n)
	}
}
 
func TestBlocklistAndTrailingSpace(t *testing.T) {
	s := NewService(5, 1, []string{"Bad Word"})
	s.Seed(map[string]int{"bad word": 100, "bad weather": 5, "new york": 50, "newsletter": 80})
	s.Rebuild()
	if got := terms(s.Suggest("bad", 5)); !slices.Equal(got, []string{"bad weather"}) {
		t.Errorf("blocked term leaked: %v", got)
	}
	if got := terms(s.Suggest("new ", 5)); !slices.Equal(got, []string{"new york"}) {
		t.Errorf(`"new " = %v`, got)
	}
	if got := terms(s.Suggest("new", 5)); !slices.Equal(got, []string{"newsletter", "new york"}) {
		t.Errorf(`"new" = %v`, got)
	}
}
 
func TestConcurrentReadsDuringRebuild(t *testing.T) {
	s := NewService(5, 0.9, nil)
	s.Seed(map[string]int{"go": 10})
	s.Rebuild()
	var wg sync.WaitGroup
	for range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for range 1000 {
				s.Suggest("g", 5)
				s.Record("golang")
			}
		}()
	}
	for range 20 {
		s.Rebuild()
	}
	wg.Wait()
	s.Rebuild()
	if got := terms(s.Suggest("go", 5)); len(got) != 2 {
		t.Errorf("after concurrent load: %v", got)
	}
}

package main
 
import (
	_ "embed"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"time"
 
	"typeahead/internal/suggest"
)
 
//go:embed web/index.html
var indexHTML []byte
 
var seedCounts = map[string]int{
	"golang tutorial": 900, "golang generics": 520, "golang channels": 610, "golang context": 480,
	"goroutine leak": 210, "go modules": 700, "go test coverage": 260, "go vs rust": 650,
	"google maps": 2000, "google translate": 1800, "gopher": 120, "graphql": 400,
	"grpc go": 310, "git rebase": 880, "git stash": 640, "github actions": 950,
	"redis": 820, "redis cluster": 300, "rate limiter": 410, "react hooks": 1200,
	"rust ownership": 380, "system design interview": 1500, "sql join": 990, "sliding window": 450,
	"consistent hashing": 520, "kafka": 760, "kubernetes": 1400, "new york weather": 1100, "news": 1600,
}
 
func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(v)
}
 
func newServer(s *suggest.Service) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /{$}", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write(indexHTML)
	})
	mux.HandleFunc("GET /api/suggest", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query().Get("q")
		w.Header().Set("Cache-Control", "public, max-age=30")
		writeJSON(w, http.StatusOK, map[string]any{"q": q, "suggestions": s.Suggest(q, 8)})
	})
	mux.HandleFunc("POST /api/search", func(w http.ResponseWriter, r *http.Request) {
		var in struct{ Q string }
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<12)).Decode(&in) != nil || !s.Record(in.Q) {
			writeJSON(w, http.StatusBadRequest, map[string]string{"error": "send {\"q\": \"...\"} with 1 to 100 characters"})
			return
		}
		w.WriteHeader(http.StatusNoContent)
	})
	return mux
}
 
func main() {
	s := suggest.NewService(8, 0.99, []string{"some blocked phrase"})
	s.Seed(seedCounts)
	s.Rebuild()
	go func() {
		for range time.Tick(10 * time.Second) { // production: every few minutes to hours
			n := s.Rebuild()
			slog.Info("rebuilt trie", "terms", n)
		}
	}()
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	slog.Info("listening", "port", port)
	srv := &http.Server{Addr: ":" + port, Handler: newServer(s), ReadHeaderTimeout: 5 * time.Second}
	if err := srv.ListenAndServe(); err != nil {
		slog.Error("server", "err", err)
		os.Exit(1)
	}
}

package main
 
import (
	"encoding/json"
	"net/http/httptest"
	"strings"
	"testing"
 
	"typeahead/internal/suggest"
)
 
func TestAPI(t *testing.T) {
	s := suggest.NewService(8, 0.9, nil)
	s.Seed(seedCounts)
	s.Rebuild()
	h := newServer(s)
 
	get := func(q string) []suggest.Entry {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest("GET", "/api/suggest?q="+q, nil))
		if cc := rec.Header().Get("Cache-Control"); cc != "public, max-age=30" {
			t.Errorf("Cache-Control = %q", cc)
		}
		var out struct{ Suggestions []suggest.Entry }
		json.NewDecoder(rec.Body).Decode(&out)
		return out.Suggestions
	}
	if got := get("gola"); len(got) != 4 || got[0].Term != "golang tutorial" {
		t.Errorf("gola = %+v", got)
	}
	if got := get("zzz"); got == nil || len(got) != 0 {
		t.Errorf("no match should be an empty list, got %#v", got)
	}
 
	post := func(body string) int {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest("POST", "/api/search", strings.NewReader(body)))
		return rec.Code
	}
	for range 3 {
		if code := post(`{"q":"golang iterators"}`); code != 204 {
			t.Fatalf("search: %d", code)
		}
	}
	if code := post(`{"q":"   "}`); code != 400 {
		t.Errorf("blank search: %d", code)
	}
	s.Rebuild()
	found := false
	for _, e := range get("golang+i") {
		found = found || e.Term == "golang iterators"
	}
	if !found {
		t.Error("a recorded search should be suggested after a rebuild")
	}
}

<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>Typeahead</title>
  <style>
    body { font-family: system-ui, sans-serif; max-width: 620px; margin: 4rem auto; padding: 0 1rem; }
    #box { position: relative; }
    #q { width: 100%; box-sizing: border-box; padding: .75rem 1rem; font-size: 18px; border: 1px solid #999; border-radius: 12px; }
    ul { list-style: none; margin: .25rem 0 0; padding: .25rem; border: 1px solid #ddd; border-radius: 12px; }
    ul:empty { display: none; }
    li { display: flex; justify-content: space-between; padding: .45rem .75rem; border-radius: 8px; cursor: pointer; }
    li.on, li:hover { background: #eee; }
    li small { color: #888; font-family: monospace; }
    #info { color: #666; font-size: 14px; margin-top: 1.5rem; }
  </style>
</head>
<body>
  <h1>Search</h1>
  <div id="box">
    <input id="q" autocomplete="off" placeholder="Try go, git, red, new…" autofocus>
    <ul id="list"></ul>
  </div>
  <p id="info">Requests sent: <b id="sent">0</b>. Searches you run are counted and show up after the next rebuild (every 10 s).</p>
  <script>
    const cache = new Map();
    let timer = null, latest = 0, active = -1, items = [];
    function render(prefix, suggestions) {
      items = suggestions; active = -1;
      list.innerHTML = "";
      for (const s of suggestions) {
        const li = document.createElement("li");
        const term = document.createElement("span");
        const b = document.createElement("b");
        b.textContent = s.term.slice(0, prefix.length);
        term.append(b, document.createTextNode(s.term.slice(prefix.length)));
        const n = document.createElement("small");
        n.textContent = s.count;
        li.append(term, n);
        li.onmousedown = () => search(s.term);
        list.append(li);
      }
    }
    async function fetchSuggestions(prefix) {
      if (cache.has(prefix)) return render(prefix, cache.get(prefix));
      const id = ++latest;
      sent.textContent = latest;
      const data = await (await fetch(`/api/suggest?q=${encodeURIComponent(prefix)}`)).json();
      cache.set(prefix, data.suggestions);
      if (id === latest && q.value.toLowerCase() === prefix) render(prefix, data.suggestions);
    }
    q.oninput = () => {
      clearTimeout(timer);
      const prefix = q.value.toLowerCase();
      if (!prefix.trim()) return render("", []);
      timer = setTimeout(() => fetchSuggestions(prefix), 100);
    };
    q.onkeydown = (e) => {
      if (e.key === "ArrowDown" || e.key === "ArrowUp") {
        e.preventDefault();
        if (!items.length) return;
        active = (active + (e.key === "ArrowDown" ? 1 : -1) + items.length) % items.length;
        [...list.children].forEach((li, i) => li.classList.toggle("on", i === active));
        q.value = items[active].term;
      } else if (e.key === "Enter") {
        search(q.value);
      } else if (e.key === "Escape") {
        render("", []);
      }
    };
    async function search(term) {
      q.value = term;
      render("", []);
      await fetch("/api/search", { method: "POST", body: JSON.stringify({ q: term }) });
      cache.clear(); // our own search may change suggestions after the rebuild
    }
  </script>
</body>
</html>

go test -race ./...
go run .
# open http://localhost:8080 and type "go"
curl -s "localhost:8080/api/suggest?q=git"
# {"q":"git","suggestions":[{"term":"github actions","count":950},{"term":"git rebase","count":880},{"term":"git stash","count":640}]}