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();

lib/SIL/IR/SILFunctionBuilder.cpp

@@ -312,7 +312,14 @@ SILFunction *SILFunctionBuilder::getOrCreateFunction(
       F->setClangNodeOwner(decl);
 
     F->setAvailabilityForLinkage(decl->getAvailabilityForLinkage());
-    F->setAlwaysWeakImported(decl->isAlwaysWeakImported());
+
+    if (decl->isAlwaysWeakImported()) {
+      F->setIsWeakImported(IsWeakImported_t::IsAlwaysWeakImported);
+    } else if (decl->isWeakImported(mod.getSwiftModule())) {
+      F->setIsWeakImported(IsWeakImported_t::IsWeakImportedByModule);
+    } else {
+      F->setIsWeakImported(IsWeakImported_t::IsNotWeakImported);
+    }
 
     if (auto *accessor = dyn_cast<AccessorDecl>(decl)) {
       auto *storage = accessor->getStorage();

lib/SIL/Parser/ParseSIL.cpp

@@ -6483,7 +6483,9 @@ bool SILParserState::parseDeclSIL(Parser &P) {
     if (!objCReplacementFor.empty())
       FunctionState.F->setObjCReplacement(objCReplacementFor);
     FunctionState.F->setSpecialPurpose(specialPurpose);
-    FunctionState.F->setAlwaysWeakImported(isWeakImported);
+    FunctionState.F->setIsWeakImported(
+        isWeakImported ? IsWeakImported_t::IsAlwaysWeakImported
+                       : IsWeakImported_t::IsNotWeakImported);
     FunctionState.F->setAvailabilityForLinkage(availability);
     FunctionState.F->setWithoutActuallyEscapingThunk(
       isWithoutActuallyEscapingThunk);

lib/Sema/ImportResolution.cpp

@@ -554,6 +554,9 @@ UnboundImport::UnboundImport(ImportDecl *ID)
     import.options |= ImportFlags::Preconcurrency;
     import.preconcurrencyRange = attr->getRangeWithAt();
   }
+
+  if (auto attr = ID->getAttrs().getAttribute<WeakLinkedAttr>())
+    import.options |= ImportFlags::WeakLinked;
 }
 
 bool UnboundImport::checkNotTautological(const SourceFile &SF) {

lib/Serialization/DeserializeSIL.cpp

@@ -661,7 +661,9 @@ SILDeserializer::readSILFunctionChecked(DeclID FID, SILFunction *existingFn,
     fn->setEffectsKind(EffectsKind(effect));
     fn->setOptimizationMode(OptimizationMode(optimizationMode));
     fn->setPerfConstraints((PerformanceConstraints)perfConstr);
-    fn->setAlwaysWeakImported(isWeakImported);
+    fn->setIsWeakImported(isWeakImported
+                              ? IsWeakImported_t::IsAlwaysWeakImported
+                              : IsWeakImported_t::IsNotWeakImported);
     fn->setClassSubclassScope(SubclassScope(subclassScope));
     fn->setHasCReferences(bool(hasCReferences));
 

test/IRGen/Inputs/usr/include/BridgeTestFoundation.h

@@ -53,10 +53,13 @@ __attribute__((availability(macosx,introduced=10.51)))
 @interface NSUserNotificationAction : NSObject
 @end
 
+void always_available_function();
+
 __attribute__((availability(macosx,introduced=10.51)))
 void future_function_should_be_weak();
 
 extern int weak_variable __attribute__((weak_import));
+extern int strong_variable;
 
 @interface NSError : NSObject