Infer @PointerBounds macro from clang __counted_by parameters

This results in an automatic wrapper function with safe pointer types
when the imported function has bounds attributes. This exercises similar
pathways as the recently added functionality for specifying macros from
swift_attr, and fixes some bugs related to macro source file management.

rdar://97942270

Author: hnrklssn

include/swift/AST/DiagnosticsClangImporter.def

@@ -87,6 +87,11 @@ NOTE(unresolvable_clang_decl_is_a_framework_bug,none,
 WARNING(clang_swift_attr_unhandled,none,
         "ignoring unknown Swift attribute or modifier '%0'", (StringRef))
 
+WARNING(clang_pointer_bounds_unhandled, none,
+        "ignoring unknown pointer bound attributes on '%0' because the "
+        "generated macro could not be parsed: '%1'",
+        (DeclName, StringRef))
+
 WARNING(clang_error_code_must_be_sendable,none,
         "cannot make error code type '%0' non-sendable because Swift errors "
         "are always sendable", (StringRef))

include/swift/AST/TypeCheckRequests.h

@@ -775,6 +775,27 @@ class PrettyPrintDeclRequest
   bool isCached() const { return true; }
 };
 
+/// Pretty-print the given attribute into a buffer and return a source
+/// location that refers to the attribute in that buffer.
+class PrettyPrintCustomAttrRequest
+    : public SimpleRequest<PrettyPrintCustomAttrRequest,
+                           SourceLoc(const CustomAttr *, const Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  SourceLoc evaluate(Evaluator &eval, const CustomAttr *attr,
+                     const Decl *decl) const;
+
+public:
+  // Caching
+  bool isCached() const { return true; }
+};
+
 // Find the type in the cache or look it up
 class DefaultTypeRequest
     : public SimpleRequest<DefaultTypeRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -81,6 +81,8 @@ SWIFT_REQUEST(TypeChecker, DefaultTypeRequest,
               NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, PrettyPrintDeclRequest,
               SourceLoc(const Decl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PrettyPrintCustomAttrRequest,
+              SourceLoc(const CustomAttr *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DifferentiableAttributeTypeCheckRequest,
               IndexSubset *(DifferentiableAttr *),
               SeparatelyCached, NoLocationInfo)

lib/ClangImporter/ClangImporter.cpp

@@ -576,6 +576,10 @@ void importer::getNormalInvocationArguments(
     }
   }
 
+#ifdef SWIFT_ENABLE_EXPERIMENTAL_POINTER_BOUNDS
+  invokationsArgStrs.push_back("-fexperimental-bounds-safety-attributes");
+#endif
+
   // Set C language options.
   if (triple.isOSDarwin()) {
     invocationArgStrs.insert(invocationArgStrs.end(), {

lib/ClangImporter/ImportDecl.cpp

@@ -125,6 +125,7 @@ createFuncOrAccessor(ClangImporter::Implementation &impl, SourceLoc funcLoc,
                                     genericParams, dc, clangNode);
   }
   impl.importSwiftAttrAttributes(decl);
+  impl.importBoundsAttributes(decl);
   return decl;
 }
 
@@ -8532,6 +8533,107 @@ ClangImporter::Implementation::importSwiftAttrAttributes(Decl *MappedDecl) {
   }
 }
 
+namespace {
+class PointerParamInfo {
+public:
+  virtual void print(clang::ASTContext &ctx, llvm::raw_ostream &out) const = 0;
+  virtual ~PointerParamInfo() {}
+};
+
+class CountedByParam : public PointerParamInfo {
+public:
+  size_t pointerIndex;
+  clang::Expr *countExpr;
+  bool isSizedBy;
+  CountedByParam(size_t idx, clang::Expr *E, bool sizedBy)
+      : pointerIndex(idx), countExpr(E), isSizedBy(sizedBy) {}
+
+  virtual void print(clang::ASTContext &ctx,
+                     llvm::raw_ostream &out) const override {
+    out << ".";
+    if (isSizedBy)
+      out << "sizedBy";
+    else
+      out << "countedBy";
+    out << "(pointer: " << pointerIndex << ", ";
+    if (isSizedBy)
+      out << "size";
+    else
+      out << "count";
+    out << ": \"";
+    countExpr->printPretty(
+        out, {}, {ctx.getLangOpts()}); // TODO: map clang::Expr to Swift Expr
+    out << "\")";
+  }
+  virtual ~CountedByParam() {}
+};
+} // namespace
+
+void ClangImporter::Implementation::importBoundsAttributes(
+    FuncDecl *MappedDecl) {
+  auto ClangDecl =
+      dyn_cast_or_null<clang::FunctionDecl>(MappedDecl->getClangDecl());
+  if (!ClangDecl)
+    return;
+
+  SmallVector<PointerParamInfo *, 4> BoundsInfo;
+  size_t parameterIndex = 1;
+  for (auto param : ClangDecl->parameters()) {
+    if (auto CAT = param->getType()->getAs<clang::CountAttributedType>()) {
+      BoundsInfo.push_back(new CountedByParam(
+          parameterIndex, CAT->getCountExpr(), CAT->isCountInBytes()));
+    }
+    parameterIndex++;
+  }
+  if (BoundsInfo.empty())
+    return;
+
+  llvm::SmallString<128> MacroString;
+  {
+    llvm::raw_svector_ostream out(MacroString);
+
+    out << "@PointerBounds(";
+    for (size_t i = 0; i < BoundsInfo.size(); i++) {
+      BoundsInfo[i]->print(getClangASTContext(), out);
+      if (i + 1 < BoundsInfo.size()) {
+        out << ", ";
+      }
+    }
+    out << ")";
+  }
+
+  // Dig out a buffer with the attribute text.
+  unsigned bufferID = getClangSwiftAttrSourceBuffer(MacroString);
+
+  // Dig out a source file we can use for parsing.
+  auto &sourceFile = getClangSwiftAttrSourceFile(
+      *MappedDecl->getDeclContext()->getParentModule(), bufferID);
+
+  // Spin up a parser.
+  swift::Parser parser(bufferID, sourceFile, &SwiftContext.Diags, nullptr,
+                       nullptr);
+  // Prime the lexer.
+  parser.consumeTokenWithoutFeedingReceiver();
+
+  bool hadError = false;
+  assert(parser.Tok.is(tok::at_sign));
+  SourceLoc atEndLoc = parser.Tok.getRange().getEnd();
+  SourceLoc atLoc = parser.consumeToken(tok::at_sign);
+  DeclContext *DC = MappedDecl->getParent();
+  auto initContext = PatternBindingInitializer::create(DC);
+  hadError = parser
+                 .parseDeclAttribute(MappedDecl->getAttrs(), atLoc, atEndLoc,
+                                     initContext,
+                                     /*isFromClangAttribute=*/true)
+                 .isError();
+  if (hadError) {
+    HeaderLoc attrLoc(ClangDecl->getLocation());
+    diagnose(attrLoc, diag::clang_pointer_bounds_unhandled,
+             MappedDecl->getName(), MacroString);
+    return;
+  }
+}
+
 static bool isUsingMacroName(clang::SourceManager &SM,
                              clang::SourceLocation loc,
                              StringRef MacroName) {

lib/ClangImporter/ImportType.cpp

@@ -2424,6 +2424,10 @@ ClangImporter::Implementation::importParameterType(
         }
       }
     }
+  } else if (auto CAT = dyn_cast<clang::CountAttributedType>(paramTy)) {
+    // Treat as a normal pointer. importBoundsAttributes() will generate a safe
+    // overload later.
+    paramTy = CAT->desugar();
   } else if (isa<clang::PointerType>(paramTy) &&
              isa<clang::TemplateTypeParmType>(paramTy->getPointeeType())) {
     auto pointeeType = paramTy->getPointeeType();

lib/ClangImporter/ImporterImpl.h

@@ -1732,6 +1732,7 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   }
 
   void importSwiftAttrAttributes(Decl *decl);
+  void importBoundsAttributes(FuncDecl *MappedDecl);
 
   /// Find the lookup table that corresponds to the given Clang module.
   ///

lib/Sema/TypeCheckDecl.cpp

@@ -3247,9 +3247,8 @@ namespace {
   };
 }
 
-SourceLoc PrettyPrintDeclRequest::evaluate(Evaluator &eval, const Decl *decl) const {
-  // Conjure a buffer name for this declaration.
-  SmallVector<std::string, 4> nameComponents;
+static void collectQualifiedDeclNameComponents(
+    const Decl *decl, SmallVectorImpl<std::string> &nameComponents) {
   DeclContext *dc;
   if (auto valueDecl = dyn_cast<ValueDecl>(decl)) {
     nameComponents.push_back(valueDecl->getBaseName().userFacingName().str());
@@ -3296,7 +3295,13 @@ SourceLoc PrettyPrintDeclRequest::evaluate(Evaluator &eval, const Decl *decl) co
 
     dc = dc->getParent();
   }
+}
 
+SourceLoc PrettyPrintDeclRequest::evaluate(Evaluator &eval,
+                                           const Decl *decl) const {
+  // Conjure a buffer name for this declaration.
+  SmallVector<std::string, 4> nameComponents;
+  collectQualifiedDeclNameComponents(decl, nameComponents);
 
   std::string bufferName;
   {
@@ -3391,3 +3396,62 @@ SourceLoc PrettyPrintDeclRequest::evaluate(Evaluator &eval, const Decl *decl) co
 
   return memBufferStartLoc.getAdvancedLoc(targetDeclOffsetInBuffer);
 }
+
+//----------------------------------------------------------------------------//
+// PrettyPrintCustomAttrRequest
+//----------------------------------------------------------------------------//
+
+SourceLoc PrettyPrintCustomAttrRequest::evaluate(Evaluator &eval,
+                                                 const CustomAttr *attr,
+                                                 const Decl *decl) const {
+  // Conjure a buffer name for this declaration.
+  SmallVector<std::string, 4> nameComponents;
+  nameComponents.push_back(attr->getAttrName().str());
+  collectQualifiedDeclNameComponents(decl, nameComponents);
+
+  std::string bufferName;
+  {
+    llvm::raw_string_ostream out(bufferName);
+    for (auto iter = nameComponents.rbegin(); iter != nameComponents.rend();
+         ++iter) {
+      out << *iter;
+
+      if (iter + 1 != nameComponents.rend())
+        out << ".";
+    }
+  }
+
+  // Render the buffer contents.
+  ASTContext &ctx = decl->getASTContext();
+  llvm::SmallString<128> bufferContents;
+  {
+    llvm::raw_svector_ostream out(bufferContents);
+    StreamPrinter P(out);
+
+    // Print this macro attribute.
+    auto options = PrintOptions::printForDiagnostics(
+        getBufferAccessLevel(decl), ctx.TypeCheckerOpts.PrintFullConvention);
+    attr->print(P, options, decl);
+  }
+
+  // Build a buffer with the pretty-printed macro attribute.
+  SourceManager &sourceMgr = ctx.SourceMgr;
+  auto bufferID = sourceMgr.addMemBufferCopy(bufferContents, bufferName);
+  auto memBufferStartLoc = sourceMgr.getLocForBufferStart(bufferID);
+
+  // Note that this is a pretty-printed buffer.
+  sourceMgr.setGeneratedSourceInfo(
+      bufferID,
+      GeneratedSourceInfo{
+          GeneratedSourceInfo::PrettyPrinted, CharSourceRange(),
+          CharSourceRange(memBufferStartLoc, bufferContents.size()),
+          ASTNode(const_cast<Decl *>(decl)).getOpaqueValue(), nullptr});
+
+  // Add a source file for the buffer.
+  auto moduleDecl = decl->getDeclContext()->getParentModule();
+  auto sourceFile =
+      new (ctx) SourceFile(*moduleDecl, SourceFileKind::Library, bufferID);
+  sourceFile->setImports({});
+
+  return memBufferStartLoc;
+}

lib/Sema/TypeCheckMacros.cpp

@@ -1017,21 +1017,21 @@ createMacroSourceFile(std::unique_ptr<llvm::MemoryBuffer> buffer,
   // Create a new source buffer with the contents of the expanded macro.
   unsigned macroBufferID = sourceMgr.addNewSourceBuffer(std::move(buffer));
   auto macroBufferRange = sourceMgr.getRangeForBuffer(macroBufferID);
-  GeneratedSourceInfo sourceInfo{generatedSourceKind,
-                                 generatedOriginalSourceRange,
-                                 macroBufferRange,
-                                 target.getOpaqueValue(),
-                                 dc,
-                                 attr,
-                                 macroName.c_str()
-  };
-  sourceMgr.setGeneratedSourceInfo(macroBufferID, sourceInfo);
+  auto macroContext = dc;
 
   // Create a source file to hold the macro buffer. This is automatically
   // registered with the enclosing module.
   auto macroSourceFile = new (ctx) SourceFile(
       *dc->getParentModule(), SourceFileKind::MacroExpansion, macroBufferID,
       /*parsingOpts=*/{}, /*isPrimary=*/false);
+  if (isa<FileUnit>(macroContext) && !isa<SourceFile>(macroContext))
+    macroContext = macroSourceFile;
+  GeneratedSourceInfo sourceInfo{
+      generatedSourceKind, generatedOriginalSourceRange,
+      macroBufferRange,    target.getOpaqueValue(),
+      macroContext,        attr,
+      macroName.c_str()};
+  sourceMgr.setGeneratedSourceInfo(macroBufferID, sourceInfo);
   if (auto parentSourceFile = dc->getParentSourceFile())
     macroSourceFile->setImports(parentSourceFile->getImports());
   else if (auto clangModuleUnit =
@@ -1359,8 +1359,20 @@ static SourceFile *evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo,
 
   auto moduleDecl = dc->getParentModule();
 
-  auto attrSourceFile =
-    moduleDecl->getSourceFileContainingLocation(attr->AtLoc);
+  // If the attribute has no source location and is attached to a declaration
+  // from a Clang module, pretty-print the attribute and use that location.
+  // This is relevant for Swift macros inferred from Clang attributes,
+  // since they don't have a source representation.
+  SourceLoc attrLoc = attr->AtLoc;
+  if (attrLoc.isInvalid() &&
+      isa<ClangModuleUnit>(dc->getModuleScopeContext())) {
+    attrLoc = evaluateOrDefault(ctx.evaluator,
+                                PrettyPrintCustomAttrRequest{attr, attachedTo},
+                                SourceLoc());
+    assert(attrLoc);
+  }
+
+  auto attrSourceFile = moduleDecl->getSourceFileContainingLocation(attrLoc);
   if (!attrSourceFile)
     return nullptr;
 
@@ -1524,7 +1536,7 @@ static SourceFile *evaluateAttachedMacro(MacroDecl *macro, Decl *attachedTo,
     swift_Macros_expandAttachedMacro(
         &ctx.Diags, externalDef.get(), discriminator->c_str(),
         extendedType.c_str(), conformanceList.c_str(), getRawMacroRole(role),
-        astGenAttrSourceFile, attr->AtLoc.getOpaquePointerValue(),
+        astGenAttrSourceFile, attrLoc.getOpaquePointerValue(),
         astGenDeclSourceFile, startLoc.getOpaquePointerValue(),
         astGenParentDeclSourceFile, parentDeclLoc, &evaluatedSourceOut);
     if (!evaluatedSourceOut.unbridged().data())

stdlib/cmake/modules/AddSwiftStdlib.cmake

@@ -478,6 +478,10 @@ function(_add_target_variant_c_compile_flags)
     list(APPEND result "-DSWIFT_USE_OS_TRACE_LAZY_INIT")
   endif()
 
+  if(SWIFT_ENABLE_EXPERIMENTAL_POINTER_BOUNDS)
+    list(APPEND result "-DSWIFT_ENABLE_EXPERIMENTAL_POINTER_BOUNDS")
+  endif()
+
   list(APPEND result ${SWIFT_STDLIB_EXTRA_C_COMPILE_FLAGS})
 
   set("${CFLAGS_RESULT_VAR_NAME}" "${result}" PARENT_SCOPE)

test/Interop/C/pointer-bounds/Inputs/counted-by.h

@@ -0,0 +1,18 @@
+#pragma once
+
+#define __counted_by(x) __attribute__((__counted_by__(x)))
+
+void simple(int len, int * __counted_by(len) p);
+
+void swiftAttr(int len, int * p) __attribute__((swift_attr("@PointerBounds(.countedBy(pointer: 2, count: \"len\"))")));
+
+void shared(int len, int * __counted_by(len) p1, int * __counted_by(len) p2);
+
+// TODO: test this when counted_by expressions land
+// 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);

test/Interop/C/pointer-bounds/Inputs/module.modulemap

@@ -0,0 +1,5 @@
+module CountedByClang {
+    header "counted-by.h"
+    export *
+}
+