[Serialization] Fast lookup of nested types in modules with overlays (#20024)

Previously, the fast path for nested types only worked when the nested
type was defined in a Swift module or a Clang module without an
overlay; this is because it was originally designed to fix circularity
issues when merging partial modules for a single target. By having a
Swift overlay module pass through requests for nested types to the
underlying Clang module, we get the fast-path behavior in more cases.
(The one case where it /won't/ kick in is if the overlay has a nested
type that shadows a nested type from the Clang module, but that's
probably pretty rare!)

Author: jrose-apple

lib/Serialization/Deserialization.cpp

@@ -44,7 +44,7 @@ STATISTIC(NumNormalProtocolConformancesLoaded,
 STATISTIC(NumNormalProtocolConformancesCompleted,
           "# of normal protocol conformances completed");
 STATISTIC(NumNestedTypeShortcuts,
-          "# of same-module nested types resolved without lookup");
+          "# of nested types resolved without full lookup");
 
 using namespace swift;
 using namespace swift::serialization;
@@ -1444,10 +1444,6 @@ ModuleFile::resolveCrossReference(ModuleDecl *baseModule, uint32_t pathLen) {
         if (!extensionModule)
           extensionModule = baseType->getModuleContext();
 
-        // FIXME: If 'importedFromClang' is true but 'extensionModule' is an
-        // overlay module, the search below will fail and we'll fall back to
-        // the slow path.
-
         // Fault in extensions, then ask every file in the module.
         (void)baseType->getExtensions();
         TypeDecl *nestedType = nullptr;
@@ -1460,10 +1456,16 @@ ModuleFile::resolveCrossReference(ModuleDecl *baseModule, uint32_t pathLen) {
         }
 
         if (nestedType) {
-          values.clear();
-          values.push_back(nestedType);
-          ++NumNestedTypeShortcuts;
-          break;
+          SmallVector<ValueDecl *, 1> singleValueBuffer{nestedType};
+          filterValues(/*expectedTy*/Type(), extensionModule, genericSig,
+                       /*isType*/true, /*inProtocolExt*/false,
+                       importedFromClang, /*isStatic*/false, /*ctorInit*/None,
+                       singleValueBuffer);
+          if (!singleValueBuffer.empty()) {
+            values.assign({nestedType});
+            ++NumNestedTypeShortcuts;
+            break;
+          }
         }
 
         pathTrace.removeLast();

lib/Serialization/ModuleFile.cpp

@@ -1575,25 +1575,29 @@ TypeDecl *ModuleFile::lookupNestedType(Identifier name,
                                        const NominalTypeDecl *parent) {
   PrettyStackTraceModuleFile stackEntry(*this);
 
-  if (!NestedTypeDecls)
-    return nullptr;
+  if (NestedTypeDecls) {
+    auto iter = NestedTypeDecls->find(name);
+    if (iter != NestedTypeDecls->end()) {
+      for (std::pair<DeclID, DeclID> entry : *iter) {
+        assert(entry.first);
+        auto declOrOffset = Decls[entry.first - 1];
+        if (!declOrOffset.isComplete())
+          continue;
 
-  auto iter = NestedTypeDecls->find(name);
-  if (iter == NestedTypeDecls->end())
-    return nullptr;
+        Decl *decl = declOrOffset;
+        if (decl != parent)
+          continue;
+        return cast<TypeDecl>(getDecl(entry.second));
+      }
+    }
+  }
 
-  auto data = *iter;
-  for (std::pair<DeclID, DeclID> entry : data) {
-    assert(entry.first);
-    auto declOrOffset = Decls[entry.first - 1];
-    if (!declOrOffset.isComplete())
-      continue;
+  if (!ShadowedModule)
+    return nullptr;
 
-    Decl *decl = declOrOffset;
-    if (decl != parent)
-      continue;
-    return cast<TypeDecl>(getDecl(entry.second));
-  }
+  for (FileUnit *file : ShadowedModule->getFiles())
+    if (auto *nestedType = file->lookupNestedType(name, parent))
+      return nestedType;
 
   return nullptr;
 }

test/Serialization/Inputs/nested-type-with-overlay/HasOverlay.h

@@ -0,0 +1,13 @@
+struct Base {
+  int dummy;
+};
+
+struct Nested {
+  int dummy;
+} __attribute((swift_name("Base.NestedFromClang")));
+
+struct NestedAndShadowed {
+  int dummy;
+} __attribute((swift_name("Base.NestedAndShadowed")));
+
+struct NestedAndShadowed getShadowedFromClang();

test/Serialization/Inputs/nested-type-with-overlay/module.modulemap

@@ -0,0 +1,3 @@
+module HasOverlay {
+  header "HasOverlay.h"
+}

test/Serialization/Inputs/nested-type-with-overlay/overlay.swift

@@ -0,0 +1,10 @@
+@_exported import HasOverlay
+
+extension Base {
+  public struct NestedFromSwift {}
+  public struct NestedAndShadowed {
+    init(dummy: ()) {}
+  }
+}
+
+public var shadowedFromSwift = Base.NestedAndShadowed(dummy: ())

test/Serialization/Inputs/nested-type-with-overlay/verify.swift

@@ -0,0 +1,12 @@
+import main
+import HasOverlay
+
+func test() {
+  var a = getShadowedFromClang()
+  a = main.shadowedFromClang // okay
+  a = HasOverlay.shadowedFromSwift // expected-error {{cannot assign}}
+
+  var b = HasOverlay.shadowedFromSwift
+  b = main.shadowedFromSwift // okay
+  b = getShadowedFromClang() // expected-error {{cannot assign}}
+}

test/Serialization/multi-file-nested-type-circularity.swift

@@ -8,7 +8,7 @@
 // REQUIRES: asserts
 
 // CHECK: Statistics
-// CHECK: 1 Serialization - # of same-module nested types resolved without lookup
+// CHECK: 1 Serialization - # of nested types resolved without full lookup
 
 // Note the Optional here and below; this was once necessary to produce a crash.
 // Without it, the type of the parameter is initialized "early" enough to not

test/Serialization/multi-file-nested-type-extension.swift

@@ -10,7 +10,7 @@
 // REQUIRES: asserts
 
 // CHECK: Statistics
-// CHECK: 1 Serialization - # of same-module nested types resolved without lookup
+// CHECK: 1 Serialization - # of nested types resolved without full lookup
 
 // Note the Optional here and below; this was once necessary to produce a crash.
 // Without it, the type of the parameter is initialized "early" enough to not

test/Serialization/multi-file-nested-type-simple.swift

@@ -15,9 +15,9 @@
 
 // REQUIRES: asserts
 
-// CHECK: 4 Serialization - # of same-module nested types resolved without lookup
+// CHECK: 4 Serialization - # of nested types resolved without full lookup
 // DISABLED: Statistics
-// DISABLED-NOT: same-module nested types resolved without lookup
+// DISABLED-NOT: nested types resolved without full lookup
 
 public func useTypes(
   _: Outer.Inner,

test/Serialization/nested-type-with-overlay.swift

@@ -0,0 +1,22 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -o %t/HasOverlay.swiftmodule %S/Inputs/nested-type-with-overlay/overlay.swift -I %S/Inputs/nested-type-with-overlay -module-name HasOverlay
+
+// RUN: %target-swift-frontend -emit-module -o %t/main~partial.swiftmodule -primary-file %s -module-name main -I %t -I %S/Inputs/nested-type-helper
+// RUN: %target-swift-frontend -merge-modules -emit-module -o %t/main.swiftmodule %t/main~partial.swiftmodule -print-stats -module-name main -I %t -I %S/Inputs/nested-type-with-overlay 2>&1 | %FileCheck %s
+
+// RUN: %target-swift-frontend -emit-silgen %S/Inputs/nested-type-with-overlay/verify.swift -I %t -I %S/Inputs/nested-type-with-overlay -verify
+
+// REQUIRES: asserts
+
+// CHECK: 3 Serialization - # of nested types resolved without full lookup
+// Unfortunately this isn't 4 because of the shadowed nested type from Clang.
+
+import HasOverlay
+
+public func resolveNestedTypes(
+  _: Base.NestedFromClang,
+  _: Base.NestedFromSwift
+) {}
+
+public var shadowedFromClang = getShadowedFromClang()
+public var shadowedFromSwift = HasOverlay.shadowedFromSwift