Coalesce adjacent scalars and characters in the DSL

Previously we would emit a series of scalars
written in the DSL as a series of individual
characters in grapheme semantic mode. Change the
behavior such that we coalesce any adjacent
scalars and characters, including those in regex
literals and nested concatenations. We then
perform grapheme breaking over the result, and can
emit character matches for scalars that coalesced
into a grapheme.

This transform subsumes a similar transform we
performed for regex literals when converting them
to a DSLTree. This has the nice side effect of
allowing us to better preserve scalar syntax in
the DSL transform.

rdar://96942688

Author: hamishknight

Sources/_StringProcessing/ByteCodeGen.swift

@@ -791,6 +791,41 @@ fileprivate extension Compiler.ByteCodeGen {
     }
   }
 
+  mutating func emitConcatenation(_ children: [DSLTree.Node]) throws {
+    // Before emitting a concatenation, we need to flatten out any nested
+    // concatenations, and coalesce any adjacent characters and scalars, forming
+    // quoted literals of their contents, over which we can perform grapheme
+    // breaking.
+    func flatten(_ node: DSLTree.Node) -> [DSLTree.Node] {
+      switch node {
+      case .concatenation(let ch):
+        return ch.flatMap(flatten)
+      case .convertedRegexLiteral(let n, _):
+        return flatten(n)
+      default:
+        return [node]
+      }
+    }
+    let children = children
+      .flatMap(flatten)
+      .coalescing(with: "", into: DSLTree.Node.quotedLiteral) { str, node in
+        switch node {
+        case .atom(let a):
+          guard let c = a.literalCharacterValue else { return false }
+          str.append(c)
+          return true
+        case .quotedLiteral(let q):
+          str += q
+          return true
+        default:
+          return false
+        }
+      }
+    for child in children {
+      try emitConcatenationComponent(child)
+    }
+  }
+
   @discardableResult
   mutating func emitNode(_ node: DSLTree.Node) throws -> ValueRegister? {
     switch node {
@@ -799,9 +834,7 @@ fileprivate extension Compiler.ByteCodeGen {
       try emitAlternation(children)
 
     case let .concatenation(children):
-      for child in children {
-        try emitConcatenationComponent(child)
-      }
+      try emitConcatenation(children)
 
     case let .capture(name, refId, child, transform):
       options.beginScope()

Sources/_StringProcessing/PrintAsPattern.swift

@@ -70,16 +70,9 @@ extension PrettyPrinter {
     for namedCapture in namedCaptures {
       print("let \(namedCapture) = Reference(Substring.self)")
     }
-    
-    switch node {
-    case .concatenation(_):
-      printAsPattern(convertedFromAST: node)
-    case .convertedRegexLiteral(.concatenation(_), _):
-      printAsPattern(convertedFromAST: node)
-    default:
-      printBlock("Regex") { printer in
-        printer.printAsPattern(convertedFromAST: node)
-      }
+
+    printBlock("Regex") { printer in
+      printer.printAsPattern(convertedFromAST: node, isTopLevel: true)
     }
   }
 
@@ -89,7 +82,7 @@ extension PrettyPrinter {
   // to have a non-backing-off pretty-printer that this
   // can defer to.
   private mutating func printAsPattern(
-    convertedFromAST node: DSLTree.Node
+    convertedFromAST node: DSLTree.Node, isTopLevel: Bool = false
   ) {
     if patternBackoff(DSLTree._Tree(node)) {
       printBackoff(node)
@@ -106,11 +99,7 @@ extension PrettyPrinter {
       }
 
     case let .concatenation(c):
-      printBlock("Regex") { printer in
-        c.forEach {
-          printer.printAsPattern(convertedFromAST: $0)
-        }
-      }
+      printConcatenationAsPattern(c, isTopLevel: isTopLevel)
 
     case let .nonCapturingGroup(kind, child):
       switch kind.ast {
@@ -263,7 +252,7 @@ extension PrettyPrinter {
       // check above, so it should work out. Need a
       // cleaner way to do this. This means the argument
       // label is a lie.
-      printAsPattern(convertedFromAST: n)
+      printAsPattern(convertedFromAST: n, isTopLevel: isTopLevel)
 
     case let .customCharacterClass(ccc):
       printAsPattern(ccc)
@@ -279,6 +268,60 @@ extension PrettyPrinter {
       print("/* TODO: absent function */")
     }
   }
+
+  enum NodeToPrint {
+    case dslNode(DSLTree.Node)
+    case stringLiteral(String)
+  }
+
+  mutating func printAsPattern(_ node: NodeToPrint) {
+    switch node {
+    case .dslNode(let n):
+      printAsPattern(convertedFromAST: n)
+    case .stringLiteral(let str):
+      print(str)
+    }
+  }
+
+  mutating func printConcatenationAsPattern(
+    _ nodes: [DSLTree.Node], isTopLevel: Bool
+  ) {
+    // We need to coalesce any adjacent character and scalar elements into a
+    // string literal, preserving scalar syntax.
+    let nodes = nodes
+      .map { NodeToPrint.dslNode($0.lookingThroughConvertedLiteral) }
+      .coalescing(
+        with: StringLiteralBuilder(), into: { .stringLiteral($0.result) }
+      ) { literal, node in
+        guard case .dslNode(let node) = node else { return false }
+        switch node {
+        case let .atom(.char(c)):
+          literal.append(c)
+          return true
+        case let .atom(.scalar(s)):
+          literal.append(unescaped: s._dslBase)
+          return true
+        case .quotedLiteral(let q):
+          literal.append(q)
+          return true
+        default:
+          return false
+        }
+      }
+    if isTopLevel || nodes.count == 1 {
+      // If we're at the top level, or we coalesced everything into a single
+      // element, we don't need to print a surrounding Regex { ... }.
+      for n in nodes {
+        printAsPattern(n)
+      }
+      return
+    }
+    printBlock("Regex") { printer in
+      for n in nodes {
+        printer.printAsPattern(n)
+      }
+    }
+  }
 
   mutating func printAsPattern(
     _ ccc: DSLTree.CustomCharacterClass,
@@ -341,8 +384,7 @@ extension PrettyPrinter {
           charMembers.append(c)
           return false
         case let .scalar(s):
-          charMembers.append(
-            unescaped: "\\u{\(String(s.value, radix: 16, uppercase: true))}")
+          charMembers.append(unescaped: s._dslBase)
           return false
         case .unconverted(_):
           return true
@@ -449,9 +491,9 @@ extension PrettyPrinter {
       case let .scalar(s):
 
         if wrap {
-          output("One(.anyOf(\"\\u{\(String(s.value, radix: 16, uppercase: true))}\"))")
+          output("One(.anyOf(\(s._dslBase._bareQuoted)))")
         } else {
-          output(".anyOf(\"\\u{\(String(s.value, radix: 16, uppercase: true))}\")")
+          output(".anyOf(\(s._dslBase._bareQuoted))")
         }
 
       case let .unconverted(a):
@@ -625,6 +667,10 @@ extension String {
   }
 }
 
+extension UnicodeScalar {
+  var _dslBase: String { "\\u{\(String(value, radix: 16, uppercase: true))}" }
+}
+
 /// A helper for building string literals, which handles escaping the contents
 /// appended.
 fileprivate struct StringLiteralBuilder {
@@ -851,19 +897,15 @@ extension AST.Atom {
   }
 
   var _dslBase: (String, canBeWrapped: Bool) {
-    func scalarLiteral(_ s: UnicodeScalar) -> String {
-      let hex = String(s.value, radix: 16, uppercase: true)
-      return "\\u{\(hex)}"
-    }
     switch kind {
     case let .char(c):
       return (String(c), false)
 
     case let .scalar(s):
-      return (scalarLiteral(s.value), false)
+      return (s.value._dslBase, false)
 
     case let .scalarSequence(seq):
-      return (seq.scalarValues.map(scalarLiteral).joined(), false)
+      return (seq.scalarValues.map(\._dslBase).joined(), false)
 
     case let .property(p):
       return (p._dslBase, true)

Sources/_StringProcessing/Regex/ASTConversion.swift

@@ -43,61 +43,7 @@ extension AST.Node {
         return .orderedChoice(children)
 
       case let .concatenation(v):
-        // Coalesce adjacent children who can produce a
-        // string literal representation
-        let astChildren = v.children
-        func coalesce(
-          _ idx: Array<AST>.Index
-        ) -> (Array<AST>.Index, String)? {
-          var result = ""
-          var idx = idx
-          while idx < astChildren.endIndex {
-            guard let atom: AST.Atom = astChildren[idx].as() else { break }
-
-            // TODO: For printing, nice to coalesce
-            // scalars literals too. We likely need a different
-            // approach even before we have a better IR.
-            if let char = atom.singleCharacter  {
-              result.append(char)
-            } else if let scalar = atom.singleScalar {
-              result.append(Character(scalar))
-            } else if case .scalarSequence(let seq) = atom.kind {
-              result += seq.scalarValues.map(Character.init)
-            } else {
-              break
-            }
-            
-            astChildren.formIndex(after: &idx)
-          }
-          return result.isEmpty ? nil : (idx, result)
-        }
-
-        // No need to nest single children concatenations
-        if astChildren.count == 1 {
-          return astChildren.first!.dslTreeNode
-        }
-
-        // Check for a single child post-coalescing
-        if let (idx, str) = coalesce(astChildren.startIndex),
-           idx == astChildren.endIndex
-        {
-          return .quotedLiteral(str)
-        }
-
-        // Coalesce adjacent string children
-        var curIdx = astChildren.startIndex
-        var children = Array<DSLTree.Node>()
-        while curIdx < astChildren.endIndex {
-          if let (nextIdx, str) = coalesce(curIdx) {
-            // TODO: Track source info...
-            children.append(.quotedLiteral(str))
-            curIdx = nextIdx
-          } else {
-            children.append(astChildren[curIdx].dslTreeNode)
-            astChildren.formIndex(after: &curIdx)
-          }
-        }
-        return .concatenation(children)
+        return .concatenation(v.children.map(\.dslTreeNode))
 
       case let .group(v):
         let child = v.child.dslTreeNode
@@ -135,10 +81,9 @@ extension AST.Node {
       case let .atom(v):
         switch v.kind {
         case .scalarSequence(let seq):
-          // Scalar sequences are splatted into concatenated scalars, which
-          // becomes a quoted literal. Sequences nested in concatenations have
-          // already been coalesced, this just handles the lone atom case.
-          return .quotedLiteral(String(seq.scalarValues.map(Character.init)))
+          // The DSL doesn't have an equivalent node for scalar sequences. Splat
+          // them into a concatenation of scalars.
+          return .concatenation(seq.scalarValues.map { .atom(.scalar($0)) })
         default:
           return .atom(v.dslTreeAtom)
         }

Sources/_StringProcessing/Regex/DSLTree.swift

@@ -334,6 +334,14 @@ extension DSLTree.Node {
     default: return nil
     }
   }
+
+  /// If this node is for a converted literal, look through it.
+  var lookingThroughConvertedLiteral: Self {
+    switch self {
+    case let .convertedRegexLiteral(n, _): return n
+    default: return self
+    }
+  }
 }
 
 extension DSLTree.Atom {

Sources/_StringProcessing/Utility/Misc.swift

@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+extension Array {
+  /// Coalesce adjacent elements using a given accumulator. The accumulator is
+  /// transformed into an element of the array by `finish`. The `accumulate`
+  /// function should return `true` if the accumulator has coalesced the
+  /// element, `false` otherwise.
+  func coalescing<T>(
+    with initialAccumulator: T, into finish: (T) -> Element,
+    accumulate: (inout T, Element) -> Bool
+  ) -> Self {
+    var didAccumulate = false
+    var accumulator = initialAccumulator
+
+    var result = Self()
+    for elt in self {
+      if accumulate(&accumulator, elt) {
+        // The element has been coalesced into accumulator, there is nothing
+        // else to do.
+        didAccumulate = true
+        continue
+      }
+      if didAccumulate {
+        // We have a leftover accumulator, which needs to be finished before we
+        // can append the next element.
+        result.append(finish(accumulator))
+        accumulator = initialAccumulator
+        didAccumulate = false
+      }
+      result.append(elt)
+    }
+    // Handle a leftover accumulation.
+    if didAccumulate {
+      result.append(finish(accumulator))
+    }
+    return result
+  }
+}