perf: reduce memory usage in command suggestions

Use a single-row Levenshtein distance implementation instead of
allocating the full DP matrix for every command comparison.

The edit distance calculation is unchanged, but temporary allocations
are reduced significantly during suggestion generation.

Signed-off-by: 0xff-dev <stevenshuang521@gmail.com>

Author: 0xff-dev

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)
+				}
+			})
+		}
+	}
+}