From 5e3544fcc3df67c5d8ea6303281b868edca089bc Mon Sep 17 00:00:00 2001
From: Nick Rolfe <nickrolfe@github.com>
Date: Wed, 21 Oct 2020 12:45:54 +0100
Subject: [PATCH] Use fmt::Display trait for writing dbscheme

---
 generator/src/dbscheme.rs | 130 ++++++++++++++++++++++----------------
 generator/src/main.rs     |  12 ++--
 2 files changed, 80 insertions(+), 62 deletions(-)

diff --git a/generator/src/dbscheme.rs b/generator/src/dbscheme.rs
index e6503707986..f4296beb605 100644
--- a/generator/src/dbscheme.rs
+++ b/generator/src/dbscheme.rs
@@ -1,10 +1,12 @@
+use std::fmt;
+
 /// Represents a distinct entry in the database schema.
 pub enum Entry {
     /// An entry defining a database table.
     Table(Table),
 
     /// An entry defining type that is a union of other types.
-    Union { name: String, members: Vec<String> },
+    Union(Union),
 }
 
 /// A table in the database schema.
@@ -14,6 +16,12 @@ pub struct Table {
     pub keysets: Option<Vec<String>>,
 }
 
+/// A union in the database schema.
+pub struct Union {
+    pub name: String,
+    pub members: Vec<String>,
+}
+
 /// A column in a table.
 pub struct Column {
     pub db_type: DbColumnType,
@@ -101,6 +109,69 @@ pub fn escape_name(name: &str) -> String {
     result
 }
 
+impl fmt::Display for Table {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for keyset in &self.keysets {
+            write!(f, "#keyset[")?;
+            for (key_index, key) in keyset.iter().enumerate() {
+                if key_index > 0 {
+                    write!(f, ", ")?;
+                }
+                write!(f, "{}", key)?;
+            }
+            write!(f, "]\n")?;
+        }
+
+        write!(f, "{}(\n", self.name)?;
+        for (column_index, column) in self.columns.iter().enumerate() {
+            write!(f, "  ")?;
+            if column.unique {
+                write!(f, "unique ")?;
+            }
+            write!(
+                f,
+                "{} ",
+                match column.db_type {
+                    DbColumnType::Int => "int",
+                    DbColumnType::String => "string",
+                }
+            )?;
+            write!(f, "{}: ", column.name)?;
+            match &column.ql_type {
+                QlColumnType::Int => write!(f, "int")?,
+                QlColumnType::String => write!(f, "string")?,
+                QlColumnType::Custom(name) => write!(f, "@{}", name)?,
+            }
+            if column.ql_type_is_ref {
+                write!(f, " ref")?;
+            }
+            if column_index + 1 != self.columns.len() {
+                write!(f, ",")?;
+            }
+            write!(f, "\n")?;
+        }
+        write!(f, ");")?;
+
+        Ok(())
+    }
+}
+
+impl fmt::Display for Union {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "@{} = ", self.name)?;
+        let mut first = true;
+        for member in &self.members {
+            if first {
+                first = false;
+            } else {
+                write!(f, " | ")?;
+            }
+            write!(f, "@{}", member)?;
+        }
+        Ok(())
+    }
+}
+
 /// Generates the dbscheme by writing the given dbscheme `entries` to the `file`.
 pub fn write(
     language_name: &str,
@@ -115,61 +186,8 @@ pub fn write(
 
     for entry in entries {
         match entry {
-            Entry::Table(table) => {
-                for keyset in &table.keysets {
-                    write!(file, "#keyset[")?;
-                    for (key_index, key) in keyset.iter().enumerate() {
-                        if key_index > 0 {
-                            write!(file, ", ")?;
-                        }
-                        write!(file, "{}", key)?;
-                    }
-                    write!(file, "]\n")?;
-                }
-
-                write!(file, "{}(\n", table.name)?;
-                for (column_index, column) in table.columns.iter().enumerate() {
-                    write!(file, "  ")?;
-                    if column.unique {
-                        write!(file, "unique ")?;
-                    }
-                    write!(
-                        file,
-                        "{} ",
-                        match column.db_type {
-                            DbColumnType::Int => "int",
-                            DbColumnType::String => "string",
-                        }
-                    )?;
-                    write!(file, "{}: ", column.name)?;
-                    match &column.ql_type {
-                        QlColumnType::Int => write!(file, "int")?,
-                        QlColumnType::String => write!(file, "string")?,
-                        QlColumnType::Custom(name) => write!(file, "@{}", name)?,
-                    }
-                    if column.ql_type_is_ref {
-                        write!(file, " ref")?;
-                    }
-                    if column_index + 1 != table.columns.len() {
-                        write!(file, ",")?;
-                    }
-                    write!(file, "\n")?;
-                }
-                write!(file, ");\n\n")?;
-            }
-            Entry::Union { name, members } => {
-                write!(file, "@{} = ", name)?;
-                let mut first = true;
-                for member in members {
-                    if first {
-                        first = false;
-                    } else {
-                        write!(file, " | ")?;
-                    }
-                    write!(file, "@{}", member)?;
-                }
-                write!(file, "\n\n")?;
-            }
+            Entry::Table(table) => write!(file, "{}\n\n", table)?,
+            Entry::Union(union) => write!(file, "{}\n\n", union)?,
         }
     }
 
diff --git a/generator/src/main.rs b/generator/src/main.rs
index 06fc7f250c6..c4a39f50192 100644
--- a/generator/src/main.rs
+++ b/generator/src/main.rs
@@ -57,10 +57,10 @@ fn make_field_type(
                 field_type.named,
             )));
         }
-        entries.push(dbscheme::Entry::Union {
+        entries.push(dbscheme::Entry::Union(dbscheme::Union {
             name: field_union_name.clone(),
             members,
-        });
+        }));
         field_union_name
     }
 }
@@ -144,10 +144,10 @@ fn convert_nodes(nodes: &[NodeInfo]) -> Vec<dbscheme::Entry> {
                     subtype.named,
                 )))
             }
-            entries.push(dbscheme::Entry::Union {
+            entries.push(dbscheme::Entry::Union(dbscheme::Union {
                 name: dbscheme::escape_name(&node_type_name(&node.kind, node.named)),
                 members,
-            });
+            }));
         } else {
             // It's a product type, defined by a table.
             let name = node_type_name(&node.kind, node.named);
@@ -207,10 +207,10 @@ fn convert_nodes(nodes: &[NodeInfo]) -> Vec<dbscheme::Entry> {
     }
 
     // Create a union of all database types.
-    entries.push(dbscheme::Entry::Union {
+    entries.push(dbscheme::Entry::Union(dbscheme::Union {
         name: "top".to_string(),
         members: top_members,
-    });
+    }));
 
     entries
 }