Branch data Line data Source code
1 : : /*
2 : : * This file is part of the MicroPython project, http://micropython.org/
3 : : *
4 : : * The MIT License (MIT)
5 : : *
6 : : * Copyright (c) 2016 Paul Sokolovsky
7 : : *
8 : : * Permission is hereby granted, free of charge, to any person obtaining a copy
9 : : * of this software and associated documentation files (the "Software"), to deal
10 : : * in the Software without restriction, including without limitation the rights
11 : : * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12 : : * copies of the Software, and to permit persons to whom the Software is
13 : : * furnished to do so, subject to the following conditions:
14 : : *
15 : : * The above copyright notice and this permission notice shall be included in
16 : : * all copies or substantial portions of the Software.
17 : : *
18 : : * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 : : * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 : : * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 : : * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 : : * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 : : * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
24 : : * THE SOFTWARE.
25 : : */
26 : :
27 : : #include "py/runtime.h"
28 : : #include "py/stream.h"
29 : :
30 : : #if MICROPY_PY_BTREE
31 : :
32 : : #include <stdio.h>
33 : : #include <errno.h> // for declaration of global errno variable
34 : : #include <fcntl.h>
35 : :
36 : : // Undefine queue macros that will be defined in berkeley-db-1.xx headers
37 : : // below, in case they clash with system ones defined in headers above.
38 : : #undef LIST_HEAD
39 : : #undef LIST_ENTRY
40 : : #undef LIST_INIT
41 : : #undef LIST_INSERT_AFTER
42 : : #undef LIST_INSERT_HEAD
43 : : #undef LIST_REMOVE
44 : : #undef TAILQ_HEAD
45 : : #undef TAILQ_ENTRY
46 : : #undef TAILQ_INIT
47 : : #undef TAILQ_INSERT_HEAD
48 : : #undef TAILQ_INSERT_TAIL
49 : : #undef TAILQ_INSERT_AFTER
50 : : #undef TAILQ_REMOVE
51 : : #undef CIRCLEQ_HEAD
52 : : #undef CIRCLEQ_ENTRY
53 : : #undef CIRCLEQ_INIT
54 : : #undef CIRCLEQ_INSERT_AFTER
55 : : #undef CIRCLEQ_INSERT_BEFORE
56 : : #undef CIRCLEQ_INSERT_HEAD
57 : : #undef CIRCLEQ_INSERT_TAIL
58 : : #undef CIRCLEQ_REMOVE
59 : :
60 : : #include "berkeley-db/db.h"
61 : : #include "berkeley-db/btree.h"
62 : :
63 : : typedef struct _mp_obj_btree_t {
64 : : mp_obj_base_t base;
65 : : mp_obj_t stream; // retain a reference to prevent GC from reclaiming it
66 : : DB *db;
67 : : mp_obj_t start_key;
68 : : mp_obj_t end_key;
69 : : #define FLAG_END_KEY_INCL 1
70 : : #define FLAG_DESC 2
71 : : #define FLAG_ITER_TYPE_MASK 0xc0
72 : : #define FLAG_ITER_KEYS 0x40
73 : : #define FLAG_ITER_VALUES 0x80
74 : : #define FLAG_ITER_ITEMS 0xc0
75 : : byte flags;
76 : : byte next_flags;
77 : : } mp_obj_btree_t;
78 : :
79 : : #if !MICROPY_ENABLE_DYNRUNTIME
80 : : static const mp_obj_type_t btree_type;
81 : : #endif
82 : :
83 : : #define CHECK_ERROR(res) \
84 : : if (res == RET_ERROR) { \
85 : : mp_raise_OSError(errno); \
86 : : }
87 : :
88 : 0 : void __dbpanic(DB *db) {
89 : 0 : mp_printf(&mp_plat_print, "__dbpanic(%p)\n", db);
90 : 0 : }
91 : :
92 : 440 : static void check_btree_is_open(mp_obj_btree_t *self) {
93 [ + + ]: 440 : if (!self->db) {
94 : 8 : mp_raise_ValueError(MP_ERROR_TEXT("database closed"));
95 : : }
96 : 432 : }
97 : :
98 : 6 : static mp_obj_btree_t *btree_new(DB *db, mp_obj_t stream) {
99 : 6 : mp_obj_btree_t *o = mp_obj_malloc(mp_obj_btree_t, (mp_obj_type_t *)&btree_type);
100 : 6 : o->stream = stream;
101 : 6 : o->db = db;
102 : 6 : o->start_key = mp_const_none;
103 : 6 : o->end_key = mp_const_none;
104 : 6 : o->next_flags = 0;
105 : 6 : return o;
106 : : }
107 : :
108 : 548 : static void buf_to_dbt(mp_obj_t obj, DBT *dbt) {
109 : 548 : mp_buffer_info_t bufinfo;
110 : 548 : mp_get_buffer_raise(obj, &bufinfo, MP_BUFFER_READ);
111 : 548 : dbt->data = bufinfo.buf;
112 : 548 : dbt->size = bufinfo.len;
113 : 548 : }
114 : :
115 : 4 : static void btree_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
116 : 4 : (void)kind;
117 : 4 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
118 : 4 : mp_printf(print, "<btree %p>", self->db);
119 : 4 : }
120 : :
121 : 4 : static mp_obj_t btree_flush(mp_obj_t self_in) {
122 : 4 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
123 : 4 : check_btree_is_open(self);
124 : 2 : return MP_OBJ_NEW_SMALL_INT(__bt_sync(self->db, 0));
125 : : }
126 : : static MP_DEFINE_CONST_FUN_OBJ_1(btree_flush_obj, btree_flush);
127 : :
128 : 8 : static mp_obj_t btree_close(mp_obj_t self_in) {
129 : 8 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
130 : 8 : int res;
131 [ + + ]: 8 : if (self->db) {
132 : 6 : res = __bt_close(self->db);
133 : 6 : self->db = NULL;
134 : : } else {
135 : : res = RET_SUCCESS; // Closing an already-closed DB always succeeds.
136 : : }
137 : 8 : return MP_OBJ_NEW_SMALL_INT(res);
138 : : }
139 : : static MP_DEFINE_CONST_FUN_OBJ_1(btree_close_obj, btree_close);
140 : :
141 : 2 : static mp_obj_t btree_put(size_t n_args, const mp_obj_t *args) {
142 : 2 : (void)n_args;
143 : 2 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
144 : 2 : check_btree_is_open(self);
145 : 2 : DBT key, val;
146 : 2 : buf_to_dbt(args[1], &key);
147 : 2 : buf_to_dbt(args[2], &val);
148 : 2 : return MP_OBJ_NEW_SMALL_INT(__bt_put(self->db, &key, &val, 0));
149 : : }
150 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_put_obj, 3, 4, btree_put);
151 : :
152 : 4 : static mp_obj_t btree_get(size_t n_args, const mp_obj_t *args) {
153 : 4 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
154 : 4 : check_btree_is_open(self);
155 : 4 : DBT key, val;
156 : 4 : buf_to_dbt(args[1], &key);
157 : 4 : int res = __bt_get(self->db, &key, &val, 0);
158 [ + - ]: 4 : if (res == RET_SPECIAL) {
159 [ + + ]: 4 : if (n_args > 2) {
160 : 2 : return args[2];
161 : : } else {
162 : : return mp_const_none;
163 : : }
164 : : }
165 [ # # ]: 0 : CHECK_ERROR(res);
166 : 0 : return mp_obj_new_bytes(val.data, val.size);
167 : : }
168 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_get_obj, 2, 3, btree_get);
169 : :
170 : 6 : static mp_obj_t btree_seq(size_t n_args, const mp_obj_t *args) {
171 : 6 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
172 : 6 : check_btree_is_open(self);
173 : 6 : int flags = MP_OBJ_SMALL_INT_VALUE(args[1]);
174 : 6 : DBT key, val;
175 [ + + ]: 6 : if (n_args > 2) {
176 : 4 : buf_to_dbt(args[2], &key);
177 : : }
178 : :
179 : 6 : int res = __bt_seq(self->db, &key, &val, flags);
180 [ + + ]: 6 : CHECK_ERROR(res);
181 [ + + ]: 4 : if (res == RET_SPECIAL) {
182 : : return mp_const_none;
183 : : }
184 : :
185 : 2 : mp_obj_t pair_o = mp_obj_new_tuple(2, NULL);
186 : 2 : mp_obj_tuple_t *pair = MP_OBJ_TO_PTR(pair_o);
187 : 2 : pair->items[0] = mp_obj_new_bytes(key.data, key.size);
188 : 2 : pair->items[1] = mp_obj_new_bytes(val.data, val.size);
189 : 2 : return pair_o;
190 : : }
191 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_seq_obj, 2, 4, btree_seq);
192 : :
193 : 20 : static mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type) {
194 : 20 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
195 : 20 : self->next_flags = type;
196 : 20 : self->start_key = mp_const_none;
197 : 20 : self->end_key = mp_const_none;
198 [ + + ]: 20 : if (n_args > 1) {
199 : 12 : self->start_key = args[1];
200 [ + + ]: 12 : if (n_args > 2) {
201 : 10 : self->end_key = args[2];
202 [ + + ]: 10 : if (n_args > 3) {
203 : 4 : self->next_flags = type | MP_OBJ_SMALL_INT_VALUE(args[3]);
204 : : }
205 : : }
206 : : }
207 : 20 : return args[0];
208 : : }
209 : :
210 : 2 : static mp_obj_t btree_keys(size_t n_args, const mp_obj_t *args) {
211 : 2 : return btree_init_iter(n_args, args, FLAG_ITER_KEYS);
212 : : }
213 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_keys_obj, 1, 4, btree_keys);
214 : :
215 : 2 : static mp_obj_t btree_values(size_t n_args, const mp_obj_t *args) {
216 : 2 : return btree_init_iter(n_args, args, FLAG_ITER_VALUES);
217 : : }
218 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_values_obj, 1, 4, btree_values);
219 : :
220 : 16 : static mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) {
221 : 16 : return btree_init_iter(n_args, args, FLAG_ITER_ITEMS);
222 : : }
223 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items);
224 : :
225 : 22 : static mp_obj_t btree_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
226 : 22 : (void)iter_buf;
227 : 22 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
228 [ + + ]: 22 : if (self->next_flags != 0) {
229 : : // If we're called immediately after keys(), values(), or items(),
230 : : // use their setup for iteration.
231 : 20 : self->flags = self->next_flags;
232 : 20 : self->next_flags = 0;
233 : : } else {
234 : : // Otherwise, iterate over all keys.
235 : 2 : self->flags = FLAG_ITER_KEYS;
236 : 2 : self->start_key = mp_const_none;
237 : 2 : self->end_key = mp_const_none;
238 : : }
239 : :
240 : 22 : return self_in;
241 : : }
242 : :
243 : 72 : static mp_obj_t btree_iternext(mp_obj_t self_in) {
244 : 72 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
245 : 72 : check_btree_is_open(self);
246 : 70 : DBT key, val;
247 : 70 : int res;
248 : 70 : bool desc = self->flags & FLAG_DESC;
249 [ + + ]: 70 : if (self->start_key != MP_OBJ_NULL) {
250 : 20 : int flags = R_FIRST;
251 [ + + ]: 20 : if (self->start_key != mp_const_none) {
252 : 6 : buf_to_dbt(self->start_key, &key);
253 : 6 : flags = R_CURSOR;
254 [ + + ]: 14 : } else if (desc) {
255 : 2 : flags = R_LAST;
256 : : }
257 : 20 : res = __bt_seq(self->db, &key, &val, flags);
258 : 20 : self->start_key = MP_OBJ_NULL;
259 : : } else {
260 [ + + ]: 94 : res = __bt_seq(self->db, &key, &val, desc ? R_PREV : R_NEXT);
261 : : }
262 : :
263 [ + + ]: 70 : if (res == RET_SPECIAL) {
264 : : return MP_OBJ_STOP_ITERATION;
265 : : }
266 [ - + ]: 54 : CHECK_ERROR(res);
267 : :
268 [ + + ]: 54 : if (self->end_key != mp_const_none) {
269 : 14 : DBT end_key;
270 : 14 : buf_to_dbt(self->end_key, &end_key);
271 : 14 : BTREE *t = self->db->internal;
272 : 14 : int cmp = t->bt_cmp(&key, &end_key);
273 [ - + ]: 14 : if (desc) {
274 : 0 : cmp = -cmp;
275 : : }
276 [ + + ]: 14 : if (self->flags & FLAG_END_KEY_INCL) {
277 : 4 : cmp--;
278 : : }
279 [ + + ]: 14 : if (cmp >= 0) {
280 : 4 : self->end_key = MP_OBJ_NULL;
281 : 4 : return MP_OBJ_STOP_ITERATION;
282 : : }
283 : : }
284 : :
285 [ + + + ]: 50 : switch (self->flags & FLAG_ITER_TYPE_MASK) {
286 : 12 : case FLAG_ITER_KEYS:
287 : 12 : return mp_obj_new_bytes(key.data, key.size);
288 : 6 : case FLAG_ITER_VALUES:
289 : 6 : return mp_obj_new_bytes(val.data, val.size);
290 : 32 : default: {
291 : 32 : mp_obj_t pair_o = mp_obj_new_tuple(2, NULL);
292 : 32 : mp_obj_tuple_t *pair = MP_OBJ_TO_PTR(pair_o);
293 : 32 : pair->items[0] = mp_obj_new_bytes(key.data, key.size);
294 : 32 : pair->items[1] = mp_obj_new_bytes(val.data, val.size);
295 : 32 : return pair_o;
296 : : }
297 : : }
298 : : }
299 : :
300 : 344 : static mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
301 : 344 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
302 : 344 : check_btree_is_open(self);
303 [ + + ]: 340 : if (value == MP_OBJ_NULL) {
304 : : // delete
305 : 4 : DBT key;
306 : 4 : buf_to_dbt(index, &key);
307 : 4 : int res = __bt_delete(self->db, &key, 0);
308 [ + + ]: 4 : if (res == RET_SPECIAL) {
309 : 2 : mp_raise_type(&mp_type_KeyError);
310 : : }
311 [ - + ]: 2 : CHECK_ERROR(res);
312 : 2 : return mp_const_none;
313 [ + + ]: 336 : } else if (value == MP_OBJ_SENTINEL) {
314 : : // load
315 : 166 : DBT key, val;
316 : 166 : buf_to_dbt(index, &key);
317 : 166 : int res = __bt_get(self->db, &key, &val, 0);
318 [ + + ]: 166 : if (res == RET_SPECIAL) {
319 : 2 : mp_raise_type(&mp_type_KeyError);
320 : : }
321 [ - + ]: 164 : CHECK_ERROR(res);
322 : 164 : return mp_obj_new_bytes(val.data, val.size);
323 : : } else {
324 : : // store
325 : 170 : DBT key, val;
326 : 170 : buf_to_dbt(index, &key);
327 : 170 : buf_to_dbt(value, &val);
328 : 170 : int res = __bt_put(self->db, &key, &val, 0);
329 [ - + ]: 170 : CHECK_ERROR(res);
330 : 170 : return mp_const_none;
331 : : }
332 : : }
333 : :
334 : 8 : static mp_obj_t btree_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
335 : 8 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(lhs_in);
336 : 8 : check_btree_is_open(self);
337 [ + + ]: 8 : switch (op) {
338 : 6 : case MP_BINARY_OP_CONTAINS: {
339 : 6 : DBT key, val;
340 : 6 : buf_to_dbt(rhs_in, &key);
341 : 6 : int res = __bt_get(self->db, &key, &val, 0);
342 [ - + ]: 6 : CHECK_ERROR(res);
343 [ + + ]: 8 : return mp_obj_new_bool(res != RET_SPECIAL);
344 : : }
345 : : default:
346 : : // op not supported
347 : : return MP_OBJ_NULL;
348 : : }
349 : : }
350 : :
351 : : #if !MICROPY_ENABLE_DYNRUNTIME
352 : : static const mp_rom_map_elem_t btree_locals_dict_table[] = {
353 : : { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&btree_close_obj) },
354 : : { MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&btree_flush_obj) },
355 : : { MP_ROM_QSTR(MP_QSTR_get), MP_ROM_PTR(&btree_get_obj) },
356 : : { MP_ROM_QSTR(MP_QSTR_put), MP_ROM_PTR(&btree_put_obj) },
357 : : { MP_ROM_QSTR(MP_QSTR_seq), MP_ROM_PTR(&btree_seq_obj) },
358 : : { MP_ROM_QSTR(MP_QSTR_keys), MP_ROM_PTR(&btree_keys_obj) },
359 : : { MP_ROM_QSTR(MP_QSTR_values), MP_ROM_PTR(&btree_values_obj) },
360 : : { MP_ROM_QSTR(MP_QSTR_items), MP_ROM_PTR(&btree_items_obj) },
361 : : };
362 : :
363 : : static MP_DEFINE_CONST_DICT(btree_locals_dict, btree_locals_dict_table);
364 : :
365 : : static const mp_getiter_iternext_custom_t btree_getiter_iternext = {
366 : : .getiter = btree_getiter,
367 : : .iternext = btree_iternext,
368 : : };
369 : :
370 : : static MP_DEFINE_CONST_OBJ_TYPE(
371 : : btree_type,
372 : : MP_QSTR_btree,
373 : : MP_TYPE_FLAG_ITER_IS_CUSTOM,
374 : : // Save on qstr's, reuse same as for module
375 : : print, btree_print,
376 : : iter, &btree_getiter_iternext,
377 : : binary_op, btree_binary_op,
378 : : subscr, btree_subscr,
379 : : locals_dict, &btree_locals_dict
380 : : );
381 : : #endif
382 : :
383 : : static const FILEVTABLE btree_stream_fvtable = {
384 : : mp_stream_posix_read,
385 : : mp_stream_posix_write,
386 : : mp_stream_posix_lseek,
387 : : mp_stream_posix_fsync
388 : : };
389 : :
390 : : #if !MICROPY_ENABLE_DYNRUNTIME
391 : 12 : static mp_obj_t mod_btree_open(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
392 : 12 : static const mp_arg_t allowed_args[] = {
393 : : { MP_QSTR_flags, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
394 : : { MP_QSTR_cachesize, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
395 : : { MP_QSTR_pagesize, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
396 : : { MP_QSTR_minkeypage, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
397 : : };
398 : :
399 : : // Make sure we got a stream object
400 : 12 : mp_get_stream_raise(pos_args[0], MP_STREAM_OP_READ | MP_STREAM_OP_WRITE | MP_STREAM_OP_IOCTL);
401 : :
402 : 12 : struct {
403 : : mp_arg_val_t flags;
404 : : mp_arg_val_t cachesize;
405 : : mp_arg_val_t pagesize;
406 : : mp_arg_val_t minkeypage;
407 : : } args;
408 : 12 : mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args,
409 : : MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t *)&args);
410 : 12 : BTREEINFO openinfo = {0};
411 : 12 : openinfo.flags = args.flags.u_int;
412 : 12 : openinfo.cachesize = args.cachesize.u_int;
413 : 12 : openinfo.psize = args.pagesize.u_int;
414 : 12 : openinfo.minkeypage = args.minkeypage.u_int;
415 : :
416 : 12 : DB *db = __bt_open(MP_OBJ_TO_PTR(pos_args[0]), &btree_stream_fvtable, &openinfo, /*dflags*/ 0);
417 [ + + ]: 12 : if (db == NULL) {
418 : 6 : mp_raise_OSError(errno);
419 : : }
420 : 6 : return MP_OBJ_FROM_PTR(btree_new(db, pos_args[0]));
421 : : }
422 : : static MP_DEFINE_CONST_FUN_OBJ_KW(mod_btree_open_obj, 1, mod_btree_open);
423 : :
424 : : static const mp_rom_map_elem_t mp_module_btree_globals_table[] = {
425 : : { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_btree) },
426 : : { MP_ROM_QSTR(MP_QSTR_open), MP_ROM_PTR(&mod_btree_open_obj) },
427 : : { MP_ROM_QSTR(MP_QSTR_INCL), MP_ROM_INT(FLAG_END_KEY_INCL) },
428 : : { MP_ROM_QSTR(MP_QSTR_DESC), MP_ROM_INT(FLAG_DESC) },
429 : : };
430 : :
431 : : static MP_DEFINE_CONST_DICT(mp_module_btree_globals, mp_module_btree_globals_table);
432 : :
433 : : const mp_obj_module_t mp_module_btree = {
434 : : .base = { &mp_type_module },
435 : : .globals = (mp_obj_dict_t *)&mp_module_btree_globals,
436 : : };
437 : :
438 : : MP_REGISTER_MODULE(MP_QSTR_btree, mp_module_btree);
439 : : #endif
440 : :
441 : : #endif // MICROPY_PY_BTREE
|
