[astgen] Use ASTNode to implement BraceStmt correctly.

Author: zoecarver

include/swift/AST/CASTBridging.h

@@ -128,8 +128,13 @@ void *SwiftVarDecl_create(void *ctx, BridgedIdentifier _Nullable name,
 void *IfStmt_create(void *ctx, void *ifLoc, void *cond, void *_Nullable then, void *_Nullable elseLoc,
                     void *_Nullable elseStmt);
 
-void *BraceStmt_createExpr(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc);
-void *BraceStmt_createStmt(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc);
+struct ASTNodeBridged {
+  void *ptr;
+  _Bool isExpr; // Must be expr or stmt.
+};
+
+void *BraceStmt_create(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc);
+void *BraceStmt_create(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc);
 
 void *BridgedSourceLoc_advanced(void *loc, long len);
 

lib/AST/CASTBridging.cpp

@@ -160,17 +160,16 @@ void *IfStmt_create(void *ctx, void *ifLoc, void *cond, void *_Nullable then, vo
       getSourceLocFromPointer(elseLoc), (Stmt *)elseStmt, None, Context);
 }
 
-void *BraceStmt_createExpr(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc) {
-  ASTContext &Context = *static_cast<ASTContext *>(ctx);
-  return BraceStmt::create(Context, getSourceLocFromPointer(lbloc),
-                           getArrayRef<ASTNode>(elements),
-                           getSourceLocFromPointer(rbloc));
-}
-
-void *BraceStmt_createStmt(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc) {
+void *BraceStmt_create(void *ctx, void *lbloc, BridgedArrayRef elements, void *rbloc) {
   llvm::SmallVector<ASTNode, 6> nodes;
-  for (auto stmt : getArrayRef<Stmt *>(elements)) {
-    nodes.push_back(stmt);
+  for (auto node : getArrayRef<ASTNodeBridged>(elements)) {
+    if (node.isExpr) {
+      auto expr = (Expr *)node.ptr;
+      nodes.push_back(expr);
+    } else {
+      auto stmt = (Stmt *)node.ptr;
+      nodes.push_back(stmt);
+    }
   }
 
   ASTContext &Context = *static_cast<ASTContext *>(ctx);

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -35,6 +35,17 @@ enum ASTNode {
       return ptr
     }
   }
+  
+  func bridged() -> ASTNodeBridged {
+    switch self {
+    case .expr(let e):
+      return ASTNodeBridged(ptr: e, isExpr: true)
+    case .stmt(let s):
+      return ASTNodeBridged(ptr: s, isExpr: false)
+    default:
+      fatalError("Must be expr or stmt.")
+    }
+  }
 }
 
 /// Little utility wrapper that lets us have some mutable state within

lib/ASTGen/Sources/ASTGen/Exprs.swift

@@ -5,11 +5,11 @@ import CASTBridging
 
 extension ASTGenVisitor {
   public func visit(_ node: ClosureExprSyntax) -> ASTNode {
-    let statements = node.statements.map(self.visit)
+    let statements = node.statements.map(self.visit).map { $0.bridged() }
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     let body = statements.withBridgedArrayRef { ref in
-        BraceStmt_createExpr(ctx, loc, ref, loc)
+        BraceStmt_create(ctx, loc, ref, loc)
     }
 
     return .expr(ClosureExpr_create(ctx, body, declContext))

lib/ASTGen/Sources/ASTGen/Stmts.swift

@@ -5,11 +5,11 @@ import CASTBridging
 
 extension ASTGenVisitor {
   public func visit(_ node: CodeBlockSyntax) -> ASTNode {
-    let statements = node.statements.map(self.visit)
+    let statements = node.statements.map(self.visit).map { $0.bridged() }
     let loc = self.base.advanced(by: node.position.utf8Offset).raw
 
     return .stmt(statements.withBridgedArrayRef { ref in
-      BraceStmt_createStmt(ctx, loc, ref, loc)
+      BraceStmt_create(ctx, loc, ref, loc)
     })
   }