Flatten optional nesting for regex literal captures

When computing the CaptureList for AST nodes,
including converted AST -> DSL nodes, only permit
at most one level of optionality. This means that
regex literal captures are now either `Substring`
or `Substring?`.

Optional nesting is however still performed in the
DSL (due to result builder limitations). If a
regex literal is nested in the DSL, it may only
add at most one extra level of optionality to the
current nesting level.

Author: hamishknight

Sources/_RegexParser/Regex/Parse/CaptureList.swift

@@ -57,63 +57,105 @@ extension CaptureList {
   }
 }
 
+extension CaptureList {
+  public struct Builder {
+    public var captures = CaptureList()
+
+    public init() {}
+
+    public struct OptionalNesting {
+      // We maintain two depths, inner and outer. These allow e.g the nesting
+      // of a regex literal in a DSL, where outside of the scope of the literal,
+      // nesting is allowed, but inside the literal at most one extra layer of
+      // optionality may be added.
+      public var outerDepth: Int
+      public var canNest: Bool
+      public var innerDepth: Int
+
+      internal init(outerDepth: Int, canNest: Bool) {
+        self.outerDepth = outerDepth
+        self.canNest = canNest
+        self.innerDepth = 0
+      }
+
+      public init(canNest: Bool) {
+        self.init(outerDepth: 0, canNest: canNest)
+      }
+
+      public var depth: Int { outerDepth + innerDepth }
+
+      public var disablingNesting: Self {
+        // If we are currently able to nest, store the current depth as the
+        // outer depth, and disable nesting for an inner scope.
+        guard canNest else { return self }
+        return .init(outerDepth: depth, canNest: false)
+      }
+
+      public var addingOptional: Self {
+        var result = self
+        result.innerDepth = canNest ? innerDepth + 1 : 1
+        return result
+      }
+    }
+  }
+}
+
 // MARK: Generating from AST
 
-extension AST.Node {
-  public func _addCaptures(
-    to list: inout CaptureList,
-    optionalNesting nesting: Int
+extension CaptureList.Builder {
+  public mutating func addCaptures(
+    of node: AST.Node, optionalNesting nesting: OptionalNesting
   ) {
-    let addOptional = nesting+1
-    switch self {
+    switch node {
     case let .alternation(a):
       for child in a.children {
-        child._addCaptures(to: &list, optionalNesting: addOptional)
+        addCaptures(of: child, optionalNesting: nesting.addingOptional)
       }
 
     case let .concatenation(c):
       for child in c.children {
-        child._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: child, optionalNesting: nesting)
       }
 
     case let .group(g):
       switch g.kind.value {
       case .capture:
-        list.append(.init(optionalDepth: nesting, g.location))
+        captures.append(.init(optionalDepth: nesting.depth, g.location))
 
       case .namedCapture(let name):
-        list.append(.init(name: name.value, optionalDepth: nesting, g.location))
+        captures.append(.init(
+          name: name.value, optionalDepth: nesting.depth, g.location))
 
       case .balancedCapture(let b):
-        list.append(.init(name: b.name?.value, optionalDepth: nesting,
-                          g.location))
+        captures.append(.init(
+          name: b.name?.value, optionalDepth: nesting.depth, g.location))
 
       default: break
       }
-      g.child._addCaptures(to: &list, optionalNesting: nesting)
+      addCaptures(of: g.child, optionalNesting: nesting)
 
     case .conditional(let c):
       switch c.condition.kind {
       case .group(let g):
-        AST.Node.group(g)._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: .group(g), optionalNesting: nesting)
       default:
         break
       }
 
-      c.trueBranch._addCaptures(to: &list, optionalNesting: addOptional)
-      c.falseBranch._addCaptures(to: &list, optionalNesting: addOptional)
+      addCaptures(of: c.trueBranch, optionalNesting: nesting.addingOptional)
+      addCaptures(of: c.falseBranch, optionalNesting: nesting.addingOptional)
 
     case .quantification(let q):
       var optNesting = nesting
       if q.amount.value.bounds.atLeast == 0 {
-        optNesting += 1
+        optNesting = optNesting.addingOptional
       }
-      q.child._addCaptures(to: &list, optionalNesting: optNesting)
+      addCaptures(of: q.child, optionalNesting: optNesting)
 
     case .absentFunction(let abs):
       switch abs.kind {
       case .expression(_, _, let child):
-        child._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: child, optionalNesting: nesting)
       case .clearer, .repeater, .stopper:
         break
       }
@@ -122,16 +164,17 @@ extension AST.Node {
       break
     }
   }
+  public static func build(_ ast: AST) -> CaptureList {
+    var builder = Self()
+    builder.captures.append(.init(optionalDepth: 0, .fake))
+    builder.addCaptures(of: ast.root, optionalNesting: .init(canNest: false))
+    return builder.captures
+  }
 }
 
 extension AST {
   /// The capture list (including the whole match) of this AST.
-  public var captureList: CaptureList {
-    var caps = CaptureList()
-    caps.append(.init(optionalDepth: 0, .fake))
-    root._addCaptures(to: &caps, optionalNesting: 0)
-    return caps
-  }
+  public var captureList: CaptureList { .Builder.build(self) }
 }
 
 // MARK: Convenience for testing and inspection

Sources/_StringProcessing/Regex/DSLTree.swift

@@ -545,68 +545,62 @@ struct CaptureTransform: Hashable, CustomStringConvertible {
   }
 }
 
-// MARK: AST wrapper types
-//
-// These wrapper types are required because even @_spi-marked public APIs can't
-// include symbols from implementation-only dependencies.
-
-extension DSLTree.Node {
-  func _addCaptures(
-    to list: inout CaptureList,
-    optionalNesting nesting: Int
+extension CaptureList.Builder {
+  mutating func addCaptures(
+    of node: DSLTree.Node, optionalNesting nesting: OptionalNesting
   ) {
-    let addOptional = nesting+1
-    switch self {
+    switch node {
     case let .orderedChoice(children):
       for child in children {
-        child._addCaptures(to: &list, optionalNesting: addOptional)
+        addCaptures(of: child, optionalNesting: nesting.addingOptional)
       }
 
     case let .concatenation(children):
       for child in children {
-        child._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: child, optionalNesting: nesting)
       }
 
     case let .capture(name, _, child, transform):
-      list.append(.init(
+      captures.append(.init(
         name: name,
         type: transform?.resultType ?? child.wholeMatchType,
-        optionalDepth: nesting, .fake))
-      child._addCaptures(to: &list, optionalNesting: nesting)
+        optionalDepth: nesting.depth, .fake))
+      addCaptures(of: child, optionalNesting: nesting)
 
     case let .nonCapturingGroup(kind, child):
       assert(!kind.ast.isCapturing)
-      child._addCaptures(to: &list, optionalNesting: nesting)
+      addCaptures(of: child, optionalNesting: nesting)
 
     case let .conditional(cond, trueBranch, falseBranch):
       switch cond.ast {
       case .group(let g):
-        AST.Node.group(g)._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: .group(g), optionalNesting: nesting)
       default:
         break
       }
 
-      trueBranch._addCaptures(to: &list, optionalNesting: addOptional)
-      falseBranch._addCaptures(to: &list, optionalNesting: addOptional)
-
+      addCaptures(of: trueBranch, optionalNesting: nesting.addingOptional)
+      addCaptures(of: falseBranch, optionalNesting: nesting.addingOptional)
 
     case let .quantification(amount, _, child):
       var optNesting = nesting
       if amount.ast.bounds.atLeast == 0 {
-        optNesting += 1
+        optNesting = optNesting.addingOptional
       }
-      child._addCaptures(to: &list, optionalNesting: optNesting)
+      addCaptures(of: child, optionalNesting: optNesting)
 
     case let .absentFunction(abs):
       switch abs.ast.kind {
       case .expression(_, _, let child):
-        child._addCaptures(to: &list, optionalNesting: nesting)
+        addCaptures(of: child, optionalNesting: nesting)
       case .clearer, .repeater, .stopper:
         break
       }
 
     case let .convertedRegexLiteral(n, _):
-      return n._addCaptures(to: &list, optionalNesting: nesting)
+      // We disable nesting for converted AST trees, as literals do not nest
+      // captures. This includes literals nested in a DSL.
+      return addCaptures(of: n, optionalNesting: nesting.disablingNesting)
 
     case .matcher:
       break
@@ -617,6 +611,16 @@ extension DSLTree.Node {
     }
   }
 
+  static func build(_ dsl: DSLTree) -> CaptureList {
+    var builder = Self()
+    builder.captures.append(
+      .init(type: dsl.root.wholeMatchType, optionalDepth: 0, .fake))
+    builder.addCaptures(of: dsl.root, optionalNesting: .init(canNest: true))
+    return builder.captures
+  }
+}
+
+extension DSLTree.Node {
   /// Returns true if the node is output-forwarding, i.e. not defining its own
   /// output but forwarding its only child's output.
   var isOutputForwarding: Bool {
@@ -651,13 +655,13 @@ extension DSLTree.Node {
   }
 }
 
+// MARK: AST wrapper types
+//
+// These wrapper types are required because even @_spi-marked public APIs can't
+// include symbols from implementation-only dependencies.
+
 extension DSLTree {
-  var captureList: CaptureList {
-    var list = CaptureList()
-    list.append(.init(type: root.wholeMatchType, optionalDepth: 0, .fake))
-    root._addCaptures(to: &list, optionalNesting: 0)
-    return list
-  }
+  var captureList: CaptureList { .Builder.build(self) }
 
   /// Presents a wrapped version of `DSLTree.Node` that can provide an internal
   /// `_TreeNode` conformance.

Tests/RegexBuilderTests/RegexDSLTests.swift

@@ -1214,6 +1214,88 @@ class RegexDSLTests: XCTestCase {
   }
 
   func testOptionalNesting() throws {
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring?).self, ==)
+    {
+      try! Regex("(?:a|(b)*)?", as: (Substring, Substring?).self)
+    }
+
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring??).self, ==)
+    {
+      Optionally {
+        try! Regex("a|(b)*", as: (Substring, Substring?).self)
+      }
+    }
+
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring???).self, ==)
+    {
+      Optionally {
+        ChoiceOf {
+          try! Regex("a", as: Substring.self)
+          try! Regex("(b)*", as: (Substring, Substring?).self)
+        }
+      }
+    }
+
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring??).self, ==)
+    {
+      ChoiceOf {
+        try! Regex("a", as: Substring.self)
+        try! Regex("(b)*", as: (Substring, Substring?).self)
+      }
+    }
+
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring??).self, ==)
+    {
+      ChoiceOf {
+        try! Regex("a", as: Substring.self)
+        ZeroOrMore {
+          try! Regex("(b)", as: (Substring, Substring).self)
+        }
+      }
+    }
+
+    try _testDSLCaptures(
+      ("a", ("a", nil)),
+      ("", ("", nil)),
+      ("b", ("b", "b")),
+      ("bb", ("bb", "b")),
+      matchType: (Substring, Substring??).self, ==)
+    {
+      ChoiceOf {
+        try! Regex("a", as: Substring.self)
+        ZeroOrMore {
+          Capture {
+            try! Regex("b", as: Substring.self)
+          }
+        }
+      }
+    }
+
     let r = Regex {
       Optionally {
         Optionally {