Merge pull request #59509 from zoecarver/frt-with-ref-counting

[cxx-interop] Add support for reference counting operations on foreign reference types.

Author: zoecarver

include/swift/AST/Decl.h

@@ -4241,15 +4241,7 @@ class ClassDecl final : public NominalTypeDecl {
 
   /// Whether the class uses the ObjC object model (reference counting,
   /// allocation, etc.), the Swift model, or has no reference counting at all.
-  ReferenceCounting getObjectModel() const {
-    if (isForeignReferenceType())
-      return ReferenceCounting::None;
-
-    if (checkAncestry(AncestryFlags::ObjCObjectModel))
-      return ReferenceCounting::ObjC;
-
-    return ReferenceCounting::Native;
-  }
+  ReferenceCounting getObjectModel() const;
 
   LayoutConstraintKind getLayoutConstraintKind() const {
     if (getObjectModel() == ReferenceCounting::ObjC)
@@ -4363,6 +4355,8 @@ class ClassDecl final : public NominalTypeDecl {
   /// non-reference-counted swift reference type that was imported from a C++
   /// record.
   bool isForeignReferenceType() const;
+
+  bool hasRefCountingAnnotations() const;
 };
 
 /// The set of known protocols for which derived conformances are supported.

include/swift/AST/DiagnosticsClangImporter.def

@@ -170,6 +170,27 @@ NOTE(dont_use_iterator_api, none, "C++ method '%0' that returns an unsafe "
                                   "APIs instead",
      (StringRef))
 
+ERROR(reference_type_must_have_retain_attr,none,
+      "reference type '%0' must have 'retain:' swift attribute.", (StringRef))
+ERROR(reference_type_must_have_release_attr,none,
+      "reference type '%0' must have 'release:' swift attribute.", (StringRef))
+ERROR(foreign_reference_types_cannot_find_retain,none,
+      "cannot find retain function '%0' for reference type '%1'.", (StringRef, StringRef))
+ERROR(foreign_reference_types_cannot_find_release,none,
+      "cannot find release function '%0' for reference type '%1'.", (StringRef, StringRef))
+ERROR(too_many_reference_type_retain_operations,none,
+      "too many functions with name '%0'. There must be exactly one retain "
+      "function for reference type '%1'.", (StringRef, StringRef))
+ERROR(too_many_reference_type_release_operations,none,
+      "too many functions with name '%0'. There must be exactly one release "
+      "function for reference type '%1'.", (StringRef, StringRef))
+ERROR(foreign_reference_types_invalid_retain,none,
+      "specified retain function '%0' is invalid. Retain must have exactly one "
+      "argument of type '%1'", (StringRef, StringRef))
+ERROR(foreign_reference_types_invalid_release,none,
+      "specified release function '%0' is invalid. Release must have exactly "
+      "one argument of type '%1'", (StringRef, StringRef))
+
 NOTE(unsupported_builtin_type, none, "built-in type '%0' not supported", (StringRef))
 NOTE(record_field_not_imported, none, "field %0 not imported", (const clang::NamedDecl*))
 NOTE(invoked_func_not_imported, none, "function %0 not imported", (const clang::NamedDecl*))

include/swift/AST/DiagnosticsIRGen.def

@@ -59,10 +59,5 @@ ERROR(temporary_allocation_alignment_not_positive,none,
 ERROR(temporary_allocation_alignment_not_power_of_2,none,
       "alignment value must be a power of two", ())
 
-ERROR(foreign_reference_types_unsupported,none,
-      "attempt to use a foreign reference type in a generic context. "
-      "Foreign reference types are currently not supported. Using foreign "
-      "reference types in a generic context is not yet implemented.", ())
-
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/ReferenceCounting.h

@@ -33,6 +33,10 @@ enum class ReferenceCounting : uint8_t {
   /// types.
   None,
 
+  /// The object uses swift_attr("retain:XXX") and "release:XXX" to implement
+  /// reference counting.
+  Custom,
+
   /// The object uses _Block_copy/_Block_release reference counting.
   ///
   /// This is a strict subset of ObjC; all blocks are also ObjC reference

include/swift/Basic/Compiler.h

@@ -180,7 +180,9 @@
 
 // Tells Swift's ClangImporter to import a C++ type as a foreign reference type.
 #if __has_attribute(swift_attr)
-#define SWIFT_IMPORT_REFERENCE __attribute__((swift_attr("import_reference")))
+#define SWIFT_IMPORT_REFERENCE __attribute__((swift_attr("import_reference"))) \
+  __attribute__((swift_attr("retain:immortal")))                               \
+  __attribute__((swift_attr("release:immortal")))
 #else
 #define SWIFT_IMPORT_REFERENCE
 #endif

include/swift/ClangImporter/ClangImporterRequests.h

@@ -285,6 +285,73 @@ class IsSafeUseOfCxxDecl
   bool evaluate(Evaluator &evaluator, SafeUseOfCxxDeclDescriptor desc) const;
 };
 
+enum class CustomRefCountingOperationKind { retain, release };
+
+struct CustomRefCountingOperationDescriptor final {
+  const ClassDecl *decl;
+  CustomRefCountingOperationKind kind;
+
+  CustomRefCountingOperationDescriptor(const ClassDecl *decl,
+                                       CustomRefCountingOperationKind kind)
+      : decl(decl), kind(kind) {}
+
+  friend llvm::hash_code
+  hash_value(const CustomRefCountingOperationDescriptor &desc) {
+    return llvm::hash_combine(desc.decl, desc.kind);
+  }
+
+  friend bool operator==(const CustomRefCountingOperationDescriptor &lhs,
+                         const CustomRefCountingOperationDescriptor &rhs) {
+    return lhs.decl == rhs.decl && lhs.kind == rhs.kind;
+  }
+
+  friend bool operator!=(const CustomRefCountingOperationDescriptor &lhs,
+                         const CustomRefCountingOperationDescriptor &rhs) {
+    return !(lhs == rhs);
+  }
+};
+
+void simple_display(llvm::raw_ostream &out,
+                    CustomRefCountingOperationDescriptor desc);
+SourceLoc extractNearestSourceLoc(CustomRefCountingOperationDescriptor desc);
+
+struct CustomRefCountingOperationResult {
+  enum CustomRefCountingOperationResultKind {
+    noAttribute,
+    immortal,
+    notFound,
+    tooManyFound,
+    foundOperation
+  };
+
+  CustomRefCountingOperationResultKind kind;
+  ValueDecl *operation;
+  std::string name;
+};
+
+class CustomRefCountingOperation
+    : public SimpleRequest<CustomRefCountingOperation,
+                           CustomRefCountingOperationResult(
+                               CustomRefCountingOperationDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+  // Caching
+  bool isCached() const { return true; }
+
+  // Source location
+  SourceLoc getNearestLoc() const { return SourceLoc(); };
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  CustomRefCountingOperationResult
+  evaluate(Evaluator &evaluator,
+           CustomRefCountingOperationDescriptor desc) const;
+};
+
 #define SWIFT_TYPEID_ZONE ClangImporter
 #define SWIFT_TYPEID_HEADER "swift/ClangImporter/ClangImporterTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/ClangImporter/ClangImporterTypeIDZone.def

@@ -30,3 +30,6 @@ SWIFT_REQUEST(ClangImporter, CxxRecordSemantics,
 SWIFT_REQUEST(ClangImporter, IsSafeUseOfCxxDecl,
               bool(SafeUseOfCxxRecordDescriptor), Cached,
               NoLocationInfo)
+SWIFT_REQUEST(ClangImporter, CustomRefCountingOperation,
+              CustomRefCountingOperationResult(CustomRefCountingOperationDescriptor), Cached,
+              NoLocationInfo)

lib/AST/Decl.cpp

@@ -15,11 +15,11 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/AST/Decl.h"
-#include "swift/AST/AccessRequests.h"
-#include "swift/AST/AccessScope.h"
 #include "swift/AST/ASTContext.h"
-#include "swift/AST/ASTWalker.h"
 #include "swift/AST/ASTMangler.h"
+#include "swift/AST/ASTWalker.h"
+#include "swift/AST/AccessRequests.h"
+#include "swift/AST/AccessScope.h"
 #include "swift/AST/Attr.h"
 #include "swift/AST/CaptureInfo.h"
 #include "swift/AST/DiagnosticEngine.h"
@@ -32,7 +32,6 @@
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/LazyResolver.h"
-#include "swift/AST/ASTMangler.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
@@ -44,11 +43,17 @@
 #include "swift/AST/ResilienceExpansion.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/Stmt.h"
+#include "swift/AST/SwiftNameTranslation.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeLoc.h"
-#include "swift/AST/SwiftNameTranslation.h"
 #include "swift/Basic/Defer.h"
+#include "swift/Basic/Range.h"
+#include "swift/Basic/Statistic.h"
+#include "swift/Basic/StringExtras.h"
+#include "swift/Basic/TypeID.h"
+#include "swift/ClangImporter/ClangImporterRequests.h"
 #include "swift/ClangImporter/ClangModule.h"
+#include "swift/Demangling/ManglingMacros.h"
 #include "swift/Parse/Lexer.h" // FIXME: Bad dependency
 #include "clang/Lex/MacroInfo.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -57,11 +62,6 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
-#include "swift/Basic/Range.h"
-#include "swift/Basic/StringExtras.h"
-#include "swift/Basic/Statistic.h"
-#include "swift/Basic/TypeID.h"
-#include "swift/Demangling/ManglingMacros.h"
 
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/Module.h"
@@ -5233,6 +5233,25 @@ bool ClassDecl::isForeignReferenceType() const {
   return getClangDecl() && isa<clang::RecordDecl>(getClangDecl());
 }
 
+bool ClassDecl::hasRefCountingAnnotations() const {
+  return evaluateOrDefault(getASTContext().evaluator,
+                           CustomRefCountingOperation(
+                               {this, CustomRefCountingOperationKind::release}),
+                           {})
+             .kind != CustomRefCountingOperationResult::immortal;
+}
+
+ReferenceCounting ClassDecl::getObjectModel() const {
+  if (isForeignReferenceType())
+    return hasRefCountingAnnotations() ? ReferenceCounting::Custom
+                                       : ReferenceCounting::None;
+
+  if (checkAncestry(AncestryFlags::ObjCObjectModel))
+    return ReferenceCounting::ObjC;
+
+  return ReferenceCounting::Native;
+}
+
 EnumCaseDecl *EnumCaseDecl::create(SourceLoc CaseLoc,
                                    ArrayRef<EnumElementDecl *> Elements,
                                    DeclContext *DC) {

lib/AST/NameLookupRequests.cpp

@@ -476,6 +476,23 @@ swift::extractNearestSourceLoc(const ClangRecordMemberLookupDescriptor &desc) {
   return extractNearestSourceLoc(desc.recordDecl);
 }
 
+//----------------------------------------------------------------------------//
+// CustomRefCountingOperation computation.
+//----------------------------------------------------------------------------//
+
+void swift::simple_display(llvm::raw_ostream &out,
+                           CustomRefCountingOperationDescriptor desc) {
+  out << "Finding custom (foreign reference) reference counting operation '"
+      << (desc.kind == CustomRefCountingOperationKind::retain ? "retain"
+                                                              : "release")
+      << "for '" << desc.decl->getNameStr() << "'.\n";
+}
+
+SourceLoc
+swift::extractNearestSourceLoc(CustomRefCountingOperationDescriptor desc) {
+  return SourceLoc();
+}
+
 // Implement the clang importer type zone.
 #define SWIFT_TYPEID_ZONE ClangImporter
 #define SWIFT_TYPEID_HEADER "swift/ClangImporter/ClangImporterTypeIDZone.def"

lib/AST/Type.cpp

@@ -6037,12 +6037,16 @@ ReferenceCounting TypeBase::getReferenceCounting() {
   CanType type = getCanonicalType();
   ASTContext &ctx = type->getASTContext();
 
+  if (isForeignReferenceType())
+    return lookThroughAllOptionalTypes()
+                   ->getClassOrBoundGenericClass()
+                   ->hasRefCountingAnnotations()
+               ? ReferenceCounting::Custom
+               : ReferenceCounting::None;
+
   // In the absence of Objective-C interoperability, everything uses native
   // reference counting or is the builtin BridgeObject.
   if (!ctx.LangOpts.EnableObjCInterop) {
-    if (isForeignReferenceType())
-      return ReferenceCounting::None;
-
     return type->getKind() == TypeKind::BuiltinBridgeObject
              ? ReferenceCounting::Bridge
              : ReferenceCounting::Native;

lib/ClangImporter/ClangImporter.cpp

@@ -6073,4 +6073,49 @@ void swift::simple_display(llvm::raw_ostream &out,
 
 SourceLoc swift::extractNearestSourceLoc(SafeUseOfCxxDeclDescriptor desc) {
   return SourceLoc();
-}
+}
+
+CustomRefCountingOperationResult CustomRefCountingOperation::evaluate(
+    Evaluator &evaluator, CustomRefCountingOperationDescriptor desc) const {
+  auto swiftDecl = desc.decl;
+  auto operation = desc.kind;
+  auto &ctx = swiftDecl->getASTContext();
+
+  std::string operationStr = operation == CustomRefCountingOperationKind::retain
+                                 ? "retain:"
+                                 : "release:";
+
+  auto decl = cast<clang::RecordDecl>(swiftDecl->getClangDecl());
+  if (!decl->hasAttrs())
+    return {CustomRefCountingOperationResult::noAttribute, nullptr, ""};
+
+  auto retainFnAttr =
+      llvm::find_if(decl->getAttrs(), [&operationStr](auto *attr) {
+        if (auto swiftAttr = dyn_cast<clang::SwiftAttrAttr>(attr))
+          return swiftAttr->getAttribute().startswith(operationStr);
+        return false;
+      });
+  if (retainFnAttr == decl->getAttrs().end()) {
+    return {CustomRefCountingOperationResult::noAttribute, nullptr, ""};
+  }
+
+  auto name = cast<clang::SwiftAttrAttr>(*retainFnAttr)
+                  ->getAttribute()
+                  .drop_front(StringRef(operationStr).size())
+                  .str();
+
+  if (name == "immortal")
+    return {CustomRefCountingOperationResult::immortal, nullptr, name};
+
+  llvm::SmallVector<ValueDecl *, 1> results;
+  ctx.lookupInModule(swiftDecl->getParentModule(), name, results);
+
+  if (results.size() == 1)
+    return {CustomRefCountingOperationResult::foundOperation, results.front(),
+            name};
+
+  if (results.empty())
+    return {CustomRefCountingOperationResult::notFound, nullptr, name};
+
+  return {CustomRefCountingOperationResult::tooManyFound, nullptr, name};
+}