[Parse] Post-process parsed AST to register decls with opaque result

Move the logic out of `Paser` so `ASTGen` can use it.

Author: rintaro

include/swift/AST/ASTBridgingImpl.h

@@ -107,7 +107,6 @@ BridgedASTContext BridgedASTContext_fromRaw(void * _Nonnull ptr) {
   return *static_cast<swift::ASTContext *>(ptr);
 }
 
-BRIDGED_INLINE
 void *_Nullable BridgedASTContext_allocate(BridgedASTContext bridged,
                                            size_t size, size_t alignment) {
   return bridged.unbridged().Allocate(size, alignment);
@@ -157,7 +156,7 @@ BridgedDeclContext_getParentSourceFile(BridgedDeclContext dc) {
 // MARK: BridgedSoureFile
 //===----------------------------------------------------------------------===//
 
-BRIDGED_INLINE bool BridgedSourceFile_isScriptMode(BridgedSourceFile sf) {
+bool BridgedSourceFile_isScriptMode(BridgedSourceFile sf) {
   return sf.unbridged()->isScriptMode();
 }
 

lib/AST/ASTDumper.cpp

@@ -2037,15 +2037,15 @@ namespace {
     void visitTypeAliasDecl(TypeAliasDecl *TAD, Label label) {
       printCommon(TAD, "typealias", label);
 
-      printWhereRequirements(TAD);
-      printAttributes(TAD);
-      
       if (auto underlying = TAD->getCachedUnderlyingType()) {
         printTypeField(underlying, Label::always("type"));
       } else {
         printRec(TAD->getUnderlyingTypeRepr(), Label::always("type_repr"));
       }
 
+      printWhereRequirements(TAD);
+      printAttributes(TAD);
+
       printFoot();
     }
 

lib/Parse/ParseDecl.cpp

@@ -8491,24 +8491,14 @@ Parser::parseDeclVar(ParseDeclOptions Flags,
 
       pattern = patternRes.get();
     }
-    
-    bool hasOpaqueReturnTy = false;
-    if (auto typedPattern = dyn_cast<TypedPattern>(pattern)) {
-      hasOpaqueReturnTy = typedPattern->getTypeRepr()->hasOpaque();
-    }
-    auto sf = CurDeclContext->getOutermostParentSourceFile();
-    
+
     // Configure all vars with attributes, 'static' and parent pattern.
     pattern->forEachVariable([&](VarDecl *VD) {
       VD->setStatic(StaticLoc.isValid());
       VD->attachParsedAttrs(Attributes);
       VD->setTopLevelGlobal(topLevelDecl);
 
       Decls.push_back(VD);
-      if (hasOpaqueReturnTy && sf && !InInactiveClauseEnvironment
-          && !InFreestandingMacroArgument) {
-        sf->addUnvalidatedDeclWithOpaqueResultType(VD);
-      }
     });
 
     // Remember this pattern/init pair for our ultimate PatternBindingDecl. The
@@ -8788,14 +8778,6 @@ ParserResult<FuncDecl> Parser::parseDeclFunc(SourceLoc StaticLoc,
                               BodyParams, FuncRetTy,
                               CurDeclContext);
 
-  // Let the source file track the opaque return type mapping, if any.
-  if (FuncRetTy && FuncRetTy->hasOpaque() &&
-      !InInactiveClauseEnvironment && !InFreestandingMacroArgument) {
-    if (auto sf = CurDeclContext->getOutermostParentSourceFile()) {
-      sf->addUnvalidatedDeclWithOpaqueResultType(FD);
-    }
-  }
-  
   // Parse a 'where' clause if present.
   if (Tok.is(tok::kw_where)) {
     ContextChange CC(*this, FD);
@@ -9682,14 +9664,6 @@ Parser::parseDeclSubscript(SourceLoc StaticLoc,
       Context, StaticLoc, StaticSpelling, SubscriptLoc, Indices.get(), ArrowLoc,
       ElementTy.get(), CurDeclContext, GenericParams);
   Subscript->attachParsedAttrs(Attributes);
-  
-  // Let the source file track the opaque return type mapping, if any.
-  if (ElementTy.get() && ElementTy.get()->hasOpaque() &&
-      !InInactiveClauseEnvironment && !InFreestandingMacroArgument) {
-    if (auto sf = CurDeclContext->getOutermostParentSourceFile()) {
-      sf->addUnvalidatedDeclWithOpaqueResultType(Subscript);
-    }
-  }
 
   DefaultArgs.setFunctionContext(Subscript, Subscript->getIndices());
 

lib/Parse/ParseRequests.cpp

@@ -15,11 +15,13 @@
 //===----------------------------------------------------------------------===//
 #include "swift/AST/ParseRequests.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/AvailabilitySpec.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/TypeRepr.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Bridging/ASTGen.h"
@@ -47,6 +49,82 @@ void swift::simple_display(llvm::raw_ostream &out,
   simple_display(out, value.members);
 }
 
+namespace {
+
+/// Find 'ValueDecl' with opaque type results and let \p SF track the opaque
+/// return type mapping.
+// FIXME: This should be done in Sema or AST?
+class OpaqueResultCollector : public ASTWalker {
+  SourceFile &SF;
+
+public:
+  OpaqueResultCollector(SourceFile &SF) : SF(SF) {}
+
+  void handle(ValueDecl *VD) const {
+    TypeRepr *tyR = nullptr;
+    if (auto *FD = dyn_cast<FuncDecl>(VD)) {
+      if (isa<AccessorDecl>(VD))
+        return;
+      tyR = FD->getResultTypeRepr();
+    } else if (auto *VarD = dyn_cast<VarDecl>(VD)) {
+      if (!VarD->getParentPatternBinding())
+        return;
+      tyR = VarD->getTypeReprOrParentPatternTypeRepr();
+    } else if (auto *SD = dyn_cast<SubscriptDecl>(VD)) {
+      tyR = SD->getElementTypeRepr();
+    }
+
+    if (!tyR || !tyR->hasOpaque())
+      return;
+
+    SF.addUnvalidatedDeclWithOpaqueResultType(VD);
+  }
+
+  PreWalkAction walkToDeclPre(Decl *D) override {
+    if (auto *VD = dyn_cast<ValueDecl>(D))
+      handle(VD);
+
+    if (auto *IDC = dyn_cast<IterableDeclContext>(D)) {
+      // Only walk into cached parsed members to avoid invoking lazy parsing.
+      ParseMembersRequest req(IDC);
+      if (SF.getASTContext().evaluator.hasCachedResult(req)) {
+        auto memberResult = evaluateOrFatal(SF.getASTContext().evaluator, req);
+        for (auto *D : memberResult.members)
+          D->walk(*this);
+      }
+      return Action::SkipChildren();
+    }
+
+    return Action::Continue();
+  }
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::None;
+  }
+
+  PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
+    return Action::SkipNode();
+  }
+};
+
+void registerDeclWithOpaqueResultType(SourceFile *SF, ArrayRef<ASTNode> items) {
+  if (!SF)
+    return;
+  OpaqueResultCollector walker(*SF);
+  for (const ASTNode &item : items)
+    const_cast<ASTNode &>(item).walk(walker);
+}
+
+void registerDeclWithOpaqueResultType(SourceFile *SF, ArrayRef<Decl *> items) {
+  if (!SF)
+    return;
+  OpaqueResultCollector walker(*SF);
+  for (Decl *item : items)
+    item->walk(walker);
+}
+
+} // namespace
+
 FingerprintAndMembers
 ParseMembersRequest::evaluate(Evaluator &evaluator,
                               IterableDeclContext *idc) const {
@@ -82,6 +160,9 @@ ParseMembersRequest::evaluate(Evaluator &evaluator,
   auto declsAndHash = parser.parseDeclListDelayed(idc);
   FingerprintAndMembers fingerprintAndMembers = {declsAndHash.second,
                                                  declsAndHash.first};
+
+  registerDeclWithOpaqueResultType(sf->getOutermostParentSourceFile(),
+                                   fingerprintAndMembers.members);
   return FingerprintAndMembers{
       fingerprintAndMembers.fingerprint,
       ctx.AllocateCopy(llvm::ArrayRef(fingerprintAndMembers.members))};
@@ -134,6 +215,8 @@ ParseAbstractFunctionBodyRequest::evaluate(Evaluator &evaluator,
     Parser parser(bufferID, *sf, /*SIL*/ nullptr);
     auto result = parser.parseAbstractFunctionBodyDelayed(afd);
     afd->setBodyKind(BodyKind::Parsed);
+    registerDeclWithOpaqueResultType(afd->getOutermostParentSourceFile(),
+                                     result.getBody()->getElements());
     return result;
   }
   }
@@ -312,6 +395,9 @@ SourceFileParsingResult parseSourceFileViaASTGen(SourceFile &SF) {
     fp = swift_ASTGen_getSourceFileFingerprint(exportedSourceFile, Ctx)
              .unbridged();
 
+  registerDeclWithOpaqueResultType(declContext->getOutermostParentSourceFile(),
+                                   items);
+
   return SourceFileParsingResult{/*TopLevelItems=*/Ctx.AllocateCopy(items),
                                  /*CollectedTokens=*/std::nullopt,
                                  /*Fingerprint=*/fp};
@@ -444,6 +530,9 @@ SourceFileParsingResult parseSourceFile(SourceFile &SF) {
   if (parser.CurrentTokenHash)
     fp = Fingerprint(std::move(*parser.CurrentTokenHash));
 
+  registerDeclWithOpaqueResultType(
+      parser.CurDeclContext->getOutermostParentSourceFile(), items);
+
   return SourceFileParsingResult{ctx.AllocateCopy(items), tokensRef, fp};
 }
 

test/ASTGen/attrs.swift

@@ -199,6 +199,5 @@ protocol OpProto {}
 struct OpStruct: OpProto {}
 struct OpTest {
   func opResult() -> some OpProto { OpStruct() }
-  // FIXME: Implement SF->addUnvalidatedDeclWithOpaqueResultType() in ASTGen
-  // typealias Result = @_opaqueReturnTypeOf("$s6ASTGen6OpTestV8opResultQryF", 0) __
+  typealias Result = @_opaqueReturnTypeOf("$s6ASTGen6OpTestV8opResultQryF", 0) __
 }