From ea0ceff7fb3e4685085c2d2555af366917fa81ba Mon Sep 17 00:00:00 2001 From: "Yukihiro \"Matz\" Matsumoto" Date: Mon, 23 Jun 2025 09:50:32 +0900 Subject: [PATCH] mruby-set: implement Set#flatten and flatten! in C The code is generated by Cursor. --- mrbgems/mruby-set/mrblib/set.rb | 42 +------- mrbgems/mruby-set/src/set.c | 167 ++++++++++++++++++++++++++++++++ 2 files changed, 169 insertions(+), 40 deletions(-) diff --git a/mrbgems/mruby-set/mrblib/set.rb b/mrbgems/mruby-set/mrblib/set.rb index 7fe8d674b..65a315f01 100644 --- a/mrbgems/mruby-set/mrblib/set.rb +++ b/mrbgems/mruby-set/mrblib/set.rb @@ -137,46 +137,8 @@ class Set end end - # Recursively merges nested sets into this set, tracking seen sets to prevent - # infinite recursion. - # - # @param [Set] set The set to flatten and merge - # @param [Set] seen Set of object IDs to track visited sets - # @return [Set] self - def flatten_merge(set, seen = Set.new) - seen.add(set.object_id) - set.each { |e| - if e.is_a?(Set) - if seen.include?(e.object_id) - raise ArgumentError, "tried to flatten recursive Set" - end - - flatten_merge(e, seen) - else - add(e) - end - } - seen.delete(set.object_id) - self - end - - # Returns a new set that is a flattened version of this set. - # - # @return [Set] A new flattened set - def flatten - self.class.new.flatten_merge(self) - end - - # Replaces the contents of this set with a flattened version of itself. - # - # @return [Set] self if flattened, nil if no changes were made - def flatten! - if detect { |e| e.is_a?(Set) } - replace(flatten()) - else - nil - end - end + # flatten and flatten! are now implemented in C + # See mrbgems/mruby-set/src/set.c # superset? and proper_superset? are now implemented in C # See mrbgems/mruby-set/src/set.c diff --git a/mrbgems/mruby-set/src/set.c b/mrbgems/mruby-set/src/set.c index 5eb292c27..97c89a57b 100644 --- a/mrbgems/mruby-set/src/set.c +++ b/mrbgems/mruby-set/src/set.c @@ -1126,6 +1126,170 @@ set_add_all(mrb_state *mrb, mrb_value self) return self; } +/* + * Optimized implementation for flattening sets + * Uses a more efficient algorithm with minimal memory usage + */ + +/* Small array for tracking seen object IDs to detect cycles */ +#define MAX_NESTED_DEPTH 16 + +/* + * Recursively flattens a set by merging nested sets into the target set. + * This is an internal helper function that does not call back to the VM. + * + * @param mrb The mruby state + * @param target_kh The target hash table to add elements to + * @param source_kh The source hash table to flatten + * @param seen_ids Array of object IDs to track seen sets + * @param seen_count Pointer to the current count of seen sets + * @return 0 on success, -1 on error + */ +static int +set_flatten_recursive(mrb_state *mrb, khash_t(set) *target_kh, khash_t(set) *source_kh, + mrb_int *seen_ids, int *seen_count) +{ + if (!source_kh || !target_kh) return 0; + if (*seen_count >= MAX_NESTED_DEPTH) return -1; + + int ai = mrb_gc_arena_save(mrb); + + /* Process each element in the source set */ + KHASH_FOREACH(mrb, source_kh, k) { + mrb_value elem = kh_key(source_kh, k); + + /* Check if element is a Set */ + if (mrb_obj_is_kind_of(mrb, elem, mrb_class_get(mrb, "Set"))) { + /* Get the object ID to track recursion */ + mrb_int obj_id = mrb_obj_id(elem); + + /* Check if we've seen this set before to prevent infinite recursion */ + for (int i = 0; i < *seen_count; i++) { + if (seen_ids[i] == obj_id) { + mrb_raise(mrb, E_ARGUMENT_ERROR, "tried to flatten recursive Set"); + } + } + + /* Mark this set as seen */ + seen_ids[(*seen_count)++] = obj_id; + + /* Recursively flatten the nested set */ + khash_t(set) *nested_kh = set_get_khash(mrb, elem); + if (nested_kh) { + set_flatten_recursive(mrb, target_kh, nested_kh, seen_ids, seen_count); + } + + /* Remove from seen array after processing */ + (*seen_count)--; + } else { + /* Add non-Set element directly */ + kh_put(set, mrb, target_kh, elem); + } + } + + mrb_gc_arena_restore(mrb, ai); + return 0; +} + +/* + * call-seq: + * set.flatten -> new_set + * + * Returns a new set that is a flattened version of this set. + * Recursively flattens nested sets. + */ +static mrb_value +set_flatten(mrb_state *mrb, mrb_value self) +{ + khash_t(set) *self_kh = set_get_khash(mrb, self); + + /* Fast path for empty sets */ + if (!self_kh || kh_size(self_kh) == 0) { + return mrb_obj_new(mrb, mrb_obj_class(mrb, self), 0, NULL); + } + + /* Fast path: check if there are any nested sets */ + mrb_bool has_nested_sets = FALSE; + int ai = mrb_gc_arena_save(mrb); + KHASH_FOREACH(mrb, self_kh, k) { + if (mrb_obj_is_kind_of(mrb, kh_key(self_kh, k), mrb_class_get(mrb, "Set"))) { + has_nested_sets = TRUE; + break; + } + } + mrb_gc_arena_restore(mrb, ai); + + /* If no nested sets, just return a duplicate */ + if (!has_nested_sets) { + return mrb_obj_dup(mrb, self); + } + + /* Create a new set of the same class */ + mrb_value result_set = mrb_obj_new(mrb, mrb_obj_class(mrb, self), 0, NULL); + khash_t(set) *result_kh = set_get_khash(mrb, result_set); + + /* Use a small array for tracking seen object IDs */ + mrb_int seen_ids[MAX_NESTED_DEPTH]; + int seen_count = 0; + + /* Flatten the set */ + set_flatten_recursive(mrb, result_kh, self_kh, seen_ids, &seen_count); + + return result_set; +} + +/* + * call-seq: + * set.flatten! -> self or nil + * + * Replaces the contents of this set with a flattened version of itself. + * Returns self if flattened, nil if no changes were made. + */ +static mrb_value +set_flatten_bang(mrb_state *mrb, mrb_value self) +{ + mrb_check_frozen_value(mrb, self); + + khash_t(set) *self_kh = set_get_khash(mrb, self); + if (!self_kh || kh_size(self_kh) == 0) { + return mrb_nil_value(); /* No changes needed for empty set */ + } + + /* First, check if there are any nested sets */ + mrb_bool has_nested_sets = FALSE; + int ai = mrb_gc_arena_save(mrb); + KHASH_FOREACH(mrb, self_kh, k) { + mrb_value elem = kh_key(self_kh, k); + if (mrb_obj_is_kind_of(mrb, elem, mrb_class_get(mrb, "Set"))) { + has_nested_sets = TRUE; + break; + } + } + mrb_gc_arena_restore(mrb, ai); + + if (!has_nested_sets) { + return mrb_nil_value(); /* No nested sets, no changes needed */ + } + + /* Create a temporary hash table for the flattened result */ + khash_t(set) *new_kh = kh_init(set, mrb); + + /* Use a small array for tracking seen object IDs */ + mrb_int seen_ids[MAX_NESTED_DEPTH]; + int seen_count = 0; + + /* Flatten the set into the new hash table */ + set_flatten_recursive(mrb, new_kh, self_kh, seen_ids, &seen_count); + + /* Replace the old hash table with the new one */ + set_set_khash(mrb, self, new_kh); + + /* Clean up the old hash table */ + kh_destroy(set, mrb, self_kh); + + return self; +} + /* * call-seq: * set.delete_all(*objects) -> self @@ -1304,6 +1468,9 @@ mrb_mruby_set_gem_init(mrb_state *mrb) mrb_define_method(mrb, set, "disjoint?", set_disjoint_p, MRB_ARGS_REQ(1)); mrb_define_method(mrb, set, "<=>", set_cmp, MRB_ARGS_REQ(1)); + + mrb_define_method(mrb, set, "flatten", set_flatten, MRB_ARGS_NONE()); + mrb_define_method(mrb, set, "flatten!", set_flatten_bang, MRB_ARGS_NONE()); } void