Fix a use-after-free caused by creating DerivedFileUnit on the fly

Sema was creating DerivedFileUnit on the fly, while something else is iterating
over FileUnits in the module.  The fix is to create DerivedFileUnit in advance.
This change immediately uncovered a lot of code that assumed that the module
consists of a single FileUnit at certain conditions.  This patch also fixes
that code (SourceKit patch is separate, not sending it).

The test change is because now operator == on NSObjects is correctly recognised
as coming from a system module.

rdar://16153700, rdar://16227621, possibly rdar://16049613


Swift SVN r14692

include/swift/AST/Module.h

@@ -26,6 +26,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -42,6 +43,7 @@ namespace swift {
   enum class DeclKind : uint8_t;
   class ExtensionDecl;
   class DebuggerClient;
+  class DerivedFileUnit;
   class FileUnit;
   class FuncDecl;
   class InfixOperatorDecl;
@@ -172,7 +174,7 @@ private:
   // FIXME: Do we really need to bloat all modules with this?
   DebuggerClient *DebugClient = nullptr;
 
-  TinyPtrVector<FileUnit *> Files;
+  SmallVector<FileUnit *, 2> Files;
 
   Module(Identifier name, ASTContext &ctx);
 public:
@@ -198,6 +200,8 @@ public:
   /// dealing with.
   FileUnit &getMainFile(FileUnitKind expectedKind) const;
 
+  DerivedFileUnit &getDerivedFileUnit() const;
+
   DebuggerClient *getDebugClient() const { return DebugClient; }
   void setDebugClient(DebuggerClient *R) {
     assert(!DebugClient && "Debugger client already set");
@@ -555,12 +559,15 @@ public:
 /// A container for a module-level definition derived as part of an implicit
 /// protocol conformance.
 class DerivedFileUnit final : public FileUnit {
-  FuncDecl *const DerivedDecl;
+  TinyPtrVector<FuncDecl *> DerivedDecls;
 
 public:
-  DerivedFileUnit(Module &M, FuncDecl *derivedDecl)
-    : FileUnit(FileUnitKind::Derived, M), DerivedDecl(derivedDecl) {}
-  ~DerivedFileUnit() {}
+  DerivedFileUnit(Module &M);
+  ~DerivedFileUnit() = default;
+
+  void addDerivedDecl(FuncDecl *FD) {
+    DerivedDecls.push_back(FD);
+  }
 
   void lookupValue(Module::AccessPathTy accessPath, Identifier name,
                    NLKind lookupKind,

lib/AST/Module.cpp

@@ -277,12 +277,35 @@ Module::Module(Identifier name, ASTContext &ctx)
 
 
 void Module::addFile(FileUnit &newFile) {
+  assert(!isa<DerivedFileUnit>(newFile) &&
+         "DerivedFileUnits are added automatically");
+
   // Require Main and REPL files to be the first file added.
   assert(Files.empty() ||
          !isa<SourceFile>(newFile) ||
          cast<SourceFile>(newFile).Kind == SourceFileKind::Library ||
          cast<SourceFile>(newFile).Kind == SourceFileKind::SIL);
   Files.push_back(&newFile);
+
+  switch (newFile.getKind()) {
+  case FileUnitKind::Source:
+  case FileUnitKind::ClangModule: {
+    for (auto File : Files) {
+      if (isa<DerivedFileUnit>(File))
+        return;
+    }
+    auto DFU = new (Ctx) DerivedFileUnit(*this);
+    Files.push_back(DFU);
+    break;
+  }
+
+  case FileUnitKind::Builtin:
+  case FileUnitKind::SerializedAST:
+    break;
+
+  case FileUnitKind::Derived:
+    llvm_unreachable("DerivedFileUnits are added automatically");
+  }
 }
 
 void Module::removeFile(FileUnit &existingFile) {
@@ -297,6 +320,15 @@ void Module::removeFile(FileUnit &existingFile) {
   Files.erase(I.base());
 }
 
+DerivedFileUnit &Module::getDerivedFileUnit() const {
+  for (auto File : Files) {
+    if (auto DFU = dyn_cast<DerivedFileUnit>(File))
+      return *DFU;
+  }
+  llvm_unreachable("the client should not be calling this function if "
+                   "there is no DerivedFileUnit");
+}
+
 #define FORWARD(name, args) \
   for (const FileUnit *file : getFiles()) \
     file->name args;
@@ -313,6 +345,11 @@ void BuiltinUnit::lookupValue(Module::AccessPathTy accessPath, Identifier name,
   getCache().lookupValue(name, lookupKind, *this, result);
 }
 
+DerivedFileUnit::DerivedFileUnit(Module &M)
+    : FileUnit(FileUnitKind::Derived, M) {
+  M.Ctx.addDestructorCleanup(*this);
+}
+
 void DerivedFileUnit::lookupValue(Module::AccessPathTy accessPath,
                                   Identifier name,
                                   NLKind lookupKind,
@@ -324,8 +361,10 @@ void DerivedFileUnit::lookupValue(Module::AccessPathTy accessPath,
   if (accessPath.size() == 1 && accessPath.front().first != name)
     return;
 
-  if (name == DerivedDecl->getName())
-    result.push_back(DerivedDecl);
+  for (auto D : DerivedDecls) {
+    if (name == D->getName())
+      result.push_back(D);
+  }
 }
 
 void DerivedFileUnit::lookupVisibleDecls(Module::AccessPathTy accessPath,
@@ -333,18 +372,20 @@ void DerivedFileUnit::lookupVisibleDecls(Module::AccessPathTy accessPath,
                                          NLKind lookupKind) const {
   assert(accessPath.size() <= 1 && "can only refer to top-level decls");
 
-  // If this import is specific to some named type or decl ("import Swift.int")
-  // then filter out any lookups that don't match.
-  if (accessPath.size() == 1 &&
-      accessPath.front().first != DerivedDecl->getName())
-    return;
+  Identifier Id;
+  if (!accessPath.empty()) {
+    Id = accessPath.front().first;
+  }
 
-  consumer.foundDecl(DerivedDecl, DeclVisibilityKind::VisibleAtTopLevel);
+  for (auto D : DerivedDecls) {
+    if (Id.empty() || D->getName() == Id)
+      consumer.foundDecl(D, DeclVisibilityKind::VisibleAtTopLevel);
+  }
 }
 
 void DerivedFileUnit::getTopLevelDecls(SmallVectorImpl<swift::Decl *> &results)
 const {
-  results.push_back(DerivedDecl);
+  results.append(DerivedDecls.begin(), DerivedDecls.end());
 }
 
 void SourceFile::lookupValue(Module::AccessPathTy accessPath, Identifier name,
@@ -972,13 +1013,25 @@ bool Module::isSameAccessPath(AccessPathTy lhs, AccessPathTy rhs) {
 StringRef Module::getModuleFilename() const {
   // FIXME: Audit uses of this function and figure out how to migrate them to
   // per-file names. Modules can consist of more than one file.
-  if (getFiles().size() == 1) {
-    if (auto SF = dyn_cast<SourceFile>(getFiles().front()))
-      return SF->getFilename();
-    if (auto LF = dyn_cast<LoadedFile>(getFiles().front()))
-      return LF->getFilename();
-  }
+  StringRef Result;
+  for (auto F : getFiles()) {
+    if (auto SF = dyn_cast<SourceFile>(F)) {
+      if (!Result.empty())
         return StringRef();
+      Result = SF->getFilename();
+      continue;
+    }
+    if (auto LF = dyn_cast<LoadedFile>(F)) {
+      if (!Result.empty())
+        return StringRef();
+      Result = LF->getFilename();
+      continue;
+    }
+    if (isa<DerivedFileUnit>(F))
+      continue;
+    return StringRef();
+  }
+  return Result;
 }
 
 bool Module::isStdlibModule() const {
@@ -988,8 +1041,8 @@ bool Module::isStdlibModule() const {
 bool Module::isSystemModule() const {
   if (isStdlibModule())
     return true;
-  if (getFiles().size() == 1) {
-    if (auto LF = dyn_cast<LoadedFile>(getFiles().front()))
+  for (auto F : getFiles()) {
+    if (auto LF = dyn_cast<LoadedFile>(F))
       return LF->isSystemModule();
   }
   return false;

lib/Frontend/Frontend.cpp

@@ -322,15 +322,15 @@ void CompilerInstance::performParse() {
 
   if (!Invocation.getParseOnly()) {
     // Type-check each top-level input besides the main source file.
-    auto InputSourceFiles = MainModule->getFiles().slice(0, BufferIDs.size());
-    for (auto File : InputSourceFiles)
+    for (auto File : MainModule->getFiles())
       if (auto SF = dyn_cast<SourceFile>(File))
         if (PrimaryBufferID == NO_SUCH_BUFFER ||
             (SF->getBufferID().hasValue() &&
              SF->getBufferID().getValue() == PrimaryBufferID))
           performTypeChecking(*SF, PersistentState.getTopLevelContext());
 
-    // If there were no source files, we should still record known protocols.
+    // Even if there were no source files, we should still record known
+    // protocols.
     if (Context->getStdlibModule())
       Context->recordKnownProtocols(Context->getStdlibModule());
   }

lib/SIL/SILPrinter.cpp

@@ -1299,9 +1299,6 @@ void SILModule::print(llvm::raw_ostream &OS, bool Verbose,
   
   OS << "\n\nimport Builtin\nimport " << STDLIB_NAME << "\n\n";
 
-  // Print the declarations and types from the origin module.
-  // FIXME: What about multi-file modules?
-  if (M && M->getFiles().size() == 1) {
   // Compute the list of emitted functions, whose AST Decls we do not need to
   // print.
   llvm::DenseSet<const Decl*> emittedFunctions;
@@ -1309,15 +1306,16 @@ void SILModule::print(llvm::raw_ostream &OS, bool Verbose,
     if (f.hasLocation())
       emittedFunctions.insert(f.getLocation().getAsASTNode<Decl>());
 
+  // Print the declarations and types from the origin module.
+  {
     PrintOptions Options;
     Options.FunctionDefinitions = false;
     Options.TypeDefinitions = true;
     Options.VarInitializers = true;
     Options.SkipImplicit = true;
 
-    // FIXME: Use some kind of visitor interface here.
     SmallVector<Decl *, 32> topLevelDecls;
-    M->getFiles().front()->getTopLevelDecls(topLevelDecls);
+    M->getTopLevelDecls(topLevelDecls);
     for (const Decl *D : topLevelDecls) {
       if ((isa<ValueDecl>(D) || isa<OperatorDecl>(D)) &&
           !emittedFunctions.count(D) &&

lib/Sema/DerivedConformanceRawRepresentable.cpp

@@ -46,7 +46,7 @@ void DerivedConformance::_insertOperatorDecl(NominalTypeDecl *scope,
   auto mod = scope->getModuleContext();
 
   // Add it to the module in a DerivedFileUnit.
-  mod->addFile(*new (C) DerivedFileUnit(*mod, cast<FuncDecl>(member)));
+  mod->getDerivedFileUnit().addDerivedDecl(cast<FuncDecl>(member));
 
   // Add it as a derived global decl to the nominal type.
   auto oldDerived = scope->getDerivedGlobalDecls();