improve RawString handling

xxchan · Oct 24, 2023 · 13884c8 · 13884c8
1 parent 09b43a1
commit 13884c8
Showing 1 changed file with 33 additions and 25 deletions.
diff --git a/src/formatting/cargo_toml.rs b/src/formatting/cargo_toml.rs
@@ -1,7 +1,8 @@
 use itertools::Itertools;
-use std::cmp::Ordering;
+use std::{cmp::Ordering, os::macos::raw};
 use toml_edit::{
-    visit_mut::*, Decor, Document, Formatted, Item, KeyMut, Table, TableLike, TomlError, Value,
+    visit_mut::*, Decor, Document, Formatted, Item, KeyMut, RawString, Table, TableLike, TomlError,
+    Value,
 };
 
 use crate::{Config, ErrorKind};
@@ -190,10 +191,12 @@ impl BlankLine {
     }
 
     fn trim_decor_blank_lines(decor: &mut Decor) {
-        let prefix = decor.prefix().cloned().unwrap_or_default();
-        let suffix = decor.suffix().cloned().unwrap_or_default();
-        decor.set_prefix(Self::trim_blank_lines(prefix.as_str().unwrap_or_default()));
-        decor.set_suffix(Self::trim_blank_lines(suffix.as_str().unwrap_or_default()));
+        if let Some(prefix) = decor.prefix().map(raw_string_as_str) {
+            decor.set_prefix(Self::trim_blank_lines(prefix));
+        }
+        if let Some(suffix) = decor.suffix().map(raw_string_as_str) {
+            decor.set_suffix(Self::trim_blank_lines(suffix));
+        }
     }
 }
 
@@ -221,8 +224,12 @@ impl VisitMut for BlankLine {
             });
         } else {
             let decor = table.decor_mut();
-            let prefix = decor.prefix().cloned().unwrap_or_default();
-            decor.set_prefix(format!("\n{}", prefix.as_str().unwrap_or_default()));
+            if let Some(prefix) = decor.prefix() {
+                decor.set_prefix(format!(
+                    "\n{}",
+                    prefix.as_str().expect("should already be despanded")
+                ));
+            }
         }
     }
 }
@@ -237,7 +244,6 @@ impl KeyValue {
 
 impl VisitMut for KeyValue {
     fn visit_table_like_kv_mut(&mut self, mut key: KeyMut<'_>, value: &mut Item) {
-        let prefix = key.decor().prefix().cloned().unwrap_or_default();
         if Self::can_be_bare_key(key.get()) {
             // will remove decors and set the key to the bare key
             key.fmt();
@@ -246,7 +252,8 @@ impl VisitMut for KeyValue {
             key.decor_mut().set_suffix(" ");
         }
         // start all key names at the start of a line, but preserve comments
-        if let Some(prefix) = prefix.as_str() {
+        if let Some(prefix) = key.decor().prefix().map(raw_string_as_str) {
+            let prefix = prefix.to_string();
             key.decor_mut()
                 .set_prefix(prefix.trim_end_matches(|c: char| c.is_whitespace() && c != '\n'));
         }
@@ -404,19 +411,14 @@ impl VisitMut for TrimSpaces {
         self.visit_table_mut(node);
 
         let set_prefix = |decor: &mut Decor, i: usize| {
-            let prefix = format!(
-                "{}{}",
-                if i == 0 { "" } else { "\n" },
-                Self::trim_block(
-                    decor
-                        .prefix()
-                        .cloned()
-                        .unwrap_or_default()
-                        .as_str()
-                        .unwrap_or_default()
-                )
-            );
-            decor.set_prefix(prefix);
+            if let Some(prefix) = decor.prefix().map(raw_string_as_str) {
+                let prefix = format!(
+                    "{}{}",
+                    if i == 0 { "" } else { "\n" },
+                    Self::trim_block(prefix)
+                );
+                decor.set_prefix(prefix);
+            }
         };
         let table = node.as_table_mut();
         for (i, (_, item)) in table.iter_mut().enumerate() {
@@ -429,9 +431,9 @@ impl VisitMut for TrimSpaces {
             }
         }
 
-        let trailing = node.trailing().as_str().unwrap_or_default();
+        let trailing = raw_string_as_str(node.trailing());
         if !trailing.trim().is_empty() {
-            let trailing: String = Self::trim_block(trailing);
+            let trailing = Self::trim_block(trailing);
             node.set_trailing(&format!("\n{trailing}"));
         } else {
             node.set_trailing("");
@@ -470,3 +472,9 @@ fn trim_suffix(decor: &mut Decor) {
         }
     }
 }
+
+/// Note: in `Document::from_str`, the document is despanned, so we can safely unwrap `as_str`
+/// when handling `RawString`.
+fn raw_string_as_str(raw_string: &RawString) -> &str {
+    raw_string.as_str().expect("should already be despanded")
+}