From ec333e9211ad8971001e2f83c78c8431f3990fa5 Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Thu, 9 Jan 2025 18:35:17 +0100
Subject: [PATCH] futhark fmt: preserve doc comments for local defs.

---
 CHANGELOG.md                         | 3 +++
 src/Futhark/Fmt/Printer.hs           | 2 +-
 src/Language/Futhark/Parser/Monad.hs | 3 ++-
 tests_fmt/expected/local.fut         | 7 +++++++
 tests_fmt/local.fut                  | 7 +++++++
 5 files changed, 20 insertions(+), 2 deletions(-)
 create mode 100644 tests_fmt/expected/local.fut
 create mode 100644 tests_fmt/local.fut

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 72b91381ac..d98239e9bf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * GPU backends did not handle array constants correctly in some cases.
 
+* `futhark fmt`: do not throw away doc comments for `local`
+  definitions.
+
 ## [0.25.25]
 
 ### Added
diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs
index e6b6ff33f2..027d37d7cd 100644
--- a/src/Futhark/Fmt/Printer.hs
+++ b/src/Futhark/Fmt/Printer.hs
@@ -614,7 +614,7 @@ instance Format UncheckedDec where
   fmt (ModTypeDec tb) = fmt tb
   fmt (ModDec tb) = fmt tb
   fmt (OpenDec tb loc) = addComments loc $ "open" <+> fmt tb
-  fmt (LocalDec tb loc) = addComments loc $ "local" <+> fmt tb
+  fmt (LocalDec tb loc) = addComments loc $ "local" </> fmt tb
   fmt (ImportDec path _tb loc) =
     addComments loc $ "import" <+> "\"" <> fmtPretty path <> "\""
 
diff --git a/src/Language/Futhark/Parser/Monad.hs b/src/Language/Futhark/Parser/Monad.hs
index 93e57e464c..e33a2c9b65 100644
--- a/src/Language/Futhark/Parser/Monad.hs
+++ b/src/Language/Futhark/Parser/Monad.hs
@@ -59,6 +59,7 @@ addDoc doc (ValDec val) = ValDec (val {valBindDoc = Just doc})
 addDoc doc (TypeDec tp) = TypeDec (tp {typeDoc = Just doc})
 addDoc doc (ModTypeDec sig) = ModTypeDec (sig {modTypeDoc = Just doc})
 addDoc doc (ModDec mod) = ModDec (mod {modDoc = Just doc})
+addDoc doc (LocalDec dec loc) = LocalDec (addDoc doc dec) loc
 addDoc _ dec = dec
 
 addDocSpec :: DocComment -> SpecBase NoInfo Name -> SpecBase NoInfo Name
@@ -66,7 +67,7 @@ addDocSpec doc (TypeAbbrSpec tpsig) = TypeAbbrSpec (tpsig {typeDoc = Just doc})
 addDocSpec doc (ValSpec name ps t NoInfo _ loc) = ValSpec name ps t NoInfo (Just doc) loc
 addDocSpec doc (TypeSpec l name ps _ loc) = TypeSpec l name ps (Just doc) loc
 addDocSpec doc (ModSpec name se _ loc) = ModSpec name se (Just doc) loc
-addDocSpec _ spec = spec
+addDocSpec _ spec@IncludeSpec {} = spec
 
 addAttr :: AttrInfo Name -> UncheckedDec -> UncheckedDec
 addAttr attr (ValDec val) =
diff --git a/tests_fmt/expected/local.fut b/tests_fmt/expected/local.fut
new file mode 100644
index 0000000000..1957c22c56
--- /dev/null
+++ b/tests_fmt/expected/local.fut
@@ -0,0 +1,7 @@
+-- | top level.
+
+local
+-- | I am a doc comment.
+module type mt = {
+
+}
diff --git a/tests_fmt/local.fut b/tests_fmt/local.fut
new file mode 100644
index 0000000000..1957c22c56
--- /dev/null
+++ b/tests_fmt/local.fut
@@ -0,0 +1,7 @@
+-- | top level.
+
+local
+-- | I am a doc comment.
+module type mt = {
+
+}