From dd240d3aa7324c5ca579302629e0c1b278182054 Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Thu, 7 Nov 2024 16:52:49 +0100
Subject: [PATCH] futhark fmt: better sum types.

---
 src/Futhark/Fmt/Printer.hs     | 13 ++++++++++---
 tests_fmt/expected/sumtype.fut |  6 ++++++
 tests_fmt/sumtype.fut          |  6 ++++++
 3 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 tests_fmt/expected/sumtype.fut
 create mode 100644 tests_fmt/sumtype.fut

diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs
index 2fd594b20f..ecae898ac8 100644
--- a/src/Futhark/Fmt/Printer.hs
+++ b/src/Futhark/Fmt/Printer.hs
@@ -55,6 +55,8 @@ fmtParamType (Just n) te =
 fmtParamType Nothing te = fmt te
 
 fmtSumTypeConstr :: (Name, [UncheckedTypeExp]) -> Fmt
+fmtSumTypeConstr (name, []) =
+  "#" <> fmtName mempty name
 fmtSumTypeConstr (name, fs) =
   "#" <> fmtName mempty name <+> sep space (map fmt fs)
 
@@ -98,9 +100,14 @@ instance Format UncheckedTypeExp where
   fmt (TEArrow name te0 te1 loc) =
     addComments loc $ fmtParamType name te0 <+> "->" </> stdIndent (fmt te1)
   fmt (TESum tes loc) =
-    addComments loc $
-      sep (line <> "|" <> space) $
-        map fmtSumTypeConstr tes -- Comments can not be inserted correctly here because names do not have a location.
+    -- Comments can not be inserted correctly here because names do not
+    -- have a location.
+    addComments loc $ fmtByLayout loc singleLine multiLine
+    where
+      singleLine = sep " | " $ map fmtSumTypeConstr tes
+      multiLine = sep line $ zipWith prefix [0 :: Int ..] tes
+      prefix 0 te = "  " <> fmtSumTypeConstr te
+      prefix _ te = "| " <> fmtSumTypeConstr te
   fmt (TEDim dims te loc) =
     addComments loc $ "?" <> dims' <> "." <> fmt te
     where
diff --git a/tests_fmt/expected/sumtype.fut b/tests_fmt/expected/sumtype.fut
new file mode 100644
index 0000000000..683ff79ca0
--- /dev/null
+++ b/tests_fmt/expected/sumtype.fut
@@ -0,0 +1,6 @@
+type t1 = #foo | #bar | #baz
+
+type t2 =
+    #foo
+  | #bar
+  | #baz
diff --git a/tests_fmt/sumtype.fut b/tests_fmt/sumtype.fut
new file mode 100644
index 0000000000..ce59e31c63
--- /dev/null
+++ b/tests_fmt/sumtype.fut
@@ -0,0 +1,6 @@
+type t1 = #foo  | #bar  | #baz
+
+type t2 =
+  #foo
+  | #bar
+  | #baz