objdict.c 14.8 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;

20
21
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);
John R. Lenton's avatar
John R. Lenton committed
22

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
    // TODO create from an iterable!
    return rt_build_map(0);
}

46
STATIC mp_obj_t dict_unary_op(int op, mp_obj_t self_in) {
47
48
    mp_obj_dict_t *self = self_in;
    switch (op) {
49
        case RT_UNARY_OP_BOOL: return MP_BOOL(self->map.used != 0);
50
        case RT_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT((machine_int_t)self->map.used);
51
52
53
54
        default: return MP_OBJ_NULL; // op not supported for None
    }
}

55
STATIC mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
56
57
58
59
60
    mp_obj_dict_t *o = lhs_in;
    switch (op) {
        case RT_BINARY_OP_SUBSCR:
        {
            // dict load
61
            mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
62
            if (elem == NULL) {
63
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
64
65
66
67
            } else {
                return elem->value;
            }
        }
68
        case RT_BINARY_OP_IN:
69
70
        {
            mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
71
            return MP_BOOL(elem != NULL);
72
        }
73
74
75
76
77
78
        default:
            // op not supported
            return NULL;
    }
}

John R. Lenton's avatar
John R. Lenton committed
79
80
81
82
83
84
85
86
87
88

/******************************************************************************/
/* 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;

89
STATIC mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    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;
    }
}

114
STATIC const mp_obj_type_t dict_it_type = {
John R. Lenton's avatar
John R. Lenton committed
115
116
117
118
119
    { &mp_const_type },
    "dict_iterator",
    .iternext = dict_it_iternext,
};

120
STATIC mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur) {
John R. Lenton's avatar
John R. Lenton committed
121
122
123
124
125
126
127
    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;
}

128
STATIC mp_obj_t dict_getiter(mp_obj_t o_in) {
John R. Lenton's avatar
John R. Lenton committed
129
130
131
132
133
134
    return mp_obj_new_dict_iterator(o_in, 0);
}

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

135
STATIC mp_obj_t dict_clear(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
136
137
138
139
140
141
142
    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;
}
143
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_clear_obj, dict_clear);
144

145
STATIC mp_obj_t dict_copy(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
146
147
148
149
150
151
152
    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;
}
153
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy);
John R. Lenton's avatar
John R. Lenton committed
154

155
// this is a classmethod
156
STATIC mp_obj_t dict_fromkeys(uint n_args, const mp_obj_t *args) {
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    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;
}
181
182
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);
183

184
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) {
185
    mp_map_elem_t *elem = mp_map_lookup(self, key, lookup_kind);
John R. Lenton's avatar
John R. Lenton committed
186
187
    mp_obj_t value;
    if (elem == NULL || elem->value == NULL) {
John R. Lenton's avatar
John R. Lenton committed
188
        if (deflt == NULL) {
189
            if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
John R. Lenton's avatar
John R. Lenton committed
190
191
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
            } else {
John R. Lenton's avatar
John R. Lenton committed
192
                value = mp_const_none;
John R. Lenton's avatar
John R. Lenton committed
193
194
            }
        } else {
John R. Lenton's avatar
John R. Lenton committed
195
            value = deflt;
John R. Lenton's avatar
John R. Lenton committed
196
197
        }
    } else {
John R. Lenton's avatar
John R. Lenton committed
198
        value = elem->value;
199
200
        if (lookup_kind == MP_MAP_LOOKUP_REMOVE_IF_FOUND) {
            // catch the leak (from mp_map_lookup)
201
            m_free(elem, sizeof(mp_map_elem_t));
John R. Lenton's avatar
John R. Lenton committed
202
203
        }
    }
204
    if (lookup_kind == MP_MAP_LOOKUP_ADD_IF_NOT_FOUND) {
John R. Lenton's avatar
John R. Lenton committed
205
206
207
        elem->value = value;
    }
    return value;
John R. Lenton's avatar
John R. Lenton committed
208
209
}

210
STATIC mp_obj_t dict_get(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
211
212
213
    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
214
215
216
    return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
                           args[1],
                           n_args == 3 ? args[2] : NULL,
217
                           MP_MAP_LOOKUP);
John R. Lenton's avatar
John R. Lenton committed
218
}
219
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);
John R. Lenton's avatar
John R. Lenton committed
220

221
STATIC mp_obj_t dict_pop(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
222
223
224
225
226
227
    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,
228
                           MP_MAP_LOOKUP_REMOVE_IF_FOUND);
John R. Lenton's avatar
John R. Lenton committed
229
}
230
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);
John R. Lenton's avatar
John R. Lenton committed
231
232


233
STATIC mp_obj_t dict_setdefault(uint n_args, const mp_obj_t *args) {
John R. Lenton's avatar
John R. Lenton committed
234
235
236
237
238
239
    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,
240
                           MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
John R. Lenton's avatar
John R. Lenton committed
241
}
242
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault);
John R. Lenton's avatar
John R. Lenton committed
243

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

245
STATIC mp_obj_t dict_popitem(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
246
247
248
249
250
251
252
253
254
255
256
    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
257
    next->value = NULL;
John R. Lenton's avatar
John R. Lenton committed
258
259
260
261
    mp_obj_t tuple = mp_obj_new_tuple(2, items);

    return tuple;
}
262
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_popitem_obj, dict_popitem);
John R. Lenton's avatar
John R. Lenton committed
263

264
STATIC mp_obj_t dict_update(mp_obj_t self_in, mp_obj_t iterable) {
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    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 {
282
            mp_map_lookup(&self->map, key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
283
284
285
286
287
        }
    }

    return mp_const_none;
}
288
STATIC MP_DEFINE_CONST_FUN_OBJ_2(dict_update_obj, dict_update);
289

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

291
292
293
/******************************************************************************/
/* dict views                                                                 */

294
295
STATIC const mp_obj_type_t dict_view_type;
STATIC const mp_obj_type_t dict_view_it_type;
296
297
298
299
300
301
302

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

303
STATIC char *mp_dict_view_names[] = {"dict_items", "dict_keys", "dict_values"};
304
305
306
307
308
309
310
311
312
313
314
315
316
317

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;

318
STATIC mp_obj_t dict_view_it_iternext(mp_obj_t self_in) {
319
320
321
322
323
324
    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) {
325
326
327
328
329
330
331
332
333
334
335
336
            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;
337
338
339
340
341
342
        }
    } else {
        return mp_const_stop_iteration;
    }
}

343
STATIC const mp_obj_type_t dict_view_it_type = {
344
345
346
347
348
349
    { &mp_const_type },
    "dict_view_iterator",
    .iternext = dict_view_it_iternext,
    .methods = NULL,            /* set operations still to come */
};

350
STATIC mp_obj_t dict_view_getiter(mp_obj_t view_in) {
351
352
353
354
355
356
357
358
359
    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;
}

360
STATIC void dict_view_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
361
362
363
364
365
366
367
368
369
370
371
372
    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;
373
        mp_obj_print_helper(print, env, next, PRINT_REPR);
374
375
376
377
    }
    print(env, "])");
}

378
STATIC mp_obj_t dict_view_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
379
380
381
    /* 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;
382
    if (op != RT_BINARY_OP_IN) return NULL;
383
384
385
386
    return dict_binary_op(op, o->dict, rhs_in);
}


387
STATIC const mp_obj_type_t dict_view_type = {
388
389
390
    { &mp_const_type },
    "dict_view",
    .print = dict_view_print,
391
    .binary_op = dict_view_binary_op,
392
393
394
395
396
397
398
399
400
401
402
    .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;
}

403
STATIC mp_obj_t dict_view(mp_obj_t self_in, mp_dict_view_kind_t kind) {
404
405
406
407
408
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    return mp_obj_new_dict_view(self, kind);
}

409
STATIC mp_obj_t dict_items(mp_obj_t self_in) {
410
411
    return dict_view(self_in, MP_DICT_VIEW_ITEMS);
}
412
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_items_obj, dict_items);
413

414
STATIC mp_obj_t dict_keys(mp_obj_t self_in) {
415
416
    return dict_view(self_in, MP_DICT_VIEW_KEYS);
}
417
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_keys_obj, dict_keys);
418

419
STATIC mp_obj_t dict_values(mp_obj_t self_in) {
420
421
    return dict_view(self_in, MP_DICT_VIEW_VALUES);
}
422
STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_values_obj, dict_values);
423

John R. Lenton's avatar
John R. Lenton committed
424
/******************************************************************************/
425
/* dict constructors & public C API                                           */
John R. Lenton's avatar
John R. Lenton committed
426

427
STATIC const mp_method_t dict_type_methods[] = {
428
429
    { "clear", &dict_clear_obj },
    { "copy", &dict_copy_obj },
430
    { "fromkeys", &dict_fromkeys_obj },
431
    { "get", &dict_get_obj },
432
433
    { "items", &dict_items_obj },
    { "keys", &dict_keys_obj },
434
435
436
437
    { "pop", &dict_pop_obj },
    { "popitem", &dict_popitem_obj },
    { "setdefault", &dict_setdefault_obj },
    { "update", &dict_update_obj },
438
    { "values", &dict_values_obj },
439
440
441
    { NULL, NULL }, // end-of-list sentinel
};

442
const mp_obj_type_t dict_type = {
443
    { &mp_const_type },
444
    "dict",
445
446
    .print = dict_print,
    .make_new = dict_make_new,
447
    .unary_op = dict_unary_op,
448
    .binary_op = dict_binary_op,
John R. Lenton's avatar
John R. Lenton committed
449
    .getiter = dict_getiter,
450
    .methods = dict_type_methods,
451
452
453
454
455
};

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;
456
    mp_map_init(&o->map, n_args);
457
458
    return o;
}
459
460

uint mp_obj_dict_len(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
461
    return ((mp_obj_dict_t *)self_in)->map.used;
462
463
464
465
466
}

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;
467
    mp_map_lookup(&self->map, key, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND)->value = value;
468
469
    return self_in;
}
470
471
472
473
474
475

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;
}