1 /*
2  * This file is part of the MicroPython project, http://micropython.org/
3  *
4  * The MIT License (MIT)
5  *
6  * Copyright (c) 2013-2017 Damien P. George
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 <stdbool.h>
28 #include <string.h>
29 #include <assert.h>
30 
31 #include "py/runtime.h"
32 #include "py/builtin.h"
33 
34 #if MICROPY_PY_BUILTINS_SET
35 
36 typedef struct _mp_obj_set_t {
37     mp_obj_base_t base;
38     mp_set_t set;
39 } mp_obj_set_t;
40 
41 typedef struct _mp_obj_set_it_t {
42     mp_obj_base_t base;
43     mp_fun_1_t iternext;
44     mp_obj_set_t *set;
45     size_t cur;
46 } mp_obj_set_it_t;
47 
is_set_or_frozenset(mp_obj_t o)48 STATIC bool is_set_or_frozenset(mp_obj_t o) {
49     return mp_obj_is_type(o, &mp_type_set)
50            #if MICROPY_PY_BUILTINS_FROZENSET
51            || mp_obj_is_type(o, &mp_type_frozenset)
52            #endif
53     ;
54 }
55 
56 // This macro is shorthand for mp_check_self to verify the argument is a set.
57 #define check_set(o) mp_check_self(mp_obj_is_type(o, &mp_type_set))
58 
59 // This macro is shorthand for mp_check_self to verify the argument is a
60 // set or frozenset for methods that operate on both of these types.
61 #define check_set_or_frozenset(o) mp_check_self(is_set_or_frozenset(o))
62 
set_print(const mp_print_t * print,mp_obj_t self_in,mp_print_kind_t kind)63 STATIC void set_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
64     (void)kind;
65     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
66     #if MICROPY_PY_BUILTINS_FROZENSET
67     bool is_frozen = mp_obj_is_type(self_in, &mp_type_frozenset);
68     #endif
69     if (self->set.used == 0) {
70         #if MICROPY_PY_BUILTINS_FROZENSET
71         if (is_frozen) {
72             mp_print_str(print, "frozen");
73         }
74         #endif
75         mp_print_str(print, "set()");
76         return;
77     }
78     bool first = true;
79     #if MICROPY_PY_BUILTINS_FROZENSET
80     if (is_frozen) {
81         mp_print_str(print, "frozenset(");
82     }
83     #endif
84     mp_print_str(print, "{");
85     for (size_t i = 0; i < self->set.alloc; i++) {
86         if (mp_set_slot_is_filled(&self->set, i)) {
87             if (!first) {
88                 mp_print_str(print, ", ");
89             }
90             first = false;
91             mp_obj_print_helper(print, self->set.table[i], PRINT_REPR);
92         }
93     }
94     mp_print_str(print, "}");
95     #if MICROPY_PY_BUILTINS_FROZENSET
96     if (is_frozen) {
97         mp_print_str(print, ")");
98     }
99     #endif
100 }
101 
set_make_new(const mp_obj_type_t * type,size_t n_args,size_t n_kw,const mp_obj_t * args)102 STATIC mp_obj_t set_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
103     mp_arg_check_num(n_args, n_kw, 0, 1, false);
104 
105     switch (n_args) {
106         case 0: {
107             // create a new, empty set
108             mp_obj_set_t *set = MP_OBJ_TO_PTR(mp_obj_new_set(0, NULL));
109             // set actual set/frozenset type
110             set->base.type = type;
111             return MP_OBJ_FROM_PTR(set);
112         }
113 
114         case 1:
115         default: { // can only be 0 or 1 arg
116             // 1 argument, an iterable from which we make a new set
117             mp_obj_t set = mp_obj_new_set(0, NULL);
118             mp_obj_t iterable = mp_getiter(args[0], NULL);
119             mp_obj_t item;
120             while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
121                 mp_obj_set_store(set, item);
122             }
123             // Set actual set/frozenset type
124             ((mp_obj_set_t *)MP_OBJ_TO_PTR(set))->base.type = type;
125             return set;
126         }
127     }
128 }
129 
set_it_iternext(mp_obj_t self_in)130 STATIC mp_obj_t set_it_iternext(mp_obj_t self_in) {
131     mp_obj_set_it_t *self = MP_OBJ_TO_PTR(self_in);
132     size_t max = self->set->set.alloc;
133     mp_set_t *set = &self->set->set;
134 
135     for (size_t i = self->cur; i < max; i++) {
136         if (mp_set_slot_is_filled(set, i)) {
137             self->cur = i + 1;
138             return set->table[i];
139         }
140     }
141 
142     return MP_OBJ_STOP_ITERATION;
143 }
144 
set_getiter(mp_obj_t set_in,mp_obj_iter_buf_t * iter_buf)145 STATIC mp_obj_t set_getiter(mp_obj_t set_in, mp_obj_iter_buf_t *iter_buf) {
146     assert(sizeof(mp_obj_set_it_t) <= sizeof(mp_obj_iter_buf_t));
147     mp_obj_set_it_t *o = (mp_obj_set_it_t *)iter_buf;
148     o->base.type = &mp_type_polymorph_iter;
149     o->iternext = set_it_iternext;
150     o->set = (mp_obj_set_t *)MP_OBJ_TO_PTR(set_in);
151     o->cur = 0;
152     return MP_OBJ_FROM_PTR(o);
153 }
154 
155 /******************************************************************************/
156 /* set methods                                                                */
157 
set_add(mp_obj_t self_in,mp_obj_t item)158 STATIC mp_obj_t set_add(mp_obj_t self_in, mp_obj_t item) {
159     check_set(self_in);
160     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
161     mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
162     return mp_const_none;
163 }
164 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_add_obj, set_add);
165 
set_clear(mp_obj_t self_in)166 STATIC mp_obj_t set_clear(mp_obj_t self_in) {
167     check_set(self_in);
168     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
169     mp_set_clear(&self->set);
170     return mp_const_none;
171 }
172 STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_clear_obj, set_clear);
173 
set_copy(mp_obj_t self_in)174 STATIC mp_obj_t set_copy(mp_obj_t self_in) {
175     check_set_or_frozenset(self_in);
176     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
177     mp_obj_set_t *other = m_new_obj(mp_obj_set_t);
178     other->base.type = self->base.type;
179     mp_set_init(&other->set, self->set.alloc);
180     other->set.used = self->set.used;
181     memcpy(other->set.table, self->set.table, self->set.alloc * sizeof(mp_obj_t));
182     return MP_OBJ_FROM_PTR(other);
183 }
184 STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_copy_obj, set_copy);
185 
set_discard(mp_obj_t self_in,mp_obj_t item)186 STATIC mp_obj_t set_discard(mp_obj_t self_in, mp_obj_t item) {
187     check_set(self_in);
188     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
189     mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_REMOVE_IF_FOUND);
190     return mp_const_none;
191 }
192 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_discard_obj, set_discard);
193 
set_diff_int(size_t n_args,const mp_obj_t * args,bool update)194 STATIC mp_obj_t set_diff_int(size_t n_args, const mp_obj_t *args, bool update) {
195     mp_obj_t self;
196     if (update) {
197         check_set(args[0]);
198         self = args[0];
199     } else {
200         self = set_copy(args[0]);
201     }
202 
203     for (size_t i = 1; i < n_args; i++) {
204         mp_obj_t other = args[i];
205         if (self == other) {
206             set_clear(self);
207         } else {
208             mp_set_t *self_set = &((mp_obj_set_t *)MP_OBJ_TO_PTR(self))->set;
209             mp_obj_t iter = mp_getiter(other, NULL);
210             mp_obj_t next;
211             while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
212                 mp_set_lookup(self_set, next, MP_MAP_LOOKUP_REMOVE_IF_FOUND);
213             }
214         }
215     }
216 
217     return self;
218 }
219 
set_diff(size_t n_args,const mp_obj_t * args)220 STATIC mp_obj_t set_diff(size_t n_args, const mp_obj_t *args) {
221     return set_diff_int(n_args, args, false);
222 }
223 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_obj, 1, set_diff);
224 
set_diff_update(size_t n_args,const mp_obj_t * args)225 STATIC mp_obj_t set_diff_update(size_t n_args, const mp_obj_t *args) {
226     set_diff_int(n_args, args, true);
227     return mp_const_none;
228 }
229 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_update_obj, 1, set_diff_update);
230 
set_intersect_int(mp_obj_t self_in,mp_obj_t other,bool update)231 STATIC mp_obj_t set_intersect_int(mp_obj_t self_in, mp_obj_t other, bool update) {
232     if (update) {
233         check_set(self_in);
234     } else {
235         check_set_or_frozenset(self_in);
236     }
237 
238     if (self_in == other) {
239         return update ? mp_const_none : set_copy(self_in);
240     }
241 
242     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
243     mp_obj_set_t *out = MP_OBJ_TO_PTR(mp_obj_new_set(0, NULL));
244 
245     mp_obj_t iter = mp_getiter(other, NULL);
246     mp_obj_t next;
247     while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
248         if (mp_set_lookup(&self->set, next, MP_MAP_LOOKUP)) {
249             set_add(MP_OBJ_FROM_PTR(out), next);
250         }
251     }
252 
253     if (update) {
254         m_del(mp_obj_t, self->set.table, self->set.alloc);
255         self->set.alloc = out->set.alloc;
256         self->set.used = out->set.used;
257         self->set.table = out->set.table;
258     }
259 
260     return update ? mp_const_none : MP_OBJ_FROM_PTR(out);
261 }
262 
set_intersect(mp_obj_t self_in,mp_obj_t other)263 STATIC mp_obj_t set_intersect(mp_obj_t self_in, mp_obj_t other) {
264     return set_intersect_int(self_in, other, false);
265 }
266 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_intersect_obj, set_intersect);
267 
set_intersect_update(mp_obj_t self_in,mp_obj_t other)268 STATIC mp_obj_t set_intersect_update(mp_obj_t self_in, mp_obj_t other) {
269     return set_intersect_int(self_in, other, true);
270 }
271 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_intersect_update_obj, set_intersect_update);
272 
set_isdisjoint(mp_obj_t self_in,mp_obj_t other)273 STATIC mp_obj_t set_isdisjoint(mp_obj_t self_in, mp_obj_t other) {
274     check_set_or_frozenset(self_in);
275     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
276 
277     mp_obj_iter_buf_t iter_buf;
278     mp_obj_t iter = mp_getiter(other, &iter_buf);
279     mp_obj_t next;
280     while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
281         if (mp_set_lookup(&self->set, next, MP_MAP_LOOKUP)) {
282             return mp_const_false;
283         }
284     }
285     return mp_const_true;
286 }
287 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_isdisjoint_obj, set_isdisjoint);
288 
set_issubset_internal(mp_obj_t self_in,mp_obj_t other_in,bool proper)289 STATIC mp_obj_t set_issubset_internal(mp_obj_t self_in, mp_obj_t other_in, bool proper) {
290     mp_obj_set_t *self;
291     bool cleanup_self = false;
292     if (is_set_or_frozenset(self_in)) {
293         self = MP_OBJ_TO_PTR(self_in);
294     } else {
295         self = MP_OBJ_TO_PTR(set_make_new(&mp_type_set, 1, 0, &self_in));
296         cleanup_self = true;
297     }
298 
299     mp_obj_set_t *other;
300     bool cleanup_other = false;
301     if (is_set_or_frozenset(other_in)) {
302         other = MP_OBJ_TO_PTR(other_in);
303     } else {
304         other = MP_OBJ_TO_PTR(set_make_new(&mp_type_set, 1, 0, &other_in));
305         cleanup_other = true;
306     }
307     mp_obj_t out = mp_const_true;
308     if (proper && self->set.used == other->set.used) {
309         out = mp_const_false;
310     } else {
311         mp_obj_iter_buf_t iter_buf;
312         mp_obj_t iter = set_getiter(MP_OBJ_FROM_PTR(self), &iter_buf);
313         mp_obj_t next;
314         while ((next = set_it_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
315             if (!mp_set_lookup(&other->set, next, MP_MAP_LOOKUP)) {
316                 out = mp_const_false;
317                 break;
318             }
319         }
320     }
321     // TODO: Should free objects altogether
322     if (cleanup_self) {
323         set_clear(MP_OBJ_FROM_PTR(self));
324     }
325     if (cleanup_other) {
326         set_clear(MP_OBJ_FROM_PTR(other));
327     }
328     return out;
329 }
330 
set_issubset(mp_obj_t self_in,mp_obj_t other_in)331 STATIC mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) {
332     return set_issubset_internal(self_in, other_in, false);
333 }
334 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_issubset_obj, set_issubset);
335 
set_issubset_proper(mp_obj_t self_in,mp_obj_t other_in)336 STATIC mp_obj_t set_issubset_proper(mp_obj_t self_in, mp_obj_t other_in) {
337     return set_issubset_internal(self_in, other_in, true);
338 }
339 
set_issuperset(mp_obj_t self_in,mp_obj_t other_in)340 STATIC mp_obj_t set_issuperset(mp_obj_t self_in, mp_obj_t other_in) {
341     return set_issubset_internal(other_in, self_in, false);
342 }
343 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_issuperset_obj, set_issuperset);
344 
set_issuperset_proper(mp_obj_t self_in,mp_obj_t other_in)345 STATIC mp_obj_t set_issuperset_proper(mp_obj_t self_in, mp_obj_t other_in) {
346     return set_issubset_internal(other_in, self_in, true);
347 }
348 
set_equal(mp_obj_t self_in,mp_obj_t other_in)349 STATIC mp_obj_t set_equal(mp_obj_t self_in, mp_obj_t other_in) {
350     assert(is_set_or_frozenset(other_in));
351     check_set_or_frozenset(self_in);
352     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
353     mp_obj_set_t *other = MP_OBJ_TO_PTR(other_in);
354     if (self->set.used != other->set.used) {
355         return mp_const_false;
356     }
357     return set_issubset(self_in, other_in);
358 }
359 
set_pop(mp_obj_t self_in)360 STATIC mp_obj_t set_pop(mp_obj_t self_in) {
361     check_set(self_in);
362     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
363     mp_obj_t obj = mp_set_remove_first(&self->set);
364     if (obj == MP_OBJ_NULL) {
365         mp_raise_msg(&mp_type_KeyError, MP_ERROR_TEXT("pop from an empty set"));
366     }
367     return obj;
368 }
369 STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_pop_obj, set_pop);
370 
set_remove(mp_obj_t self_in,mp_obj_t item)371 STATIC mp_obj_t set_remove(mp_obj_t self_in, mp_obj_t item) {
372     check_set(self_in);
373     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
374     if (mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_REMOVE_IF_FOUND) == MP_OBJ_NULL) {
375         mp_raise_type_arg(&mp_type_KeyError, item);
376     }
377     return mp_const_none;
378 }
379 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_remove_obj, set_remove);
380 
set_symmetric_difference_update(mp_obj_t self_in,mp_obj_t other_in)381 STATIC mp_obj_t set_symmetric_difference_update(mp_obj_t self_in, mp_obj_t other_in) {
382     check_set_or_frozenset(self_in); // can be frozenset due to call from set_symmetric_difference
383     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
384     mp_obj_t iter = mp_getiter(other_in, NULL);
385     mp_obj_t next;
386     while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
387         mp_set_lookup(&self->set, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND_OR_REMOVE_IF_FOUND);
388     }
389     return mp_const_none;
390 }
391 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_symmetric_difference_update_obj, set_symmetric_difference_update);
392 
set_symmetric_difference(mp_obj_t self_in,mp_obj_t other_in)393 STATIC mp_obj_t set_symmetric_difference(mp_obj_t self_in, mp_obj_t other_in) {
394     mp_obj_t self_out = set_copy(self_in);
395     set_symmetric_difference_update(self_out, other_in);
396     return self_out;
397 }
398 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_symmetric_difference_obj, set_symmetric_difference);
399 
set_update_int(mp_obj_set_t * self,mp_obj_t other_in)400 STATIC void set_update_int(mp_obj_set_t *self, mp_obj_t other_in) {
401     mp_obj_t iter = mp_getiter(other_in, NULL);
402     mp_obj_t next;
403     while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
404         mp_set_lookup(&self->set, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
405     }
406 }
407 
set_update(size_t n_args,const mp_obj_t * args)408 STATIC mp_obj_t set_update(size_t n_args, const mp_obj_t *args) {
409     check_set(args[0]);
410     for (size_t i = 1; i < n_args; i++) {
411         set_update_int(MP_OBJ_TO_PTR(args[0]), args[i]);
412     }
413 
414     return mp_const_none;
415 }
416 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_update_obj, 1, set_update);
417 
set_union(mp_obj_t self_in,mp_obj_t other_in)418 STATIC mp_obj_t set_union(mp_obj_t self_in, mp_obj_t other_in) {
419     check_set_or_frozenset(self_in);
420     mp_obj_t self = set_copy(self_in);
421     set_update_int(MP_OBJ_TO_PTR(self), other_in);
422     return self;
423 }
424 STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_union_obj, set_union);
425 
set_unary_op(mp_unary_op_t op,mp_obj_t self_in)426 STATIC mp_obj_t set_unary_op(mp_unary_op_t op, mp_obj_t self_in) {
427     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
428     switch (op) {
429         case MP_UNARY_OP_BOOL:
430             return mp_obj_new_bool(self->set.used != 0);
431         case MP_UNARY_OP_LEN:
432             return MP_OBJ_NEW_SMALL_INT(self->set.used);
433         #if MICROPY_PY_BUILTINS_FROZENSET
434         case MP_UNARY_OP_HASH:
435             if (mp_obj_is_type(self_in, &mp_type_frozenset)) {
436                 // start hash with unique value
437                 mp_int_t hash = (mp_int_t)(uintptr_t)&mp_type_frozenset;
438                 size_t max = self->set.alloc;
439                 mp_set_t *set = &self->set;
440 
441                 for (size_t i = 0; i < max; i++) {
442                     if (mp_set_slot_is_filled(set, i)) {
443                         hash += MP_OBJ_SMALL_INT_VALUE(mp_unary_op(MP_UNARY_OP_HASH, set->table[i]));
444                     }
445                 }
446                 return MP_OBJ_NEW_SMALL_INT(hash);
447             }
448             MP_FALLTHROUGH
449         #endif
450         default:
451             return MP_OBJ_NULL;      // op not supported
452     }
453 }
454 
set_binary_op(mp_binary_op_t op,mp_obj_t lhs,mp_obj_t rhs)455 STATIC mp_obj_t set_binary_op(mp_binary_op_t op, mp_obj_t lhs, mp_obj_t rhs) {
456     mp_obj_t args[] = {lhs, rhs};
457     #if MICROPY_PY_BUILTINS_FROZENSET
458     bool update = mp_obj_is_type(lhs, &mp_type_set);
459     #else
460     bool update = true;
461     #endif
462     if (op != MP_BINARY_OP_CONTAINS && !is_set_or_frozenset(rhs)) {
463         // For all ops except containment the RHS must be a set/frozenset
464         return MP_OBJ_NULL;
465     }
466     switch (op) {
467         case MP_BINARY_OP_OR:
468             return set_union(lhs, rhs);
469         case MP_BINARY_OP_XOR:
470             return set_symmetric_difference(lhs, rhs);
471         case MP_BINARY_OP_AND:
472             return set_intersect(lhs, rhs);
473         case MP_BINARY_OP_SUBTRACT:
474             return set_diff(2, args);
475         case MP_BINARY_OP_INPLACE_OR:
476             if (update) {
477                 set_update(2, args);
478                 return lhs;
479             } else {
480                 return set_union(lhs, rhs);
481             }
482         case MP_BINARY_OP_INPLACE_XOR:
483             if (update) {
484                 set_symmetric_difference_update(lhs, rhs);
485                 return lhs;
486             } else {
487                 return set_symmetric_difference(lhs, rhs);
488             }
489         case MP_BINARY_OP_INPLACE_AND:
490             rhs = set_intersect_int(lhs, rhs, update);
491             if (update) {
492                 return lhs;
493             } else {
494                 return rhs;
495             }
496         case MP_BINARY_OP_INPLACE_SUBTRACT:
497             return set_diff_int(2, args, update);
498         case MP_BINARY_OP_LESS:
499             return set_issubset_proper(lhs, rhs);
500         case MP_BINARY_OP_MORE:
501             return set_issuperset_proper(lhs, rhs);
502         case MP_BINARY_OP_EQUAL:
503             return set_equal(lhs, rhs);
504         case MP_BINARY_OP_LESS_EQUAL:
505             return set_issubset(lhs, rhs);
506         case MP_BINARY_OP_MORE_EQUAL:
507             return set_issuperset(lhs, rhs);
508         case MP_BINARY_OP_CONTAINS: {
509             mp_obj_set_t *o = MP_OBJ_TO_PTR(lhs);
510             mp_obj_t elem = mp_set_lookup(&o->set, rhs, MP_MAP_LOOKUP);
511             return mp_obj_new_bool(elem != MP_OBJ_NULL);
512         }
513         default:
514             return MP_OBJ_NULL; // op not supported
515     }
516 }
517 
518 /******************************************************************************/
519 /* set constructors & public C API                                            */
520 
521 STATIC const mp_rom_map_elem_t set_locals_dict_table[] = {
522     { MP_ROM_QSTR(MP_QSTR_add), MP_ROM_PTR(&set_add_obj) },
523     { MP_ROM_QSTR(MP_QSTR_clear), MP_ROM_PTR(&set_clear_obj) },
524     { MP_ROM_QSTR(MP_QSTR_copy), MP_ROM_PTR(&set_copy_obj) },
525     { MP_ROM_QSTR(MP_QSTR_discard), MP_ROM_PTR(&set_discard_obj) },
526     { MP_ROM_QSTR(MP_QSTR_difference), MP_ROM_PTR(&set_diff_obj) },
527     { MP_ROM_QSTR(MP_QSTR_difference_update), MP_ROM_PTR(&set_diff_update_obj) },
528     { MP_ROM_QSTR(MP_QSTR_intersection), MP_ROM_PTR(&set_intersect_obj) },
529     { MP_ROM_QSTR(MP_QSTR_intersection_update), MP_ROM_PTR(&set_intersect_update_obj) },
530     { MP_ROM_QSTR(MP_QSTR_isdisjoint), MP_ROM_PTR(&set_isdisjoint_obj) },
531     { MP_ROM_QSTR(MP_QSTR_issubset), MP_ROM_PTR(&set_issubset_obj) },
532     { MP_ROM_QSTR(MP_QSTR_issuperset), MP_ROM_PTR(&set_issuperset_obj) },
533     { MP_ROM_QSTR(MP_QSTR_pop), MP_ROM_PTR(&set_pop_obj) },
534     { MP_ROM_QSTR(MP_QSTR_remove), MP_ROM_PTR(&set_remove_obj) },
535     { MP_ROM_QSTR(MP_QSTR_symmetric_difference), MP_ROM_PTR(&set_symmetric_difference_obj) },
536     { MP_ROM_QSTR(MP_QSTR_symmetric_difference_update), MP_ROM_PTR(&set_symmetric_difference_update_obj) },
537     { MP_ROM_QSTR(MP_QSTR_union), MP_ROM_PTR(&set_union_obj) },
538     { MP_ROM_QSTR(MP_QSTR_update), MP_ROM_PTR(&set_update_obj) },
539     { MP_ROM_QSTR(MP_QSTR___contains__), MP_ROM_PTR(&mp_op_contains_obj) },
540 };
541 STATIC MP_DEFINE_CONST_DICT(set_locals_dict, set_locals_dict_table);
542 
543 const mp_obj_type_t mp_type_set = {
544     { &mp_type_type },
545     .name = MP_QSTR_set,
546     .print = set_print,
547     .make_new = set_make_new,
548     .unary_op = set_unary_op,
549     .binary_op = set_binary_op,
550     .getiter = set_getiter,
551     .locals_dict = (mp_obj_dict_t *)&set_locals_dict,
552 };
553 
554 #if MICROPY_PY_BUILTINS_FROZENSET
555 STATIC const mp_rom_map_elem_t frozenset_locals_dict_table[] = {
556     { MP_ROM_QSTR(MP_QSTR_copy), MP_ROM_PTR(&set_copy_obj) },
557     { MP_ROM_QSTR(MP_QSTR_difference), MP_ROM_PTR(&set_diff_obj) },
558     { MP_ROM_QSTR(MP_QSTR_intersection), MP_ROM_PTR(&set_intersect_obj) },
559     { MP_ROM_QSTR(MP_QSTR_isdisjoint), MP_ROM_PTR(&set_isdisjoint_obj) },
560     { MP_ROM_QSTR(MP_QSTR_issubset), MP_ROM_PTR(&set_issubset_obj) },
561     { MP_ROM_QSTR(MP_QSTR_issuperset), MP_ROM_PTR(&set_issuperset_obj) },
562     { MP_ROM_QSTR(MP_QSTR_symmetric_difference), MP_ROM_PTR(&set_symmetric_difference_obj) },
563     { MP_ROM_QSTR(MP_QSTR_union), MP_ROM_PTR(&set_union_obj) },
564     { MP_ROM_QSTR(MP_QSTR___contains__), MP_ROM_PTR(&mp_op_contains_obj) },
565 };
566 STATIC MP_DEFINE_CONST_DICT(frozenset_locals_dict, frozenset_locals_dict_table);
567 
568 const mp_obj_type_t mp_type_frozenset = {
569     { &mp_type_type },
570     .flags = MP_TYPE_FLAG_EQ_CHECKS_OTHER_TYPE,
571     .name = MP_QSTR_frozenset,
572     .print = set_print,
573     .make_new = set_make_new,
574     .unary_op = set_unary_op,
575     .binary_op = set_binary_op,
576     .getiter = set_getiter,
577     .locals_dict = (mp_obj_dict_t *)&frozenset_locals_dict,
578 };
579 #endif
580 
mp_obj_new_set(size_t n_args,mp_obj_t * items)581 mp_obj_t mp_obj_new_set(size_t n_args, mp_obj_t *items) {
582     mp_obj_set_t *o = m_new_obj(mp_obj_set_t);
583     o->base.type = &mp_type_set;
584     mp_set_init(&o->set, n_args);
585     for (size_t i = 0; i < n_args; i++) {
586         mp_set_lookup(&o->set, items[i], MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
587     }
588     return MP_OBJ_FROM_PTR(o);
589 }
590 
mp_obj_set_store(mp_obj_t self_in,mp_obj_t item)591 void mp_obj_set_store(mp_obj_t self_in, mp_obj_t item) {
592     mp_check_self(mp_obj_is_type(self_in, &mp_type_set));
593     mp_obj_set_t *self = MP_OBJ_TO_PTR(self_in);
594     mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
595 }
596 
597 #endif // MICROPY_PY_BUILTINS_SET
598