Improve assert statement generation and fix token choices

Sources/SwiftSyntaxBuilderGeneration/Child.swift

@@ -26,6 +26,7 @@ class Child {
   let classification: SyntaxClassification?
   /// A restricted set of token kinds that will be accepted for this child.
   let tokenChoices: [Token]
+  let tokenCanContainArbitraryText: Bool
 
   var swiftName: String {
     return lowercaseFirstWord(name: name)
@@ -95,14 +96,18 @@ class Child {
     self.requiresLeadingNewline = requiresLeadingNewline
     self.isOptional = isOptional
 
+    let isToken = syntaxKind.hasSuffix("Token")
     var mappedTokenChoices = [Token]()
 
-    if syntaxKind.hasSuffix("Token"), let token = SYNTAX_TOKEN_MAP[syntaxKind] {
+    if isToken, let token = SYNTAX_TOKEN_MAP[syntaxKind] {
       mappedTokenChoices.append(token)
     }
 
-    mappedTokenChoices.append(contentsOf: tokenChoices.compactMap { SYNTAX_TOKEN_MAP[$0] })
+    mappedTokenChoices.append(contentsOf: tokenChoices.compactMap { SYNTAX_TOKEN_MAP["\($0)Token"] })
     self.tokenChoices = mappedTokenChoices
+
+    // If mappedTokenChoices contains `nil`, the token can contain arbitrary text
+    self.tokenCanContainArbitraryText = mappedTokenChoices.contains { $0.text == nil }
 
     // A list of valid text for tokens, if specified.
     // This will force validation logic to check the text passed into the

Sources/SwiftSyntaxBuilderGeneration/SyntaxBuildableChild.swift

@@ -71,34 +71,41 @@ extension Child {
   /// If this node is a token that can't contain arbitrary text, generate a Swift
   /// `assert` statement that verifies the variable with name var_name and of type
   /// `TokenSyntax` contains one of the supported text options. Otherwise return `nil`.
-  func generateAssertStmtTextChoices(varName: String) -> String? {
+  func generateAssertStmtTextChoices(varName: String) -> ExpressibleAsExprBuildable? {
     guard type.isToken else {
       return nil
     }
 
     let choices: [String]
     if !textChoices.isEmpty {
       choices = textChoices
+    } else if tokenCanContainArbitraryText {
+      // Don't generate an assert statement if token can contain arbitrary text.
+      return nil
     } else if !tokenChoices.isEmpty {
-      let optionalChoices = tokenChoices.map { SYNTAX_TOKEN_MAP["\($0.name)Token"]?.text }
-      choices = optionalChoices.compactMap { $0 }
-      guard choices.count == optionalChoices.count else {
-        // If `nil` is in the text choices, one of the token options can contain arbitrary
-        // text. Don't generate an assert statement.
-        return nil
-      }
+      choices = tokenChoices.compactMap(\.text)
     } else {
       return nil
     }
 
-    var assertChoices: [String] = []
+    var assertChoices: [ExpressibleAsExprBuildable] = []
     if type.isOptional {
-      assertChoices.append("\(varName) == nil")
+      assertChoices.append(SequenceExpr {
+        varName
+        BinaryOperatorExpr("==")
+        NilLiteralExpr()
+      })
     }
-    let unwrap = type.isOptional ? "!" : ""
     for textChoice in choices {
-      assertChoices.append("\(varName)\(unwrap) == \"\(textChoice)\"")
+      assertChoices.append(SequenceExpr {
+        MemberAccessExpr(base: type.forceUnwrappedIfNeeded(expr: varName), name: "text")
+        BinaryOperatorExpr("==")
+        StringLiteralExpr(raw: textChoice)
+      })
+    }
+    let disjunction = ExprList(assertChoices.flatMap { [$0, BinaryOperatorExpr("||")] }.dropLast())
+    return FunctionCallExpr("assert") {
+      SequenceExpr(elements: disjunction)
     }
-    return "assert(\(assertChoices.joined(separator: " || ")))"
   }
 }

Sources/SwiftSyntaxBuilderGeneration/SyntaxBuildableType.swift

@@ -256,6 +256,15 @@ struct SyntaxBuildableType: Hashable {
     }
   }
 
+  /// Wraps a type in a force unwrap expression depending on whether `isOptional` is true.
+  func forceUnwrappedIfNeeded(expr: ExpressibleAsExprBuildable) -> ExpressibleAsExprBuildable {
+    if isOptional {
+      return ForcedValueExpr(expression: expr)
+    } else {
+      return expr
+    }
+  }
+
   /// Generate an expression that converts a variable named `varName`
   /// which is of `expressibleAs` type to an object of type `buildable`.
   func generateExprConvertParamTypeToStorageType(varName: String) -> ExpressibleAsExprBuildable {