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

The effect of declaring an import `@_weakLinked` is to treat every declaration from the module as if it were declared with `@_weakLinked`. This is useful in environments where entire modules may not be present at runtime. Although it is already possible to instruct the linker to weakly link an entire dylib, a Swift attribute provides a way to declare intent in source code and also opens the door to diagnostics and other compiler behaviors that depend on knowing that all the module's symbols will be weakly linked.

rdar://96098097

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

@@ -686,6 +686,9 @@ class ModuleDecl
   // Is \p spiGroup accessible as an explicitly imported SPI from this module?
   bool isImportedAsSPI(Identifier spiGroup, const ModuleDecl *fromModule) const;
 
+  /// Is \p importedModule imported as \c @_weakLinked from this module?
+  bool importsModuleAsWeakLinked(const ModuleDecl *importedModule) 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->importsModuleAsWeakLinked(getModuleContext()))
+    return true;
+
   auto availability = getAvailabilityForLinkage();
   if (availability.isAlwaysAvailable())
     return false;

lib/AST/Module.cpp

@@ -2574,6 +2574,15 @@ 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) &&
+        module == import.module.importedModule)
+      return true;
+  }
+  return false;
+}
+
 bool ModuleDecl::isImportedAsSPI(const SpecializeAttr *attr,
                                  const ValueDecl *targetDecl) const {
   auto targetModule = targetDecl->getModuleContext();
@@ -2599,6 +2608,15 @@ bool ModuleDecl::isImportedAsSPI(Identifier spiGroup,
   return importedSPIGroups.count(spiGroup);
 }
 
+bool ModuleDecl::importsModuleAsWeakLinked(
+    const ModuleDecl *importedModule) const {
+  for (auto file : getFiles()) {
+    if (file->importsModuleAsWeakLinked(importedModule))
+      return true;
+  }
+  return false;
+}
+
 bool Decl::isSPI() const {
   return !getSPIGroups().empty();
 }

lib/SIL/IR/SILFunction.cpp

@@ -462,7 +462,7 @@ bool SILFunction::isWeakImported() const {
   if (!isAvailableExternally())
     return false;
 
-  if (isAlwaysWeakImported())
+  if (isAlwaysWeakImported() || isWeakImportedByModule())
     return true;
 
   if (Availability.isAlwaysAvailable())

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