[2.4] Add support for chomp kwarg to IO/StringIO#gets.

jruby · Jul 20, 2017 · 5e68ee4 · 5e68ee4
1 parent 63ea2ac
commit 5e68ee4
Showing 3 changed files with 130 additions and 36 deletions.
diff --git a/core/src/main/java/org/jruby/RubyIO.java b/core/src/main/java/org/jruby/RubyIO.java
@@ -680,26 +680,27 @@ public IRubyObject reopen(ThreadContext context, IRubyObject[] args) {
     }
 
     public IRubyObject getline(ThreadContext context, IRubyObject separator) {
-        return getline(context, separator, -1, null);
+        return getlineInner(context, separator, -1, false, null);
     }
 
     /**
      * getline using logic of gets.  If limit is -1 then read unlimited amount.
      *
      */
     public IRubyObject getline(ThreadContext context, IRubyObject separator, long limit) {
-        return getline(context, separator, limit, null);
+        return getlineInner(context, separator, (int) limit, false, null);
     }
 
     private IRubyObject getline(ThreadContext context, IRubyObject separator, long limit, ByteListCache cache) {
-        return getlineInner(context, separator, (int)limit, cache);
+        return getlineInner(context, separator, (int) limit, false, cache);
     }
 
+
     /**
      * getline using logic of gets.  If limit is -1 then read unlimited amount.
      * mri: rb_io_getline_1 (mostly)
      */
-    private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _limit, ByteListCache cache) {
+    private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _limit, boolean chomp, ByteListCache cache) {
         Ruby runtime = context.runtime;
         IRubyObject str = context.nil;
         boolean noLimit = false;
@@ -714,6 +715,7 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
             if (rs.isNil() && _limit < 0) {
                 str = fptr.readAll(context, 0, context.nil);
                 if (((RubyString) str).size() == 0) return context.nil;
+                if (chomp) ((RubyString) str).chomp_bang(context, runtime.getGlobalVariables().getDefaultSeparator());
             } else if (_limit == 0) {
                 return RubyString.newEmptyString(runtime, fptr.readEncoding(runtime));
             } else if (
@@ -722,7 +724,7 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
                             && !fptr.needsReadConversion()
                             && (enc = fptr.readEncoding(runtime)).isAsciiCompatible()) {
                 fptr.NEED_NEWLINE_DECORATOR_ON_READ_CHECK();
-                return fptr.getlineFast(context, enc, this);
+                return fptr.getlineFast(context, enc, this, chomp);
             }
 
             // slow path logic
@@ -732,6 +734,7 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
             int rslen = 0;
             boolean rspara = false;
             int extraLimit = 16;
+            boolean chompCR = chomp;
 
             fptr.SET_BINARY_MODE();
             enc = getReadEncoding();
@@ -761,6 +764,7 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
                     rsptr = rsByteList.getBegin();
                 }
                 newline = rsptrBytes[rsptr + rslen - 1] & 0xFF;
+                chompCR = chomp && rslen == 1 && newline == '\n';
             }
 
             ByteList buf = cache != null ? cache.allocate(0) : new ByteList(0);
@@ -772,21 +776,31 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
                 while ((c = fptr.appendline(context, newline, strPtr, limit_p)) != OpenFile.EOF) {
                     int s, p, pp, e;
 
+                    byte[] strBytes = strPtr[0].getUnsafeBytes();
+                    int realSize = strPtr[0].getRealSize();
+                    int begin = strPtr[0].getBegin();
+
                     if (c == newline) {
-                        if (strPtr[0].getRealSize() < rslen) continue;
-                        s = strPtr[0].getBegin();
-                        e = s + strPtr[0].getRealSize();
+                        if (realSize < rslen) continue;
+                        s = begin;
+                        e = s + realSize;
                         p = e - rslen;
-                        pp = enc.leftAdjustCharHead(strPtr[0].getUnsafeBytes(), s, p, e);
+                        pp = enc.leftAdjustCharHead(strBytes, s, p, e);
                         if (pp != p) continue;
-                        if (ByteList.memcmp(strPtr[0].getUnsafeBytes(), p, rsptrBytes, rsptr, rslen) == 0) break;
+                        if (ByteList.memcmp(strBytes, p, rsptrBytes, rsptr, rslen) == 0) {
+                            if (chomp) {
+                                if (chompCR && p > s && strBytes[p-1] == '\r') --p;
+                                strPtr[0].length(p - s);
+                            }
+                            break;
+                        }
                     }
                     if (limit_p[0] == 0) {
-                        s = strPtr[0].getBegin();
-                        p = s + strPtr[0].getRealSize();
-                        pp = enc.leftAdjustCharHead(strPtr[0].getUnsafeBytes(), s, p - 1, p);
+                        s = begin;
+                        p = s + realSize;
+                        pp = enc.leftAdjustCharHead(strBytes, s, p - 1, p);
                         if (extraLimit != 0 &&
-                                StringSupport.MBCLEN_NEEDMORE_P(StringSupport.preciseLength(enc, strPtr[0].getUnsafeBytes(), pp, p))) {
+                                StringSupport.MBCLEN_NEEDMORE_P(StringSupport.preciseLength(enc, strBytes, pp, p))) {
                             limit_p[0] = 1;
                             extraLimit--;
                         } else {
@@ -827,6 +841,7 @@ private IRubyObject getlineInner(ThreadContext context, IRubyObject rs, int _lim
             if (locked) fptr.unlock();
         }
 
+
         return str;
     }
 
@@ -2230,7 +2245,7 @@ protected RubyIO flushRaw(ThreadContext context, boolean sync) {
     @JRubyMethod(name = "gets", writes = FrameField.LASTLINE)
     public IRubyObject gets(ThreadContext context) {
         IRubyObject separator = prepareGetsSeparator(context, null, null);
-        IRubyObject result = getline(context, separator);
+        IRubyObject result = getlineInner(context, separator, -1, false, null);
 
         if (!result.isNil()) context.setLastLine(result);
 
@@ -2240,10 +2255,22 @@ public IRubyObject gets(ThreadContext context) {
     // rb_io_gets_m
     @JRubyMethod(name = "gets", writes = FrameField.LASTLINE)
     public IRubyObject gets(ThreadContext context, IRubyObject arg) {
-        IRubyObject separator = prepareGetsSeparator(context, arg, null);
-        long limit = prepareGetsLimit(context, arg, null);
+        boolean chomp = false;
+        IRubyObject rs = null;
+        IRubyObject opt = ArgsUtil.getOptionsArg(context.runtime, arg);
+        long limit = -1;
+        if (opt.isNil()) {
+            rs = prepareGetsSeparator(context, arg, null);
+            limit = prepareGetsLimit(context, arg, null);
+        } else {
+            IRubyObject chompKwarg = ArgsUtil.extractKeywordArg(context, "chomp", opt);
+            if (chompKwarg != null) {
+                chomp = chompKwarg.isTrue();
+            }
+            rs = prepareGetsSeparator(context, null, null);
+        }
 
-        IRubyObject result = getline(context, separator, limit);
+        IRubyObject result = getlineInner(context, rs, (int) limit, chomp, null);
 
         if (!result.isNil()) context.setLastLine(result);
 
@@ -2253,9 +2280,38 @@ public IRubyObject gets(ThreadContext context, IRubyObject arg) {
     // rb_io_gets_m
     @JRubyMethod(name = "gets", writes = FrameField.LASTLINE)
     public IRubyObject gets(ThreadContext context, IRubyObject rs, IRubyObject limit_arg) {
+        boolean chomp = false;
+        IRubyObject opt = ArgsUtil.getOptionsArg(context.runtime, limit_arg);
+        long limit;
+        if (opt.isNil()) {
+            rs = prepareGetsSeparator(context, rs, limit_arg);
+            limit = prepareGetsLimit(context, rs, limit_arg);
+        } else {
+            IRubyObject chompKwarg = ArgsUtil.extractKeywordArg(context, "chomp", opt);
+            if (chompKwarg != null) {
+                chomp = chompKwarg.isTrue();
+            }
+            rs = prepareGetsSeparator(context, rs, null);
+            limit = prepareGetsLimit(context, rs, null);
+        }
+        IRubyObject result = getlineInner(context, rs, (int) limit, chomp, null);
+
+        if (!result.isNil()) context.setLastLine(result);
+
+        return result;
+    }
+
+    @JRubyMethod(name = "gets", writes = FrameField.LASTLINE)
+    public IRubyObject gets(ThreadContext context, IRubyObject rs, IRubyObject limit_arg, IRubyObject opt) {
+        boolean chomp = false;
+        long limit = -1;
+        IRubyObject chompKwarg = ArgsUtil.extractKeywordArg(context, "chomp", opt);
+        if (chompKwarg != null) {
+            chomp = chompKwarg.isTrue();
+        }
         rs = prepareGetsSeparator(context, rs, limit_arg);
-        long limit = prepareGetsLimit(context, rs, limit_arg);
-        IRubyObject result = getline(context, rs, limit);
+        limit = prepareGetsLimit(context, rs, limit_arg);
+        IRubyObject result = getlineInner(context, rs, (int) limit, chomp, null);
 
         if (!result.isNil()) context.setLastLine(result);
 

diff --git a/core/src/main/java/org/jruby/ext/stringio/StringIO.java b/core/src/main/java/org/jruby/ext/stringio/StringIO.java
@@ -577,8 +577,19 @@ private IRubyObject getline(ThreadContext context, IRubyObject[] args) {
 
         IRubyObject str = context.nil;
         int n, limit = -1;
+        boolean chomp = false;
 
-        switch (args.length) {
+        int argc = args.length;
+
+        IRubyObject opt = ArgsUtil.getOptionsArg(runtime, args);
+        if (!opt.isNil()) {
+            argc--;
+            IRubyObject chompKwarg = ArgsUtil.extractKeywordArg(context, "chomp", opt);
+            if (chompKwarg != null) {
+                chomp = chompKwarg.isTrue();
+            }
+        }
+        switch (argc) {
             case 0:
                 str = runtime.getGlobalVariables().get("$/");
                 break;
@@ -621,39 +632,52 @@ private IRubyObject getline(ThreadContext context, IRubyObject[] args) {
             int s = begin + ptr.pos;
             int e = begin + string.getRealSize();
             int p;
+            int w = 0;
 
             if (limit > 0 && s + limit < e) {
                 e = ptr.enc.rightAdjustCharHead(stringBytes, s, s + limit, e);
             }
             if (str.isNil()) {
-                str = strioSubstr(runtime, ptr.pos, e - s);
+                if (chomp) {
+                    w = chompNewlineWidth(stringBytes, s, e);
+                }
+                str = strioSubstr(runtime, ptr.pos, e - s - w);
             } else if ((n = ((RubyString) str).size()) == 0) {
                 // this is not an exact port; the original confused me
                 p = s;
                 // remove leading \n
-                while (stringBytes[p] == '\n') {
+                while (stringBytes[p + (((p + 1 < e) && (stringBytes[p] == '\r') && false)?1:0)] == '\n') {
+                    p += (stringBytes[p] == '\r')?1:0;
                     if (++p == e) {
                         return context.nil;
                     }
                 }
                 s = p;
                 // find next \n or end; if followed by \n, include it too
-                p = StringSupport.memchr(stringBytes, p, '\n', e - p);
-                if (p != -1) {
-                    if (++p < e && stringBytes[p] == '\n') {
+                while ((p = StringSupport.memchr(stringBytes, p, '\n', e - p)) != -1 && (p != e)) {
+                    if (stringBytes[++p] == '\n') {
                         e = p + 1;
-                    } else {
-                        e = p;
+                        w = (chomp ? 1 : 0);
+                        break;
                     }
+            	    else if (stringBytes[p] == '\r' && p < e && stringBytes[p + 1] == '\n') {
+                        e = p + 2;
+                        w = (chomp ? 2 : 0);
+                        break;
+                    }
+                }
+                if (w == 0 && chomp) {
+                    w = chompNewlineWidth(stringBytes, s, e);
                 }
-                str = strioSubstr(runtime, s - begin, e - s);
+                str = strioSubstr(runtime, s - begin, e - s - w);
             } else if (n == 1) {
                 RubyString strStr = (RubyString) str;
                 ByteList strByteList = strStr.getByteList();
                 if ((p = StringSupport.memchr(stringBytes, s, strByteList.get(0), e - s)) != -1) {
                     e = p + 1;
+                    w = (chomp ? ((p > s && stringBytes[p-1] == '\r')?1:0) + 1 : 0);
                 }
-                str = strioSubstr(runtime, ptr.pos, e - s);
+                str = strioSubstr(runtime, ptr.pos, e - s - w);
             } else {
                 if (n < e - s) {
                     RubyString strStr = (RubyString) str;
@@ -668,7 +692,7 @@ private IRubyObject getline(ThreadContext context, IRubyObject[] args) {
                         e = s + pos + n;
                     }
                 }
-                str = strioSubstr(runtime, ptr.pos, e - s);
+                str = strioSubstr(runtime, ptr.pos, e - s - w);
             }
             ptr.pos = e - begin;
             ptr.lineno++;
@@ -677,6 +701,14 @@ private IRubyObject getline(ThreadContext context, IRubyObject[] args) {
         return str;
     }
 
+    private static int chompNewlineWidth(byte[] bytes, int s, int e) {
+        if (e > s && bytes[--e] == '\n') {
+            if (e > s && bytes[--e] == '\r') return 2;
+            return 1;
+        }
+        return 0;
+    }
+
     @JRubyMethod(name = {"length", "size"})
     public IRubyObject length() {
         checkInitialized();

diff --git a/core/src/main/java/org/jruby/util/io/OpenFile.java b/core/src/main/java/org/jruby/util/io/OpenFile.java
@@ -1547,7 +1547,7 @@ else if (cbuf.capa / 2 < cbuf.off) {
     }
 
     // rb_io_getline_fast
-    public IRubyObject getlineFast(ThreadContext context, Encoding enc, RubyIO io) {
+    public IRubyObject getlineFast(ThreadContext context, Encoding enc, RubyIO io, boolean chomp) {
         Ruby runtime = context.runtime;
         IRubyObject str = null;
         ByteList strByteList;
@@ -1564,22 +1564,28 @@ public IRubyObject getlineFast(ThreadContext context, Encoding enc, RubyIO io) {
                     byte[] pBytes = READ_DATA_PENDING_PTR();
                     int p = READ_DATA_PENDING_OFF();
                     int e;
+                    int chomplen = 0;
 
                     e = memchr(pBytes, p, '\n', pending);
                     if (e != -1) {
                         pending = (int) (e - p + 1);
+                        if (chomp) {
+                            chomplen = ((pending > 1 &&  pBytes[e - 1] == '\r')?1:0) + 1;
+                        }
                     }
                     if (str == null) {
-                        str = RubyString.newString(runtime, pBytes, p, pending);
+                        str = RubyString.newString(runtime, pBytes, p, pending - chomplen);
                         strByteList = ((RubyString) str).getByteList();
                         rbuf.off += pending;
                         rbuf.len -= pending;
                     } else {
-                        ((RubyString) str).resize(len + pending);
+                        ((RubyString) str).resize(len + pending - chomplen);
                         strByteList = ((RubyString) str).getByteList();
-                        readBufferedData(strByteList.unsafeBytes(), strByteList.begin() + len, pending);
+                        readBufferedData(strByteList.unsafeBytes(), strByteList.begin() + len, pending - chomplen);
+                        rbuf.off += chomplen;
+                        rbuf.len -= chomplen;
                     }
-                    len += pending;
+                    len += pending - chomplen;
                     if (cr != StringSupport.CR_BROKEN)
                         pos += StringSupport.codeRangeScanRestartable(enc, strByteList.unsafeBytes(), strByteList.begin() + pos, strByteList.begin() + len, cr);
                     if (e != -1) break;