From 7092673c45bf02c6c618c63db5e6a5dad5177211 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Sat, 5 Nov 2022 17:53:11 -0400
Subject: [PATCH] feat: if expression formatting

---
 crates/lang/src/format.rs           | 43 ++++++++++++++++++++++++++++-
 examples/sample/src/scripts/swap.ak |  9 ++++--
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/crates/lang/src/format.rs b/crates/lang/src/format.rs
index e0189fc1..765a9b1c 100644
--- a/crates/lang/src/format.rs
+++ b/crates/lang/src/format.rs
@@ -601,7 +601,48 @@ impl<'comments> Formatter<'comments> {
 
         let document = match expr {
             UntypedExpr::ByteArray { .. } => todo!(),
-            UntypedExpr::If { .. } => todo!(),
+            UntypedExpr::If {
+                branches,
+                final_else,
+                ..
+            } => {
+                let first = branches.first();
+
+                break_("if", "if ")
+                    .append(self.wrap_expr(&first.condition))
+                    .nest(INDENT)
+                    .append(break_("", " "))
+                    .append("{")
+                    .group()
+                    .append(line())
+                    .nest(INDENT)
+                    .append(self.expr(&first.body))
+                    .append(line())
+                    .append("} ")
+                    .append(join(
+                        branches[1..].iter().map(|branch| {
+                            break_("else if", "else if ")
+                                .append(self.wrap_expr(&branch.condition))
+                                .nest(INDENT)
+                                .append(break_("", " "))
+                                .append("{")
+                                .group()
+                                .append(line())
+                                .nest(INDENT)
+                                .append(self.expr(&branch.body))
+                                .append(line())
+                                .append("} ")
+                        }),
+                        nil(),
+                    ))
+                    .append("else {")
+                    .group()
+                    .append(line().nest(INDENT))
+                    .append(self.expr(final_else))
+                    .append(line())
+                    .append("}")
+                    .force_break()
+            }
             UntypedExpr::Todo { label: None, .. } => "todo".to_doc(),
 
             UntypedExpr::Todo { label: Some(l), .. } => docvec!["todo(\"", l, "\")"],
diff --git a/examples/sample/src/scripts/swap.ak b/examples/sample/src/scripts/swap.ak
index a36d525e..d0f1c4e0 100644
--- a/examples/sample/src/scripts/swap.ak
+++ b/examples/sample/src/scripts/swap.ak
@@ -11,8 +11,11 @@ pub type Redeemer {
 }
 
 pub fn spend(datum: Datum, rdmr: Redeemer, ctx: spend.ScriptContext) -> Bool {
-  when rdmr is {
-    Buy -> True
-    Sell -> datum.something == "Aiken"
+  if datum.something == "Aiken" {
+    True
+  } else if datum.something == "Wow" {
+    True
+  } else {
+    False
   }
 }