[AutoDiff] Add generated implicit declarations to SynthesizedFileUnit.

Add implicit declarations generated by the differentiation transform to a
`SynthesizedFileUnit` instead of an ad-hoc pre-existing `SourceFile`.

Resolves TF-1232: type reconstruction for AutoDiff-generated declarations.

Previously, type reconstruction failed because retroactively adding declarations
to a `SourceFile` did not update name lookup caches.

Author: dan-zheng

include/swift/AST/SynthesizedFileUnit.h

@@ -33,7 +33,7 @@ class SynthesizedFileUnit final : public FileUnit {
   ~SynthesizedFileUnit() = default;
 
   /// Add a synthesized top-level declaration.
-  void addTopLevelDecl(ValueDecl *FD) { TopLevelDecls.push_back(FD); }
+  void addTopLevelDecl(ValueDecl *D) { TopLevelDecls.push_back(D); }
 
   virtual void lookupValue(DeclName name, NLKind lookupKind,
                            SmallVectorImpl<ValueDecl *> &result) const override;

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;
+
   /// The worklist (stack) of `differentiable_function` instructions to be
   /// processed.
   llvm::SmallVector<DifferentiableFunctionInst *, 32>
@@ -120,6 +124,7 @@ class ADContext {
   SILModule &getModule() const { return module; }
   ASTContext &getASTContext() const { return module.getASTContext(); }
   SILPassManager &getPassManager() const { return passManager; }
+  SynthesizedFileUnit &getSynthesizedFile() { return synthesizedFile; }
   Lowering::TypeConverter &getTypeConverter() { return module.Types; }
 
   /// Returns true if the `differentiable_function` instruction worklist is

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/Module.cpp

@@ -1180,6 +1180,8 @@ lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
     // 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:
@@ -1527,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;

lib/IRGen/IRGen.cpp

@@ -1241,7 +1241,7 @@ static void performParallelIRGeneration(
       CurrentIGMPtr IGM = irgen.getGenModule(SF);
       IGM->emitSourceFile(*SF);
     } else if (auto *nextSFU = dyn_cast<SynthesizedFileUnit>(File)) {
-      CurrentIGMPtr IGM = irgen.getGenModule(SF);
+      CurrentIGMPtr IGM = irgen.getGenModule(nextSFU);
       IGM->emitSynthesizedFileUnit(*nextSFU);
     } else {
       File->collectLinkLibraries([&](LinkLibrary LinkLib) {

lib/SILOptimizer/Utils/Differentiation/ADContext.cpp

@@ -54,9 +54,17 @@ static FuncDecl *findOperatorDeclInProtocol(DeclName operatorName,
 // ADContext methods
 //===----------------------------------------------------------------------===//
 
+static SynthesizedFileUnit *createNewSynthesizedFileUnit(ModuleDecl *module) {
+  auto &ctx = module->getASTContext();
+  auto *file = new (ctx) SynthesizedFileUnit(*module);
+  module->addFile(*file);
+  return file;
+}
+
 ADContext::ADContext(SILModuleTransform &transform)
     : transform(transform), module(*transform.getModule()),
-      passManager(*transform.getPassManager()) {}
+      passManager(*transform.getPassManager()),
+      synthesizedFile(*createNewSynthesizedFileUnit(module.getSwiftModule())) {}
 
 FuncDecl *ADContext::getPlusDecl() const {
   if (!cachedPlusFn) {

lib/SILOptimizer/Utils/Differentiation/LinearMapInfo.cpp

@@ -58,6 +58,7 @@ LinearMapInfo::LinearMapInfo(ADContext &context, AutoDiffLinearMapKind kind,
                              const DifferentiableActivityInfo &activityInfo)
     : kind(kind), original(original), derivative(derivative),
       activityInfo(activityInfo), indices(indices),
+      synthesizedFile(context.getSynthesizedFile()),
       typeConverter(context.getTypeConverter()) {
   generateDifferentiationDataStructures(context, derivative);
 }
@@ -84,17 +85,6 @@ VarDecl *LinearMapInfo::addVarDecl(NominalTypeDecl *nominal, StringRef name,
   return varDecl;
 }
 
-SourceFile &LinearMapInfo::getDeclarationFileUnit() {
-  if (original->hasLocation())
-    if (auto *declContext = original->getLocation().getAsDeclContext())
-      if (auto *parentSourceFile = declContext->getParentSourceFile())
-        return *parentSourceFile;
-  for (auto *file : original->getModule().getSwiftModule()->getFiles())
-    if (auto *src = dyn_cast<SourceFile>(file))
-      return *src;
-  llvm_unreachable("No files?");
-}
-
 void LinearMapInfo::computeAccessLevel(NominalTypeDecl *nominal,
                                        SILLinkage originalLinkage) {
   auto &astCtx = nominal->getASTContext();
@@ -129,7 +119,7 @@ LinearMapInfo::createBranchingTraceDecl(SILBasicBlock *originalBB,
   assert(originalBB->getParent() == original);
   auto &astCtx = original->getASTContext();
   auto *moduleDecl = original->getModule().getSwiftModule();
-  auto &file = getDeclarationFileUnit();
+  auto &file = getSynthesizedFile();
   // Create a branching trace enum.
   Mangle::ASTMangler mangler;
   auto *resultIndices = IndexSubset::get(
@@ -200,7 +190,7 @@ LinearMapInfo::createLinearMapStruct(SILBasicBlock *originalBB,
   assert(originalBB->getParent() == original);
   auto *original = originalBB->getParent();
   auto &astCtx = original->getASTContext();
-  auto &file = getDeclarationFileUnit();
+  auto &file = getSynthesizedFile();
   // Create a linear map struct.
   Mangle::ASTMangler mangler;
   auto *resultIndices = IndexSubset::get(

test/AutoDiff/compiler_crashers/tf1232-autodiff-generated-declaration-mangling.swift renamed to test/AutoDiff/compiler_crashers_fixed/tf1232-autodiff-generated-declaration-mangling.swift

@@ -1,4 +1,4 @@
-// RUN: not %target-build-swift -g %s
+// RUN: %target-build-swift -g %s
 // REQUIRES: asserts
 
 // TF-1232: IRGenDebugInfo crash due to lack of proper mangling for