╭ search.zig ───────────────────────────────────────────────────╮ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ╰───────────────────────────────────────────────────────────────╯
╭────╮ ╭ clone ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ╮ │http│ ssh │ │ ╰────╯ ╰ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ╯
1 const std = @import("std"); 2 3 // an inverted index over any xitdb cursor: a sorted map from word to the 4 // sorted set of doc keys whose text holds it. a doc key is opaque bytes whose 5 // sort order is the order results come back in. 6 7 // the shortest word indexed or queried. a single character would expand to 8 // hundreds of words under prefix matching. 9 const min_word_len = 2; 10 // a longer word is truncated to this rather than skipped, so a prefix query 11 // still reaches it. 12 const max_word_len = 32; 13 // how many distinct words one document contributes. 14 const max_words_per_doc = 256; 15 // how much of a document is tokenized. 16 pub const max_indexed_bytes = 16 * 1024; 17 // how many indexed words one typed word expands to. 18 const max_expansions = 50; 19 // the largest doc key, so results are read into a buffer rather than allocated. 20 pub const max_doc_key_len = 64; 21 22 // index `text` under `doc_key`. 23 pub fn add( 24 comptime DB: type, 25 index: DB.SortedMap(.read_write), 26 allocator: std.mem.Allocator, 27 doc_key: []const u8, 28 text: []const u8, 29 ) !void { 30 var arena = std.heap.ArenaAllocator.init(allocator); 31 defer arena.deinit(); 32 for (try tokenize(arena.allocator(), text)) |word| { 33 try addWord(DB, index, doc_key, word); 34 } 35 } 36 37 // drop `doc_key`'s postings. the caller supplies the text again, so no forward 38 // index is stored. 39 pub fn remove( 40 comptime DB: type, 41 index: DB.SortedMap(.read_write), 42 allocator: std.mem.Allocator, 43 doc_key: []const u8, 44 text: []const u8, 45 ) !void { 46 var arena = std.heap.ArenaAllocator.init(allocator); 47 defer arena.deinit(); 48 for (try tokenize(arena.allocator(), text)) |word| { 49 try removeWord(DB, index, doc_key, word); 50 } 51 } 52 53 // move `doc_key` from `old_text`'s words to `new_text`'s. the words they share 54 // are left alone, so editing a long description costs what the edit changed 55 // rather than the whole text. 56 pub fn replace( 57 comptime DB: type, 58 index: DB.SortedMap(.read_write), 59 allocator: std.mem.Allocator, 60 doc_key: []const u8, 61 old_text: []const u8, 62 new_text: []const u8, 63 ) !void { 64 var arena = std.heap.ArenaAllocator.init(allocator); 65 defer arena.deinit(); 66 const aa = arena.allocator(); 67 const old_words = try tokenize(aa, old_text); 68 const new_words = try tokenize(aa, new_text); 69 70 for (old_words) |word| { 71 if (!contains(new_words, word)) try removeWord(DB, index, doc_key, word); 72 } 73 for (new_words) |word| { 74 if (!contains(old_words, word)) try addWord(DB, index, doc_key, word); 75 } 76 } 77 78 // a document contributes at most max_words_per_doc words, so a scan costs less 79 // than building a second map 80 fn contains(words: []const []const u8, word: []const u8) bool { 81 for (words) |candidate| { 82 if (std.mem.eql(u8, candidate, word)) return true; 83 } 84 return false; 85 } 86 87 fn addWord(comptime DB: type, index: DB.SortedMap(.read_write), doc_key: []const u8, word: []const u8) !void { 88 const set = try DB.SortedSet(.read_write).init(try index.putCursor(word)); 89 try set.put(doc_key); 90 } 91 92 fn removeWord(comptime DB: type, index: DB.SortedMap(.read_write), doc_key: []const u8, word: []const u8) !void { 93 if (null == try index.getCursor(word)) return; 94 const set = try DB.SortedSet(.read_write).init(try index.putCursor(word)); 95 _ = try set.remove(doc_key); 96 if (0 == try set.count()) _ = try index.remove(word); 97 } 98 99 // the doc keys holding every typed word, in key order. a stream, so callers 100 // stop at their page size. 101 pub fn Query(comptime DB: type) type { 102 return struct { 103 // one per typed word; empty when nothing can match 104 terms: []Term, 105 // the key the next result starts from, plus room for the one byte that 106 // steps past a returned key 107 pos: [max_doc_key_len + 1]u8 = undefined, 108 pos_len: usize = 0, 109 // whether `pos` holds a key that was already returned 110 returned: bool = false, 111 112 const Self = @This(); 113 114 // one typed word: the doc keys of every indexed word it expands to, 115 // merged into one ascending stream. 116 const Term = struct { 117 postings: []Posting, 118 119 // the smallest doc key at or after `key`, or null when the term is 120 // exhausted. 121 fn seek(self: *Term, key: []const u8) !?[]const u8 { 122 var best: ?[]const u8 = null; 123 for (self.postings) |*posting| { 124 try posting.advance(key); 125 if (posting.done) continue; 126 const current = posting.key[0..posting.key_len]; 127 if (best) |found| { 128 if (std.mem.order(u8, current, found) == .lt) best = current; 129 } else best = current; 130 } 131 return best; 132 } 133 }; 134 135 // one indexed word's doc keys, positioned at the current one. 136 const Posting = struct { 137 set: DB.SortedSet(.read_only), 138 iter: DB.Cursor(.read_only).Iter = undefined, 139 key: [max_doc_key_len]u8 = undefined, 140 key_len: usize = 0, 141 done: bool = false, 142 started: bool = false, 143 144 // move to the first key at or after `key`. one step reaches it in 145 // a dense posting, and a seek skips the rest of a sparse one. 146 fn advance(self: *Posting, key: []const u8) !void { 147 if (self.done) return; 148 if (self.started) { 149 if (std.mem.order(u8, self.key[0..self.key_len], key) != .lt) return; 150 try self.step(); 151 if (self.done or std.mem.order(u8, self.key[0..self.key_len], key) != .lt) return; 152 } 153 self.iter = try self.set.iteratorFrom(key); 154 self.started = true; 155 try self.step(); 156 } 157 158 fn step(self: *Posting) !void { 159 const cursor = (try self.iter.next()) orelse { 160 self.done = true; 161 return; 162 }; 163 const pair = try cursor.readKeyValuePair(); 164 self.key_len = (try pair.key_cursor.readBytes(&self.key)).len; 165 } 166 }; 167 168 pub fn init(index: DB.SortedMap(.read_only), aa: std.mem.Allocator, text: []const u8) !Self { 169 var terms: std.ArrayList(Term) = .empty; 170 for (try tokenize(aa, text)) |word| { 171 var postings: std.ArrayList(Posting) = .empty; 172 var iter = try index.iteratorFrom(word); 173 while (try iter.next()) |cursor| { 174 if (postings.items.len == max_expansions) break; 175 const pair = try cursor.readKeyValuePair(); 176 var buffer: [max_word_len]u8 = undefined; 177 const indexed = try pair.key_cursor.readBytes(&buffer); 178 if (!std.mem.startsWith(u8, indexed, word)) break; 179 try postings.append(aa, .{ .set = try DB.SortedSet(.read_only).init(pair.value_cursor) }); 180 } 181 // a word nothing starts with rules out every document 182 if (postings.items.len == 0) return .{ .terms = &.{} }; 183 try terms.append(aa, .{ .postings = postings.items }); 184 } 185 return .{ .terms = terms.items }; 186 } 187 188 // narrow the results to the doc keys in `set`. it is one more term, 189 // with that set as its only posting. a query nothing can match stays 190 // empty, since requiring a set would turn it into every doc key in it. 191 pub fn require(self: *Self, aa: std.mem.Allocator, set: DB.SortedSet(.read_only)) !void { 192 if (self.terms.len == 0) return; 193 var terms: std.ArrayList(Term) = .empty; 194 try terms.appendSlice(aa, self.terms); 195 const postings = try aa.alloc(Posting, 1); 196 postings[0] = .{ .set = set }; 197 try terms.append(aa, .{ .postings = postings }); 198 self.terms = terms.items; 199 } 200 201 // start the next result at `key` rather than the first one. 202 pub fn seek(self: *Self, key: []const u8) void { 203 @memcpy(self.pos[0..key.len], key); 204 self.pos_len = key.len; 205 self.returned = false; 206 } 207 208 pub fn next(self: *Self) !?[]const u8 { 209 if (self.terms.len == 0) return null; 210 if (self.returned) { 211 // doc keys are at most max_doc_key_len, so appending a byte 212 // lands strictly after the one just returned 213 self.pos[self.pos_len] = 0; 214 self.pos_len += 1; 215 self.returned = false; 216 } 217 // intersect by seeking: a term that lands past the candidate makes 218 // its key the new candidate and the others seek up to it. 219 var i: usize = 0; 220 while (i < self.terms.len) { 221 const key = (try self.terms[i].seek(self.pos[0..self.pos_len])) orelse return null; 222 if (std.mem.order(u8, key, self.pos[0..self.pos_len]) == .gt) { 223 @memcpy(self.pos[0..key.len], key); 224 self.pos_len = key.len; 225 i = 0; 226 continue; 227 } 228 i += 1; 229 } 230 self.returned = true; 231 return self.pos[0..self.pos_len]; 232 } 233 }; 234 } 235 236 // the unique lowercased words of `text`, in the order they first appear. 237 fn tokenize(aa: std.mem.Allocator, text: []const u8) ![]const []const u8 { 238 var words: std.StringArrayHashMapUnmanaged(void) = .empty; 239 const limit = @min(text.len, max_indexed_bytes); 240 var i: usize = 0; 241 while (i < limit) { 242 if (!isWordByte(text[i])) { 243 i += 1; 244 continue; 245 } 246 const start = i; 247 while (i < limit and isWordByte(text[i])) i += 1; 248 const word = text[start..@min(i, start + max_word_len)]; 249 if (word.len < min_word_len) continue; 250 const lowered = try aa.alloc(u8, word.len); 251 _ = std.ascii.lowerString(lowered, word); 252 try words.put(aa, lowered, {}); 253 if (words.count() == max_words_per_doc) break; 254 } 255 return words.keys(); 256 } 257 258 // a word is a run of ascii alphanumerics and non-ascii bytes. 259 fn isWordByte(byte: u8) bool { 260 return std.ascii.isAlphanumeric(byte) or byte >= 0x80; 261 } 262