From af43cd4233a1d49e989c69bed89a1dce78ab81c5 Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Thu, 9 Nov 2023 11:50:44 +0100
Subject: [PATCH] futhark literate: do not choke on malformed files.

---
 CHANGELOG.md                    |  3 +++
 src/Futhark/CLI/Literate.hs     |  8 ++++----
 tests_literate/expected/nonl.md | 14 ++++++++++++++
 tests_literate/nonl.fut         |  6 ++++++
 4 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 tests_literate/expected/nonl.md
 create mode 100644 tests_literate/nonl.fut

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 73759d9ada..ff298d7fdb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ### Fixed
 
+* `futhark literate` no longer fails if the final line is a directive
+  without a trailing newline.
+
 ## [0.25.7]
 
 ### Added
diff --git a/src/Futhark/CLI/Literate.hs b/src/Futhark/CLI/Literate.hs
index 1d7d2250d6..643cd35ea4 100644
--- a/src/Futhark/CLI/Literate.hs
+++ b/src/Futhark/CLI/Literate.hs
@@ -329,7 +329,7 @@ atStartOfLine = do
   when (col /= pos1) empty
 
 afterExp :: Parser ()
-afterExp = choice [atStartOfLine, void eol]
+afterExp = choice [atStartOfLine, choice [void eol, eof]]
 
 withParsedSource :: Parser a -> (a -> T.Text -> b) -> Parser b
 withParsedSource p f = do
@@ -369,12 +369,12 @@ parseBlock =
             $> DirectiveImg
             <*> parseExp postlexeme
             <*> parseImgParams
-            <* eol,
+            <* choice [void eol, eof],
           directiveName "plot2d"
             $> DirectivePlot
             <*> parseExp postlexeme
             <*> parsePlotParams
-            <* eol,
+            <* choice [void eol, eof],
           directiveName "gnuplot"
             $> DirectiveGnuplot
             <*> parseExp postlexeme
@@ -388,7 +388,7 @@ parseBlock =
             $> DirectiveAudio
             <*> parseExp postlexeme
             <*> parseAudioParams
-            <* eol
+            <* choice [void eol, eof]
         ]
     directiveName s = try $ token (":" <> s)
 
diff --git a/tests_literate/expected/nonl.md b/tests_literate/expected/nonl.md
new file mode 100644
index 0000000000..744e65acd9
--- /dev/null
+++ b/tests_literate/expected/nonl.md
@@ -0,0 +1,14 @@
+Note: it is intentional that this file does not include a final
+newline character. Do not add it!
+
+```futhark
+def main x = x + 1
+```
+
+```
+> main 1
+```
+
+```
+2i32
+```
diff --git a/tests_literate/nonl.fut b/tests_literate/nonl.fut
new file mode 100644
index 0000000000..8a4cd6d53f
--- /dev/null
+++ b/tests_literate/nonl.fut
@@ -0,0 +1,6 @@
+-- Note: it is intentional that this file does not include a final
+-- newline character. Do not add it!
+
+def main x = x + 1
+
+-- > main 1
\ No newline at end of file