objdict.c 14.8 KB
Newer Older
xbe's avatar
xbe committed
1
#include <stdbool.h>
2
3
4
5
6
7
#include <string.h>
#include <assert.h>

#include "nlr.h"
#include "misc.h"
#include "mpconfig.h"
8
#include "qstr.h"
9
10
11
12
13
14
15
16
17
18
#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;

19
20
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
21

22
STATIC void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
23
24
25
    mp_obj_dict_t *self = self_in;
    bool first = true;
    print(env, "{");
John R. Lenton's avatar
John R. Lenton committed
26
27
28
29
30
    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, ", ");
31
        }
John R. Lenton's avatar
John R. Lenton committed
32
        first = false;
33
        mp_obj_print_helper(print, env, next->key, PRINT_REPR);
John R. Lenton's avatar
John R. Lenton committed
34
        print(env, ": ");
35
        mp_obj_print_helper(print, env, next->value, PRINT_REPR);
36
37
38
39
    }
    print(env, "}");
}

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

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

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

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

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

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

113
STATIC const mp_obj_type_t dict_it_type = {
114
    { &mp_type_type },
115
    .name = MP_QSTR_iterator,
John R. Lenton's avatar
John R. Lenton committed
116
117
118
    .iternext = dict_it_iternext,
};

119
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
120
121
122
123
124
125
126
    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;
}

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

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

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

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

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

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

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

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


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

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

244
STATIC mp_obj_t dict_popitem(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
245
246
247
    assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
    mp_obj_dict_t *self = self_in;
    if (self->map.used == 0) {
248
        nlr_jump(mp_obj_new_exception_msg(&mp_type_KeyError, "popitem(): dictionary is empty"));
John R. Lenton's avatar
John R. Lenton committed
249
250
251
252
253
254
255
    }
    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
256
    next->value = NULL;
John R. Lenton's avatar
John R. Lenton committed
257
258
259
260
    mp_obj_t tuple = mp_obj_new_tuple(2, items);

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

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

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

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

290
291
292
/******************************************************************************/
/* dict views                                                                 */

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

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

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

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;

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

342
STATIC const mp_obj_type_t dict_view_it_type = {
343
    { &mp_type_type },
344
    .name = MP_QSTR_iterator,
345
346
347
348
    .iternext = dict_view_it_iternext,
    .methods = NULL,            /* set operations still to come */
};

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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