[ASTGen] Eliminate ASTGen.ASTNode

include/swift/AST/ASTBridging.h

@@ -335,12 +335,6 @@ bool BridgedASTContext_canImport(BridgedASTContext cContext,
 // MARK: AST nodes
 //===----------------------------------------------------------------------===//
 
-enum ENUM_EXTENSIBILITY_ATTR(open) ASTNodeKind : size_t {
-  ASTNodeKindExpr,
-  ASTNodeKindStmt,
-  ASTNodeKindDecl
-};
-
 void registerBridgedDecl(BridgedStringRef bridgedClassName, SwiftMetatype metatype);
 
 struct OptionalBridgedDeclObj {
@@ -376,16 +370,53 @@ struct BridgedDeclObj {
   BRIDGED_INLINE bool Destructor_isIsolated() const;
 };
 
-struct BridgedASTNode {
-  SWIFT_NAME("raw")
-  void *_Nonnull Raw;
+enum ENUM_EXTENSIBILITY_ATTR(open) BridgedASTNodeKind : uint8_t {
+  BridgedASTNodeKindExpr,
+  BridgedASTNodeKindStmt,
+  BridgedASTNodeKindDecl
+};
+
+class BridgedASTNode {
+  intptr_t opaque;
 
-  SWIFT_NAME("kind")
-  ASTNodeKind Kind;
+  BRIDGED_INLINE BridgedASTNode(void *_Nonnull pointer,
+                                BridgedASTNodeKind kind);
+
+  void *_Nonnull getPointer() const {
+    return reinterpret_cast<void *>(opaque & ~0x7);
+  }
+
+public:
+  SWIFT_NAME("decl(_:)")
+  static BridgedASTNode createDecl(BridgedDecl d) {
+    return BridgedASTNode(d.unbridged(), BridgedASTNodeKindDecl);
+  }
+  SWIFT_NAME("stmt(_:)")
+  static BridgedASTNode createStmt(BridgedStmt s) {
+    return BridgedASTNode(s.unbridged(), BridgedASTNodeKindStmt);
+  }
+  SWIFT_NAME("expr(_:)")
+  static BridgedASTNode createExor(BridgedExpr e) {
+    return BridgedASTNode(e.unbridged(), BridgedASTNodeKindExpr);
+  }
+
+  SWIFT_UNAVAILABLE("use .kind")
+  BridgedASTNodeKind getKind() const {
+    return static_cast<BridgedASTNodeKind>(opaque & 0x7);
+  }
+
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedExpr castToExpr() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedStmt castToStmt() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDecl castToDecl() const;
 
   BRIDGED_INLINE swift::ASTNode unbridged() const;
 };
 
+SWIFT_NAME("getter:BridgedASTNode.kind(self:)")
+inline BridgedASTNodeKind BridgedASTNode_getKind(BridgedASTNode node) {
+  return node.getKind();
+}
+
 // Declare `.asDecl` on each BridgedXXXDecl type, which upcasts a wrapper for
 // a Decl subclass to a BridgedDecl.
 #define DECL(Id, Parent)                                                       \

include/swift/AST/ASTBridgingImpl.h

@@ -224,14 +224,32 @@ bool BridgedDeclObj::Destructor_isIsolated() const {
 // MARK: BridgedASTNode
 //===----------------------------------------------------------------------===//
 
+BridgedASTNode::BridgedASTNode(void *_Nonnull pointer, BridgedASTNodeKind kind)
+    : opaque(intptr_t(pointer) | kind) {
+  assert(getPointer() == pointer && getKind() == kind);
+}
+
+BridgedExpr BridgedASTNode::castToExpr() const {
+  assert(getKind() == BridgedASTNodeKindExpr);
+  return static_cast<swift::Expr *>(getPointer());
+}
+BridgedStmt BridgedASTNode::castToStmt() const {
+  assert(getKind() == BridgedASTNodeKindStmt);
+  return static_cast<swift::Stmt *>(getPointer());
+}
+BridgedDecl BridgedASTNode::castToDecl() const {
+  assert(getKind() == BridgedASTNodeKindDecl);
+  return static_cast<swift::Decl *>(getPointer());
+}
+
 swift::ASTNode BridgedASTNode::unbridged() const {
-  switch (Kind) {
-  case ASTNodeKindExpr:
-    return swift::ASTNode(static_cast<swift::Expr *>(Raw));
-  case ASTNodeKindStmt:
-    return swift::ASTNode(static_cast<swift::Stmt *>(Raw));
-  case ASTNodeKindDecl:
-    return swift::ASTNode(static_cast<swift::Decl *>(Raw));
+  switch (getKind()) {
+  case BridgedASTNodeKindExpr:
+    return castToExpr().unbridged();
+  case BridgedASTNodeKindStmt:
+    return castToStmt().unbridged();
+  case BridgedASTNodeKindDecl:
+    return castToDecl().unbridged();
   }
 }
 

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -18,44 +18,6 @@ import SwiftIfConfig
 
 import struct SwiftDiagnostics.Diagnostic
 
-enum ASTNode {
-  case decl(BridgedDecl)
-  case stmt(BridgedStmt)
-  case expr(BridgedExpr)
-
-  var castToExpr: BridgedExpr {
-    guard case .expr(let bridged) = self else {
-      fatalError("Expected an expr")
-    }
-    return bridged
-  }
-
-  var castToStmt: BridgedStmt {
-    guard case .stmt(let bridged) = self else {
-      fatalError("Expected a stmt")
-    }
-    return bridged
-  }
-
-  var castToDecl: BridgedDecl {
-    guard case .decl(let bridged) = self else {
-      fatalError("Expected a decl")
-    }
-    return bridged
-  }
-
-  var bridged: BridgedASTNode {
-    switch self {
-    case .expr(let e):
-      return BridgedASTNode(raw: e.raw, kind: .expr)
-    case .stmt(let s):
-      return BridgedASTNode(raw: s.raw, kind: .stmt)
-    case .decl(let d):
-      return BridgedASTNode(raw: d.raw, kind: .decl)
-    }
-  }
-}
-
 /// Little utility wrapper that lets us have some mutable state within
 /// immutable structs, and is therefore pretty evil.
 @propertyWrapper
@@ -95,8 +57,8 @@ struct ASTGenVisitor {
     self.configuredRegions = configuredRegions
   }
 
-  func generate(sourceFile node: SourceFileSyntax) -> [ASTNode] {
-    var out = [ASTNode]()
+  func generate(sourceFile node: SourceFileSyntax) -> [BridgedASTNode] {
+    var out = [BridgedASTNode]()
     let isTopLevel = self.declContext.isModuleScopeContext
 
     visitIfConfigElements(
@@ -105,7 +67,7 @@ struct ASTGenVisitor {
       split: Self.splitCodeBlockItemIfConfig
     ) { element in
 
-      func generateStmtOrExpr(_ body: () -> ASTNode) -> ASTNode {
+      func generateStmtOrExpr(_ body: () -> BridgedASTNode) -> BridgedASTNode {
         if !isTopLevel {
           return body()
         }
@@ -117,7 +79,7 @@ struct ASTGenVisitor {
 
         // Diagnose top level code in non-script file.
         if (!declContext.parentSourceFile.isScriptMode) {
-          switch astNode {
+          switch element.item {
           case .stmt:
             self.diagnose(.illegalTopLevelStmt(element))
           case .expr:
@@ -131,7 +93,7 @@ struct ASTGenVisitor {
         let body = BridgedBraceStmt.createImplicit(
           self.ctx,
           lBraceLoc: bodyRange.start,
-          element: astNode.bridged,
+          element: astNode,
           rBraceLoc: bodyRange.end
         )
         topLevelDecl.setBody(body: body)
@@ -147,7 +109,7 @@ struct ASTGenVisitor {
           // Hoist 'VarDecl' to the top-level.
           withBridgedSwiftClosure { ptr in
             let hoisted = ptr!.load(as: BridgedDecl.self)
-            out.append(ASTNode.decl(hoisted))
+            out.append(.decl(hoisted))
           } call: { handle in
             d.forEachDeclToHoist(handle)
           }
@@ -476,11 +438,11 @@ public func buildTopLevelASTNodes(
   switch sourceFile.pointee.syntax.as(SyntaxEnum.self) {
   case .sourceFile(let node):
     for elem in visitor.generate(sourceFile: node) {
-      callback(elem.bridged, outputContext)
+      callback(elem, outputContext)
     }
   case .memberBlockItemList(let node):
     for elem in visitor.generate(memberBlockItemList: node) {
-      callback(ASTNode.decl(elem).bridged, outputContext)
+      callback(.decl(elem), outputContext)
     }
   default:
     fatalError("invalid syntax for a source file")

lib/ASTGen/Sources/ASTGen/Decls.swift

@@ -596,7 +596,7 @@ extension ASTGenVisitor {
       let body = BridgedBraceStmt.createImplicit(
         self.ctx,
         lBraceLoc: range.start,
-        element: ASTNode.decl(decl.asDecl).bridged,
+        element: .decl(decl.asDecl),
         rBraceLoc: range.end
       )
       topLevelDecl.setBody(body: body);

lib/ASTGen/Sources/ASTGen/Literals.swift

@@ -67,7 +67,7 @@ extension ASTGenVisitor {
     }()
 
     // 'stmts' is a list of body elements of 'TapExpr' aka "appendingExpr" for the 'InterpolatedStringLiteralExpr'.
-    var stmts: [ASTNode] = []
+    var stmts: [BridgedASTNode] = []
 
     // The first element is a 'VarDecl'.
     let interpolationVar = BridgedVarDecl.createImplicitStringInterpolationVar(self.declContext)
@@ -199,7 +199,7 @@ extension ASTGenVisitor {
     let body = BridgedBraceStmt.createParsed(
       self.ctx,
       lBraceLoc: nil,
-      elements: stmts.lazy.map({ $0.bridged }).bridgedArray(in: self),
+      elements: stmts.lazy.bridgedArray(in: self),
       rBraceLoc: nil
     )
     let appendingExpr = BridgedTapExpr.create(

lib/ASTGen/Sources/ASTGen/Stmts.swift

@@ -66,7 +66,7 @@ extension ASTGenVisitor {
     }
   }
 
-  func generate(codeBlockItem node: CodeBlockItemSyntax) -> ASTNode? {
+  func generate(codeBlockItem node: CodeBlockItemSyntax) -> BridgedASTNode? {
     // TODO: Set semicolon loc.
     switch node.item {
     case .decl(let node):
@@ -89,15 +89,15 @@ extension ASTGenVisitor {
       guard let item = self.generate(codeBlockItem: codeBlockItem) else {
         return
       }
-      allItems.append(item.bridged)
+      allItems.append(item)
 
       // Hoist 'VarDecl' to the block.
-      if case .decl(let decl) = item {
+      if item.kind == .decl {
         withBridgedSwiftClosure { ptr in
           let d = ptr!.load(as: BridgedDecl.self)
-          allItems.append(ASTNode.decl(d).bridged)
+          allItems.append(.decl(d))
         } call: { handle in
-          decl.forEachDeclToHoist(handle)
+          item.castToDecl().forEachDeclToHoist(handle)
         }
       }
     }
@@ -504,7 +504,7 @@ extension ASTGenVisitor {
       }
     } body: { caseNode in
       allBridgedCases.append(
-        ASTNode.stmt(self.generate(switchCase: caseNode).asStmt).bridged
+        .stmt(self.generate(switchCase: caseNode).asStmt)
       )
     }