Allow lambdas to be passed as last argument.

CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * `futhark fmt`, by William Due and Therese Lyngby.
 
+* Lambdas can now be passed as the last argument to a function application.
+
 ### Removed
 
 ### Changed

src/Futhark/Fmt/Monad.hs

@@ -60,7 +60,7 @@ import Control.Monad.State
   )
 import Data.ByteString qualified as BS
 import Data.List.NonEmpty qualified as NE
-import Data.Loc (Loc (..), Located (..), posCoff, posLine)
+import Data.Loc (Loc (..), Located (..), locStart, posCoff, posLine)
 import Data.Maybe (fromMaybe)
 import Data.String
 import Data.Text qualified as T
@@ -150,16 +150,11 @@ fmtByLayout a s m =
             _any -> m
         )
 
--- | This function determines the Layout of @a@ and if it is singleline then it
--- updates the monads enviroment to format singleline style otherwise format using
--- multiline style. It determines this by checking if the location of @a@ spans
--- over two or more lines.
+-- | This function determines the Layout of @a@ and updates the monads
+-- environment to format in the appropriate style. It determines this
+-- by checking if the location of @a@ spans over two or more lines.
 localLayout :: (Located a) => a -> FmtM b -> FmtM b
-localLayout a m = do
-  lo <- ask
-  case lo of
-    MultiLine -> local (const $ fromMaybe lo $ lineLayout a) m
-    SingleLine -> m
+localLayout a = local (\lo -> fromMaybe lo $ lineLayout a)
 
 -- | This function determines the Layout of @[a]@ and if it is singleline then it
 -- updates the monads enviroment to format singleline style otherwise format using
@@ -416,11 +411,19 @@ sepLineComments :: (a -> Loc) -> (a -> Fmt) -> Fmt -> [a] -> Fmt
 sepLineComments floc fmt s =
   sepComments floc fmt (s <> space <|> hardline <> s)
 
--- | This is used for function arguments. It seperates multiline arguments by
--- lines and singleline arguments by spaces.
+-- | This is used for function arguments. It seperates multiline
+-- arguments by lines and singleline arguments by spaces. We specially
+-- handle the case where all the arguments are on a single line except
+-- for the last one, which may continue to the next line.
 sepArgs :: (Located a) => (a -> Fmt) -> NE.NonEmpty a -> Fmt
-sepArgs fmt ls = localLayout ls' $ align $ sep line $ map fmt ls'
+sepArgs fmt ls =
+  localLayout locs $ align' $ sep line $ map fmtArg ls'
   where
+    locs = map (locStart . locOf) ls'
+    align' = case lineLayout locs of
+      Just SingleLine -> id
+      _ -> align
+    fmtArg x = localLayout x $ fmt x
     ls' = NE.toList ls
 
 -- | Nest but with the standard value of two spaces.

src/Futhark/Fmt/Printer.hs

@@ -396,7 +396,7 @@ instance Format (AppExpBase NoInfo Name) where
           AppExp If {} _ -> space <> fmt f
           _ -> space <> align (fmt f)
   fmt (Apply f args loc) =
-    addComments loc $ fmt f <+> align fmt_args
+    addComments loc $ fmt f <+> fmt_args
     where
       fmt_args = sepArgs fmt $ fmap snd args
 

src/Language/Futhark/Parser/Parser.y

@@ -580,17 +580,21 @@ Exp2 :: { UncheckedExp }
      | Exp2 with FieldAccesses_ '=' Exp2
        { RecordUpdate $1 (map unLoc $3) $5 NoInfo (srcspan $1 $>) }
 
-     | '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp %prec letprec
-       { Lambda (fst $2 : snd $2) $5 $3 NoInfo (srcspan $1 $>) }
-
      | ApplyList {% applyExp $1 }
 
 ApplyList :: { NE.NonEmpty UncheckedExp }
           : Atom ApplyList %prec juxtprec
             { NE.cons $1 $2 }
-          | Atom %prec juxtprec
+          | LastArg
             { NE.singleton $1 }
 
+LastArg :: { UncheckedExp }
+        : '\\' FunParams1 maybeAscription(TypeExpTerm) '->' Exp %prec letprec
+          { Lambda (fst $2 : snd $2) $5 $3 NoInfo (srcspan $1 $>) }
+        | Atom %prec juxtprec
+          { $1 }
+
+
 Atom :: { UncheckedExp }
 Atom : PrimLit        { Literal (fst $1) (srclocOf (snd $1)) }
      | Constr         { Constr (fst $1) [] NoInfo (srclocOf (snd $1)) }

tests_fmt/expected/lastarg.fut

@@ -0,0 +1,5 @@
+def f = tabulate_2d 2 3 \i j -> i + j
+
+def g =
+  tabulate_2d 2 3 \i j ->
+    i + j

tests_fmt/lastarg.fut

@@ -0,0 +1,4 @@
+def f = tabulate_2d 2 3 \i j -> i + j
+
+def g = tabulate_2d 2 3 \i j ->
+                          i + j