Merge pull request #82940 from tshortli/silgen-custom-availability-domain-query

tshortli · web-flow · commit 96d0450d62a5 · 2025-07-10T13:16:14.000-07:00
ClangImporter: Generate and call custom availability domain predicates
diff --git a/include/swift/AST/AvailabilityDomain.h b/include/swift/AST/AvailabilityDomain.h
@@ -34,6 +34,7 @@ class ASTContext;
 class CustomAvailabilityDomain;
 class Decl;
 class DeclContext;
+class FuncDecl;
 class ModuleDecl;
 
 /// Represents a dimension of availability (e.g. macOS platform or Swift
@@ -336,20 +337,27 @@ class CustomAvailabilityDomain : public llvm::FoldingSetNode {
   Kind kind;
   ModuleDecl *mod;
   Decl *decl;
+  FuncDecl *predicateFunc;
 
   CustomAvailabilityDomain(Identifier name, Kind kind, ModuleDecl *mod,
-                           Decl *decl);
+                           Decl *decl, FuncDecl *predicateFunc);
 
 public:
   static const CustomAvailabilityDomain *get(StringRef name, Kind kind,
                                              ModuleDecl *mod, Decl *decl,
+                                             FuncDecl *predicateFunc,
                                              const ASTContext &ctx);
 
   Identifier getName() const { return name; }
   Kind getKind() const { return kind; }
   ModuleDecl *getModule() const { return mod; }
   Decl *getDecl() const { return decl; }
 
+  /// Returns the function that should be called at runtime to determine whether
+  /// the domain is available, or `nullptr` if the domain's availability is not
+  /// determined at runtime.
+  FuncDecl *getPredicateFunc() const { return predicateFunc; }
+
   /// Uniquing for `ASTContext`.
   static void Profile(llvm::FoldingSetNodeID &ID, Identifier name,
                       ModuleDecl *mod);
diff --git a/include/swift/AST/ClangModuleLoader.h b/include/swift/AST/ClangModuleLoader.h
@@ -298,6 +298,9 @@ class ClangModuleLoader : public ModuleLoader {
 
   virtual FuncDecl *getDefaultArgGenerator(const clang::ParmVarDecl *param) = 0;
 
+  virtual FuncDecl *
+  getAvailabilityDomainPredicate(const clang::VarDecl *var) = 0;
+
   virtual std::optional<Type>
   importFunctionReturnType(const clang::FunctionDecl *clangDecl,
                            DeclContext *dc) = 0;
diff --git a/include/swift/ClangImporter/ClangImporter.h b/include/swift/ClangImporter/ClangImporter.h
@@ -630,6 +630,8 @@ class ClangImporter final : public ClangModuleLoader {
 
   FuncDecl *getDefaultArgGenerator(const clang::ParmVarDecl *param) override;
 
+  FuncDecl *getAvailabilityDomainPredicate(const clang::VarDecl *var) override;
+
   bool isAnnotatedWith(const clang::CXXMethodDecl *method, StringRef attr);
 
   /// Find the lookup table that corresponds to the given Clang module.
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -5905,7 +5905,8 @@ const AvailabilityContext::Storage *AvailabilityContext::Storage::get(
 
 const CustomAvailabilityDomain *
 CustomAvailabilityDomain::get(StringRef name, Kind kind, ModuleDecl *mod,
-                              Decl *decl, const ASTContext &ctx) {
+                              Decl *decl, FuncDecl *predicateFunc,
+                              const ASTContext &ctx) {
   auto identifier = ctx.getIdentifier(name);
   llvm::FoldingSetNodeID id;
   CustomAvailabilityDomain::Profile(id, identifier, mod);
@@ -5918,8 +5919,8 @@ CustomAvailabilityDomain::get(StringRef name, Kind kind, ModuleDecl *mod,
 
   void *mem = ctx.Allocate(sizeof(CustomAvailabilityDomain),
                            alignof(CustomAvailabilityDomain));
-  auto *newNode =
-      ::new (mem) CustomAvailabilityDomain(identifier, kind, mod, decl);
+  auto *newNode = ::new (mem)
+      CustomAvailabilityDomain(identifier, kind, mod, decl, predicateFunc);
   foldingSet.InsertNode(newNode, insertPos);
 
   return newNode;
diff --git a/lib/AST/AvailabilityDomain.cpp b/lib/AST/AvailabilityDomain.cpp
@@ -17,6 +17,7 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
+#include "swift/ClangImporter/ClangImporter.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "llvm/ADT/StringSwitch.h"
@@ -42,8 +43,12 @@ customDomainForClangDecl(Decl *decl, const ASTContext &ctx) {
   auto *clangDecl = decl->getClangDecl();
   ASSERT(clangDecl);
 
+  auto *varDecl = dyn_cast<clang::VarDecl>(clangDecl);
+  if (!varDecl)
+    return nullptr;
+
   auto featureInfo = clangDecl->getASTContext().getFeatureAvailInfo(
-      const_cast<clang::Decl *>(clangDecl));
+      const_cast<clang::VarDecl *>(varDecl));
 
   // Ensure the decl actually represents an availability domain.
   if (featureInfo.first.empty())
@@ -52,9 +57,14 @@ customDomainForClangDecl(Decl *decl, const ASTContext &ctx) {
   if (featureInfo.second.Kind == clang::FeatureAvailKind::None)
     return nullptr;
 
+  FuncDecl *predicate = nullptr;
+  if (featureInfo.second.Kind == clang::FeatureAvailKind::Dynamic)
+    predicate =
+        ctx.getClangModuleLoader()->getAvailabilityDomainPredicate(varDecl);
+
   return CustomAvailabilityDomain::get(
       featureInfo.first, getCustomDomainKind(featureInfo.second.Kind),
-      decl->getModuleContext(), decl, ctx);
+      decl->getModuleContext(), decl, predicate, ctx);
 }
 
 std::optional<AvailabilityDomain>
@@ -359,10 +369,14 @@ bool StableAvailabilityDomainComparator::operator()(
 }
 
 CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name, Kind kind,
-                                                   ModuleDecl *mod, Decl *decl)
-    : name(name), kind(kind), mod(mod), decl(decl) {
+                                                   ModuleDecl *mod, Decl *decl,
+                                                   FuncDecl *predicateFunc)
+    : name(name), kind(kind), mod(mod), decl(decl),
+      predicateFunc(predicateFunc) {
   ASSERT(!name.empty());
   ASSERT(mod);
+  if (predicateFunc)
+    ASSERT(kind == Kind::Dynamic);
 }
 
 void CustomAvailabilityDomain::Profile(llvm::FoldingSetNodeID &ID,
diff --git a/lib/AST/AvailabilityQuery.cpp b/lib/AST/AvailabilityQuery.cpp
@@ -129,7 +129,8 @@ getOSAvailabilityDeclAndArguments(const AvailabilityQuery &query,
 
 FuncDecl *AvailabilityQuery::getDynamicQueryDeclAndArguments(
     llvm::SmallVectorImpl<unsigned> &arguments, ASTContext &ctx) const {
-  switch (getDomain().getKind()) {
+  auto domain = getDomain();
+  switch (domain.getKind()) {
   case AvailabilityDomain::Kind::Universal:
   case AvailabilityDomain::Kind::SwiftLanguage:
   case AvailabilityDomain::Kind::PackageDescription:
@@ -138,7 +139,9 @@ FuncDecl *AvailabilityQuery::getDynamicQueryDeclAndArguments(
   case AvailabilityDomain::Kind::Platform:
     return getOSAvailabilityDeclAndArguments(*this, arguments, ctx);
   case AvailabilityDomain::Kind::Custom:
-    // FIXME: [availability] Support custom domains.
-    return nullptr;
+    auto customDomain = domain.getCustomDomain();
+    ASSERT(customDomain);
+
+    return customDomain->getPredicateFunc();
   }
 }
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -7646,6 +7646,14 @@ ClangImporter::getDefaultArgGenerator(const clang::ParmVarDecl *param) {
   return nullptr;
 }
 
+FuncDecl *
+ClangImporter::getAvailabilityDomainPredicate(const clang::VarDecl *var) {
+  auto it = Impl.availabilityDomainPredicates.find(var);
+  if (it != Impl.availabilityDomainPredicates.end())
+    return it->second;
+  return nullptr;
+}
+
 SwiftLookupTable *
 ClangImporter::findLookupTable(const clang::Module *clangModule) {
   return Impl.findLookupTable(clangModule);
diff --git a/lib/ClangImporter/ImportDecl.cpp b/lib/ClangImporter/ImportDecl.cpp
@@ -4675,6 +4675,14 @@ namespace {
       if (correctSwiftName)
         markAsVariant(result, *correctSwiftName);
 
+      // If the decl represents an availability domain, eagerly synthesize its
+      // `if #available` predicate function.
+      if (decl->hasAttrs() &&
+          llvm::any_of(decl->getAttrs(), [](clang::Attr *attr) {
+            return isa<clang::AvailabilityDomainAttr>(attr);
+          }))
+        (void)synthesizer.makeAvailabilityDomainPredicate(decl);
+
       return result;
     }
 
diff --git a/lib/ClangImporter/ImporterImpl.h b/lib/ClangImporter/ImporterImpl.h
@@ -818,6 +818,11 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
 public:
   importer::PlatformAvailability platformAvailability;
 
+  /// The synthesized predicate functions for imported `VarDecl`s that represent
+  /// availability domains.
+  llvm::DenseMap<const clang::VarDecl *, FuncDecl *>
+      availabilityDomainPredicates;
+
 private:
   /// For importing names. This is initialized by the ClangImporter::create()
   /// after having set up a suitable Clang instance.
diff --git a/lib/ClangImporter/SwiftDeclSynthesizer.cpp b/lib/ClangImporter/SwiftDeclSynthesizer.cpp
@@ -2747,3 +2747,107 @@ SwiftDeclSynthesizer::synthesizeStaticFactoryForCXXForeignRef(
 
   return synthesizedFactories;
 }
+
+static std::pair<BraceStmt *, bool>
+synthesizeAvailabilityDomainPredicateBody(AbstractFunctionDecl *afd,
+                                          void *context) {
+  auto clangVarDecl = static_cast<const clang::VarDecl *>(context);
+  clang::ASTContext &clangCtx = clangVarDecl->getASTContext();
+  auto domainInfo =
+      clangCtx.getFeatureAvailInfo(const_cast<clang::VarDecl *>(clangVarDecl));
+  ASSERT(domainInfo.second.Call);
+
+  auto funcDecl = cast<FuncDecl>(afd);
+  ASTContext &ctx = funcDecl->getASTContext();
+
+  // FIXME: The need for an intermediate function to call could be eliminated if
+  // Clang provided the predicate function decl directly, rather than a call
+  // expression that must be wrapped in a function.
+  // Synthesize `return {domain predicate expression}`.
+  auto clangHelperReturnStmt = clang::ReturnStmt::Create(
+      clangCtx, clang::SourceLocation(), domainInfo.second.Call, nullptr);
+
+  // Synthesize `int __XYZ_isAvailable() { return {predicate expr}; }`.
+  auto clangDeclName = clang::DeclarationName(
+      &clangCtx.Idents.get("__" + domainInfo.first.str() + "_isAvailable"));
+  auto clangDeclContext = clangCtx.getTranslationUnitDecl();
+  clang::QualType funcTy =
+      clangCtx.getFunctionType(domainInfo.second.Call->getType(), {},
+                               clang::FunctionProtoType::ExtProtoInfo());
+
+  auto clangHelperFuncDecl = clang::FunctionDecl::Create(
+      clangCtx, clangDeclContext, clang::SourceLocation(),
+      clang::SourceLocation(), clangDeclName, funcTy,
+      clangCtx.getTrivialTypeSourceInfo(funcTy),
+      clang::StorageClass::SC_Static);
+  clangHelperFuncDecl->setImplicit();
+  clangHelperFuncDecl->setImplicitlyInline();
+  clangHelperFuncDecl->setBody(clangHelperReturnStmt);
+
+  // Import `func __XYZ_isAvailable() -> Bool` into Swift.
+  auto helperFuncDecl = dyn_cast_or_null<FuncDecl>(
+      ctx.getClangModuleLoader()->importDeclDirectly(clangHelperFuncDecl));
+  if (!helperFuncDecl)
+    return {nullptr, /*isTypeChecked=*/true};
+
+  auto helperFuncRef = new (ctx) DeclRefExpr(ConcreteDeclRef(helperFuncDecl),
+                                             DeclNameLoc(), /*Implicit=*/true);
+  helperFuncRef->setType(helperFuncDecl->getInterfaceType());
+
+  // Synthesize `__XYZ_isAvailable()`.
+  auto helperCall = CallExpr::createImplicit(
+      ctx, helperFuncRef, ArgumentList::createImplicit(ctx, {}));
+  helperCall->setType(helperFuncDecl->getResultInterfaceType());
+  helperCall->setThrows(nullptr);
+
+  // Synthesize `__XYZ_isAvailable()._value`.
+  auto *memberRef =
+      UnresolvedDotExpr::createImplicit(ctx, helperCall, ctx.Id_value_);
+
+  // Synthesize `return __XYZ_isAvailable()._value`.
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, memberRef);
+  auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
+                                /*implicit=*/true);
+
+  return {body, /*isTypeChecked=*/false};
+}
+
+FuncDecl *SwiftDeclSynthesizer::makeAvailabilityDomainPredicate(
+    const clang::VarDecl *var) {
+  ASTContext &ctx = ImporterImpl.SwiftContext;
+  clang::ASTContext &clangCtx = var->getASTContext();
+  auto featureInfo =
+      clangCtx.getFeatureAvailInfo(const_cast<clang::VarDecl *>(var));
+
+  // If the decl doesn't represent and availability domain, skip it.
+  if (featureInfo.first.empty())
+    return nullptr;
+
+  // Only dynamic availability domains require a predicate function.
+  if (featureInfo.second.Kind != clang::FeatureAvailKind::Dynamic)
+    return nullptr;
+
+  if (!featureInfo.second.Call)
+    return nullptr;
+
+  // Synthesize `func __swift_XYZ_isAvailable() -> Builtin.Int1 { ... }`.
+  std::string s;
+  llvm::raw_string_ostream os(s);
+  os << "__swift_" << featureInfo.first << "_isAvailable";
+  DeclName funcName(ctx, DeclBaseName(ctx.getIdentifier(s)),
+                    ParameterList::createEmpty(ctx));
+
+  auto funcDecl = FuncDecl::createImplicit(
+      ctx, StaticSpellingKind::None, funcName, SourceLoc(), /*Async=*/false,
+      /*Throws=*/false, Type(), {}, ParameterList::createEmpty(ctx),
+      BuiltinIntegerType::get(1, ctx), ImporterImpl.ImportedHeaderUnit);
+  funcDecl->setBodySynthesizer(synthesizeAvailabilityDomainPredicateBody,
+                               (void *)var);
+  funcDecl->setAccess(AccessLevel::Public);
+  funcDecl->getAttrs().add(new (ctx)
+                               AlwaysEmitIntoClientAttr(/*IsImplicit=*/true));
+
+  ImporterImpl.availabilityDomainPredicates[var] = funcDecl;
+
+  return funcDecl;
+}
diff --git a/lib/ClangImporter/SwiftDeclSynthesizer.h b/lib/ClangImporter/SwiftDeclSynthesizer.h
@@ -341,6 +341,10 @@ class SwiftDeclSynthesizer {
   synthesizeStaticFactoryForCXXForeignRef(
       const clang::CXXRecordDecl *cxxRecordDecl);
 
+  /// Synthesize a Swift function that calls the Clang runtime predicate
+  /// function for the availability domain represented by `var`.
+  FuncDecl *makeAvailabilityDomainPredicate(const clang::VarDecl *var);
+
 private:
   Type getConstantLiteralType(Type type, ConstantConvertKind convertKind);
 };
diff --git a/lib/Frontend/Frontend.cpp b/lib/Frontend/Frontend.cpp
@@ -1440,8 +1440,8 @@ static void configureAvailabilityDomains(const ASTContext &ctx,
   llvm::SmallDenseMap<Identifier, const CustomAvailabilityDomain *> domainMap;
   auto createAndInsertDomain = [&](const std::string &name,
                                    CustomAvailabilityDomain::Kind kind) {
-    auto *domain =
-        CustomAvailabilityDomain::get(name, kind, mainModule, nullptr, ctx);
+    auto *domain = CustomAvailabilityDomain::get(name, kind, mainModule,
+                                                 nullptr, nullptr, ctx);
     bool inserted = domainMap.insert({domain->getName(), domain}).second;
     ASSERT(inserted); // Domains must be unique.
   };
diff --git a/test/IRGen/Inputs/AvailabilityDomains.h b/test/IRGen/Inputs/AvailabilityDomains.h
@@ -0,0 +1,34 @@
+#include <feature-availability.h>
+
+static struct __AvailabilityDomain __EnabledDomain __attribute__((
+    availability_domain(EnabledDomain))) = {__AVAILABILITY_DOMAIN_ENABLED, 0};
+
+static struct __AvailabilityDomain __DisabledDomain __attribute__((
+    availability_domain(DisabledDomain))) = {__AVAILABILITY_DOMAIN_DISABLED, 0};
+
+int dynamic_domain_pred();
+
+static struct __AvailabilityDomain __DynamicDomain
+    __attribute__((availability_domain(DynamicDomain))) = {
+        __AVAILABILITY_DOMAIN_DYNAMIC, dynamic_domain_pred};
+
+#define AVAIL 0
+#define UNAVAIL 1
+
+__attribute__((availability(domain : EnabledDomain, AVAIL))) void
+available_in_enabled_domain(void);
+
+__attribute__((availability(domain : EnabledDomain, UNAVAIL))) void
+unavailable_in_enabled_domain(void);
+
+__attribute__((availability(domain : DisabledDomain, AVAIL))) void
+available_in_disabled_domain(void);
+
+__attribute__((availability(domain : DisabledDomain, UNAVAIL))) void
+unavailable_in_disabled_domain(void);
+
+__attribute__((availability(domain : DynamicDomain, AVAIL))) void
+available_in_dynamic_domain(void);
+
+__attribute__((availability(domain : DynamicDomain, UNAVAIL))) void
+unavailable_in_dynamic_domain(void);
diff --git a/test/IRGen/availability_custom_domains.swift b/test/IRGen/availability_custom_domains.swift
diff --git a/test/IRGen/availability_custom_domains_clang.swift b/test/IRGen/availability_custom_domains_clang.swift
diff --git a/test/IRGen/availability_custom_domains_maccatalyst_zippered.swift b/test/IRGen/availability_custom_domains_maccatalyst_zippered.swift
diff --git a/test/SILGen/availability_query_custom_domains_clang.swift b/test/SILGen/availability_query_custom_domains_clang.swift
