perf: reduce memory usage in command suggestions

cobra.go

@@ -189,37 +189,50 @@ func tmpl(text string) *tmplFunc {
 }
 
 // ld compares two strings and returns the levenshtein distance between them.
+//
+// Only a single row of the distance matrix is kept: computing row i needs
+// nothing but row i-1, so the full matrix never has to be materialized. This
+// keeps the allocation at O(min(len(s), len(t))) instead of O(len(s)*len(t)).
 func ld(s, t string, ignoreCase bool) int {
 	if ignoreCase {
 		s = strings.ToLower(s)
 		t = strings.ToLower(t)
 	}
-	d := make([][]int, len(s)+1)
-	for i := range d {
-		d[i] = make([]int, len(t)+1)
-		d[i][0] = i
+	// The distance is symmetric, so iterate with the shorter string along the
+	// row to keep the allocation as small as possible.
+	if len(t) > len(s) {
+		s, t = t, s
 	}
-	for j := range d[0] {
-		d[0][j] = j
+	// row[j] holds d[i][j] once column j of the current row has been written,
+	// and d[i-1][j] until then.
+	row := make([]int, len(t)+1)
+	for j := range row {
+		row[j] = j
 	}
-	for j := 1; j <= len(t); j++ {
-		for i := 1; i <= len(s); i++ {
+	var diag, above int
+	for i := 1; i <= len(s); i++ {
+		// diag carries d[i-1][j-1], which row[j-1] is about to overwrite.
+		diag = row[0]
+		row[0] = i
+		for j := 1; j <= len(t); j++ {
+			// d[i-1][j], the next iteration's diagonal
+			above = row[j]
 			if s[i-1] == t[j-1] {
-				d[i][j] = d[i-1][j-1]
+				row[j] = diag
 			} else {
-				min := d[i-1][j]
-				if d[i][j-1] < min {
-					min = d[i][j-1]
+				min := above
+				if row[j-1] < min {
+					min = row[j-1]
 				}
-				if d[i-1][j-1] < min {
-					min = d[i-1][j-1]
+				if diag < min {
+					min = diag
 				}
-				d[i][j] = min + 1
+				row[j] = min + 1
 			}
+			diag = above
 		}
-
 	}
-	return d[len(s)][len(t)]
+	return row[len(t)]
 }
 
 func stringInSlice(a string, list []string) bool {

cobra_test.go

@@ -15,6 +15,7 @@
 package cobra
 
 import (
+	"math/rand"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -299,3 +300,223 @@ func main() {
 		t.Error("compiled programs contains MethodByName symbol")
 	}
 }
+
+// ldReference is the straightforward full-matrix levenshtein implementation.
+// It is intentionally kept as dumb as possible so that it can be trusted as
+// the oracle for the single-row implementation used by ld.
+func ldReference(s, t string, ignoreCase bool) int {
+	if ignoreCase {
+		s = strings.ToLower(s)
+		t = strings.ToLower(t)
+	}
+	d := make([][]int, len(s)+1)
+	for i := range d {
+		d[i] = make([]int, len(t)+1)
+		d[i][0] = i
+	}
+	for j := range d[0] {
+		d[0][j] = j
+	}
+	for j := 1; j <= len(t); j++ {
+		for i := 1; i <= len(s); i++ {
+			if s[i-1] == t[j-1] {
+				d[i][j] = d[i-1][j-1]
+			} else {
+				min := d[i-1][j]
+				if d[i][j-1] < min {
+					min = d[i][j-1]
+				}
+				if d[i-1][j-1] < min {
+					min = d[i-1][j-1]
+				}
+				d[i][j] = min + 1
+			}
+		}
+	}
+	return d[len(s)][len(t)]
+}
+
+func TestLdBasicCases(t *testing.T) {
+	tests := []struct {
+		name       string
+		s          string
+		t          string
+		ignoreCase bool
+		want       int
+	}{
+		{
+			name: "both empty",
+			s:    "",
+			t:    "",
+			want: 0,
+		},
+		{
+			name: "empty source",
+			s:    "",
+			t:    "abc",
+			want: 3,
+		},
+		{
+			name: "empty target",
+			s:    "abc",
+			t:    "",
+			want: 3,
+		},
+		{
+			name: "same string",
+			s:    "hello",
+			t:    "hello",
+			want: 0,
+		},
+		{
+			name: "single replace",
+			s:    "a",
+			t:    "b",
+			want: 1,
+		},
+		{
+			name: "insert",
+			s:    "abc",
+			t:    "abcd",
+			want: 1,
+		},
+		{
+			name: "delete",
+			s:    "abcd",
+			t:    "abc",
+			want: 1,
+		},
+		{
+			name: "replace",
+			s:    "kitten",
+			t:    "sitten",
+			want: 1,
+		},
+		{
+			name: "swap to shorter row",
+			s:    "a",
+			t:    "abcdefgh",
+			want: 7,
+		},
+		{
+			name: "longer source",
+			s:    "abcdefgh",
+			t:    "a",
+			want: 7,
+		},
+		{
+			name:       "ignore case",
+			s:          "Hello",
+			t:          "hello",
+			ignoreCase: true,
+			want:       0,
+		},
+		{
+			name: "multi byte string",
+			s:    "café",
+			t:    "cafe",
+			want: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := ld(tt.s, tt.t, tt.ignoreCase)
+			if got != tt.want {
+				t.Fatalf("ld(%q, %q, %v) = %v, want %v",
+					tt.s, tt.t, tt.ignoreCase, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestLdMatchesReference(t *testing.T) {
+	alphabet := []string{"a", "A", "b", "B", "c", "-", "é"}
+
+	// Fixed seed: any failure reported below is reproducible as-is.
+	rng := rand.New(rand.NewSource(1))
+	randString := func(n int) string {
+		var sb strings.Builder
+		for i := 0; i < n; i++ {
+			sb.WriteString(alphabet[rng.Intn(len(alphabet))])
+		}
+		return sb.String()
+	}
+
+	for i := 0; i < 20000; i++ {
+		s := randString(rng.Intn(13))
+		u := randString(rng.Intn(13))
+
+		for _, ignoreCase := range []bool{false, true} {
+			want := ldReference(s, u, ignoreCase)
+			got := ld(s, u, ignoreCase)
+			if got != want {
+				t.Fatalf("ld(%q, %q, %v) = %v, want %v", s, u, ignoreCase, got, want)
+			}
+
+			// The distance is symmetric; ld swaps its arguments internally, so
+			// assert the caller cannot observe that.
+			if rev := ld(u, s, ignoreCase); rev != want {
+				t.Fatalf("ld(%q, %q, %v) = %v, want %v (asymmetric)", u, s, ignoreCase, rev, want)
+			}
+
+			// Distance is bounded by the length difference from below and by
+			// the longer string from above.
+			lo, hi := len(s)-len(u), len(s)
+			if lo < 0 {
+				lo = -lo
+			}
+			if len(u) > hi {
+				hi = len(u)
+			}
+			// Bounds hold for the byte lengths only when no case folding
+			// changed them.
+			if !ignoreCase && (got < lo || got > hi) {
+				t.Fatalf("ld(%q, %q, false) = %v, want within [%v, %v]", s, u, got, lo, hi)
+			}
+		}
+
+		if d := ld(s, s, false); d != 0 {
+			t.Fatalf("ld(%q, %q, false) = %v, want 0", s, s, d)
+		}
+	}
+}
+
+// ldSink keeps the compiler from optimizing the benchmarked calls away.
+var ldSink int
+
+// BenchmarkLd compares the single-row implementation against the full-matrix
+// one over the input sizes cobra actually sees: SuggestionsFor measures a
+// mistyped command name against every sibling command name, so both operands
+// are command names rather than arbitrary user input. The "long" case is well
+// past anything realistic and is only there to show how the two implementations
+// diverge as the inputs grow.
+func BenchmarkLd(b *testing.B) {
+	benchmarks := []struct {
+		name string
+		s, t string
+	}{
+		{"short", "get", "set"},
+		{"typical", "kubectl-config", "kubectl-configs"},
+		{"long", strings.Repeat("abcde-", 8), strings.Repeat("abdce-", 8)},
+	}
+
+	impls := []struct {
+		name string
+		fn   func(s, t string, ignoreCase bool) int
+	}{
+		{"optimized", ld},
+		{"reference", ldReference},
+	}
+
+	for _, bm := range benchmarks {
+		for _, impl := range impls {
+			b.Run(bm.name+"/"+impl.name, func(b *testing.B) {
+				b.ReportAllocs()
+				for i := 0; i < b.N; i++ {
+					ldSink = impl.fn(bm.s, bm.t, true)
+				}
+			})
+		}
+	}
+}