fix: Preserve dotted-key ordering

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};
+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,
@@ -75,6 +75,7 @@ impl InlineTable {
                 _ => {}
             }
         }
+        sort_values_by_position(values);
     }
 
     /// Auto formats the table.
@@ -86,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`.
@@ -122,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.

crates/toml_edit/src/key.rs

@@ -32,6 +32,7 @@ pub struct Key {
     pub(crate) repr: Option<Repr>,
     pub(crate) leaf_decor: Decor,
     pub(crate) dotted_decor: Decor,
+    pub(crate) position: Option<usize>,
 }
 
 impl Key {
@@ -42,6 +43,7 @@ impl Key {
             repr: None,
             leaf_decor: Default::default(),
             dotted_decor: Default::default(),
+            position: Default::default(),
         }
     }
 
@@ -76,6 +78,12 @@ impl Key {
         self
     }
 
+    /// While creating the `Key`, add a table position to it
+    pub fn with_position(mut self, position: Option<usize>) -> Self {
+        self.position = position;
+        self
+    }
+
     /// Access a mutable proxy for the `Key`.
     pub fn as_mut(&mut self) -> KeyMut<'_> {
         KeyMut { key: self }
@@ -158,6 +166,16 @@ impl Key {
         }
     }
 
+    /// Get the position relative to other keys in parent table
+    pub fn position(&self) -> Option<usize> {
+        self.position
+    }
+
+    /// Set the position relative to other keys in parent table
+    pub fn set_position(&mut self, position: Option<usize>) {
+        self.position = position;
+    }
+
     /// Auto formats the key.
     pub fn fmt(&mut self) {
         self.repr = None;
@@ -190,6 +208,7 @@ impl Clone for Key {
             repr: self.repr.clone(),
             leaf_decor: self.leaf_decor.clone(),
             dotted_decor: self.dotted_decor.clone(),
+            position: self.position,
         }
     }
 }

crates/toml_edit/src/parser/inline_table.rs

@@ -40,7 +40,8 @@ fn table_from_pairs(
     // Assuming almost all pairs will be directly in `root`
     root.items.reserve(v.len());
 
-    for (path, (key, value)) in v {
+    for (position, (path, (mut key, value))) in v.into_iter().enumerate() {
+        key.set_position(Some(position));
         let table = descend_path(&mut root, &path)?;
 
         // "Likewise, using dotted keys to redefine tables already defined in [table] form is not allowed"
@@ -162,6 +163,7 @@ mod test {
             r#"{a = 1e165}"#,
             r#"{ hello = "world", a = 1}"#,
             r#"{ hello.world = "a" }"#,
+            r#"{ hello.world = "a", goodbye = "b", hello.moon = "c" }"#,
         ];
         for input in inputs {
             dbg!(input);

crates/toml_edit/src/parser/mod.rs

@@ -216,6 +216,10 @@ that
 key = "value"
 "#,
             r#"hello.world = "a"
+"#,
+            r#"hello.world = "a"
+goodbye = "b"
+hello.moon = "c"
 "#,
             r#"foo = 1979-05-27 # Comment
 "#,

crates/toml_edit/src/parser/state.rs

@@ -7,6 +7,8 @@ pub(crate) struct ParseState {
     root: Table,
     trailing: Option<std::ops::Range<usize>>,
     current_table_position: usize,
+    // Current position within a table, to help order dotted keys
+    current_value_position: usize,
     current_table: Table,
     current_is_array: bool,
     current_table_path: Vec<Key>,
@@ -20,6 +22,7 @@ impl ParseState {
             root: Table::new(),
             trailing: None,
             current_table_position: 0,
+            current_value_position: 0,
             current_table: root,
             current_is_array: false,
             current_table_path: Vec::new(),
@@ -74,6 +77,9 @@ impl ParseState {
         if let (Some(existing), Some(value)) = (self.current_table.span(), value.span()) {
             self.current_table.span = Some((existing.start)..(value.end));
         }
+        key.set_position(Some(self.current_value_position));
+        self.current_value_position += 1;
+
         let table = &mut self.current_table;
         let table = Self::descend_path(table, &path, true)?;
 
@@ -161,6 +167,7 @@ impl ParseState {
         }
 
         self.current_table_position += 1;
+        self.current_value_position = 0;
         self.current_table.decor = decor;
         self.current_table.set_implicit(false);
         self.current_table.set_dotted(false);

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;
@@ -96,6 +96,7 @@ impl Table {
                 _ => {}
             }
         }
+        sort_values_by_position(values);
     }
 
     /// Auto formats the table.
@@ -109,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`.
@@ -141,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.
@@ -518,6 +519,36 @@ fn decorate_table(table: &mut Table) {
     }
 }
 
+/// `Vec::sort_by_key` works because we add the position to _every_ item's key during parsing,
+/// so keys without positions would be either:
+///    1. non-leaf keys (i.e. "foo" or "bar" in dotted key "foo.bar.baz")
+///    2. custom keys added to the doc without an explicit position
+///
+/// In the case of (1), we'd never see it since we only look at the last
+/// key in a dotted-key. So, we can safely return a constant value for these cases.
+///
+/// To support the most intuitive behavior, we return the maximum usize, placing
+/// position=None items at the end, so when you insert it without position, it
+/// appends it to the end.
+pub(crate) fn sort_values_by_position<'s>(values: &mut [(Vec<&'s Key>, &'s Value)]) {
+    values.sort_by_key(|(left_kp, _)| {
+        return match left_kp.last().map(|x| x.position) {
+            Some(Some(pos)) => pos,
+            _ => usize::MAX
+        }
+    });
+}
+
+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) = ("", " ");

crates/toml_edit/tests/testsuite/edit.rs

@@ -26,8 +26,12 @@ struct Test {
 
 fn given(input: &str) -> Test {
     let doc = input.parse::<DocumentMut>();
-    assert!(doc.is_ok());
-    Test { doc: doc.unwrap() }
+    match doc {
+        Err(e) => panic!("{}", e),
+        _ => Test { doc: doc.unwrap() }
+    }
+    // assert!(doc.is_ok());
+    // Test { doc: doc.unwrap() }
 }
 
 impl Test {
@@ -579,6 +583,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(
@@ -1026,3 +1066,68 @@ name = "test.swf"
 "#]]
     );
 }
+
+#[test]
+fn insert_key_with_position() {
+    given(r#"foo.bar = 1
+
+foo.baz.qwer = "a"
+goodbye = { forever="no" }
+
+foo.pub = 2
+foo.baz.asdf = "b"
+"#)
+    .running_on_doc(|doc| {
+        let root = doc.as_table_mut();
+        let (_, foo_v) = root.get_key_value_mut("foo").unwrap();
+        let foo = as_table!(foo_v);
+        let (_, foo_baz_v) = foo.get_key_value_mut("baz").unwrap();
+        let foo_baz = as_table!(foo_baz_v);
+
+        let foo_baz_asdf_v = foo_baz.remove("asdf").unwrap();
+        let foo_baz_asdf_k = Key::new("asdf").with_position(Some(0));
+
+        foo_baz.insert_formatted(&foo_baz_asdf_k, foo_baz_asdf_v);
+    })
+    .produces_display(str![[r#"foo.bar = 1
+foo.baz.asdf = "b"
+
+foo.baz.qwer = "a"
+goodbye = { forever="no" }
+
+foo.pub = 2
+
+"#]]);
+}
+
+#[test]
+fn insert_key_with_position_none() {
+    given(r#"foo.bar = 1
+
+foo.baz.qwer = "a"
+goodbye = { forever="no" }
+
+foo.pub = 2
+foo.baz.asdf = "b"
+"#)
+    .running_on_doc(|doc| {
+        let root = doc.as_table_mut();
+        let (_, foo_v) = root.get_key_value_mut("foo").unwrap();
+        let foo = as_table!(foo_v);
+        let (_, foo_baz_v) = foo.get_key_value_mut("baz").unwrap();
+        let foo_baz = as_table!(foo_baz_v);
+
+        let foo_baz_asdf_v = foo_baz.remove("qwer").unwrap();
+        let foo_baz_asdf_k = Key::new("qwer").with_position(None);
+
+        foo_baz.insert_formatted(&foo_baz_asdf_k, foo_baz_asdf_v);
+    })
+    .produces_display(str![[r#"foo.bar = 1
+goodbye = { forever="no" }
+
+foo.pub = 2
+foo.baz.asdf = "b"
+foo.baz.qwer = "a"
+
+"#]]);
+}