Fix diku-dk#2038.

CHANGELOG.md

@@ -22,6 +22,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * Very large `iota`s now work.
 
+* Lambda lifting in `while` conditions (#2038).
+
 ## [0.25.6]
 
 ### Added

src/Futhark/Internalise/LiftLambdas.hs

@@ -169,9 +169,10 @@ transformExp (AppExp (Match e cases loc) appres) = do
         <*> pure case_loc
 transformExp (AppExp (Loop sizes pat args form body loc) appres) = do
   args' <- transformExp args
-  form' <- astMap transformSubExps form
-  body' <- bindingParams sizes [pat] $ bindingForm form' $ transformExp body
-  pure $ AppExp (Loop sizes pat args' form' body' loc) appres
+  bindingParams sizes [pat] $ do
+    form' <- astMap transformSubExps form
+    body' <- bindingForm form' $ transformExp body
+    pure $ AppExp (Loop sizes pat args' form' body' loc) appres
 transformExp e@(Var v _ _) =
   -- Note that function-typed variables can only occur in expressions,
   -- not in other places where VNames/QualNames can occur.

src/Language/Futhark/Pretty.hs

@@ -232,8 +232,8 @@ prettyAppExp _ (Loop sizeparams pat initexp form loopbody _) =
       ( hsep (map (brackets . prettyName) sizeparams ++ [pretty pat])
           <+> equals
           <+> pretty initexp
-          <+> pretty form
-          <+> "do"
+          </> pretty form
+          </> "do"
       )
     </> indent 2 (pretty loopbody)
 prettyAppExp _ (Index e idxs _) =

tests/issue2038.fut

@@ -0,0 +1,11 @@
+entry all_work_indices (n:i64) (m:i64) =
+  let block i prog prog' =
+    (i, prog, i64.max 0 (i64.min (n-i) (if i == 0 then n else prog')))
+  let size (_, a, b) = b-a >> 1
+  let (iter, _) = loop (iter, progress) = (0, replicate m 0)
+                  while not (all id (tabulate m (\i -> progress[i] >= n-i-1))) do
+                  let blockrow = map3 block
+                                      (iota m) (progress) (rotate (-1) progress)
+                  let sizes = map size blockrow
+                  in (iter + 1i32, map2 (+) progress sizes)
+  in iter