@@ -132,44 +132,77 @@ void mp_set_init(mp_set_t *set, int n) {
132132 set -> table = m_new0 (mp_obj_t , set -> alloc );
133133}
134134
135- mp_obj_t mp_set_lookup (mp_set_t * set , mp_obj_t index , bool add_if_not_found ) {
136- int hash = mp_obj_hash (index );
137- assert (set -> alloc ); /* FIXME: if alloc is ever 0 when doing a lookup, this'll fail: */
138- int pos = hash % set -> alloc ;
135+ static void mp_set_rehash (mp_set_t * set ) {
136+ int old_alloc = set -> alloc ;
137+ mp_obj_t * old_table = set -> table ;
138+ set -> alloc = get_doubling_prime_greater_or_equal_to (set -> alloc + 1 );
139+ set -> used = 0 ;
140+ set -> table = m_new0 (mp_obj_t , set -> alloc );
141+ for (int i = 0 ; i < old_alloc ; i ++ ) {
142+ if (old_table [i ] != NULL ) {
143+ mp_set_lookup (set , old_table [i ], true);
144+ }
145+ }
146+ m_del (mp_obj_t , old_table , old_alloc );
147+ }
148+
149+ mp_obj_t mp_set_lookup (mp_set_t * set , mp_obj_t index , mp_map_lookup_kind_t lookup_kind ) {
150+ int hash ;
151+ int pos ;
152+ if (set -> alloc == 0 ) {
153+ if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND ) {
154+ mp_set_rehash (set );
155+ } else {
156+ return NULL ;
157+ }
158+ }
159+ if (lookup_kind & MP_MAP_LOOKUP_FIRST ) {
160+ hash = 0 ;
161+ pos = 0 ;
162+ } else {
163+ hash = mp_obj_hash (index );;
164+ pos = hash % set -> alloc ;
165+ }
139166 for (;;) {
140167 mp_obj_t elem = set -> table [pos ];
141168 if (elem == MP_OBJ_NULL ) {
142169 // not in table
143- if (add_if_not_found ) {
170+ if (lookup_kind & MP_MAP_LOOKUP_ADD_IF_NOT_FOUND ) {
144171 if (set -> used + 1 >= set -> alloc ) {
145172 // not enough room in table, rehash it
146- int old_alloc = set -> alloc ;
147- mp_obj_t * old_table = set -> table ;
148- set -> alloc = get_doubling_prime_greater_or_equal_to (set -> alloc + 1 );
149- set -> used = 0 ;
150- set -> table = m_new (mp_obj_t , set -> alloc );
151- for (int i = 0 ; i < old_alloc ; i ++ ) {
152- if (old_table [i ] != NULL ) {
153- mp_set_lookup (set , old_table [i ], true);
154- }
155- }
156- m_del (mp_obj_t , old_table , old_alloc );
173+ mp_set_rehash (set );
157174 // restart the search for the new element
158175 pos = hash % set -> alloc ;
159176 } else {
160177 set -> used += 1 ;
161178 set -> table [pos ] = index ;
162179 return index ;
163180 }
181+ } else if (lookup_kind & MP_MAP_LOOKUP_FIRST ) {
182+ pos ++ ;
164183 } else {
165184 return MP_OBJ_NULL ;
166185 }
167- } else if (mp_obj_equal (elem , index )) {
186+ } else if (lookup_kind & MP_MAP_LOOKUP_FIRST || mp_obj_equal (elem , index )) {
168187 // found it
188+ if (lookup_kind & MP_MAP_LOOKUP_REMOVE_IF_FOUND ) {
189+ set -> used -- ;
190+ set -> table [pos ] = NULL ;
191+ }
169192 return elem ;
170193 } else {
171194 // not yet found, keep searching in this table
172195 pos = (pos + 1 ) % set -> alloc ;
173196 }
174197 }
175198}
199+
200+ void mp_set_clear (mp_set_t * set ) {
201+ set -> used = 0 ;
202+ machine_uint_t a = set -> alloc ;
203+ set -> alloc = 0 ;
204+ set -> table = m_renew (mp_obj_t , set -> table , a , set -> alloc );
205+ for (uint i = 0 ; i < set -> alloc ; i ++ ) {
206+ set -> table [i ] = NULL ;
207+ }
208+ }
0 commit comments