Merge pull request #60414 from tshortli/weak-linked-import

AST/IRGen: Accept `@_weakLinked` on import decls to force weak linkage of symbols from a module

Author: tshortli

include/swift/AST/Attr.def

@@ -456,7 +456,8 @@ DECL_ATTR(_clangImporterSynthesizedType, ClangImporterSynthesizedType,
   74)
 SIMPLE_DECL_ATTR(_weakLinked, WeakLinked,
   OnNominalType | OnAssociatedType | OnFunc | OnAccessor | OnVar |
-  OnSubscript | OnConstructor | OnEnumElement | OnExtension | UserInaccessible |
+  OnSubscript | OnConstructor | OnEnumElement | OnExtension | OnImport |
+  UserInaccessible |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   75)
 SIMPLE_DECL_ATTR(frozen, Frozen,

include/swift/AST/FileUnit.h

@@ -123,6 +123,17 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
                             const ModuleDecl *importedModule,
                             SmallSetVector<Identifier, 4> &spiGroups) const {};
 
+  /// Checks whether this file imports \c module as \c @_weakLinked.
+  virtual bool importsModuleAsWeakLinked(const ModuleDecl *module) const {
+    // For source files, this should be overridden to inspect the import
+    // declarations in the file. Other kinds of file units, like serialized
+    // modules, can just use this default implementation since the @_weakLinked
+    // attribute is not transitive. If module C is imported @_weakLinked by
+    // module B, that does not imply that module A imports module C @_weakLinked
+    // if it imports module B.
+    return false;
+  }
+
   virtual Optional<Fingerprint>
   loadFingerprint(const IterableDeclContext *IDC) const { return None; }
 

include/swift/AST/Import.h

@@ -84,6 +84,9 @@ enum class ImportFlags {
   /// concurrency.
   Preconcurrency = 0x20,
 
+  /// The module's symbols are linked weakly.
+  WeakLinked = 0x40,
+
   /// Used for DenseMap.
   Reserved = 0x80
 };

include/swift/AST/Module.h

@@ -414,12 +414,11 @@ class ModuleDecl
   /// present overlays as if they were part of their underlying module.
   std::pair<ModuleDecl *, Identifier> getDeclaringModuleAndBystander();
 
+public:
   ///  If this is a traditional (non-cross-import) overlay, get its underlying
   ///  module if one exists.
   ModuleDecl *getUnderlyingModuleIfOverlay() const;
 
-public:
-
   /// Returns true if this module is an underscored cross import overlay
   /// declared by \p other or its underlying clang module, either directly or
   /// transitively (via intermediate cross-import overlays - for cross-imports
@@ -686,6 +685,10 @@ class ModuleDecl
   // Is \p spiGroup accessible as an explicitly imported SPI from this module?
   bool isImportedAsSPI(Identifier spiGroup, const ModuleDecl *fromModule) const;
 
+  /// Is \p targetDecl from a module that is imported as \c @_weakLinked from
+  /// this module?
+  bool isImportedAsWeakLinked(const Decl *targetDecl) const;
+
   /// \sa getImportedModules
   enum class ImportFilterKind {
     /// Include imports declared with `@_exported`.

include/swift/AST/SourceFile.h

@@ -356,6 +356,9 @@ class SourceFile final : public FileUnit {
                 const ModuleDecl *importedModule,
                 llvm::SmallSetVector<Identifier, 4> &spiGroups) const override;
 
+  /// Is \p module imported as \c @_weakLinked by this file?
+  bool importsModuleAsWeakLinked(const ModuleDecl *module) const override;
+
   // Is \p targetDecl accessible as an explicitly imported SPI from this file?
   bool isImportedAsSPI(const ValueDecl *targetDecl) const;
 

include/swift/SIL/SILFunction.h

@@ -64,6 +64,11 @@ enum IsDistributed_t {
   IsNotDistributed,
   IsDistributed,
 };
+enum IsWeakImported_t {
+  IsNotWeakImported,
+  IsWeakImportedByModule,
+  IsAlwaysWeakImported,
+};
 
 enum class PerformanceConstraints : uint8_t {
   None = 0,
@@ -317,9 +322,8 @@ class SILFunction
   /// would indicate.
   unsigned HasCReferences : 1;
 
-  /// Whether cross-module references to this function should always use
-  /// weak linking.
-  unsigned IsWeakImported : 1;
+  /// Whether cross-module references to this function should use weak linking.
+  unsigned IsWeakImported : 2;
 
   /// Whether the implementation can be dynamically replaced.
   unsigned IsDynamicReplaceable : 1;
@@ -797,12 +801,18 @@ class SILFunction
 
   /// Returns whether this function's symbol must always be weakly referenced
   /// across module boundaries.
-  bool isAlwaysWeakImported() const { return IsWeakImported; }
+  bool isAlwaysWeakImported() const {
+    return IsWeakImported == IsWeakImported_t::IsAlwaysWeakImported;
+  }
 
-  void setAlwaysWeakImported(bool value) {
-    IsWeakImported = value;
+  /// Returns whether this function's symbol was referenced by a module that
+  /// imports the defining module \c @_weakLinked.
+  bool isWeakImportedByModule() const {
+    return IsWeakImported == IsWeakImported_t::IsWeakImportedByModule;
   }
 
+  void setIsWeakImported(IsWeakImported_t value) { IsWeakImported = value; }
+
   bool isWeakImported() const;
 
   /// Returns whether this function implementation can be dynamically replaced.

lib/AST/Decl.cpp

@@ -1002,6 +1002,9 @@ bool Decl::isWeakImported(ModuleDecl *fromModule) const {
   if (isAlwaysWeakImported())
     return true;
 
+  if (fromModule->isImportedAsWeakLinked(this))
+    return true;
+
   auto availability = getAvailabilityForLinkage();
   if (availability.isAlwaysAvailable())
     return false;

lib/AST/Module.cpp

@@ -2574,6 +2574,25 @@ bool SourceFile::isImportedAsSPI(const ValueDecl *targetDecl) const {
   return false;
 }
 
+bool SourceFile::importsModuleAsWeakLinked(const ModuleDecl *module) const {
+  for (auto &import : *Imports) {
+    if (!import.options.contains(ImportFlags::WeakLinked))
+      continue;
+
+    const ModuleDecl *importedModule = import.module.importedModule;
+    if (module == importedModule)
+      return true;
+
+    // Also check whether the target module is actually the underlyingClang
+    // module for this @_weakLinked import.
+    const ModuleDecl *clangModule =
+        importedModule->getUnderlyingModuleIfOverlay();
+    if (module == clangModule)
+      return true;
+  }
+  return false;
+}
+
 bool ModuleDecl::isImportedAsSPI(const SpecializeAttr *attr,
                                  const ValueDecl *targetDecl) const {
   auto targetModule = targetDecl->getModuleContext();
@@ -2599,6 +2618,15 @@ bool ModuleDecl::isImportedAsSPI(Identifier spiGroup,
   return importedSPIGroups.count(spiGroup);
 }
 
+bool ModuleDecl::isImportedAsWeakLinked(const Decl *targetDecl) const {
+  const auto *declaringModule = targetDecl->getModuleContext();
+  for (auto file : getFiles()) {
+    if (file->importsModuleAsWeakLinked(declaringModule))
+      return true;
+  }
+  return false;
+}
+
 bool Decl::isSPI() const {
   return !getSPIGroups().empty();
 }

lib/IRGen/GenDecl.cpp

@@ -2487,10 +2487,23 @@ void IRGenModule::emitGlobalDecl(Decl *D) {
 Address IRGenModule::getAddrOfSILGlobalVariable(SILGlobalVariable *var,
                                                 const TypeInfo &ti,
                                                 ForDefinition_t forDefinition) {
+  LinkEntity entity = LinkEntity::forSILGlobalVariable(var, *this);
+  LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
+
   if (auto clangDecl = var->getClangDecl()) {
     auto addr = getAddrOfClangGlobalDecl(cast<clang::VarDecl>(clangDecl),
                                          forDefinition);
 
+    // Override the linkage computed by Clang if the decl is from another
+    // module that imported @_weakLinked.
+    //
+    // FIXME: We should be able to set the linkage unconditionally here but
+    //        some fixes are needed for Cxx interop.
+    if (auto globalVar = dyn_cast<llvm::GlobalVariable>(addr)) {
+      if (getSwiftModule()->isImportedAsWeakLinked(var->getDecl()))
+        globalVar->setLinkage(link.getLinkage());
+    }
+
     // If we're not emitting this to define it, make sure we cast it to the
     // right type.
     if (!forDefinition) {
@@ -2503,7 +2516,6 @@ Address IRGenModule::getAddrOfSILGlobalVariable(SILGlobalVariable *var,
     return Address(addr, alignment);
   }
 
-  LinkEntity entity = LinkEntity::forSILGlobalVariable(var, *this);
   ResilienceExpansion expansion = getResilienceExpansionForLayout(var);
 
   llvm::Type *storageType;
@@ -2549,7 +2561,6 @@ Address IRGenModule::getAddrOfSILGlobalVariable(SILGlobalVariable *var,
 
   // Check whether we've created the global variable already.
   // FIXME: We should integrate this into the LinkEntity cache more cleanly.
-  LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
   auto gvar = Module.getGlobalVariable(link.getName(), /*allowInternal*/ true);
   if (gvar) {
     if (forDefinition)
@@ -3222,6 +3233,7 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
     }
   }
 
+  LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
   bool isDefinition = f->isDefinition();
   bool hasOrderNumber =
       isDefinition && !shouldCallPreviousImplementation;
@@ -3242,8 +3254,16 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
   if (clangAddr) {
     fn = dyn_cast<llvm::Function>(clangAddr->stripPointerCasts());
 
-    // If we have a function, move it to the appropriate position.
     if (fn) {
+      // Override the linkage computed by Clang if the decl is from another
+      // module that imported @_weakLinked.
+      //
+      // FIXME: We should be able to set the linkage unconditionally here but
+      //        some fixes are needed for Cxx interop.
+      if (!forDefinition && f->isWeakImportedByModule())
+        fn->setLinkage(link.getLinkage());
+
+      // If we have a function, move it to the appropriate position.
       if (hasOrderNumber) {
         auto &fnList = Module.getFunctionList();
         fnList.remove(fn);
@@ -3264,8 +3284,6 @@ llvm::Function *IRGenModule::getAddrOfSILFunction(
       getSignature(f->getLoweredFunctionType(), fpKind);
   addLLVMFunctionAttributes(f, signature);
 
-  LinkInfo link = LinkInfo::get(*this, entity, forDefinition);
-
   fn = createFunction(*this, link, signature, insertBefore,
                       f->getOptimizationMode(), shouldEmitStackProtector(f));
 

lib/SIL/IR/SILFunction.cpp

@@ -453,6 +453,9 @@ bool SILFunction::isTypeABIAccessible(SILType type) const {
 }
 
 bool SILFunction::isWeakImported() const {
+  if (isWeakImportedByModule())
+    return true;
+
   // For imported functions check the Clang declaration.
   if (ClangNodeOwner)
     return ClangNodeOwner->getClangDecl()->isWeakImported();