TRegex: further restrict quantifier unrolling to avoid NFA transition

explosions.

Author: djoooooe

regex/src/com.oracle.truffle.regex.test/src/com/oracle/truffle/regex/tregex/test/JsTests.java

@@ -688,6 +688,11 @@ public void generatedTests() {
         /* GENERATED CODE END - KEEP THIS MARKER FOR AUTOMATIC UPDATES */
     }
 
+    @Test
+    public void quantifierUnrollNFAExplosion() {
+        test("(?:\\3((?:[^]|.{0}|\\B|\ub9b5\\b|$){4,}){1,4})", "yim", "\ua3bb\n\n\u00a1\n\na\u2f77\n\n\ua3bb\n\n\u00a1\n\na\u2f77\n\n", 0, false);
+    }
+
     @Test
     public void overlappingBq() {
         testBoolean("(?=a{2,4})[ab]{4,68}c", "", NEVER_UNROLL_OPT, "aabbbbbbbbbbbbbbbbbbbbbbc", 0, true);

regex/src/com.oracle.truffle.regex.test/src/com/oracle/truffle/regex/tregex/test/OracleDBTests.java

@@ -1700,4 +1700,10 @@ public void orcl38190286() {
         test("[[:alpha:]]", "", Collections.emptyMap(), Encodings.UTF_16, "\uDDF2", 0, false);
         test("[[:alpha:]]", "", Collections.emptyMap(), Encodings.UTF_16, "\uD839\uDDF2", 0, false);
     }
+
+    @Test
+    public void bqTransitionExplosion() {
+        test("(a(b(b(b(b(b(b(b(b(b(b(b(b(b(b(b(b(b(b(b|)|)|)|)|)|)|)|)|)|)|)|)|)|)|)|)|)|)|){2,2}c)de", "", Map.of("regexDummyLang.QuantifierUnrollLimitGroup", "1"),
+                        "abbbbbbbcdebbbbbbbf", 0, true, 0, 11, 0, 9, 8, 8, 2, 8, 3, 8, 4, 8, 5, 8, 6, 8, 7, 8, 8, 8);
+    }
 }

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/TRegexOptions.java

@@ -183,6 +183,12 @@ public class TRegexOptions {
      */
     public static final int TRegexQuantifierUnrollLimitGroup = 6;
 
+    /**
+     * Quantified groups whose node count is greater that this threshold will not be considered for
+     * quantifier unrolling.
+     */
+    public static final int TRegexQuantifierUnrollLimitGroupNodeCount = 100;
+
     /**
      * Bailout threshold for number of capture groups.
      */

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/automaton/StateTransitionCanonicalizer.java

@@ -42,6 +42,9 @@
 
 import java.util.Arrays;
 import java.util.Iterator;
+import java.util.Objects;
+
+import org.graalvm.collections.EconomicSet;
 
 import com.oracle.truffle.regex.charset.CodePointSet;
 import com.oracle.truffle.regex.tregex.buffer.CompilationBuffer;
@@ -69,6 +72,7 @@ public abstract class StateTransitionCanonicalizer<SI extends StateIndex<? super
     private boolean anyArgConstraints;
     private final IntArrayBuffer stack = new IntArrayBuffer();
     private final IntArrayBuffer skipStack = new IntArrayBuffer();
+    private final EconomicSet<ConstraintDeduplicationKey> constraintDeduplicationSet = EconomicSet.create();
 
     private static final int INITIAL_CAPACITY = 8;
 
@@ -287,6 +291,7 @@ private void calcDisjointTransitions(CompilationBuffer compilationBuffer) {
             for (int i : intersectingArgs[iStage1]) {
                 addToStack(argTransitions.get(i), matcher, argConstraints.get(i), argOperations.get(i), resultLength);
             }
+            constraintDeduplicationSet.clear();
             outer: while (!stack.isEmpty()) {
                 int i = stack.pop();
                 T argTransition = argTransitions.get(i);
@@ -301,30 +306,36 @@ private void calcDisjointTransitions(CompilationBuffer compilationBuffer) {
                     if (constraintMerge == null) {
                         continue;
                     }
-                    assert TransitionConstraint.isNormalized(constraintMerge.middle());
 
                     // if the slot already contains an unconditional transition to the final state,
                     // we don't have to check any other constraints.
                     if (!(shouldPruneAfterFinalState() && leadsToFinalState.get(j))) {
                         if (constraintMerge.lhs().length == 0) {
                             addTransitionTo(j, argTransition, argOperation);
                         } else {
+                            assert TransitionConstraint.isNormalized(constraintMerge.lhs()[0]);
                             constraintBuilder[j] = constraintMerge.lhs()[0];
                             for (int k = 1; k < constraintMerge.lhs().length; ++k) {
                                 duplicateSlot(j, matcher, constraintMerge.lhs()[k]);
                                 resultLength++;
                             }
-
                             duplicateSlot(j, matcher, constraintMerge.middle());
                             addTransitionTo(resultLength, argTransition, argOperation);
                             resultLength++;
                         }
                     }
                     for (long[] rhs : constraintMerge.rhs()) {
-                        addToStack(argTransition, matcher, rhs, argOperation, j + 1);
+                        // TransitionConstraint.intersectAndSubtract may yield the same constraint
+                        // formula for the same transition multiple times when intersecting a large
+                        // set of transitions. We drop those duplicates here, otherwise they would
+                        // increase this stage's run time exponentially.
+                        if (constraintDeduplicationSet.add(new ConstraintDeduplicationKey(argTransition.getId(), rhs))) {
+                            addToStack(argTransition, matcher, rhs, argOperation, j + 1);
+                        }
                     }
                     continue outer;
                 }
+                assert TransitionConstraint.isNormalized(argConstraint);
                 createSlot();
                 matcherBuilders[resultLength] = matcher;
                 constraintBuilder[resultLength] = argConstraint;
@@ -336,6 +347,7 @@ private void calcDisjointTransitions(CompilationBuffer compilationBuffer) {
     }
 
     private void duplicateSlot(int i, CodePointSet matcher, long[] constraints) {
+        assert TransitionConstraint.isNormalized(constraints);
         createSlot();
         targetStateSets[resultLength] = targetStateSets[i].copy();
         transitionLists[resultLength].addAll(transitionLists[i]);
@@ -473,6 +485,22 @@ private TB[] mergeSameTargets(CompilationBuffer compilationBuffer) {
         return resultBuffer2.toArray(createResultArray(resultBuffer2.length()));
     }
 
+    private record ConstraintDeduplicationKey(int transitionID, long[] constraints) {
+
+        @Override
+        public boolean equals(Object obj) {
+            if (!(obj instanceof ConstraintDeduplicationKey o)) {
+                return false;
+            }
+            return transitionID == o.transitionID && Arrays.equals(constraints, o.constraints);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(transitionID, Arrays.hashCode(constraints));
+        }
+    }
+
     protected abstract TB createTransitionBuilder(T[] transitions, StateSet<SI, S> targetStateSet, CodePointSet matcherBuilder, long[] constraints, long[] operations);
 
     /**

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/nfa/NFAStateTransition.java

@@ -148,6 +148,11 @@ public int hashCode() {
         return getId();
     }
 
+    @Override
+    public String toString() {
+        return String.format("%d - %s -> %d", source.getId(), codePointSet, target.getId());
+    }
+
     @TruffleBoundary
     @Override
     public JsonValue toJson() {

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/parser/RegexASTPostProcessor.java

@@ -178,23 +178,25 @@ private boolean shouldUnroll(QuantifiableTerm term) {
             return ast.getNumberOfNodes() <= TRegexOptions.TRegexMaxParseTreeSizeForDFA && (term.getQuantifier().isUnrollTrivial() || term.isUnrollingCandidate(ast.getOptions()));
         }
 
-        private static final class ShouldUnrollQuantifierVisitor extends DepthFirstTraversalRegexASTVisitor {
+        private static final class ShouldUnrollQuantifierVisitor extends NodeCountVisitor {
 
-            private boolean result;
+            private boolean containsBackReference;
 
             boolean shouldUnroll(Group group) {
                 assert group.hasQuantifier();
                 if (group.getQuantifier().isUnrollTrivial()) {
                     return true;
                 }
-                result = true;
+                count = 0;
+                containsBackReference = false;
                 run(group);
-                return result;
+                return count <= TRegexOptions.TRegexQuantifierUnrollLimitGroupNodeCount && !containsBackReference;
             }
 
             @Override
             protected void visit(BackReference backReference) {
-                result = false;
+                super.visit(backReference);
+                containsBackReference = true;
             }
         }
 

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/parser/ast/CharacterClass.java

@@ -128,7 +128,28 @@ public void setWasSingleChar(boolean value) {
 
     @Override
     public boolean isUnrollingCandidate(RegexOptions options) {
-        return hasQuantifier() && getQuantifier().isWithinThreshold(options.quantifierUnrollLimitSingleCC);
+        if (!hasQuantifier()) {
+            return false;
+        }
+        if (getQuantifier().isWithinThreshold(Math.min(options.quantifierUnrollLimitGroup, options.quantifierUnrollLimitSingleCC))) {
+            return true;
+        }
+        if (getQuantifier().isWithinThreshold(Math.max(options.quantifierUnrollLimitGroup, options.quantifierUnrollLimitSingleCC))) {
+            // If the quantifier on this character class is large, but still within unroll
+            // threshold, we take parent group quantifiers into account; If there is a bounded
+            // quantifier on any parent group, unrolling the current quantifier may lead to either
+            // an explosion of NFA states if the parent quantifier is unrolled as well, or an
+            // explosion of guarded NFA transitions if the parent quantifier is not unrolled.
+            Group parent = getQuantifiedParentGroup();
+            while (parent != null) {
+                if (!parent.getQuantifier().isUnrollTrivial()) {
+                    return false;
+                }
+                parent = parent.getQuantifiedParentGroup();
+            }
+            return true;
+        }
+        return false;
     }
 
     public void addLookBehindEntry(RegexAST ast, LookBehindAssertion lookBehindEntry) {

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/parser/ast/RegexASTNode.java

@@ -706,26 +706,6 @@ public SubexpressionCall asSubexpressionCall() {
         return (SubexpressionCall) this;
     }
 
-    public Group getParentQuantifier() {
-        RegexASTNode current = this;
-        while (current.getParent() != null) {
-            assert current instanceof Term;
-            if (current.getParent() instanceof RegexASTSubtreeRootNode) {
-                return null;
-            } else {
-                assert current.getParent().getParent() instanceof Group;
-                var nextParent = current.getParent().getParent().asGroup();
-                if (nextParent.hasQuantifier() && !nextParent.isExpandedQuantifier()) {
-                    return nextParent;
-                }
-                // structure is always Group -> Sequence -> Term
-                current = nextParent;
-            }
-        }
-        // this should only be reached by nodes generated by RegexAST#createNFAInitialStates()!
-        return null;
-    }
-
     protected JsonObject toJson(String typeName) {
         return Json.obj(Json.prop("id", id),
                         Json.prop("type", typeName),

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/parser/ast/Term.java

@@ -95,4 +95,23 @@ public RegexASTSubtreeRootNode getSubTreeParent() {
         // this should only be reached by nodes generated by RegexAST#createNFAInitialStates()!
         return null;
     }
+
+    public Group getQuantifiedParentGroup() {
+        RegexASTNode current = this;
+        while (current.getParent() != null) {
+            if (current.getParent() instanceof RegexASTSubtreeRootNode) {
+                return null;
+            } else {
+                assert current.getParent().getParent() instanceof Group;
+                Group nextParent = current.getParent().getParent().asGroup();
+                if (nextParent.hasNotUnrolledQuantifier()) {
+                    return nextParent;
+                }
+                // structure is always Group -> Sequence -> Term
+                current = nextParent;
+            }
+        }
+        // this should only be reached by nodes generated by RegexAST#createNFAInitialStates()!
+        return null;
+    }
 }

regex/src/com.oracle.truffle.regex/src/com/oracle/truffle/regex/tregex/parser/ast/visitors/NFATraversalRegexASTVisitor.java

@@ -408,8 +408,9 @@ protected long[] getTransitionGuardsOnPath() {
      * @return the quantifier index or -1 if no guards are needed.
      */
     protected int needsMaintainGuard() {
-        Group parentQuant = root.getParentQuantifier();
-        Group otherParent = cur.getParentQuantifier();
+        assert cur instanceof Term;
+        Group parentQuant = root.getQuantifiedParentGroup();
+        Group otherParent = ((Term) cur).getQuantifiedParentGroup();
         if (parentQuant != null && parentQuant.equals(otherParent)) {
             for (long guard : transitionGuardsCanonicalized) {
                 if (TransitionGuard.isQuantifierOp(guard)) {