Merge pull request #77352 from rintaro/astgen-macroexpansion

[ASTGen] Support macro expanded source files

Author: rintaro

include/swift/AST/ASTBridging.h

@@ -504,7 +504,19 @@ void BridgedDiagnostic_finish(BridgedDiagnostic cDiag);
 //===----------------------------------------------------------------------===//
 
 SWIFT_NAME("getter:BridgedDeclContext.isLocalContext(self:)")
-bool BridgedDeclContext_isLocalContext(BridgedDeclContext cDeclContext);
+BRIDGED_INLINE bool
+BridgedDeclContext_isLocalContext(BridgedDeclContext cDeclContext);
+
+SWIFT_NAME("getter:BridgedDeclContext.isTypeContext(self:)")
+BRIDGED_INLINE bool BridgedDeclContext_isTypeContext(BridgedDeclContext dc);
+
+SWIFT_NAME("getter:BridgedDeclContext.isModuleScopeContext(self:)")
+BRIDGED_INLINE bool
+BridgedDeclContext_isModuleScopeContext(BridgedDeclContext dc);
+
+SWIFT_NAME("getter:BridgedDeclContext.astContext(self:)")
+BRIDGED_INLINE BridgedASTContext
+BridgedDeclContext_getASTContext(BridgedDeclContext dc);
 
 SWIFT_NAME("BridgedPatternBindingInitializer.create(declContext:)")
 BridgedPatternBindingInitializer
@@ -1365,6 +1377,12 @@ BridgedPackExpansionExpr_createParsed(BridgedASTContext cContext,
                                       BridgedSourceLoc cRepeatLoc,
                                       BridgedExpr cPatternExpr);
 
+SWIFT_NAME("BridgedParenExpr.createParsed(_:leftParenLoc:expr:rightParenLoc:)")
+BridgedParenExpr BridgedParenExpr_createParsed(BridgedASTContext cContext,
+                                               BridgedSourceLoc cLParen,
+                                               BridgedExpr cExpr,
+                                               BridgedSourceLoc cRParen);
+
 SWIFT_NAME("BridgedPostfixUnaryExpr.createParsed(_:operator:operand:)")
 BridgedPostfixUnaryExpr
 BridgedPostfixUnaryExpr_createParsed(BridgedASTContext cContext,

include/swift/AST/ASTBridgingImpl.h

@@ -18,9 +18,10 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/IfConfigClauseRangeInfo.h"
-#include "swift/AST/Stmt.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/ProtocolConformanceRef.h"
+#include "swift/AST/SourceFile.h"
+#include "swift/AST/Stmt.h"
 #include "swift/Basic/Assertions.h"
 
 SWIFT_BEGIN_NULLABILITY_ANNOTATIONS
@@ -97,6 +98,26 @@ BridgedStringRef BridgedASTContext_allocateCopyString(BridgedASTContext bridged,
   return bridged.unbridged().AllocateCopy(cStr.unbridged());
 }
 
+//===----------------------------------------------------------------------===//
+// MARK: BridgedDeclContext
+//===----------------------------------------------------------------------===//
+
+bool BridgedDeclContext_isLocalContext(BridgedDeclContext dc) {
+  return dc.unbridged()->isLocalContext();
+}
+
+bool BridgedDeclContext_isTypeContext(BridgedDeclContext dc) {
+  return dc.unbridged()->isTypeContext();
+}
+
+bool BridgedDeclContext_isModuleScopeContext(BridgedDeclContext dc) {
+  return dc.unbridged()->isModuleScopeContext();
+}
+
+BridgedASTContext BridgedDeclContext_getASTContext(BridgedDeclContext dc) {
+  return dc.unbridged()->getASTContext();
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: BridgedDeclObj
 //===----------------------------------------------------------------------===//

include/swift/AST/SourceFile.h

@@ -595,6 +595,8 @@ class SourceFile final : public FileUnit {
     return BufferID;
   }
 
+  const GeneratedSourceInfo *getGeneratedSourceFileInfo() const;
+
   /// For source files created to hold the source code created by expanding
   /// a macro, this is the AST node that describes the macro expansion.
   ///
@@ -641,6 +643,9 @@ class SourceFile final : public FileUnit {
   /// Otherwise, return an empty string.
   StringRef getFilename() const;
 
+  /// Retrieve the source text buffer.
+  StringRef getBuffer() const;
+
   /// Retrieve the scope that describes this source file.
   ASTScope &getScope();
 

include/swift/Basic/BasicBridging.h

@@ -432,6 +432,23 @@ class BridgedSwiftVersion {
   unsigned getMinor() const { return Minor; }
 };
 
+//===----------------------------------------------------------------------===//
+// MARK: GeneratedSourceInfo
+//===----------------------------------------------------------------------===//
+
+enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedGeneratedSourceFileKind {
+#define MACRO_ROLE(Name, Description)                                          \
+  BridgedGeneratedSourceFileKind##Name##MacroExpansion,
+#include "swift/Basic/MacroRoles.def"
+#undef MACRO_ROLE
+
+  BridgedGeneratedSourceFileKindReplacedFunctionBody,
+  BridgedGeneratedSourceFileKindPrettyPrinted,
+  BridgedGeneratedSourceFileKindDefaultArgument,
+
+  BridgedGeneratedSourceFileKindNone,
+};
+
 SWIFT_END_NULLABILITY_ANNOTATIONS
 
 #ifndef PURE_BRIDGING_MODE

include/swift/Basic/Features.def

@@ -299,11 +299,6 @@ EXPERIMENTAL_FEATURE(ParserRoundTrip, false)
 /// Swift parser.
 EXPERIMENTAL_FEATURE(ParserValidation, false)
 
-/// Whether to emit diagnostics from the new parser first, and only emit
-/// diagnostics from the existing parser when there are none from the new
-/// parser.
-EXPERIMENTAL_FEATURE(ParserDiagnostics, false)
-
 /// Enables implicit some while also enabling existential `any`
 EXPERIMENTAL_FEATURE(ImplicitSome, false)
 

include/swift/Bridging/ASTGen.h

@@ -37,11 +37,11 @@ void swift_ASTGen_renderQueuedDiagnostics(
 
 // FIXME: Hack because we cannot easily get to the already-parsed source
 // file from here. Fix this egregious oversight!
-void *_Nullable swift_ASTGen_parseSourceFile(const char *_Nonnull buffer,
-                                             size_t bufferLength,
-                                             const char *_Nonnull moduleName,
-                                             const char *_Nonnull filename,
-                                             void *_Nullable ctx);
+void *_Nullable swift_ASTGen_parseSourceFile(BridgedStringRef buffer,
+                                             BridgedStringRef moduleName,
+                                             BridgedStringRef filename,
+                                             void *_Nullable declContextPtr,
+                                             BridgedGeneratedSourceFileKind);
 void swift_ASTGen_destroySourceFile(void *_Nonnull sourceFile);
 
 /// Check whether the given source file round-trips correctly. Returns 0 if
@@ -60,7 +60,7 @@ void swift_ASTGen_buildTopLevelASTNodes(
     BridgedDiagnosticEngine diagEngine, void *_Nonnull sourceFile,
     BridgedDeclContext declContext, BridgedASTContext astContext,
     BridgedLegacyParser legacyParser, void *_Nonnull outputContext,
-    void (*_Nonnull)(void *_Nonnull, void *_Nonnull));
+    void (*_Nonnull)(BridgedASTNode, void *_Nonnull));
 
 void swift_ASTGen_freeBridgedString(BridgedStringRef);
 

include/swift/Parse/Parser.h

@@ -930,12 +930,6 @@ class Parser {
   /// Each item will be a declaration, statement, or expression.
   void parseTopLevelItems(SmallVectorImpl<ASTNode> &items);
 
-  /// Parse the source file via the Swift Parser using the ASTGen library.
-  void
-  parseSourceFileViaASTGen(SmallVectorImpl<ASTNode> &items,
-                           std::optional<DiagnosticTransaction> &transaction,
-                           bool suppressDiagnostics = false);
-
   /// Parse the top-level SIL decls into the SIL module.
   /// \returns \c true if there was a parsing error.
   bool parseTopLevelSIL();

lib/AST/Bridging/DeclContextBridging.cpp

@@ -20,10 +20,6 @@ using namespace swift;
 // MARK: DeclContexts
 //===----------------------------------------------------------------------===//
 
-bool BridgedDeclContext_isLocalContext(BridgedDeclContext cDeclContext) {
-  return cDeclContext.unbridged()->isLocalContext();
-}
-
 BridgedPatternBindingInitializer
 BridgedPatternBindingInitializer_create(BridgedDeclContext cDeclContext) {
   return PatternBindingInitializer::create(cDeclContext.unbridged());

lib/AST/Bridging/ExprBridging.cpp

@@ -383,6 +383,14 @@ BridgedRegexLiteralExpr_createParsed(BridgedASTContext cContext,
                                         cRegexText.unbridged());
 }
 
+BridgedParenExpr BridgedParenExpr_createParsed(BridgedASTContext cContext,
+                                               BridgedSourceLoc cLParen,
+                                               BridgedExpr cExpr,
+                                               BridgedSourceLoc cRParen) {
+  ASTContext &context = cContext.unbridged();
+  return new (context)
+      ParenExpr(cLParen.unbridged(), cExpr.unbridged(), cRParen.unbridged());
+}
 BridgedSequenceExpr BridgedSequenceExpr_createParsed(BridgedASTContext cContext,
                                                      BridgedArrayRef exprs) {
   return SequenceExpr::create(cContext.unbridged(), exprs.unbridged<Expr *>());

lib/AST/DiagnosticBridge.cpp

@@ -114,8 +114,8 @@ void *DiagnosticBridge::getSourceFileSyntax(SourceManager &sourceMgr,
 
   auto bufferContents = sourceMgr.getEntireTextForBuffer(bufferID);
   auto sourceFile = swift_ASTGen_parseSourceFile(
-      bufferContents.data(), bufferContents.size(), "module",
-      displayName.str().c_str(), /*ctx*/ nullptr);
+      bufferContents, StringRef{"module"}, displayName,
+      /*declContextPtr=*/nullptr, BridgedGeneratedSourceFileKindNone);
 
   sourceFileSyntax[{&sourceMgr, bufferID}] = sourceFile;
   return sourceFile;

lib/AST/Module.cpp

@@ -1073,6 +1073,10 @@ void SourceFile::lookupClassMembers(ImportPath::Access accessPath,
   cache.lookupClassMembers(accessPath, consumer);
 }
 
+const GeneratedSourceInfo *SourceFile::getGeneratedSourceFileInfo() const {
+  return getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
+}
+
 ASTNode SourceFile::getMacroExpansion() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;
@@ -1084,28 +1088,22 @@ SourceRange SourceFile::getMacroInsertionRange() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return SourceRange();
 
-  auto generatedInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
-  auto origRange = generatedInfo.originalSourceRange;
+  auto origRange = getGeneratedSourceFileInfo()->originalSourceRange;
   return {origRange.getStart(), origRange.getEnd()};
 }
 
 CustomAttr *SourceFile::getAttachedMacroAttribute() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;
 
-  auto genInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
-  return genInfo.attachedMacroCustomAttr;
+  return getGeneratedSourceFileInfo()->attachedMacroCustomAttr;
 }
 
 std::optional<MacroRole> SourceFile::getFulfilledMacroRole() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return std::nullopt;
 
-  auto genInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
-  switch (genInfo.kind) {
+  switch (getGeneratedSourceFileInfo()->kind) {
 #define MACRO_ROLE(Name, Description)               \
   case GeneratedSourceInfo::Name##MacroExpansion: \
     return MacroRole::Name;
@@ -1123,9 +1121,7 @@ SourceFile *SourceFile::getEnclosingSourceFile() const {
       Kind != SourceFileKind::DefaultArgument)
     return nullptr;
 
-  auto genInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
-  auto sourceLoc = genInfo.originalSourceRange.getStart();
+  auto sourceLoc = getGeneratedSourceFileInfo()->originalSourceRange.getStart();
   return getParentModule()->getSourceFileContainingLocation(sourceLoc);
 }
 
@@ -1134,9 +1130,7 @@ ASTNode SourceFile::getNodeInEnclosingSourceFile() const {
       Kind != SourceFileKind::DefaultArgument)
     return nullptr;
 
-  auto genInfo =
-      *getASTContext().SourceMgr.getGeneratedSourceInfo(getBufferID());
-  return ASTNode::getFromOpaqueValue(genInfo.astNode);
+  return ASTNode::getFromOpaqueValue(getGeneratedSourceFileInfo()->astNode);
 }
 
 void ModuleDecl::lookupClassMember(ImportPath::Access accessPath,
@@ -3573,6 +3567,11 @@ StringRef SourceFile::getFilename() const {
   return SM.getIdentifierForBuffer(BufferID);
 }
 
+StringRef SourceFile::getBuffer() const {
+  SourceManager &SM = getASTContext().SourceMgr;
+  return SM.getEntireTextForBuffer(BufferID);
+}
+
 ASTScope &SourceFile::getScope() {
   if (!Scope)
     Scope = new (getASTContext()) ASTScope(this);

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -101,52 +101,59 @@ struct ASTGenVisitor {
     self.legacyParse = legacyParser
   }
 
-  func generate(sourceFile node: SourceFileSyntax) -> [BridgedDecl] {
-    var out = [BridgedDecl]()
+  func generate(sourceFile node: SourceFileSyntax) -> [ASTNode] {
+    var out = [ASTNode]()
+    let isTopLevel = self.declContext.isModuleScopeContext
 
     visitIfConfigElements(
       node.statements,
       of: CodeBlockItemSyntax.self,
       split: Self.splitCodeBlockItemIfConfig
     ) { element in
-      let loc = self.generateSourceLoc(element)
+      let astNode = generate(codeBlockItem: element)
+      if !isTopLevel {
+        out.append(astNode)
+        return
+      }
 
-      func endLoc() -> BridgedSourceLoc {
+      func getRange() -> (start: BridgedSourceLoc, end: BridgedSourceLoc) {
+        let loc = self.generateSourceLoc(element)
         if let endTok = element.lastToken(viewMode: .sourceAccurate) {
           switch endTok.parent?.kind {
           case .stringLiteralExpr, .regexLiteralExpr:
             // string/regex literal are single token in AST.
-            return self.generateSourceLoc(endTok.parent)
+            return (loc, self.generateSourceLoc(endTok.parent))
           default:
-            return self.generateSourceLoc(endTok)
+            return (loc, self.generateSourceLoc(endTok))
           }
         } else {
-          return loc
+          return (loc, loc)
         }
       }
 
-      let swiftASTNodes = generate(codeBlockItem: element)
-      switch swiftASTNodes {
+      switch astNode {
       case .decl(let d):
-        out.append(d)
+        out.append(.decl(d))
       case .stmt(let s):
+        let range = getRange()
         let topLevelDecl = BridgedTopLevelCodeDecl.createParsed(
           self.ctx,
           declContext: self.declContext,
-          startLoc: loc,
+          startLoc: range.start,
           stmt: s,
-          endLoc: endLoc()
+          endLoc: range.end
         )
-        out.append(topLevelDecl.asDecl)
+        out.append(.decl(topLevelDecl.asDecl))
       case .expr(let e):
+        let range = getRange()
         let topLevelDecl = BridgedTopLevelCodeDecl.createParsed(
           self.ctx,
           declContext: self.declContext,
-          startLoc: loc,
+          startLoc: range.start,
           expr: e,
-          endLoc: endLoc()
+          endLoc: range.end
         )
-        out.append(topLevelDecl.asDecl)
+        out.append(.decl(topLevelDecl.asDecl))
       }
     }
 
@@ -442,7 +449,7 @@ public func buildTopLevelASTNodes(
   ctx: BridgedASTContext,
   legacyParser: BridgedLegacyParser,
   outputContext: UnsafeMutableRawPointer,
-  callback: @convention(c) (UnsafeMutableRawPointer, UnsafeMutableRawPointer) -> Void
+  callback: @convention(c) (BridgedASTNode, UnsafeMutableRawPointer) -> Void
 ) {
   let sourceFile = sourceFilePtr.assumingMemoryBound(to: ExportedSourceFile.self)
   let visitor = ASTGenVisitor(
@@ -454,8 +461,18 @@ public func buildTopLevelASTNodes(
     legacyParser: legacyParser
   )
 
-  visitor.generate(sourceFile: sourceFile.pointee.syntax)
-    .forEach { callback($0.raw, outputContext) }
+  switch sourceFile.pointee.syntax.as(SyntaxEnum.self) {
+  case .sourceFile(let node):
+    for elem in visitor.generate(sourceFile: node) {
+      callback(elem.bridged, outputContext)
+    }
+  case .memberBlockItemList(let node):
+    for elem in visitor.generate(memberBlockItemList: node) {
+      callback(ASTNode.decl(elem).bridged, outputContext)
+    }
+  default:
+    fatalError("invalid syntax for a source file")
+  }
 
   // Diagnose any errors from evaluating #ifs.
   visitor.diagnoseAll(visitor.configuredRegions.diagnostics)