[llvm] [Support] Fix bugs in formatv automatic index assignment (PR #108384)

[Rahul Joshi via llvm-commits]

https://github.com/jurahul created https://github.com/llvm/llvm-project/pull/108384

Fix bugs found when actually trying to use formatv() automatic index assignment in IntrinsicEmitter.cpp:
- Assign automatic index only for `ReplacementType::Format`.
- Make the check for all replacement indices being either automatic or explicit more accurate. The existing check fails for format formatv("{}{0}{}", 0, 1) (added as a unit test). Explicitly track if we have seen any explicit and any automatic index instead.

>From 82bafe9154e5e4941d4bce51816e7fe5ed0d37be Mon Sep 17 00:00:00 2001
From: Rahul Joshi <rjoshi at nvidia.com>
Date: Thu, 12 Sep 2024 06:12:39 -0700
Subject: [PATCH] [Support] Fix bugs in formatv automatic index assignment

Fix bugs found when actually trying to use formatv() automatic index
assignment in IntrinsicEmitter.cpp:
- Assign automatic index only for `ReplacementType::Format`.
- Make the check for all replacement indices being either automatic
  or explicit more accurate. The existing check fails for format
  formatv("{}{0}{}", 0, 1) (added as a unit test). Explicitly track
  if we have seen any explicit and any automatic index instead.
---
 llvm/lib/Support/FormatVariadic.cpp           | 19 +++++++++++--------
 llvm/unittests/Support/FormatVariadicTest.cpp | 12 ++++++++++++
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Support/FormatVariadic.cpp b/llvm/lib/Support/FormatVariadic.cpp
index 3240b1c0f4e413..7c4b265b2ee228 100644
--- a/llvm/lib/Support/FormatVariadic.cpp
+++ b/llvm/lib/Support/FormatVariadic.cpp
@@ -144,6 +144,7 @@ formatv_object_base::parseFormatString(StringRef Fmt, size_t NumArgs,
                                        bool Validate) {
   SmallVector<ReplacementItem, 2> Replacements;
   unsigned NextAutomaticIndex = 0;
+  bool HasExplicitIndex = false;
 
 #if ENABLE_VALIDATION
   const StringRef SavedFmtStr = Fmt;
@@ -155,14 +156,17 @@ formatv_object_base::parseFormatString(StringRef Fmt, size_t NumArgs,
     std::tie(I, Fmt) = splitLiteralAndReplacement(Fmt);
     if (!I)
       continue;
-    if (I->Index == ~0U)
-      I->Index = NextAutomaticIndex++;
-
-    Replacements.emplace_back(*I);
+    if (I->Type == ReplacementType::Format) {
+      if (I->Index == ~0U)
+        I->Index = NextAutomaticIndex++;
+      else
+        HasExplicitIndex = true;
 #if ENABLE_VALIDATION
-    if (I->Type == ReplacementType::Format)
       NumExpectedArgs = std::max(NumExpectedArgs, I->Index + 1);
 #endif
+    }
+
+    Replacements.emplace_back(*I);
   }
 
 #if ENABLE_VALIDATION
@@ -208,9 +212,8 @@ formatv_object_base::parseFormatString(StringRef Fmt, size_t NumArgs,
     return getErrorReplacements("Replacement indices have holes");
   }
 
-  // If we had automatic numbering of replacement indices, verify that all
-  // indices used automatic numbering.
-  if (NextAutomaticIndex != 0 && NextAutomaticIndex != Count) {
+  // Fail validation if we see both automatic index and explicit index.
+  if (NextAutomaticIndex != 0 && HasExplicitIndex) {
     errs() << formatv(
         "Cannot mix automatic and explicit indices for format string '{}'\n",
         SavedFmtStr);
diff --git a/llvm/unittests/Support/FormatVariadicTest.cpp b/llvm/unittests/Support/FormatVariadicTest.cpp
index 68938480ecfe20..34bfbd16bd08dd 100644
--- a/llvm/unittests/Support/FormatVariadicTest.cpp
+++ b/llvm/unittests/Support/FormatVariadicTest.cpp
@@ -283,6 +283,15 @@ TEST(FormatVariadicTest, AutomaticIndices) {
   EXPECT_EQ(ReplacementType::Format, Replacements[0].Type);
   EXPECT_EQ(0u, Replacements[0].Index);
 
+  Replacements = parseFormatString("{}bar{}");
+  ASSERT_EQ(3u, Replacements.size());
+  EXPECT_EQ(ReplacementType::Format, Replacements[0].Type);
+  EXPECT_EQ(0u, Replacements[0].Index);
+  EXPECT_EQ(ReplacementType::Literal, Replacements[1].Type);
+  EXPECT_EQ("bar", Replacements[1].Spec);
+  EXPECT_EQ(ReplacementType::Format, Replacements[2].Type);
+  EXPECT_EQ(1u, Replacements[2].Index);
+
   Replacements = parseFormatString("{}{}");
   ASSERT_EQ(2u, Replacements.size());
   EXPECT_EQ(ReplacementType::Format, Replacements[0].Type);
@@ -761,6 +770,8 @@ TEST(FormatVariadicTest, Validate) {
                "Replacement field indices cannot have holes");
   EXPECT_DEATH(formatv("{}{1}", 0, 1).str(),
                "Cannot mix automatic and explicit indices");
+  EXPECT_DEATH(formatv("{}{0}{}", 0, 1).str(),
+               "Cannot mix automatic and explicit indices");
 #else  // GTEST_HAS_DEATH_TEST
   GTEST_SKIP() << "No support for EXPECT_DEATH";
 #endif // GTEST_HAS_DEATH_TEST
@@ -769,6 +780,7 @@ TEST(FormatVariadicTest, Validate) {
   EXPECT_EQ(formatv("{0}", 1, 2).str(), "1");
   EXPECT_EQ(formatv("{0} {2}", 1, 2, 3).str(), "1 3");
   EXPECT_EQ(formatv("{}{1}", 0, 1).str(), "01");
+  EXPECT_EQ(formatv("{}{0}{}", 0, 1).str(), "001");
 #endif // NDEBUG
 }