[AutoDiff] Add generated implicit declarations to SynthesizedFileUnit.

include/swift/AST/Module.h

@@ -95,6 +95,8 @@ enum class FileUnitKind {
   Builtin,
   /// A serialized Swift AST.
   SerializedAST,
+  /// A synthesized file.
+  Synthesized,
   /// An imported Clang module.
   ClangModule,
   /// A Clang module imported from DWARF.

include/swift/AST/SourceFile.h

@@ -394,8 +394,6 @@ class SourceFile final : public FileUnit {
   void cacheVisibleDecls(SmallVectorImpl<ValueDecl *> &&globals) const;
   const SmallVectorImpl<ValueDecl *> &getCachedVisibleDecls() const;
 
-  void addVisibleDecl(ValueDecl *decl);
-
   virtual void lookupValue(DeclName name, NLKind lookupKind,
                            SmallVectorImpl<ValueDecl*> &result) const override;
 

include/swift/AST/SynthesizedFileUnit.h

@@ -0,0 +1,63 @@
+//===--- SynthesizedFileUnit.h - A synthesized file unit --------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_SYNTHESIZEDFILEUNIT_H
+#define SWIFT_AST_SYNTHESIZEDFILEUNIT_H
+
+#include "swift/AST/FileUnit.h"
+#include "swift/Basic/Debug.h"
+
+namespace swift {
+
+/// A container for synthesized module-level declarations.
+class SynthesizedFileUnit final : public FileUnit {
+  /// Synthesized top level declarations.
+  TinyPtrVector<ValueDecl *> TopLevelDecls;
+
+  /// A unique identifier representing this file; used to mark private decls
+  /// within the file to keep them from conflicting with other files in the
+  /// same module.
+  mutable Identifier PrivateDiscriminator;
+
+public:
+  SynthesizedFileUnit(ModuleDecl &M);
+  ~SynthesizedFileUnit() = default;
+
+  /// Add a synthesized top-level declaration.
+  void addTopLevelDecl(ValueDecl *D) { TopLevelDecls.push_back(D); }
+
+  virtual void lookupValue(DeclName name, NLKind lookupKind,
+                           SmallVectorImpl<ValueDecl *> &result) const override;
+
+  void lookupObjCMethods(
+      ObjCSelector selector,
+      SmallVectorImpl<AbstractFunctionDecl *> &results) const override;
+
+  Identifier getDiscriminatorForPrivateValue(const ValueDecl *D) const override;
+
+  void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
+
+  ArrayRef<ValueDecl *> getTopLevelDecls() const {
+    return TopLevelDecls;
+  };
+
+  static bool classof(const FileUnit *file) {
+    return file->getKind() == FileUnitKind::Synthesized;
+  }
+  static bool classof(const DeclContext *DC) {
+    return isa<FileUnit>(DC) && classof(cast<FileUnit>(DC));
+  }
+};
+
+} // namespace swift
+
+#endif // SWIFT_AST_SYNTHESIZEDFILEUNIT_H

include/swift/SILOptimizer/Utils/Differentiation/ADContext.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/DiagnosticsSIL.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SILOptimizer/Utils/Differentiation/Common.h"
 #include "swift/SILOptimizer/Utils/Differentiation/DifferentiationInvoker.h"
@@ -66,6 +67,9 @@ class ADContext {
   /// Shared pass manager.
   SILPassManager &passManager;
 
+  /// A synthesized file unit.
+  SynthesizedFileUnit *synthesizedFile = nullptr;
+
   /// The worklist (stack) of `differentiable_function` instructions to be
   /// processed.
   llvm::SmallVector<DifferentiableFunctionInst *, 32>
@@ -122,6 +126,10 @@ class ADContext {
   SILPassManager &getPassManager() const { return passManager; }
   Lowering::TypeConverter &getTypeConverter() { return module.Types; }
 
+  /// Get or create a synthesized file for adding generated linear map structs
+  /// and branching trace enums. Used by `LinearMapInfo`.
+  SynthesizedFileUnit &getOrCreateSynthesizedFile();
+
   /// Returns true if the `differentiable_function` instruction worklist is
   /// empty.
   bool isDifferentiableFunctionInstsWorklistEmpty() const {

include/swift/SILOptimizer/Utils/Differentiation/LinearMapInfo.h

@@ -18,6 +18,7 @@
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_LINEARMAPINFO_H
 
 #include "swift/AST/AutoDiff.h"
+#include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/SIL/ApplySite.h"
 #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
 #include "llvm/ADT/DenseMap.h"
@@ -85,6 +86,9 @@ class LinearMapInfo {
   /// Mapping from linear map structs to their branching trace enum fields.
   llvm::DenseMap<StructDecl *, VarDecl *> linearMapStructEnumFields;
 
+  /// A synthesized file unit.
+  SynthesizedFileUnit &synthesizedFile;
+
   /// A type converter, used to compute struct/enum SIL types.
   Lowering::TypeConverter &typeConverter;
 
@@ -97,13 +101,8 @@ class LinearMapInfo {
   VarDecl *addVarDecl(NominalTypeDecl *nominal, StringRef name, Type type);
 
   /// Retrieves the file unit that contains implicit declarations in the
-  /// current Swift module. If it does not exist, create one.
-  ///
-  // FIXME: Currently it defaults to the file containing `original`, if it can
-  // be determined. Otherwise, it defaults to any file unit in the module. To
-  // handle this more properly, we could revive the DerivedFileUnit class to
-  // contain all synthesized implicit type declarations.
-  SourceFile &getDeclarationFileUnit();
+  /// current Swift module.
+  SynthesizedFileUnit &getSynthesizedFile() { return synthesizedFile; }
 
   /// Computes and sets the access level for the given nominal type, given the
   /// original function linkage.

lib/AST/Decl.cpp

@@ -639,6 +639,7 @@ static_assert(sizeof(checkSourceLocType(&ID##Decl::getLoc)) == 2, \
     return getSerializedLocs()->Loc;
   }
   case FileUnitKind::Builtin:
+  case FileUnitKind::Synthesized:
   case FileUnitKind::ClangModule:
   case FileUnitKind::DWARFModule:
     return SourceLoc();

lib/AST/DeclContext.cpp

@@ -623,6 +623,9 @@ unsigned DeclContext::printContext(raw_ostream &OS, const unsigned indent,
     case FileUnitKind::Source:
       OS << " file=\"" << cast<SourceFile>(this)->getFilename() << "\"";
       break;
+    case FileUnitKind::Synthesized:
+      OS << " synthesized file";
+      break;
     case FileUnitKind::SerializedAST:
     case FileUnitKind::ClangModule:
     case FileUnitKind::DWARFModule:

lib/AST/Module.cpp

@@ -38,6 +38,7 @@
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/ReferencedNameTracker.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Compiler.h"
 #include "swift/Basic/SourceManager.h"
@@ -218,6 +219,7 @@ class swift::SourceLookupCache {
 
   SmallVector<ValueDecl *, 0> AllVisibleValues;
 };
+
 SourceLookupCache &SourceFile::getCache() const {
   if (!Cache) {
     const_cast<SourceFile *>(this)->Cache =
@@ -317,6 +319,10 @@ SourceLookupCache::SourceLookupCache(const ModuleDecl &M) {
   FrontendStatsTracer tracer(M.getASTContext().Stats,
                              "module-populate-cache");
   for (const FileUnit *file : M.getFiles()) {
+    if (auto *SFU = dyn_cast<SynthesizedFileUnit>(file)) {
+      addToUnqualifiedLookupCache(SFU->getTopLevelDecls(), false);
+      continue;
+    }
     auto &SF = *cast<SourceFile>(file);
     addToUnqualifiedLookupCache(SF.getTopLevelDecls(), false);
   }
@@ -1173,6 +1179,9 @@ lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
   case FileUnitKind::Builtin:
     // The Builtin module declares no operators.
     return nullptr;
+  case FileUnitKind::Synthesized:
+    // Synthesized files currently declare no operators.
+    return nullptr;
   case FileUnitKind::Source:
     break;
   case FileUnitKind::SerializedAST:
@@ -1520,6 +1529,9 @@ StringRef ModuleDecl::getModuleFilename() const {
       Result = LF->getFilename();
       continue;
     }
+    // Skip synthesized files.
+    if (auto *SFU = dyn_cast<SynthesizedFileUnit>(F))
+      continue;
     return StringRef();
   }
   return Result;
@@ -2232,11 +2244,6 @@ SourceFile::getCachedVisibleDecls() const {
   return getCache().AllVisibleValues;
 }
 
-void SourceFile::addVisibleDecl(ValueDecl *decl) {
-  Decls->push_back(decl);
-  getCache().AllVisibleValues.push_back(decl);
-}
-
 static void performAutoImport(
     SourceFile &SF,
     SourceFile::ImplicitModuleImportKind implicitModuleImportKind) {
@@ -2662,6 +2669,70 @@ SourceFile::lookupOpaqueResultType(StringRef MangledName) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// SynthesizedFileUnit Implementation
+//===----------------------------------------------------------------------===//
+
+SynthesizedFileUnit::SynthesizedFileUnit(ModuleDecl &M)
+    : FileUnit(FileUnitKind::Synthesized, M) {
+  M.getASTContext().addDestructorCleanup(*this);
+}
+
+Identifier
+SynthesizedFileUnit::getDiscriminatorForPrivateValue(const ValueDecl *D) const {
+  assert(D->getDeclContext()->getModuleScopeContext() == this);
+
+  // Use cached primitive discriminator if it exists.
+  if (!PrivateDiscriminator.empty())
+    return PrivateDiscriminator;
+
+  assert(1 == count_if(getParentModule()->getFiles(),
+                       [](const FileUnit *FU) -> bool {
+                         return isa<SynthesizedFileUnit>(FU);
+                       }) &&
+         "Cannot promise uniqueness if multiple synthesized file units exist");
+
+  // Use a discriminator invariant across Swift version: a hash of the module
+  // name and a special string.
+  llvm::MD5 hash;
+  hash.update(getParentModule()->getName().str());
+  // TODO: Use a more robust discriminator for synthesized files. Pick something
+  // that cannot conflict with `SourceFile` discriminators.
+  hash.update("SYNTHESIZED FILE");
+  llvm::MD5::MD5Result result;
+  hash.final(result);
+
+  // Use the hash as a hex string, prefixed with an underscore to make sure
+  // it is a valid identifier.
+  // FIXME: There are more compact ways to encode a 16-byte value.
+  SmallString<33> buffer{"_"};
+  SmallString<32> hashString;
+  llvm::MD5::stringifyResult(result, hashString);
+  buffer += hashString;
+  PrivateDiscriminator = getASTContext().getIdentifier(buffer.str().upper());
+  return PrivateDiscriminator;
+}
+
+void SynthesizedFileUnit::lookupValue(
+    DeclName name, NLKind lookupKind,
+    SmallVectorImpl<ValueDecl *> &result) const {
+  for (auto *decl : TopLevelDecls) {
+    if (decl->getFullName().matchesRef(name))
+      result.push_back(decl);
+  }
+}
+
+void SynthesizedFileUnit::lookupObjCMethods(
+    ObjCSelector selector,
+    SmallVectorImpl<AbstractFunctionDecl *> &results) const {
+  // Synthesized files only contain top-level declarations, no `@objc` methods.
+}
+
+void SynthesizedFileUnit::getTopLevelDecls(
+    SmallVectorImpl<swift::Decl *> &results) const {
+  results.append(TopLevelDecls.begin(), TopLevelDecls.end());
+}
+
 //===----------------------------------------------------------------------===//
 // Miscellaneous
 //===----------------------------------------------------------------------===//
@@ -2700,6 +2771,9 @@ void swift::simple_display(llvm::raw_ostream &out, const FileUnit *file) {
   case FileUnitKind::Builtin:
     out << "(Builtin)";
     return;
+  case FileUnitKind::Synthesized:
+    out << "(synthesized)";
+    return;
   case FileUnitKind::DWARFModule:
   case FileUnitKind::ClangModule:
   case FileUnitKind::SerializedAST:

lib/IRGen/GenDecl.cpp

@@ -415,6 +415,7 @@ class ObjCProtocolInitializerVisitor
 } // end anonymous namespace
 
 namespace {
+
 class PrettySourceFileEmission : public llvm::PrettyStackTraceEntry {
   const SourceFile &SF;
 public:
@@ -424,6 +425,19 @@ class PrettySourceFileEmission : public llvm::PrettyStackTraceEntry {
     os << "While emitting IR for source file " << SF.getFilename() << '\n';
   }
 };
+
+class PrettySynthesizedFileUnitEmission : public llvm::PrettyStackTraceEntry {
+  const SynthesizedFileUnit &SFU;
+
+public:
+  explicit PrettySynthesizedFileUnitEmission(const SynthesizedFileUnit &SFU)
+      : SFU(SFU) {}
+
+  void print(raw_ostream &os) const override {
+    os << "While emitting IR for synthesized file" << &SFU << "\n";
+  }
+};
+
 } // end anonymous namespace
 
 /// Emit all the top-level code in the source file.
@@ -475,6 +489,13 @@ void IRGenModule::emitSourceFile(SourceFile &SF) {
   }
 }
 
+/// Emit all the top-level code in the synthesized file unit.
+void IRGenModule::emitSynthesizedFileUnit(SynthesizedFileUnit &SFU) {
+  PrettySynthesizedFileUnitEmission StackEntry(SFU);
+  for (auto *decl : SFU.getTopLevelDecls())
+    emitGlobalDecl(decl);
+}
+
 /// Collect elements of an already-existing global list with the given
 /// \c name into \c list.
 ///

lib/IRGen/IRGen.cpp

@@ -990,6 +990,11 @@ performIRGeneration(const IRGenOptions &Opts, ModuleDecl *M,
       }
     }
 
+    // Emit synthesized file units.
+    for (auto *File : M->getFiles())
+      if (auto *nextSFU = dyn_cast<SynthesizedFileUnit>(File))
+        IGM.emitSynthesizedFileUnit(*nextSFU);
+
     // Okay, emit any definitions that we suddenly need.
     irgen.emitLazyDefinitions();
 
@@ -1235,6 +1240,9 @@ static void performParallelIRGeneration(
     if (auto *SF = dyn_cast<SourceFile>(File)) {
       CurrentIGMPtr IGM = irgen.getGenModule(SF);
       IGM->emitSourceFile(*SF);
+    } else if (auto *nextSFU = dyn_cast<SynthesizedFileUnit>(File)) {
+      CurrentIGMPtr IGM = irgen.getGenModule(nextSFU);
+      IGM->emitSynthesizedFileUnit(*nextSFU);
     } else {
       File->collectLinkLibraries([&](LinkLibrary LinkLib) {
         irgen.getPrimaryIGM()->addLinkLibrary(LinkLib);

lib/IRGen/IRGenModule.h

@@ -25,6 +25,7 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/ReferenceCounting.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/Basic/ClusteredBitVector.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/OptimizationMode.h"
@@ -1289,6 +1290,7 @@ private:                            \
   llvm::LLVMContext &getLLVMContext() const { return LLVMContext; }
 
   void emitSourceFile(SourceFile &SF);
+  void emitSynthesizedFileUnit(SynthesizedFileUnit &SFU);
   void addLinkLibrary(const LinkLibrary &linkLib);
 
   /// Attempt to finalize the module.