fix(edit): Preserve previous line decor on leaf key

We tracked one decor per key.
When dealing with dotted keys,
we used the decor prefix from the first key.
The problem is that we then re-encode that key for each dotted key,
duplicating any comments or whitespace associated with it.

We solved this by tracking decor between dotted keys separate from the
decor used per final, or leaf, entry.

Fixes #673

crates/toml_edit/src/encode.rs

@@ -34,25 +34,26 @@ fn encode_key_path(
     input: Option<&str>,
     default_decor: (&str, &str),
 ) -> Result {
+    let leaf_decor = this.last().expect("always at least one key").leaf_decor();
     for (i, key) in this.iter().enumerate() {
-        let decor = key.decor();
+        let dotted_decor = key.dotted_decor();
 
         let first = i == 0;
         let last = i + 1 == this.len();
 
         if first {
-            decor.prefix_encode(buf, input, default_decor.0)?;
+            leaf_decor.prefix_encode(buf, input, default_decor.0)?;
         } else {
             write!(buf, ".")?;
-            decor.prefix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.0)?;
+            dotted_decor.prefix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.0)?;
         }
 
         encode_key(key, buf, input)?;
 
         if last {
-            decor.suffix_encode(buf, input, default_decor.1)?;
+            leaf_decor.suffix_encode(buf, input, default_decor.1)?;
         } else {
-            decor.suffix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.1)?;
+            dotted_decor.suffix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.1)?;
         }
     }
     Ok(())
@@ -64,25 +65,26 @@ pub(crate) fn encode_key_path_ref(
     input: Option<&str>,
     default_decor: (&str, &str),
 ) -> Result {
+    let leaf_decor = this.last().expect("always at least one key").leaf_decor();
     for (i, key) in this.iter().enumerate() {
-        let decor = key.decor();
+        let dotted_decor = key.dotted_decor();
 
         let first = i == 0;
         let last = i + 1 == this.len();
 
         if first {
-            decor.prefix_encode(buf, input, default_decor.0)?;
+            leaf_decor.prefix_encode(buf, input, default_decor.0)?;
         } else {
             write!(buf, ".")?;
-            decor.prefix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.0)?;
+            dotted_decor.prefix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.0)?;
         }
 
         encode_key(key, buf, input)?;
 
         if last {
-            decor.suffix_encode(buf, input, default_decor.1)?;
+            leaf_decor.suffix_encode(buf, input, default_decor.1)?;
         } else {
-            decor.suffix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.1)?;
+            dotted_decor.suffix_encode(buf, input, DEFAULT_KEY_PATH_DECOR.1)?;
         }
     }
     Ok(())
@@ -240,11 +242,7 @@ where
     for kv in table.items.values() {
         match kv.value {
             Item::Table(ref t) => {
-                let mut key = kv.key.clone();
-                if t.is_dotted() {
-                    // May have newlines and generally isn't written for standard tables
-                    key.decor_mut().clear();
-                }
+                let key = kv.key.clone();
                 path.push(key);
                 visit_nested_tables(t, path, false, callback)?;
                 path.pop();

crates/toml_edit/src/inline_table.rs

@@ -185,14 +185,23 @@ impl InlineTable {
         &self.decor
     }
 
+    /// Returns an accessor to a key's formatting
+    pub fn key_mut(&mut self, key: &str) -> Option<KeyMut<'_>> {
+        self.items.get_mut(key).map(|kv| kv.key.as_mut())
+    }
+
     /// Returns the decor associated with a given key of the table.
+    #[deprecated(since = "0.21.1", note = "Replaced with `key_mut`")]
     pub fn key_decor_mut(&mut self, key: &str) -> Option<&mut Decor> {
-        self.items.get_mut(key).map(|kv| &mut kv.key.decor)
+        #![allow(deprecated)]
+        self.items.get_mut(key).map(|kv| kv.key.leaf_decor_mut())
     }
 
     /// Returns the decor associated with a given key of the table.
+    #[deprecated(since = "0.21.1", note = "Replaced with `key_mut`")]
     pub fn key_decor(&self, key: &str) -> Option<&Decor> {
-        self.items.get(key).map(|kv| &kv.key.decor)
+        #![allow(deprecated)]
+        self.items.get(key).map(|kv| kv.key.leaf_decor())
     }
 
     /// Set whitespace after before element
@@ -469,13 +478,14 @@ impl<'s> IntoIterator for &'s InlineTable {
 }
 
 fn decorate_inline_table(table: &mut InlineTable) {
-    for (key_decor, value) in table
+    for (mut key, value) in table
         .items
         .iter_mut()
         .filter(|(_, kv)| kv.value.is_value())
-        .map(|(_, kv)| (&mut kv.key.decor, kv.value.as_value_mut().unwrap()))
+        .map(|(_, kv)| (kv.key.as_mut(), kv.value.as_value_mut().unwrap()))
     {
-        key_decor.clear();
+        key.leaf_decor_mut().clear();
+        key.dotted_decor_mut().clear();
         value.decor_mut().clear();
     }
 }
@@ -563,10 +573,15 @@ impl TableLike for InlineTable {
         self.is_dotted()
     }
 
+    fn key_mut(&mut self, key: &str) -> Option<KeyMut<'_>> {
+        self.key_mut(key)
+    }
     fn key_decor_mut(&mut self, key: &str) -> Option<&mut Decor> {
+        #![allow(deprecated)]
         self.key_decor_mut(key)
     }
     fn key_decor(&self, key: &str) -> Option<&Decor> {
+        #![allow(deprecated)]
         self.key_decor(key)
     }
 }

crates/toml_edit/src/key.rs

@@ -30,7 +30,8 @@ use crate::InternalString;
 pub struct Key {
     key: InternalString,
     pub(crate) repr: Option<Repr>,
-    pub(crate) decor: Decor,
+    pub(crate) leaf_decor: Decor,
+    pub(crate) dotted_decor: Decor,
 }
 
 impl Key {
@@ -39,7 +40,8 @@ impl Key {
         Self {
             key: key.into(),
             repr: None,
-            decor: Default::default(),
+            leaf_decor: Default::default(),
+            dotted_decor: Default::default(),
         }
     }
 
@@ -57,8 +59,20 @@ impl Key {
     }
 
     /// While creating the `Key`, add `Decor` to it
-    pub fn with_decor(mut self, decor: Decor) -> Self {
-        self.decor = decor;
+    #[deprecated(since = "0.21.1", note = "Replaced with `with_leaf_decor`")]
+    pub fn with_decor(self, decor: Decor) -> Self {
+        self.with_leaf_decor(decor)
+    }
+
+    /// While creating the `Key`, add `Decor` to it for the line entry
+    pub fn with_leaf_decor(mut self, decor: Decor) -> Self {
+        self.leaf_decor = decor;
+        self
+    }
+
+    /// While creating the `Key`, add `Decor` to it for between dots
+    pub fn with_dotted_decor(mut self, decor: Decor) -> Self {
+        self.dotted_decor = decor;
         self
     }
 
@@ -99,13 +113,35 @@ impl Key {
     }
 
     /// Returns the surrounding whitespace
+    #[deprecated(since = "0.21.1", note = "Replaced with `decor_mut`")]
     pub fn decor_mut(&mut self) -> &mut Decor {
-        &mut self.decor
+        self.leaf_decor_mut()
+    }
+
+    /// Returns the surrounding whitespace for the line entry
+    pub fn leaf_decor_mut(&mut self) -> &mut Decor {
+        &mut self.leaf_decor
+    }
+
+    /// Returns the surrounding whitespace for between dots
+    pub fn dotted_decor_mut(&mut self) -> &mut Decor {
+        &mut self.dotted_decor
     }
 
     /// Returns the surrounding whitespace
+    #[deprecated(since = "0.21.1", note = "Replaced with `decor`")]
     pub fn decor(&self) -> &Decor {
-        &self.decor
+        self.leaf_decor()
+    }
+
+    /// Returns the surrounding whitespace for the line entry
+    pub fn leaf_decor(&self) -> &Decor {
+        &self.leaf_decor
+    }
+
+    /// Returns the surrounding whitespace for between dots
+    pub fn dotted_decor(&self) -> &Decor {
+        &self.dotted_decor
     }
 
     /// Returns the location within the original document
@@ -115,7 +151,8 @@ impl Key {
     }
 
     pub(crate) fn despan(&mut self, input: &str) {
-        self.decor.despan(input);
+        self.leaf_decor.despan(input);
+        self.dotted_decor.despan(input);
         if let Some(repr) = &mut self.repr {
             repr.despan(input)
         }
@@ -124,7 +161,8 @@ impl Key {
     /// Auto formats the key.
     pub fn fmt(&mut self) {
         self.repr = None;
-        self.decor.clear();
+        self.leaf_decor.clear();
+        self.dotted_decor.clear();
     }
 
     #[cfg(feature = "parse")]
@@ -150,7 +188,8 @@ impl Clone for Key {
         Self {
             key: self.key.clone(),
             repr: self.repr.clone(),
-            decor: self.decor.clone(),
+            leaf_decor: self.leaf_decor.clone(),
+            dotted_decor: self.dotted_decor.clone(),
         }
     }
 }
@@ -291,7 +330,7 @@ impl From<Key> for InternalString {
     }
 }
 
-/// A mutable reference to a `Key`
+/// A mutable reference to a `Key`'s formatting
 #[derive(Debug, Eq, PartialEq, PartialOrd, Ord, Hash)]
 pub struct KeyMut<'k> {
     key: &'k mut Key,
@@ -321,15 +360,39 @@ impl<'k> KeyMut<'k> {
     }
 
     /// Returns the surrounding whitespace
+    #[deprecated(since = "0.21.1", note = "Replaced with `decor_mut`")]
     pub fn decor_mut(&mut self) -> &mut Decor {
+        #![allow(deprecated)]
         self.key.decor_mut()
     }
 
+    /// Returns the surrounding whitespace for the line entry
+    pub fn leaf_decor_mut(&mut self) -> &mut Decor {
+        self.key.leaf_decor_mut()
+    }
+
+    /// Returns the surrounding whitespace for between dots
+    pub fn dotted_decor_mut(&mut self) -> &mut Decor {
+        self.key.dotted_decor_mut()
+    }
+
     /// Returns the surrounding whitespace
+    #[deprecated(since = "0.21.1", note = "Replaced with `decor`")]
     pub fn decor(&self) -> &Decor {
+        #![allow(deprecated)]
         self.key.decor()
     }
 
+    /// Returns the surrounding whitespace for the line entry
+    pub fn leaf_decor(&self) -> &Decor {
+        self.key.leaf_decor()
+    }
+
+    /// Returns the surrounding whitespace for between dots
+    pub fn dotted_decor(&self) -> &Decor {
+        self.key.dotted_decor()
+    }
+
     /// Auto formats the key.
     pub fn fmt(&mut self) {
         self.key.fmt()

crates/toml_edit/src/parser/key.rs

@@ -18,13 +18,13 @@ use crate::RawString;
 // key = simple-key / dotted-key
 // dotted-key = simple-key 1*( dot-sep simple-key )
 pub(crate) fn key(input: &mut Input<'_>) -> PResult<Vec<Key>> {
-    trace(
+    let mut key_path = trace(
         "dotted-key",
         separated1(
             (ws.span(), simple_key, ws.span()).map(|(pre, (raw, key), suffix)| {
                 Key::new(key)
                     .with_repr_unchecked(Repr::new_unchecked(raw))
-                    .with_decor(Decor::new(
+                    .with_dotted_decor(Decor::new(
                         RawString::with_span(pre),
                         RawString::with_span(suffix),
                     ))
@@ -38,7 +38,31 @@ pub(crate) fn key(input: &mut Input<'_>) -> PResult<Vec<Key>> {
             Ok::<_, CustomError>(k)
         }),
     )
-    .parse_next(input)
+    .parse_next(input)?;
+
+    let mut leaf_decor = Decor::new("", "");
+    {
+        let first_dotted_decor = key_path
+            .first_mut()
+            .expect("always at least one key")
+            .dotted_decor_mut();
+        if let Some(prefix) = first_dotted_decor.prefix().cloned() {
+            leaf_decor.set_prefix(prefix);
+            first_dotted_decor.set_prefix("");
+        }
+    }
+    let last_key = &mut key_path.last_mut().expect("always at least one key");
+    {
+        let last_dotted_decor = last_key.dotted_decor_mut();
+        if let Some(suffix) = last_dotted_decor.suffix().cloned() {
+            leaf_decor.set_suffix(suffix);
+            last_dotted_decor.set_suffix("");
+        }
+    }
+
+    *last_key.leaf_decor_mut() = leaf_decor;
+
+    Ok(key_path)
 }
 
 // simple-key = quoted-key / unquoted-key

crates/toml_edit/src/parser/state.rs

@@ -39,26 +39,21 @@ impl ParseState {
 
     pub(crate) fn on_keyval(
         &mut self,
-        mut path: Vec<Key>,
+        path: Vec<Key>,
         mut kv: TableKeyValue,
     ) -> Result<(), CustomError> {
         {
             let mut prefix = self.trailing.take();
-            let first_key = if path.is_empty() {
-                &mut kv.key
-            } else {
-                &mut path[0]
-            };
             let prefix = match (
                 prefix.take(),
-                first_key.decor.prefix().and_then(|d| d.span()),
+                kv.key.leaf_decor.prefix().and_then(|d| d.span()),
             ) {
                 (Some(p), Some(k)) => Some(p.start..k.end),
                 (Some(p), None) | (None, Some(p)) => Some(p),
                 (None, None) => None,
             };
-            first_key
-                .decor
+            kv.key
+                .leaf_decor
                 .set_prefix(prefix.map(RawString::with_span).unwrap_or_default());
         }