If a node’s child has node_choices or a collection has element_choices use an enum to ensure there are only matching child nodes

Previously, we relied on runtime checks to make sure only supported node types were used for a child with `node_choices`. If we explicitly model the set of supported nodes as an enum, we get compile-time safety.

rdar://101200491

Co-authored-by: Rintaro Ishizaki <[email protected]>

Author: ahoppen

CodeGeneration/Sources/Utils/SyntaxBuildableChild.swift

@@ -26,6 +26,18 @@ public extension Child {
     )
   }
 
+  var parameterBaseType: String {
+    if !self.nodeChoices.isEmpty {
+      return self.name
+    } else {
+      return type.parameterBaseType
+    }
+  }
+
+  var parameterType: Type {
+    return self.type.optionalWrapped(type: Type(parameterBaseType))
+  }
+
   /// If the child node has documentation associated with it, return it as single
   /// line string. Otherwise return an empty string.
   var documentation: String {

CodeGeneration/Sources/generate-swiftsyntaxbuilder/BuildableCollectionNodesFile.swift

@@ -35,25 +35,31 @@ let buildableCollectionNodesFile = SourceFile {
       inheritanceClause: TypeInheritanceClause { InheritedType(typeName: Type("ExpressibleByArrayLiteral")) }
     ) {
       // Generate initializers
-      if elementType.isBaseType {
+      if elementType.isBaseType && node.collectionElementChoices?.isEmpty ?? true {
         InitializerDecl(
           """
-          /// Creates a `\(node.type.shorthandName)` with the provided list of elements.
-          /// - Parameters:
-          ///   - elements: A list of `\(elementType.parameterType)`
           public init(_ elements: \(ArrayType(elementType: elementType.parameterType))) {
             self = \(node.type.syntaxBaseName)(elements.map { \(elementType.syntax)(fromProtocol: $0) })
           }
           """
         )
+
+        InitializerDecl(
+          """
+          public init(arrayLiteral elements: \(elementType.parameterType)...) {
+            self.init(elements)
+          }
+          """
+        )
+      } else {
+        InitializerDecl(
+          """
+          public init(arrayLiteral elements: Element...) {
+            self.init(elements)
+          }
+          """
+        )
       }
-      InitializerDecl(
-        """
-        public init(arrayLiteral elements: \(elementType.parameterType)...) {
-          self.init(elements)
-        }
-        """
-      )
     }
   }
 }

CodeGeneration/Sources/generate-swiftsyntaxbuilder/BuildableNodesFile.swift

@@ -63,7 +63,7 @@ let buildableNodesFile = SourceFile {
 }
 
 private func convertFromSyntaxProtocolToSyntaxType(child: Child) -> Expr {
-  if child.type.isBaseType {
+  if child.type.isBaseType && child.nodeChoices.isEmpty {
     return Expr(FunctionCallExpr("\(child.type.syntaxBaseName)(fromProtocol: \(child.swiftName))"))
   } else {
     return Expr(IdentifierExpr(child.swiftName))
@@ -100,7 +100,7 @@ private func createDefaultInitializer(node: Node) -> InitializerDecl {
           FunctionParameter(
             firstName: .identifier(child.swiftName),
             colon: .colon,
-            type: child.type.parameterType,
+            type: child.parameterType,
             defaultArgument: child.type.defaultInitialization.map { InitializerClause(value: $0) }
           )
         }
@@ -192,7 +192,7 @@ private func createConvenienceInitializer(node: Node) -> InitializerDecl? {
       normalParameters.append(FunctionParameter(
         firstName: .identifier(child.swiftName),
         colon: .colon,
-        type: child.type.parameterType,
+        type: child.parameterType,
         defaultArgument: child.type.defaultInitialization.map { InitializerClause(value: $0) }
       ))
     }

CodeGeneration/Sources/generate-swiftsyntaxbuilder/ResultBuildersFile.swift

@@ -24,25 +24,26 @@ let resultBuildersFile = SourceFile {
   for node in SYNTAX_NODES where node.isSyntaxCollection {
     let type = SyntaxBuildableType(syntaxKind: node.syntaxKind)
     let elementType = node.collectionElementType
+    let expressionType: Type = (node.collectionElementChoices?.isEmpty ?? true) ? elementType.parameterType : Type(MemberTypeIdentifier("\(type.buildable).Element"))
 
     StructDecl(
-      attributes: [CustomAttribute(leadingTrivia: .newline, trailingTrivia: .newline, attributeName: Type("resultBuilder"))],
+      attributes: [CustomAttribute(trailingTrivia: .newline, attributeName: Type("resultBuilder"))],
       modifiers: [DeclModifier(name: .public)],
       identifier: "\(type.syntaxKind)Builder") {
 
         TypealiasDecl(
           """
           /// The type of individual statement expressions in the transformed function,
           /// which defaults to Component if buildExpression() is not provided.
-          public typealias Expression = \(elementType.parameterType)
+          public typealias Expression = \(expressionType)
           """
         )
 
         TypealiasDecl(
           """
           /// The type of a partial result, which will be carried through all of the
           /// build methods.
-          public typealias Component = [\(elementType.parameterType)]
+          public typealias Component = [Expression]
           """
         )
 
@@ -58,7 +59,7 @@ let resultBuildersFile = SourceFile {
           """
           /// Required by every result builder to build combined results from
           /// statement blocks.
-          public static func buildBlock(_ components: Component...) -> Component {
+          public static func buildBlock(_ components: Self.Component...) -> Self.Component {
             return components.flatMap { $0 }
           }
           """
@@ -68,16 +69,28 @@ let resultBuildersFile = SourceFile {
           """
           /// If declared, provides contextual type information for statement
           /// expressions to translate them into partial results.
-          public static func buildExpression(_ expression: Expression) -> Component {
+          public static func buildExpression(_ expression: Self.Expression) -> Self.Component {
             return [expression]
           }
           """
         )
 
+        for elementChoice in node.collectionElementChoices ?? [] {
+          FunctionDecl(
+            """
+            /// If declared, provides contextual type information for statement
+            /// expressions to translate them into partial results.
+            public static func buildExpression(_ expression: \(elementChoice)) -> Self.Component {
+              return buildExpression(.init(expression))
+            }
+            """
+          )
+        }
+        
         FunctionDecl(
           """
           /// Add all the elements of `expression` to this result builder, effectively flattening them.
-          public static func buildExpression(_ expression: FinalResult) -> Component {
+          public static func buildExpression(_ expression: Self.FinalResult) -> Self.Component {
             return expression.map { $0 }
           }
           """
@@ -86,7 +99,7 @@ let resultBuildersFile = SourceFile {
         FunctionDecl(
           """
           /// Enables support for `if` statements that do not have an `else`.
-          public static func buildOptional(_ component: Component?) -> Component {
+          public static func buildOptional(_ component: Self.Component?) -> Self.Component {
             return component ?? []
           }
           """
@@ -96,7 +109,7 @@ let resultBuildersFile = SourceFile {
           """
           /// With buildEither(second:), enables support for 'if-else' and 'switch'
           /// statements by folding conditional results into a single result.
-          public static func buildEither(first component: Component) -> Component {
+          public static func buildEither(first component: Self.Component) -> Self.Component {
             return component
           }
           """
@@ -106,7 +119,7 @@ let resultBuildersFile = SourceFile {
           """
           /// With buildEither(first:), enables support for 'if-else' and 'switch'
           /// statements by folding conditional results into a single result.
-          public static func buildEither(second component: Component) -> Component {
+          public static func buildEither(second component: Self.Component) -> Self.Component {
             return component
           }
           """
@@ -116,7 +129,7 @@ let resultBuildersFile = SourceFile {
           """
           /// Enables support for 'for..in' loops by combining the
           /// results of all iterations into a single result.
-          public static func buildArray(_ components: [Component]) -> Component {
+          public static func buildArray(_ components: [Self.Component]) -> Self.Component {
             return components.flatMap { $0 }
           }
           """
@@ -127,7 +140,7 @@ let resultBuildersFile = SourceFile {
           /// If declared, this will be called on the partial result of an 'if'
           /// #available' block to allow the result builder to erase type
           /// information.
-          public static func buildLimitedAvailability(_ component: Component) -> Component {
+          public static func buildLimitedAvailability(_ component: Self.Component) -> Self.Component {
             return component
           }
           """

Examples/MigrateToNewIfLetSyntax.swift

@@ -35,7 +35,7 @@ class MigrateToNewIfLetSyntax: SyntaxRewriter {
         if index != node.conditions.count - 1 {
           binding.pattern = binding.pattern.withoutTrailingTrivia()
         }
-        conditionCopy.condition = Syntax(binding)
+        conditionCopy.condition = .optionalBinding(binding)
       }
       return conditionCopy
     }

Sources/SwiftOperators/OperatorTable+Semantics.swift

@@ -20,7 +20,7 @@ extension PrecedenceGroup {
     self.syntax = syntax
 
     for attr in syntax.groupAttributes {
-      switch attr.as(SyntaxEnum.self) {
+      switch attr {
       // Relation (lowerThan, higherThan)
       case .precedenceGroupRelation(let relation):
         let isLowerThan = relation.higherThanOrLowerThan.text == "lowerThan"
@@ -52,9 +52,6 @@ extension PrecedenceGroup {
         default:
           break
         }
-
-      default:
-        break
       }
     }
   }

Sources/SwiftOperators/SyntaxSynthesis.swift

@@ -74,12 +74,12 @@ extension PrecedenceGroup {
       TokenSyntax.identifier(name, leadingTrivia: .space)
     let leftBrace = TokenSyntax.leftBraceToken(leadingTrivia: .space)
 
-    var groupAttributes: [Syntax] = []
+    var groupAttributes: [PrecedenceGroupAttributeListSyntax.Element] = []
 
     switch associativity {
     case .left, .right:
       groupAttributes.append(
-        Syntax(
+        .init(
           PrecedenceGroupAssociativitySyntax(
             associativityKeyword:
                 .identifier(
@@ -99,7 +99,7 @@ extension PrecedenceGroup {
 
     if assignment {
       groupAttributes.append(
-        Syntax(
+        .init(
           PrecedenceGroupAssignmentSyntax(
             assignmentKeyword:
                 .identifier(
@@ -115,7 +115,7 @@ extension PrecedenceGroup {
 
     for relation in relations {
       groupAttributes.append(
-        Syntax(
+        .init(
           relation.synthesizedSyntax()
         )
       )

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -589,7 +589,7 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     }
     if node.unknownAttr?.isMissingAllTokens != false && node.label.isMissingAllTokens {
       addDiagnostic(node.statements, .allStatmentsInSwitchMustBeCoveredByCase, fixIts: [
-        FixIt(message: InsertTokenFixIt(missingNodes: [node.label]), changes: .makePresent(node.label, leadingTrivia: .newline))
+        FixIt(message: InsertTokenFixIt(missingNodes: [Syntax(node.label)]), changes: .makePresent(node.label, leadingTrivia: .newline))
       ], handledNodes: [node.label.id])
     }
     return .visitChildren

Sources/SwiftSyntax/Syntax.swift

@@ -59,13 +59,13 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
 }
 
 // Casting functions to specialized syntax nodes.
-extension Syntax {
+extension SyntaxProtocol {
   public func `is`<S: SyntaxProtocol>(_ syntaxType: S.Type) -> Bool {
     return self.as(syntaxType) != nil
   }
 
   public func `as`<S: SyntaxProtocol>(_ syntaxType: S.Type) -> S? {
-    return S.init(self)
+    return S.init(self._syntaxNode)
   }
 
   func cast<S: SyntaxProtocol>(_ syntaxType: S.Type) -> S {
@@ -625,6 +625,16 @@ public extension SyntaxProtocol {
   }
 }
 
+/// Protocol for the enums nested inside `Syntax` nodes that enumerate all the
+/// possible types a child node might have.
+public protocol SyntaxChildChoices: SyntaxProtocol {}
+
+public extension SyntaxChildChoices {
+  func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {
+    return Syntax(self).childNameForDiagnostics(index)
+  }
+}
+
 /// Sequence of tokens that are part of the provided Syntax node.
 public struct TokenSequence: Sequence {
   public struct Iterator: IteratorProtocol {