|
| 1 | +// Package searchrank reranks FTS-ranked code-search candidates using |
| 2 | +// dependency-free structural signals (name fuzzy similarity, path proximity), |
| 3 | +// so both the CLI `search` command and the MCP `search` tool share one ranking. |
| 4 | +package searchrank |
| 5 | + |
| 6 | +import ( |
| 7 | + "sort" |
| 8 | + "strings" |
| 9 | + "unicode" |
| 10 | + |
| 11 | + "github.com/tae2089/code-context-graph/internal/model" |
| 12 | +) |
| 13 | + |
| 14 | +// Reciprocal Rank Fusion constants. |
| 15 | +// rrfK dampens how sharply rank position affects the fused score (60 is the |
| 16 | +// conventional value). rrfStructWeight (>1) lets the structural signal override |
| 17 | +// a small FTS-rank gap; at weight 1 RRF is symmetric and a two-item rank swap |
| 18 | +// always ties, leaving FTS order untouched. |
| 19 | +const ( |
| 20 | + rrfK = 60.0 |
| 21 | + rrfStructWeight = 2.0 |
| 22 | +) |
| 23 | + |
| 24 | +// Structural signal weights. Name match is more discriminating than path |
| 25 | +// proximity: every node in one file shares the path, so an unweighted max lets |
| 26 | +// a path hit saturate the score and bury the exact-name match. Weighting name |
| 27 | +// above path keeps name the dominant signal with path as a nudge. |
| 28 | +const ( |
| 29 | + nameSignalWeight = 1.0 |
| 30 | + pathSignalWeight = 0.25 |
| 31 | +) |
| 32 | + |
| 33 | +// Candidate-pool sizing. Callers over-fetch a wider pool than the requested |
| 34 | +// limit so reranking has more than `limit` rows to reorder before bounding. |
| 35 | +const ( |
| 36 | + fetchFactor = 5 |
| 37 | + fetchFloor = 50 |
| 38 | + fetchCap = 500 |
| 39 | +) |
| 40 | + |
| 41 | +// FetchLimit widens the candidate pool pulled from FTS so structural reranking |
| 42 | +// (and any path filtering) has more than the caller's `limit` rows to reorder; |
| 43 | +// the final slice is bounded back to `limit` after reranking. |
| 44 | +func FetchLimit(limit int) int { |
| 45 | + return min(max(limit*fetchFactor, fetchFloor), fetchCap) |
| 46 | +} |
| 47 | + |
| 48 | +// Rerank reorders FTS-ranked search candidates using structural signals fused |
| 49 | +// with the backend rank via Reciprocal Rank Fusion. |
| 50 | +// |
| 51 | +// @requires nodes is the backend's rank-ordered candidate slice (index == FTS rank). |
| 52 | +// @ensures deterministic output; empty query or empty nodes returns the input |
| 53 | +// bounded by limit, preserving FTS order. |
| 54 | +func Rerank(query string, nodes []model.Node, limit int) []model.Node { |
| 55 | + if strings.TrimSpace(query) == "" || len(nodes) == 0 { |
| 56 | + return applyLimit(nodes, limit) |
| 57 | + } |
| 58 | + qTokens := tokenize(query) |
| 59 | + if len(qTokens) == 0 { |
| 60 | + return applyLimit(nodes, limit) |
| 61 | + } |
| 62 | + |
| 63 | + structScores := make([]float64, len(nodes)) |
| 64 | + for i := range nodes { |
| 65 | + structScores[i] = structScore(qTokens, nodes[i]) |
| 66 | + } |
| 67 | + structRank := rankDesc(structScores) |
| 68 | + |
| 69 | + final := make([]float64, len(nodes)) |
| 70 | + for i := range nodes { |
| 71 | + final[i] = 1.0/(rrfK+float64(i)) + rrfStructWeight/(rrfK+float64(structRank[i])) |
| 72 | + } |
| 73 | + |
| 74 | + order := make([]int, len(nodes)) |
| 75 | + for i := range order { |
| 76 | + order[i] = i |
| 77 | + } |
| 78 | + // Stable sort by fused score descending; equal scores keep the original |
| 79 | + // FTS order, so ranking stays deterministic. |
| 80 | + sort.SliceStable(order, func(a, b int) bool { |
| 81 | + return final[order[a]] > final[order[b]] |
| 82 | + }) |
| 83 | + |
| 84 | + out := make([]model.Node, len(nodes)) |
| 85 | + for pos, idx := range order { |
| 86 | + out[pos] = nodes[idx] |
| 87 | + } |
| 88 | + return applyLimit(out, limit) |
| 89 | +} |
| 90 | + |
| 91 | +// applyLimit bounds the result slice, treating a non-positive limit as unbounded. |
| 92 | +func applyLimit(nodes []model.Node, limit int) []model.Node { |
| 93 | + if limit > 0 && len(nodes) > limit { |
| 94 | + return nodes[:limit] |
| 95 | + } |
| 96 | + return nodes |
| 97 | +} |
| 98 | + |
| 99 | +// structScore is a weighted sum of the name and path signals, with name |
| 100 | +// dominant so an exact-name match outranks a mere same-file (path) match. |
| 101 | +func structScore(qTokens []string, node model.Node) float64 { |
| 102 | + return nameSignalWeight*nameSim(qTokens, node) + pathSignalWeight*pathScore(qTokens, node) |
| 103 | +} |
| 104 | + |
| 105 | +// nameSim scores fuzzy similarity of the query against the node name and the |
| 106 | +// last segment of its qualified name. For each target it takes the stronger of: |
| 107 | +// - token-level: every query token's best match against the whole name or any |
| 108 | +// of its identifier sub-tokens (so "user" or "id" matches getUserById), and |
| 109 | +// - joined-whole: the run-together query vs the whole name (so a typo like |
| 110 | +// "getUsrById" still matches getUserById). |
| 111 | +func nameSim(qTokens []string, node model.Node) float64 { |
| 112 | + joined := strings.Join(qTokens, "") |
| 113 | + rawTargets := []string{node.Name, lastSegment(node.QualifiedName, '.')} |
| 114 | + best := 0.0 |
| 115 | + for _, raw := range rawTargets { |
| 116 | + if raw == "" { |
| 117 | + continue |
| 118 | + } |
| 119 | + lower := strings.ToLower(raw) |
| 120 | + subs := splitIdentifier(raw) // original case: camelCase boundaries matter |
| 121 | + sum := 0.0 |
| 122 | + for _, tok := range qTokens { |
| 123 | + b := normLevSim(tok, lower) |
| 124 | + for _, st := range subs { |
| 125 | + b = max(b, normLevSim(tok, st)) |
| 126 | + } |
| 127 | + sum += b |
| 128 | + } |
| 129 | + cand := max(sum/float64(len(qTokens)), normLevSim(joined, lower)) |
| 130 | + best = max(best, cand) |
| 131 | + } |
| 132 | + return best |
| 133 | +} |
| 134 | + |
| 135 | +// splitIdentifier breaks an identifier into lowercased sub-tokens on separators, |
| 136 | +// camelCase boundaries, and letter/digit transitions ("getUserById" -> get, |
| 137 | +// user, by, id; "HTTPServer" -> http, server; "parseHTML5" -> parse, html, 5). |
| 138 | +func splitIdentifier(s string) []string { |
| 139 | + runes := []rune(s) |
| 140 | + var tokens []string |
| 141 | + var cur []rune |
| 142 | + flush := func() { |
| 143 | + if len(cur) > 0 { |
| 144 | + tokens = append(tokens, strings.ToLower(string(cur))) |
| 145 | + cur = cur[:0] |
| 146 | + } |
| 147 | + } |
| 148 | + for i, r := range runes { |
| 149 | + if !isAlnum(r) { |
| 150 | + flush() |
| 151 | + continue |
| 152 | + } |
| 153 | + if len(cur) > 0 { |
| 154 | + prev := cur[len(cur)-1] |
| 155 | + switch { |
| 156 | + case unicode.IsUpper(r) && (unicode.IsLower(prev) || unicode.IsDigit(prev)): |
| 157 | + flush() // lower/digit -> Upper: new word |
| 158 | + case unicode.IsUpper(r) && unicode.IsUpper(prev) && i+1 < len(runes) && unicode.IsLower(runes[i+1]): |
| 159 | + flush() // acronym tail begins a new word (HTTPServer -> http, server) |
| 160 | + case unicode.IsDigit(r) != unicode.IsDigit(prev): |
| 161 | + flush() // letter/digit transition |
| 162 | + } |
| 163 | + } |
| 164 | + cur = append(cur, r) |
| 165 | + } |
| 166 | + flush() |
| 167 | + return tokens |
| 168 | +} |
| 169 | + |
| 170 | +func isAlnum(r rune) bool { |
| 171 | + return unicode.IsLetter(r) || unicode.IsDigit(r) |
| 172 | +} |
| 173 | + |
| 174 | +// pathScore is the fraction of query tokens that appear as file-path segments. |
| 175 | +func pathScore(qTokens []string, node model.Node) float64 { |
| 176 | + segs := map[string]struct{}{} |
| 177 | + for _, seg := range strings.FieldsFunc(strings.ToLower(node.FilePath), isPathSep) { |
| 178 | + segs[seg] = struct{}{} |
| 179 | + } |
| 180 | + if len(segs) == 0 { |
| 181 | + return 0 |
| 182 | + } |
| 183 | + hits := 0 |
| 184 | + for _, tok := range qTokens { |
| 185 | + if _, ok := segs[tok]; ok { |
| 186 | + hits++ |
| 187 | + } |
| 188 | + } |
| 189 | + return float64(hits) / float64(len(qTokens)) |
| 190 | +} |
| 191 | + |
| 192 | +// tokenize lowercases and splits input into alphanumeric tokens. |
| 193 | +func tokenize(s string) []string { |
| 194 | + return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool { |
| 195 | + return !unicode.IsLetter(r) && !unicode.IsDigit(r) |
| 196 | + }) |
| 197 | +} |
| 198 | + |
| 199 | +func isPathSep(r rune) bool { |
| 200 | + switch r { |
| 201 | + case '/', '.', '_', '-': |
| 202 | + return true |
| 203 | + default: |
| 204 | + return false |
| 205 | + } |
| 206 | +} |
| 207 | + |
| 208 | +func lastSegment(s string, sep rune) string { |
| 209 | + if i := strings.LastIndexByte(s, byte(sep)); i >= 0 { |
| 210 | + return s[i+1:] |
| 211 | + } |
| 212 | + return s |
| 213 | +} |
| 214 | + |
| 215 | +// rankDesc returns each index's position when scores are ordered descending; |
| 216 | +// ties resolve by original index so the ranking is deterministic. |
| 217 | +func rankDesc(scores []float64) []int { |
| 218 | + order := make([]int, len(scores)) |
| 219 | + for i := range order { |
| 220 | + order[i] = i |
| 221 | + } |
| 222 | + sort.SliceStable(order, func(a, b int) bool { |
| 223 | + return scores[order[a]] > scores[order[b]] |
| 224 | + }) |
| 225 | + rank := make([]int, len(scores)) |
| 226 | + for pos, idx := range order { |
| 227 | + rank[idx] = pos |
| 228 | + } |
| 229 | + return rank |
| 230 | +} |
| 231 | + |
| 232 | +// normLevSim is the Levenshtein distance normalized to a [0,1] similarity. |
| 233 | +func normLevSim(a, b string) float64 { |
| 234 | + ra, rb := []rune(a), []rune(b) |
| 235 | + maxLen := max(len(ra), len(rb)) |
| 236 | + if maxLen == 0 { |
| 237 | + return 0 |
| 238 | + } |
| 239 | + return 1.0 - float64(levenshtein(ra, rb))/float64(maxLen) |
| 240 | +} |
| 241 | + |
| 242 | +// levenshtein computes edit distance with a rolling single-row DP. |
| 243 | +func levenshtein(a, b []rune) int { |
| 244 | + prev := make([]int, len(b)+1) |
| 245 | + for j := range prev { |
| 246 | + prev[j] = j |
| 247 | + } |
| 248 | + for i := 1; i <= len(a); i++ { |
| 249 | + curr := make([]int, len(b)+1) |
| 250 | + curr[0] = i |
| 251 | + for j := 1; j <= len(b); j++ { |
| 252 | + cost := 1 |
| 253 | + if a[i-1] == b[j-1] { |
| 254 | + cost = 0 |
| 255 | + } |
| 256 | + curr[j] = min(prev[j]+1, curr[j-1]+1, prev[j-1]+cost) |
| 257 | + } |
| 258 | + prev = curr |
| 259 | + } |
| 260 | + return prev[len(b)] |
| 261 | +} |
0 commit comments