Merge pull request #3447 from dabroz/fix-ljust

String#ljust and String#rjust reimplementation (fix #3445)

3 files changed, 112 insertions, 46 deletions

diff --git a/mrbgems/mruby-string-ext/mrblib/string.rb b/mrbgems/mruby-string-ext/mrblib/string.rb
index 6e5f3c73d..610a462a7 100644
--- a/mrbgems/mruby-string-ext/mrblib/string.rb
+++ b/mrbgems/mruby-string-ext/mrblib/string.rb
@@ -263,52 +263,6 @@ class String
     self
   end
 
-  ##
-  #  call-seq:
-  #     str.ljust(integer, padstr=' ')   -> new_str
-  #
-  #  If <i>integer</i> is greater than the length of <i>str</i>, returns a new
-  #  <code>String</code> of length <i>integer</i> with <i>str</i> left justified
-  #  and padded with <i>padstr</i>; otherwise, returns <i>str</i>.
-  #
-  #     "hello".ljust(4)            #=> "hello"
-  #     "hello".ljust(20)           #=> "hello               "
-  #     "hello".ljust(20, '1234')   #=> "hello123412341234123"
-  def ljust(idx, padstr = ' ')
-    if idx <= self.size
-      return self
-    end
-    newstr = self.dup
-    newstr << padstr
-    while newstr.size <= idx
-      newstr << padstr
-    end
-    return newstr.slice(0,idx)
-  end
-
-  ##
-  #  call-seq:
-  #     str.rjust(integer, padstr=' ')   -> new_str
-  #
-  #  If <i>integer</i> is greater than the length of <i>str</i>, returns a new
-  #  <code>String</code> of length <i>integer</i> with <i>str</i> right justified
-  #  and padded with <i>padstr</i>; otherwise, returns <i>str</i>.
-  #
-  #     "hello".rjust(4)            #=> "hello"
-  #     "hello".rjust(20)           #=> "               hello"
-  #     "hello".rjust(20, '1234')   #=> "123412341234123hello"
-  def rjust(idx, padstr = ' ')
-    if idx <= self.size
-      return self
-    end
-      padsize = idx - self.size
-      newstr = padstr.dup
-      while newstr.size <= padsize
-        newstr << padstr
-      end
-    return newstr.slice(0,padsize) + self
-  end
-
   #     str.upto(other_str, exclusive=false) {|s| block }   -> str
   #     str.upto(other_str, exclusive=false)                -> an_enumerator
   #
diff --git a/mrbgems/mruby-string-ext/src/string.c b/mrbgems/mruby-string-ext/src/string.c
index 7e87b3db4..45d406f6f 100644
--- a/mrbgems/mruby-string-ext/src/string.c
+++ b/mrbgems/mruby-string-ext/src/string.c
@@ -5,6 +5,9 @@
 #include <mruby/string.h>
 #include <mruby/range.h>
 
+#define MAX(a,b) ((a)>(b) ? (a) : (b))
+#define MIN(a,b) ((a)<(b) ? (a) : (b))
+
 static mrb_value
 mrb_str_getbyte(mrb_state *mrb, mrb_value str)
 {
@@ -509,6 +512,93 @@ mrb_str_ord(mrb_state* mrb, mrb_value str)
 }
 #endif
 
+static mrb_value
+mrb_str_just(mrb_state* mrb, mrb_value str, mrb_bool right_just)
+{
+  mrb_value new_str;
+  mrb_int idx = 0, i = 0, bytes_to_copy = 0, start_pos = 0, final_pos = 0,
+    pad_str_length = 0;
+  mrb_int str_length = RSTRING_LEN(str);
+  const char *pad_str = NULL;
+  char *new_str_ptr = NULL;
+
+  mrb_get_args(mrb, "i|s!", &idx, &pad_str, &pad_str_length);
+
+  if (pad_str == NULL)
+  {
+    pad_str = " ";
+    pad_str_length = 1;
+  }
+
+  if (pad_str_length == 0)
+  {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "zero width padding");
+  }
+
+  if (idx <= str_length)
+  {
+    return str;
+  }
+
+  new_str = mrb_str_dup(mrb, str);
+  mrb_str_resize(mrb, new_str, idx);
+
+  new_str_ptr = RSTRING_PTR(new_str);
+
+  if (right_just)
+  {
+    memcpy(new_str_ptr + idx - str_length, RSTRING_PTR(str), str_length);
+  }
+
+  start_pos = right_just ? 0 : str_length;
+  final_pos = idx - (right_just ? str_length : 0);
+
+  for (i = start_pos; i < final_pos; i += pad_str_length)
+  {
+    bytes_to_copy = idx - i - (right_just ? str_length : 0);
+    bytes_to_copy = MIN(pad_str_length, bytes_to_copy);
+    memcpy(new_str_ptr + i, pad_str, bytes_to_copy);
+  }
+
+  return new_str;
+}
+
+/*
+ *  call-seq:
+ *     str.ljust(integer, padstr=' ')   -> new_str
+ *
+ *  If <i>integer</i> is greater than the length of <i>str</i>, returns a new
+ *  <code>String</code> of length <i>integer</i> with <i>str</i> left justified
+ *  and padded with <i>padstr</i>; otherwise, returns <i>str</i>.
+ *
+ *     "hello".ljust(4)            #=> "hello"
+ *     "hello".ljust(20)           #=> "hello               "
+ *     "hello".ljust(20, '1234')   #=> "hello123412341234123"
+ */
+static mrb_value
+mrb_str_ljust(mrb_state* mrb, mrb_value str)
+{
+  return mrb_str_just(mrb, str, FALSE);
+}
+
+/*
+ *  call-seq:
+ *     str.rjust(integer, padstr=' ')   -> new_str
+ *
+ *  If <i>integer</i> is greater than the length of <i>str</i>, returns a new
+ *  <code>String</code> of length <i>integer</i> with <i>str</i> right justified
+ *  and padded with <i>padstr</i>; otherwise, returns <i>str</i>.
+ *
+ *     "hello".rjust(4)            #=> "hello"
+ *     "hello".rjust(20)           #=> "               hello"
+ *     "hello".rjust(20, '1234')   #=> "123412341234123hello"
+ */
+static mrb_value
+mrb_str_rjust(mrb_state* mrb, mrb_value str)
+{
+  return mrb_str_just(mrb, str, TRUE);
+}
+
 void
 mrb_mruby_string_ext_gem_init(mrb_state* mrb)
 {
@@ -530,6 +620,8 @@ mrb_mruby_string_ext_gem_init(mrb_state* mrb)
   mrb_define_method(mrb, s, "lines",           mrb_str_lines,           MRB_ARGS_NONE());
   mrb_define_method(mrb, s, "succ",            mrb_str_succ,            MRB_ARGS_NONE());
   mrb_define_method(mrb, s, "succ!",           mrb_str_succ_bang,       MRB_ARGS_NONE());
+  mrb_define_method(mrb, s, "ljust",           mrb_str_ljust,           MRB_ARGS_REQ(1)|MRB_ARGS_OPT(1));
+  mrb_define_method(mrb, s, "rjust",           mrb_str_rjust,           MRB_ARGS_REQ(1)|MRB_ARGS_OPT(1));
   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());
diff --git a/mrbgems/mruby-string-ext/test/string.rb b/mrbgems/mruby-string-ext/test/string.rb
index 228a236af..2b2c02b8b 100644
--- a/mrbgems/mruby-string-ext/test/string.rb
+++ b/mrbgems/mruby-string-ext/test/string.rb
@@ -444,6 +444,26 @@ assert('String#rjust') do
   assert_equal "hello", "hello".rjust(-3)
 end
 
+assert('String#ljust should not change string') do
+  a = "hello"
+  a.ljust(20)
+  assert_equal "hello", a
+end
+
+assert('String#rjust should not change string') do
+  a = "hello"
+  a.rjust(20)
+  assert_equal "hello", a
+end
+
+assert('String#ljust should raise on zero width padding') do
+  assert_raise(ArgumentError) { "foo".ljust(10, '') }
+end
+
+assert('String#rjust should raise on zero width padding') do
+  assert_raise(ArgumentError) { "foo".rjust(10, '') }
+end
+
 assert('String#upto') do
   a     = "aa"
   start = "aa"