[ClangImporter] Attach _SwiftifyImportProtocol to imported protocols

...with bounds attributes

This creates safe overloads for any methods in the protocol annotated
with bounds information. Updates _SwiftifyImportProtocol to make the
added overloads in the protocol public.

rdar://144335990

Author: hnrklssn

lib/ClangImporter/ImportDecl.cpp

@@ -5881,6 +5881,8 @@ namespace {
 
       result->setMemberLoader(&Impl, 0);
 
+      Impl.swiftifyProtocol(result);
+
       return result;
     }
 

lib/ClangImporter/ImporterImpl.h

@@ -1827,6 +1827,7 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   void addOptionSetTypealiases(NominalTypeDecl *nominal);
 
   void swiftify(AbstractFunctionDecl *MappedDecl);
+  void swiftifyProtocol(NominalTypeDecl *MappedDecl);
 
   /// Find the lookup table that corresponds to the given Clang module.
   ///

lib/ClangImporter/SwiftifyDecl.cpp

@@ -16,6 +16,7 @@
 
 #include "ImporterImpl.h"
 #include "swift/AST/ASTContext.h"
+#include "swift/AST/ASTPrinter.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsClangImporter.h"
 #include "swift/AST/ParameterList.h"
@@ -151,6 +152,7 @@ class SwiftifyInfoPrinter {
     out << "\"";
   }
 
+protected:
   void printSeparator() {
     if (!firstParam) {
       out << ", ";
@@ -159,6 +161,7 @@ class SwiftifyInfoPrinter {
     }
   }
 
+private:
   void printParamOrReturn(ssize_t pointerIndex) {
     if (pointerIndex == SELF_PARAM_INDEX)
       out << ".self";
@@ -296,12 +299,25 @@ static StringRef getAttributeName(const clang::CountAttributedType *CAT) {
       llvm_unreachable("CountAttributedType cannot be ended_by");
   }
 }
+
+template<typename T>
+static bool getImplicitObjectParamAnnotation(const clang::ObjCMethodDecl* D) {
+    return false; // Only C++ methods have implicit params
+}
+
+static size_t getNumParams(const clang::ObjCMethodDecl* D) {
+    return D->param_size();
+}
+static size_t getNumParams(const clang::FunctionDecl* D) {
+    return D->getNumParams();
+}
 } // namespace
 
+template<typename T>
 static bool swiftifyImpl(ClangImporter::Implementation &Self,
                          SwiftifyInfoPrinter &printer,
                          const AbstractFunctionDecl *MappedDecl,
-                         const clang::FunctionDecl *ClangDecl) {
+                         const T *ClangDecl) {
   DLOG("Checking " << *ClangDecl << " for bounds and lifetime info\n");
 
   // FIXME: for private macro generated functions we do not serialize the
@@ -310,13 +326,6 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
       ClangDecl->getAccess() == clang::AS_private)
     return false;
 
-  {
-    UnaliasedInstantiationVisitor visitor;
-    visitor.TraverseType(ClangDecl->getType());
-    if (visitor.hasUnaliasedInstantiation)
-      return false;
-  }
-
   clang::ASTContext &clangASTContext = Self.getClangASTContext();
 
   // We only attach the macro if it will produce an overload. Any __counted_by
@@ -332,9 +341,10 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
       swiftReturnTy = ctorDecl->getResultInterfaceType();
     else
       ABORT("Unexpected AbstractFunctionDecl subclass.");
+    clang::QualType clangReturnTy = ClangDecl->getReturnType();
     bool returnIsStdSpan = printer.registerStdSpanTypeMapping(
-        swiftReturnTy, ClangDecl->getReturnType());
-    auto *CAT = ClangDecl->getReturnType()->getAs<clang::CountAttributedType>();
+        swiftReturnTy, clangReturnTy);
+    auto *CAT = clangReturnTy->getAs<clang::CountAttributedType>();
     if (SwiftifiableCAT(clangASTContext, CAT, swiftReturnTy)) {
       printer.printCountedBy(CAT, SwiftifyInfoPrinter::RETURN_VALUE_INDEX);
       DLOG("  Found bounds info '" << clang::QualType(CAT, 0) << "' on return value\n");
@@ -361,13 +371,13 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
     size_t swiftNumParams = MappedDecl->getParameters()->size() -
                             (ClangDecl->isVariadic() ? 1 : 0);
     ASSERT((MappedDecl->isImportAsInstanceMember() == isClangInstanceMethod) ==
-           (ClangDecl->getNumParams() == swiftNumParams));
+           (getNumParams(ClangDecl) == swiftNumParams));
 
     size_t selfParamIndex = MappedDecl->isImportAsInstanceMember()
                                 ? MappedDecl->getSelfIndex()
-                                : ClangDecl->getNumParams();
+                                : getNumParams(ClangDecl);
     for (auto [index, clangParam] : llvm::enumerate(ClangDecl->parameters())) {
-      auto clangParamTy = clangParam->getType();
+      clang::QualType clangParamTy = clangParam->getType();
       int mappedIndex = index < selfParamIndex ? index :
         index > selfParamIndex ? index - 1 :
         SwiftifyInfoPrinter::SELF_PARAM_INDEX;
@@ -383,13 +393,12 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
       auto *CAT = clangParamTy->getAs<clang::CountAttributedType>();
       if (CAT && mappedIndex == SwiftifyInfoPrinter::SELF_PARAM_INDEX) {
         Self.diagnose(HeaderLoc(clangParam->getLocation()),
-                      diag::warn_clang_ignored_bounds_on_self,
-                      getAttributeName(CAT));
-        auto swiftName = ClangDecl->getAttr<clang::SwiftNameAttr>();
+                 diag::warn_clang_ignored_bounds_on_self, getAttributeName(CAT));
+        auto swiftName = ClangDecl->template getAttr<clang::SwiftNameAttr>();
         ASSERT(swiftName &&
                "free function mapped to instance method without swift_name??");
         Self.diagnose(HeaderLoc(swiftName->getLocation()),
-                      diag::note_swift_name_instance_method);
+                 diag::note_swift_name_instance_method);
       } else if (SwiftifiableCAT(clangASTContext, CAT, swiftParamTy)) {
         printer.printCountedBy(CAT, mappedIndex);
         DLOG("  Found bounds info '" << clangParamTy << "' on parameter '"
@@ -401,7 +410,7 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
       paramHasBoundsInfo |= paramIsStdSpan;
 
       bool paramHasLifetimeInfo = false;
-      if (clangParam->hasAttr<clang::NoEscapeAttr>()) {
+      if (clangParam->template hasAttr<clang::NoEscapeAttr>()) {
         DLOG("  Found noescape attribute on parameter '" << *clangParam << "'\n");
         printer.printNonEscaping(mappedIndex);
         paramHasLifetimeInfo = true;
@@ -435,6 +444,47 @@ static bool swiftifyImpl(ClangImporter::Implementation &Self,
   return attachMacro;
 }
 
+class SwiftifyProtocolInfoPrinter : public SwiftifyInfoPrinter {
+private:
+  ClangImporter::Implementation &ImporterImpl;
+
+public:
+  SwiftifyProtocolInfoPrinter(ClangImporter::Implementation &ImporterImpl,
+                              clang::ASTContext &ctx, ASTContext &SwiftContext,
+                              llvm::raw_ostream &out,
+                              MacroDecl &SwiftifyImportDecl)
+      : SwiftifyInfoPrinter(ctx, SwiftContext, out, SwiftifyImportDecl),
+        ImporterImpl(ImporterImpl) {}
+
+  bool printMethod(const FuncDecl *Method) {
+    // don't emit .method() if we know it's going to be empty
+    auto ClangDecl = dyn_cast_or_null<clang::ObjCMethodDecl>(Method->getClangDecl());
+    if (!ClangDecl)
+      return false;
+
+    printSeparator();
+    out << ".method(signature: \"";
+    printMethodSignature(Method);
+    out << "\", paramInfo: [";
+    // reset firstParam inside paramInfo array. At this point firstParam will
+    // always be false, so no need to save the current value.
+    assert(!firstParam);
+    firstParam = true;
+    bool hadAttributes = swiftifyImpl(ImporterImpl, *this, Method, ClangDecl);
+    firstParam = false;
+    out << "])";
+    return hadAttributes;
+  }
+
+private:
+  void printMethodSignature(const FuncDecl *Method) {
+    auto options =
+        PrintOptions::printForDiagnostics(AccessLevel::Private, true);
+    StreamPrinter printer(out);
+    Method->print(printer, options);
+  }
+};
+
 void ClangImporter::Implementation::swiftify(AbstractFunctionDecl *MappedDecl) {
   if (!SwiftContext.LangOpts.hasFeature(Feature::SafeInteropWrappers) ||
       SwiftContext.ClangImporterOpts.DisableSafeInteropWrappers)
@@ -443,6 +493,13 @@ void ClangImporter::Implementation::swiftify(AbstractFunctionDecl *MappedDecl) {
   if (!ClangDecl)
     return;
 
+  {
+    UnaliasedInstantiationVisitor visitor;
+    visitor.TraverseType(ClangDecl->getType());
+    if (visitor.hasUnaliasedInstantiation)
+      return;
+  }
+
   MacroDecl *SwiftifyImportDecl = dyn_cast_or_null<MacroDecl>(getKnownSingleDecl(SwiftContext, "_SwiftifyImport"));
   if (!SwiftifyImportDecl)
     return;
@@ -479,3 +536,43 @@ void ClangImporter::Implementation::swiftify(AbstractFunctionDecl *MappedDecl) {
   }
 }
 
+void ClangImporter::Implementation::swiftifyProtocol(
+    NominalTypeDecl *MappedDecl) {
+  if (!SwiftContext.LangOpts.hasFeature(Feature::SafeInteropWrappers) ||
+      SwiftContext.ClangImporterOpts.DisableSafeInteropWrappers)
+    return;
+  if (!isa<ProtocolDecl, ClassDecl>(MappedDecl))
+    return;
+
+  MacroDecl *SwiftifyImportDecl = dyn_cast_or_null<MacroDecl>(
+      getKnownSingleDecl(SwiftContext, "_SwiftifyImportProtocol"));
+  if (!SwiftifyImportDecl)
+    return;
+
+  DLOG("Checking " << MappedDecl->getName() << " protocol for methods with bounds and lifetime info\n");
+  llvm::SmallString<128> MacroString;
+  {
+    llvm::raw_svector_ostream out(MacroString);
+    out << "@_SwiftifyImportProtocol";
+
+    bool hasBoundsAttributes = false;
+    SwiftifyProtocolInfoPrinter printer(*this, getClangASTContext(),
+                                        SwiftContext, out, *SwiftifyImportDecl);
+    for (Decl *SwiftMember :
+         cast<IterableDeclContext>(MappedDecl)->getAllMembers()) {
+      FuncDecl *SwiftDecl = dyn_cast<FuncDecl>(SwiftMember);
+      if (!SwiftDecl)
+        continue;
+      hasBoundsAttributes |= printer.printMethod(SwiftDecl);
+    }
+
+    if (!hasBoundsAttributes)
+      return;
+    printer.printAvailability();
+    printer.printTypeMapping();
+  }
+
+  DLOG("Attaching safe interop macro: " << MacroString << "\n");
+  importNontrivialAttribute(MappedDecl, MacroString);
+}
+

lib/Sema/TypeCheckMacros.cpp

@@ -1360,6 +1360,7 @@ swift::expandFreestandingMacro(MacroExpansionDecl *med) {
   return macroSourceFile->getBufferID();
 }
 
+
 static SourceFile *evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo,
                                          CustomAttr *attr,
                                          bool passParentContext, MacroRole role,
@@ -1370,7 +1371,11 @@ static SourceFile *evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo,
     dc = attachedTo->getDeclContext();
   } else if (role == MacroRole::Conformance || role == MacroRole::Extension) {
     // Conformance macros always expand to extensions at file-scope.
-    dc = attachedTo->getDeclContext()->getParentSourceFile();
+    dc = attachedTo->getDeclContext();
+    if (!isa<ClangModuleUnit>(dc->getModuleScopeContext()))
+      dc = dc->getParentSourceFile();
+    else // decls imported from clang do not have a SourceFile
+      ASSERT(isa<FileUnit>(dc) && !isa<SourceFile>(dc));
   } else {
     dc = attachedTo->getInnermostDeclContext();
   }

test/Interop/ObjC/swiftify-import/counted-by-protocol.swift

@@ -0,0 +1,133 @@
+// REQUIRES: swift_feature_SafeInteropWrappers
+
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/CountedByProtocol.swiftmodule -I %t/Inputs -enable-experimental-feature SafeInteropWrappers %t/counted-by-protocol.swift -verify -Xcc -Wno-nullability-completeness
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/CountedByProtocol.swiftmodule -I %t/Inputs -enable-experimental-feature SafeInteropWrappers %t/counted-by-protocol.swift -dump-macro-expansions 2> %t/expansions.out
+// RUN: %diff %t/expansions.out %t/expansions.expected
+
+//--- Inputs/module.modulemap
+module CountedByProtocolClang {
+    header "counted-by-protocol.h"
+    export *
+}
+
+//--- Inputs/counted-by-protocol.h
+#pragma once
+
+#define __counted_by(x) __attribute__((__counted_by__(x)))
+
+int foo(int len, int * __counted_by(len) p);
+@protocol CountedByProtocol
+ - (void) simple:(int)len :(int * __counted_by(len))p;
+ - (void) shared:(int)len :(int * __counted_by(len))p1 :(int * __counted_by(len))p2;
+ - (void) complexExpr:(int)len :(int) offset :(int * __counted_by(len - offset))p;
+ - (void) nullUnspecified:(int)len :(int * __counted_by(len) _Null_unspecified)p;
+ - (void) nonnull:(int)len :(int * __counted_by(len) _Nonnull)p;
+ - (void) nullable:(int)len :(int * __counted_by(len) _Nullable)p;
+ - (int * __counted_by(len)) returnPointer:(int)len;
+
+ + (void) staticMethod:(int)len :(int * __counted_by(len))p;
+@end
+
+__attribute__((swift_attr("@_SwiftifyImportProtocol(.method(signature: \"func swiftAttr(_ len: Int32, _ p: UnsafeMutablePointer<Int32>!)\", paramInfo: [.countedBy(pointer: .param(2), count: \"len\")]))")))
+@protocol SwiftAttrProtocol
+ - (void)swiftAttr:(int)len :(int *)p;
+@end
+
+//--- expansions.expected
+@__swiftmacro_So17CountedByProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension CountedByProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func simple(_ p: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p.count)!
+        return unsafe simple(len, p.baseAddress!)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func shared(_ p1: UnsafeMutableBufferPointer<Int32>, _ p2: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p1.count)!
+        if p2.count != len {
+          fatalError("bounds check failure in shared: expected \(len) but got \(p2.count)")
+        }
+        return unsafe shared(len, p1.baseAddress!, p2.baseAddress!)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func complexExpr(_ len: Int32, _ offset: Int32, _ p: UnsafeMutableBufferPointer<Int32>) {
+        let _pCount = p.count
+        if _pCount != len - offset {
+          fatalError("bounds check failure in complexExpr: expected \(len - offset) but got \(_pCount)")
+        }
+        return unsafe complexExpr(len, offset, p.baseAddress!)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func nullUnspecified(_ p: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p.count)!
+        return unsafe nullUnspecified(len, p.baseAddress!)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func nonnull(_ p: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p.count)!
+        return unsafe nonnull(len, p.baseAddress!)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func nullable(_ p: UnsafeMutableBufferPointer<Int32>?) {
+        let len = Int32(exactly: unsafe p?.count ?? 0)!
+        return unsafe nullable(len, p?.baseAddress)
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func returnPointer(_ len: Int32) -> UnsafeMutableBufferPointer<Int32> {
+        return unsafe UnsafeMutableBufferPointer<Int32>(start: unsafe returnPointer(len), count: Int(len))
+    }
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload
+        static public func staticMethod(_ p: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p.count)!
+        return unsafe staticMethod(len, p.baseAddress!)
+    }
+}
+------------------------------
+@__swiftmacro_So17SwiftAttrProtocol015_SwiftifyImportC0fMe_.swift
+------------------------------
+extension SwiftAttrProtocol {
+    /// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public
+        func swiftAttr(_ p: UnsafeMutableBufferPointer<Int32>) {
+        let len = Int32(exactly: p.count)!
+        return unsafe swiftAttr(len, p.baseAddress!)
+    }
+}
+------------------------------
+@__swiftmacro_So3foo15_SwiftifyImportfMp_.swift
+------------------------------
+/// This is an auto-generated wrapper for safer interop
+@_alwaysEmitIntoClient @_disfavoredOverload public func foo(_ p: UnsafeMutableBufferPointer<Int32>) -> Int32 {
+    let len = Int32(exactly: p.count)!
+    return unsafe foo(len, p.baseAddress!)
+}
+------------------------------
+//--- counted-by-protocol.swift
+import CountedByProtocolClang
+
+@inlinable
+public func call(p: UnsafeMutableBufferPointer<CInt>, x: CInt, y: CInt, a: CountedByProtocol, b: SwiftAttrProtocol) {
+  a.simple(p)
+  a.shared(p, p)
+  a.complexExpr(x, y, p)
+  a.nullUnspecified(p)
+  a.nonnull(p)
+  a.nullable(p)
+  let _: UnsafeMutableBufferPointer<CInt> = a.returnPointer(x)
+  let r2 = a.returnPointer(x)
+  let _: UnsafeMutablePointer<CInt>? = r2 // make sure the original is the favored overload
+  b.swiftAttr(p)
+  let _ = foo(p)
+}