objdict.c 14.5 KB
Newer Older
1
2
3
4
5
6
7
8
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <assert.h>

#include "nlr.h"
#include "misc.h"
#include "mpconfig.h"
9
#include "qstr.h"
10
11
12
13
14
15
16
17
18
19
#include "obj.h"
#include "runtime0.h"
#include "runtime.h"
#include "map.h"

typedef struct _mp_obj_dict_t {
    mp_obj_base_t base;
    mp_map_t map;
} mp_obj_dict_t;

John R. Lenton's avatar
John R. Lenton committed
20
21
22
static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur);
static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in);

23
static void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
24
25
26
    mp_obj_dict_t *self = self_in;
    bool first = true;
    print(env, "{");
John R. Lenton's avatar
John R. Lenton committed
27
28
29
30
31
    mp_obj_t *dict_iter = mp_obj_new_dict_iterator(self, 0);
    mp_map_elem_t *next = NULL;
    while ((next = dict_it_iternext_elem(dict_iter)) != NULL) {
        if (!first) {
            print(env, ", ");
32
        }
John R. Lenton's avatar
John R. Lenton committed
33
        first = false;
34
        mp_obj_print_helper(print, env, next->key, PRINT_REPR);
John R. Lenton's avatar
John R. Lenton committed
35
        print(env, ": ");
36
        mp_obj_print_helper(print, env, next->value, PRINT_REPR);
37
38
39
40
    }
    print(env, "}");
}

41
static mp_obj_t dict_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
42
43
44
45
46
    // TODO create from an iterable!
    return rt_build_map(0);
}

static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
47
48
49
50
51
    mp_obj_dict_t *o = lhs_in;
    switch (op) {
        case RT_BINARY_OP_SUBSCR:
        {
            // dict load
52
            mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
53
            if (elem == NULL) {
54
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
55
56
57
58
            } else {
                return elem->value;
            }
        }
59
60
61
62
63
64
        case RT_COMPARE_OP_IN:
        case RT_COMPARE_OP_NOT_IN:
        {
            mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
            return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
        }
65
66
67
68
69
70
        default:
            // op not supported
            return NULL;
    }
}

John R. Lenton's avatar
John R. Lenton committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

/******************************************************************************/
/* dict iterator                                                              */

typedef struct _mp_obj_dict_it_t {
    mp_obj_base_t base;
    mp_obj_dict_t *dict;
    machine_uint_t cur;
} mp_obj_dict_it_t;

static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) {
    mp_obj_dict_it_t *self = self_in;
    machine_uint_t max = self->dict->map.alloc;
    mp_map_elem_t *table = self->dict->map.table;

    for (int i = self->cur; i < max; i++) {
        if (table[i].key != NULL) {
            self->cur = i + 1;
            return &(table[i]);
        }
    }

    return NULL;
}

mp_obj_t dict_it_iternext(mp_obj_t self_in) {
    mp_map_elem_t *next = dict_it_iternext_elem(self_in);

    if (next != NULL) {
        return next->key;
    } else {
        return mp_const_stop_iteration;
    }
}

static const mp_obj_type_t dict_it_type = {
    { &mp_const_type },
    "dict_iterator",
    .iternext = dict_it_iternext,
};

static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur) {
    mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t);
    o->base.type = &dict_it_type;
    o->dict = dict;
    o->cur = cur;
    return o;
}

static mp_obj_t dict_getiter(mp_obj_t o_in) {
    return mp_obj_new_dict_iterator(o_in, 0);
}

/******************************************************************************/
/* dict methods                                                               */

John R. Lenton's avatar
John R. Lenton committed
127
128
129
130
131
132
133
134
static mp_obj_t dict_clear(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;

    mp_map_clear(&self->map);

    return mp_const_none;
}
135
136
static MP_DEFINE_CONST_FUN_OBJ_1(dict_clear_obj, dict_clear);

John R. Lenton's avatar
John R. Lenton committed
137
138
139
140
141
142
143
144
145
static mp_obj_t dict_copy(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    mp_obj_dict_t *other = mp_obj_new_dict(self->map.alloc);
    other->map.used = self->map.used;
    memcpy(other->map.table, self->map.table, self->map.alloc * sizeof(mp_map_elem_t));
    return other;
}
static MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy);
John R. Lenton's avatar
John R. Lenton committed
146

147
// this is a classmethod
148
static mp_obj_t dict_fromkeys(uint n_args, const mp_obj_t *args) {
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    assert(2 <= n_args && n_args <= 3);
    mp_obj_t iter = rt_getiter(args[1]);
    mp_obj_t len = mp_obj_len_maybe(iter);
    mp_obj_t value = mp_const_none;
    mp_obj_t next = NULL;
    mp_obj_dict_t *self = NULL;

    if (n_args > 2) {
        value = args[2];
    }

    if (len == MP_OBJ_NULL) {
        /* object's type doesn't have a __len__ slot */
        self = mp_obj_new_dict(0);
    } else {
        self = mp_obj_new_dict(MP_OBJ_SMALL_INT_VALUE(len));
    }

    while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
        mp_map_lookup(&self->map, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
    }

    return self;
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_fromkeys_fun_obj, 2, 3, dict_fromkeys);
static MP_DEFINE_CONST_CLASSMETHOD_OBJ(dict_fromkeys_obj, (const mp_obj_t)&dict_fromkeys_fun_obj);

176
177
static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, mp_map_lookup_kind_t lookup_kind) {
    mp_map_elem_t *elem = mp_map_lookup(self, key, lookup_kind);
John R. Lenton's avatar
John R. Lenton committed
178
179
    mp_obj_t value;
    if (elem == NULL || elem->value == NULL) {
John R. Lenton's avatar
John R. Lenton committed
180
        if (deflt == NULL) {
181
            if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
John R. Lenton's avatar
John R. Lenton committed
182
183
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
            } else {
John R. Lenton's avatar
John R. Lenton committed
184
                value = mp_const_none;
John R. Lenton's avatar
John R. Lenton committed
185
186
            }
        } else {
John R. Lenton's avatar
John R. Lenton committed
187
            value = deflt;
John R. Lenton's avatar
John R. Lenton committed
188
189
        }
    } else {
John R. Lenton's avatar
John R. Lenton committed
190
        value = elem->value;
191
192
        if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
            // catch the leak (from mp_map_lookup)
193
            m_free(elem, sizeof(mp_map_elem_t));
John R. Lenton's avatar
John R. Lenton committed
194
195
        }
    }
196
    if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
John R. Lenton's avatar
John R. Lenton committed
197
198
199
        elem->value = value;
    }
    return value;
John R. Lenton's avatar
John R. Lenton committed
200
201
}

202
static mp_obj_t dict_get(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
203
204
205
    assert(2 <= n_args && n_args <= 3);
    assert(MP_OBJ_IS_TYPE(args[0], &dict_type));

John R. Lenton's avatar
John R. Lenton committed
206
207
208
    return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
                           args[1],
                           n_args == 3 ? args[2] : NULL,
209
                           MP_MAP_LOOKUP);
John R. Lenton's avatar
John R. Lenton committed
210
211
212
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);

213
static mp_obj_t dict_pop(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
214
215
216
217
218
219
    assert(2 <= n_args && n_args <= 3);
    assert(MP_OBJ_IS_TYPE(args[0], &dict_type));

    return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
                           args[1],
                           n_args == 3 ? args[2] : NULL,
220
                           MP_MAP_LOOKUP_REMOVE_IF_FOUND);
John R. Lenton's avatar
John R. Lenton committed
221
222
223
224
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);


225
static mp_obj_t dict_setdefault(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
226
227
228
229
230
231
    assert(2 <= n_args && n_args <= 3);
    assert(MP_OBJ_IS_TYPE(args[0], &dict_type));

    return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
                           args[1],
                           n_args == 3 ? args[2] : NULL,
232
                           MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
John R. Lenton's avatar
John R. Lenton committed
233
234
235
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault);

John R. Lenton's avatar
John R. Lenton committed
236
237
238
239
240
241
242
243
244
245
246
247
248

static mp_obj_t dict_popitem(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    if (self->map.used == 0) {
        nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "popitem(): dictionary is empty"));
    }
    mp_obj_dict_it_t *iter = mp_obj_new_dict_iterator(self, 0);

    mp_map_elem_t *next = dict_it_iternext_elem(iter);
    self->map.used--;
    mp_obj_t items[] = {next->key, next->value};
    next->key = NULL;
John R. Lenton's avatar
John R. Lenton committed
249
    next->value = NULL;
John R. Lenton's avatar
John R. Lenton committed
250
251
252
253
254
255
    mp_obj_t tuple = mp_obj_new_tuple(2, items);

    return tuple;
}
static MP_DEFINE_CONST_FUN_OBJ_1(dict_popitem_obj, dict_popitem);

256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
static mp_obj_t dict_update(mp_obj_t self_in, mp_obj_t iterable) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    /* TODO: check for the "keys" method */
    mp_obj_t iter = rt_getiter(iterable);
    mp_obj_t next = NULL;
    while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
        mp_obj_t inneriter = rt_getiter(next);
        mp_obj_t key = rt_iternext(inneriter);
        mp_obj_t value = rt_iternext(inneriter);
        mp_obj_t stop = rt_iternext(inneriter);
        if (key == mp_const_stop_iteration
            || value == mp_const_stop_iteration
            || stop != mp_const_stop_iteration) {
            nlr_jump(mp_obj_new_exception_msg(
                         MP_QSTR_ValueError,
                         "dictionary update sequence has the wrong length"));
        } else {
274
            mp_map_lookup(&self->map, key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
275
276
277
278
279
280
281
        }
    }

    return mp_const_none;
}
static MP_DEFINE_CONST_FUN_OBJ_2(dict_update_obj, dict_update);

John R. Lenton's avatar
John R. Lenton committed
282

283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
/******************************************************************************/
/* dict views                                                                 */

static const mp_obj_type_t dict_view_type;
static const mp_obj_type_t dict_view_it_type;

typedef enum _mp_dict_view_kind_t {
    MP_DICT_VIEW_ITEMS,
    MP_DICT_VIEW_KEYS,
    MP_DICT_VIEW_VALUES,
} mp_dict_view_kind_t;

static char *mp_dict_view_names[] = {"dict_items", "dict_keys", "dict_values"};

typedef struct _mp_obj_dict_view_it_t {
    mp_obj_base_t base;
    mp_dict_view_kind_t kind;
    mp_obj_dict_it_t *iter;
    machine_uint_t cur;
} mp_obj_dict_view_it_t;

typedef struct _mp_obj_dict_view_t {
    mp_obj_base_t base;
    mp_obj_dict_t *dict;
    mp_dict_view_kind_t kind;
} mp_obj_dict_view_t;

static mp_obj_t dict_view_it_iternext(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_view_it_type));
    mp_obj_dict_view_it_t *self = self_in;
    mp_map_elem_t *next = dict_it_iternext_elem(self->iter);

    if (next != NULL) {
        switch (self->kind) {
317
318
319
320
321
322
323
324
325
326
327
328
            case MP_DICT_VIEW_ITEMS:
            {
                mp_obj_t items[] = {next->key, next->value};
                return mp_obj_new_tuple(2, items);
            }
            case MP_DICT_VIEW_KEYS:
                return next->key;
            case MP_DICT_VIEW_VALUES:
                return next->value;
            default:
                assert(0);          /* can't happen */
                return mp_const_none;
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
        }
    } else {
        return mp_const_stop_iteration;
    }
}

static const mp_obj_type_t dict_view_it_type = {
    { &mp_const_type },
    "dict_view_iterator",
    .iternext = dict_view_it_iternext,
    .methods = NULL,            /* set operations still to come */
};

static mp_obj_t dict_view_getiter(mp_obj_t view_in) {
    assert(MP_OBJ_IS_TYPE(view_in, &dict_view_type));
    mp_obj_dict_view_t *view = view_in;
    mp_obj_dict_view_it_t *o = m_new_obj(mp_obj_dict_view_it_t);
    o->base.type = &dict_view_it_type;
    o->kind = view->kind;
    o->iter = mp_obj_new_dict_iterator(view->dict, 0);
    return o;
}

352
static void dict_view_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
353
354
355
356
357
358
359
360
361
362
363
364
    assert(MP_OBJ_IS_TYPE(self_in, &dict_view_type));
    mp_obj_dict_view_t *self = self_in;
    bool first = true;
    print(env, mp_dict_view_names[self->kind]);
    print(env, "([");
    mp_obj_t *self_iter = dict_view_getiter(self);
    mp_obj_t *next = NULL;
    while ((next = dict_view_it_iternext(self_iter)) != mp_const_stop_iteration) {
        if (!first) {
            print(env, ", ");
        }
        first = false;
365
        mp_obj_print_helper(print, env, next, PRINT_REPR);
366
367
368
369
    }
    print(env, "])");
}

370
371
372
373
374
375
376
377
378
static mp_obj_t dict_view_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
    /* only supported for the 'keys' kind until sets and dicts are refactored */
    mp_obj_dict_view_t *o = lhs_in;
    if (o->kind != MP_DICT_VIEW_KEYS) return NULL;
    if (op != RT_COMPARE_OP_IN && op != RT_COMPARE_OP_NOT_IN) return NULL;
    return dict_binary_op(op, o->dict, rhs_in);
}


379
380
381
382
static const mp_obj_type_t dict_view_type = {
    { &mp_const_type },
    "dict_view",
    .print = dict_view_print,
383
    .binary_op = dict_view_binary_op,
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
    .getiter = dict_view_getiter,
};

mp_obj_t mp_obj_new_dict_view(mp_obj_dict_t *dict, mp_dict_view_kind_t kind) {
    mp_obj_dict_view_t *o = m_new_obj(mp_obj_dict_view_t);
    o->base.type = &dict_view_type;
    o->dict = dict;
    o->kind = kind;
    return o;
}

static mp_obj_t dict_view(mp_obj_t self_in, mp_dict_view_kind_t kind) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    return mp_obj_new_dict_view(self, kind);
}

static mp_obj_t dict_items(mp_obj_t self_in) {
    return dict_view(self_in, MP_DICT_VIEW_ITEMS);
}
static MP_DEFINE_CONST_FUN_OBJ_1(dict_items_obj, dict_items);

static mp_obj_t dict_keys(mp_obj_t self_in) {
    return dict_view(self_in, MP_DICT_VIEW_KEYS);
}
static MP_DEFINE_CONST_FUN_OBJ_1(dict_keys_obj, dict_keys);

static mp_obj_t dict_values(mp_obj_t self_in) {
    return dict_view(self_in, MP_DICT_VIEW_VALUES);
}
static MP_DEFINE_CONST_FUN_OBJ_1(dict_values_obj, dict_values);

John R. Lenton's avatar
John R. Lenton committed
416
/******************************************************************************/
417
/* dict constructors & public C API                                           */
John R. Lenton's avatar
John R. Lenton committed
418

419
420
421
static const mp_method_t dict_type_methods[] = {
    { "clear", &dict_clear_obj },
    { "copy", &dict_copy_obj },
422
    { "fromkeys", &dict_fromkeys_obj },
423
    { "get", &dict_get_obj },
424
425
    { "items", &dict_items_obj },
    { "keys", &dict_keys_obj },
426
427
428
429
    { "pop", &dict_pop_obj },
    { "popitem", &dict_popitem_obj },
    { "setdefault", &dict_setdefault_obj },
    { "update", &dict_update_obj },
430
    { "values", &dict_values_obj },
431
432
433
    { NULL, NULL }, // end-of-list sentinel
};

434
const mp_obj_type_t dict_type = {
435
    { &mp_const_type },
436
    "dict",
437
438
439
    .print = dict_print,
    .make_new = dict_make_new,
    .binary_op = dict_binary_op,
John R. Lenton's avatar
John R. Lenton committed
440
    .getiter = dict_getiter,
441
    .methods = dict_type_methods,
442
443
444
445
446
};

mp_obj_t mp_obj_new_dict(int n_args) {
    mp_obj_dict_t *o = m_new_obj(mp_obj_dict_t);
    o->base.type = &dict_type;
447
    mp_map_init(&o->map, n_args);
448
449
    return o;
}
450
451

uint mp_obj_dict_len(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
452
    return ((mp_obj_dict_t *)self_in)->map.used;
453
454
455
456
457
}

mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
458
    mp_map_lookup(&self->map, key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
459
460
    return self_in;
}
461
462
463
464
465
466

mp_map_t *mp_obj_dict_get_map(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    return &self->map;
}