Fix #2197.

CHANGELOG.md

@@ -25,6 +25,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
   ranges with start==end and negative step size, e.g. `1..0...1`
   produces `[1]` rather than an invalid range error.
 
+* Inconsistent handling of types in lambda lifting (#2197).
+
 ## [0.25.24]
 
 ### Added

src/Futhark/Internalise/LiftLambdas.hs

@@ -6,6 +6,7 @@ module Futhark.Internalise.LiftLambdas (transformProg) where
 import Control.Monad.Reader
 import Control.Monad.State
 import Data.Bifunctor
+import Data.Bitraversable
 import Data.Foldable
 import Data.List (partition)
 import Data.Map.Strict qualified as M
@@ -144,6 +145,12 @@ liftFunction fname tparams params (RetType dims ret) funbody = do
 transformSubExps :: ASTMapper LiftM
 transformSubExps = identityMapper {mapOnExp = transformExp}
 
+transformType :: TypeBase Exp u -> LiftM (TypeBase Exp u)
+transformType = bitraverse transformExp pure
+
+transformPat :: PatBase Info VName (TypeBase Exp u) -> LiftM (PatBase Info VName (TypeBase Exp u))
+transformPat = traverse transformType
+
 transformExp :: Exp -> LiftM Exp
 transformExp (AppExp (LetFun fname (tparams, params, _, Info ret, funbody) body _) _) = do
   funbody' <- bindingParams (map typeParamName tparams) params $ transformExp funbody
@@ -156,8 +163,9 @@ transformExp e@(Lambda params body _ (Info ret) _) = do
   liftFunction fname [] params ret body' <*> pure (typeOf e)
 transformExp (AppExp (LetPat sizes pat e body loc) appres) = do
   e' <- transformExp e
-  body' <- bindingLetPat (map sizeName sizes) pat $ transformExp body
-  pure $ AppExp (LetPat sizes pat e' body' loc) appres
+  pat' <- transformPat pat
+  body' <- bindingLetPat (map sizeName sizes) pat' $ transformExp body
+  pure $ AppExp (LetPat sizes pat' e' body' loc) appres
 transformExp (AppExp (Match e cases loc) appres) = do
   e' <- transformExp e
   cases' <- mapM transformCase cases
@@ -173,10 +181,11 @@ transformExp (AppExp (Loop sizes pat args form body loc) appres) = do
     form' <- astMap transformSubExps form
     body' <- bindingForm form' $ transformExp body
     pure $ AppExp (Loop sizes pat (LoopInitExplicit args') form' body' loc) appres
-transformExp e@(Var v (Info t) _) =
+transformExp (Var v (Info t) loc) = do
+  t' <- transformType t
   -- Note that function-typed variables can only occur in expressions,
   -- not in other places where VNames/QualNames can occur.
-  asks $ maybe e ($ t) . M.lookup (qualLeaf v) . envReplace
+  asks $ maybe (Var v (Info t') loc) ($ t') . M.lookup (qualLeaf v) . envReplace
 transformExp e = astMap transformSubExps e
 
 transformValBind :: ValBind -> LiftM ()

tests/issue2197.fut

@@ -0,0 +1,8 @@
+def index_of_first p xs =
+  loop i = 0 while i < length xs && !p xs[i] do i + 1
+
+def span p xs = let i = index_of_first p xs in (take i xs, drop i xs)
+
+entry part1 [l] (ls: [][l]i32) =
+  let blank (l: [l]i32) = null l
+  in span blank ls |> \(x, y) -> (id x, tail y)