[llvm] [YAML] Fix incorrect dash output in nested sequences (PR #116488)

[NAKAMURA Takumi via llvm-commits]

https://github.com/chapuni created https://github.com/llvm/llvm-project/pull/116488

Nested sequences could be defined but the YAML output was incorrect. `Output::newLineCheck()` was not able to emit multiple dashes `- ` and YAML parser sometimes didn't accept its output as the result.

This fixes for emitting corresponding dashes for consecutive `inSeqFirstElement`, but suppresses emission to the top `inSeqFirstElement`.

This also fixes for emitting flow elements onto nested sequences.

>From d04d8b4ca279eb7570b1baca8aae0c57ae7824f1 Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic at gmail.com>
Date: Sun, 17 Nov 2024 00:25:24 +0900
Subject: [PATCH] [YAML] Fix incorrect dash output in nested sequences

---
 llvm/lib/Support/YAMLTraits.cpp       | 46 +++++++++-----
 llvm/unittests/Support/YAMLIOTest.cpp | 89 +++++++++++++++++++++++++++
 2 files changed, 119 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Support/YAMLTraits.cpp b/llvm/lib/Support/YAMLTraits.cpp
index 56b557646100b1..0d694d5b9d45b5 100644
--- a/llvm/lib/Support/YAMLTraits.cpp
+++ b/llvm/lib/Support/YAMLTraits.cpp
@@ -838,26 +838,40 @@ void Output::newLineCheck(bool EmptySequence) {
     return;
 
   unsigned Indent = StateStack.size() - 1;
-  bool OutputDash = false;
-
-  if (StateStack.back() == inSeqFirstElement ||
-      StateStack.back() == inSeqOtherElement) {
-    OutputDash = true;
-  } else if ((StateStack.size() > 1) &&
-             ((StateStack.back() == inMapFirstKey) ||
-              inFlowSeqAnyElement(StateStack.back()) ||
-              (StateStack.back() == inFlowMapFirstKey)) &&
-             inSeqAnyElement(StateStack[StateStack.size() - 2])) {
-    --Indent;
-    OutputDash = true;
+  bool PossiblyNestedSeq = false;
+  auto I = StateStack.rbegin(), E = StateStack.rend();
+
+  if (inSeqAnyElement(*I)) {
+    PossiblyNestedSeq = true; // Not possibly but always.
+    ++Indent;
+  } else if (*I == inMapFirstKey || *I == inFlowMapFirstKey ||
+             inFlowSeqAnyElement(*I)) {
+    PossiblyNestedSeq = true;
+    ++I; // Skip back().
   }
 
-  for (unsigned i = 0; i < Indent; ++i) {
-    output("  ");
+  unsigned OutputDashCount = 0;
+  if (PossiblyNestedSeq) {
+    // Count up consecutive inSeqFirstElement from the end, unless
+    // inSeqFirstElement is the top of nested sequence.
+    while (I != E) {
+      // Don't count the top of nested sequence.
+      if (!inSeqAnyElement(*I))
+        break;
+
+      ++OutputDashCount;
+
+      // Stop counting if consecutive inSeqFirstElement ends.
+      if (*I++ != inSeqFirstElement)
+        break;
+    }
   }
-  if (OutputDash) {
+
+  for (unsigned I = OutputDashCount; I < Indent; ++I)
+    output("  ");
+
+  for (unsigned I = 0; I < OutputDashCount; ++I)
     output("- ");
-  }
 }
 
 void Output::paddedKey(StringRef key) {
diff --git a/llvm/unittests/Support/YAMLIOTest.cpp b/llvm/unittests/Support/YAMLIOTest.cpp
index e10fe099a30adb..be592559e03cda 100644
--- a/llvm/unittests/Support/YAMLIOTest.cpp
+++ b/llvm/unittests/Support/YAMLIOTest.cpp
@@ -1538,6 +1538,95 @@ TEST(YAMLIO, TestReadWriteMySecondsSequence) {
 }
 
 
+//===----------------------------------------------------------------------===//
+//  Test nested sequence
+//===----------------------------------------------------------------------===//
+using NestedStringSeq1 = std::array<std::string, 2>;
+using NestedStringSeq2 = std::array<NestedStringSeq1, 2>;
+using NestedStringSeq3 = std::array<NestedStringSeq2, 2>;
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(NestedStringSeq1)
+LLVM_YAML_IS_SEQUENCE_VECTOR(NestedStringSeq2)
+
+struct MappedStringSeq3 {
+  NestedStringSeq3 Seq3;
+};
+
+template <> struct llvm::yaml::MappingTraits<MappedStringSeq3> {
+  static void mapping(IO &io, MappedStringSeq3 &seq) {
+    io.mapRequired("Seq3", seq.Seq3);
+  }
+};
+
+using NestedIntSeq1 = std::array<int, 2>;
+using NestedIntSeq2 = std::array<NestedIntSeq1, 2>;
+using NestedIntSeq3 = std::array<NestedIntSeq2, 2>;
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(NestedIntSeq1)
+LLVM_YAML_IS_SEQUENCE_VECTOR(NestedIntSeq2)
+LLVM_YAML_IS_SEQUENCE_VECTOR(NestedIntSeq3)
+
+template <typename Ty> std::string ParseAndEmit(llvm::StringRef YAML) {
+  Ty seq3;
+  Input yin(YAML);
+  yin >> seq3;
+  std::string out;
+  llvm::raw_string_ostream ostr(out);
+  Output yout(ostr);
+  yout << seq3;
+  return out;
+}
+
+TEST(YAMLIO, TestNestedSequence) {
+  {
+    llvm::StringRef Seq3YAML(R"YAML(---
+- - [ 1000, 1001 ]
+  - [ 1010, 1011 ]
+- - [ 1100, 1101 ]
+  - [ 1110, 1111 ]
+...
+)YAML");
+
+    std::string out = ParseAndEmit<NestedIntSeq3>(Seq3YAML);
+    EXPECT_EQ(out, Seq3YAML);
+  }
+
+  {
+    llvm::StringRef Seq3YAML(R"YAML(---
+- - - '000'
+    - '001'
+  - - '010'
+    - '011'
+- - - '100'
+    - '101'
+  - - '110'
+    - '111'
+...
+)YAML");
+
+    std::string out = ParseAndEmit<NestedStringSeq3>(Seq3YAML);
+    EXPECT_EQ(out, Seq3YAML);
+  }
+
+  {
+    llvm::StringRef Seq3YAML(R"YAML(---
+Seq3:
+  - - - '000'
+      - '001'
+    - - '010'
+      - '011'
+  - - - '100'
+      - '101'
+    - - '110'
+      - '111'
+...
+)YAML");
+
+    std::string out = ParseAndEmit<MappedStringSeq3>(Seq3YAML);
+    EXPECT_EQ(out, Seq3YAML);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //  Test dynamic typing
 //===----------------------------------------------------------------------===//