release memory for children ropes when a concat rope becomes flattened

Author: Nicolas Laurent

src/main/java/org/truffleruby/core/rope/ConcatRope.java

@@ -9,68 +9,116 @@
  */
 package org.truffleruby.core.rope;
 
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.CompilerDirectives.ValueType;
 import org.jcodings.Encoding;
 import org.jcodings.specific.ASCIIEncoding;
 
 public class ConcatRope extends ManagedRope {
 
-    private final ManagedRope left;
-    private final ManagedRope right;
+    @ValueType
+    public static class ConcatChildren {
+        public final ManagedRope left, right;
+
+        public ConcatChildren(ManagedRope left, ManagedRope right) {
+            this.left = left;
+            this.right = right;
+        }
+    }
+
+    private ManagedRope left;
+    private ManagedRope right;
 
     public ConcatRope(
             ManagedRope left,
             ManagedRope right,
             Encoding encoding,
             CodeRange codeRange) {
-        this(left, right, encoding, codeRange, left.characterLength() + right.characterLength(), null);
+        this(
+                left,
+                right,
+                encoding,
+                codeRange,
+                left.byteLength() + right.byteLength(),
+                left.characterLength() + right.characterLength(),
+                null);
     }
 
     private ConcatRope(
             ManagedRope left,
             ManagedRope right,
             Encoding encoding,
             CodeRange codeRange,
+            int byteLength,
             int characterLength,
             byte[] bytes) {
-        super(
-                encoding,
-                codeRange,
-                left.byteLength() + right.byteLength(),
-                characterLength,
-                bytes);
+        super(encoding, codeRange, byteLength, characterLength, bytes);
         this.left = left;
         this.right = right;
     }
 
     @Override
     Rope withEncoding7bit(Encoding newEncoding) {
         assert getCodeRange() == CodeRange.CR_7BIT;
-        return new ConcatRope(
-                getLeft(),
-                getRight(),
-                newEncoding,
-                CodeRange.CR_7BIT,
-                characterLength(),
-                getRawBytes());
+        return withEncoding(newEncoding, CodeRange.CR_7BIT, characterLength());
     }
 
     @Override
     Rope withBinaryEncoding() {
         assert getCodeRange() == CodeRange.CR_VALID;
-        return new ConcatRope(
-                getLeft(),
-                getRight(),
-                ASCIIEncoding.INSTANCE,
-                CodeRange.CR_VALID,
-                byteLength(),
-                getRawBytes());
+        return withEncoding(ASCIIEncoding.INSTANCE, CodeRange.CR_VALID, byteLength());
+    }
+
+    private ConcatRope withEncoding(Encoding encoding, CodeRange codeRange, int characterLength) {
+        final ConcatChildren children = getChildrenOrNull();
+        final ManagedRope left, right;
+        final byte[] bytes;
+        if (children != null) {
+            left = children.left;
+            right = children.right;
+            bytes = null;
+        } else {
+            left = null;
+            right = null;
+            bytes = getRawBytes();
+        }
+        return new ConcatRope(left, right, encoding, codeRange, byteLength(), characterLength, bytes);
+    }
+
+    @Override
+    protected byte[] getBytesSlow() {
+        final byte[] out = RopeOperations.flattenBytes(this);
+        this.left = null;
+        this.right = null;
+        return out;
     }
 
-    public ManagedRope getLeft() {
-        return left;
+    public ConcatChildren getChildrenOrNull() {
+        final ManagedRope left = this.left;
+        final ManagedRope right = this.right;
+
+        if (left != null && right != null) {
+            return new ConcatChildren(left, right);
+        } else if (this.bytes != null) {
+            return null;
+        } else {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            assert this.bytes != null;
+            return null;
+        }
     }
 
-    public ManagedRope getRight() {
-        return right;
+    public byte[] getBytesOrNull() {
+        final byte[] bytes = this.bytes;
+
+        if (bytes != null) {
+            return bytes;
+        } else if (this.left != null && this.right != null) {
+            return null;
+        } else {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            assert this.bytes != null;
+            return this.bytes;
+        }
     }
 }

src/main/java/org/truffleruby/core/rope/RopeNodes.java

@@ -21,13 +21,15 @@
 
 import java.util.Arrays;
 
+import com.oracle.truffle.api.dsl.Bind;
 import org.jcodings.Encoding;
 import org.jcodings.specific.ASCIIEncoding;
 import org.jcodings.specific.USASCIIEncoding;
 import org.jcodings.specific.UTF8Encoding;
 import org.truffleruby.SuppressFBWarnings;
 import org.truffleruby.core.array.ArrayUtils;
 import org.truffleruby.core.encoding.EncodingNodes;
+import org.truffleruby.core.rope.ConcatRope.ConcatChildren;
 import org.truffleruby.core.rope.RopeNodesFactory.AreComparableRopesNodeGen;
 import org.truffleruby.core.rope.RopeNodesFactory.CompareRopesNodeGen;
 import org.truffleruby.core.rope.RopeNodesFactory.SetByteNodeGen;
@@ -804,26 +806,46 @@ protected Object debugPrintSubstringRope(SubstringRope rope, int currentLevel, b
         }
 
         @TruffleBoundary
-        @Specialization
-        protected Object debugPrintConcatRope(ConcatRope rope, int currentLevel, boolean printString) {
+        @Specialization(guards = "bytes != null")
+        protected Object debugPrintConcatRope(ConcatRope rope, int currentLevel, boolean printString,
+                @Bind("rope.getBytesOrNull()") byte[] bytes) {
+            printPreamble(currentLevel);
+
+            System.err
+                    .println(StringUtils.format(
+                            "%s (%s; BN: %b; BL: %d; CL: %d; CR: %s; E: %s)",
+                            printString ? rope.toString() : "<skipped>",
+                            rope.getClass().getSimpleName(),
+                            false,
+                            rope.byteLength(),
+                            rope.characterLength(),
+                            rope.getCodeRange(),
+                            rope.getEncoding()));
+
+            return nil;
+        }
+
+        @TruffleBoundary
+        @Specialization(guards = "children != null")
+        protected Object debugPrintConcatRope(ConcatRope rope, int currentLevel, boolean printString,
+                @Bind("rope.getChildrenOrNull()") ConcatChildren children) {
             printPreamble(currentLevel);
 
-            // Converting a rope to a java.lang.String may populate the byte[], so we need to query for the array status beforehand.
-            final boolean bytesAreNull = rope.getRawBytes() == null;
 
             System.err
                     .println(StringUtils.format(
                             "%s (%s; BN: %b; BL: %d; CL: %d; CR: %s; E: %s)",
                             printString ? rope.toString() : "<skipped>",
                             rope.getClass().getSimpleName(),
-                            bytesAreNull,
+                            // Note: converting a rope to a java.lang.String may populate the byte[].
+                            true, // are bytes null?
                             rope.byteLength(),
                             rope.characterLength(),
                             rope.getCodeRange(),
                             rope.getEncoding()));
 
-            executeDebugPrint(rope.getLeft(), currentLevel + 1, printString);
-            executeDebugPrint(rope.getRight(), currentLevel + 1, printString);
+            executeDebugPrint(children.left, currentLevel + 1, printString);
+            executeDebugPrint(children.right, currentLevel + 1, printString);
 
             return nil;
         }
@@ -996,28 +1018,36 @@ protected int getByteRepeatingRope(RepeatingRope rope, int index,
             return rope.getChild().getRawBytes()[index % rope.getChild().byteLength()] & 0xff;
         }
 
-        @Specialization(guards = "rope.getRawBytes() == null")
+        @Specialization(guards = "children != null")
         protected int getByteConcatRope(ConcatRope rope, int index,
+                @Bind("rope.getChildrenOrNull()") ConcatChildren children,
                 @Cached ConditionProfile chooseLeftChildProfile,
                 @Cached ConditionProfile leftChildRawBytesNullProfile,
                 @Cached ConditionProfile rightChildRawBytesNullProfile,
                 @Cached ByteSlowNode byteSlowLeft,
                 @Cached ByteSlowNode byteSlowRight) {
-            if (chooseLeftChildProfile.profile(index < rope.getLeft().byteLength())) {
-                if (leftChildRawBytesNullProfile.profile(rope.getLeft().getRawBytes() == null)) {
-                    return byteSlowLeft.execute(rope.getLeft(), index) & 0xff;
+            if (chooseLeftChildProfile.profile(index < children.left.byteLength())) {
+                if (leftChildRawBytesNullProfile.profile(children.left.getRawBytes() == null)) {
+                    return byteSlowLeft.execute(children.left, index) & 0xff;
                 }
 
-                return rope.getLeft().getRawBytes()[index] & 0xff;
+                return children.left.getRawBytes()[index] & 0xff;
             }
 
-            if (rightChildRawBytesNullProfile.profile(rope.getRight().getRawBytes() == null)) {
-                return byteSlowRight.execute(rope.getRight(), index - rope.getLeft().byteLength()) & 0xff;
+            if (rightChildRawBytesNullProfile.profile(children.right.getRawBytes() == null)) {
+                return byteSlowRight.execute(children.right, index - children.left.byteLength()) & 0xff;
             }
 
-            return rope.getRight().getRawBytes()[index - rope.getLeft().byteLength()] & 0xff;
+            return children.right.getRawBytes()[index - children.left.byteLength()] & 0xff;
         }
 
+        // Necessary because getRawBytes() might return null, but then be populated and the children nulled
+        // before we get to run the other getByteConcatRope.
+        @Specialization(guards = "bytes != null")
+        protected int getByteConcatRope(ConcatRope rope, int index,
+                @Bind("rope.getBytesOrNull()") byte[] bytes) {
+            return bytes[index] & 0xff;
+        }
     }
 
     public abstract static class SetByteNode extends RubyContextNode {

src/main/java/org/truffleruby/core/rope/RopeOperations.java

@@ -39,6 +39,7 @@
 import org.truffleruby.collections.IntStack;
 import org.truffleruby.core.Hashing;
 import org.truffleruby.core.encoding.EncodingManager;
+import org.truffleruby.core.rope.ConcatRope.ConcatChildren;
 import org.truffleruby.core.rope.RopeNodesFactory.WithEncodingNodeGen;
 import org.truffleruby.core.string.StringAttributes;
 import org.truffleruby.core.string.StringOperations;
@@ -384,28 +385,36 @@ public static byte[] flattenBytes(Rope rope) {
             if (current instanceof ConcatRope) {
                 final ConcatRope concatRope = (ConcatRope) current;
 
+                final ConcatChildren children = concatRope.getChildrenOrNull();
+                if (children == null) {
+                    // The rope got concurrently flattened between entering the iteration and reaching here,
+                    // restart the iteration from the top.
+                    workStack.push(concatRope);
+                    continue;
+                }
+
                 // In the absence of any SubstringRopes, we always take the full contents of the ConcatRope.
                 if (substringLengths.isEmpty()) {
-                    workStack.push(concatRope.getRight());
-                    workStack.push(concatRope.getLeft());
+                    workStack.push(children.right);
+                    workStack.push(children.left);
                 } else {
-                    final int leftLength = concatRope.getLeft().byteLength();
+                    final int leftLength = children.left.byteLength();
 
                     // If we reach here, this ConcatRope is a descendant of a SubstringRope at some level. Based on
                     // the currently calculated byte[] offset and the number of bytes to extract, determine which of
                     // the ConcatRope's children we need to visit.
                     if (byteOffset < leftLength) {
                         if ((byteOffset + substringLengths.peek()) > leftLength) {
-                            workStack.push(concatRope.getRight());
-                            workStack.push(concatRope.getLeft());
+                            workStack.push(children.right);
+                            workStack.push(children.left);
                         } else {
-                            workStack.push(concatRope.getLeft());
+                            workStack.push(children.left);
                         }
                     } else {
                         // If we can skip the left child entirely, we need to update the offset so it's accurate for
                         // the right child as each child's starting point is 0.
                         byteOffset -= leftLength;
-                        workStack.push(concatRope.getRight());
+                        workStack.push(children.right);
                     }
                 }
             } else if (current instanceof SubstringRope) {
@@ -505,12 +514,16 @@ static byte getByteSlow(Rope rope, int index) {
                 return rawBytes[index];
             } else if (rope instanceof ConcatRope) {
                 final ConcatRope concatRope = (ConcatRope) rope;
-                final Rope left = concatRope.getLeft();
-                if (index < left.byteLength()) {
-                    rope = left;
+                final ConcatChildren children = concatRope.getChildrenOrNull();
+                if (children == null) {
+                    // Rope got concurrently flattened.
+                    return concatRope.getRawBytes()[index];
+                }
+                if (index < children.left.byteLength()) {
+                    rope = children.left;
                 } else {
-                    rope = concatRope.getRight();
-                    index -= left.byteLength();
+                    rope = children.right;
+                    index -= children.left.byteLength();
                 }
             } else if (rope instanceof SubstringRope) {
                 final SubstringRope substringRope = (SubstringRope) rope;
@@ -573,21 +586,27 @@ class Params {
                 workStack.push(new Params(child, startingHashCode, newOffset, length, false));
             } else if (rope instanceof ConcatRope) {
                 final ConcatRope concatRope = (ConcatRope) rope;
-                final Rope left = concatRope.getLeft();
-                final Rope right = concatRope.getRight();
-                final int leftLength = left.byteLength();
-
-                if (offset >= leftLength) {
-                    // range fully contained in right child
-                    workStack.push(new Params(right, startingHashCode, offset - leftLength, length, false));
-                } else if (offset + length <= leftLength) {
-                    // range fully contained in left child
-                    workStack.push(new Params(left, startingHashCode, offset, length, false));
+                final ConcatChildren children = concatRope.getChildrenOrNull();
+                if (children == null) {
+                    // Rope got concurrently flattened.
+                    resultHash = Hashing.stringHash(bytes, startingHashCode, offset, length);
                 } else {
-                    final int coveredByLeft = leftLength - offset;
-                    // push right node first, starting hash is the result from the left node
-                    workStack.push(new Params(right, 0, 0, length - coveredByLeft, true));
-                    workStack.push(new Params(left, startingHashCode, offset, coveredByLeft, false));
+                    final Rope left = children.left;
+                    final Rope right = children.right;
+                    final int leftLength = left.byteLength();
+
+                    if (offset >= leftLength) {
+                        // range fully contained in right child
+                        workStack.push(new Params(right, startingHashCode, offset - leftLength, length, false));
+                    } else if (offset + length <= leftLength) {
+                        // range fully contained in left child
+                        workStack.push(new Params(left, startingHashCode, offset, length, false));
+                    } else {
+                        final int coveredByLeft = leftLength - offset;
+                        // push right node first, starting hash is the result from the left node
+                        workStack.push(new Params(right, 0, 0, length - coveredByLeft, true));
+                        workStack.push(new Params(left, startingHashCode, offset, coveredByLeft, false));
+                    }
                 }
             } else if (rope instanceof RepeatingRope) {
                 final RepeatingRope repeatingRope = (RepeatingRope) rope;
@@ -668,14 +687,11 @@ public static boolean anyChildContains(Rope rope, String value) {
         if (rope instanceof SubstringRope) {
             return anyChildContains(((SubstringRope) rope).getChild(), value);
         } else if (rope instanceof ConcatRope) {
-            return anyChildContains(((ConcatRope) rope).getLeft(), value) ||
-                    anyChildContains(((ConcatRope) rope).getRight(), value);
-        } else {
-            if (rope.byteLength() < value.length()) {
-                return false;
+            ConcatChildren children = ((ConcatRope) rope).getChildrenOrNull();
+            if (children != null) {
+                return anyChildContains(children.left, value) || anyChildContains(children.right, value);
             }
-            return RopeOperations.decodeRope(rope).contains(value);
         }
+        return rope.byteLength() < value.length() && RopeOperations.decodeRope(rope).contains(value);
     }
-
 }