From 2d0d736c505f1f3d6af6a44ce652081c842b828b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Bia=C5=82y?= <lukasz.marcin.bialy@gmail.com>
Date: Tue, 30 Jul 2024 14:35:49 +0200
Subject: [PATCH] issue 331 - preserve order of keys when parsing and printing

---
 .../src/main/scala/org/virtuslab/yaml/Node.scala  |  7 +++++++
 .../yaml/internal/load/compose/Composer.scala     |  3 ++-
 .../org/virtuslab/yaml/decoder/DecoderSuite.scala |  2 ++
 .../org/virtuslab/yaml/parser/ParserSuite.scala   | 15 +++++++++++++++
 4 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/core/shared/src/main/scala/org/virtuslab/yaml/Node.scala b/core/shared/src/main/scala/org/virtuslab/yaml/Node.scala
index 92118205a..6ca819397 100644
--- a/core/shared/src/main/scala/org/virtuslab/yaml/Node.scala
+++ b/core/shared/src/main/scala/org/virtuslab/yaml/Node.scala
@@ -23,6 +23,13 @@ sealed trait Node {
 }
 
 object Node {
+
+  // Only for MappingNode key order, only ScalarNodes can be keys.
+  implicit val nodeOrdering: Ordering[Node] = Ordering.by[Node, String] {
+    case ScalarNode(value, _) => value
+    case _                    => ""
+  }
+
   final case class ScalarNode private[yaml] (value: String, tag: Tag, pos: Option[Range] = None)
       extends Node
 
diff --git a/core/shared/src/main/scala/org/virtuslab/yaml/internal/load/compose/Composer.scala b/core/shared/src/main/scala/org/virtuslab/yaml/internal/load/compose/Composer.scala
index 0d8df42b1..3b4030c96 100644
--- a/core/shared/src/main/scala/org/virtuslab/yaml/internal/load/compose/Composer.scala
+++ b/core/shared/src/main/scala/org/virtuslab/yaml/internal/load/compose/Composer.scala
@@ -2,6 +2,7 @@ package org.virtuslab.yaml.internal.load.compose
 
 import scala.annotation.tailrec
 import scala.collection.mutable
+import scala.collection.immutable.ListMap
 
 import org.virtuslab.yaml.ComposerError
 import org.virtuslab.yaml.Node
@@ -127,7 +128,7 @@ object ComposerImpl extends Composer {
     }
 
     parseMappings(events, Nil).map { case (Result(nodes, rest), pos) =>
-      val mapping = Node.MappingNode(nodes.toMap, Tag.map, pos)
+      val mapping = Node.MappingNode(ListMap.from(nodes), Tag.map, pos)
       anchorOpt.foreach(anchor => aliases.put(anchor, mapping))
       Result(
         mapping,
diff --git a/core/shared/src/test/scala-3/org/virtuslab/yaml/decoder/DecoderSuite.scala b/core/shared/src/test/scala-3/org/virtuslab/yaml/decoder/DecoderSuite.scala
index b3da48179..6ba972f3f 100644
--- a/core/shared/src/test/scala-3/org/virtuslab/yaml/decoder/DecoderSuite.scala
+++ b/core/shared/src/test/scala-3/org/virtuslab/yaml/decoder/DecoderSuite.scala
@@ -572,6 +572,8 @@ class DecoderSuite extends munit.FunSuite:
     yaml.as[Float] match
       case Left(e: ConstructError) =>
         assertEquals(e.expected, Some("Float"))
+      case Left(e) =>
+        fail(s"Should fail, but got $e", e)
       case Right(value) =>
         fail(s"Should fail, but got $value")
   }
diff --git a/core/shared/src/test/scala/org/virtuslab/yaml/parser/ParserSuite.scala b/core/shared/src/test/scala/org/virtuslab/yaml/parser/ParserSuite.scala
index b2fea840f..d40a4085f 100644
--- a/core/shared/src/test/scala/org/virtuslab/yaml/parser/ParserSuite.scala
+++ b/core/shared/src/test/scala/org/virtuslab/yaml/parser/ParserSuite.scala
@@ -179,4 +179,19 @@ class ParserSuite extends BaseYamlSuite {
 
     assertEquals(yaml.events, Right(expectedEvents))
   }
+
+  test("parsing keeps order of keys") {
+    val yaml = """
+                 |P: 
+                 |  a: 0
+                 |  b: 1
+                 |  c: 2
+                 |  d: 3
+                 |  e: 4
+                 |""".stripMargin
+
+    val node = yaml.asNode.toOption.get
+
+    assertEquals(node.asYaml.trim, yaml.trim)
+  }
 }