{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE Trustworthy #-}

-- | 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 qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import Futhark.IR.Pretty ()
import qualified Futhark.Internalise.FreeVars as FV
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals

newtype Env = Env {Env -> Map VName Exp
envReplace :: M.Map VName Exp}

initialEnv :: Env
initialEnv :: Env
initialEnv = Map VName Exp -> Env
Env Map VName Exp
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 [ValBind]
forall a. Monoid a => a
mempty (Set VName -> LiftState) -> Set VName -> LiftState
forall a b. (a -> b) -> a -> b
$ [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [VName]
forall k a. Map k a -> [k]
M.keys Map VName Intrinsic
intrinsics

newtype LiftM a = LiftM (ReaderT Env (State LiftState) a)
  deriving (a -> LiftM b -> LiftM a
(a -> b) -> LiftM a -> LiftM b
(forall a b. (a -> b) -> LiftM a -> LiftM b)
-> (forall a b. a -> LiftM b -> LiftM a) -> Functor LiftM
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
<$ :: a -> LiftM b -> LiftM a
$c<$ :: forall a b. a -> LiftM b -> LiftM a
fmap :: (a -> b) -> LiftM a -> LiftM b
$cfmap :: forall a b. (a -> b) -> LiftM a -> LiftM b
Functor, Functor LiftM
a -> LiftM a
Functor LiftM
-> (forall a. a -> LiftM a)
-> (forall a b. LiftM (a -> b) -> LiftM a -> LiftM b)
-> (forall a b c. (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c)
-> (forall a b. LiftM a -> LiftM b -> LiftM b)
-> (forall a b. LiftM a -> LiftM b -> LiftM a)
-> Applicative LiftM
LiftM a -> LiftM b -> LiftM b
LiftM a -> LiftM b -> LiftM a
LiftM (a -> b) -> LiftM a -> LiftM b
(a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
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
<* :: LiftM a -> LiftM b -> LiftM a
$c<* :: forall a b. LiftM a -> LiftM b -> LiftM a
*> :: LiftM a -> LiftM b -> LiftM b
$c*> :: forall a b. LiftM a -> LiftM b -> LiftM b
liftA2 :: (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
$cliftA2 :: forall a b c. (a -> b -> c) -> LiftM a -> LiftM b -> LiftM c
<*> :: LiftM (a -> b) -> LiftM a -> LiftM b
$c<*> :: forall a b. LiftM (a -> b) -> LiftM a -> LiftM b
pure :: a -> LiftM a
$cpure :: forall a. a -> LiftM a
$cp1Applicative :: Functor LiftM
Applicative, Applicative LiftM
a -> LiftM a
Applicative LiftM
-> (forall a b. LiftM a -> (a -> LiftM b) -> LiftM b)
-> (forall a b. LiftM a -> LiftM b -> LiftM b)
-> (forall a. a -> LiftM a)
-> Monad LiftM
LiftM a -> (a -> LiftM b) -> LiftM b
LiftM a -> LiftM b -> LiftM b
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 :: a -> LiftM a
$creturn :: forall a. a -> LiftM a
>> :: LiftM a -> LiftM b -> LiftM b
$c>> :: forall a b. LiftM a -> LiftM b -> LiftM b
>>= :: LiftM a -> (a -> LiftM b) -> LiftM b
$c>>= :: forall a b. LiftM a -> (a -> LiftM b) -> LiftM b
$cp1Monad :: Applicative LiftM
Monad, MonadReader Env, MonadState LiftState)

instance MonadFreshNames LiftM where
  putNameSource :: VNameSource -> LiftM ()
putNameSource VNameSource
src = (LiftState -> LiftState) -> LiftM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((LiftState -> LiftState) -> LiftM ())
-> (LiftState -> LiftState) -> LiftM ()
forall a b. (a -> b) -> a -> b
$ \LiftState
s -> LiftState
s {stateNameSource :: VNameSource
stateNameSource = VNameSource
src}
  getNameSource :: LiftM VNameSource
getNameSource = (LiftState -> VNameSource) -> LiftM VNameSource
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 = State LiftState () -> LiftState -> LiftState
forall s a. State s a -> s -> s
execState (ReaderT Env (State LiftState) () -> Env -> State LiftState ()
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 ([ValBind] -> [ValBind]
forall a. [a] -> [a]
reverse (LiftState -> [ValBind]
stateValBinds LiftState
s), LiftState -> VNameSource
stateNameSource LiftState
s)

addValBind :: ValBind -> LiftM ()
addValBind :: ValBind -> LiftM ()
addValBind ValBind
vb = (LiftState -> LiftState) -> LiftM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((LiftState -> LiftState) -> LiftM ())
-> (LiftState -> LiftState) -> LiftM ()
forall a b. (a -> b) -> a -> b
$ \LiftState
s ->
  LiftState
s
    { stateValBinds :: [ValBind]
stateValBinds = ValBind
vb ValBind -> [ValBind] -> [ValBind]
forall a. a -> [a] -> [a]
: LiftState -> [ValBind]
stateValBinds LiftState
s,
      stateGlobal :: Set VName
stateGlobal = (Set VName -> VName -> Set VName)
-> Set VName -> [VName] -> Set VName
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((VName -> Set VName -> Set VName)
-> Set VName -> VName -> Set VName
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert) (LiftState -> Set VName
stateGlobal LiftState
s) [VName]
names
    }
  where
    names :: [VName]
names = ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
vb VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: (StructType, [VName]) -> [VName]
forall a b. (a, b) -> b
snd (Info (StructType, [VName]) -> (StructType, [VName])
forall a. Info a -> a
unInfo (ValBind -> Info (StructType, [VName])
forall (f :: * -> *) vn.
ValBindBase f vn -> f (StructType, [VName])
valBindRetType ValBind
vb))

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

existentials :: Exp -> S.Set VName
existentials :: Exp -> Set VName
existentials Exp
e =
  let here :: Set VName
here = case Exp
e of
        Apply Exp
_ Exp
_ (Info (Diet
_, Maybe VName
pdim)) (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (Maybe VName -> [VName]
forall a. Maybe a -> [a]
maybeToList Maybe VName
pdim [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
ext)
        If Exp
_ Exp
_ Exp
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        LetPat PatternBase Info VName
_ Exp
_ Exp
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        Coerce Exp
_ TypeDeclBase Info VName
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        Range Exp
_ Maybe Exp
_ Inclusiveness Exp
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        Index Exp
_ [DimIndexBase Info VName]
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        Match Exp
_ NonEmpty (CaseBase Info VName)
_ (Info PatternType
_, Info [VName]
ext) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        DoLoop [VName]
_ PatternBase Info VName
_ Exp
_ LoopFormBase Info VName
_ Exp
_ (Info (PatternType
_, [VName]
ext)) SrcLoc
_ ->
          [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
ext
        Exp
_ ->
          Set VName
forall a. Monoid a => a
mempty

      m :: ASTMapper (StateT (Set VName) Identity)
m = ASTMapper (StateT (Set VName) Identity)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> StateT (Set VName) Identity Exp
mapOnExp = \Exp
e' -> (Set VName -> Set VName) -> StateT (Set VName) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Exp -> Set VName
existentials Exp
e') StateT (Set VName) Identity ()
-> StateT (Set VName) Identity Exp
-> StateT (Set VName) Identity Exp
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Exp -> StateT (Set VName) Identity Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e'}
   in StateT (Set VName) Identity Exp -> Set VName -> Set VName
forall s a. State s a -> s -> s
execState (ASTMapper (StateT (Set VName) Identity)
-> Exp -> StateT (Set VName) Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT (Set VName) Identity)
m Exp
e) Set VName
here

liftFunction :: VName -> [TypeParam] -> [Pattern] -> StructType -> Exp -> LiftM Exp
liftFunction :: VName
-> [TypeParam]
-> [PatternBase Info VName]
-> StructType
-> Exp
-> LiftM Exp
liftFunction VName
fname [TypeParam]
tparams [PatternBase Info VName]
params StructType
ret Exp
funbody = do
  -- Find free variables
  Set VName
global <- (LiftState -> Set VName) -> LiftM (Set VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets LiftState -> Set VName
stateGlobal
  let bound :: Set VName
bound =
        Set VName
global
          Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames [PatternBase Info VName]
params
          Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((TypeParam -> VName) -> [TypeParam] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParam]
tparams)
          Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Exp -> Set VName
existentials Exp
funbody

      free :: [(VName, StructType)]
free =
        let immediate_free :: NameSet
immediate_free = Exp -> NameSet
FV.freeVars Exp
funbody NameSet -> Set VName -> NameSet
`FV.without` Set VName
bound
            sizes_in_free :: Set VName
sizes_in_free =
              (StructType -> Set VName) -> [StructType] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames ([StructType] -> Set VName) -> [StructType] -> Set VName
forall a b. (a -> b) -> a -> b
$
                Map VName StructType -> [StructType]
forall k a. Map k a -> [a]
M.elems (Map VName StructType -> [StructType])
-> Map VName StructType -> [StructType]
forall a b. (a -> b) -> a -> b
$ NameSet -> Map VName StructType
FV.unNameSet NameSet
immediate_free
            sizes :: NameSet
sizes =
              Set VName -> NameSet
FV.sizes (Set VName -> NameSet) -> Set VName -> NameSet
forall a b. (a -> b) -> a -> b
$
                Set VName
sizes_in_free
                  Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternDimNames [PatternBase Info VName]
params
                  Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames StructType
ret
         in Map VName StructType -> [(VName, StructType)]
forall k a. Map k a -> [(k, a)]
M.toList (Map VName StructType -> [(VName, StructType)])
-> Map VName StructType -> [(VName, StructType)]
forall a b. (a -> b) -> a -> b
$ NameSet -> Map VName StructType
FV.unNameSet (NameSet -> Map VName StructType)
-> NameSet -> Map VName StructType
forall a b. (a -> b) -> a -> b
$ NameSet
immediate_free NameSet -> NameSet -> NameSet
forall a. Semigroup a => a -> a -> a
<> (NameSet
sizes NameSet -> Set VName -> NameSet
`FV.without` Set VName
bound)

      -- Those parameters that correspond to sizes must come first.
      sizes_in_types :: Set VName
sizes_in_types =
        (StructType -> Set VName) -> [StructType] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames (StructType
ret StructType -> [StructType] -> [StructType]
forall a. a -> [a] -> [a]
: ((VName, StructType) -> StructType)
-> [(VName, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (VName, StructType) -> StructType
forall a b. (a, b) -> b
snd [(VName, StructType)]
free [StructType] -> [StructType] -> [StructType]
forall a. [a] -> [a] -> [a]
++ (PatternBase Info VName -> StructType)
-> [PatternBase Info VName] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> StructType
patternStructType [PatternBase Info VName]
params)
      isSize :: (VName, b) -> Bool
isSize (VName
v, b
_) = VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes_in_types
      ([(VName, StructType)]
free_dims, [(VName, StructType)]
free_nondims) = ((VName, StructType) -> Bool)
-> [(VName, StructType)]
-> ([(VName, StructType)], [(VName, StructType)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (VName, StructType) -> Bool
forall b. (VName, b) -> Bool
isSize [(VName, StructType)]
free

      free_params :: [PatternBase Info VName]
free_params =
        ((VName, StructType) -> PatternBase Info VName)
-> [(VName, StructType)] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map ((VName, StructType) -> PatternBase Info VName
forall vn as.
(vn, TypeBase (DimDecl VName) as) -> PatternBase Info vn
mkParam ((VName, StructType) -> PatternBase Info VName)
-> ((VName, StructType) -> (VName, StructType))
-> (VName, StructType)
-> PatternBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StructType -> StructType)
-> (VName, StructType) -> (VName, StructType)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (StructType -> Uniqueness -> StructType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique)) ([(VName, StructType)] -> [PatternBase Info VName])
-> [(VName, StructType)] -> [PatternBase Info VName]
forall a b. (a -> b) -> a -> b
$
          [(VName, StructType)]
free_dims [(VName, StructType)]
-> [(VName, StructType)] -> [(VName, StructType)]
forall a. [a] -> [a] -> [a]
++ [(VName, StructType)]
free_nondims

  ValBind -> LiftM ()
addValBind (ValBind -> LiftM ()) -> ValBind -> LiftM ()
forall a b. (a -> b) -> a -> b
$
    ValBind :: forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructType, [VName])
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
ValBind
      { valBindName :: VName
valBindName = VName
fname,
        valBindTypeParams :: [TypeParam]
valBindTypeParams = [TypeParam]
tparams,
        valBindParams :: [PatternBase Info VName]
valBindParams = [PatternBase Info VName]
free_params [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. [a] -> [a] -> [a]
++ [PatternBase Info VName]
params,
        valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
forall a. Maybe a
Nothing,
        valBindRetType :: Info (StructType, [VName])
valBindRetType = (StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info (StructType
ret, [VName]
forall a. Monoid a => a
mempty),
        valBindBody :: Exp
valBindBody = Exp
funbody,
        valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing,
        valBindAttrs :: [AttrInfo]
valBindAttrs = [AttrInfo]
forall a. Monoid a => a
mempty,
        valBindLocation :: SrcLoc
valBindLocation = SrcLoc
forall a. Monoid a => a
mempty,
        valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing
      }

  Exp -> LiftM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> LiftM Exp) -> Exp -> LiftM Exp
forall a b. (a -> b) -> a -> b
$
    Exp -> [(VName, StructType)] -> Exp
apply
      (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
fname) (PatternType -> Info PatternType
forall a. a -> Info a
Info ([(VName, StructType)] -> PatternType
forall as. [(VName, TypeBase (DimDecl VName) as)] -> PatternType
augType ([(VName, StructType)] -> PatternType)
-> [(VName, StructType)] -> PatternType
forall a b. (a -> b) -> a -> b
$ [(VName, StructType)]
free_dims [(VName, StructType)]
-> [(VName, StructType)] -> [(VName, StructType)]
forall a. [a] -> [a] -> [a]
++ [(VName, StructType)]
free_nondims)) SrcLoc
forall a. Monoid a => a
mempty)
      ([(VName, StructType)] -> Exp) -> [(VName, StructType)] -> Exp
forall a b. (a -> b) -> a -> b
$ [(VName, StructType)]
free_dims [(VName, StructType)]
-> [(VName, StructType)] -> [(VName, StructType)]
forall a. [a] -> [a] -> [a]
++ [(VName, StructType)]
free_nondims
  where
    orig_type :: StructType
orig_type = [PatternBase Info VName] -> StructType -> StructType
funType [PatternBase Info VName]
params StructType
ret
    mkParam :: (vn, TypeBase (DimDecl VName) as) -> PatternBase Info vn
mkParam (vn
v, TypeBase (DimDecl VName) as
t) = vn -> Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id vn
v (PatternType -> Info PatternType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) as -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) as
t)) SrcLoc
forall a. Monoid a => a
mempty
    freeVar :: (vn, TypeBase (DimDecl VName) as) -> ExpBase Info vn
freeVar (vn
v, TypeBase (DimDecl VName) as
t) = QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
v) (PatternType -> Info PatternType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) as -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) as
t)) SrcLoc
forall a. Monoid a => a
mempty
    augType :: [(VName, TypeBase (DimDecl VName) as)] -> PatternType
augType [(VName, TypeBase (DimDecl VName) as)]
rem_free = StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (StructType -> PatternType) -> StructType -> PatternType
forall a b. (a -> b) -> a -> b
$ [PatternBase Info VName] -> StructType -> StructType
funType (((VName, TypeBase (DimDecl VName) as) -> PatternBase Info VName)
-> [(VName, TypeBase (DimDecl VName) as)]
-> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName, TypeBase (DimDecl VName) as) -> PatternBase Info VName
forall vn as.
(vn, TypeBase (DimDecl VName) as) -> PatternBase Info vn
mkParam [(VName, TypeBase (DimDecl VName) as)]
rem_free) 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 :: (Info PatternType, Info [VName])
inner_ret = (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct ([(VName, StructType)] -> PatternType
forall as. [(VName, TypeBase (DimDecl VName) as)] -> PatternType
augType [(VName, StructType)]
rem_ps)), [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty)
          inner :: Exp
inner = Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply Exp
f ((VName, StructType) -> Exp
forall vn as. (vn, TypeBase (DimDecl VName) as) -> ExpBase Info vn
freeVar (VName, StructType)
p) ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing)) (Info PatternType, Info [VName])
inner_ret SrcLoc
forall a. Monoid a => a
mempty
       in Exp -> [(VName, StructType)] -> Exp
apply Exp
inner [(VName, StructType)]
rem_ps

transformExp :: Exp -> LiftM Exp
transformExp :: Exp -> LiftM Exp
transformExp (LetFun VName
fname ([TypeParam]
tparams, [PatternBase Info VName]
params, Maybe (TypeExp VName)
_, Info StructType
ret, Exp
funbody) Exp
body Info PatternType
_ SrcLoc
_) = do
  Exp
funbody' <- Exp -> LiftM Exp
transformExp Exp
funbody
  VName
fname' <- String -> LiftM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (String -> LiftM VName) -> String -> LiftM VName
forall a b. (a -> b) -> a -> b
$ String
"lifted_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ VName -> String
baseString VName
fname
  Exp
lifted_call <- VName
-> [TypeParam]
-> [PatternBase Info VName]
-> StructType
-> Exp
-> LiftM Exp
liftFunction VName
fname' [TypeParam]
tparams [PatternBase Info VName]
params StructType
ret Exp
funbody'
  VName -> Exp -> LiftM Exp -> LiftM Exp
forall a. VName -> Exp -> LiftM a -> LiftM a
replacing VName
fname Exp
lifted_call (LiftM Exp -> LiftM Exp) -> LiftM Exp -> LiftM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> LiftM Exp
transformExp Exp
body
transformExp (Lambda [PatternBase Info VName]
params Exp
body Maybe (TypeExp VName)
_ (Info (Aliasing
_, StructType
ret)) SrcLoc
_) = do
  Exp
body' <- Exp -> LiftM Exp
transformExp Exp
body
  VName
fname <- String -> LiftM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"lifted_lambda"
  VName
-> [TypeParam]
-> [PatternBase Info VName]
-> StructType
-> Exp
-> LiftM Exp
liftFunction VName
fname [] [PatternBase Info VName]
params StructType
ret Exp
body'
transformExp e :: Exp
e@(Var QualName VName
v Info PatternType
_ SrcLoc
_) =
  -- Note that function-typed variables can only occur in expressions,
  -- not in other places where VNames/QualNames can occur.
  (Env -> Exp) -> LiftM Exp
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Exp -> Maybe Exp -> Exp
forall a. a -> Maybe a -> a
fromMaybe Exp
e (Maybe Exp -> Exp) -> (Env -> Maybe Exp) -> Env -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Map VName Exp -> Maybe Exp
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) (Map VName Exp -> Maybe Exp)
-> (Env -> Map VName Exp) -> Env -> Maybe Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName Exp
envReplace)
transformExp Exp
e =
  ASTMapper LiftM -> Exp -> LiftM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper LiftM
m Exp
e
  where
    m :: ASTMapper LiftM
m = ASTMapper LiftM
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> LiftM Exp
mapOnExp = Exp -> LiftM Exp
transformExp}

transformValBind :: ValBind -> LiftM ()
transformValBind :: ValBind -> LiftM ()
transformValBind ValBind
vb = do
  Exp
e <- Exp -> LiftM Exp
transformExp (Exp -> LiftM Exp) -> Exp -> LiftM Exp
forall a b. (a -> b) -> a -> b
$ ValBind -> Exp
forall (f :: * -> *) vn. ValBindBase f vn -> ExpBase f vn
valBindBody ValBind
vb
  ValBind -> LiftM ()
addValBind (ValBind -> LiftM ()) -> ValBind -> LiftM ()
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 :: [ValBind] -> m [ValBind]
transformProg [ValBind]
vbinds =
  (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([ValBind], VNameSource)) -> m [ValBind])
-> (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
    VNameSource -> LiftM () -> ([ValBind], VNameSource)
runLiftM VNameSource
namesrc (LiftM () -> ([ValBind], VNameSource))
-> LiftM () -> ([ValBind], VNameSource)
forall a b. (a -> b) -> a -> b
$ (ValBind -> LiftM ()) -> [ValBind] -> LiftM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ValBind -> LiftM ()
transformValBind [ValBind]
vbinds