From 670ff9a252f6875a60923fa35a4676134f2b2179 Mon Sep 17 00:00:00 2001
From: mxt <mxnd.taylor@gmail.com>
Date: Sat, 9 Nov 2024 18:04:31 -0600
Subject: [PATCH] fix: Preserve dotted-key ordering while allowing key order
 modifications

Fixes #163
---
 crates/toml_edit/src/inline_table.rs     | 10 +++----
 crates/toml_edit/src/table.rs            | 20 +++++++++----
 crates/toml_edit/tests/testsuite/edit.rs | 36 ++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/crates/toml_edit/src/inline_table.rs b/crates/toml_edit/src/inline_table.rs
index ccf77212..a821ceaf 100644
--- a/crates/toml_edit/src/inline_table.rs
+++ b/crates/toml_edit/src/inline_table.rs
@@ -2,7 +2,7 @@ use std::iter::FromIterator;
 
 use crate::key::Key;
 use crate::repr::Decor;
-use crate::table::{Iter, IterMut, KeyValuePairs, TableLike, sort_values_by_position};
+use crate::table::{Iter, IterMut, KeyValuePairs, TableLike, sort_values_by_position, modify_key_position};
 use crate::{InternalString, Item, KeyMut, RawString, Table, Value};
 
 /// Type representing a TOML inline table,
@@ -87,14 +87,14 @@ impl InlineTable {
     pub fn sort_values(&mut self) {
         // Assuming standard tables have their position set and this won't negatively impact them
         self.items.sort_keys();
-        for value in self.items.values_mut() {
+        modify_key_position(&mut self.items,|value| {
             match value {
                 Item::Value(Value::InlineTable(table)) if table.is_dotted() => {
                     table.sort_values();
                 }
                 _ => {}
             }
-        }
+        });
     }
 
     /// Sort Key/Value Pairs of the table using the using the comparison function `compare`.
@@ -123,14 +123,14 @@ impl InlineTable {
             };
 
         self.items.sort_by(modified_cmp);
-        for value in self.items.values_mut() {
+        modify_key_position(&mut self.items,|value| {
             match value {
                 Item::Value(Value::InlineTable(table)) if table.is_dotted() => {
                     table.sort_values_by_internal(compare);
                 }
                 _ => {}
             }
-        }
+        });
     }
 
     /// If a table has no key/value pairs and implicit, it will not be displayed.
diff --git a/crates/toml_edit/src/table.rs b/crates/toml_edit/src/table.rs
index 0d0c5b73..0427f9dc 100644
--- a/crates/toml_edit/src/table.rs
+++ b/crates/toml_edit/src/table.rs
@@ -1,6 +1,6 @@
 use std::iter::FromIterator;
 
-use indexmap::map::IndexMap;
+use indexmap::map::{IndexMap, MutableKeys};
 
 use crate::key::Key;
 use crate::repr::Decor;
@@ -110,14 +110,14 @@ impl Table {
     pub fn sort_values(&mut self) {
         // Assuming standard tables have their doc_position set and this won't negatively impact them
         self.items.sort_keys();
-        for value in self.items.values_mut() {
+        modify_key_position(&mut self.items,|value| {
             match value {
                 Item::Table(table) if table.is_dotted() => {
                     table.sort_values();
                 }
                 _ => {}
             }
-        }
+        });
     }
 
     /// Sort Key/Value Pairs of the table using the using the comparison function `compare`.
@@ -142,14 +142,14 @@ impl Table {
 
         self.items.sort_by(modified_cmp);
 
-        for value in self.items.values_mut() {
+       modify_key_position(&mut self.items,|value| {
             match value {
                 Item::Table(table) if table.is_dotted() => {
                     table.sort_values_by_internal(compare);
                 }
                 _ => {}
             }
-        }
+        });
     }
 
     /// If a table has no key/value pairs and implicit, it will not be displayed.
@@ -531,6 +531,16 @@ pub(crate) fn sort_values_by_position<'s>(values: &mut Vec<(Vec<&'s Key>, &'s Va
     });
 }
 
+pub(crate) fn modify_key_position<F>(items: &mut KeyValuePairs, mut recursive_step: F)
+where
+    F: FnMut(&mut Item),
+{
+    for (pos, (key, value)) in items.iter_mut2().enumerate() {
+        key.set_position(Some(pos));
+        recursive_step(value);
+    }
+}
+
 // `key1 = value1`
 pub(crate) const DEFAULT_ROOT_DECOR: (&str, &str) = ("", "");
 pub(crate) const DEFAULT_KEY_DECOR: (&str, &str) = ("", " ");
diff --git a/crates/toml_edit/tests/testsuite/edit.rs b/crates/toml_edit/tests/testsuite/edit.rs
index 84c53d05..befc7f2e 100644
--- a/crates/toml_edit/tests/testsuite/edit.rs
+++ b/crates/toml_edit/tests/testsuite/edit.rs
@@ -579,6 +579,42 @@ fn test_sort_values() {
 "#]]);
 }
 
+#[test]
+fn test_sort_dotted_values() {
+    given(
+        r#"
+        [a.z]
+
+        [a]
+        a.b = 2
+        # this comment is attached to b
+        b = 3 # as well as this
+        a.a = 1
+        c = 4
+
+        [a.y]"#,
+    )
+        .running(|root| {
+            let a = root.get_mut("a").unwrap();
+            let a = as_table!(a);
+            a.sort_values();
+        })
+        .produces_display(str![[r#"
+
+        [a.z]
+
+        [a]
+        a.a = 1
+        a.b = 2
+        # this comment is attached to b
+        b = 3 # as well as this
+        c = 4
+
+        [a.y]
+
+"#]]);
+}
+
 #[test]
 fn test_sort_values_by() {
     given(