1 files changed, 31 insertions, 5 deletions

diff --git a/src/parser/key.rs b/src/parser/key.rs
index 12715da..e72b195 100644
--- a/src/parser/key.rs
+++ b/src/parser/key.rs
@@ -7,7 +7,7 @@ use winnow::token::take_while;
 use winnow::trace::trace;
 
 use crate::key::Key;
-use crate::parser::errors::CustomError;
+use crate::parser::error::CustomError;
 use crate::parser::prelude::*;
 use crate::parser::strings::{basic_string, literal_string};
 use crate::parser::trivia::{from_utf8_unchecked, ws};
@@ -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
@@ -88,6 +112,8 @@ const UNQUOTED_CHAR: (
 const DOT_SEP: u8 = b'.';
 
 #[cfg(test)]
+#[cfg(feature = "parse")]
+#[cfg(feature = "display")]
 mod test {
     use super::*;
 
@@ -96,7 +122,7 @@ mod test {
         let cases = [
             ("a", "a"),
             (r#""hello\n ""#, "hello\n "),
-            (r#"'hello\n '"#, "hello\\n "),
+            (r"'hello\n '", "hello\\n "),
         ];
 
         for (input, expected) in cases {