Reimplement String#upto

mrbgems/mruby-string-ext/src/string.c

@@ -521,6 +521,119 @@ mrb_str_ord(mrb_state* mrb, mrb_value str)
 }
 #endif
 
+static mrb_bool
+all_digits_p(const char *s, mrb_int len)
+{
+  while (len-- > 0) {
+    if (!ISDIGIT(*s)) return FALSE;
+    s++;
+  }
+  return TRUE;
+}
+
+/*
+ *  call-seq:
+ *     str.upto(other_str, exclusive=false) {|s| block }   -> str
+ *     str.upto(other_str, exclusive=false)                -> an_enumerator
+ *
+ *  Iterates through successive values, starting at <i>str</i> and
+ *  ending at <i>other_str</i> inclusive, passing each value in turn to
+ *  the block. The <code>String#succ</code> method is used to generate
+ *  each value.  If optional second argument exclusive is omitted or is false,
+ *  the last value will be included; otherwise it will be excluded.
+ *
+ *  If no block is given, an enumerator is returned instead.
+ *
+ *     "a8".upto("b6") {|s| print s, ' ' }
+ *     for s in "a8".."b6"
+ *       print s, ' '
+ *     end
+ *
+ *  <em>produces:</em>
+ *
+ *     a8 a9 b0 b1 b2 b3 b4 b5 b6
+ *     a8 a9 b0 b1 b2 b3 b4 b5 b6
+ *
+ *  If <i>str</i> and <i>other_str</i> contains only ascii numeric characters,
+ *  both are recognized as decimal numbers. In addition, the width of
+ *  string (e.g. leading zeros) is handled appropriately.
+ *
+ *     "9".upto("11").to_a   #=> ["9", "10", "11"]
+ *     "25".upto("5").to_a   #=> []
+ *     "07".upto("11").to_a  #=> ["07", "08", "09", "10", "11"]
+ */
+static mrb_value
+mrb_str_upto(mrb_state *mrb, mrb_value beg)
+{
+  mrb_value end;
+  mrb_value exclusive = mrb_false_value();
+  mrb_value block = mrb_nil_value();
+  mrb_value current, after_end;
+  mrb_int n;
+  mrb_bool excl;
+
+  mrb_get_args(mrb, "o|o&", &end, &exclusive, &block);
+
+  if (mrb_nil_p(block)) {
+    return mrb_funcall(mrb, beg, "to_enum", 3, mrb_symbol_value(mrb_intern_lit(mrb, "upto")), end, exclusive);
+  }
+  end = mrb_string_type(mrb, end);
+  excl = mrb_test(exclusive);
+
+  /* single character */
+  if (RSTRING_LEN(beg) == 1 && RSTRING_LEN(end) == 1 &&
+  ISASCII(RSTRING_PTR(beg)[0]) && ISASCII(RSTRING_PTR(end)[0])) {
+    char c = RSTRING_PTR(beg)[0];
+    char e = RSTRING_PTR(end)[0];
+
+    if (c > e || (excl && c == e)) return beg;
+    for (;;) {
+      mrb_yield(mrb, block, mrb_str_new(mrb, &c, 1));
+      if (!excl && c == e) break;
+      c++;
+      if (excl && c == e) break;
+    }
+    return beg;
+  }
+  /* both edges are all digits */
+  if (ISDIGIT(RSTRING_PTR(beg)[0]) && ISDIGIT(RSTRING_PTR(end)[0]) &&
+  all_digits_p(RSTRING_PTR(beg), RSTRING_LEN(beg)) &&
+  all_digits_p(RSTRING_PTR(end), RSTRING_LEN(end))) {
+    mrb_int min_width = RSTRING_LEN(beg);
+    mrb_int max_width = RSTRING_LEN(end);
+    mrb_int bi = mrb_int(mrb, mrb_str_to_inum(mrb, beg, 10, FALSE));
+    mrb_int ei = mrb_int(mrb, mrb_str_to_inum(mrb, end, 10, FALSE));
+    char buf[max_width+1];
+
+    while (bi <= ei) {
+      if (excl && bi == ei) break;
+      snprintf(buf, sizeof(buf), "%.*d", min_width, bi);
+      mrb_yield(mrb, block, mrb_str_new(mrb, buf, strlen(buf)));
+      bi++;
+    }
+    return beg;
+  }
+  /* normal case */
+  n = mrb_int(mrb, mrb_funcall(mrb, beg, "<=>", 1, end));
+  if (n > 0 || (excl && n == 0)) return beg;
+
+  after_end = mrb_funcall(mrb, end, "succ", 0);
+  current = mrb_str_dup(mrb, beg);
+  while (!mrb_str_equal(mrb, current, after_end)) {
+    mrb_value next = mrb_nil_value();
+    if (excl || !mrb_str_equal(mrb, current, end))
+      next = mrb_funcall(mrb, current, "succ", 0);
+    mrb_yield(mrb, block, current);
+    if (mrb_nil_p(next)) break;
+    current = mrb_str_to_str(mrb, next);
+    if (excl && mrb_str_equal(mrb, current, end)) break;
+    if (RSTRING_LEN(current) > RSTRING_LEN(end) || RSTRING_LEN(current) == 0)
+      break;
+  }
+
+  return beg;
+}
+
 void
 mrb_mruby_string_ext_gem_init(mrb_state* mrb)
 {
@@ -545,6 +658,7 @@ mrb_mruby_string_ext_gem_init(mrb_state* mrb)
   mrb_alias_method(mrb, s, mrb_intern_lit(mrb, "next"), mrb_intern_lit(mrb, "succ"));
   mrb_alias_method(mrb, s, mrb_intern_lit(mrb, "next!"), mrb_intern_lit(mrb, "succ!"));
   mrb_define_method(mrb, s, "ord", mrb_str_ord, MRB_ARGS_NONE());
+  mrb_define_method(mrb, s, "upto", mrb_str_upto, MRB_ARGS_ANY());
 
   mrb_define_method(mrb, mrb->fixnum_class, "chr", mrb_fixnum_chr, MRB_ARGS_NONE());
 }