From 6aae18484897cb6bb6b882f5f0e1c962b18b6916 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Sat, 18 Jun 2022 22:53:02 -0400
Subject: [PATCH] feat: new fmt command and pretty printing works

---
 crates/cli/src/args.rs                        |  3 +++
 crates/cli/src/main.rs                        |  9 ++++++++
 crates/uplc/src/debruijn.rs                   | 14 +++++++----
 crates/uplc/src/pretty.rs                     | 23 ++++++++++++++-----
 crates/uplc/src/test.rs                       |  4 ++++
 .../uplc/test_data/basic/integer/integer.uplc |  5 ++--
 6 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/crates/cli/src/args.rs b/crates/cli/src/args.rs
index bf149545..543a397b 100644
--- a/crates/cli/src/args.rs
+++ b/crates/cli/src/args.rs
@@ -31,6 +31,9 @@ pub enum UplcCommand {
         #[clap(short, long)]
         out: Option<String>,
     },
+    Fmt {
+        input: PathBuf,
+    },
 }
 
 impl Default for Args {
diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
index c383e00f..98b81c7e 100644
--- a/crates/cli/src/main.rs
+++ b/crates/cli/src/main.rs
@@ -47,6 +47,15 @@ fn main() -> anyhow::Result<()> {
                     fs::write(&out_name, &bytes)?;
                 }
             }
+            UplcCommand::Fmt { input } => {
+                let code = std::fs::read_to_string(&input)?;
+
+                let program = parser::program(&code)?;
+
+                let pretty = program.to_pretty();
+
+                fs::write(&input, pretty)?;
+            }
             UplcCommand::Unflat { input, print, out } => {
                 let bytes = std::fs::read(&input)?;
 
diff --git a/crates/uplc/src/debruijn.rs b/crates/uplc/src/debruijn.rs
index 3c4fa7f3..45416f8a 100644
--- a/crates/uplc/src/debruijn.rs
+++ b/crates/uplc/src/debruijn.rs
@@ -162,10 +162,14 @@ impl Converter {
 
     pub fn debruijn_to_name(&mut self, term: Term<DeBruijn>) -> Result<Term<Name>, Error> {
         let converted_term = match term {
-            Term::Var(index) => Term::Var(Name {
-                text: String::from("i"),
-                unique: self.get_unique(index)?,
-            }),
+            Term::Var(index) => {
+                let unique = self.get_unique(index)?;
+
+                Term::Var(Name {
+                    text: format!("i_{}", unique),
+                    unique,
+                })
+            }
             Term::Delay(term) => Term::Delay(Box::new(self.debruijn_to_name(*term)?)),
             Term::Lambda {
                 parameter_name,
@@ -176,7 +180,7 @@ impl Converter {
                 let unique = self.get_unique(parameter_name)?;
 
                 let name = Name {
-                    text: String::from("i"),
+                    text: format!("i_{}", unique),
                     unique,
                 };
 
diff --git a/crates/uplc/src/pretty.rs b/crates/uplc/src/pretty.rs
index 04259f68..90eb08d5 100644
--- a/crates/uplc/src/pretty.rs
+++ b/crates/uplc/src/pretty.rs
@@ -8,7 +8,21 @@ impl Program<Name> {
 
         self.to_doc().render(80, &mut w).unwrap();
 
-        String::from_utf8(w).unwrap()
+        String::from_utf8(w)
+            .unwrap()
+            .lines()
+            // This is a hack to deal with blank newlines
+            // that end up with a bunch of useless whitespace
+            // because of the nesting
+            .map(|l| {
+                if l.chars().all(|c| c.is_whitespace()) {
+                    "".to_string()
+                } else {
+                    l.to_string()
+                }
+            })
+            .collect::<Vec<_>>()
+            .join("\n")
     }
 
     fn to_doc(&self) -> RcDoc<()> {
@@ -29,7 +43,7 @@ impl Program<Name> {
 impl Term<Name> {
     fn to_doc(&self) -> RcDoc<()> {
         match self {
-            Term::Var(name) => RcDoc::text(format!("{}_{}", name.text, name.unique)),
+            Term::Var(name) => RcDoc::text(&name.text),
             Term::Delay(term) => RcDoc::text("(")
                 .append(
                     RcDoc::text("delay")
@@ -46,10 +60,7 @@ impl Term<Name> {
                 .append(
                     RcDoc::text("lam")
                         .append(RcDoc::line())
-                        .append(RcDoc::text(format!(
-                            "{}_{}",
-                            parameter_name.text, parameter_name.unique
-                        )))
+                        .append(RcDoc::text(&parameter_name.text))
                         .append(RcDoc::line())
                         .append(body.to_doc())
                         .nest(2),
diff --git a/crates/uplc/src/test.rs b/crates/uplc/src/test.rs
index 294c11bb..ae5f3f1f 100644
--- a/crates/uplc/src/test.rs
+++ b/crates/uplc/src/test.rs
@@ -24,6 +24,10 @@ fn integer() {
     let name_program: Program<Name> = decoded_program.try_into().unwrap();
 
     assert_eq!(parsed_program, name_program);
+
+    let pretty = name_program.to_pretty();
+
+    assert_eq!(pretty, code);
 }
 
 #[test]
diff --git a/crates/uplc/test_data/basic/integer/integer.uplc b/crates/uplc/test_data/basic/integer/integer.uplc
index 7e4c5e38..3ba0e4a6 100644
--- a/crates/uplc/test_data/basic/integer/integer.uplc
+++ b/crates/uplc/test_data/basic/integer/integer.uplc
@@ -1,3 +1,4 @@
-(program 11.22.33
-    (con integer 11)
+(program
+  11.22.33
+  (con integer 11)
 )
\ No newline at end of file