objdict.c 8.72 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 "mpqstr.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) {
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
34
35
36
        first = false;
        mp_obj_print_helper(print, env, next->key);
        print(env, ": ");
        mp_obj_print_helper(print, env, next->value);
37
38
39
40
    }
    print(env, "}");
}

41
42
43
44
45
46
47
// args are reverse in the array
static mp_obj_t dict_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) {
    // 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) {
48
49
50
51
52
    mp_obj_dict_t *o = lhs_in;
    switch (op) {
        case RT_BINARY_OP_SUBSCR:
        {
            // dict load
John R. Lenton's avatar
John R. Lenton committed
53
            mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false, false);
54
            if (elem == NULL) {
55
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
56
57
58
59
60
61
62
63
64
65
            } else {
                return elem->value;
            }
        }
        default:
            // op not supported
            return NULL;
    }
}

John R. Lenton's avatar
John R. Lenton committed
66
67
68
69
70
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

/******************************************************************************/
/* 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,
    .methods = { { NULL, NULL }, },
};

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
123
124
125
126
127
128
129
130
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;
}
131
132
static MP_DEFINE_CONST_FUN_OBJ_1(dict_clear_obj, dict_clear);

John R. Lenton's avatar
John R. Lenton committed
133
134
135
136
137
138
139
140
141
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
142

John R. Lenton's avatar
John R. Lenton committed
143
144
145
146
static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, bool pop, bool set) {
    mp_map_elem_t *elem = mp_map_lookup_helper(self, key, set, pop);
    mp_obj_t value;
    if (elem == NULL || elem->value == NULL) {
John R. Lenton's avatar
John R. Lenton committed
147
148
149
150
        if (deflt == NULL) {
            if (pop) {
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
            } else {
John R. Lenton's avatar
John R. Lenton committed
151
                value = mp_const_none;
John R. Lenton's avatar
John R. Lenton committed
152
153
            }
        } else {
John R. Lenton's avatar
John R. Lenton committed
154
            value = deflt;
John R. Lenton's avatar
John R. Lenton committed
155
156
        }
    } else {
John R. Lenton's avatar
John R. Lenton committed
157
        value = elem->value;
John R. Lenton's avatar
John R. Lenton committed
158
159
        if (pop) {
            /* catch the leak (from mp_map_lookup_helper) */
160
            m_free(elem, sizeof(mp_map_elem_t));
John R. Lenton's avatar
John R. Lenton committed
161
162
        }
    }
John R. Lenton's avatar
John R. Lenton committed
163
164
165
166
    if (set) {
        elem->value = value;
    }
    return value;
John R. Lenton's avatar
John R. Lenton committed
167
168
}

John R. Lenton's avatar
John R. Lenton committed
169
170
171
172
static mp_obj_t dict_get(int n_args, const mp_obj_t *args) {
    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
173
174
175
    return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
                           args[1],
                           n_args == 3 ? args[2] : NULL,
John R. Lenton's avatar
John R. Lenton committed
176
                           false, false);
John R. Lenton's avatar
John R. Lenton committed
177
178
179
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);

John R. Lenton's avatar
John R. Lenton committed
180
181
182
183
184
185
186
static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) {
    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,
John R. Lenton's avatar
John R. Lenton committed
187
                           true, false);
John R. Lenton's avatar
John R. Lenton committed
188
189
190
191
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);


John R. Lenton's avatar
John R. Lenton committed
192
193
194
195
196
197
198
199
200
201
202
static mp_obj_t dict_setdefault(int n_args, const mp_obj_t *args) {
    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,
                           false, true);
}
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault);

John R. Lenton's avatar
John R. Lenton committed
203
204
205
206
207
208
209
210
211
212
213
214
215

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
216
    next->value = NULL;
John R. Lenton's avatar
John R. Lenton committed
217
218
219
220
221
222
    mp_obj_t tuple = mp_obj_new_tuple(2, items);

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

223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
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 {
            mp_map_lookup_helper(&self->map, key, true, false)->value = value;
        }
    }

    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
249

John R. Lenton's avatar
John R. Lenton committed
250
251
252
/******************************************************************************/
/* dict constructors & etc                                                    */

253
254
255
const mp_obj_type_t dict_type = {
    { &mp_const_type },
    "dict",
256
257
258
    .print = dict_print,
    .make_new = dict_make_new,
    .binary_op = dict_binary_op,
John R. Lenton's avatar
John R. Lenton committed
259
    .getiter = dict_getiter,
John R. Lenton's avatar
John R. Lenton committed
260
261
    .methods = {
        { "clear", &dict_clear_obj },
John R. Lenton's avatar
John R. Lenton committed
262
        { "copy", &dict_copy_obj },
John R. Lenton's avatar
John R. Lenton committed
263
        { "get", &dict_get_obj },
John R. Lenton's avatar
John R. Lenton committed
264
        { "pop", &dict_pop_obj },
John R. Lenton's avatar
John R. Lenton committed
265
        { "popitem", &dict_popitem_obj },
John R. Lenton's avatar
John R. Lenton committed
266
        { "setdefault", &dict_setdefault_obj },
267
        { "update", &dict_update_obj },
John R. Lenton's avatar
John R. Lenton committed
268
269
        { NULL, NULL }, // end-of-list sentinel
    },
270
271
272
273
274
275
276
277
};

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;
    mp_map_init(&o->map, MP_MAP_OBJ, n_args);
    return o;
}
278
279

uint mp_obj_dict_len(mp_obj_t self_in) {
John R. Lenton's avatar
John R. Lenton committed
280
    return ((mp_obj_dict_t *)self_in)->map.used;
281
282
283
284
285
}

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;
John R. Lenton's avatar
John R. Lenton committed
286
    mp_map_lookup_helper(&self->map, key, true, false)->value = value;
287
288
    return self_in;
}