From 43ac4faac4780e45c55708dbaf55d5d95614e858 Mon Sep 17 00:00:00 2001
From: Yuya Nishihara <yuya@tcha.org>
Date: Sat, 28 Dec 2024 11:18:40 +0900
Subject: [PATCH] config: preserve key formatting on set_value()

I just noticed toml_edit::Table has a format-preserving insert() API. I think
it's better to retain the user-specified quoting style.
---
 cli/tests/test_config_command.rs |  6 +++---
 lib/src/config.rs                | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/cli/tests/test_config_command.rs b/cli/tests/test_config_command.rs
index 55b1fe9187..563a8f54d2 100644
--- a/cli/tests/test_config_command.rs
+++ b/cli/tests/test_config_command.rs
@@ -522,13 +522,13 @@ fn test_config_set_for_user() {
     // Ensure test-key successfully written to user config.
     let user_config_toml = std::fs::read_to_string(&user_config_path)
         .unwrap_or_else(|_| panic!("Failed to read file {}", user_config_path.display()));
-    insta::assert_snapshot!(user_config_toml, @r###"
+    insta::assert_snapshot!(user_config_toml, @r#"
     test-key = "test-val"
 
     [test-table]
     foo = true
-    "bar()" = 0
-    "###);
+    'bar()' = 0
+    "#);
 }
 
 #[test]
diff --git a/lib/src/config.rs b/lib/src/config.rs
index 00f208c933..9d79ea79a3 100644
--- a/lib/src/config.rs
+++ b/lib/src/config.rs
@@ -397,7 +397,7 @@ impl ConfigLayer {
             .map_err(|keys| ConfigUpdateError::WouldOverwriteValue {
                 name: keys.join("."),
             })?;
-        match parent_table.entry(leaf_key) {
+        match parent_table.entry_format(leaf_key) {
             toml_edit::Entry::Occupied(mut entry) => {
                 if !entry.get().is_value() {
                     return Err(ConfigUpdateError::WouldOverwriteTable {
@@ -409,6 +409,9 @@ impl ConfigLayer {
             }
             toml_edit::Entry::Vacant(entry) => {
                 entry.insert(toml_edit::value(new_value));
+                // Reset whitespace formatting (i.e. insert space before '=')
+                let mut new_key = parent_table.key_mut(leaf_key).unwrap();
+                new_key.leaf_decor_mut().clear();
                 Ok(None)
             }
         }
@@ -476,7 +479,7 @@ fn ensure_parent_table<'a, 'b>(
     let mut keys = name.components();
     let leaf_key = keys.next_back().ok_or(&name.0[..])?;
     let parent_table = keys.enumerate().try_fold(root_table, |table, (i, key)| {
-        let sub_item = table.entry(key).or_insert_with(new_implicit_table);
+        let sub_item = table.entry_format(key).or_insert_with(new_implicit_table);
         sub_item.as_table_mut().ok_or(&name.0[..=i])
     })?;
     Ok((parent_table, leaf_key))
@@ -918,6 +921,31 @@ mod tests {
         "#);
     }
 
+    #[test]
+    fn test_config_layer_set_value_formatting() {
+        let mut layer = ConfigLayer::empty(ConfigSource::User);
+        // Quoting style should be preserved on insertion
+        layer
+            .set_value(
+                "'foo' . bar . 'baz'",
+                ConfigValue::from_str("'value'").unwrap(),
+            )
+            .unwrap();
+        insta::assert_snapshot!(layer.data, @r"
+        ['foo' . bar]
+        'baz' = 'value'
+        ");
+
+        // Style of existing keys isn't updated
+        layer.set_value("foo.bar.baz", "new value").unwrap();
+        layer.set_value("foo.'bar'.blah", 0).unwrap();
+        insta::assert_snapshot!(layer.data, @r#"
+        ['foo' . bar]
+        'baz' = "new value"
+        blah = 0
+        "#);
+    }
+
     #[test]
     fn test_config_layer_delete_value() {
         let mut layer = ConfigLayer::empty(ConfigSource::User);