lib/sql/src/properties/search/model.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const field_separator: u8 = 0x1f;
4 pub const field_names = [_][]const u8{ "title", "body" };
5
6 pub const Token = struct {
7 text: []const u8,
8 field: usize,
9 position: usize,
10 };
11
12 pub const Atom = struct {
13 terms: []const []const u8,
14 prefix: bool = false,
15 phrase: bool = false,
16 field: ?usize = null,
17 negated: bool = false,
18 };
19
20 pub const Clause = struct {
21 atoms: []const Atom,
22 };
23
24 pub const Query = struct {
25 clauses: []const Clause,
26 };
27
28 pub const Group = struct {
29 atoms: []const Atom,
30 negated: bool = false,
31 };
32
33 pub const Hit = struct {
34 rowid: i64,
35 score: usize,
36 };
37
38 pub const Document = struct {
39 rowid: i64,
40 text: []const u8,
41 };
42
43 pub fn tokens(allocator: std.mem.Allocator, text: []const u8) ![]Token {
44 var out: std.ArrayList(Token) = .empty;
45 errdefer out.deinit(allocator);
46 var field: usize = 0;
47 var position: usize = 0;
48 var start: ?usize = null;
49 for (text, 0..) |byte, index| {
50 if (std.ascii.isAlphanumeric(byte)) {
51 if (start == null) start = index;
52 continue;
53 }
54 if (start) |begin| {
55 try out.append(allocator, .{ .text = text[begin..index], .field = field, .position = position });
56 position += 1;
57 start = null;
58 }
59 if (byte == field_separator) {
60 field += 1;
61 position = 0;
62 }
63 }
64 if (start) |begin| try out.append(allocator, .{ .text = text[begin..], .field = field, .position = position });
65 return try out.toOwnedSlice(allocator);
66 }
67
68 fn foldedEql(token: []const u8, lowercase: []const u8) bool {
69 if (token.len != lowercase.len) return false;
70 for (token, lowercase) |left, right| if (std.ascii.toLower(left) != right) return false;
71 return true;
72 }
73
74 fn foldedStartsWith(token: []const u8, lowercase: []const u8) bool {
75 if (token.len < lowercase.len) return false;
76 return foldedEql(token[0..lowercase.len], lowercase);
77 }
78
79 fn termCount(document: []const Token, term: []const u8, field: ?usize) usize {
80 var count: usize = 0;
81 for (document) |token| {
82 if (field) |filter| if (token.field != filter) continue;
83 if (foldedEql(token.text, term)) count += 1;
84 }
85 return count;
86 }
87
88 fn prefixCount(document: []const Token, prefix: []const u8, field: ?usize) usize {
89 var count: usize = 0;
90 for (document) |token| {
91 if (field) |filter| if (token.field != filter) continue;
92 if (foldedStartsWith(token.text, prefix)) count += 1;
93 }
94 return count;
95 }
96
97 fn phraseCount(document: []const Token, phrase: []const []const u8, field: ?usize) usize {
98 if (phrase.len == 0) return 0;
99 var count: usize = 0;
100 for (document, 0..) |first, start| {
101 if (field) |filter| if (first.field != filter) continue;
102 if (!foldedEql(first.text, phrase[0])) continue;
103 if (start + phrase.len > document.len) continue;
104 var index: usize = 1;
105 while (index < phrase.len) : (index += 1) {
106 const next = document[start + index];
107 if (next.field != first.field) break;
108 if (next.position != first.position + index) break;
109 if (!foldedEql(next.text, phrase[index])) break;
110 }
111 if (index == phrase.len) count += 1;
112 }
113 return count;
114 }
115
116 pub fn atomScore(document: []const Token, atom: Atom) usize {
117 if (atom.terms.len == 0) return 0;
118 if (atom.prefix) return prefixCount(document, atom.terms[0], atom.field);
119 if (atom.phrase) return phraseCount(document, atom.terms, atom.field) * atom.terms.len;
120 var score: usize = 0;
121 for (atom.terms) |term| {
122 const count = termCount(document, term, atom.field);
123 if (count == 0) return 0;
124 score += count;
125 }
126 return score;
127 }
128
129 pub fn clauseScore(document: []const Token, clause: Clause) usize {
130 var positive: usize = 0;
131 var score: usize = 0;
132 for (clause.atoms) |atom| {
133 const value = atomScore(document, atom);
134 if (atom.negated) {
135 if (value != 0) return 0;
136 continue;
137 }
138 positive += 1;
139 if (value == 0) return 0;
140 score += value;
141 }
142 return if (positive == 0) 1 else score;
143 }
144
145 pub fn queryScore(document: []const Token, query: Query) usize {
146 var score: usize = 0;
147 for (query.clauses) |clause| score += clauseScore(document, clause);
148 return score;
149 }
150
151 pub fn groupsMatch(document: []const Token, groups: []const Group) bool {
152 if (groups.len == 0) return false;
153 for (groups) |group| {
154 var matched = false;
155 for (group.atoms) |atom| {
156 if (atomScore(document, atom) != 0) {
157 matched = true;
158 break;
159 }
160 }
161 if (matched == group.negated) return false;
162 }
163 return true;
164 }
165
166 pub fn hitLess(_: void, left: Hit, right: Hit) bool {
167 if (left.score != right.score) return left.score > right.score;
168 return left.rowid < right.rowid;
169 }
170
171 pub fn hits(allocator: std.mem.Allocator, corpus: []const Document, query: Query) ![]Hit {
172 var out: std.ArrayList(Hit) = .empty;
173 errdefer out.deinit(allocator);
174 for (corpus) |document| {
175 const document_tokens = try tokens(allocator, document.text);
176 defer allocator.free(document_tokens);
177 const score = queryScore(document_tokens, query);
178 if (score != 0) try out.append(allocator, .{ .rowid = document.rowid, .score = score });
179 }
180 const owned = try out.toOwnedSlice(allocator);
181 std.mem.sort(Hit, owned, {}, hitLess);
182 return owned;
183 }
184
185 pub fn groupHits(allocator: std.mem.Allocator, corpus: []const Document, groups: []const Group) ![]i64 {
186 var out: std.ArrayList(i64) = .empty;
187 errdefer out.deinit(allocator);
188 for (corpus) |document| {
189 const document_tokens = try tokens(allocator, document.text);
190 defer allocator.free(document_tokens);
191 if (groupsMatch(document_tokens, groups)) try out.append(allocator, document.rowid);
192 }
193 const owned = try out.toOwnedSlice(allocator);
194 std.mem.sort(i64, owned, {}, std.sort.asc(i64));
195 return owned;
196 }
197
198 pub fn sortedRowids(allocator: std.mem.Allocator, hit_list: anytype) ![]i64 {
199 const rowids = try allocator.alloc(i64, hit_list.len);
200 for (hit_list, 0..) |hit, index| rowids[index] = hit.rowid;
201 std.mem.sort(i64, rowids, {}, std.sort.asc(i64));
202 return rowids;
203 }
204
205 pub fn band(full: anytype, limit: usize) @TypeOf(full) {
206 if (limit == 0 or full.len <= limit) return full;
207 const cutoff = full[limit - 1].score;
208 var count = limit;
209 while (count < full.len and full[count].score >= cutoff) count += 1;
210 return full[0..count];
211 }
212
213 test "model tokens split fields and reset positions" {
214 const allocator = std.testing.allocator;
215 const text = "Alpha, beta-2" ++ [_]u8{field_separator} ++ "Gamma";
216 const document_tokens = try tokens(allocator, text);
217 defer allocator.free(document_tokens);
218 try std.testing.expectEqual(@as(usize, 4), document_tokens.len);
219 try std.testing.expectEqualStrings("Alpha", document_tokens[0].text);
220 try std.testing.expectEqual(@as(usize, 0), document_tokens[0].position);
221 try std.testing.expectEqualStrings("2", document_tokens[2].text);
222 try std.testing.expectEqual(@as(usize, 2), document_tokens[2].position);
223 try std.testing.expectEqual(@as(usize, 1), document_tokens[3].field);
224 try std.testing.expectEqual(@as(usize, 0), document_tokens[3].position);
225 }
226
227 test "model phrase counting stays within one field" {
228 const allocator = std.testing.allocator;
229 const split = "alpha" ++ [_]u8{field_separator} ++ "beta";
230 const split_tokens = try tokens(allocator, split);
231 defer allocator.free(split_tokens);
232 const phrase = [_][]const u8{ "alpha", "beta" };
233 try std.testing.expectEqual(@as(usize, 0), phraseCount(split_tokens, &phrase, null));
234
235 const joined_tokens = try tokens(allocator, "alpha beta alpha beta");
236 defer allocator.free(joined_tokens);
237 try std.testing.expectEqual(@as(usize, 2), phraseCount(joined_tokens, &phrase, null));
238 }
239
240 test "model clause score folds case and sums positive atoms" {
241 const allocator = std.testing.allocator;
242 const document_tokens = try tokens(allocator, "ALPHA alpha beta");
243 defer allocator.free(document_tokens);
244 const alpha = Atom{ .terms = &.{"alpha"} };
245 const beta = Atom{ .terms = &.{"beta"} };
246 const gamma = Atom{ .terms = &.{"gamma"} };
247 try std.testing.expectEqual(@as(usize, 3), clauseScore(document_tokens, .{ .atoms = &.{ alpha, beta } }));
248 try std.testing.expectEqual(@as(usize, 0), clauseScore(document_tokens, .{ .atoms = &.{ alpha, gamma } }));
249 const negated_gamma = Atom{ .terms = &.{"gamma"}, .negated = true };
250 try std.testing.expectEqual(@as(usize, 2), clauseScore(document_tokens, .{ .atoms = &.{ alpha, negated_gamma } }));
251 try std.testing.expectEqual(@as(usize, 1), clauseScore(document_tokens, .{ .atoms = &.{negated_gamma} }));
252 }