reimplement String#chomp

2 files changed, 57 insertions, 1 deletions

diff --git a/src/string.c b/src/string.c
index c000a9d7c..1511c4f88 100644
--- a/src/string.c
+++ b/src/string.c
@@ -830,6 +830,62 @@ mrb_str_capitalize(mrb_state *mrb, mrb_value self)
 static mrb_value
 mrb_str_chomp_bang(mrb_state *mrb, mrb_value str)
 {
+  mrb_value rs;
+  mrb_int newline;
+  char *p, *pp;
+  long len, rslen;
+
+  len = RSTRING_LEN(str);
+  if (mrb_get_args(mrb, "|S", &rs) == 0) {
+    if (len == 0) return mrb_nil_value();
+  smart_chomp:
+    if (RSTRING_PTR(str)[len-1] == '\n') {
+      STR_DEC_LEN(str);
+      if (RSTRING_LEN(str) > 0 &&
+	  RSTRING_PTR(str)[RSTRING_LEN(str)-1] == '\r') {
+	STR_DEC_LEN(str);
+      }
+    }
+    else if (RSTRING_PTR(str)[len-1] == '\r') {
+      STR_DEC_LEN(str);
+    }
+    else {
+      return mrb_nil_value();
+    }
+    return str;
+  }
+
+  if (len == 0 || mrb_nil_p(rs)) return mrb_nil_value();
+  p = RSTRING_PTR(str);
+  rslen = RSTRING_LEN(rs);
+  if (rslen == 0) {
+    while (len>0 && p[len-1] == '\n') {
+      len--;
+      if (len>0 && p[len-1] == '\r')
+        len--;
+    }
+    if (len < RSTRING_LEN(str)) {
+      STR_SET_LEN(str, len);
+      p[len] = '\0';
+      return str;
+    }
+    return mrb_nil_value();
+  }
+  if (rslen > len) return mrb_nil_value();
+  newline = RSTRING_PTR(rs)[rslen-1];
+  if (rslen == 1 && newline == '\n')
+    newline = RSTRING_PTR(rs)[rslen-1];
+  if (rslen == 1 && newline == '\n')
+    goto smart_chomp;
+
+  pp = p + len - rslen;
+  if (p[len-1] == newline &&
+     (rslen <= 1 ||
+     memcmp(RSTRING_PTR(rs), pp, rslen) == 0)) {
+    STR_SET_LEN(str, len - rslen);
+    p[len] = '\0';
+    return str;
+  }
   return mrb_nil_value();
 }
 
diff --git a/test/t/string.rb b/test/t/string.rb
index 76df18aaf..f38790c17 100644
--- a/test/t/string.rb
+++ b/test/t/string.rb
@@ -85,7 +85,7 @@ assert('String#chomp', '15.2.10.5.9') do
     d == "abc\n" and e == 'abc' and f == "abc\n"
 end
 
-assert('String#chomp', '15.2.10.5.10') do
+assert('String#chomp!', '15.2.10.5.10') do
   a = 'abc'
   b = ''
   c = "abc\n"