[APINotes] Add support for bounds safety annotations

This adds support for annotating function parameters with __counted_by,
__sized_by, __counted_by_or_null, __sized_by_or_null, and __ended_by,
using API notes. The main content of handlePtrCountedByEndedByAttr is
extracted to applyPtrCountedByEndedByAttr and decoupled from ParsedAttr.
The helper function ParseBoundsAttributeArgFromString is added to make
it possible to parse count expressions from SemaAPINotes.

The current implementation of __terminated_by/__null_terminated makes it
harder to extract from the iterative type processing, but since it
doesn't require any extra context to parse the attribute, it can be
applied using the normal Type override instead.

rdar://139830881

Author: hnrklssn

clang/include/clang/APINotes/Types.h

@@ -300,6 +300,68 @@ inline bool operator!=(const ContextInfo &LHS, const ContextInfo &RHS) {
   return !(LHS == RHS);
 }
 
+class BoundsSafetyInfo {
+public:
+  enum class BoundsSafetyKind {
+    CountedBy,
+    CountedByOrNull,
+    SizedBy,
+    SizedByOrNull,
+    EndedBy,
+  };
+
+private:
+  /// Whether this property has been audited for nullability.
+  LLVM_PREFERRED_TYPE(bool)
+  unsigned KindAudited : 1;
+
+  /// The kind of nullability for this property. Only valid if the nullability
+  /// has been audited.
+  LLVM_PREFERRED_TYPE(BoundsSafetyKind)
+  unsigned Kind : 3;
+
+  LLVM_PREFERRED_TYPE(bool)
+  unsigned LevelAudited : 1;
+
+  unsigned Level : 3;
+
+public:
+  std::string ExternalBounds;
+
+  BoundsSafetyInfo()
+      : KindAudited(false), Kind(0), LevelAudited(false), Level(0),
+        ExternalBounds("") {}
+
+  std::optional<BoundsSafetyKind> getKind() const {
+    return KindAudited ? std::optional<BoundsSafetyKind>(
+                             static_cast<BoundsSafetyKind>(Kind))
+                       : std::nullopt;
+  }
+
+  void setKindAudited(BoundsSafetyKind kind) {
+    KindAudited = true;
+    Kind = static_cast<unsigned>(kind);
+  }
+
+  std::optional<unsigned> getLevel() const {
+    return LevelAudited ? std::optional<unsigned>(Level) : std::nullopt;
+  }
+
+  void setLevelAudited(unsigned level) {
+    LevelAudited = true;
+    Level = level;
+  }
+
+  friend bool operator==(const BoundsSafetyInfo &, const BoundsSafetyInfo &);
+};
+
+inline bool operator==(const BoundsSafetyInfo &LHS,
+                       const BoundsSafetyInfo &RHS) {
+  return LHS.KindAudited == RHS.KindAudited && LHS.Kind == RHS.Kind &&
+         LHS.LevelAudited == RHS.LevelAudited && LHS.Level == RHS.Level &&
+         LHS.ExternalBounds == RHS.ExternalBounds;
+}
+
 /// API notes for a variable/property.
 class VariableInfo : public CommonEntityInfo {
   /// Whether this property has been audited for nullability.
@@ -439,10 +501,12 @@ class ParamInfo : public VariableInfo {
   unsigned RawRetainCountConvention : 3;
 
 public:
+  std::optional<BoundsSafetyInfo> BoundsSafety;
+
   ParamInfo()
       : NoEscapeSpecified(false), NoEscape(false),
         LifetimeboundSpecified(false), Lifetimebound(false),
-        RawRetainCountConvention() {}
+        RawRetainCountConvention(), BoundsSafety(std::nullopt) {}
 
   std::optional<bool> isNoEscape() const {
     return NoEscapeSpecified ? std::optional<bool>(NoEscape) : std::nullopt;
@@ -488,6 +552,9 @@ class ParamInfo : public VariableInfo {
     if (!RawRetainCountConvention)
       RawRetainCountConvention = RHS.RawRetainCountConvention;
 
+    if (!BoundsSafety)
+      BoundsSafety = RHS.BoundsSafety;
+
     return *this;
   }
 
@@ -502,7 +569,8 @@ inline bool operator==(const ParamInfo &LHS, const ParamInfo &RHS) {
          LHS.NoEscape == RHS.NoEscape &&
          LHS.LifetimeboundSpecified == RHS.LifetimeboundSpecified &&
          LHS.Lifetimebound == RHS.Lifetimebound &&
-         LHS.RawRetainCountConvention == RHS.RawRetainCountConvention;
+         LHS.RawRetainCountConvention == RHS.RawRetainCountConvention &&
+         LHS.BoundsSafety == RHS.BoundsSafety;
 }
 
 inline bool operator!=(const ParamInfo &LHS, const ParamInfo &RHS) {

clang/include/clang/Parse/Parser.h

@@ -3924,6 +3924,10 @@ class Parser : public CodeCompletionHandler {
   /// \param IncludeLoc The location at which this parse was triggered.
   TypeResult ParseTypeFromString(StringRef TypeStr, StringRef Context,
                                  SourceLocation IncludeLoc);
+  ExprResult ParseBoundsAttributeArgFromString(StringRef ExprStr,
+                                               StringRef Context,
+                                               Decl *ParentDecl,
+                                               SourceLocation IncludeLoc);
 
   //===--------------------------------------------------------------------===//
   // Modules

clang/include/clang/Sema/Sema.h

@@ -1131,6 +1131,10 @@ class Sema final : public SemaBase {
   std::function<TypeResult(StringRef, StringRef, SourceLocation)>
       ParseTypeFromStringCallback;
 
+  /// Callback to the parser to parse a type expressed as a string.
+  std::function<ExprResult(StringRef, StringRef, Decl *, SourceLocation)>
+      ParseBoundsAttributeArgFromStringCallback;
+
   /// VAListTagName - The declaration name corresponding to __va_list_tag.
   /// This is used as part of a hack to omit that class from ADL results.
   DeclarationName VAListTagName;
@@ -15504,6 +15508,11 @@ class Sema final : public SemaBase {
       llvm::SmallVectorImpl<TypeCoupledDeclRefInfo> &Decls, bool CountInBytes,
       bool OrNull);
 
+  void applyPtrCountedByEndedByAttr(Decl *D, unsigned Level,
+                                    AttributeCommonInfo::Kind Kind,
+                                    Expr *AttrArg, SourceLocation Loc,
+                                    SourceRange Range, StringRef DiagName);
+
   /* TO_UPSTREAM(BoundsSafety) ON*/
   /// Perform Bounds Safety Semantic checks for assigning to a `__counted_by` or
   /// `__counted_by_or_null` pointer type \param LHSTy.

clang/lib/APINotes/APINotesFormat.h

@@ -24,7 +24,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// API notes file minor version number.
 ///
 /// When the format changes IN ANY WAY, this number should be incremented.
-const uint16_t VERSION_MINOR = 34; // SwiftReturnOwnership
+const uint16_t VERSION_MINOR = 35; // BoundsSafety
 
 const uint8_t kSwiftConforms = 1;
 const uint8_t kSwiftDoesNotConform = 2;

clang/lib/APINotes/APINotesReader.cpp

@@ -322,6 +322,44 @@ class FieldTableInfo
   }
 };
 
+BoundsSafetyInfo::BoundsSafetyKind readKindFlag(uint8_t kind) {
+  switch (kind) {
+  case 0x00:
+    return BoundsSafetyInfo::BoundsSafetyKind::CountedBy;
+  case 0x01:
+    return BoundsSafetyInfo::BoundsSafetyKind::CountedByOrNull;
+  case 0x02:
+    return BoundsSafetyInfo::BoundsSafetyKind::SizedBy;
+  case 0x03:
+    return BoundsSafetyInfo::BoundsSafetyKind::SizedByOrNull;
+  case 0x04:
+    return BoundsSafetyInfo::BoundsSafetyKind::EndedBy;
+  default:
+    llvm_unreachable("invalid bounds safety kind");
+  }
+}
+
+/// Read serialized BoundsSafetyInfo.
+void ReadBoundsSafetyInfo(const uint8_t *&Data, BoundsSafetyInfo &Info) {
+  uint8_t Payload = endian::readNext<uint8_t, llvm::endianness::little>(Data);
+
+  if (Payload & 0x01) {
+    uint8_t Level = (Payload >> 1) & 0x7;
+    Info.setLevelAudited(Level);
+  }
+  Payload >>= 4;
+
+  if (Payload & 0x01) {
+    uint8_t Kind = (Payload >> 1) & 0x7;
+    Info.setKindAudited(readKindFlag(Kind));
+  }
+
+  uint16_t ExternalBoundsLen =
+      endian::readNext<uint16_t, llvm::endianness::little>(Data);
+  Info.ExternalBounds = std::string(Data, Data + ExternalBoundsLen);
+  Data += ExternalBoundsLen;
+}
+
 /// Read serialized ParamInfo.
 void ReadParamInfo(const uint8_t *&Data, ParamInfo &Info) {
   ReadVariableInfo(Data, Info);
@@ -338,7 +376,11 @@ void ReadParamInfo(const uint8_t *&Data, ParamInfo &Info) {
   if (Payload & 0x01)
     Info.setNoEscape(Payload & 0x02);
   Payload >>= 2;
-  assert(Payload == 0 && "Bad API notes");
+  if (Payload & 0x01) {
+    BoundsSafetyInfo BSI;
+    ReadBoundsSafetyInfo(Data, BSI);
+    Info.BoundsSafety = BSI;
+  }
 }
 
 /// Read serialized FunctionInfo.

clang/lib/APINotes/APINotesWriter.cpp

@@ -1075,14 +1075,58 @@ void APINotesWriter::Implementation::writeGlobalVariableBlock(
 }
 
 namespace {
+uint8_t getKindFlag(BoundsSafetyInfo::BoundsSafetyKind kind) {
+  switch (kind) {
+  case BoundsSafetyInfo::BoundsSafetyKind::CountedBy:
+    return 0x00;
+  case BoundsSafetyInfo::BoundsSafetyKind::CountedByOrNull:
+    return 0x01;
+  case BoundsSafetyInfo::BoundsSafetyKind::SizedBy:
+    return 0x02;
+  case BoundsSafetyInfo::BoundsSafetyKind::SizedByOrNull:
+    return 0x03;
+  case BoundsSafetyInfo::BoundsSafetyKind::EndedBy:
+    return 0x04;
+  }
+}
+
+void emitBoundsSafetyInfo(raw_ostream &OS, const BoundsSafetyInfo &BSI) {
+  llvm::support::endian::Writer writer(OS, llvm::endianness::little);
+  uint8_t flags = 0;
+  if (auto kind = BSI.getKind()) {
+    flags |= 0x01;                    // 1 bit
+    flags |= getKindFlag(*kind) << 1; // 3 bits
+  }
+  flags <<= 4;
+  if (auto level = BSI.getLevel()) {
+    flags |= 0x01;        // 1 bit
+    flags |= *level << 1; // 3 bits
+  }
+
+  writer.write<uint8_t>(flags);
+  writer.write<uint16_t>(BSI.ExternalBounds.size());
+  writer.write(
+      ArrayRef<char>{BSI.ExternalBounds.data(), BSI.ExternalBounds.size()});
+}
+
+unsigned getBoundsSafetyInfoSize(const BoundsSafetyInfo &BSI) {
+  return 1 + sizeof(uint16_t) + BSI.ExternalBounds.size();
+}
+
 unsigned getParamInfoSize(const ParamInfo &PI) {
-  return getVariableInfoSize(PI) + 1;
+  unsigned BSISize = 0;
+  if (auto BSI = PI.BoundsSafety)
+    BSISize = getBoundsSafetyInfoSize(*BSI);
+  return getVariableInfoSize(PI) + 1 + BSISize;
 }
 
 void emitParamInfo(raw_ostream &OS, const ParamInfo &PI) {
   emitVariableInfo(OS, PI);
 
   uint8_t flags = 0;
+  if (PI.BoundsSafety)
+    flags |= 0x01;
+  flags <<= 2;
   if (auto noescape = PI.isNoEscape()) {
     flags |= 0x01;
     if (*noescape)
@@ -1100,6 +1144,8 @@ void emitParamInfo(raw_ostream &OS, const ParamInfo &PI) {
 
   llvm::support::endian::Writer writer(OS, llvm::endianness::little);
   writer.write<uint8_t>(flags);
+  if (auto BSI = PI.BoundsSafety)
+    emitBoundsSafetyInfo(OS, *PI.BoundsSafety);
 }
 
 /// Retrieve the serialized size of the given FunctionInfo, for use in on-disk

clang/lib/APINotes/APINotesYAMLCompiler.cpp

@@ -150,6 +150,31 @@ enum class APIAvailability {
 };
 } // namespace
 
+namespace {
+struct BoundsSafety {
+  BoundsSafetyInfo::BoundsSafetyKind Kind;
+  unsigned Level = 0;
+  StringRef BoundsExpr = "";
+};
+} // namespace
+
+namespace llvm {
+namespace yaml {
+template <> struct ScalarEnumerationTraits<BoundsSafetyInfo::BoundsSafetyKind> {
+  static void enumeration(IO &IO, BoundsSafetyInfo::BoundsSafetyKind &AA) {
+    IO.enumCase(AA, "counted_by",
+                BoundsSafetyInfo::BoundsSafetyKind::CountedBy);
+    IO.enumCase(AA, "counted_by_or_null",
+                BoundsSafetyInfo::BoundsSafetyKind::CountedByOrNull);
+    IO.enumCase(AA, "sized_by", BoundsSafetyInfo::BoundsSafetyKind::SizedBy);
+    IO.enumCase(AA, "sized_by_or_null",
+                BoundsSafetyInfo::BoundsSafetyKind::SizedByOrNull);
+    IO.enumCase(AA, "ended_by", BoundsSafetyInfo::BoundsSafetyKind::EndedBy);
+  }
+};
+} // namespace yaml
+} // namespace llvm
+
 namespace llvm {
 namespace yaml {
 template <> struct ScalarEnumerationTraits<APIAvailability> {
@@ -188,6 +213,7 @@ struct Param {
   std::optional<NullabilityKind> Nullability;
   std::optional<RetainCountConventionKind> RetainCountConvention;
   StringRef Type;
+  BoundsSafety BoundsSafety;
 };
 
 typedef std::vector<Param> ParamsSeq;
@@ -238,6 +264,15 @@ template <> struct MappingTraits<Param> {
     IO.mapOptional("NoEscape", P.NoEscape);
     IO.mapOptional("Lifetimebound", P.Lifetimebound);
     IO.mapOptional("Type", P.Type, StringRef(""));
+    IO.mapOptional("BoundsSafety", P.BoundsSafety);
+  }
+};
+
+template <> struct MappingTraits<BoundsSafety> {
+  static void mapping(IO &IO, BoundsSafety &BS) {
+    IO.mapRequired("Kind", BS.Kind);
+    IO.mapRequired("BoundedBy", BS.BoundsExpr);
+    IO.mapOptional("Level", BS.Level, 0);
   }
 };
 } // namespace yaml
@@ -862,6 +897,11 @@ class YAMLConverter {
       PI.setLifetimebound(P.Lifetimebound);
       PI.setType(std::string(P.Type));
       PI.setRetainCountConvention(P.RetainCountConvention);
+      BoundsSafetyInfo BSI;
+      BSI.setKindAudited(P.BoundsSafety.Kind);
+      BSI.setLevelAudited(P.BoundsSafety.Level);
+      BSI.ExternalBounds = P.BoundsSafety.BoundsExpr.str();
+      PI.BoundsSafety = BSI;
       if (static_cast<int>(OutInfo.Params.size()) <= P.Position)
         OutInfo.Params.resize(P.Position + 1);
       if (P.Position == -1)