From 44f59f415a7f2beafce4db925077581c41a2d561 Mon Sep 17 00:00:00 2001
From: Lukas Kalbertodt <lukas.kalbertodt@gmail.com>
Date: Wed, 19 Oct 2022 17:10:17 +0200
Subject: [PATCH] Implement `Serialize` for `meta::Expr` and use it in
 `PrintExpr`

This probably fixes a few edge case bugs (float infinity and stuff) and
results in less code.
---
 src/meta.rs |  9 ++++++---
 src/toml.rs | 27 +++++++++------------------
 src/yaml.rs | 28 +++++++++++++++-------------
 3 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/src/meta.rs b/src/meta.rs
index c5c9531..533268f 100644
--- a/src/meta.rs
+++ b/src/meta.rs
@@ -47,7 +47,8 @@ pub enum LeafKind {
     Optional,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize)]
+#[serde(untagged)]
 pub enum Expr {
     Str(&'static str),
     Float(Float),
@@ -56,13 +57,15 @@ pub enum Expr {
     Array(&'static [Expr]),
 }
 
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize)]
+#[serde(untagged)]
 pub enum Float {
     F32(f32),
     F64(f64),
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
+#[serde(untagged)]
 pub enum Integer {
     U8(u8),
     U16(u16),
diff --git a/src/toml.rs b/src/toml.rs
index 6fbb000..68be572 100644
--- a/src/toml.rs
+++ b/src/toml.rs
@@ -216,26 +216,17 @@ fn format_impl(
 /// Helper to emit `meta::Expr` into TOML.
 struct PrintExpr(&'static Expr);
 
+impl From<&'static Expr> for PrintExpr {
+    fn from(expr: &'static Expr) -> Self {
+        Self(expr)
+    }
+}
+
 impl fmt::Display for PrintExpr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match *self.0 {
-            Expr::Str(v) => toml::Value::String(v.to_owned()).fmt(f),
-            Expr::Float(v) => v.fmt(f),
-            Expr::Integer(v) => v.fmt(f),
-            Expr::Bool(v) => v.fmt(f),
-            Expr::Array(items) => {
-                // TODO: pretty printing of long arrays onto multiple lines?
-                f.write_char('[')?;
-                for (i, item) in items.iter().enumerate() {
-                    if i != 0 {
-                        f.write_str(", ")?;
-                    }
-                    PrintExpr(item).fmt(f)?;
-                }
-                f.write_char(']')?;
-                Ok(())
-            },
-        }
+        toml::to_string(&self.0)
+            .expect("string serialization to TOML failed")
+            .fmt(f)
     }
 }
 
diff --git a/src/yaml.rs b/src/yaml.rs
index eb3a15a..a27f27d 100644
--- a/src/yaml.rs
+++ b/src/yaml.rs
@@ -201,19 +201,8 @@ struct PrintExpr(&'static Expr);
 impl fmt::Display for PrintExpr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match *self.0 {
-            Expr::Str(v) => {
-                // This is a bit ugly. Sadly, no YAML crate in our dependency
-                // tree has an API to serialize a string only, without emitting
-                // the `---` at the start of the document. But instead of
-                // implementing the quoting logic ourselves (which is really
-                // complicated as it turns out!), we use this hack.
-                let value = serde_yaml::Value::String(v.to_owned());
-                let serialized = serde_yaml::to_string(&value).unwrap();
-                serialized[4..].trim_end_matches('\n').fmt(f)
-            },
-            Expr::Float(v) => v.fmt(f),
-            Expr::Integer(v) => v.fmt(f),
-            Expr::Bool(v) => v.fmt(f),
+            // We have to special case arrays as the normal formatter only emits
+            // multi line lists.
             Expr::Array(items) => {
                 // TODO: pretty printing of long arrays onto multiple lines?
                 f.write_char('[')?;
@@ -226,6 +215,19 @@ impl fmt::Display for PrintExpr {
                 f.write_char(']')?;
                 Ok(())
             }
+
+            // All these other types can simply be serialized as is.
+            Expr::Str(_) | Expr::Float(_) | Expr::Integer(_) | Expr::Bool(_) => {
+                let out = serde_yaml::to_string(&self.0).expect("string serialization to YAML failed");
+
+                // Unfortunately, `serde_yaml` cannot serialize these values on its own
+                // without embedding them in a full document (starting with `---` and
+                // ending with a newline). So we need to cleanup.
+                out.strip_prefix("---\n")
+                    .unwrap_or(&out)
+                    .trim_matches('\n')
+                    .fmt(f)
+            }
         }
     }
 }