From 7f8856eaa83602c794fc710bbcbcc665768ff94e Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Thu, 7 Nov 2024 20:37:16 +0100
Subject: [PATCH] futhark fmt: improve specs and sum types.

---
 src/Futhark/Fmt/Printer.hs     | 16 +++++++++++-----
 tests_fmt/expected/modules.fut | 11 +++++++++++
 tests_fmt/modules.fut          | 13 +++++++++++++
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs
index ecae898ac8..a9d3ffa21b 100644
--- a/src/Futhark/Fmt/Printer.hs
+++ b/src/Futhark/Fmt/Printer.hs
@@ -98,7 +98,10 @@ instance Format UncheckedTypeExp where
   fmt (TEUnique te loc) = addComments loc $ "*" <> fmt te
   fmt (TEApply te tArgE loc) = addComments loc $ fmt te <+> fmt tArgE
   fmt (TEArrow name te0 te1 loc) =
-    addComments loc $ fmtParamType name te0 <+> "->" </> stdIndent (fmt te1)
+    addComments loc $
+      fmtParamType name te0 </> "->" <+> case te1 of
+        TEArrow {} -> fmt te1
+        _ -> align (fmt te1)
   fmt (TESum tes loc) =
     -- Comments can not be inserted correctly here because names do not
     -- have a location.
@@ -170,9 +173,10 @@ instance Format (UncheckedPat t) where
   fmt (Wildcard _t loc) = addComments loc "_"
   fmt (PatAscription pat t loc) = addComments loc $ fmt pat <> ":" <+> fmt t
   fmt (PatLit _e _ loc) = addComments loc $ fmtCopyLoc constantStyle loc
+  fmt (PatConstr n _ [] loc) =
+    addComments loc $ "#" <> fmt n
   fmt (PatConstr n _ pats loc) =
-    addComments loc $
-      "#" <> fmt n </> align (sep line (map fmt pats))
+    addComments loc $ "#" <> fmt n </> align (sep line (map fmt pats))
   fmt (PatAttr attr pat loc) = addComments loc $ fmt attr <+> fmt pat
 
 instance Format (FieldBase NoInfo Name) where
@@ -272,6 +276,8 @@ instance Format UncheckedExp where
     addComments loc $ parens ("." <> idxs')
     where
       idxs' = brackets $ sep ("," <> space) $ map fmt idxs
+  fmt (Constr n [] _ loc) =
+    addComments loc $ "#" <> fmt n
   fmt (Constr n cs _ loc) =
     addComments loc $ "#" <> fmt n <+> align (sep line $ map fmt cs)
   fmt (Attr attr e loc) = addComments loc $ align $ fmt attr </> fmt e
@@ -477,11 +483,11 @@ instance Format UncheckedSpec where
         | null ps = fmtName bindingStyle name
         | otherwise = fmtName bindingStyle name </> align (sep line $ map fmt ps)
   fmt (ValSpec name ps te _ doc loc) =
-    addComments loc $ fmt doc <> "val" <+> sub <+> ":" <+> fmt te
+    addComments loc $ fmt doc <> "val" <+> sub <+> ":" </> stdIndent (fmt te)
     where
       sub
         | null ps = fmtName bindingStyle name
-        | otherwise = fmtName bindingStyle name </> align (sep line $ map fmt ps)
+        | otherwise = fmtName bindingStyle name <+> align (sep line $ map fmt ps)
   fmt (ModSpec name mte doc loc) =
     addComments loc $ fmt doc <> "module" <+> fmtName bindingStyle name <> ":" <+> fmt mte
   fmt (IncludeSpec mte loc) = addComments loc $ "include" <+> fmt mte
diff --git a/tests_fmt/expected/modules.fut b/tests_fmt/expected/modules.fut
index 22d8c18ddb..f10e08e41d 100644
--- a/tests_fmt/expected/modules.fut
+++ b/tests_fmt/expected/modules.fut
@@ -24,3 +24,14 @@ module m
   def x = 123
   def y = 321
 }
+
+module type mt4 = {
+  val f [n] : [n]i32 -> [n]i32
+
+  val g [n] :
+    [n]i32 -> [n]i32
+
+  val g [n] :
+    [n]i32
+    -> [n]i32
+}
diff --git a/tests_fmt/modules.fut b/tests_fmt/modules.fut
index d83aeef536..32d5870dba 100644
--- a/tests_fmt/modules.fut
+++ b/tests_fmt/modules.fut
@@ -18,3 +18,16 @@ module m : mt1 with a = i32
   def x = 123
   def y = 321
 }
+
+module type mt4 = {
+  val f [n] : [n]i32 -> [n]i32
+
+  val g [n] :
+  [n]i32 -> [n]i32
+
+  val g [n] :
+    [n]i32
+ -> [n]i32
+
+
+}