diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb index 335f1e35b..f5f342421 100644 --- a/mrbgems/mruby-array-ext/mrblib/array.rb +++ b/mrbgems/mruby-array-ext/mrblib/array.rb @@ -99,18 +99,18 @@ class Array def fetch(n, ifnone=NONE, &block) #warn "block supersedes default value argument" if !n.nil? && ifnone != NONE && block - idx = n.__to_int - if idx < 0 - idx += size - end - if idx < 0 || size <= idx - return block.call(n) if block - if NONE.equal?(ifnone) - raise IndexError, "index #{n} outside of array bounds: #{-size}...#{size}" + if block + # Block case: use shared index helper + Ruby block handling + normalized_index = __normalize_index(n) + if normalized_index + self[normalized_index] + else + block.call(n) end - return ifnone + else + # Fast C implementation for non-block cases + __fetch(n, ifnone, NONE) end - self[idx] end ## diff --git a/mrbgems/mruby-array-ext/src/array.c b/mrbgems/mruby-array-ext/src/array.c index a05c3268e..33a07a21a 100644 --- a/mrbgems/mruby-array-ext/src/array.c +++ b/mrbgems/mruby-array-ext/src/array.c @@ -1103,6 +1103,75 @@ ary_flatten(mrb_state *mrb, mrb_value self) return flatten_internal(mrb, self, level, &modified); } +/* + * Shared helper for index normalization and bounds checking. + * Returns normalized index if in bounds, nil if out of bounds. + */ + +static mrb_value +ary_normalize_index(mrb_state *mrb, mrb_value self) +{ + mrb_value index_val; + mrb_get_args(mrb, "o", &index_val); + + mrb_int index = mrb_as_int(mrb, index_val); + struct RArray *ary = mrb_ary_ptr(self); + mrb_int len = ARY_LEN(ary); + + // Handle negative indices + if (index < 0) { + index += len; + } + + // Check bounds + if (index >= 0 && index < len) { + return mrb_fixnum_value(index); + } + else { + return mrb_nil_value(); + } +} + +/* + * Fast C implementation for Array#fetch without blocks. + * Returns the element at index, or default if out of bounds. + * Raises IndexError if out of bounds and default equals none. + */ + +static mrb_value +ary_fetch(mrb_state *mrb, mrb_value self) +{ + mrb_value index_val, default_val, none; + mrb_get_args(mrb, "ooo", &index_val, &default_val, &none); + + // Convert index to integer + mrb_int index = mrb_as_int(mrb, index_val); + mrb_int original_index = index; // Keep original for error message + + struct RArray *ary = mrb_ary_ptr(self); + mrb_int len = ARY_LEN(ary); + + // Handle negative indices + if (index < 0) { + index += len; + } + + // Check bounds + if (index < 0 || index >= len) { + // Check if default is the NONE sentinel (means no default provided) + if (mrb_obj_equal(mrb, default_val, none)) { + // No default provided - raise IndexError + mrb_raisef(mrb, E_INDEX_ERROR, + "index %i outside of array bounds: %i...%i", + original_index, -len, len); + } + return default_val; + } + + // Return element at index + return ARY_PTR(ary)[index]; +} + /* * call-seq: * ary.flatten! -> ary or nil @@ -1219,6 +1288,8 @@ mrb_mruby_array_ext_gem_init(mrb_state* mrb) mrb_define_method_id(mrb, a, MRB_SYM_B(__uniq), ary_uniq_bang, MRB_ARGS_NONE()); mrb_define_method_id(mrb, a, MRB_SYM(flatten), ary_flatten, MRB_ARGS_OPT(1)); mrb_define_method_id(mrb, a, MRB_SYM_B(flatten), ary_flatten_bang, MRB_ARGS_OPT(1)); + mrb_define_method_id(mrb, a, MRB_SYM(__normalize_index), ary_normalize_index, MRB_ARGS_REQ(1)); + mrb_define_method_id(mrb, a, MRB_SYM(__fetch), ary_fetch, MRB_ARGS_REQ(3)); mrb_define_method_id(mrb, a, MRB_SYM(insert), ary_insert, MRB_ARGS_ARG(1, -1)); } diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb index a2185a1d5..b6d6f1c26 100644 --- a/mrbgems/mruby-array-ext/test/array.rb +++ b/mrbgems/mruby-array-ext/test/array.rb @@ -358,6 +358,19 @@ assert("Array#fetch") do a.fetch(100) { |i| ret = i } assert_equal 100, ret assert_raise(IndexError) { a.fetch(100) } + + # Additional edge cases + assert_equal "default", [].fetch(0, "default") + assert_equal "missing 5", ["a"].fetch(5) { |i| "missing #{i}" } + assert_equal "from block", ["a"].fetch(5, "default") { "from block" } + + # Error message format + begin + ["a", "b"].fetch(5) + assert_false true + rescue IndexError => e + assert_true e.message.include?("index 5 outside of array bounds: -2...2") + end end assert("Array#fetch_values") do