WIP enumerate node_choices and element_choices

Author: rintaro

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)
   }
 }
 

Sources/SwiftSyntax/SyntaxCollections.swift.gyb

@@ -38,6 +38,40 @@ public protocol SyntaxCollection: SyntaxProtocol, Sequence {
 /// versions of the collection with different children.
 % end
 public struct ${node.name}: SyntaxCollection, SyntaxHashable {
+% if node.collection_element_choices:
+  public enum Element: SyntaxProtocol {
+%   for choice_name in node.collection_element_choices:
+%     choice = NODE_MAP[choice_name]
+    case `${choice.swift_syntax_kind}`(${choice.name})
+%   end
+    public var _syntaxNode: Syntax {
+      switch self {
+%   for choice_name in node.collection_element_choices:
+%     choice = NODE_MAP[choice_name]
+      case .${choice.swift_syntax_kind}(let node): return node._syntaxNode
+%   end
+      }
+    }
+    init(_ data: SyntaxData) { self.init(Syntax(data))! }
+%   for choice_name in node.collection_element_choices:
+%     choice = NODE_MAP[choice_name]
+    public init(${choice.swift_syntax_kind} node: ${choice.name}) { self = .${choice.swift_syntax_kind}(node) }
+%   end
+    public init?<Node: SyntaxProtocol>(_ syntaxNode: Node) {
+%   for choice_name in node.collection_element_choices:
+%     choice = NODE_MAP[choice_name]
+      if let node = syntaxNode.as(${choice.name}.self) {
+        self = .${choice.swift_syntax_kind}(node)
+        return
+      }
+%   end
+      return nil
+    }
+  }
+% else:
+  public typealias Element = ${node.collection_element_type}
+% end
+
   public let _syntaxNode: Syntax
 
   var layoutView: RawSyntaxLayoutView {
@@ -59,7 +93,7 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
     self._syntaxNode = Syntax(data)
   }
 
-  public init(_ children: [${node.collection_element_type}]) {
+  public init(_ children: [Element]) {
     let raw = RawSyntax.makeLayout(kind: SyntaxKind.${node.swift_syntax_kind},
       from: children.map { $0.raw }, arena: .default)
     let data = SyntaxData.forRoot(raw)
@@ -87,8 +121,7 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
   ///
   /// - Parameter syntax: The element to append.
   /// - Returns: A new `${node.name}` with that element appended to the end.
-  public func appending(
-    _ syntax: ${node.collection_element_type}) -> ${node.name} {
+  public func appending(_ syntax: Element) -> ${node.name} {
     var newLayout = layoutView.formLayoutArray()
     newLayout.append(syntax.raw)
     return replacingLayout(newLayout)
@@ -100,8 +133,7 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
   /// - Parameter syntax: The element to prepend.
   /// - Returns: A new `${node.name}` with that element prepended to the
   ///            beginning.
-  public func prepending(
-    _ syntax: ${node.collection_element_type}) -> ${node.name} {
+  public func prepending(_ syntax: Element) -> ${node.name} {
     return inserting(syntax, at: 0)
   }
 
@@ -113,8 +145,7 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
   ///   - index: The index at which to insert the element in the collection.
   ///
   /// - Returns: A new `${node.name}` with that element appended to the end.
-  public func inserting(_ syntax: ${node.collection_element_type},
-                        at index: Int) -> ${node.name} {
+  public func inserting(_ syntax: Element, at index: Int) -> ${node.name} {
     var newLayout = layoutView.formLayoutArray()
     /// Make sure the index is a valid insertion index (0 to 1 past the end)
     precondition((newLayout.startIndex...newLayout.endIndex).contains(index),
@@ -131,8 +162,7 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
   ///   - syntax: The element to replace with.
   ///
   /// - Returns: A new `${node.name}` with the new element at the provided index.
-  public func replacing(childAt index: Int,
-                        with syntax: ${node.collection_element_type}) -> ${node.name} {
+  public func replacing(childAt index: Int, with syntax: Element) -> ${node.name} {
     var newLayout = layoutView.formLayoutArray()
     /// Make sure the index is a valid index for replacing
     precondition((newLayout.startIndex..<newLayout.endIndex).contains(index),
@@ -221,7 +251,6 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
 
 /// Conformance for `${node.name}` to the `BidirectionalCollection` protocol.
 extension ${node.name}: BidirectionalCollection {
-  public typealias Element = ${node.collection_element_type}
   public typealias Index = SyntaxChildrenIndex
 
   public struct Iterator: IteratorProtocol {
@@ -233,13 +262,13 @@ extension ${node.name}: BidirectionalCollection {
       self.iterator = rawChildren.makeIterator()
     }
 
-    public mutating func next() -> ${node.collection_element_type}? {
+    public mutating func next() -> Element? {
       guard let (raw, info) = self.iterator.next() else {
         return nil
       }
       let absoluteRaw = AbsoluteRawSyntax(raw: raw!, info: info)
       let data = SyntaxData(absoluteRaw, parent: parent)
-      return ${node.collection_element_type}(data)
+      return Element(data)
     }
   }
 
@@ -274,11 +303,11 @@ extension ${node.name}: BidirectionalCollection {
     return rawChildren.distance(from: start, to: end)
   }
 
-  public subscript(position: SyntaxChildrenIndex) -> ${node.collection_element_type} {
+  public subscript(position: SyntaxChildrenIndex) -> Element {
     let (raw, info) = rawChildren[position]
     let absoluteRaw = AbsoluteRawSyntax(raw: raw!, info: info)
     let data = SyntaxData(absoluteRaw, parent: Syntax(self))
-    return ${node.collection_element_type}(data)
+    return Element(data)
   }
 }
 %   end

Sources/SwiftSyntax/SyntaxNodes.swift.gyb.template

@@ -47,6 +47,36 @@ nodes whose base kind are that specified kind.
 /// ${line}
 %     end
 public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
+%     for child in node.children:
+%       if child.node_choices:
+  public enum ${child.name}: SyntaxProtocol {
+%         for choice in child.node_choices:
+    case `${choice.swift_name}`(${choice.type_name})
+%         end
+    public var _syntaxNode: Syntax {
+      switch self {
+%         for choice in child.node_choices:
+      case .${choice.swift_name}(let node): return node._syntaxNode
+%         end
+      }
+    }
+    init(_ data: SyntaxData) { self.init(Syntax(data))! }
+%         for choice in child.node_choices:
+    public init(${choice.swift_name} node: ${choice.type_name}) { self = .${choice.swift_name}(node) }
+%         end
+    public init?<Node: SyntaxProtocol>(_ syntaxNode: Node) {
+%         for choice in child.node_choices:
+      if let node = syntaxNode.as(${choice.type_name}.self) {
+        self = .${choice.swift_name}(node)
+        return
+      }
+%         end
+      return nil
+    }
+  }
+
+%       end
+%     end
 %     # ==============
 %     # Initialization
 %     # ==============
@@ -71,7 +101,7 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
   public init(
 %     for (index, child) in enumerate(node.children):
 %       comma = ',' if index != len(node.children) - 1 else ''
-%       param_type = child.type_name
+%       param_type = child.name if child.node_choices else child.type_name
 %       if child.is_optional:
 %         param_type = param_type + "?"
 %       if child.is_unexpected_nodes():
@@ -102,21 +132,21 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
 %       # Children properties
 %       # ===================
 %
-%       ret_type = child.type_name
-%       if child.is_optional:
-%         ret_type += '?'
-%       end
+%       child_type = child.type_name
+%       if child.node_choices:
+%         child_type = child.name
+%       optional_mark = "?" if child.is_optional else ""
 
 %       for line in dedented_lines(child.description):
   /// ${line}
 %       end
-  public var ${child.swift_name}: ${ret_type} {
+  public var ${child.swift_name}: ${child_type}${optional_mark} {
     get {
       let childData = data.child(at: ${idx}, parent: Syntax(self))
 %       if child.is_optional:
       if childData == nil { return nil }
 %       end
-      return ${child.type_name}(childData!)
+      return ${child_type}(childData!)
     }
     set(value) {
       self = with${child.name}(value)
@@ -161,7 +191,7 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
   /// - param newChild: The new `${child.swift_name}` to replace the node's
   ///                   current `${child.swift_name}`, if present.
   public func with${child.name}(
-    _ newChild: ${child.type_name}?) -> ${node.name} {
+    _ newChild: ${child_type}?) -> ${node.name} {
 %       if child.is_optional:
     let raw = newChild?.raw
 %       else: