Replace boundary on StringCharacterIndexPrimitiveNode

with faster specializations for boundary checks and single-byte-optimizable strings. (Mirroring StringByteIndexPrimitiveNode after improvements in oracle#2380).
Shopify · Jun 22, 2021 · 30e6cad · 30e6cad
1 parent b964e28
commit 30e6cad
Showing 1 changed file with 70 additions and 24 deletions.
diff --git a/src/main/java/org/truffleruby/core/string/StringNodes.java b/src/main/java/org/truffleruby/core/string/StringNodes.java
@@ -4399,48 +4399,94 @@ protected int notValidUtf8(Object string, int byteIndex,
     }
 
     @Primitive(name = "string_character_index", lowerFixnum = 2)
-    public abstract static class StringCharacterIndexPrimitiveNode extends PrimitiveArrayArgumentsNode {
+    @NodeChild(value = "string", type = RubyNode.class)
+    @NodeChild(value = "pattern", type = RubyNode.class)
+    @NodeChild(value = "offset", type = RubyNode.class)
+    public abstract static class StringCharacterIndexPrimitiveNode extends PrimitiveNode {
 
-        @TruffleBoundary
-        @Specialization
-        protected Object stringCharacterIndex(Object string, Object pattern, int offset,
-                @CachedLibrary(limit = "2") RubyStringLibrary stringLibrary,
-                @CachedLibrary(limit = "2") RubyStringLibrary patternLibrary,
-                @Cached RopeNodes.CalculateCharacterLengthNode calculateCharacterLengthNode) {
-            if (offset < 0) {
-                return nil;
-            }
+        @Child SingleByteOptimizableNode singleByteOptimizableNode = SingleByteOptimizableNode.create();
 
-            final Rope stringRope = stringLibrary.getRope(string);
-            final Rope patternRope = patternLibrary.getRope(pattern);
+        @CreateCast("string")
+        protected RubyNode coerceStringToRope(RubyNode string) {
+            return ToRopeNodeGen.create(string);
+        }
 
-            final int total = stringRope.byteLength();
-            int p = 0;
-            final int e = p + total;
+        @CreateCast("pattern")
+        protected RubyNode coercePatternToRope(RubyNode pattern) {
+            return ToRopeNodeGen.create(pattern);
+        }
+
+        @Specialization(guards = "offset < 0")
+        protected Object stringCharacterIndexNegativeOffset(Rope stringRope, Rope patternRope, int offset) {
+            return nil;
+        }
+
+        @Specialization(
+                guards = {
+                        "offset >= 0",
+                        "singleByteOptimizableNode.execute(stringRope)",
+                        "patternRope.byteLength() > stringRope.byteLength()" })
+        protected Object stringCharacterIndexPatternTooLarge(Rope stringRope, Rope patternRope, int offset) {
+            return nil;
+        }
+
+        @Specialization(
+                guards = {
+                        "offset >= 0",
+                        "singleByteOptimizableNode.execute(stringRope)",
+                        "patternRope.byteLength() <= stringRope.byteLength()" })
+        protected Object stringCharacterIndexSingleByteOptimizable(Rope stringRope, Rope patternRope, int offset,
+                @Cached RopeNodes.BytesNode stringBytesNode,
+                @Cached RopeNodes.BytesNode patternBytesNode,
+                @Cached LoopConditionProfile loopProfile,
+                @Cached("createCountingProfile()") ConditionProfile matchProfile) {
+
+            int p = offset;
+            final int e = stringRope.byteLength();
             final int pe = patternRope.byteLength();
             final int l = e - pe + 1;
 
-            final byte[] stringBytes = stringRope.getBytes();
-            final byte[] patternBytes = patternRope.getBytes();
+            final byte[] stringBytes = stringBytesNode.execute(stringRope);
+            final byte[] patternBytes = patternBytesNode.execute(patternRope);
 
-            if (stringRope.isSingleByteOptimizable()) {
-                for (p += offset; p < l; p++) {
-                    if (ArrayUtils.memcmp(stringBytes, p, patternBytes, 0, pe) == 0) {
+            try {
+                for (; loopProfile.profile(p < l); p++) {
+                    if (matchProfile.profile(ArrayUtils.memcmp(stringBytes, p, patternBytes, 0, pe) == 0)) {
                         return p;
                     }
                 }
-
-                return nil;
+            } finally {
+                LoopNode.reportLoopCount(this, p - offset);
             }
 
+            return nil;
+        }
+
+        @TruffleBoundary
+        @Specialization(
+                guards = {
+                        "offset >= 0",
+                        "!singleByteOptimizableNode.execute(stringRope)" })
+        protected Object stringCharacterIndex(Rope stringRope, Rope patternRope, int offset,
+                @Cached RopeNodes.CalculateCharacterLengthNode calculateCharacterLengthNode,
+                @Cached RopeNodes.BytesNode stringBytesNode,
+                @Cached RopeNodes.BytesNode patternBytesNode) {
+
+            int p = offset;
+            final int e = stringRope.byteLength();
+            final int pe = patternRope.byteLength();
+            final int l = e - pe + 1;
+
+            final byte[] stringBytes = stringBytesNode.execute(stringRope);
+            final byte[] patternBytes = patternBytesNode.execute(patternRope);
+
             final Encoding enc = stringRope.getEncoding();
             final CodeRange cr = stringRope.getCodeRange();
-            int index = 0;
             int c = 0;
+            int index = 0;
 
             while (p < e && index < offset) {
                 c = calculateCharacterLengthNode.characterLength(enc, cr, Bytes.fromRange(stringBytes, p, e));
-
                 if (StringSupport.MBCLEN_CHARFOUND_P(c)) {
                     p += c;
                     index++;