[clang] [clang-format] Fix a bug in lexing C++ UDL ending in $ (PR #136476)

[Owen Pan via cfe-commits]

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

None

>From 1507898f4b68d6644aae1449eb8e7fec2b903283 Mon Sep 17 00:00:00 2001
From: Owen Pan <owenpiano at gmail.com>
Date: Sun, 20 Apr 2025 00:47:22 -0700
Subject: [PATCH] [clang-format] Fix a bug in lexing C++ UDL ending in $

---
 clang/lib/Format/FormatTokenLexer.cpp         | 26 +++++++++++++++++++
 clang/lib/Format/FormatTokenLexer.h           |  1 +
 clang/unittests/Format/TokenAnnotatorTest.cpp |  6 +++++
 3 files changed, 33 insertions(+)

diff --git a/clang/lib/Format/FormatTokenLexer.cpp b/clang/lib/Format/FormatTokenLexer.cpp
index 5c4e1f814d9b7..b67623d84be24 100644
--- a/clang/lib/Format/FormatTokenLexer.cpp
+++ b/clang/lib/Format/FormatTokenLexer.cpp
@@ -128,6 +128,12 @@ void FormatTokenLexer::tryMergePreviousTokens() {
   if (Style.isCpp() && tryTransformTryUsageForC())
     return;
 
+  if ((Style.Language == FormatStyle::LK_Cpp ||
+       Style.Language == FormatStyle::LK_ObjC) &&
+      tryMergeUserDefinedLiteral()) {
+    return;
+  }
+
   if (Style.isJavaScript() || Style.isCSharp()) {
     static const tok::TokenKind NullishCoalescingOperator[] = {tok::question,
                                                                tok::question};
@@ -559,6 +565,26 @@ bool FormatTokenLexer::tryMergeGreaterGreater() {
   return true;
 }
 
+bool FormatTokenLexer::tryMergeUserDefinedLiteral() {
+  if (Tokens.size() < 2)
+    return false;
+
+  auto *First = Tokens.end() - 2;
+  auto &Suffix = First[1];
+  if (Suffix->hasWhitespaceBefore() || Suffix->TokenText != "$")
+    return false;
+
+  auto &Literal = First[0];
+  auto &Text = Literal->TokenText;
+  if (!Literal->Tok.isLiteral() || !Text.ends_with("_"))
+    return false;
+
+  Text = StringRef(Text.data(), Text.size() + 1);
+  ++Literal->ColumnWidth;
+  Tokens.erase(&Suffix);
+  return true;
+}
+
 bool FormatTokenLexer::tryMergeTokens(ArrayRef<tok::TokenKind> Kinds,
                                       TokenType NewType) {
   if (Tokens.size() < Kinds.size())
diff --git a/clang/lib/Format/FormatTokenLexer.h b/clang/lib/Format/FormatTokenLexer.h
index 61474a3f9ada8..3f001bc69415d 100644
--- a/clang/lib/Format/FormatTokenLexer.h
+++ b/clang/lib/Format/FormatTokenLexer.h
@@ -48,6 +48,7 @@ class FormatTokenLexer {
 
   bool tryMergeLessLess();
   bool tryMergeGreaterGreater();
+  bool tryMergeUserDefinedLiteral();
   bool tryMergeNSStringLiteral();
   bool tryMergeJSPrivateIdentifier();
   bool tryMergeCSharpStringLiteral();
diff --git a/clang/unittests/Format/TokenAnnotatorTest.cpp b/clang/unittests/Format/TokenAnnotatorTest.cpp
index 043ee2e088ddb..c0bdbb6bad09f 100644
--- a/clang/unittests/Format/TokenAnnotatorTest.cpp
+++ b/clang/unittests/Format/TokenAnnotatorTest.cpp
@@ -3969,6 +3969,12 @@ TEST_F(TokenAnnotatorTest, IdentifierPackage) {
   EXPECT_FALSE(Tokens[0]->isObjCAccessSpecifier());
 }
 
+TEST_F(TokenAnnotatorTest, UserDefinedLiteral) {
+  auto Tokens = annotate("auto dollars = 2_$;");
+  ASSERT_EQ(Tokens.size(), 6u) << Tokens;
+  EXPECT_EQ(Tokens[3]->TokenText, "2_$");
+}
+
 } // namespace
 } // namespace format
 } // namespace clang