Blame


1 e2b5610c 2022-04-11 op /*
2 e2b5610c 2022-04-11 op * Copyright (c) 2022 Omar Polo <op@openbsd.org>
3 e2b5610c 2022-04-11 op *
4 e2b5610c 2022-04-11 op * Permission to use, copy, modify, and distribute this software for any
5 e2b5610c 2022-04-11 op * purpose with or without fee is hereby granted, provided that the above
6 e2b5610c 2022-04-11 op * copyright notice and this permission notice appear in all copies.
7 e2b5610c 2022-04-11 op *
8 e2b5610c 2022-04-11 op * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9 e2b5610c 2022-04-11 op * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10 e2b5610c 2022-04-11 op * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11 e2b5610c 2022-04-11 op * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12 e2b5610c 2022-04-11 op * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13 e2b5610c 2022-04-11 op * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14 e2b5610c 2022-04-11 op * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15 e2b5610c 2022-04-11 op */
16 e2b5610c 2022-04-11 op
17 e2b5610c 2022-04-11 op #include <sys/mman.h>
18 e2b5610c 2022-04-11 op
19 e2b5610c 2022-04-11 op #include <stdio.h>
20 e2b5610c 2022-04-11 op #include <stdint.h>
21 e2b5610c 2022-04-11 op #include <stdlib.h>
22 e2b5610c 2022-04-11 op #include <string.h>
23 e2b5610c 2022-04-11 op #include <unistd.h>
24 e2b5610c 2022-04-11 op
25 e2b5610c 2022-04-11 op #include "db.h"
26 e2b5610c 2022-04-11 op #include "dictionary.h"
27 e2b5610c 2022-04-11 op
28 e2b5610c 2022-04-11 op #define IDX_ENTRY_SIZE (DB_WORDLEN + sizeof(int64_t))
29 e2b5610c 2022-04-11 op
30 e2b5610c 2022-04-11 op static int
31 e2b5610c 2022-04-11 op write_dictionary(FILE *fp, struct dictionary *dict)
32 e2b5610c 2022-04-11 op {
33 e2b5610c 2022-04-11 op off_t start;
34 e2b5610c 2022-04-11 op uint64_t pos;
35 e2b5610c 2022-04-11 op uint32_t n;
36 e2b5610c 2022-04-11 op size_t i, len;
37 e2b5610c 2022-04-11 op
38 e2b5610c 2022-04-11 op if ((uint64_t)dict->len > UINT32_MAX)
39 e2b5610c 2022-04-11 op return -1;
40 e2b5610c 2022-04-11 op
41 e2b5610c 2022-04-11 op n = dict->len;
42 e2b5610c 2022-04-11 op if (fwrite(&n, sizeof(n), 1, fp) != 1)
43 e2b5610c 2022-04-11 op return -1;
44 e2b5610c 2022-04-11 op
45 e2b5610c 2022-04-11 op if ((start = ftello(fp)) == -1)
46 e2b5610c 2022-04-11 op return -1;
47 e2b5610c 2022-04-11 op
48 e2b5610c 2022-04-11 op len = DB_WORDLEN + sizeof(int64_t);
49 e2b5610c 2022-04-11 op pos = start + (n * len);
50 e2b5610c 2022-04-11 op for (i = 0; i < dict->len; ++i) {
51 e2b5610c 2022-04-11 op char word[DB_WORDLEN];
52 e2b5610c 2022-04-11 op
53 e2b5610c 2022-04-11 op memset(word, 0, sizeof(word));
54 e2b5610c 2022-04-11 op strlcpy(word, dict->entries[i].word, sizeof(word));
55 e2b5610c 2022-04-11 op if (fwrite(word, sizeof(word), 1, fp) != 1)
56 e2b5610c 2022-04-11 op return -1;
57 e2b5610c 2022-04-11 op
58 e2b5610c 2022-04-11 op if (fwrite(&pos, sizeof(pos), 1, fp) != 1)
59 e2b5610c 2022-04-11 op return -1;
60 e2b5610c 2022-04-11 op
61 e2b5610c 2022-04-11 op /* one for the len */
62 e2b5610c 2022-04-11 op pos += sizeof(uint32_t) * (dict->entries[i].len + 1);
63 e2b5610c 2022-04-11 op }
64 e2b5610c 2022-04-11 op
65 e2b5610c 2022-04-11 op for (i = 0; i < dict->len; ++i) {
66 e2b5610c 2022-04-11 op size_t j;
67 e2b5610c 2022-04-11 op uint32_t t, x;
68 e2b5610c 2022-04-11 op
69 e2b5610c 2022-04-11 op x = dict->entries[i].len;
70 e2b5610c 2022-04-11 op if (fwrite(&x, sizeof(x), 1, fp) != 1)
71 e2b5610c 2022-04-11 op return -1;
72 e2b5610c 2022-04-11 op
73 e2b5610c 2022-04-11 op for (j = 0; j < x; ++j) {
74 e2b5610c 2022-04-11 op t = dict->entries[i].ids[j];
75 e2b5610c 2022-04-11 op if (fwrite(&t, sizeof(t), 1, fp) != 1)
76 e2b5610c 2022-04-11 op return -1;
77 e2b5610c 2022-04-11 op }
78 e2b5610c 2022-04-11 op }
79 e2b5610c 2022-04-11 op
80 e2b5610c 2022-04-11 op return 0;
81 e2b5610c 2022-04-11 op }
82 e2b5610c 2022-04-11 op
83 e2b5610c 2022-04-11 op int
84 e2b5610c 2022-04-11 op db_create(FILE *fp, struct dictionary *dict, struct db_entry *entries,
85 e2b5610c 2022-04-11 op size_t n)
86 e2b5610c 2022-04-11 op {
87 e2b5610c 2022-04-11 op int64_t endidx;
88 e2b5610c 2022-04-11 op size_t i;
89 e2b5610c 2022-04-11 op uint32_t version = DB_VERSION;
90 e2b5610c 2022-04-11 op
91 e2b5610c 2022-04-11 op if (n > INT32_MAX)
92 e2b5610c 2022-04-11 op return -1;
93 e2b5610c 2022-04-11 op
94 e2b5610c 2022-04-11 op if (fwrite(&version, sizeof(version), 1, fp) != 1)
95 e2b5610c 2022-04-11 op return -1;
96 e2b5610c 2022-04-11 op
97 e2b5610c 2022-04-11 op /* reserve space for the start pointer -- filled later */
98 e2b5610c 2022-04-11 op if (fseek(fp, sizeof(int64_t), SEEK_CUR) == -1)
99 e2b5610c 2022-04-11 op return -1;
100 e2b5610c 2022-04-11 op
101 e2b5610c 2022-04-11 op if (write_dictionary(fp, dict) == -1)
102 e2b5610c 2022-04-11 op return -1;
103 e2b5610c 2022-04-11 op
104 e2b5610c 2022-04-11 op if ((endidx = ftello(fp)) == -1)
105 e2b5610c 2022-04-11 op return -1;
106 e2b5610c 2022-04-11 op
107 e2b5610c 2022-04-11 op for (i = 0; i < n; ++i) {
108 e2b5610c 2022-04-11 op uint16_t namelen, descrlen = 0;
109 e2b5610c 2022-04-11 op
110 e2b5610c 2022-04-11 op namelen = strlen(entries[i].name);
111 e2b5610c 2022-04-11 op if (entries[i].descr != NULL)
112 e2b5610c 2022-04-11 op descrlen = strlen(entries[i].descr);
113 e2b5610c 2022-04-11 op
114 e2b5610c 2022-04-11 op if (fwrite(&namelen, sizeof(namelen), 1, fp) != 1)
115 e2b5610c 2022-04-11 op return -1;
116 e2b5610c 2022-04-11 op if (fwrite(entries[i].name, namelen+1, 1, fp) != 1)
117 e2b5610c 2022-04-11 op return -1;
118 e2b5610c 2022-04-11 op
119 e2b5610c 2022-04-11 op if (fwrite(&descrlen, sizeof(descrlen), 1, fp) != 1)
120 e2b5610c 2022-04-11 op return -1;
121 e2b5610c 2022-04-11 op if (descrlen > 0 &&
122 e2b5610c 2022-04-11 op fwrite(entries[i].descr, descrlen, 1, fp) != 1)
123 e2b5610c 2022-04-11 op return -1;
124 e2b5610c 2022-04-11 op if (fwrite("", 1, 1, fp) != 1)
125 e2b5610c 2022-04-11 op return -1;
126 e2b5610c 2022-04-11 op }
127 e2b5610c 2022-04-11 op
128 e2b5610c 2022-04-11 op if (fseek(fp, sizeof(version), SEEK_SET) == -1)
129 e2b5610c 2022-04-11 op return -1;
130 e2b5610c 2022-04-11 op
131 e2b5610c 2022-04-11 op if (fwrite(&endidx, sizeof(endidx), 1, fp) != 1)
132 e2b5610c 2022-04-11 op return -1;
133 e2b5610c 2022-04-11 op
134 e2b5610c 2022-04-11 op return 0;
135 e2b5610c 2022-04-11 op }
136 e2b5610c 2022-04-11 op
137 e2b5610c 2022-04-11 op static int
138 e2b5610c 2022-04-11 op initdb(struct db *db)
139 e2b5610c 2022-04-11 op {
140 e2b5610c 2022-04-11 op off_t hdrlen = sizeof(uint32_t) + sizeof(int64_t) + sizeof(uint32_t);
141 e2b5610c 2022-04-11 op int64_t end_off;
142 e2b5610c 2022-04-11 op uint8_t *p = db->m;
143 e2b5610c 2022-04-11 op
144 e2b5610c 2022-04-11 op if (hdrlen > db->len)
145 e2b5610c 2022-04-11 op return -1;
146 e2b5610c 2022-04-11 op
147 e2b5610c 2022-04-11 op memcpy(&db->version, p, sizeof(db->version));
148 e2b5610c 2022-04-11 op p += sizeof(db->version);
149 e2b5610c 2022-04-11 op
150 e2b5610c 2022-04-11 op memcpy(&end_off, p, sizeof(end_off));
151 e2b5610c 2022-04-11 op p += sizeof(end_off);
152 e2b5610c 2022-04-11 op
153 e2b5610c 2022-04-11 op memcpy(&db->nwords, p, sizeof(db->nwords));
154 e2b5610c 2022-04-11 op p += sizeof(db->nwords);
155 e2b5610c 2022-04-11 op
156 e2b5610c 2022-04-11 op db->idx_start = p;
157 e2b5610c 2022-04-11 op db->idx_end = p + db->nwords * IDX_ENTRY_SIZE;
158 e2b5610c 2022-04-11 op db->list_start = db->idx_end;
159 e2b5610c 2022-04-11 op db->list_end = db->m + end_off;
160 e2b5610c 2022-04-11 op db->docs_start = db->list_end;
161 e2b5610c 2022-04-11 op db->docs_end = db->m + db->len;
162 e2b5610c 2022-04-11 op
163 e2b5610c 2022-04-11 op if (db->idx_end > db->docs_end)
164 e2b5610c 2022-04-11 op return -1;
165 e2b5610c 2022-04-11 op if (db->list_end > db->docs_end)
166 e2b5610c 2022-04-11 op return -1;
167 e2b5610c 2022-04-11 op
168 e2b5610c 2022-04-11 op return 0;
169 e2b5610c 2022-04-11 op }
170 e2b5610c 2022-04-11 op
171 e2b5610c 2022-04-11 op int
172 e2b5610c 2022-04-11 op db_open(struct db *db, int fd)
173 e2b5610c 2022-04-11 op {
174 e2b5610c 2022-04-11 op memset(db, 0, sizeof(*db));
175 e2b5610c 2022-04-11 op
176 e2b5610c 2022-04-11 op if ((db->len = lseek(fd, 0, SEEK_END)) == -1)
177 e2b5610c 2022-04-11 op return -1;
178 e2b5610c 2022-04-11 op
179 e2b5610c 2022-04-11 op if (lseek(fd, 0, SEEK_SET) == -1)
180 e2b5610c 2022-04-11 op return -1;
181 e2b5610c 2022-04-11 op
182 e2b5610c 2022-04-11 op db->m = mmap(NULL, db->len, PROT_READ, MAP_PRIVATE, fd, 0);
183 e2b5610c 2022-04-11 op if (db->m == MAP_FAILED)
184 e2b5610c 2022-04-11 op return -1;
185 e2b5610c 2022-04-11 op
186 e2b5610c 2022-04-11 op if (initdb(db) == -1) {
187 e2b5610c 2022-04-11 op db_close(db);
188 e2b5610c 2022-04-11 op return -1;
189 e2b5610c 2022-04-11 op }
190 e2b5610c 2022-04-11 op
191 e2b5610c 2022-04-11 op return 0;
192 e2b5610c 2022-04-11 op }
193 e2b5610c 2022-04-11 op
194 e2b5610c 2022-04-11 op static int
195 e2b5610c 2022-04-11 op db_countdocs(struct db *db, struct db_entry *e, void *d)
196 e2b5610c 2022-04-11 op {
197 e2b5610c 2022-04-11 op struct db_stats *stats = d;
198 e2b5610c 2022-04-11 op
199 e2b5610c 2022-04-11 op stats->ndocs++;
200 e2b5610c 2022-04-11 op return 0;
201 e2b5610c 2022-04-11 op }
202 e2b5610c 2022-04-11 op
203 e2b5610c 2022-04-11 op static int
204 e2b5610c 2022-04-11 op db_idx_compar(const void *key, const void *elem)
205 e2b5610c 2022-04-11 op {
206 e2b5610c 2022-04-11 op const char *word = key;
207 e2b5610c 2022-04-11 op const char *idx_entry = elem;
208 e2b5610c 2022-04-11 op
209 e2b5610c 2022-04-11 op if (idx_entry[DB_WORDLEN-1] != '\0')
210 e2b5610c 2022-04-11 op return -1;
211 e2b5610c 2022-04-11 op return strcmp(word, idx_entry);
212 e2b5610c 2022-04-11 op }
213 e2b5610c 2022-04-11 op
214 e2b5610c 2022-04-11 op static inline uint32_t *
215 e2b5610c 2022-04-11 op db_getdocs(struct db *db, const uint8_t *entry, size_t *len)
216 e2b5610c 2022-04-11 op {
217 e2b5610c 2022-04-11 op int64_t pos;
218 e2b5610c 2022-04-11 op uint32_t l;
219 e2b5610c 2022-04-11 op
220 e2b5610c 2022-04-11 op entry += DB_WORDLEN;
221 e2b5610c 2022-04-11 op memcpy(&pos, entry, sizeof(pos));
222 e2b5610c 2022-04-11 op
223 e2b5610c 2022-04-11 op entry = db->m + pos;
224 e2b5610c 2022-04-11 op if (entry < db->list_start || entry > db->list_end)
225 e2b5610c 2022-04-11 op return NULL;
226 e2b5610c 2022-04-11 op
227 e2b5610c 2022-04-11 op memcpy(&l, entry, sizeof(l));
228 e2b5610c 2022-04-11 op entry += sizeof(l);
229 e2b5610c 2022-04-11 op *len = l;
230 e2b5610c 2022-04-11 op return (uint32_t *)entry;
231 e2b5610c 2022-04-11 op }
232 e2b5610c 2022-04-11 op
233 e2b5610c 2022-04-11 op uint32_t *
234 e2b5610c 2022-04-11 op db_word_docs(struct db *db, const char *word, size_t *len)
235 e2b5610c 2022-04-11 op {
236 e2b5610c 2022-04-11 op uint8_t *e;
237 e2b5610c 2022-04-11 op
238 e2b5610c 2022-04-11 op *len = 0;
239 e2b5610c 2022-04-11 op
240 e2b5610c 2022-04-11 op e = bsearch(word, db->idx_start, db->nwords, IDX_ENTRY_SIZE,
241 e2b5610c 2022-04-11 op db_idx_compar);
242 e2b5610c 2022-04-11 op if (e == NULL)
243 e2b5610c 2022-04-11 op return NULL;
244 e2b5610c 2022-04-11 op return db_getdocs(db, e, len);
245 e2b5610c 2022-04-11 op }
246 e2b5610c 2022-04-11 op
247 e2b5610c 2022-04-11 op int
248 e2b5610c 2022-04-11 op db_stats(struct db *db, struct db_stats *stats)
249 e2b5610c 2022-04-11 op {
250 e2b5610c 2022-04-11 op const uint8_t *p;
251 e2b5610c 2022-04-11 op size_t l, maxl = 0, idlen;
252 e2b5610c 2022-04-11 op
253 e2b5610c 2022-04-11 op memset(stats, 0, sizeof(*stats));
254 e2b5610c 2022-04-11 op
255 e2b5610c 2022-04-11 op if (db_listall(db, db_countdocs, stats) == -1)
256 e2b5610c 2022-04-11 op return -1;
257 e2b5610c 2022-04-11 op
258 e2b5610c 2022-04-11 op stats->nwords = db->nwords;
259 e2b5610c 2022-04-11 op
260 e2b5610c 2022-04-11 op p = db->idx_start;
261 e2b5610c 2022-04-11 op while (p < db->idx_end) {
262 e2b5610c 2022-04-11 op if (p + DB_WORDLEN > db->idx_end)
263 e2b5610c 2022-04-11 op return -1;
264 e2b5610c 2022-04-11 op
265 e2b5610c 2022-04-11 op if (p[DB_WORDLEN-1] != '\0')
266 e2b5610c 2022-04-11 op return -1;
267 e2b5610c 2022-04-11 op
268 e2b5610c 2022-04-11 op l = strlen(p);
269 e2b5610c 2022-04-11 op if (l > maxl) {
270 e2b5610c 2022-04-11 op maxl = l;
271 e2b5610c 2022-04-11 op stats->longest_word = p;
272 e2b5610c 2022-04-11 op }
273 e2b5610c 2022-04-11 op
274 e2b5610c 2022-04-11 op if (db_getdocs(db, p, &idlen) == NULL)
275 e2b5610c 2022-04-11 op return -1;
276 e2b5610c 2022-04-11 op
277 e2b5610c 2022-04-11 op if (idlen > stats->most_popular_ndocs) {
278 e2b5610c 2022-04-11 op stats->most_popular_ndocs = idlen;
279 e2b5610c 2022-04-11 op stats->most_popular = p;
280 e2b5610c 2022-04-11 op }
281 e2b5610c 2022-04-11 op
282 e2b5610c 2022-04-11 op p += IDX_ENTRY_SIZE;
283 e2b5610c 2022-04-11 op }
284 e2b5610c 2022-04-11 op
285 e2b5610c 2022-04-11 op return 0;
286 e2b5610c 2022-04-11 op }
287 e2b5610c 2022-04-11 op
288 e2b5610c 2022-04-11 op static inline uint8_t *
289 e2b5610c 2022-04-11 op db_extract_doc(struct db *db, uint8_t *p, struct db_entry *e)
290 e2b5610c 2022-04-11 op {
291 e2b5610c 2022-04-11 op uint16_t namelen, descrlen;
292 e2b5610c 2022-04-11 op
293 e2b5610c 2022-04-11 op /*
294 e2b5610c 2022-04-11 op * namelen[2] name[namelen]
295 e2b5610c 2022-04-11 op * descrlen[2] descr[descrlen]
296 e2b5610c 2022-04-11 op */
297 e2b5610c 2022-04-11 op
298 e2b5610c 2022-04-11 op if (p + 2 > db->docs_end)
299 e2b5610c 2022-04-11 op return NULL;
300 e2b5610c 2022-04-11 op memcpy(&namelen, p, sizeof(namelen));
301 e2b5610c 2022-04-11 op p += sizeof(namelen);
302 e2b5610c 2022-04-11 op
303 e2b5610c 2022-04-11 op if (p + namelen > db->docs_end || p[namelen] != '\0')
304 e2b5610c 2022-04-11 op return NULL;
305 e2b5610c 2022-04-11 op e->name = p;
306 e2b5610c 2022-04-11 op p += namelen + 1;
307 e2b5610c 2022-04-11 op
308 e2b5610c 2022-04-11 op if (p + 2 > db->docs_end)
309 e2b5610c 2022-04-11 op return NULL;
310 e2b5610c 2022-04-11 op memcpy(&descrlen, p, sizeof(descrlen));
311 e2b5610c 2022-04-11 op p += sizeof(descrlen);
312 e2b5610c 2022-04-11 op
313 e2b5610c 2022-04-11 op if (p + descrlen > db->docs_end || p[descrlen] != '\0')
314 e2b5610c 2022-04-11 op return NULL;
315 e2b5610c 2022-04-11 op e->descr = p;
316 e2b5610c 2022-04-11 op p += descrlen + 1;
317 e2b5610c 2022-04-11 op
318 e2b5610c 2022-04-11 op return p;
319 e2b5610c 2022-04-11 op }
320 e2b5610c 2022-04-11 op
321 e2b5610c 2022-04-11 op int
322 e2b5610c 2022-04-11 op db_listall(struct db *db, db_hit_cb cb, void *data)
323 e2b5610c 2022-04-11 op {
324 e2b5610c 2022-04-11 op uint8_t *p = db->docs_start;
325 e2b5610c 2022-04-11 op
326 e2b5610c 2022-04-11 op while (p < db->docs_end) {
327 e2b5610c 2022-04-11 op struct db_entry e;
328 e2b5610c 2022-04-11 op
329 e2b5610c 2022-04-11 op if ((p = db_extract_doc(db, p, &e)) == NULL)
330 e2b5610c 2022-04-11 op return -1;
331 e2b5610c 2022-04-11 op
332 e2b5610c 2022-04-11 op if (cb(db, &e, data) == -1)
333 e2b5610c 2022-04-11 op return -1;
334 e2b5610c 2022-04-11 op }
335 e2b5610c 2022-04-11 op
336 e2b5610c 2022-04-11 op return 0;
337 e2b5610c 2022-04-11 op }
338 e2b5610c 2022-04-11 op
339 e2b5610c 2022-04-11 op int
340 e2b5610c 2022-04-11 op db_doc_by_id(struct db *db, int docid, struct db_entry *e)
341 e2b5610c 2022-04-11 op {
342 e2b5610c 2022-04-11 op uint8_t *p = db->docs_start;
343 e2b5610c 2022-04-11 op int n = 0;
344 e2b5610c 2022-04-11 op
345 e2b5610c 2022-04-11 op while (p < db->docs_end) {
346 e2b5610c 2022-04-11 op if ((p = db_extract_doc(db, p, e)) == NULL)
347 e2b5610c 2022-04-11 op return -1;
348 e2b5610c 2022-04-11 op
349 e2b5610c 2022-04-11 op if (n == docid)
350 e2b5610c 2022-04-11 op return 0;
351 e2b5610c 2022-04-11 op
352 e2b5610c 2022-04-11 op n++;
353 e2b5610c 2022-04-11 op }
354 e2b5610c 2022-04-11 op
355 e2b5610c 2022-04-11 op return -1;
356 e2b5610c 2022-04-11 op }
357 e2b5610c 2022-04-11 op
358 e2b5610c 2022-04-11 op void
359 e2b5610c 2022-04-11 op db_close(struct db *db)
360 e2b5610c 2022-04-11 op {
361 e2b5610c 2022-04-11 op munmap(db->m, db->len);
362 e2b5610c 2022-04-11 op memset(db, 0, sizeof(*db));
363 e2b5610c 2022-04-11 op }