reimplement Array#delete with a helper method

mruby · Apr 13, 2024 · bb78c2c · bb78c2c
1 parent 312d49c
commit bb78c2c
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 24 deletions.
diff --git a/mrblib/array.rb b/mrblib/array.rb
@@ -191,31 +191,9 @@ def <=>(other)
   #
   # Delete element with index +key+
   def delete(key, &block)
-    ret = key
-    i = 0
-    j = 0
     len = self.length
-    while i < len
-      elem = self[i]
-
-      if key == elem
-        ret = elem
-        i += 1
-        next
-      end
-
-      self[j] = elem if i != j
-
-      i += 1
-      j += 1
-    end
-
-    if i == j
-      return block.call if block
-      return nil
-    end
-
-    self.replace(self[0...j])
+    ret = self.__delete(key)
+    return block.call() if len == self.length
 
     ret
   end

diff --git a/src/array.c b/src/array.c
@@ -1397,6 +1397,48 @@ mrb_ary_svalue(mrb_state *mrb, mrb_value ary)
   }
 }
 
+static mrb_value
+mrb_ary_delete(mrb_state *mrb, mrb_value self)
+{
+  struct RArray *ary = RARRAY(self);
+  mrb_value *val_ptr = ARY_PTR(ary);
+  size_t len = ARY_LEN(ary);
+  mrb_bool modified = FALSE;
+
+  mrb_value obj = mrb_get_arg1(mrb);
+  mrb_value ret = obj;
+
+  size_t i = 0;
+  size_t j = 0;
+  for (; i < len; ++i) {
+    mrb_value elem = val_ptr[i];
+
+    if (mrb_equal(mrb, elem, obj)) {
+      ret = elem;
+      continue;
+    }
+
+    if (i != j) {
+      if (!modified) {
+        ary_modify(mrb, ary);
+        val_ptr = ARY_PTR(ary);
+        modified = TRUE;
+      }
+      val_ptr[j] = elem;
+    }
+
+    ++j;
+  }
+
+  if (i == j) {
+    return mrb_nil_value();
+  }
+
+  ARY_SET_LEN(ary, j);
+
+  return ret;
+}
+
 void
 mrb_init_array(mrb_state *mrb)
 {
@@ -1438,5 +1480,6 @@ mrb_init_array(mrb_state *mrb)
   mrb_define_method_id(mrb, a, MRB_SYM(__ary_eq),        mrb_ary_eq,           MRB_ARGS_REQ(1));
   mrb_define_method_id(mrb, a, MRB_SYM(__ary_cmp),       mrb_ary_cmp,          MRB_ARGS_REQ(1));
   mrb_define_method_id(mrb, a, MRB_SYM(__ary_index),     mrb_ary_index_m,      MRB_ARGS_REQ(1));   /* kept for mruby-array-ext */
+  mrb_define_method_id(mrb, a, MRB_SYM(__delete),        mrb_ary_delete,       MRB_ARGS_REQ(1));
   mrb_define_method_id(mrb, a, MRB_SYM(__svalue),        mrb_ary_svalue,       MRB_ARGS_NONE());
 }