[BoundsSafety][APINotes] Add support for bounds attribute on return value

Previously you could only attach bounds attributes to function
parameters. Now return values are also supported, although not yet for
ObjectiveC methods. Adds test case for bounds safety API notes in C++,
and fixes bug where late parsed bounds attributes on C++ method
parameters are parsed twice, due to not being removed from the late
attributes list after parsing.

rdar://146332875

Author: hnrklssn

clang/include/clang/APINotes/Types.h

@@ -623,6 +623,11 @@ class FunctionInfo : public CommonEntityInfo {
   /// Ownership convention for return value
   std::string SwiftReturnOwnership;
 
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  /// Bounds annotations for the return value
+  std::optional<BoundsSafetyInfo> ReturnBoundsSafety;
+  /* TO_UPSTREAM(BoundsSafety) OFF */
+
   /// The function parameters.
   std::vector<ParamInfo> Params;
 

clang/include/clang/Parse/Parser.h

@@ -1655,6 +1655,9 @@ class Parser : public CodeCompletionHandler {
                                 ParsedAttributes *OutAttrs = nullptr);
   void ParseLexedAttribute(LateParsedAttribute &LA,
                            bool EnterScope, bool OnDefinition);
+  /// Note: it's important the the LateParsedAttribute is removed after parsing
+  /// to prevent being reparsed later in the wrong context. This is handled
+  /// automatically by Parser::ParseLexedCAttributeList.
   void ParseLexedCAttribute(LateParsedAttribute &LA, bool EnterScope,
                             ParsedAttributes *OutAttrs = nullptr);
   void ParseLexedMethodDeclarations(ParsingClass &Class);
@@ -3944,10 +3947,11 @@ class Parser : public CodeCompletionHandler {
   /// \param ParentDecl If a function or method is provided, the parameters are
   ///   added to the current parsing context.
   /// \param IncludeLoc The location at which this parse was triggered.
-  ExprResult ParseBoundsAttributeArgFromString(StringRef ExprStr,
-                                               StringRef Context,
-                                               Decl *ParentDecl,
-                                               SourceLocation IncludeLoc);
+  /// \param Payload Info to be passed along to the completion handler before
+  /// leaving the current context
+  void ParseBoundsAttributeArgFromString(
+      StringRef ExprStr, StringRef Context, Decl *ParentDecl,
+      SourceLocation IncludeLoc, Sema::IncompleteBoundsAttributeInfo Payload);
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
   //===--------------------------------------------------------------------===//

clang/include/clang/Sema/Sema.h

@@ -1132,8 +1132,19 @@ class Sema final : public SemaBase {
       ParseTypeFromStringCallback;
 
   /* TO_UPSTREAM(BoundsSafety) ON */
+  // Info for completing the processing of the attribute after parsing its
+  // argument
+  struct IncompleteBoundsAttributeInfo {
+    AttributeCommonInfo::Kind Kind;
+    unsigned Level;
+    Decl *D;
+  };
+  /// Complete the suspended bounds attribute after argument has been parsed
+  void CompleteBoundsAttribute(Expr *ParsedExpr,
+                               IncompleteBoundsAttributeInfo Info);
   /// Callback to the parser to parse a type expressed as a string.
-  std::function<ExprResult(StringRef, StringRef, Decl *, SourceLocation)>
+  std::function<void(StringRef, StringRef, Decl *, SourceLocation,
+                     IncompleteBoundsAttributeInfo Info)>
       ParseBoundsAttributeArgFromStringCallback;
   /* TO_UPSTREAM(BoundsSafety) OFF */
 

clang/lib/APINotes/APINotesReader.cpp

@@ -384,6 +384,14 @@ void ReadFunctionInfo(const uint8_t *&Data, FunctionInfo &Info) {
   Payload >>= 3;
   Info.NullabilityAudited = Payload & 0x1;
   Payload >>= 1;
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  if (Payload & 0x01) {
+    BoundsSafetyInfo BSI;
+    ReadBoundsSafetyInfo(Data, BSI);
+    Info.ReturnBoundsSafety = BSI;
+  }
+  /* TO_UPSTREAM(BoundsSafety) OFF */
+  Payload >>= 1;
   assert(Payload == 0 && "Bad API notes");
 
   Info.NumAdjustedNullable =

clang/lib/APINotes/APINotesTypes.cpp

@@ -61,6 +61,7 @@ LLVM_DUMP_METHOD void ObjCPropertyInfo::dump(llvm::raw_ostream &OS) const {
   OS << '\n';
 }
 
+/* TO_UPSTREAM(BoundsSafety) ON */
 LLVM_DUMP_METHOD void BoundsSafetyInfo::dump(llvm::raw_ostream &OS) const {
   if (KindAudited) {
     assert((BoundsSafetyKind)Kind >= BoundsSafetyKind::CountedBy);
@@ -88,6 +89,7 @@ LLVM_DUMP_METHOD void BoundsSafetyInfo::dump(llvm::raw_ostream &OS) const {
   OS << "ExternalBounds: "
      << (ExternalBounds.empty() ? "<missing>" : ExternalBounds) << '\n';
 }
+/* TO_UPSTREAM(BoundsSafety) OFF */
 
 LLVM_DUMP_METHOD void ParamInfo::dump(llvm::raw_ostream &OS) const {
   static_cast<const VariableInfo &>(*this).dump(OS);
@@ -97,8 +99,10 @@ LLVM_DUMP_METHOD void ParamInfo::dump(llvm::raw_ostream &OS) const {
     OS << (Lifetimebound ? "[Lifetimebound] " : "");
   OS << "RawRetainCountConvention: " << RawRetainCountConvention << ' ';
   OS << '\n';
+  /* TO_UPSTREAM(BoundsSafety) ON */
   if (BoundsSafety)
     BoundsSafety->dump(OS);
+  /* TO_UPSTREAM(BoundsSafety) OFF */
 }
 
 LLVM_DUMP_METHOD void FunctionInfo::dump(llvm::raw_ostream &OS) const {
@@ -109,6 +113,12 @@ LLVM_DUMP_METHOD void FunctionInfo::dump(llvm::raw_ostream &OS) const {
     OS << "Result Type: " << ResultType << ' ';
   if (!SwiftReturnOwnership.empty())
     OS << "SwiftReturnOwnership: " << SwiftReturnOwnership << ' ';
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  if (ReturnBoundsSafety) {
+    OS << "Result bounds: ";
+    ReturnBoundsSafety->dump(OS);
+  }
+  /* TO_UPSTREAM(BoundsSafety) OFF */
   if (!Params.empty())
     OS << '\n';
   for (auto &PI : Params)

clang/lib/APINotes/APINotesWriter.cpp

@@ -1148,6 +1148,10 @@ unsigned getFunctionInfoSize(const FunctionInfo &FI) {
     size += getParamInfoSize(P);
   size += sizeof(uint16_t) + FI.ResultType.size();
   size += sizeof(uint16_t) + FI.SwiftReturnOwnership.size();
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  if (auto BSI = FI.ReturnBoundsSafety)
+    size += getBoundsSafetyInfoSize(*BSI);
+  /* TO_UPSTREAM(BoundsSafety) OFF */
   return size;
 }
 
@@ -1156,6 +1160,11 @@ void emitFunctionInfo(raw_ostream &OS, const FunctionInfo &FI) {
   emitCommonEntityInfo(OS, FI);
 
   uint8_t flags = 0;
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  if (FI.ReturnBoundsSafety)
+    flags |= 0x01;
+  flags <<= 0x01;
+  /* TO_UPSTREAM(BoundsSafety) OFF */
   flags |= FI.NullabilityAudited;
   flags <<= 3;
   if (auto RCC = FI.getRetainCountConvention())
@@ -1164,6 +1173,10 @@ void emitFunctionInfo(raw_ostream &OS, const FunctionInfo &FI) {
   llvm::support::endian::Writer writer(OS, llvm::endianness::little);
 
   writer.write<uint8_t>(flags);
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  if (auto BSI = FI.ReturnBoundsSafety)
+    emitBoundsSafetyInfo(OS, *FI.ReturnBoundsSafety);
+  /* TO_UPSTREAM(BoundsSafety) OFF */
   writer.write<uint8_t>(FI.NumAdjustedNullable);
   writer.write<uint64_t>(FI.NullabilityPayload);
 

clang/lib/APINotes/APINotesYAMLCompiler.cpp

@@ -319,6 +319,9 @@ struct Method {
   bool Required = false;
   StringRef ResultType;
   StringRef SwiftReturnOwnership;
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  std::optional<BoundsSafetyNotes> ReturnBoundsSafety;
+  /* TO_UPSTREAM(BoundsSafety) OFF */
 };
 
 typedef std::vector<Method> MethodsSeq;
@@ -355,6 +358,9 @@ template <> struct MappingTraits<Method> {
     IO.mapOptional("ResultType", M.ResultType, StringRef(""));
     IO.mapOptional("SwiftReturnOwnership", M.SwiftReturnOwnership,
                    StringRef(""));
+    /* TO_UPSTREAM(BoundsSafety) ON */
+    IO.mapOptional("BoundsSafety", M.ReturnBoundsSafety);
+    /* TO_UPSTREAM(BoundsSafety) OFF */
   }
 };
 } // namespace yaml
@@ -451,6 +457,9 @@ struct Function {
   StringRef Type;
   StringRef ResultType;
   StringRef SwiftReturnOwnership;
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  std::optional<BoundsSafetyNotes> ReturnBoundsSafety;
+  /* TO_UPSTREAM(BoundsSafety) OFF */
 };
 
 typedef std::vector<Function> FunctionsSeq;
@@ -475,6 +484,9 @@ template <> struct MappingTraits<Function> {
     IO.mapOptional("ResultType", F.ResultType, StringRef(""));
     IO.mapOptional("SwiftReturnOwnership", F.SwiftReturnOwnership,
                    StringRef(""));
+    /* TO_UPSTREAM(BoundsSafety) ON */
+    IO.mapOptional("BoundsSafety", F.ReturnBoundsSafety);
+    /* TO_UPSTREAM(BoundsSafety) OFF */
   }
 };
 } // namespace yaml
@@ -905,14 +917,14 @@ class YAMLConverter {
       PI.setLifetimebound(P.Lifetimebound);
       PI.setType(std::string(P.Type));
       PI.setRetainCountConvention(P.RetainCountConvention);
-      BoundsSafetyInfo BSI;
       /* TO_UPSTREAM(BoundsSafety) ON */
       if (P.BoundsSafety) {
+        BoundsSafetyInfo BSI;
         BSI.setKindAudited(P.BoundsSafety->Kind);
         BSI.setLevelAudited(P.BoundsSafety->Level);
         BSI.ExternalBounds = P.BoundsSafety->BoundsExpr.str();
+        PI.BoundsSafety = BSI;
       }
-      PI.BoundsSafety = BSI;
       /* TO_UPSTREAM(BoundsSafety) OFF */
       if (static_cast<int>(OutInfo.Params.size()) <= P.Position)
         OutInfo.Params.resize(P.Position + 1);
@@ -1112,6 +1124,15 @@ class YAMLConverter {
     convertAvailability(Function.Availability, FI, Function.Name);
     FI.setSwiftPrivate(Function.SwiftPrivate);
     FI.SwiftName = std::string(Function.SwiftName);
+    /* TO_UPSTREAM(BoundsSafety) ON */
+    if (Function.ReturnBoundsSafety) {
+      BoundsSafetyInfo BSI;
+      BSI.setKindAudited(Function.ReturnBoundsSafety->Kind);
+      BSI.setLevelAudited(Function.ReturnBoundsSafety->Level);
+      BSI.ExternalBounds = Function.ReturnBoundsSafety->BoundsExpr.str();
+      FI.ReturnBoundsSafety = BSI;
+    }
+    /* TO_UPSTREAM(BoundsSafety) OFF */
     std::optional<ParamInfo> This;
     convertParams(Function.Params, FI, This);
     if constexpr (std::is_same_v<FuncOrMethodInfo, CXXMethodInfo>)

clang/lib/Parse/ParseDecl.cpp

@@ -7893,7 +7893,7 @@ void Parser::ParseFunctionDeclarator(Declarator &D,
   StartLoc = LParenLoc;
 
   /* TO_UPSTREAM(BoundsSafety) ON */
-  LateParsedAttrList LateParamAttrs(/*PSoon=*/false,
+  LateParsedAttrList LateParamAttrs(/*PSoon=*/true,
                                     /*LateAttrParseExperimentalExtOnly=*/true);
   /* TO_UPSTREAM(BoundsSafety) OFF */
 
@@ -8018,9 +8018,7 @@ void Parser::ParseFunctionDeclarator(Declarator &D,
   }
 
   /* TO_UPSTREAM(BoundsSafety) ON*/
-  for (auto LateParamAttr : LateParamAttrs) {
-    ParseLexedCAttribute(*LateParamAttr, true);
-  }
+  ParseLexedCAttributeList(LateParamAttrs, true, nullptr);
   /* TO_UPSTREAM(BoundsSafety) OFF*/
 
   // Collect non-parameter declarations from the prototype if this is a function
@@ -8996,10 +8994,9 @@ TypeResult Parser::ParseTypeFromString(StringRef TypeStr, StringRef Context,
 }
 
 /* TO_UPSTREAM(BoundsSafety) ON */
-ExprResult
-Parser::ParseBoundsAttributeArgFromString(StringRef ExprStr, StringRef Context,
-                                          Decl *ParentDecl,
-                                          SourceLocation IncludeLoc) {
+void Parser::ParseBoundsAttributeArgFromString(
+    StringRef ExprStr, StringRef Context, Decl *ParentDecl,
+    SourceLocation IncludeLoc, Sema::IncompleteBoundsAttributeInfo Payload) {
   // Consume (unexpanded) tokens up to the end-of-directive.
   SmallVector<Token, 4> Tokens;
   {
@@ -9073,9 +9070,11 @@ Parser::ParseBoundsAttributeArgFromString(StringRef ExprStr, StringRef Context,
   // Consume the end token.
   if (Tok.is(tok::eof) && Tok.getEofData() == ExprStr.data())
     ConsumeAnyToken();
+  // Need to create the type before leaving function context
+  if (Result.isUsable())
+    Actions.CompleteBoundsAttribute(Result.get(), Payload);
   if (EnterScope)
     Actions.ActOnExitFunctionContext();
-  return Result;
 }
 /* TO_UPSTREAM(BoundsSafety) OFF */
 

clang/lib/Parse/Parser.cpp

@@ -79,9 +79,10 @@ Parser::Parser(Preprocessor &pp, Sema &actions, bool skipFunctionBodies)
   /* TO_UPSTREAM(BoundsSafety) ON */
   Actions.ParseBoundsAttributeArgFromStringCallback =
       [this](StringRef ExprStr, StringRef Context, Decl *Parent,
-             SourceLocation IncludeLoc) {
+             SourceLocation IncludeLoc,
+             Sema::IncompleteBoundsAttributeInfo Payload) {
         return this->ParseBoundsAttributeArgFromString(ExprStr, Context, Parent,
-                                                       IncludeLoc);
+                                                       IncludeLoc, Payload);
       };
   /* TO_UPSTREAM(BoundsSafety) OFF */
 }

clang/lib/Sema/SemaAPINotes.cpp

@@ -410,6 +410,34 @@ static void ProcessAPINotes(Sema &S, Decl *D,
 }
 
 /* TO_UPSTREAM(BoundsSafety) ON */
+void Sema::CompleteBoundsAttribute(Expr *ParsedExpr,
+                                   IncompleteBoundsAttributeInfo Info) {
+  assert(ParsedExpr);
+  std::string AttrName;
+  switch (Info.Kind) {
+  case AttributeCommonInfo::AT_CountedBy:
+    AttrName = "__counted_by";
+    break;
+  case AttributeCommonInfo::AT_CountedByOrNull:
+    AttrName = "__counted_by_or_null";
+    break;
+  case AttributeCommonInfo::AT_SizedBy:
+    AttrName = "__sized_by";
+    break;
+  case AttributeCommonInfo::AT_SizedByOrNull:
+    AttrName = "__sized_by_or_null";
+    break;
+  case AttributeCommonInfo::AT_PtrEndedBy:
+    AttrName = "__ended_by";
+    break;
+  default:
+    llvm_unreachable("invalid bounds attribute kind in API notes");
+  }
+  applyPtrCountedByEndedByAttr(Info.D, Info.Level, Info.Kind, ParsedExpr,
+                               Info.D->getLocation(), Info.D->getSourceRange(),
+                               AttrName,
+                               /* originates in API notes */ true);
+}
 static void applyBoundsSafety(Sema &S, ValueDecl *D,
                               const api_notes::BoundsSafetyInfo &Info,
                               VersionedInfoMetadata Metadata) {
@@ -420,41 +448,29 @@ static void applyBoundsSafety(Sema &S, ValueDecl *D,
     if (auto ParmDecl = dyn_cast<ParmVarDecl>(ScopeDecl)) {
       ScopeDecl = dyn_cast<FunctionDecl>(ParmDecl->getDeclContext());
     }
-    auto ParsedExpr = S.ParseBoundsAttributeArgFromStringCallback(
-        Info.ExternalBounds, "<API Notes>", ScopeDecl, D->getLocation());
-    if (ParsedExpr.isInvalid())
-      return;
-
-    std::string AttrName;
-    AttributeCommonInfo::Kind Kind;
     assert(*Info.getKind() >= BoundsSafetyKind::CountedBy);
     assert(*Info.getKind() <= BoundsSafetyKind::EndedBy);
+    AttributeCommonInfo::Kind Kind;
     switch (*Info.getKind()) {
     case BoundsSafetyKind::CountedBy:
-      AttrName = "__counted_by";
       Kind = AttributeCommonInfo::AT_CountedBy;
       break;
     case BoundsSafetyKind::CountedByOrNull:
-      AttrName = "__counted_by_or_null";
       Kind = AttributeCommonInfo::AT_CountedByOrNull;
       break;
     case BoundsSafetyKind::SizedBy:
-      AttrName = "__sized_by";
       Kind = AttributeCommonInfo::AT_SizedBy;
       break;
     case BoundsSafetyKind::SizedByOrNull:
-      AttrName = "__sized_by_or_null";
       Kind = AttributeCommonInfo::AT_SizedByOrNull;
       break;
     case BoundsSafetyKind::EndedBy:
-      AttrName = "__ended_by";
       Kind = AttributeCommonInfo::AT_PtrEndedBy;
       break;
     }
-
-    S.applyPtrCountedByEndedByAttr(
-        D, *Info.getLevel(), Kind, ParsedExpr.get(), D->getLocation(),
-        D->getSourceRange(), AttrName, /* originates in API notes */ true);
+    S.ParseBoundsAttributeArgFromStringCallback(
+        Info.ExternalBounds, "<API Notes>", ScopeDecl, D->getLocation(),
+        {Kind, *Info.getLevel(), D});
   }
 }
 /* TO_UPSTREAM(BoundsSafety) OFF */
@@ -597,6 +613,14 @@ static void ProcessAPINotes(Sema &S, FunctionOrMethod AnyFunc,
     }
   }
 
+  /* TO_UPSTREAM(BoundsSafety) ON */
+  // FIXME(hnrklssn): apply to ObjC methods when support has landed
+  if (FD && Info.ReturnBoundsSafety.has_value())
+    // applyPtrCountedByEndedByAttr rebuilds FunctionType,
+    // no need to set AnyTypeChanged
+    applyBoundsSafety(S, FD, *Info.ReturnBoundsSafety, Metadata);
+  /* TO_UPSTREAM(BoundsSafety) OFF */
+
   // If the result type or any of the parameter types changed for a function
   // declaration, we have to rebuild the type.
   if (FD && AnyTypeChanged) {