lib/sql/src/search/snippet.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const engine_mod = @import("engine.zig");
3 const query_mod = @import("query.zig");
4 const text_mod = @import("text.zig");
5
6 const Allocator = std.mem.Allocator;
7
8 pub const SnippetOptions = struct {
9 open_marker: []const u8 = "",
10 close_marker: []const u8 = "",
11 ellipsis: []const u8 = "...",
12 max_bytes: usize = 240,
13 context_bytes: usize = 80,
14 };
15
16 pub fn snippet(allocator: Allocator, text: []const u8, query: []const u8, options: SnippetOptions) engine_mod.Error![]u8 {
17 var parsed = try query_mod.parseQuery(allocator, query, &.{});
18 defer parsed.deinit(allocator);
19 return try snippetParsed(allocator, text, parsed, options);
20 }
21
22 pub const SnippetMatch = struct {
23 begin: usize,
24 end: usize,
25 term: usize,
26 };
27
28 pub const SnippetWindow = struct {
29 start: usize,
30 end: usize,
31 };
32
33 pub const SnippetWindowScore = struct {
34 distinct_terms: usize,
35 matches: usize,
36 first_match: usize,
37 start: usize,
38 };
39
40 pub fn snippetParsed(allocator: Allocator, text: []const u8, parsed: query_mod.Query, options: SnippetOptions) engine_mod.Error![]u8 {
41 if (text.len == 0 or options.max_bytes == 0) return try allocator.dupe(u8, "");
42 const tokens = try snippetTokens(allocator, text);
43 defer allocator.free(tokens);
44 const match_set = try snippetMatches(allocator, tokens, parsed);
45 defer allocator.free(match_set.matches);
46 const window = try chooseSnippetWindow(allocator, text.len, match_set.matches, match_set.terms, options);
47 return try renderSnippet(allocator, text, match_set.matches, window, options);
48 }
49
50 pub fn snippetTokens(allocator: Allocator, text: []const u8) engine_mod.Error![]text_mod.TextToken {
51 var tokens: std.ArrayList(text_mod.TextToken) = .empty;
52 errdefer tokens.deinit(allocator);
53 var offset: usize = 0;
54 var position: usize = 0;
55 var field: usize = 0;
56 while (text_mod.nextTextToken(text, &offset, &position, &field)) |token| try tokens.append(allocator, token);
57 return try tokens.toOwnedSlice(allocator);
58 }
59
60 pub const SnippetMatchSet = struct {
61 matches: []SnippetMatch,
62 terms: usize,
63 };
64
65 pub fn snippetMatches(allocator: Allocator, tokens: []const text_mod.TextToken, parsed: query_mod.Query) engine_mod.Error!SnippetMatchSet {
66 var matches: std.ArrayList(SnippetMatch) = .empty;
67 errdefer matches.deinit(allocator);
68 var term: usize = 0;
69 for (parsed.clauses) |clause| {
70 for (clause.positive) |atom| {
71 try appendAtomSnippetMatches(allocator, &matches, tokens, atom, term);
72 term += 1;
73 }
74 }
75 std.mem.sort(SnippetMatch, matches.items, {}, snippetMatchLess);
76 return .{ .matches = try matches.toOwnedSlice(allocator), .terms = term };
77 }
78
79 pub fn appendAtomSnippetMatches(allocator: Allocator, matches: *std.ArrayList(SnippetMatch), tokens: []const text_mod.TextToken, atom: query_mod.Atom, term: usize) engine_mod.Error!void {
80 if (atom.tokens.len == 0) return;
81 if (atom.phrase) {
82 var index: usize = 0;
83 while (index < tokens.len) : (index += 1) {
84 if (atom.field) |field| if (tokens[index].field != field) continue;
85 const end_index = phraseTokenEnd(tokens, index, atom.tokens) orelse continue;
86 try matches.append(allocator, .{ .begin = tokens[index].begin, .end = tokens[end_index].end, .term = term });
87 }
88 return;
89 }
90 if (atom.prefix) {
91 for (tokens) |token| {
92 if (atom.field) |field| if (token.field != field) continue;
93 if (text_mod.tokenStartsWith(token.text, atom.tokens[0])) try matches.append(allocator, .{ .begin = token.begin, .end = token.end, .term = term });
94 }
95 return;
96 }
97 for (tokens) |token| {
98 if (atom.field) |field| if (token.field != field) continue;
99 for (atom.tokens) |needle| {
100 if (text_mod.tokenEql(token.text, needle)) try matches.append(allocator, .{ .begin = token.begin, .end = token.end, .term = term });
101 }
102 }
103 }
104
105 pub fn phraseTokenEnd(tokens: []const text_mod.TextToken, start: usize, phrase: []const []const u8) ?usize {
106 if (phrase.len == 0 or start + phrase.len > tokens.len) return null;
107 const first_field = tokens[start].field;
108 const first_position = tokens[start].position;
109 for (phrase, 0..) |needle, index| {
110 const token = tokens[start + index];
111 if (token.field != first_field) return null;
112 if (token.position != first_position + index) return null;
113 if (!text_mod.tokenEql(token.text, needle)) return null;
114 }
115 return start + phrase.len - 1;
116 }
117
118 pub fn chooseSnippetWindow(allocator: Allocator, text_len: usize, matches: []const SnippetMatch, terms: usize, options: SnippetOptions) engine_mod.Error!SnippetWindow {
119 const max_bytes = @min(options.max_bytes, text_len);
120 if (matches.len == 0 or terms == 0) return .{ .start = 0, .end = max_bytes };
121 if (terms == 1) return snippetWindowForMatch(text_len, matches[0], options);
122 const seen = try allocator.alloc(usize, terms);
123 defer allocator.free(seen);
124 @memset(seen, 0);
125 var generation: usize = 0;
126 var best = snippetWindowForMatch(text_len, matches[0], options);
127 var best_score = scoreSnippetWindow(matches, 0, best, seen, &generation);
128 var first_match: usize = 0;
129 var index: usize = 1;
130 while (index < matches.len) : (index += 1) {
131 const window = snippetWindowForMatch(text_len, matches[index], options);
132 while (first_match < matches.len and matches[first_match].begin < window.start) : (first_match += 1) {}
133 const score = scoreSnippetWindow(matches, first_match, window, seen, &generation);
134 if (snippetScoreLess(best_score, score)) {
135 best = window;
136 best_score = score;
137 }
138 }
139 return best;
140 }
141
142 pub fn snippetWindowForMatch(text_len: usize, match: SnippetMatch, options: SnippetOptions) SnippetWindow {
143 const max_bytes = @min(options.max_bytes, text_len);
144 if (max_bytes == text_len) return .{ .start = 0, .end = text_len };
145 var start: usize = if (match.begin > options.context_bytes) match.begin - options.context_bytes else 0;
146 var end: usize = @min(text_len, start + max_bytes);
147 if (match.end > end) {
148 end = @min(text_len, match.end + options.context_bytes);
149 start = if (end > max_bytes) end - max_bytes else 0;
150 }
151 return .{ .start = start, .end = end };
152 }
153
154 pub fn scoreSnippetWindow(matches: []const SnippetMatch, first_match: usize, window: SnippetWindow, seen: []usize, generation: *usize) SnippetWindowScore {
155 generation.* += 1;
156 if (generation.* == 0) {
157 @memset(seen, 0);
158 generation.* = 1;
159 }
160 var score = SnippetWindowScore{
161 .distinct_terms = 0,
162 .matches = 0,
163 .first_match = std.math.maxInt(usize),
164 .start = window.start,
165 };
166 var index = first_match;
167 while (index < matches.len and matches[index].begin < window.end) : (index += 1) {
168 const match = matches[index];
169 if (match.end > window.end) continue;
170 score.matches += 1;
171 score.first_match = @min(score.first_match, match.begin);
172 if (seen[match.term] != generation.*) {
173 seen[match.term] = generation.*;
174 score.distinct_terms += 1;
175 }
176 }
177 return score;
178 }
179
180 pub fn snippetScoreLess(left: SnippetWindowScore, right: SnippetWindowScore) bool {
181 if (left.distinct_terms != right.distinct_terms) return left.distinct_terms < right.distinct_terms;
182 if (left.matches != right.matches) return left.matches < right.matches;
183 if (left.first_match != right.first_match) return left.first_match > right.first_match;
184 return left.start > right.start;
185 }
186
187 pub fn renderSnippet(allocator: Allocator, text: []const u8, matches: []const SnippetMatch, window: SnippetWindow, options: SnippetOptions) engine_mod.Error![]u8 {
188 var out: std.ArrayList(u8) = .empty;
189 errdefer out.deinit(allocator);
190 if (window.start != 0) try out.appendSlice(allocator, options.ellipsis);
191 var cursor = window.start;
192 var index: usize = 0;
193 while (index < matches.len and matches[index].end <= window.start) : (index += 1) {}
194 while (index < matches.len and matches[index].begin < window.end) {
195 const first = matches[index];
196 if (first.begin < window.start or first.end > window.end or first.end <= cursor) {
197 index += 1;
198 continue;
199 }
200 var begin = first.begin;
201 var end = first.end;
202 index += 1;
203 while (index < matches.len and matches[index].begin <= end) {
204 const next = matches[index];
205 if (next.begin >= window.end) break;
206 if (next.begin >= window.start and next.end <= window.end) end = @max(end, next.end);
207 index += 1;
208 }
209 if (begin < cursor) begin = cursor;
210 try out.appendSlice(allocator, text[cursor..begin]);
211 try out.appendSlice(allocator, options.open_marker);
212 try out.appendSlice(allocator, text[begin..end]);
213 try out.appendSlice(allocator, options.close_marker);
214 cursor = end;
215 }
216 try out.appendSlice(allocator, text[cursor..window.end]);
217 if (window.end != text.len) try out.appendSlice(allocator, options.ellipsis);
218 return try out.toOwnedSlice(allocator);
219 }
220
221 pub fn snippetMatchLess(_: void, left: SnippetMatch, right: SnippetMatch) bool {
222 if (left.begin != right.begin) return left.begin < right.begin;
223 if (left.end != right.end) return left.end > right.end;
224 return left.term < right.term;
225 }
226
227 test "search snippet highlights exact OR terms" {
228 const out = try snippet(std.testing.allocator, "zero one alpha two beta three", "alpha OR beta", .{
229 .open_marker = "[",
230 .close_marker = "]",
231 .ellipsis = "...",
232 .max_bytes = 128,
233 .context_bytes = 8,
234 });
235 defer std.testing.allocator.free(out);
236 try std.testing.expectEqualStrings("zero one [alpha] two [beta] three", out);
237 }
238
239 test "search snippet honors field-constrained atoms" {
240 const field_names = [_][]const u8{ "title", "body" };
241 var parsed = try query_mod.parseQuery(std.testing.allocator, "title:alpha", &field_names);
242 defer parsed.deinit(std.testing.allocator);
243 const out = try snippetParsed(std.testing.allocator, "alpha intro\x1falpha outro", parsed, .{
244 .open_marker = "[",
245 .close_marker = "]",
246 .ellipsis = "...",
247 .max_bytes = 128,
248 .context_bytes = 8,
249 });
250 defer std.testing.allocator.free(out);
251 try std.testing.expectEqualStrings("[alpha] intro\x1falpha outro", out);
252 }
253
254 test "search snippet highlights phrase and prefix matches" {
255 const phrase = try snippet(std.testing.allocator, "before alpha beta after", "\"alpha beta\"", .{
256 .open_marker = "[",
257 .close_marker = "]",
258 .ellipsis = "...",
259 .max_bytes = 128,
260 .context_bytes = 8,
261 });
262 defer std.testing.allocator.free(phrase);
263 try std.testing.expectEqualStrings("before [alpha beta] after", phrase);
264
265 const prefix = try snippet(std.testing.allocator, "Fix authentication and authorization", "auth*", .{
266 .open_marker = "[",
267 .close_marker = "]",
268 .ellipsis = "...",
269 .max_bytes = 128,
270 .context_bytes = 8,
271 });
272 defer std.testing.allocator.free(prefix);
273 try std.testing.expectEqualStrings("Fix [authentication] and [authorization]", prefix);
274 }
275
276 test "search snippet bounds fallback text" {
277 const out = try snippet(std.testing.allocator, "abcdefghijklmnopqrstuvwxyz", "missing", .{
278 .ellipsis = "...",
279 .max_bytes = 10,
280 });
281 defer std.testing.allocator.free(out);
282 try std.testing.expectEqualStrings("abcdefghij...", out);
283 }
284
285 test "search snippet merges overlapping phrase matches" {
286 const out = try snippet(std.testing.allocator, "a b c d e", "\"a b c\" OR \"c d e\"", .{
287 .open_marker = "[",
288 .close_marker = "]",
289 .ellipsis = "...",
290 .max_bytes = 128,
291 });
292 defer std.testing.allocator.free(out);
293 try std.testing.expectEqualStrings("[a b c d e]", out);
294 }