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 : 4 : static mp_obj_btree_t *btree_new(DB *db, mp_obj_t stream) {
93 : 4 : mp_obj_btree_t *o = mp_obj_malloc(mp_obj_btree_t, (mp_obj_type_t *)&btree_type);
94 : 4 : o->stream = stream;
95 : 4 : o->db = db;
96 : 4 : o->start_key = mp_const_none;
97 : 4 : o->end_key = mp_const_none;
98 : 4 : o->next_flags = 0;
99 : 4 : return o;
100 : : }
101 : :
102 : 544 : static void buf_to_dbt(mp_obj_t obj, DBT *dbt) {
103 : 544 : mp_buffer_info_t bufinfo;
104 : 544 : mp_get_buffer_raise(obj, &bufinfo, MP_BUFFER_READ);
105 : 544 : dbt->data = bufinfo.buf;
106 : 544 : dbt->size = bufinfo.len;
107 : 544 : }
108 : :
109 : 2 : static void btree_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
110 : 2 : (void)kind;
111 : 2 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
112 : 2 : mp_printf(print, "<btree %p>", self->db);
113 : 2 : }
114 : :
115 : 2 : static mp_obj_t btree_flush(mp_obj_t self_in) {
116 : 2 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
117 : 2 : return MP_OBJ_NEW_SMALL_INT(__bt_sync(self->db, 0));
118 : : }
119 : : static MP_DEFINE_CONST_FUN_OBJ_1(btree_flush_obj, btree_flush);
120 : :
121 : 4 : static mp_obj_t btree_close(mp_obj_t self_in) {
122 : 4 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
123 : 4 : return MP_OBJ_NEW_SMALL_INT(__bt_close(self->db));
124 : : }
125 : : static MP_DEFINE_CONST_FUN_OBJ_1(btree_close_obj, btree_close);
126 : :
127 : 2 : static mp_obj_t btree_put(size_t n_args, const mp_obj_t *args) {
128 : 2 : (void)n_args;
129 : 2 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
130 : 2 : DBT key, val;
131 : 2 : buf_to_dbt(args[1], &key);
132 : 2 : buf_to_dbt(args[2], &val);
133 : 2 : return MP_OBJ_NEW_SMALL_INT(__bt_put(self->db, &key, &val, 0));
134 : : }
135 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_put_obj, 3, 4, btree_put);
136 : :
137 : 4 : static mp_obj_t btree_get(size_t n_args, const mp_obj_t *args) {
138 : 4 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
139 : 4 : DBT key, val;
140 : 4 : buf_to_dbt(args[1], &key);
141 : 4 : int res = __bt_get(self->db, &key, &val, 0);
142 [ + - ]: 4 : if (res == RET_SPECIAL) {
143 [ + + ]: 4 : if (n_args > 2) {
144 : 2 : return args[2];
145 : : } else {
146 : : return mp_const_none;
147 : : }
148 : : }
149 [ # # ]: 0 : CHECK_ERROR(res);
150 : 0 : return mp_obj_new_bytes(val.data, val.size);
151 : : }
152 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_get_obj, 2, 3, btree_get);
153 : :
154 : 6 : static mp_obj_t btree_seq(size_t n_args, const mp_obj_t *args) {
155 : 6 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
156 : 6 : int flags = MP_OBJ_SMALL_INT_VALUE(args[1]);
157 : 6 : DBT key, val;
158 [ + + ]: 6 : if (n_args > 2) {
159 : 4 : buf_to_dbt(args[2], &key);
160 : : }
161 : :
162 : 6 : int res = __bt_seq(self->db, &key, &val, flags);
163 [ + + ]: 6 : CHECK_ERROR(res);
164 [ + + ]: 4 : if (res == RET_SPECIAL) {
165 : : return mp_const_none;
166 : : }
167 : :
168 : 2 : mp_obj_t pair_o = mp_obj_new_tuple(2, NULL);
169 : 2 : mp_obj_tuple_t *pair = MP_OBJ_TO_PTR(pair_o);
170 : 2 : pair->items[0] = mp_obj_new_bytes(key.data, key.size);
171 : 2 : pair->items[1] = mp_obj_new_bytes(val.data, val.size);
172 : 2 : return pair_o;
173 : : }
174 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_seq_obj, 2, 4, btree_seq);
175 : :
176 : 18 : static mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type) {
177 : 18 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
178 : 18 : self->next_flags = type;
179 : 18 : self->start_key = mp_const_none;
180 : 18 : self->end_key = mp_const_none;
181 [ + + ]: 18 : if (n_args > 1) {
182 : 12 : self->start_key = args[1];
183 [ + + ]: 12 : if (n_args > 2) {
184 : 10 : self->end_key = args[2];
185 [ + + ]: 10 : if (n_args > 3) {
186 : 4 : self->next_flags = type | MP_OBJ_SMALL_INT_VALUE(args[3]);
187 : : }
188 : : }
189 : : }
190 : 18 : return args[0];
191 : : }
192 : :
193 : 2 : static mp_obj_t btree_keys(size_t n_args, const mp_obj_t *args) {
194 : 2 : return btree_init_iter(n_args, args, FLAG_ITER_KEYS);
195 : : }
196 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_keys_obj, 1, 4, btree_keys);
197 : :
198 : 2 : static mp_obj_t btree_values(size_t n_args, const mp_obj_t *args) {
199 : 2 : return btree_init_iter(n_args, args, FLAG_ITER_VALUES);
200 : : }
201 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_values_obj, 1, 4, btree_values);
202 : :
203 : 14 : static mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) {
204 : 14 : return btree_init_iter(n_args, args, FLAG_ITER_ITEMS);
205 : : }
206 : : static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items);
207 : :
208 : 20 : static mp_obj_t btree_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
209 : 20 : (void)iter_buf;
210 : 20 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
211 [ + + ]: 20 : if (self->next_flags != 0) {
212 : : // If we're called immediately after keys(), values(), or items(),
213 : : // use their setup for iteration.
214 : 18 : self->flags = self->next_flags;
215 : 18 : self->next_flags = 0;
216 : : } else {
217 : : // Otherwise, iterate over all keys.
218 : 2 : self->flags = FLAG_ITER_KEYS;
219 : 2 : self->start_key = mp_const_none;
220 : 2 : self->end_key = mp_const_none;
221 : : }
222 : :
223 : 20 : return self_in;
224 : : }
225 : :
226 : 70 : static mp_obj_t btree_iternext(mp_obj_t self_in) {
227 : 70 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
228 : 70 : DBT key, val;
229 : 70 : int res;
230 : 70 : bool desc = self->flags & FLAG_DESC;
231 [ + + ]: 70 : if (self->start_key != MP_OBJ_NULL) {
232 : 20 : int flags = R_FIRST;
233 [ + + ]: 20 : if (self->start_key != mp_const_none) {
234 : 6 : buf_to_dbt(self->start_key, &key);
235 : 6 : flags = R_CURSOR;
236 [ + + ]: 14 : } else if (desc) {
237 : 2 : flags = R_LAST;
238 : : }
239 : 20 : res = __bt_seq(self->db, &key, &val, flags);
240 : 20 : self->start_key = MP_OBJ_NULL;
241 : : } else {
242 [ + + ]: 94 : res = __bt_seq(self->db, &key, &val, desc ? R_PREV : R_NEXT);
243 : : }
244 : :
245 [ + + ]: 70 : if (res == RET_SPECIAL) {
246 : : return MP_OBJ_STOP_ITERATION;
247 : : }
248 [ - + ]: 54 : CHECK_ERROR(res);
249 : :
250 [ + + ]: 54 : if (self->end_key != mp_const_none) {
251 : 14 : DBT end_key;
252 : 14 : buf_to_dbt(self->end_key, &end_key);
253 : 14 : BTREE *t = self->db->internal;
254 : 14 : int cmp = t->bt_cmp(&key, &end_key);
255 [ - + ]: 14 : if (desc) {
256 : 0 : cmp = -cmp;
257 : : }
258 [ + + ]: 14 : if (self->flags & FLAG_END_KEY_INCL) {
259 : 4 : cmp--;
260 : : }
261 [ + + ]: 14 : if (cmp >= 0) {
262 : 4 : self->end_key = MP_OBJ_NULL;
263 : 4 : return MP_OBJ_STOP_ITERATION;
264 : : }
265 : : }
266 : :
267 [ + + + ]: 50 : switch (self->flags & FLAG_ITER_TYPE_MASK) {
268 : 12 : case FLAG_ITER_KEYS:
269 : 12 : return mp_obj_new_bytes(key.data, key.size);
270 : 6 : case FLAG_ITER_VALUES:
271 : 6 : return mp_obj_new_bytes(val.data, val.size);
272 : 32 : default: {
273 : 32 : mp_obj_t pair_o = mp_obj_new_tuple(2, NULL);
274 : 32 : mp_obj_tuple_t *pair = MP_OBJ_TO_PTR(pair_o);
275 : 32 : pair->items[0] = mp_obj_new_bytes(key.data, key.size);
276 : 32 : pair->items[1] = mp_obj_new_bytes(val.data, val.size);
277 : 32 : return pair_o;
278 : : }
279 : : }
280 : : }
281 : :
282 : 338 : static mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
283 : 338 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
284 [ + + ]: 338 : if (value == MP_OBJ_NULL) {
285 : : // delete
286 : 4 : DBT key;
287 : 4 : buf_to_dbt(index, &key);
288 : 4 : int res = __bt_delete(self->db, &key, 0);
289 [ + + ]: 4 : if (res == RET_SPECIAL) {
290 : 2 : mp_raise_type(&mp_type_KeyError);
291 : : }
292 [ - + ]: 2 : CHECK_ERROR(res);
293 : 2 : return mp_const_none;
294 [ + + ]: 334 : } else if (value == MP_OBJ_SENTINEL) {
295 : : // load
296 : 166 : DBT key, val;
297 : 166 : buf_to_dbt(index, &key);
298 : 166 : int res = __bt_get(self->db, &key, &val, 0);
299 [ + + ]: 166 : if (res == RET_SPECIAL) {
300 : 2 : mp_raise_type(&mp_type_KeyError);
301 : : }
302 [ - + ]: 164 : CHECK_ERROR(res);
303 : 164 : return mp_obj_new_bytes(val.data, val.size);
304 : : } else {
305 : : // store
306 : 168 : DBT key, val;
307 : 168 : buf_to_dbt(index, &key);
308 : 168 : buf_to_dbt(value, &val);
309 : 168 : int res = __bt_put(self->db, &key, &val, 0);
310 [ - + ]: 168 : CHECK_ERROR(res);
311 : 168 : return mp_const_none;
312 : : }
313 : : }
314 : :
315 : 8 : static mp_obj_t btree_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
316 : 8 : mp_obj_btree_t *self = MP_OBJ_TO_PTR(lhs_in);
317 [ + + ]: 8 : switch (op) {
318 : 6 : case MP_BINARY_OP_CONTAINS: {
319 : 6 : DBT key, val;
320 : 6 : buf_to_dbt(rhs_in, &key);
321 : 6 : int res = __bt_get(self->db, &key, &val, 0);
322 [ - + ]: 6 : CHECK_ERROR(res);
323 [ + + ]: 8 : return mp_obj_new_bool(res != RET_SPECIAL);
324 : : }
325 : : default:
326 : : // op not supported
327 : : return MP_OBJ_NULL;
328 : : }
329 : : }
330 : :
331 : : #if !MICROPY_ENABLE_DYNRUNTIME
332 : : static const mp_rom_map_elem_t btree_locals_dict_table[] = {
333 : : { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&btree_close_obj) },
334 : : { MP_ROM_QSTR(MP_QSTR_flush), MP_ROM_PTR(&btree_flush_obj) },
335 : : { MP_ROM_QSTR(MP_QSTR_get), MP_ROM_PTR(&btree_get_obj) },
336 : : { MP_ROM_QSTR(MP_QSTR_put), MP_ROM_PTR(&btree_put_obj) },
337 : : { MP_ROM_QSTR(MP_QSTR_seq), MP_ROM_PTR(&btree_seq_obj) },
338 : : { MP_ROM_QSTR(MP_QSTR_keys), MP_ROM_PTR(&btree_keys_obj) },
339 : : { MP_ROM_QSTR(MP_QSTR_values), MP_ROM_PTR(&btree_values_obj) },
340 : : { MP_ROM_QSTR(MP_QSTR_items), MP_ROM_PTR(&btree_items_obj) },
341 : : };
342 : :
343 : : static MP_DEFINE_CONST_DICT(btree_locals_dict, btree_locals_dict_table);
344 : :
345 : : static const mp_getiter_iternext_custom_t btree_getiter_iternext = {
346 : : .getiter = btree_getiter,
347 : : .iternext = btree_iternext,
348 : : };
349 : :
350 : : static MP_DEFINE_CONST_OBJ_TYPE(
351 : : btree_type,
352 : : MP_QSTR_btree,
353 : : MP_TYPE_FLAG_ITER_IS_CUSTOM,
354 : : // Save on qstr's, reuse same as for module
355 : : print, btree_print,
356 : : iter, &btree_getiter_iternext,
357 : : binary_op, btree_binary_op,
358 : : subscr, btree_subscr,
359 : : locals_dict, &btree_locals_dict
360 : : );
361 : : #endif
362 : :
363 : : static const FILEVTABLE btree_stream_fvtable = {
364 : : mp_stream_posix_read,
365 : : mp_stream_posix_write,
366 : : mp_stream_posix_lseek,
367 : : mp_stream_posix_fsync
368 : : };
369 : :
370 : : #if !MICROPY_ENABLE_DYNRUNTIME
371 : 10 : static mp_obj_t mod_btree_open(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
372 : 10 : static const mp_arg_t allowed_args[] = {
373 : : { MP_QSTR_flags, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
374 : : { MP_QSTR_cachesize, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
375 : : { MP_QSTR_pagesize, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
376 : : { MP_QSTR_minkeypage, MP_ARG_KW_ONLY | MP_ARG_INT, {.u_int = 0} },
377 : : };
378 : :
379 : : // Make sure we got a stream object
380 : 10 : mp_get_stream_raise(pos_args[0], MP_STREAM_OP_READ | MP_STREAM_OP_WRITE | MP_STREAM_OP_IOCTL);
381 : :
382 : 10 : struct {
383 : : mp_arg_val_t flags;
384 : : mp_arg_val_t cachesize;
385 : : mp_arg_val_t pagesize;
386 : : mp_arg_val_t minkeypage;
387 : : } args;
388 : 10 : mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args,
389 : : MP_ARRAY_SIZE(allowed_args), allowed_args, (mp_arg_val_t *)&args);
390 : 10 : BTREEINFO openinfo = {0};
391 : 10 : openinfo.flags = args.flags.u_int;
392 : 10 : openinfo.cachesize = args.cachesize.u_int;
393 : 10 : openinfo.psize = args.pagesize.u_int;
394 : 10 : openinfo.minkeypage = args.minkeypage.u_int;
395 : :
396 : 10 : DB *db = __bt_open(MP_OBJ_TO_PTR(pos_args[0]), &btree_stream_fvtable, &openinfo, /*dflags*/ 0);
397 [ + + ]: 10 : if (db == NULL) {
398 : 6 : mp_raise_OSError(errno);
399 : : }
400 : 4 : return MP_OBJ_FROM_PTR(btree_new(db, pos_args[0]));
401 : : }
402 : : static MP_DEFINE_CONST_FUN_OBJ_KW(mod_btree_open_obj, 1, mod_btree_open);
403 : :
404 : : static const mp_rom_map_elem_t mp_module_btree_globals_table[] = {
405 : : { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_btree) },
406 : : { MP_ROM_QSTR(MP_QSTR_open), MP_ROM_PTR(&mod_btree_open_obj) },
407 : : { MP_ROM_QSTR(MP_QSTR_INCL), MP_ROM_INT(FLAG_END_KEY_INCL) },
408 : : { MP_ROM_QSTR(MP_QSTR_DESC), MP_ROM_INT(FLAG_DESC) },
409 : : };
410 : :
411 : : static MP_DEFINE_CONST_DICT(mp_module_btree_globals, mp_module_btree_globals_table);
412 : :
413 : : const mp_obj_module_t mp_module_btree = {
414 : : .base = { &mp_type_module },
415 : : .globals = (mp_obj_dict_t *)&mp_module_btree_globals,
416 : : };
417 : :
418 : : MP_REGISTER_MODULE(MP_QSTR_btree, mp_module_btree);
419 : : #endif
420 : :
421 : : #endif // MICROPY_PY_BTREE
|
