Fix cursor bookkeeping in auto-possessification

Sources/_StringProcessing/Optimizations/AutoPossessification.swift

@@ -12,6 +12,16 @@
 internal import _RegexParser
 
 extension DSLList {
+  /// Returns the required first atom of the node at `position`, leaving
+  /// `position` just past that node's entire subtree, so that a caller walking
+  /// a list of nodes reads the next sibling.
+  ///
+  /// A `nil` result means this node has no effect on what has to match first,
+  /// and `.some(nil)` means that no required first atom can be determined.
+  ///
+  /// Some of the cases below stop looking at children as soon as the answer is
+  /// known, and others don't look at them at all, so the end of the subtree is
+  /// found up front rather than left to the traversal to reach.
   private func _requiredAtomImpl(
     _ position: inout Int,
     options: inout MatchingOptions,
@@ -20,7 +30,12 @@ extension DSLList {
     guard position < nodes.count else {
       return nil
     }
-
+
+    var endOfSubtree = position
+    skipNode(&endOfSubtree)
+    endOfSubtree += 1
+    defer { position = endOfSubtree }
+
     switch nodes[position] {
     case .atom(let atom):
       switch atom {
@@ -36,26 +51,15 @@ extension DSLList {
         return atom
       }
 
-    // In a concatenation, the first definitive child provides the answer,
-    // and then we need to skip past (in some cases at least) the remaining
-    // concatenation elements.
+    // In a concatenation, the first definitive child provides the answer.
     case .concatenation(let count):
-      var result: DSLTree.Atom?? = nil
-      var i = 0
-      while i < count {
-        i += 1
-        position += 1
+      position += 1
+      for _ in 0..<count {
         if let r = _requiredAtomImpl(&position, options: &options, allowOptionsChanges: allowOptionsChanges) {
-          result = r
-          break
+          return r
         }
       }
-
-      for _ in i..<count {
-        position += 1
-        skipNode(&position)
-      }
-      return result
+      return nil
 
     // For a quoted literal, we can look at the first char
     // TODO: matching semantics???
@@ -131,14 +135,30 @@ extension DSLList {
     return _requiredAtomImpl(&position, options: &options, allowOptionsChanges: allowOptionsChanges) ?? nil
   }
 
+  /// Auto-possessifies the quantifications inside the node at `position`,
+  /// leaving `position` just past that node's entire subtree, so that a caller
+  /// walking a list of nodes reads the next sibling.
+  ///
+  /// Returns a single-atom quantification that ends the node at `position`, if
+  /// there is one, so that an enclosing concatenation can try to possessify it
+  /// against whatever follows.
+  ///
+  /// Some of the cases below stop looking at children as soon as the answer is
+  /// known, and others don't look at them at all, so the end of the subtree is
+  /// found up front rather than left to the traversal to reach.
   internal mutating func autoPossessifyNextQuantification(
     _ position: inout Int,
     options: inout MatchingOptions
   ) -> (Int, DSLTree.Atom)? {
     guard position < nodes.count else {
       return nil
     }
-
+
+    var endOfSubtree = position
+    skipNode(&endOfSubtree)
+    endOfSubtree += 1
+    defer { position = endOfSubtree }
+
     switch nodes[position] {
     case .quantification(_, _):
       let quantPosition = position
@@ -153,57 +173,49 @@ extension DSLList {
       case .atom(let atom) where atom.isMatchable:
         return (quantPosition, atom)
       default:
-        var innerPosition = position
-        _ = autoPossessifyNextQuantification(&innerPosition, options: &options)
+        _ = autoPossessifyNextQuantification(&position, options: &options)
         return nil
       }
 
     case .concatenation(let count):
-      // If we find a valid quantification among this concatenation's components,
-      // we must look for a required atom in the sibling. If a definitive result
-      // is not found, pop up the recursion stack to find a sibling at a higher
-      // level.
-      var foundQuantification: (Int, DSLTree.Atom)? = nil
-      var foundNextAtom: DSLTree.Atom? = nil
-      var i = 0
+      // A quantification found among this concatenation's components can be
+      // possessified against the first component after it that has a required
+      // first atom. If no component provides a definitive answer, pop up the
+      // recursion stack to find a sibling at a higher level.
+      var pendingQuantification: (Int, DSLTree.Atom)? = nil
       position += 1
-      while i < count {
-        i += 1
-        if let result = autoPossessifyNextQuantification(&position, options: &options) {
-          foundQuantification = result
-          break
+      for _ in 0..<count {
+        if let (quantIndex, quantAtom) = pendingQuantification {
+          // Look for a required atom in this component without consuming it,
+          // so that its own quantifications still get possessified below.
+          var lookahead = position
+          switch _requiredAtomImpl(&lookahead, options: &options, allowOptionsChanges: false) {
+          case .some(.some(let nextAtom)):
+            // We found a quantifier with a required first atom and a required
+            // following atom. If the second is excluded by the first, we can
+            // safely convert the quantifier to possessive.
+            if quantAtom.excludes(nextAtom, options: options),
+                case .quantification(let amount, _) = nodes[quantIndex]
+            {
+              nodes[quantIndex] = .quantification(amount, .explicit(.possessive))
+            }
+            pendingQuantification = nil
+          case .some(nil):
+            // Nothing is known about what has to match here, so the pending
+            // quantification can't be possessified at any level.
+            pendingQuantification = nil
+          case nil:
+            // This component doesn't match anything on its own; keep looking.
+            break
+          }
         }
-      }
 
-      while i < count {
-        i += 1
-        position += 1
-        if let result = _requiredAtomImpl(&position, options: &options, allowOptionsChanges: false) {
-          foundNextAtom = result
-          break
+        if let result = autoPossessifyNextQuantification(&position, options: &options) {
+          pendingQuantification = result
         }
       }
 
-      for _ in i..<count {
-        position += 1
-        skipNode(&position)
-      }
-
-      guard let (quantIndex, firstAtom) = foundQuantification,
-            let nextAtom = foundNextAtom
-      else { return foundQuantification }
-
-      // We found a quantifier with a required first atom and a required
-      // following atom. If the second is excluded by the first, we can
-      // safely convert the quantifier to possessive.
-
-      if firstAtom.excludes(nextAtom, options: options),
-          case .quantification(let amount, _) = nodes[quantIndex]
-      {
-        nodes[quantIndex] = .quantification(amount, .explicit(.possessive))
-      }
-
-      return nil
+      return pendingQuantification
 
     // For alternations, we need to explore / auto-possessify in the different
     // branches, but quantifications inside an alternation don't
@@ -240,7 +252,6 @@ extension DSLList {
       return autoPossessifyNextQuantification(&position, options: &options)
 
     case .atom(let atom):
-      position += 1
       switch atom {
       case .changeMatchingOptions(let seq):
         options.apply(seq.ast)
@@ -251,10 +262,10 @@ extension DSLList {
     default:
       // Multi-child nodes are handled above, just handle 0 and 1 here.
       let childCount = nodes[position].directChildren
-      position += 1
 
       assert(childCount <= 1)
       if childCount == 1 {
+        position += 1
         return autoPossessifyNextQuantification(&position, options: &options)
       }
     }

Tests/RegexTests/MatchTests.swift

@@ -748,6 +748,28 @@ extension RegexTests {
       "a+(?:(?i)A)",
       ("aaaaA", "aaaaA"),
       ("aaaaa", "aaaaa"))
+    // - a quantification that can match nothing must not be mistaken for the
+    //   atom that follows the quantification being tested (issue #898)
+    firstMatchTests(
+      #"(?i)(?:a*\+*)A{2}"#,
+      ("aA", "aA"))
+    firstMatchTests(
+      #"(?i)(a*\+*)A{2}"#,
+      ("aA", "aA"))
+    firstMatchTests(
+      #"(b{1,3}\s?)b"#,
+      ("bb", "bb"))
+    firstMatchTests(
+      #"(?:a+a?|a*|\+.\+)[ab]*."#,
+      ("a", "a"))
+    // - a component with no determinable required first atom rules out
+    //   possessifying an earlier quantification, at any nesting level
+    firstMatchTests(
+      "(?:a*[ab])b",
+      ("aab", "aab"))
+    firstMatchTests(
+      "(?:(?:a*[ab])x?)b",
+      ("aab", "aab"))
 
     // XFAIL'd possessive tests
     firstMatchTests(

Tests/RegexTests/OptimizationTests.swift

@@ -51,6 +51,7 @@ import Testing
     #/a?/#, #/a+a/#, #/a+(?:b|c)/#, #/(?:a+|b+)/#, #/[a]/#, #/a?a/#,
     #/(?i)a+A/#, #/(?i:a+A)/#,  // case insensitivity when checking exclusion
     #/(?:(?:ab)+b)/#,           // single atom quantifications only
+    #/(?:a*[ab])b/#,            // a component that rules out a required atom
   ])
   func noAutoPossessify(pattern: Regex<Substring>) throws {
     var list = pattern.program.list
@@ -65,4 +66,56 @@ import Testing
       }
     }
   }
+
+  /// The number of quantifications in `pattern` that auto-possessification
+  /// converts to possessive, and the total number of quantifications.
+  @available(macOS 9999, *)
+  func possessificationCounts(
+    in pattern: String
+  ) throws -> (possessive: Int, total: Int) {
+    var list = try Regex(pattern).program.list
+    list.autoPossessify()
+    var possessive = 0
+    var total = 0
+    for node in list.nodes {
+      guard case .quantification(_, let kind) = node else { continue }
+      total += 1
+      if kind.isExplicit && kind.quantificationKind?.ast == .possessive {
+        possessive += 1
+      }
+    }
+    return (possessive, total)
+  }
+
+  /// Auto-possessification walks a flat list of nodes, so a case that stops
+  /// looking at a node's children has to leave the cursor past the whole
+  /// subtree anyway. Losing track of it possessifies the wrong quantification,
+  /// which drops matches, or stops the traversal short, which silently gives
+  /// up the optimization.
+  @available(macOS 9999, *)
+  @Test func partialAutoPossessify() throws {
+    // pattern, expected possessive count, expected quantification count
+    let expectations: [(String, Int, Int)] = [
+      // Every quantification of a concatenation participates, not just the
+      // first one.
+      (#"a+b\d+c"#, 2, 2),
+      (#"x*a+b"#, 2, 2),
+      (#"\w+\s\d+a"#, 2, 2),
+      // `\s?` is excluded by the following `b`, but `b{1,3}` isn't
+      // possessifiable, because `\s?` is allowed to match nothing.
+      (#"(b{1,3}\s?)b"#, 1, 2),
+      // Same, one level down: only `\+*` is excluded by what follows it.
+      (#"(?i)(?:a*\+*)A{2}"#, 1, 3),
+      // Nothing is known about what has to match after `[ab]`, so `a*` can't
+      // be possessified at any level.
+      (#"(?:a*[ab])b"#, 0, 1),
+      (#"(?:(?:a*[ab])x?)b"#, 1, 2),
+    ]
+    for (pattern, possessive, total) in expectations {
+      let counts = try possessificationCounts(in: pattern)
+      #expect(
+        counts == (possessive, total),
+        "Unexpected possessification in '\(pattern)'")
+    }
+  }
 }