-- | Lambda-lifting of typed, monomorphic Futhark programs without
-- modules.  After this pass, the program will no longer contain any
-- 'LetFun's or 'Lambda's.
module Futhark.Internalise.LiftLambdas (transformProg) where

import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Foldable
import Data.List (partition)
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Pretty ()
import Futhark.MonadFreshNames
import Futhark.Util (nubOrd)
import Language.Futhark
import Language.Futhark.Traversals

data Env = Env
  { Env -> Map VName Exp
envReplace :: M.Map VName Exp,
    Env -> Map VName StructType
envVtable :: M.Map VName StructType
  }

initialEnv :: Env
initialEnv :: Env
initialEnv = Map VName Exp -> Map VName StructType -> Env
Env forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty

data LiftState = State
  { LiftState -> VNameSource
stateNameSource :: VNameSource,
    LiftState -> [ValBind]
stateValBinds :: [ValBind],
    LiftState -> Set VName
stateGlobal :: S.Set VName
  }

initialState :: VNameSource -> LiftState
initialState :: VNameSource -> LiftState
initialState VNameSource
src = VNameSource -> [ValBind] -> Set VName -> LiftState
State VNameSource
src forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [k]
M.keys Map VName Intrinsic
intrinsics

newtype LiftM a = LiftM (ReaderT Env (State LiftState) a)
  deriving (forall a b. a -> LiftM b -> LiftM a
forall a b. (a -> b) -> LiftM a -> LiftM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> LiftM b -> LiftM a
$c<$ :: forall a b. a -> LiftM b -> LiftM a
fmap :: forall a b. (a -> b) -> LiftM a -> LiftM b
$cfmap :: forall a b. (a -> b) -> LiftM a -> LiftM b
Functor, Functor LiftM
forall a. a -> LiftM a
forall a b. LiftM a -> LiftM b -> LiftM a
forall a b. LiftM a -> LiftM b -> LiftM b
forall a b. LiftM (a -> b) -> LiftM a -> LiftM b
forall a b c. (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. LiftM a -> LiftM b -> LiftM a
$c<* :: forall a b. LiftM a -> LiftM b -> LiftM a
*> :: forall a b. LiftM a -> LiftM b -> LiftM b
$c*> :: forall a b. LiftM a -> LiftM b -> LiftM b
liftA2 :: forall a b c. (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
$cliftA2 :: forall a b c. (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
<*> :: forall a b. LiftM (a -> b) -> LiftM a -> LiftM b
$c<*> :: forall a b. LiftM (a -> b) -> LiftM a -> LiftM b
pure :: forall a. a -> LiftM a
$cpure :: forall a. a -> LiftM a
Applicative, Applicative LiftM
forall a. a -> LiftM a
forall a b. LiftM a -> LiftM b -> LiftM b
forall a b. LiftM a -> (a -> LiftM b) -> LiftM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> LiftM a
$creturn :: forall a. a -> LiftM a
>> :: forall a b. LiftM a -> LiftM b -> LiftM b
$c>> :: forall a b. LiftM a -> LiftM b -> LiftM b
>>= :: forall a b. LiftM a -> (a -> LiftM b) -> LiftM b
$c>>= :: forall a b. LiftM a -> (a -> LiftM b) -> LiftM b
Monad, MonadReader Env, MonadState LiftState)

instance MonadFreshNames LiftM where
  putNameSource :: VNameSource -> LiftM ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \LiftState
s -> LiftState
s {stateNameSource :: VNameSource
stateNameSource = VNameSource
src}
  getNameSource :: LiftM VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets LiftState -> VNameSource
stateNameSource

runLiftM :: VNameSource -> LiftM () -> ([ValBind], VNameSource)
runLiftM :: VNameSource -> LiftM () -> ([ValBind], VNameSource)
runLiftM VNameSource
src (LiftM ReaderT Env (State LiftState) ()
m) =
  let s :: LiftState
s = forall s a. State s a -> s -> s
execState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Env (State LiftState) ()
m Env
initialEnv) (VNameSource -> LiftState
initialState VNameSource
src)
   in (forall a. [a] -> [a]
reverse (LiftState -> [ValBind]
stateValBinds LiftState
s), LiftState -> VNameSource
stateNameSource LiftState
s)

addValBind :: ValBind -> LiftM ()
addValBind :: ValBind -> LiftM ()
addValBind ValBind
vb = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \LiftState
s ->
  LiftState
s
    { stateValBinds :: [ValBind]
stateValBinds = ValBind
vb forall a. a -> [a] -> [a]
: LiftState -> [ValBind]
stateValBinds LiftState
s,
      stateGlobal :: Set VName
stateGlobal = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Ord a => a -> Set a -> Set a
S.insert) (LiftState -> Set VName
stateGlobal LiftState
s) (ValBind -> [VName]
valBindBound ValBind
vb)
    }

replacing :: VName -> Exp -> LiftM a -> LiftM a
replacing :: forall a. VName -> Exp -> LiftM a -> LiftM a
replacing VName
v Exp
e = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env {envReplace :: Map VName Exp
envReplace = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v Exp
e forall a b. (a -> b) -> a -> b
$ Env -> Map VName Exp
envReplace Env
env}

bindingParams :: [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams :: forall a. [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams [VName]
sizes [Pat ParamType]
params = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env
    { envVtable :: Map VName StructType
envVtable =
        forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall t. Pat t -> [(VName, t)]
patternMap [Pat ParamType]
params) forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map (,forall {dim} {u}. TypeBase dim u
i64) [VName]
sizes)
          forall a. Semigroup a => a -> a -> a
<> Env -> Map VName StructType
envVtable Env
env
    }
  where
    i64 :: TypeBase dim u
i64 = forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

bindingLetPat :: [VName] -> Pat StructType -> LiftM a -> LiftM a
bindingLetPat :: forall a. [VName] -> Pat StructType -> LiftM a -> LiftM a
bindingLetPat [VName]
sizes Pat StructType
pat = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env
    { envVtable :: Map VName StructType
envVtable =
        forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct) (forall t. Pat t -> [(VName, t)]
patternMap Pat StructType
pat) forall a. Semigroup a => a -> a -> a
<> forall a b. (a -> b) -> [a] -> [b]
map (,forall {dim} {u}. TypeBase dim u
i64) [VName]
sizes)
          forall a. Semigroup a => a -> a -> a
<> Env -> Map VName StructType
envVtable Env
env
    }
  where
    i64 :: TypeBase dim u
i64 = forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar forall a b. (a -> b) -> a -> b
$ forall dim u. PrimType -> ScalarTypeBase dim u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

bindingForm :: LoopFormBase Info VName -> LiftM a -> LiftM a
bindingForm :: forall a. LoopFormBase Info VName -> LiftM a -> LiftM a
bindingForm (For IdentBase Info VName StructType
i Exp
_) = forall a. [VName] -> Pat StructType -> LiftM a -> LiftM a
bindingLetPat [] (forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id (forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
i) (forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> f t
identType IdentBase Info VName StructType
i) forall a. Monoid a => a
mempty)
bindingForm (ForIn Pat StructType
p Exp
_) = forall a. [VName] -> Pat StructType -> LiftM a -> LiftM a
bindingLetPat [] Pat StructType
p
bindingForm While {} = forall a. a -> a
id

toRet :: TypeBase Size u -> TypeBase Size Uniqueness
toRet :: forall u. TypeBase Exp u -> TypeBase Exp Uniqueness
toRet = forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const Uniqueness
Nonunique)

liftFunction :: VName -> [TypeParam] -> [Pat ParamType] -> ResRetType -> Exp -> LiftM Exp
liftFunction :: VName
-> [TypeParam] -> [Pat ParamType] -> ResRetType -> Exp -> LiftM Exp
liftFunction VName
fname [TypeParam]
tparams [Pat ParamType]
params (RetType [VName]
dims TypeBase Exp Uniqueness
ret) Exp
funbody = do
  -- Find free variables
  Map VName StructType
vtable <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Map VName StructType
envVtable
  let isFree :: VName -> Maybe (VName, StructType)
isFree VName
v = (VName
v,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName StructType
vtable
      withTypes :: FV -> [(VName, StructType)]
withTypes = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe (VName, StructType)
isFree forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
S.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars

  let free :: [(VName, StructType)]
free =
        let immediate_free :: [(VName, StructType)]
immediate_free = FV -> [(VName, StructType)]
withTypes forall a b. (a -> b) -> a -> b
$ Exp -> FV
freeInExp Exp
funbody
            sizes_in_free :: FV
sizes_in_free = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall u. TypeBase Exp u -> FV
freeInType forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(VName, StructType)]
immediate_free
            sizes :: [(VName, StructType)]
sizes =
              FV -> [(VName, StructType)]
withTypes forall a b. (a -> b) -> a -> b
$
                FV
sizes_in_free
                  forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall u. Pat (TypeBase Exp u) -> FV
freeInPat [Pat ParamType]
params
                  forall a. Semigroup a => a -> a -> a
<> forall u. TypeBase Exp u -> FV
freeInType TypeBase Exp Uniqueness
ret
         in forall a. Ord a => [a] -> [a]
nubOrd forall a b. (a -> b) -> a -> b
$ [(VName, StructType)]
immediate_free forall a. Semigroup a => a -> a -> a
<> [(VName, StructType)]
sizes

      -- Those parameters that correspond to sizes must come first.
      sizes_in_types :: FV
sizes_in_types =
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall u. TypeBase Exp u -> FV
freeInType (forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
ret forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(VName, StructType)]
free forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map forall u. Pat (TypeBase Exp u) -> StructType
patternStructType [Pat ParamType]
params)
      isSize :: (VName, b) -> Bool
isSize (VName
v, b
_) = VName
v forall a. Ord a => a -> Set a -> Bool
`S.member` FV -> Set VName
fvVars FV
sizes_in_types
      ([(VName, StructType)]
free_dims, [(VName, StructType)]
free_nondims) = forall a. (a -> Bool) -> [a] -> ([a], [a])
partition forall {b}. (VName, b) -> Bool
isSize [(VName, StructType)]
free

      free_ts :: [(VName, TypeBase Exp Uniqueness)]
free_ts = forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall dim u1 u2. TypeBase dim u1 -> u2 -> TypeBase dim u2
`setUniqueness` Uniqueness
Nonunique)) forall a b. (a -> b) -> a -> b
$ [(VName, StructType)]
free_dims forall a. [a] -> [a] -> [a]
++ [(VName, StructType)]
free_nondims

  ValBind -> LiftM ()
addValBind forall a b. (a -> b) -> a -> b
$
    ValBind
      { valBindName :: VName
valBindName = VName
fname,
        valBindTypeParams :: [TypeParam]
valBindTypeParams = [TypeParam]
tparams,
        valBindParams :: [Pat ParamType]
valBindParams = forall a b. (a -> b) -> [a] -> [b]
map forall {vn} {u}. (vn, TypeBase Exp u) -> PatBase Info vn ParamType
mkParam [(VName, TypeBase Exp Uniqueness)]
free_ts forall a. [a] -> [a] -> [a]
++ [Pat ParamType]
params,
        valBindRetDecl :: Maybe (TypeExp Info VName)
valBindRetDecl = forall a. Maybe a
Nothing,
        valBindRetType :: Info ResRetType
valBindRetType = forall a. a -> Info a
Info (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret),
        valBindBody :: Exp
valBindBody = Exp
funbody,
        valBindDoc :: Maybe DocComment
valBindDoc = forall a. Maybe a
Nothing,
        valBindAttrs :: [AttrInfo VName]
valBindAttrs = forall a. Monoid a => a
mempty,
        valBindLocation :: SrcLoc
valBindLocation = forall a. Monoid a => a
mempty,
        valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = forall a. Maybe a
Nothing
      }

  forall (f :: * -> *) a. Applicative f => a -> f a
pure
    forall a b. (a -> b) -> a -> b
$ Exp -> [(VName, StructType)] -> Exp
apply
      (forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
fname) (forall a. a -> Info a
Info (forall {u}. [(VName, TypeBase Exp u)] -> StructType
augType [(VName, TypeBase Exp Uniqueness)]
free_ts)) forall a. Monoid a => a
mempty)
    forall a b. (a -> b) -> a -> b
$ [(VName, StructType)]
free_dims forall a. [a] -> [a] -> [a]
++ [(VName, StructType)]
free_nondims
  where
    orig_type :: StructType
orig_type = [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
params forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
ret
    mkParam :: (vn, TypeBase Exp u) -> PatBase Info vn ParamType
mkParam (vn
v, TypeBase Exp u
t) = forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id vn
v (forall a. a -> Info a
Info (forall u. Diet -> TypeBase Exp u -> ParamType
toParam Diet
Observe TypeBase Exp u
t)) forall a. Monoid a => a
mempty
    freeVar :: (vn, StructType) -> ExpBase Info vn
freeVar (vn
v, StructType
t) = forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName vn
v) (forall a. a -> Info a
Info StructType
t) forall a. Monoid a => a
mempty
    augType :: [(VName, TypeBase Exp u)] -> StructType
augType [(VName, TypeBase Exp u)]
rem_free = [Pat ParamType] -> ResRetType -> StructType
funType (forall a b. (a -> b) -> [a] -> [b]
map forall {vn} {u}. (vn, TypeBase Exp u) -> PatBase Info vn ParamType
mkParam [(VName, TypeBase Exp u)]
rem_free) forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall u. TypeBase Exp u -> TypeBase Exp Uniqueness
toRet StructType
orig_type

    apply :: Exp -> [(VName, StructType)] -> Exp
    apply :: Exp -> [(VName, StructType)] -> Exp
apply Exp
f [] = Exp
f
    apply Exp
f ((VName, StructType)
p : [(VName, StructType)]
rem_ps) =
      let inner_ret :: AppRes
inner_ret = StructType -> [VName] -> AppRes
AppRes (forall {u}. [(VName, TypeBase Exp u)] -> StructType
augType [(VName, StructType)]
rem_ps) forall a. Monoid a => a
mempty
          inner :: Exp
inner = forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
f [(Diet
Observe, forall a. Maybe a
Nothing, forall {vn}. (vn, StructType) -> ExpBase Info vn
freeVar (VName, StructType)
p)] AppRes
inner_ret
       in Exp -> [(VName, StructType)] -> Exp
apply Exp
inner [(VName, StructType)]
rem_ps

transformSubExps :: ASTMapper LiftM
transformSubExps :: ASTMapper LiftM
transformSubExps = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> LiftM Exp
mapOnExp = Exp -> LiftM Exp
transformExp}

transformExp :: Exp -> LiftM Exp
transformExp :: Exp -> LiftM Exp
transformExp (AppExp (LetFun VName
fname ([TypeParam]
tparams, [Pat ParamType]
params, Maybe (TypeExp Info VName)
_, Info ResRetType
ret, Exp
funbody) Exp
body SrcLoc
_) Info AppRes
_) = do
  Exp
funbody' <- forall a. [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
tparams) [Pat ParamType]
params forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
funbody
  VName
fname' <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName forall a b. (a -> b) -> a -> b
$ [Char]
"lifted_" forall a. [a] -> [a] -> [a]
++ VName -> [Char]
baseString VName
fname
  Exp
lifted_call <- VName
-> [TypeParam] -> [Pat ParamType] -> ResRetType -> Exp -> LiftM Exp
liftFunction VName
fname' [TypeParam]
tparams [Pat ParamType]
params ResRetType
ret Exp
funbody'
  forall a. VName -> Exp -> LiftM a -> LiftM a
replacing VName
fname Exp
lifted_call forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
body
transformExp (Lambda [Pat ParamType]
params Exp
body Maybe (TypeExp Info VName)
_ (Info ResRetType
ret) SrcLoc
_) = do
  Exp
body' <- forall a. [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams [] [Pat ParamType]
params forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
body
  VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"lifted_lambda"
  VName
-> [TypeParam] -> [Pat ParamType] -> ResRetType -> Exp -> LiftM Exp
liftFunction VName
fname [] [Pat ParamType]
params ResRetType
ret Exp
body'
transformExp (AppExp (LetPat [SizeBinder VName]
sizes Pat StructType
pat Exp
e Exp
body SrcLoc
loc) Info AppRes
appres) = do
  Exp
e' <- Exp -> LiftM Exp
transformExp Exp
e
  Exp
body' <- forall a. [VName] -> Pat StructType -> LiftM a -> LiftM a
bindingLetPat (forall a b. (a -> b) -> [a] -> [b]
map forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
sizes) Pat StructType
pat forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
body
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes Pat StructType
pat Exp
e' Exp
body' SrcLoc
loc) Info AppRes
appres
transformExp (AppExp (DoLoop [VName]
sizes Pat ParamType
pat Exp
args LoopFormBase Info VName
form Exp
body SrcLoc
loc) Info AppRes
appres) = do
  Exp
args' <- Exp -> LiftM Exp
transformExp Exp
args
  LoopFormBase Info VName
form' <- forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper LiftM
transformSubExps LoopFormBase Info VName
form
  Exp
body' <- forall a. [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams [VName]
sizes [Pat ParamType
pat] forall a b. (a -> b) -> a -> b
$ forall a. LoopFormBase Info VName -> LiftM a -> LiftM a
bindingForm LoopFormBase Info VName
form' forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
body
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop [VName]
sizes Pat ParamType
pat Exp
args' LoopFormBase Info VName
form' Exp
body' SrcLoc
loc) Info AppRes
appres
transformExp e :: Exp
e@(Var QualName VName
v Info StructType
_ SrcLoc
_) =
  -- Note that function-typed variables can only occur in expressions,
  -- not in other places where VNames/QualNames can occur.
  forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (forall a. a -> Maybe a -> a
fromMaybe Exp
e forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName Exp
envReplace)
transformExp Exp
e = forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper LiftM
transformSubExps Exp
e

transformValBind :: ValBind -> LiftM ()
transformValBind :: ValBind -> LiftM ()
transformValBind ValBind
vb = do
  Exp
e <-
    forall a. [VName] -> [Pat ParamType] -> LiftM a -> LiftM a
bindingParams (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. ValBindBase f vn -> [TypeParamBase vn]
valBindTypeParams ValBind
vb) (forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
vb) forall a b. (a -> b) -> a -> b
$
      Exp -> LiftM Exp
transformExp (forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBind
vb)
  ValBind -> LiftM ()
addValBind forall a b. (a -> b) -> a -> b
$ ValBind
vb {valBindBody :: Exp
valBindBody = Exp
e}

{-# NOINLINE transformProg #-}

-- | Perform the transformation.
transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
vbinds =
  forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
    VNameSource -> LiftM () -> ([ValBind], VNameSource)
runLiftM VNameSource
namesrc forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ValBind -> LiftM ()
transformValBind [ValBind]
vbinds