blob: 55aedde3eb18448a68bb82783f28974cfe1ace79 [file] [log] [blame]
Damien George04b91472014-05-03 23:27:38 +01001/*
2 * This file is part of the Micro Python project, http://micropython.org/
3 *
4 * The MIT License (MIT)
5 *
6 * Copyright (c) 2013, 2014 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
xbeefe34222014-03-16 00:14:26 -070027#include <stdbool.h>
John R. Lenton3b0bd872014-01-12 15:56:25 +000028#include <string.h>
Damiend99b0522013-12-21 18:17:45 +000029#include <assert.h>
30
Damien George51dfcb42015-01-01 20:27:54 +000031#include "py/nlr.h"
32#include "py/runtime.h"
33#include "py/runtime0.h"
34#include "py/builtin.h"
Damiend99b0522013-12-21 18:17:45 +000035
Damien George3ebd4d02014-06-01 13:46:47 +010036#if MICROPY_PY_BUILTINS_SET
37
Damiend99b0522013-12-21 18:17:45 +000038typedef struct _mp_obj_set_t {
39 mp_obj_base_t base;
40 mp_set_t set;
41} mp_obj_set_t;
42
John R. Lenton0ce03b42014-01-12 15:17:42 +000043typedef struct _mp_obj_set_it_t {
44 mp_obj_base_t base;
45 mp_obj_set_t *set;
Damien George40f3c022014-07-03 13:25:24 +010046 mp_uint_t cur;
John R. Lenton0ce03b42014-01-12 15:17:42 +000047} mp_obj_set_it_t;
48
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020049STATIC mp_obj_t set_it_iternext(mp_obj_t self_in);
John R. Lenton0ce03b42014-01-12 15:17:42 +000050
Paul Sokolovskyb181b582014-05-10 16:02:17 +030051STATIC bool is_set_or_frozenset(mp_obj_t o) {
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030052 return MP_OBJ_IS_TYPE(o, &mp_type_set)
Damien Georgefb510b32014-06-01 13:32:54 +010053#if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030054 || MP_OBJ_IS_TYPE(o, &mp_type_frozenset)
55#endif
56 ;
Paul Sokolovskyb181b582014-05-10 16:02:17 +030057}
58
Damien Georgefb510b32014-06-01 13:32:54 +010059#if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +030060STATIC void check_set_or_frozenset(mp_obj_t o) {
61 if (!is_set_or_frozenset(o)) {
62 nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'set' object required"));
63 }
64}
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030065#else
66#define check_set_or_frozenset(o) check_set(o)
67#endif
Paul Sokolovskyb181b582014-05-10 16:02:17 +030068
69STATIC void check_set(mp_obj_t o) {
70 if (!MP_OBJ_IS_TYPE(o, &mp_type_set)) {
71 // Emulate CPython behavior
72 // AttributeError: 'frozenset' object has no attribute 'add'
Damien Georgefb510b32014-06-01 13:32:54 +010073 #if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +030074 if (MP_OBJ_IS_TYPE(o, &mp_type_frozenset)) {
75 nlr_raise(mp_obj_new_exception_msg(&mp_type_AttributeError, "'frozenset' has no such attribute"));
76 }
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030077 #endif
Paul Sokolovskyb181b582014-05-10 16:02:17 +030078 nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "'set' object required"));
79 }
80}
81
Damien George7f9d1d62015-04-09 23:56:15 +010082STATIC void set_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
Damien Georgeff8dd3f2015-01-20 12:47:20 +000083 (void)kind;
Damiend99b0522013-12-21 18:17:45 +000084 mp_obj_set_t *self = self_in;
Damien Georgefb510b32014-06-01 13:32:54 +010085 #if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +030086 bool is_frozen = MP_OBJ_IS_TYPE(self_in, &mp_type_frozenset);
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030087 #endif
John R. Lenton7244a142014-01-12 23:37:45 +000088 if (self->set.used == 0) {
Damien Georgefb510b32014-06-01 13:32:54 +010089 #if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +030090 if (is_frozen) {
Damien George7f9d1d62015-04-09 23:56:15 +010091 mp_print_str(print, "frozen");
Paul Sokolovskyb181b582014-05-10 16:02:17 +030092 }
Paul Sokolovskyd80e2472014-05-10 16:11:04 +030093 #endif
Damien George7f9d1d62015-04-09 23:56:15 +010094 mp_print_str(print, "set()");
John R. Lenton7244a142014-01-12 23:37:45 +000095 return;
96 }
Damiend99b0522013-12-21 18:17:45 +000097 bool first = true;
Damien Georgefb510b32014-06-01 13:32:54 +010098 #if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +030099 if (is_frozen) {
Damien George7f9d1d62015-04-09 23:56:15 +0100100 mp_print_str(print, "frozenset(");
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300101 }
Paul Sokolovskyd80e2472014-05-10 16:11:04 +0300102 #endif
Damien George7f9d1d62015-04-09 23:56:15 +0100103 mp_print_str(print, "{");
Damien George93965e72014-08-30 13:23:35 +0100104 for (mp_uint_t i = 0; i < self->set.alloc; i++) {
Damien George8b0535e2014-04-05 21:53:54 +0100105 if (MP_SET_SLOT_IS_FILLED(&self->set, i)) {
Damiend99b0522013-12-21 18:17:45 +0000106 if (!first) {
Damien George7f9d1d62015-04-09 23:56:15 +0100107 mp_print_str(print, ", ");
Damiend99b0522013-12-21 18:17:45 +0000108 }
109 first = false;
Damien George7f9d1d62015-04-09 23:56:15 +0100110 mp_obj_print_helper(print, self->set.table[i], PRINT_REPR);
Damiend99b0522013-12-21 18:17:45 +0000111 }
112 }
Damien George7f9d1d62015-04-09 23:56:15 +0100113 mp_print_str(print, "}");
Damien Georgefb510b32014-06-01 13:32:54 +0100114 #if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300115 if (is_frozen) {
Damien George7f9d1d62015-04-09 23:56:15 +0100116 mp_print_str(print, ")");
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300117 }
Paul Sokolovskyd80e2472014-05-10 16:11:04 +0300118 #endif
Damiend99b0522013-12-21 18:17:45 +0000119}
120
Damien Georgeecc88e92014-08-30 00:35:11 +0100121STATIC mp_obj_t set_make_new(mp_obj_t type_in, mp_uint_t n_args, mp_uint_t n_kw, const mp_obj_t *args) {
Damien George1d34e322014-05-11 18:28:48 +0100122 mp_arg_check_num(n_args, n_kw, 0, 1, false);
Damien George20006db2014-01-18 14:10:48 +0000123
Damien George71c51812014-01-04 20:21:15 +0000124 switch (n_args) {
Damien George1d34e322014-05-11 18:28:48 +0100125 case 0: {
126 // create a new, empty set
127 mp_obj_set_t *set = mp_obj_new_set(0, NULL);
128 // set actual set/frozenset type
129 set->base.type = type_in;
130 return set;
131 }
Damien George71c51812014-01-04 20:21:15 +0000132
133 case 1:
Damien George1d34e322014-05-11 18:28:48 +0100134 default: { // can only be 0 or 1 arg
Damien George71c51812014-01-04 20:21:15 +0000135 // 1 argument, an iterable from which we make a new set
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300136 mp_obj_set_t *set = mp_obj_new_set(0, NULL);
Damien Georged17926d2014-03-30 13:35:08 +0100137 mp_obj_t iterable = mp_getiter(args[0]);
Damien George71c51812014-01-04 20:21:15 +0000138 mp_obj_t item;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100139 while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
Damien George71c51812014-01-04 20:21:15 +0000140 mp_obj_set_store(set, item);
141 }
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300142 // Set actual set/frozenset type
143 set->base.type = type_in;
Damien George71c51812014-01-04 20:21:15 +0000144 return set;
145 }
Damien George71c51812014-01-04 20:21:15 +0000146 }
147}
148
Damien George3e1a5c12014-03-29 13:43:38 +0000149const mp_obj_type_t mp_type_set_it = {
Damien Georgec5966122014-02-15 16:10:44 +0000150 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000151 .name = MP_QSTR_iterator,
Paul Sokolovskyf7eaf602014-03-30 22:00:12 +0300152 .getiter = mp_identity,
John R. Lenton0ce03b42014-01-12 15:17:42 +0000153 .iternext = set_it_iternext,
154};
155
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200156STATIC mp_obj_t set_it_iternext(mp_obj_t self_in) {
Damien George3e1a5c12014-03-29 13:43:38 +0000157 assert(MP_OBJ_IS_TYPE(self_in, &mp_type_set_it));
John R. Lenton0ce03b42014-01-12 15:17:42 +0000158 mp_obj_set_it_t *self = self_in;
Damien George40f3c022014-07-03 13:25:24 +0100159 mp_uint_t max = self->set->set.alloc;
Damien George8b0535e2014-04-05 21:53:54 +0100160 mp_set_t *set = &self->set->set;
John R. Lenton0ce03b42014-01-12 15:17:42 +0000161
Damien George40f3c022014-07-03 13:25:24 +0100162 for (mp_uint_t i = self->cur; i < max; i++) {
Damien George8b0535e2014-04-05 21:53:54 +0100163 if (MP_SET_SLOT_IS_FILLED(set, i)) {
John R. Lenton0ce03b42014-01-12 15:17:42 +0000164 self->cur = i + 1;
Damien George8b0535e2014-04-05 21:53:54 +0100165 return set->table[i];
John R. Lenton0ce03b42014-01-12 15:17:42 +0000166 }
167 }
168
Damien Georgeea8d06c2014-04-17 23:19:36 +0100169 return MP_OBJ_STOP_ITERATION;
John R. Lenton0ce03b42014-01-12 15:17:42 +0000170}
171
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200172STATIC mp_obj_t set_getiter(mp_obj_t set_in) {
John R. Lenton0ce03b42014-01-12 15:17:42 +0000173 mp_obj_set_it_t *o = m_new_obj(mp_obj_set_it_t);
Damien George3e1a5c12014-03-29 13:43:38 +0000174 o->base.type = &mp_type_set_it;
John R. Lenton0ce03b42014-01-12 15:17:42 +0000175 o->set = (mp_obj_set_t *)set_in;
176 o->cur = 0;
177 return o;
178}
179
John R. Lenton19b14d32014-01-12 15:29:11 +0000180
181/******************************************************************************/
182/* set methods */
183
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200184STATIC mp_obj_t set_add(mp_obj_t self_in, mp_obj_t item) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300185 check_set(self_in);
John R. Lenton19b14d32014-01-12 15:29:11 +0000186 mp_obj_set_t *self = self_in;
John R. Lenton2a241722014-01-12 16:39:39 +0000187 mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
John R. Lenton19b14d32014-01-12 15:29:11 +0000188 return mp_const_none;
189}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200190STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_add_obj, set_add);
John R. Lenton19b14d32014-01-12 15:29:11 +0000191
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200192STATIC mp_obj_t set_clear(mp_obj_t self_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300193 check_set(self_in);
John R. Lenton1d7fb2f2014-01-12 15:44:26 +0000194 mp_obj_set_t *self = self_in;
195
196 mp_set_clear(&self->set);
197
198 return mp_const_none;
199}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200200STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_clear_obj, set_clear);
John R. Lenton1d7fb2f2014-01-12 15:44:26 +0000201
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300202STATIC mp_obj_t set_copy_as_mutable(mp_obj_t self_in) {
John R. Lenton3b0bd872014-01-12 15:56:25 +0000203 mp_obj_set_t *self = self_in;
204
205 mp_obj_set_t *other = m_new_obj(mp_obj_set_t);
Damien George3e1a5c12014-03-29 13:43:38 +0000206 other->base.type = &mp_type_set;
Paul Sokolovsky46bd12d2014-04-07 03:07:21 +0300207 mp_set_init(&other->set, self->set.alloc);
John R. Lenton3b0bd872014-01-12 15:56:25 +0000208 other->set.used = self->set.used;
209 memcpy(other->set.table, self->set.table, self->set.alloc * sizeof(mp_obj_t));
210
211 return other;
212}
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300213
214STATIC mp_obj_t set_copy(mp_obj_t self_in) {
215 check_set_or_frozenset(self_in);
216 mp_obj_set_t *self = self_in;
217
218 mp_obj_set_t *other = set_copy_as_mutable(self);
219 other->base.type = self->base.type;
220
221 return other;
222}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200223STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_copy_obj, set_copy);
John R. Lenton3b0bd872014-01-12 15:56:25 +0000224
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200225STATIC mp_obj_t set_discard(mp_obj_t self_in, mp_obj_t item) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300226 check_set(self_in);
John R. Lenton2a241722014-01-12 16:39:39 +0000227 mp_obj_set_t *self = self_in;
228 mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_REMOVE_IF_FOUND);
229 return mp_const_none;
230}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200231STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_discard_obj, set_discard);
John R. Lenton19b14d32014-01-12 15:29:11 +0000232
Damien Georgeecc88e92014-08-30 00:35:11 +0100233STATIC mp_obj_t set_diff_int(mp_uint_t n_args, const mp_obj_t *args, bool update) {
John R. Lenton032129f2014-01-12 17:07:17 +0000234 assert(n_args > 0);
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300235
John R. Lenton032129f2014-01-12 17:07:17 +0000236 mp_obj_set_t *self;
237 if (update) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300238 check_set(args[0]);
John R. Lenton032129f2014-01-12 17:07:17 +0000239 self = args[0];
240 } else {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300241 check_set_or_frozenset(args[0]);
242 self = set_copy_as_mutable(args[0]);
John R. Lenton032129f2014-01-12 17:07:17 +0000243 }
244
245
Damien George93965e72014-08-30 13:23:35 +0100246 for (mp_uint_t i = 1; i < n_args; i++) {
John R. Lenton032129f2014-01-12 17:07:17 +0000247 mp_obj_t other = args[i];
248 if (self == other) {
249 set_clear(self);
250 } else {
Damien Georged17926d2014-03-30 13:35:08 +0100251 mp_obj_t iter = mp_getiter(other);
John R. Lenton032129f2014-01-12 17:07:17 +0000252 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100253 while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
John R. Lenton032129f2014-01-12 17:07:17 +0000254 set_discard(self, next);
255 }
256 }
257 }
258
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300259 self->base.type = ((mp_obj_set_t*)args[0])->base.type;
John R. Lenton032129f2014-01-12 17:07:17 +0000260 return self;
261}
262
Damien Georgeecc88e92014-08-30 00:35:11 +0100263STATIC mp_obj_t set_diff(mp_uint_t n_args, const mp_obj_t *args) {
John R. Lenton032129f2014-01-12 17:07:17 +0000264 return set_diff_int(n_args, args, false);
265}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200266STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_obj, 1, set_diff);
John R. Lenton032129f2014-01-12 17:07:17 +0000267
Damien Georgeecc88e92014-08-30 00:35:11 +0100268STATIC mp_obj_t set_diff_update(mp_uint_t n_args, const mp_obj_t *args) {
John R. Lenton032129f2014-01-12 17:07:17 +0000269 set_diff_int(n_args, args, true);
270 return mp_const_none;
271}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200272STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_update_obj, 1, set_diff_update);
John R. Lenton032129f2014-01-12 17:07:17 +0000273
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200274STATIC mp_obj_t set_intersect_int(mp_obj_t self_in, mp_obj_t other, bool update) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300275 if (update) {
276 check_set(self_in);
277 } else {
278 check_set_or_frozenset(self_in);
279 }
280
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000281 if (self_in == other) {
282 return update ? mp_const_none : set_copy(self_in);
283 }
284
285 mp_obj_set_t *self = self_in;
286 mp_obj_set_t *out = mp_obj_new_set(0, NULL);
287
Damien Georged17926d2014-03-30 13:35:08 +0100288 mp_obj_t iter = mp_getiter(other);
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000289 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100290 while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000291 if (mp_set_lookup(&self->set, next, MP_MAP_LOOKUP)) {
292 set_add(out, next);
293 }
294 }
295
296 if (update) {
297 m_del(mp_obj_t, self->set.table, self->set.alloc);
298 self->set.alloc = out->set.alloc;
299 self->set.used = out->set.used;
300 self->set.table = out->set.table;
301 }
302
303 return update ? mp_const_none : out;
304}
305
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200306STATIC mp_obj_t set_intersect(mp_obj_t self_in, mp_obj_t other) {
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000307 return set_intersect_int(self_in, other, false);
308}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200309STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_intersect_obj, set_intersect);
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000310
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200311STATIC mp_obj_t set_intersect_update(mp_obj_t self_in, mp_obj_t other) {
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000312 return set_intersect_int(self_in, other, true);
313}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200314STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_intersect_update_obj, set_intersect_update);
John R. Lentonf1ae6b42014-01-12 17:54:03 +0000315
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200316STATIC mp_obj_t set_isdisjoint(mp_obj_t self_in, mp_obj_t other) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300317 check_set_or_frozenset(self_in);
John R. Lenton4a080672014-01-12 18:03:21 +0000318 mp_obj_set_t *self = self_in;
319
Damien Georged17926d2014-03-30 13:35:08 +0100320 mp_obj_t iter = mp_getiter(other);
John R. Lenton4a080672014-01-12 18:03:21 +0000321 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100322 while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
John R. Lenton4a080672014-01-12 18:03:21 +0000323 if (mp_set_lookup(&self->set, next, MP_MAP_LOOKUP)) {
324 return mp_const_false;
325 }
326 }
327 return mp_const_true;
328}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200329STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_isdisjoint_obj, set_isdisjoint);
John R. Lenton4a080672014-01-12 18:03:21 +0000330
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200331STATIC mp_obj_t set_issubset_internal(mp_obj_t self_in, mp_obj_t other_in, bool proper) {
John R. Lentonae00d332014-01-12 18:23:36 +0000332 mp_obj_set_t *self;
333 bool cleanup_self = false;
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300334 if (is_set_or_frozenset(self_in)) {
John R. Lentonae00d332014-01-12 18:23:36 +0000335 self = self_in;
336 } else {
Damien George3e1a5c12014-03-29 13:43:38 +0000337 self = set_make_new((mp_obj_t)&mp_type_set, 1, 0, &self_in);
John R. Lentonae00d332014-01-12 18:23:36 +0000338 cleanup_self = true;
339 }
340
341 mp_obj_set_t *other;
342 bool cleanup_other = false;
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300343 if (is_set_or_frozenset(other_in)) {
John R. Lentonae00d332014-01-12 18:23:36 +0000344 other = other_in;
345 } else {
Damien George3e1a5c12014-03-29 13:43:38 +0000346 other = set_make_new((mp_obj_t)&mp_type_set, 1, 0, &other_in);
John R. Lentonae00d332014-01-12 18:23:36 +0000347 cleanup_other = true;
348 }
John R. Lentonbe790f92014-01-12 23:09:10 +0000349 bool out = true;
350 if (proper && self->set.used == other->set.used) {
351 out = false;
352 } else {
353 mp_obj_t iter = set_getiter(self);
354 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100355 while ((next = set_it_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000356 if (!mp_set_lookup(&other->set, next, MP_MAP_LOOKUP)) {
357 out = false;
358 break;
359 }
John R. Lentonae00d332014-01-12 18:23:36 +0000360 }
361 }
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300362 // TODO: Should free objects altogether
John R. Lentonae00d332014-01-12 18:23:36 +0000363 if (cleanup_self) {
364 set_clear(self);
365 }
366 if (cleanup_other) {
367 set_clear(other);
368 }
Paul Sokolovsky1b586f32015-10-11 12:09:43 +0300369 return mp_obj_new_bool(out);
John R. Lentonbe790f92014-01-12 23:09:10 +0000370}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200371STATIC mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000372 return set_issubset_internal(self_in, other_in, false);
John R. Lentonae00d332014-01-12 18:23:36 +0000373}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200374STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_issubset_obj, set_issubset);
John R. Lentonae00d332014-01-12 18:23:36 +0000375
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200376STATIC mp_obj_t set_issubset_proper(mp_obj_t self_in, mp_obj_t other_in) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000377 return set_issubset_internal(self_in, other_in, true);
378}
379
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200380STATIC mp_obj_t set_issuperset(mp_obj_t self_in, mp_obj_t other_in) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000381 return set_issubset_internal(other_in, self_in, false);
John R. Lentonae00d332014-01-12 18:23:36 +0000382}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200383STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_issuperset_obj, set_issuperset);
John R. Lentonae00d332014-01-12 18:23:36 +0000384
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200385STATIC mp_obj_t set_issuperset_proper(mp_obj_t self_in, mp_obj_t other_in) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000386 return set_issubset_internal(other_in, self_in, true);
387}
388
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200389STATIC mp_obj_t set_equal(mp_obj_t self_in, mp_obj_t other_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300390 check_set_or_frozenset(self_in);
John R. Lentonbe790f92014-01-12 23:09:10 +0000391 mp_obj_set_t *self = self_in;
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300392 if (!is_set_or_frozenset(other_in)) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000393 return mp_const_false;
394 }
395 mp_obj_set_t *other = other_in;
396 if (self->set.used != other->set.used) {
397 return mp_const_false;
398 }
399 return set_issubset(self_in, other_in);
400}
401
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200402STATIC mp_obj_t set_pop(mp_obj_t self_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300403 check_set(self_in);
John R. Lentonae00d332014-01-12 18:23:36 +0000404 mp_obj_set_t *self = self_in;
Damien George95004e52014-04-05 17:17:19 +0100405 mp_obj_t obj = mp_set_remove_first(&self->set);
406 if (obj == MP_OBJ_NULL) {
Damien Georgeea13f402014-04-05 18:32:08 +0100407 nlr_raise(mp_obj_new_exception_msg(&mp_type_KeyError, "pop from an empty set"));
John R. Lentonae00d332014-01-12 18:23:36 +0000408 }
John R. Lentonae00d332014-01-12 18:23:36 +0000409 return obj;
410}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200411STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_pop_obj, set_pop);
John R. Lentonae00d332014-01-12 18:23:36 +0000412
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200413STATIC mp_obj_t set_remove(mp_obj_t self_in, mp_obj_t item) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300414 check_set(self_in);
John R. Lentonae00d332014-01-12 18:23:36 +0000415 mp_obj_set_t *self = self_in;
416 if (mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_REMOVE_IF_FOUND) == MP_OBJ_NULL) {
Damien Georgeea13f402014-04-05 18:32:08 +0100417 nlr_raise(mp_obj_new_exception(&mp_type_KeyError));
John R. Lentonae00d332014-01-12 18:23:36 +0000418 }
419 return mp_const_none;
420}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200421STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_remove_obj, set_remove);
John R. Lenton032129f2014-01-12 17:07:17 +0000422
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200423STATIC mp_obj_t set_symmetric_difference_update(mp_obj_t self_in, mp_obj_t other_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300424 check_set(self_in);
John R. Lenton0de386b2014-01-12 19:39:48 +0000425 mp_obj_set_t *self = self_in;
Damien Georged17926d2014-03-30 13:35:08 +0100426 mp_obj_t iter = mp_getiter(other_in);
John R. Lenton0de386b2014-01-12 19:39:48 +0000427 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100428 while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
Damien Georged1cee022015-03-20 17:41:37 +0000429 mp_set_lookup(&self->set, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND_OR_REMOVE_IF_FOUND);
John R. Lenton0de386b2014-01-12 19:39:48 +0000430 }
431 return mp_const_none;
432}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200433STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_symmetric_difference_update_obj, set_symmetric_difference_update);
John R. Lenton0de386b2014-01-12 19:39:48 +0000434
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200435STATIC mp_obj_t set_symmetric_difference(mp_obj_t self_in, mp_obj_t other_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300436 check_set_or_frozenset(self_in);
437 mp_obj_set_t *self_out = set_copy_as_mutable(self_in);
438 set_symmetric_difference_update(self_out, other_in);
439 self_out->base.type = ((mp_obj_set_t*)self_in)->base.type;
440 return self_out;
John R. Lenton0de386b2014-01-12 19:39:48 +0000441}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200442STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_symmetric_difference_obj, set_symmetric_difference);
John R. Lenton0de386b2014-01-12 19:39:48 +0000443
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200444STATIC void set_update_int(mp_obj_set_t *self, mp_obj_t other_in) {
Damien Georged17926d2014-03-30 13:35:08 +0100445 mp_obj_t iter = mp_getiter(other_in);
John R. Lenton0de386b2014-01-12 19:39:48 +0000446 mp_obj_t next;
Damien Georgeea8d06c2014-04-17 23:19:36 +0100447 while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
John R. Lenton0de386b2014-01-12 19:39:48 +0000448 mp_set_lookup(&self->set, next, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
449 }
450}
451
Damien Georgeecc88e92014-08-30 00:35:11 +0100452STATIC mp_obj_t set_update(mp_uint_t n_args, const mp_obj_t *args) {
John R. Lenton0de386b2014-01-12 19:39:48 +0000453 assert(n_args > 0);
John R. Lenton0de386b2014-01-12 19:39:48 +0000454
Damien George93965e72014-08-30 13:23:35 +0100455 for (mp_uint_t i = 1; i < n_args; i++) {
John R. Lenton0de386b2014-01-12 19:39:48 +0000456 set_update_int(args[0], args[i]);
457 }
458
459 return mp_const_none;
460}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200461STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(set_update_obj, 1, set_update);
John R. Lenton0de386b2014-01-12 19:39:48 +0000462
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200463STATIC mp_obj_t set_union(mp_obj_t self_in, mp_obj_t other_in) {
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300464 check_set_or_frozenset(self_in);
John R. Lenton0de386b2014-01-12 19:39:48 +0000465 mp_obj_set_t *self = set_copy(self_in);
466 set_update_int(self, other_in);
467 return self;
468}
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200469STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_union_obj, set_union);
John R. Lenton0de386b2014-01-12 19:39:48 +0000470
Damien Georgeecc88e92014-08-30 00:35:11 +0100471STATIC mp_obj_t set_unary_op(mp_uint_t op, mp_obj_t self_in) {
Damien George95004e52014-04-05 17:17:19 +0100472 mp_obj_set_t *self = self_in;
473 switch (op) {
Paul Sokolovsky1b586f32015-10-11 12:09:43 +0300474 case MP_UNARY_OP_BOOL: return mp_obj_new_bool(self->set.used != 0);
Damien Georgebb4c6f32014-07-31 10:49:14 +0100475 case MP_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT(self->set.used);
Paul Sokolovsky8b3b2d02015-08-28 22:31:31 +0300476#if MICROPY_PY_BUILTINS_FROZENSET
477 case MP_UNARY_OP_HASH:
478 if (MP_OBJ_IS_TYPE(self_in, &mp_type_frozenset)) {
479 // start hash with unique value
480 mp_int_t hash = (mp_int_t)&mp_type_frozenset;
481 mp_uint_t max = self->set.alloc;
482 mp_set_t *set = &self->set;
483
484 for (mp_uint_t i = 0; i < max; i++) {
485 if (MP_SET_SLOT_IS_FILLED(set, i)) {
486 hash += MP_OBJ_SMALL_INT_VALUE(mp_unary_op(MP_UNARY_OP_HASH, set->table[i]));
487 }
488 }
489 return MP_OBJ_NEW_SMALL_INT(hash);
490 }
491#endif
Damien George6ac5dce2014-05-21 19:42:43 +0100492 default: return MP_OBJ_NULL; // op not supported
Damien George95004e52014-04-05 17:17:19 +0100493 }
494}
John R. Lenton0de386b2014-01-12 19:39:48 +0000495
Damien Georgeecc88e92014-08-30 00:35:11 +0100496STATIC mp_obj_t set_binary_op(mp_uint_t op, mp_obj_t lhs, mp_obj_t rhs) {
John R. Lentonbe790f92014-01-12 23:09:10 +0000497 mp_obj_t args[] = {lhs, rhs};
498 switch (op) {
Damien Georged0a5bf32014-05-10 13:55:11 +0100499 case MP_BINARY_OP_OR:
500 return set_union(lhs, rhs);
501 case MP_BINARY_OP_XOR:
502 return set_symmetric_difference(lhs, rhs);
503 case MP_BINARY_OP_AND:
504 return set_intersect(lhs, rhs);
505 case MP_BINARY_OP_SUBTRACT:
506 return set_diff(2, args);
507 case MP_BINARY_OP_INPLACE_OR:
508 return set_union(lhs, rhs);
509 case MP_BINARY_OP_INPLACE_XOR:
510 return set_symmetric_difference(lhs, rhs);
511 case MP_BINARY_OP_INPLACE_AND:
512 return set_intersect(lhs, rhs);
513 case MP_BINARY_OP_INPLACE_SUBTRACT:
514 return set_diff(2, args);
515 case MP_BINARY_OP_LESS:
516 return set_issubset_proper(lhs, rhs);
517 case MP_BINARY_OP_MORE:
518 return set_issuperset_proper(lhs, rhs);
519 case MP_BINARY_OP_EQUAL:
520 return set_equal(lhs, rhs);
521 case MP_BINARY_OP_LESS_EQUAL:
522 return set_issubset(lhs, rhs);
523 case MP_BINARY_OP_MORE_EQUAL:
524 return set_issuperset(lhs, rhs);
525 case MP_BINARY_OP_IN: {
526 mp_obj_set_t *o = lhs;
527 mp_obj_t elem = mp_set_lookup(&o->set, rhs, MP_MAP_LOOKUP);
Paul Sokolovsky1b586f32015-10-11 12:09:43 +0300528 return mp_obj_new_bool(elem != NULL);
Damien Georged0a5bf32014-05-10 13:55:11 +0100529 }
530 default:
Damien George6ac5dce2014-05-21 19:42:43 +0100531 return MP_OBJ_NULL; // op not supported
John R. Lentonbe790f92014-01-12 23:09:10 +0000532 }
533}
John R. Lenton0de386b2014-01-12 19:39:48 +0000534
John R. Lenton19b14d32014-01-12 15:29:11 +0000535/******************************************************************************/
536/* set constructors & public C API */
537
538
Damien George9b196cd2014-03-26 21:47:19 +0000539STATIC const mp_map_elem_t set_locals_dict_table[] = {
540 { MP_OBJ_NEW_QSTR(MP_QSTR_add), (mp_obj_t)&set_add_obj },
541 { MP_OBJ_NEW_QSTR(MP_QSTR_clear), (mp_obj_t)&set_clear_obj },
542 { MP_OBJ_NEW_QSTR(MP_QSTR_copy), (mp_obj_t)&set_copy_obj },
543 { MP_OBJ_NEW_QSTR(MP_QSTR_discard), (mp_obj_t)&set_discard_obj },
544 { MP_OBJ_NEW_QSTR(MP_QSTR_difference), (mp_obj_t)&set_diff_obj },
545 { MP_OBJ_NEW_QSTR(MP_QSTR_difference_update), (mp_obj_t)&set_diff_update_obj },
546 { MP_OBJ_NEW_QSTR(MP_QSTR_intersection), (mp_obj_t)&set_intersect_obj },
547 { MP_OBJ_NEW_QSTR(MP_QSTR_intersection_update), (mp_obj_t)&set_intersect_update_obj },
548 { MP_OBJ_NEW_QSTR(MP_QSTR_isdisjoint), (mp_obj_t)&set_isdisjoint_obj },
549 { MP_OBJ_NEW_QSTR(MP_QSTR_issubset), (mp_obj_t)&set_issubset_obj },
550 { MP_OBJ_NEW_QSTR(MP_QSTR_issuperset), (mp_obj_t)&set_issuperset_obj },
551 { MP_OBJ_NEW_QSTR(MP_QSTR_pop), (mp_obj_t)&set_pop_obj },
552 { MP_OBJ_NEW_QSTR(MP_QSTR_remove), (mp_obj_t)&set_remove_obj },
553 { MP_OBJ_NEW_QSTR(MP_QSTR_symmetric_difference), (mp_obj_t)&set_symmetric_difference_obj },
554 { MP_OBJ_NEW_QSTR(MP_QSTR_symmetric_difference_update), (mp_obj_t)&set_symmetric_difference_update_obj },
555 { MP_OBJ_NEW_QSTR(MP_QSTR_union), (mp_obj_t)&set_union_obj },
556 { MP_OBJ_NEW_QSTR(MP_QSTR_update), (mp_obj_t)&set_update_obj },
Paul Sokolovsky68e7c512014-04-13 11:53:15 +0300557 { MP_OBJ_NEW_QSTR(MP_QSTR___contains__), (mp_obj_t)&mp_op_contains_obj },
John R. Lenton19b14d32014-01-12 15:29:11 +0000558};
559
Damien George9b196cd2014-03-26 21:47:19 +0000560STATIC MP_DEFINE_CONST_DICT(set_locals_dict, set_locals_dict_table);
561
Damien George3e1a5c12014-03-29 13:43:38 +0000562const mp_obj_type_t mp_type_set = {
Damien Georgec5966122014-02-15 16:10:44 +0000563 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000564 .name = MP_QSTR_set,
Damien George97209d32014-01-07 15:58:30 +0000565 .print = set_print,
566 .make_new = set_make_new,
Damien George95004e52014-04-05 17:17:19 +0100567 .unary_op = set_unary_op,
John R. Lentonbe790f92014-01-12 23:09:10 +0000568 .binary_op = set_binary_op,
John R. Lenton0ce03b42014-01-12 15:17:42 +0000569 .getiter = set_getiter,
Damien George9b196cd2014-03-26 21:47:19 +0000570 .locals_dict = (mp_obj_t)&set_locals_dict,
Damiend99b0522013-12-21 18:17:45 +0000571};
572
Damien Georgefb510b32014-06-01 13:32:54 +0100573#if MICROPY_PY_BUILTINS_FROZENSET
Paul Sokolovskyb181b582014-05-10 16:02:17 +0300574const mp_obj_type_t mp_type_frozenset = {
575 { &mp_type_type },
576 .name = MP_QSTR_frozenset,
577 .print = set_print,
578 .make_new = set_make_new,
579 .unary_op = set_unary_op,
580 .binary_op = set_binary_op,
581 .getiter = set_getiter,
582 .locals_dict = (mp_obj_t)&set_locals_dict,
583};
584#endif
585
Damien George93965e72014-08-30 13:23:35 +0100586mp_obj_t mp_obj_new_set(mp_uint_t n_args, mp_obj_t *items) {
Damiend99b0522013-12-21 18:17:45 +0000587 mp_obj_set_t *o = m_new_obj(mp_obj_set_t);
Damien George3e1a5c12014-03-29 13:43:38 +0000588 o->base.type = &mp_type_set;
Damiend99b0522013-12-21 18:17:45 +0000589 mp_set_init(&o->set, n_args);
Damien George93965e72014-08-30 13:23:35 +0100590 for (mp_uint_t i = 0; i < n_args; i++) {
John R. Lenton2a241722014-01-12 16:39:39 +0000591 mp_set_lookup(&o->set, items[i], MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
Damiend99b0522013-12-21 18:17:45 +0000592 }
593 return o;
594}
595
596void mp_obj_set_store(mp_obj_t self_in, mp_obj_t item) {
Damien George3e1a5c12014-03-29 13:43:38 +0000597 assert(MP_OBJ_IS_TYPE(self_in, &mp_type_set));
Damiend99b0522013-12-21 18:17:45 +0000598 mp_obj_set_t *self = self_in;
John R. Lenton2a241722014-01-12 16:39:39 +0000599 mp_set_lookup(&self->set, item, MP_MAP_LOOKUP_ADD_IF_NOT_FOUND);
Damiend99b0522013-12-21 18:17:45 +0000600}
Damien George3ebd4d02014-06-01 13:46:47 +0100601
602#endif // MICROPY_PY_BUILTINS_SET