Fixes for opaque return types on local functions.

- In Sema, don't traverse nested declarations while deducing the opaque return type. This would
  cause returns inside nested functions to clobber the return type of the outer function.
- In IRGen, walk the list of opaque return types we keep in the SourceFile already for type
  reconstruction, instead of trying to visit them ad-hoc as part of walking the AST, since
  IRGen doesn't normally walk the bodies of function decls directly.

Fixes rdar://problem/50459091

Author: jckarter

include/swift/AST/Module.h

@@ -1000,6 +1000,7 @@ class SourceFile final : public FileUnit {
   llvm::SetVector<TypeDecl *> LocalTypeDecls;
 
   /// The set of validated opaque return type decls in the source file.
+  llvm::SmallVector<OpaqueTypeDecl *, 4> OpaqueReturnTypes;
   llvm::StringMap<OpaqueTypeDecl *> ValidatedOpaqueReturnTypes;
   /// The set of parsed decls with opaque return types that have not yet
   /// been validated.

lib/AST/Module.cpp

@@ -1792,8 +1792,11 @@ SourceFile::lookupOpaqueResultType(StringRef MangledName,
 void SourceFile::markDeclWithOpaqueResultTypeAsValidated(ValueDecl *vd) {
   UnvalidatedDeclsWithOpaqueReturnTypes.erase(vd);
   if (auto opaqueDecl = vd->getOpaqueResultTypeDecl()) {
-    ValidatedOpaqueReturnTypes.insert(
+    auto inserted = ValidatedOpaqueReturnTypes.insert(
               {opaqueDecl->getOpaqueReturnTypeIdentifier().str(), opaqueDecl});
+    if (inserted.second) {
+      OpaqueReturnTypes.push_back(opaqueDecl);
+    }
   }
 }
 

lib/IRGen/GenDecl.cpp

@@ -436,6 +436,8 @@ void IRGenModule::emitSourceFile(SourceFile &SF) {
     emitGlobalDecl(decl);
   for (auto *localDecl : SF.LocalTypeDecls)
     emitGlobalDecl(localDecl);
+  for (auto *opaqueDecl : SF.OpaqueReturnTypes)
+    maybeEmitOpaqueTypeDecl(opaqueDecl);
 
   SF.collectLinkLibraries([this](LinkLibrary linkLib) {
       this->addLinkLibrary(linkLib);
@@ -1884,17 +1886,10 @@ void IRGenModule::emitGlobalDecl(Decl *D) {
     return;
 
   case DeclKind::Var:
-    emitAbstractStorageDecl(cast<AbstractStorageDecl>(D));
-    return;
-
   case DeclKind::Accessor:
+  case DeclKind::Func:
     // Handled in SIL.
     return;
-
-  case DeclKind::Func:
-    // The function body itself is lowered to SIL, but there may be associated
-    // decls to emit.
-    return emitFuncDecl(cast<FuncDecl>(D));
 
   case DeclKind::TopLevelCode:
     // All the top-level code will be lowered separately.
@@ -3881,12 +3876,9 @@ void IRGenModule::emitNestedTypeDecls(DeclRange members) {
       continue;
 
     case DeclKind::Func:
-      emitFuncDecl(cast<FuncDecl>(member));
-      continue;
-
     case DeclKind::Var:
     case DeclKind::Subscript:
-      emitAbstractStorageDecl(cast<AbstractStorageDecl>(member));
+      // Handled in SIL.
       continue;
 
     case DeclKind::PatternBinding:

lib/IRGen/GenStruct.cpp

@@ -891,20 +891,6 @@ void IRGenModule::maybeEmitOpaqueTypeDecl(OpaqueTypeDecl *opaque) {
   }
 }
 
-void IRGenModule::emitFuncDecl(FuncDecl *fd) {
-  // If there's an opaque return type for this function, emit its descriptor.
-  if (auto opaque = fd->getOpaqueResultTypeDecl()) {
-    maybeEmitOpaqueTypeDecl(opaque);
-  }
-}
-
-void IRGenModule::emitAbstractStorageDecl(AbstractStorageDecl *fd) {
-  // If there's an opaque return type for this function, emit its descriptor.
-  if (auto opaque = fd->getOpaqueResultTypeDecl()) {
-    maybeEmitOpaqueTypeDecl(opaque);
-  }
-}
-
 namespace {
   /// A type implementation for resilient struct types. This is not a
   /// StructTypeInfoBase at all, since we don't know anything about

lib/IRGen/IRGenModule.h

@@ -1204,8 +1204,6 @@ private:                            \
   void emitStructDecl(StructDecl *D);
   void emitClassDecl(ClassDecl *D);
   void emitExtension(ExtensionDecl *D);
-  void emitFuncDecl(FuncDecl *D);
-  void emitAbstractStorageDecl(AbstractStorageDecl *D);
   void emitOpaqueTypeDecl(OpaqueTypeDecl *D);
   void emitSILGlobalVariable(SILGlobalVariable *gv);
   void emitCoverageMapping();

lib/Sema/MiscDiagnostics.cpp

@@ -2333,6 +2333,11 @@ class OpaqueUnderlyingTypeChecker : public ASTWalker {
 
     return {true, S};
   }
+  
+  // Don't descend into nested decls.
+  bool walkToDeclPre(Decl *D) override {
+    return false;
+  }
 };
 
 } // end anonymous namespace

test/IRGen/opaque_result_type.swift

@@ -132,6 +132,16 @@ func baz<T: P & Q>(z: T) -> some P & Q {
   return z
 }
 
+// Ensure the local type's opaque descriptor gets emitted.
+// CHECK-LABEL: @"$s18opaque_result_type11localOpaqueQryF0D0L_QryFQOMQ" = 
+func localOpaque() -> some P {
+  func local() -> some P {
+    return "local"
+  }
+
+  return local()
+}
+
 public func useFoo(x: String, y: C) {
   let p = foo(x: x)
   let pa = p.poo()
@@ -181,3 +191,4 @@ struct X<T: R, U: R>: R {
     return Wrapper(wrapped: u)
   }
 }
+