[Frontend] Bail early if the stdlib is missing

Rather than trying to continue the compilation
with an empty main module, let's bail out early if
we expect an implicit stdlib import and fail to
load in the stdlib.

Author: hamishknight

include/swift/Frontend/Frontend.h

@@ -646,10 +646,11 @@ class CompilerInstance {
 public:
   void freeASTContext();
 
-private:
-  /// Load stdlib & return true if should continue, i.e. no error
-  bool loadStdlib();
+  /// If an implicit standard library import is expected, loads the standard
+  /// library, returning \c false if we should continue, i.e. no error.
+  bool loadStdlibIfNeeded();
 
+private:
   /// Retrieve a description of which modules should be implicitly imported.
   ImplicitImportInfo getImplicitImportInfo() const;
 

lib/Frontend/Frontend.cpp

@@ -740,17 +740,6 @@ void CompilerInstance::performSemaUpTo(SourceFile::ASTStage_t LimitStage,
   ModuleDecl *mainModule = getMainModule();
   Context->LoadedModules[mainModule->getName()] = mainModule;
 
-  // If we aren't in a parse-only context, load the standard library.
-  if (LimitStage > SourceFile::Unprocessed &&
-      Invocation.getImplicitStdlibKind() == ImplicitStdlibKind::Stdlib
-      && !loadStdlib()) {
-    // If we failed to load the stdlib, mark the main module as having
-    // "failed to load", as it will contain no files.
-    // FIXME: We need to better handle a missing stdlib.
-    mainModule->setFailedToLoad();
-    return;
-  }
-
   // Make sure the main file is the first file in the module, so do this now.
   if (MainBufferID != NO_SUCH_BUFFER) {
     auto *mainFile = createSourceFileForMainModule(
@@ -813,23 +802,27 @@ void CompilerInstance::performSemaUpTo(SourceFile::ASTStage_t LimitStage,
   finishTypeChecking();
 }
 
-bool CompilerInstance::loadStdlib() {
+bool CompilerInstance::loadStdlibIfNeeded() {
+  // If we aren't expecting an implicit stdlib import, there's nothing to do.
+  if (getImplicitImportInfo().StdlibKind != ImplicitStdlibKind::Stdlib)
+    return false;
+
   FrontendStatsTracer tracer(getStatsReporter(), "load-stdlib");
-  ModuleDecl *M = Context->getStdlibModule(true);
+  ModuleDecl *M = Context->getStdlibModule(/*loadIfAbsent*/ true);
 
   if (!M) {
     Diagnostics.diagnose(SourceLoc(), diag::error_stdlib_not_found,
                          Invocation.getTargetTriple());
-    return false;
+    return true;
   }
 
-  // If we failed to load, we should have already diagnosed
+  // If we failed to load, we should have already diagnosed.
   if (M->failedToLoad()) {
     assert(Diagnostics.hadAnyError() &&
            "Module failed to load but nothing was diagnosed?");
-    return false;
+    return true;
   }
-  return true;
+  return false;
 }
 
 bool CompilerInstance::loadPartialModulesAndImplicitImports() {

lib/FrontendTool/FrontendTool.cpp

@@ -1250,6 +1250,17 @@ static bool performCompile(CompilerInstance &Instance,
   if (Invocation.getInputKind() == InputFileKind::LLVM)
     return compileLLVMIR(Instance);
 
+  // If we aren't in a parse-only context and expect an implicit stdlib import,
+  // load in the standard library. If we either fail to find it or encounter an
+  // error while loading it, bail early. Continuing the compilation will at best
+  // trigger a bunch of other errors due to the stdlib being missing, or at
+  // worst crash downstream as many call sites don't currently handle a missing
+  // stdlib.
+  if (!FrontendOptions::shouldActionOnlyParse(Action)) {
+    if (Instance.loadStdlibIfNeeded())
+      return true;
+  }
+
   if (FrontendOptions::shouldActionOnlyParse(Action)) {
     // Disable delayed parsing of type and function bodies when we've been
     // asked to dump the resulting AST.

lib/IDE/CompletionInstance.cpp

@@ -527,15 +527,16 @@ bool CompletionInstance::performNewOperation(
     }
     registerIDERequestFunctions(CI.getASTContext().evaluator);
 
-    CI.performParseAndResolveImportsOnly();
-
-    // If we didn't create a source file for completion, bail. This can happen
-    // if for example we fail to load the stdlib.
-    auto completionFile = CI.getCodeCompletionFile();
-    if (!completionFile)
+    // If we're expecting a standard library, but there either isn't one, or it
+    // failed to load, let's bail early and hand back an empty completion
+    // result to avoid any downstream crashes.
+    if (CI.loadStdlibIfNeeded())
       return true;
 
+    CI.performParseAndResolveImportsOnly();
+
     // If we didn't find a code completion token, bail.
+    auto completionFile = CI.getCodeCompletionFile();
     auto *state = completionFile.get()->getDelayedParserState();
     if (!state->hasCodeCompletionDelayedDeclState())
       return true;