[Serialization] Only allow loading modules during import resolution (swiftlang#21218)

Very early groundwork for private imports. Should not affect anything
today.

Author: jrose-apple

include/swift/Serialization/ModuleFile.h

@@ -589,7 +589,8 @@ class ModuleFile
   /// because it reads from the cursor, it is not possible to reset the cursor
   /// after reading. Nothing should ever follow an XREF record except
   /// XREF_PATH_PIECE records.
-  llvm::Expected<Decl *> resolveCrossReference(ModuleDecl *M, uint32_t pathLen);
+  llvm::Expected<Decl *> resolveCrossReference(serialization::ModuleID MID,
+                                               uint32_t pathLen);
 
   /// Populates TopLevelIDs for name lookup.
   void buildTopLevelDeclMap();
@@ -856,7 +857,7 @@ class ModuleFile
   ///
   /// If the name matches the name of the current module, a shadowed module
   /// is loaded instead.
-  ModuleDecl *getModule(ArrayRef<Identifier> name);
+  ModuleDecl *getModule(ArrayRef<Identifier> name, bool allowLoading = false);
 
   /// Returns the generic signature for the given ID.
   GenericSignature *getGenericSignature(serialization::GenericSignatureID ID);

lib/Serialization/Deserialization.cpp

@@ -122,6 +122,8 @@ const char DeclDeserializationError::ID = '\0';
 void DeclDeserializationError::anchor() {}
 const char XRefError::ID = '\0';
 void XRefError::anchor() {}
+const char XRefNonLoadedModuleError::ID = '\0';
+void XRefNonLoadedModuleError::anchor() {}
 const char OverrideError::ID = '\0';
 void OverrideError::anchor() {}
 const char TypeError::ID = '\0';
@@ -557,6 +559,11 @@ ProtocolConformanceRef ModuleFile::readConformance(
     PrettyStackTraceDecl traceTo("... to", proto);
     auto module = getModule(moduleID);
 
+    // FIXME: If the module hasn't been loaded, we probably don't want to fall
+    // back to the current module like this.
+    if (!module)
+      module = getAssociatedModule();
+
     SmallVector<ProtocolConformance *, 2> conformances;
     nominal->lookupConformance(module, proto, conformances);
     PrettyStackTraceModuleFile traceMsg(
@@ -1229,8 +1236,14 @@ static void filterValues(Type expectedTy, ModuleDecl *expectedModule,
 }
 
 Expected<Decl *>
-ModuleFile::resolveCrossReference(ModuleDecl *baseModule, uint32_t pathLen) {
+ModuleFile::resolveCrossReference(ModuleID MID, uint32_t pathLen) {
   using namespace decls_block;
+
+  ModuleDecl *baseModule = getModule(MID);
+  if (!baseModule) {
+    return llvm::make_error<XRefNonLoadedModuleError>(getIdentifier(MID));
+  }
+
   assert(baseModule && "missing dependency");
   PrettyXRefTrace pathTrace(*baseModule);
 
@@ -1550,6 +1563,10 @@ ModuleFile::resolveCrossReference(ModuleDecl *baseModule, uint32_t pathLen) {
       GenericSignatureID rawGenericSig;
       XRefExtensionPathPieceLayout::readRecord(scratch, ownerID, rawGenericSig);
       M = getModule(ownerID);
+      if (!M) {
+        return llvm::make_error<XRefError>("module is not loaded",
+                                           pathTrace, getIdentifier(ownerID));
+      }
       pathTrace.addExtension(M);
 
       // Read the generic signature, if we have one.
@@ -1933,14 +1950,15 @@ ModuleDecl *ModuleFile::getModule(ModuleID MID) {
   return getModule(getIdentifier(MID));
 }
 
-ModuleDecl *ModuleFile::getModule(ArrayRef<Identifier> name) {
+ModuleDecl *ModuleFile::getModule(ArrayRef<Identifier> name,
+                                  bool allowLoading) {
   if (name.empty() || name.front().empty())
     return getContext().TheBuiltinModule;
 
   // FIXME: duplicated from NameBinder::getModule
   if (name.size() == 1 &&
       name.front() == FileContext->getParentModule()->getName()) {
-    if (!ShadowedModule) {
+    if (!ShadowedModule && allowLoading) {
       auto importer = getContext().getClangModuleLoader();
       assert(importer && "no way to import shadowed module");
       ShadowedModule = importer->loadModule(SourceLoc(),
@@ -1953,7 +1971,10 @@ ModuleDecl *ModuleFile::getModule(ArrayRef<Identifier> name) {
   SmallVector<ImportDecl::AccessPathElement, 4> importPath;
   for (auto pathElem : name)
     importPath.push_back({ pathElem, SourceLoc() });
-  return getContext().getModule(importPath);
+
+  if (allowLoading)
+    return getContext().getModule(importPath);
+  return getContext().getLoadedModule(importPath);
 }
 
 
@@ -4038,7 +4059,7 @@ ModuleFile::getDeclCheckedImpl(DeclID DID) {
     ModuleID baseModuleID;
     uint32_t pathLen;
     decls_block::XRefLayout::readRecord(scratch, baseModuleID, pathLen);
-    auto resolved = resolveCrossReference(getModule(baseModuleID), pathLen);
+    auto resolved = resolveCrossReference(baseModuleID, pathLen);
     if (!resolved)
       return resolved;
     declOrOffset = resolved.get();

lib/Serialization/DeserializationErrors.h

@@ -282,6 +282,26 @@ class XRefError : public llvm::ErrorInfo<XRefError, DeclDeserializationError> {
   }
 };
 
+class XRefNonLoadedModuleError :
+    public llvm::ErrorInfo<XRefNonLoadedModuleError, DeclDeserializationError> {
+  friend ErrorInfo;
+  static const char ID;
+  void anchor() override;
+
+public:
+  explicit XRefNonLoadedModuleError(Identifier name) {
+    this->name = name;
+  }
+
+  void log(raw_ostream &OS) const override {
+    OS << "module '" << name << "' was not loaded";
+  }
+
+  std::error_code convertToErrorCode() const override {
+    return llvm::inconvertibleErrorCode();
+  }
+};
+
 class OverrideError : public llvm::ErrorInfo<OverrideError,
                                              DeclDeserializationError> {
 private:

lib/Serialization/ModuleFile.cpp

@@ -455,7 +455,8 @@ class ModuleFile::ExtensionTableInfo {
       StringRef moduleNameOrMangledBase;
       if (nameIDOrLength < 0) {
         const ModuleDecl *module = File.getModule(-nameIDOrLength);
-        moduleNameOrMangledBase = module->getName().str();
+        if (module)
+          moduleNameOrMangledBase = module->getName().str();
       } else {
         moduleNameOrMangledBase =
             StringRef(reinterpret_cast<const char *>(data), nameIDOrLength);
@@ -1497,7 +1498,7 @@ Status ModuleFile::associateWithFileContext(FileUnit *file,
       assert(!modulePath.back().empty() &&
              "invalid module name (submodules not yet supported)");
     }
-    auto module = getModule(modulePath);
+    auto module = getModule(modulePath, /*allowLoading*/true);
     if (!module || module->failedToLoad()) {
       // If we're missing the module we're shadowing, treat that specially.
       if (modulePath.size() == 1 &&
@@ -1735,7 +1736,7 @@ void ModuleFile::getImportDecls(SmallVectorImpl<Decl *> &Results) {
       if (AccessPath.size() == 1 && AccessPath[0].first == Ctx.StdlibModuleName)
         continue;
 
-      ModuleDecl *M = Ctx.getModule(AccessPath);
+      ModuleDecl *M = Ctx.getLoadedModule(AccessPath);
 
       auto Kind = ImportKind::Module;
       if (!ScopePath.empty()) {

test/Serialization/Recovery/Inputs/custom-modules/IndirectImport.h

@@ -0,0 +1,3 @@
+#if !BAD
+# import <IndirectlyImported.h>
+#endif

test/Serialization/Recovery/Inputs/custom-modules/IndirectlyImported.h

@@ -0,0 +1,3 @@
+struct IndirectlyImportedStruct {
+  int value;
+};

test/Serialization/Recovery/Inputs/custom-modules/module.modulemap

@@ -1,3 +1,11 @@
+module IndirectImport {
+  header "IndirectImport.h"
+  export *
+}
+module IndirectlyImported {
+  header "IndirectlyImported.h"
+}
+
 module Overrides { header "Overrides.h" }
 module ProtocolInheritance { header "ProtocolInheritance.h" }
 module RenameAcrossVersions { header "RenameAcrossVersions.h" }

test/Serialization/Recovery/indirect-import-removal.swift

@@ -0,0 +1,16 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -o %t -module-name Lib -I %S/Inputs/custom-modules %s
+
+// RUN: %target-swift-ide-test -source-filename=x -print-module -module-to-print Lib -I %t -I %S/Inputs/custom-modules | %FileCheck %s
+
+// RUN: %target-swift-ide-test -source-filename=x -print-module -module-to-print Lib -I %t -I %S/Inputs/custom-modules -Xcc -DBAD | %FileCheck -check-prefix CHECK-RECOVERY %s
+
+import IndirectImport
+
+// CHECK: func baseline()
+// CHECK-RECOVERY: func baseline()
+public func baseline() {}
+
+// CHECK: func test(_: IndirectlyImportedStruct)
+// CHECK-RECOVERY-NOT: IndirectlyImportedStruct
+public func test(_: IndirectlyImportedStruct) {}

validation-test/Serialization/Inputs/conformance-removed/module.modulemap

@@ -1,3 +1,4 @@
 module ObjCLib {
   header "ObjCLib.h"
+  export *
 }