Generate annotation edits for types with `const` qualifier.

PiperOrigin-RevId: 586990556
Change-Id: I44889e76e67cb415365ee54a89f00c7716d30ace
diff --git a/nullability/inference/BUILD b/nullability/inference/BUILD
index 480165e..d778b28 100644
--- a/nullability/inference/BUILD
+++ b/nullability/inference/BUILD
@@ -131,7 +131,9 @@
         "@llvm-project//clang:ast",
         "@llvm-project//clang:basic",
         "@llvm-project//clang:index",
+        "@llvm-project//clang:lex",
         "@llvm-project//clang:transformer",
+        "@llvm-project//clang-tools-extra/clang-tidy:tidy_framework",
         "@llvm-project//llvm:Support",
     ],
 )
diff --git a/nullability/inference/eligible_ranges.cc b/nullability/inference/eligible_ranges.cc
index 38a2d72..c4009a6 100644
--- a/nullability/inference/eligible_ranges.cc
+++ b/nullability/inference/eligible_ranges.cc
@@ -5,11 +5,13 @@
 #include "nullability/inference/eligible_ranges.h"
 
 #include <cassert>
+#include <cstdint>
 #include <optional>
 #include <string_view>
 
 #include "nullability/inference/inference.proto.h"
 #include "nullability/type_nullability.h"
+#include "third_party/llvm/llvm-project/clang-tools-extra/clang-tidy/utils/LexerUtils.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/Type.h"
@@ -18,17 +20,22 @@
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TokenKinds.h"
 #include "clang/Index/USRGeneration.h"
+#include "clang/Lex/Lexer.h"
+#include "clang/Lex/Token.h"
 #include "clang/Tooling/Transformer/SourceCode.h"
+#include "llvm/Support/ErrorHandling.h"
 
 namespace clang::tidy::nullability {
 
-namespace {
-bool isEligibleTypeLoc(TypeLoc TyLoc) {
+// TODO: incorporate the predicate used by inference to identify relevant
+// top-level slots.
+static bool isEligibleTypeLoc(TypeLoc TyLoc) {
   QualType Ty = TyLoc.getType();
   if (!isSupportedPointerType(Ty)) return false;
-  // Excludes function-pointer types, which have an odd syntax, and
-  // `const`-qualified pointees, whose range is not correctly identified.
+  // Excludes function-pointer types and other corner cases that use declarator
+  // suffix syntax.
   // TODO(sammccall): In general, a declaration looks like BEGIN * MIDDLE NAME
   // END e.g. `int * const (*a) []`.  We need to rewrite as `Nullable<BEGIN *
   // MIDDLE END> NAME, e.g. `Nullable<int * const (*) []> a`. But, if we're not
@@ -39,24 +46,61 @@
   // "is the param unnamed" and the latter is "is PointerTypeLoc.getEndLoc() >
   // Decl.getLocation()?".
   if (const auto *PT = Ty.getNonReferenceType()->getAs<PointerType>()) {
-    // TODO: Allow local `const` qualifiers, function-pointer types and other
-    // corner cases that are currently excluded, because they use declarator
-    // suffix syntax. These aren't inherently bad, just tricky to edit
-    // correctly. This change will require generalizing this function to take
-    // more than just a `TypeLoc`, the declarator may be needed as well.
+    // TODO: Allow function-pointer types and other corner cases that are
+    // currently excluded, because they use declarator suffix syntax. These
+    // aren't inherently bad, just tricky to edit correctly. This change will
+    // require generalizing this function to take more than just a `TypeLoc`,
+    // the declarator may be needed as well.
     QualType PointeeTy = PT->getPointeeType();
-    // We do not consider the `const`-ness of `PT` itself, because the edit will
-    // do the correct thing implicitly: the `const` will be left out of the
-    // associated `TypeLoc`, leaving `const` outside the nullability annotation,
-    // which is the preferred spelling.
-    if (PointeeTy->isFunctionType() || PointeeTy->isArrayType() ||
-        PointeeTy.isLocalConstQualified())
-      return false;
+    if (PointeeTy->isFunctionType() || PointeeTy->isArrayType()) return false;
   }
   return true;
 }
 
-std::optional<TypeLocRanges> getEligibleRanges(const FunctionDecl &Fun) {
+// Extracts the source range of `Loc`, accounting for (nested) qualifiers.
+// Guarantees that `Loc` is eligible for editing, including that its begin and
+// end locations are in the same file.
+//
+// We do not consider the `const`-ness of `Loc` itself, because the edit will do
+// the correct thing implicitly: the `const` will be left out of `Loc`'s range,
+// leaving `const` outside the nullability annotation, which is the preferred
+// spelling.
+static std::optional<CharSourceRange> getRangeQualifierAware(
+    TypeLoc Loc, const ASTContext &Context) {
+  if (!isEligibleTypeLoc(Loc)) return std::nullopt;
+
+  auto R = tooling::getFileRange(
+      CharSourceRange::getTokenRange(Loc.getSourceRange()), Context,
+      /*IncludeMacroExpansion=*/true);
+  if (!R) return std::nullopt;
+
+  const auto &SM = Context.getSourceManager();
+  const auto &LangOpts = Context.getLangOpts();
+
+  // The start of the new range.
+  SourceLocation Begin = R->getBegin();
+
+  // Update `Begin` as we search backwards and find qualifier tokens.
+  auto PrevTok = utils::lexer::getPreviousToken(Begin, SM, LangOpts);
+  while (PrevTok.getKind() != tok::unknown) {
+    if (!PrevTok.is(tok::raw_identifier)) break;
+    StringRef RawID = PrevTok.getRawIdentifier();
+    if (RawID != "const" && RawID != "volatile" && RawID != "restrict") break;
+    Begin = PrevTok.getLocation();
+    PrevTok = utils::lexer::getPreviousToken(Begin, SM, LangOpts);
+  }
+
+  return CharSourceRange::getCharRange(Begin, R->getEnd());
+}
+
+static void initSlotRange(SlotRange &R, uint32_t Slot, unsigned Begin,
+                          unsigned End) {
+  R.set_slot(Slot);
+  R.set_begin(Begin);
+  R.set_end(End);
+}
+
+static std::optional<TypeLocRanges> getEligibleRanges(const FunctionDecl &Fun) {
   FunctionTypeLoc TyLoc = Fun.getFunctionTypeLoc();
   if (TyLoc.isNull()) return std::nullopt;
 
@@ -64,40 +108,27 @@
   const SourceManager &SrcMgr = Context.getSourceManager();
   TypeLocRanges Result;
 
-  // Guarantees that `Loc` is eligible for editing, including that its begin and
-  // end locations are in the same file.
-  auto GetLocRange = [&Context](TypeLoc Loc) -> std::optional<CharSourceRange> {
-    if (!isEligibleTypeLoc(Loc)) return std::nullopt;
-    return clang::tooling::getFileRange(
-        CharSourceRange::getTokenRange(Loc.getSourceRange()), Context,
-        /*IncludeMacroExpansion=*/true);
-  };
-
   FileID DeclFID = SrcMgr.getFileID(SrcMgr.getExpansionLoc(Fun.getLocation()));
 
-  if (auto CSR = GetLocRange(TyLoc.getReturnLoc())) {
+  if (auto CSR = getRangeQualifierAware(TyLoc.getReturnLoc(), Context)) {
     auto [FID, Begin] = SrcMgr.getDecomposedLoc(CSR->getBegin());
     // If the type comes from a different file, then don't attempt to edit -- it
     // might need manual intervention.
-    if (FID == DeclFID) {
-      auto *R = Result.add_range();
-      R->set_slot(SLOT_RETURN_TYPE);
-      R->set_begin(Begin);
-      R->set_end(SrcMgr.getFileOffset(CSR->getEnd()));
-    }
+    if (FID == DeclFID)
+      initSlotRange(*Result.add_range(), SLOT_RETURN_TYPE, Begin,
+                    SrcMgr.getFileOffset(CSR->getEnd()));
   }
 
   for (int I = 0, N = Fun.getNumParams(); I < N; ++I) {
     const ParmVarDecl *P = Fun.getParamDecl(I);
-    if (auto CSR = GetLocRange(P->getTypeSourceInfo()->getTypeLoc())) {
+    if (auto CSR = getRangeQualifierAware(P->getTypeSourceInfo()->getTypeLoc(),
+                                          Context)) {
       auto [FID, Begin] = SrcMgr.getDecomposedLoc(CSR->getBegin());
       // If the type comes from a different file, then don't attempt to edit --
       // it might need manual intervention.
-      if (FID != DeclFID) continue;
-      auto *R = Result.add_range();
-      R->set_slot(SLOT_PARAM + I);
-      R->set_begin(Begin);
-      R->set_end(SrcMgr.getFileOffset(CSR->getEnd()));
+      if (FID == DeclFID)
+        initSlotRange(*Result.add_range(), SLOT_PARAM + I, Begin,
+                      SrcMgr.getFileOffset(CSR->getEnd()));
     }
   }
 
@@ -110,8 +141,6 @@
   return Result;
 }
 
-}  // namespace
-
 std::optional<TypeLocRanges> getEligibleRanges(const Decl &D) {
   // Only functions are currently supported.
   if (const auto *Fun = clang::dyn_cast<FunctionDecl>(&D))
diff --git a/nullability/inference/eligible_ranges_test.cc b/nullability/inference/eligible_ranges_test.cc
index adf92c3..226b0cc 100644
--- a/nullability/inference/eligible_ranges_test.cc
+++ b/nullability/inference/eligible_ranges_test.cc
@@ -71,7 +71,8 @@
           MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
 }
 
-TEST(GenEditsTest, DeclCovered) {
+// Checks that a function decl without a body is handled correctly.
+TEST(GenEditsTest, DeclHandled) {
   auto Input = Annotations("void foo([[int *]]p1, int p2);");
   EXPECT_THAT(
       getEligibleRanges(Input.code()),
@@ -79,18 +80,52 @@
           MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
 }
 
-TEST(GenEditsTest, ConstIncludeInType) {
-  std::string Input = "void foo(const int *p1, int p2);";
-  // TODO: fix, once `const` is supported.
+TEST(GenEditsTest, NestedPointerEligible) {
+  auto Input = Annotations("void foo([[int **]]p);");
+  EXPECT_THAT(
+      getEligibleRanges(Input.code()),
+      Optional(TypeLocRanges(
+          MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
+}
+
+TEST(GenEditsTest, DeclConstExcluded) {
+  auto Input = Annotations("void foo([[int *]] const p1, int p2);");
+  EXPECT_THAT(
+      getEligibleRanges(Input.code()),
+      Optional(TypeLocRanges(
+          MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
+}
+
+TEST(GenEditsTest, PointeeConstIncluded) {
+  auto Input = Annotations("void foo([[const int *]]p1, int p2);");
+  EXPECT_THAT(
+      getEligibleRanges(Input.code()),
+      Optional(TypeLocRanges(
+          MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
+}
+
+TEST(GenEditsTest, NestedPointeeConstIncluded) {
+  auto Input = Annotations("void foo([[const int **]]p1, int p2);");
+  EXPECT_THAT(
+      getEligibleRanges(Input.code()),
+      Optional(TypeLocRanges(
+          MainFileName, UnorderedElementsAre(SlotRange(1, Input.range())))));
+}
+
+TEST(GenEditsTest, FunctionPointerTypeIgnored) {
+  std::string Input = "void foo(int (*p)(int));";
   EXPECT_EQ(getEligibleRanges(Input), std::nullopt);
 }
 
-// TODO: Add additional tests for type corner cases, e.g.
-//
-// int (*foo)(int (*a))
-// void foo(int **x())
-// void foo(int x[])
-// void foo(int (*z[3])(float))
+TEST(GenEditsTest, ArrayTypeIgnored) {
+  std::string Input = "void foo(int p[]);";
+  EXPECT_EQ(getEligibleRanges(Input), std::nullopt);
+}
+
+TEST(GenEditsTest, FunctionAndArrayTypeIgnored) {
+  std::string Input = "void foo(int (*z[3])(float));";
+  EXPECT_EQ(getEligibleRanges(Input), std::nullopt);
+}
 
 }  // namespace
 }  // namespace clang::tidy::nullability