[clang] [clang-format] Correctly annotate pointer/reference in range-for loop (PR #109361)

[Owen Pan via cfe-commits]

https://github.com/owenca created https://github.com/llvm/llvm-project/pull/109361

Fixes #109358.

>From d954c56f41a96ace6b507f6e6f4531c885db9f38 Mon Sep 17 00:00:00 2001
From: Owen Pan <owenpiano at gmail.com>
Date: Thu, 19 Sep 2024 18:56:38 -0700
Subject: [PATCH] [clang-format] Correctly annotate pointer/reference in
 range-for loop

Fixes #109358.
---
 clang/lib/Format/TokenAnnotator.cpp           | 6 ++++++
 clang/unittests/Format/TokenAnnotatorTest.cpp | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/clang/lib/Format/TokenAnnotator.cpp b/clang/lib/Format/TokenAnnotator.cpp
index 6f09835bad3a83..3f7ac6c7776f40 100644
--- a/clang/lib/Format/TokenAnnotator.cpp
+++ b/clang/lib/Format/TokenAnnotator.cpp
@@ -1407,6 +1407,12 @@ class AnnotatingParser {
         }
       } else if (Contexts.back().ColonIsForRangeExpr) {
         Tok->setType(TT_RangeBasedForLoopColon);
+        for (auto *Prev = Tok->Previous;
+             Prev && !Prev->isOneOf(tok::semi, tok::l_paren);
+             Prev = Prev->Previous) {
+          if (Prev->isPointerOrReference())
+            Prev->setFinalizedType(TT_PointerOrReference);
+        }
       } else if (Contexts.back().ContextType == Context::C11GenericSelection) {
         Tok->setType(TT_GenericSelectionColon);
       } else if (CurrentToken && CurrentToken->is(tok::numeric_constant)) {
diff --git a/clang/unittests/Format/TokenAnnotatorTest.cpp b/clang/unittests/Format/TokenAnnotatorTest.cpp
index 34c03d668a9a0a..dfb6c060d32094 100644
--- a/clang/unittests/Format/TokenAnnotatorTest.cpp
+++ b/clang/unittests/Format/TokenAnnotatorTest.cpp
@@ -333,6 +333,12 @@ TEST_F(TokenAnnotatorTest, UnderstandsUsesOfStarAndAmp) {
   ASSERT_EQ(Tokens.size(), 17u) << Tokens;
   EXPECT_TOKEN(Tokens[11], tok::star, TT_BinaryOperator);
 
+  Tokens = annotate("for (int i; Foo *&foo : foos)");
+  ASSERT_EQ(Tokens.size(), 13u) << Tokens;
+  EXPECT_TOKEN(Tokens[6], tok::star, TT_PointerOrReference);
+  EXPECT_TOKEN(Tokens[7], tok::amp, TT_PointerOrReference);
+  EXPECT_TOKEN(Tokens[9], tok::colon, TT_RangeBasedForLoopColon);
+
   Tokens = annotate("#define FOO auto Foo = [] { f(a * b); };");
   ASSERT_EQ(Tokens.size(), 19u) << Tokens;
   EXPECT_TOKEN(Tokens[12], tok::star, TT_BinaryOperator);