Reimplement String#upto

mrbgems/mruby-string-ext/mrbgem.rake

@@ -2,4 +2,5 @@ MRuby::Gem::Specification.new('mruby-string-ext') do |spec|
   spec.license = 'MIT'
   spec.author  = 'mruby developers'
   spec.summary = 'String class extension'
+  spec.add_test_dependency 'mruby-enumerator', core: 'mruby-enumerator'
 end

mrbgems/mruby-string-ext/mrblib/string.rb

@@ -305,50 +305,6 @@ def rjust(idx, padstr = ' ')
     padding + self
   end
 
-  #     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"]
-  #
-  def upto(other_str, excl=false, &block)
-    return to_enum :upto, other_str, excl unless block
-
-    str = self
-    n = self.<=>other_str
-    return self if n > 0 || (self == other_str && excl)
-    while true
-      block.call(str)
-      return self if !excl && str == other_str
-      str = str.succ
-      return self if excl && str == other_str
-      return self if str.size > other_str.size
-    end
-  end
-
   def chars(&block)
     if block_given?
       self.split('').each do |i|

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());
 }

mrbgems/mruby-string-ext/test/string.rb

@@ -495,6 +495,17 @@ def o.to_str
 end
 
 assert('String#upto') do
+  assert_equal %w(a8 a9 b0 b1 b2 b3 b4 b5 b6), "a8".upto("b6").to_a
+  assert_equal ["9", "10", "11"], "9".upto("11").to_a
+  assert_equal [], "25".upto("5").to_a
+  assert_equal ["07", "08", "09", "10", "11"], "07".upto("11").to_a
+
+if UTF8STRING
+  assert_equal ["あ", "ぃ", "い", "ぅ", "う", "ぇ", "え", "ぉ", "お"], "あ".upto("お").to_a
+end
+
+  assert_equal ["9", ":", ";", "<", "=", ">", "?", "@", "A"], "9".upto("A").to_a
+
   a     = "aa"
   start = "aa"
   count = 0
@@ -554,6 +565,8 @@ def o.to_str
     count += 1
   })
   assert_equal(2, count)
+
+  assert_raise(TypeError) { "a".upto(:c) {} }
 end
 
 assert('String#ord') do