module Futhark.Internalise.Defunctionalise (transformProg) where
import Control.Arrow qualified as Arrow
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition, sortOn, tails)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Pretty ()
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals
data ExtExp
= ExtLambda [Pat] Exp StructRetType SrcLoc
| ExtExp Exp
deriving (Int -> ExtExp -> ShowS
[ExtExp] -> ShowS
ExtExp -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [ExtExp] -> ShowS
$cshowList :: [ExtExp] -> ShowS
show :: ExtExp -> [Char]
$cshow :: ExtExp -> [Char]
showsPrec :: Int -> ExtExp -> ShowS
$cshowsPrec :: Int -> ExtExp -> ShowS
Show)
data StaticVal
= Dynamic PatType
| LambdaSV Pat StructRetType ExtExp Env
| RecordSV [(Name, StaticVal)]
|
SumSV Name [StaticVal] [(Name, [PatType])]
|
DynamicFun (Exp, StaticVal) StaticVal
| IntrinsicSV
deriving (Int -> StaticVal -> ShowS
[StaticVal] -> ShowS
StaticVal -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [StaticVal] -> ShowS
$cshowList :: [StaticVal] -> ShowS
show :: StaticVal -> [Char]
$cshow :: StaticVal -> [Char]
showsPrec :: Int -> StaticVal -> ShowS
$cshowsPrec :: Int -> StaticVal -> ShowS
Show)
data Binding = Binding (Maybe ([VName], StructType)) StaticVal
deriving (Int -> Binding -> ShowS
[Binding] -> ShowS
Binding -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [Binding] -> ShowS
$cshowList :: [Binding] -> ShowS
show :: Binding -> [Char]
$cshow :: Binding -> [Char]
showsPrec :: Int -> Binding -> ShowS
$cshowsPrec :: Int -> Binding -> ShowS
Show)
bindingSV :: Binding -> StaticVal
bindingSV :: Binding -> StaticVal
bindingSV (Binding Maybe ([VName], StructType)
_ StaticVal
sv) = StaticVal
sv
type Env = M.Map VName Binding
localEnv :: Env -> DefM a -> DefM a
localEnv :: forall a. Env -> DefM a -> DefM a
localEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
Arrow.second (Env
env <>)
localNewEnv :: Env -> DefM a -> DefM a
localNewEnv :: forall a. Env -> DefM a -> DefM a
localNewEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ \(Set VName
globals, Env
old_env) ->
(Set VName
globals, forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\VName
k Binding
_ -> VName
k forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals) Env
old_env forall a. Semigroup a => a -> a -> a
<> Env
env)
askEnv :: DefM Env
askEnv :: DefM Env
askEnv = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> b
snd
areGlobal :: [VName] -> DefM a -> DefM a
areGlobal :: forall a. [VName] -> DefM a -> DefM a
areGlobal [VName]
vs = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
Arrow.first (forall a. Ord a => [a] -> Set a
S.fromList [VName]
vs <>)
replaceTypeSizes ::
M.Map VName SizeSubst ->
TypeBase Size als ->
TypeBase Size als
replaceTypeSizes :: forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim
where
onDim :: Size -> Size
onDim (NamedSize QualName VName
v) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> QualName VName -> Size
NamedSize QualName VName
v'
Just (SubstConst Int
d) -> Int -> Size
ConstSize Int
d
Maybe SizeSubst
Nothing -> QualName VName -> Size
NamedSize QualName VName
v
onDim Size
d = Size
d
replaceStaticValSizes ::
S.Set VName ->
M.Map VName SizeSubst ->
StaticVal ->
StaticVal
replaceStaticValSizes :: Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv =
case StaticVal
sv of
StaticVal
_ | forall k a. Map k a -> Bool
M.null Map VName SizeSubst
orig_substs -> StaticVal
sv
LambdaSV PatBase Info VName
param (RetType [VName]
t_dims StructType
t) ExtExp
e Env
closure_env ->
let substs :: Map VName SizeSubst
substs =
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 k a. Ord k => k -> Map k a -> Map k a
M.delete) Map VName SizeSubst
orig_substs forall a b. (a -> b) -> a -> b
$
forall a. Ord a => [a] -> Set a
S.fromList (forall k a. Map k a -> [k]
M.keys Env
closure_env)
in PatBase Info VName -> StructRetType -> ExtExp -> Env -> StaticVal
LambdaSV
(forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs PatBase Info VName
param)
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs StructType
t))
(Map VName SizeSubst -> ExtExp -> ExtExp
onExtExp Map VName SizeSubst
substs ExtExp
e)
(forall {k}.
Ord k =>
Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
orig_substs Env
closure_env)
Dynamic PatType
t ->
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
orig_substs PatType
t
RecordSV [(Name, StaticVal)]
fs ->
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs)) [(Name, StaticVal)]
fs
SumSV Name
c [StaticVal]
svs [(Name, [PatType])]
ts ->
Name -> [StaticVal] -> [(Name, [PatType])] -> StaticVal
SumSV Name
c (forall a b. (a -> b) -> [a] -> [b]
map (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs) [StaticVal]
svs) forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
orig_substs) [(Name, [PatType])]
ts
DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2 ->
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
orig_substs Exp
e, Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv1) forall a b. (a -> b) -> a -> b
$
Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv2
StaticVal
IntrinsicSV ->
StaticVal
IntrinsicSV
where
tv :: Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs =
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper
{ mapOnPatType :: PatType -> m PatType
mapOnPatType = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnStructType :: StructType -> m StructType
mapOnStructType = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs,
mapOnExp :: Exp -> m Exp
mapOnExp = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs,
mapOnName :: VName -> m VName
mapOnName = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs
}
onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
Maybe SizeSubst
_ -> VName
v
onExp :: Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs (Var QualName VName
v Info PatType
t SrcLoc
loc) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') ->
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info PatType
t SrcLoc
loc
Just (SubstConst Int
d) ->
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (Int64 -> IntValue
Int64Value (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d))) SrcLoc
loc
Maybe SizeSubst
Nothing ->
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info PatType
t) SrcLoc
loc
onExp Map VName SizeSubst
substs (AppExp (Coerce Exp
e TypeExp VName
te SrcLoc
loc) (Info (AppRes PatType
t [VName]
ext))) =
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> AppExpBase f vn
Coerce (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e) TypeExp VName
te' SrcLoc
loc) (forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs PatType
t) [VName]
ext))
where
te' :: TypeExp VName
te' = Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te
onExp Map VName SizeSubst
substs Exp
e = forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs Exp
e
onTypeExpDim :: Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs d :: SizeExp VName
d@(SizeExpNamed QualName VName
v SrcLoc
loc) =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') ->
forall vn. QualName vn -> SrcLoc -> SizeExp vn
SizeExpNamed QualName VName
v' SrcLoc
loc
Just (SubstConst Int
x) ->
forall vn. Int -> SrcLoc -> SizeExp vn
SizeExpConst Int
x SrcLoc
loc
Maybe SizeSubst
Nothing ->
SizeExp VName
d
onTypeExpDim Map VName SizeSubst
_ SizeExp VName
d = SizeExp VName
d
onTypeArgExp :: Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs (TypeArgExpDim SizeExp VName
d SrcLoc
loc) =
forall vn. SizeExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim (Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs SizeExp VName
d) SrcLoc
loc
onTypeArgExp Map VName SizeSubst
substs (TypeArgExpType TypeExp VName
te) =
forall vn. TypeExp vn -> TypeArgExp vn
TypeArgExpType (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te)
onTypeExp :: Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs (TEArray SizeExp VName
d TypeExp VName
te SrcLoc
loc) =
forall vn. SizeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArray (Map VName SizeSubst -> SizeExp VName -> SizeExp VName
onTypeExpDim Map VName SizeSubst
substs SizeExp VName
d) (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEUnique TypeExp VName
t SrcLoc
loc) =
forall vn. TypeExp vn -> SrcLoc -> TypeExp vn
TEUnique (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEApply TypeExp VName
t1 TypeArgExp VName
t2 SrcLoc
loc) =
forall vn. TypeExp vn -> TypeArgExp vn -> SrcLoc -> TypeExp vn
TEApply (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs TypeArgExp VName
t2) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEArrow Maybe VName
p TypeExp VName
t1 TypeExp VName
t2 SrcLoc
loc) =
forall vn.
Maybe vn -> TypeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArrow Maybe VName
p (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t2) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TETuple [TypeExp VName]
ts SrcLoc
loc) =
forall vn. [TypeExp vn] -> SrcLoc -> TypeExp vn
TETuple (forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [TypeExp VName]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TERecord [(Name, TypeExp VName)]
ts SrcLoc
loc) =
forall vn. [(Name, TypeExp vn)] -> SrcLoc -> TypeExp vn
TERecord (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, TypeExp VName)]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TESum [(Name, [TypeExp VName])]
ts SrcLoc
loc) =
forall vn. [(Name, [TypeExp vn])] -> SrcLoc -> TypeExp vn
TESum (forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, [TypeExp VName])]
ts) SrcLoc
loc
onTypeExp Map VName SizeSubst
substs (TEDim [VName]
dims TypeExp VName
t SrcLoc
loc) =
forall vn. [vn] -> TypeExp vn -> SrcLoc -> TypeExp vn
TEDim [VName]
dims (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t) SrcLoc
loc
onTypeExp Map VName SizeSubst
_ (TEVar QualName VName
v SrcLoc
loc) =
forall vn. QualName vn -> SrcLoc -> TypeExp vn
TEVar QualName VName
v SrcLoc
loc
onExtExp :: Map VName SizeSubst -> ExtExp -> ExtExp
onExtExp Map VName SizeSubst
substs (ExtExp Exp
e) =
Exp -> ExtExp
ExtExp forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e
onExtExp Map VName SizeSubst
substs (ExtLambda [PatBase Info VName]
params Exp
e (RetType [VName]
t_dims StructType
t) SrcLoc
loc) =
[PatBase Info VName] -> Exp -> StructRetType -> SrcLoc -> ExtExp
ExtLambda
(forall a b. (a -> b) -> [a] -> [b]
map (forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs) [PatBase Info VName]
params)
(Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
t_dims (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs StructType
t))
SrcLoc
loc
onEnv :: Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
substs =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList
onBinding :: Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs (Binding Maybe ([VName], StructType)
t StaticVal
bsv) =
Maybe ([VName], StructType) -> StaticVal -> Binding
Binding
(forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
substs) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ([VName], StructType)
t)
(Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
bsv)
onAST :: ASTMappable x => M.Map VName SizeSubst -> x -> x
onAST :: forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (forall {m :: * -> *}. Monad m => Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs)
restrictEnvTo :: FV -> DefM Env
restrictEnvTo :: FV -> DefM Env
restrictEnvTo (FV Map VName StructType
m) = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Env
restrict
where
restrict :: (Set VName, Env) -> Env
restrict (Set VName
globals, Env
env) = forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey VName -> Binding -> Maybe Binding
keep Env
env
where
keep :: VName -> Binding -> Maybe Binding
keep VName
k (Binding Maybe ([VName], StructType)
t StaticVal
sv) = do
forall (f :: * -> *). Alternative f => Bool -> f ()
guard forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ VName
k forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals
Uniqueness
u <- forall shape as. TypeBase shape as -> Uniqueness
uniqueness 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
k Map VName StructType
m
forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv
restrict' :: Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
Nonunique (Dynamic PatType
t) =
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ PatType
t forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
restrict' Uniqueness
_ (Dynamic PatType
t) =
PatType -> StaticVal
Dynamic PatType
t
restrict' Uniqueness
u (LambdaSV PatBase Info VName
pat StructRetType
t ExtExp
e Env
env) =
PatBase Info VName -> StructRetType -> ExtExp -> Env -> StaticVal
LambdaSV PatBase Info VName
pat StructRetType
t ExtExp
e forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u) Env
env
restrict' Uniqueness
u (RecordSV [(Name, StaticVal)]
fields) =
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [(Name, StaticVal)]
fields
restrict' Uniqueness
u (SumSV Name
c [StaticVal]
svs [(Name, [PatType])]
fields) =
Name -> [StaticVal] -> [(Name, [PatType])] -> StaticVal
SumSV Name
c (forall a b. (a -> b) -> [a] -> [b]
map (Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [StaticVal]
svs) [(Name, [PatType])]
fields
restrict' Uniqueness
u (DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2) =
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp
e, Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv1) forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv2
restrict' Uniqueness
_ StaticVal
IntrinsicSV = StaticVal
IntrinsicSV
restrict'' :: Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u (Binding Maybe ([VName], StructType)
t StaticVal
sv) = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv
newtype DefM a
= DefM (ReaderT (S.Set VName, Env) (State ([ValBind], VNameSource)) a)
deriving
( forall a b. a -> DefM b -> DefM a
forall a b. (a -> b) -> DefM a -> DefM 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 -> DefM b -> DefM a
$c<$ :: forall a b. a -> DefM b -> DefM a
fmap :: forall a b. (a -> b) -> DefM a -> DefM b
$cfmap :: forall a b. (a -> b) -> DefM a -> DefM b
Functor,
Functor DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM 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. DefM a -> DefM b -> DefM a
$c<* :: forall a b. DefM a -> DefM b -> DefM a
*> :: forall a b. DefM a -> DefM b -> DefM b
$c*> :: forall a b. DefM a -> DefM b -> DefM b
liftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
$cliftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
$c<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
pure :: forall a. a -> DefM a
$cpure :: forall a. a -> DefM a
Applicative,
Applicative DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM a -> (a -> DefM b) -> DefM 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 -> DefM a
$creturn :: forall a. a -> DefM a
>> :: forall a b. DefM a -> DefM b -> DefM b
$c>> :: forall a b. DefM a -> DefM b -> DefM b
>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
$c>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
Monad,
MonadReader (S.Set VName, Env),
MonadState ([ValBind], VNameSource)
)
instance MonadFreshNames DefM where
putNameSource :: VNameSource -> DefM ()
putNameSource VNameSource
src = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \([ValBind]
x, VNameSource
_) -> ([ValBind]
x, VNameSource
src)
getNameSource :: DefM VNameSource
getNameSource = forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets forall a b. (a, b) -> b
snd
runDefM :: VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM :: forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
src (DefM ReaderT (Set VName, Env) (State ([ValBind], VNameSource)) a
m) =
let (a
x, ([ValBind]
vbs, VNameSource
src')) = forall s a. State s a -> s -> (a, s)
runState (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Set VName, Env) (State ([ValBind], VNameSource)) a
m forall a. Monoid a => a
mempty) (forall a. Monoid a => a
mempty, VNameSource
src)
in (a
x, VNameSource
src', forall a. [a] -> [a]
reverse [ValBind]
vbs)
addValBind :: ValBind -> DefM ()
addValBind :: ValBind -> DefM ()
addValBind ValBind
vb = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (ValBind
vb :)
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar StructType
t VName
x = do
Env
env <- DefM Env
askEnv
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Env
env of
Just (Binding (Just ([VName]
dims, StructType
sv_t)) StaticVal
sv) -> do
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [VName] -> StructType -> StructType -> StaticVal -> m StaticVal
instStaticVal Set VName
globals [VName]
dims StructType
t StructType
sv_t StaticVal
sv
Just (Binding Maybe ([VName], StructType)
Nothing StaticVal
sv) ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
sv
Maybe Binding
Nothing
| VName -> Int
baseTag VName
x forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag -> forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
IntrinsicSV
| Bool
otherwise ->
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
arraySizes :: StructType -> S.Set VName
arraySizes :: StructType -> Set VName
arraySizes (Scalar Arrow {}) = forall a. Monoid a => a
mempty
arraySizes (Scalar (Record Map Name StructType
fields)) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes Map Name StructType
fields
arraySizes (Scalar (Sum Map Name [StructType]
cs)) = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes) Map Name [StructType]
cs
arraySizes (Scalar (TypeVar ()
_ Uniqueness
_ QualName VName
_ [TypeArg Size]
targs)) =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map TypeArg Size -> Set VName
f [TypeArg Size]
targs
where
f :: TypeArg Size -> Set VName
f (TypeArgDim (NamedSize QualName VName
d) SrcLoc
_) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
d
f TypeArgDim {} = forall a. Monoid a => a
mempty
f (TypeArgType StructType
t SrcLoc
_) = StructType -> Set VName
arraySizes StructType
t
arraySizes (Scalar Prim {}) = forall a. Monoid a => a
mempty
arraySizes (Array ()
_ Uniqueness
_ Shape Size
shape ScalarTypeBase Size ()
t) =
StructType -> Set VName
arraySizes (forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase Size ()
t) forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Size -> Set VName
dimName (forall dim. Shape dim -> [dim]
shapeDims Shape Size
shape)
where
dimName :: Size -> S.Set VName
dimName :: Size -> Set VName
dimName (NamedSize QualName VName
qn) = forall a. a -> Set a
S.singleton forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
dimName Size
_ = forall a. Monoid a => a
mempty
patternArraySizes :: Pat -> S.Set VName
patternArraySizes :: PatBase Info VName -> Set VName
patternArraySizes = StructType -> Set VName
arraySizes forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> StructType
patternStructType
data SizeSubst
= SubstNamed (QualName VName)
| SubstConst Int
deriving (SizeSubst -> SizeSubst -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SizeSubst -> SizeSubst -> Bool
$c/= :: SizeSubst -> SizeSubst -> Bool
== :: SizeSubst -> SizeSubst -> Bool
$c== :: SizeSubst -> SizeSubst -> Bool
Eq, Eq SizeSubst
SizeSubst -> SizeSubst -> Bool
SizeSubst -> SizeSubst -> Ordering
SizeSubst -> SizeSubst -> SizeSubst
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: SizeSubst -> SizeSubst -> SizeSubst
$cmin :: SizeSubst -> SizeSubst -> SizeSubst
max :: SizeSubst -> SizeSubst -> SizeSubst
$cmax :: SizeSubst -> SizeSubst -> SizeSubst
>= :: SizeSubst -> SizeSubst -> Bool
$c>= :: SizeSubst -> SizeSubst -> Bool
> :: SizeSubst -> SizeSubst -> Bool
$c> :: SizeSubst -> SizeSubst -> Bool
<= :: SizeSubst -> SizeSubst -> Bool
$c<= :: SizeSubst -> SizeSubst -> Bool
< :: SizeSubst -> SizeSubst -> Bool
$c< :: SizeSubst -> SizeSubst -> Bool
compare :: SizeSubst -> SizeSubst -> Ordering
$ccompare :: SizeSubst -> SizeSubst -> Ordering
Ord, Int -> SizeSubst -> ShowS
[SizeSubst] -> ShowS
SizeSubst -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [SizeSubst] -> ShowS
$cshowList :: [SizeSubst] -> ShowS
show :: SizeSubst -> [Char]
$cshow :: SizeSubst -> [Char]
showsPrec :: Int -> SizeSubst -> ShowS
$cshowsPrec :: Int -> SizeSubst -> ShowS
Show)
dimMapping ::
Monoid a =>
TypeBase Size a ->
TypeBase Size a ->
M.Map VName SizeSubst
dimMapping :: forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping TypeBase Size a
t1 TypeBase Size a
t2 = forall s a. State s a -> s -> s
execState (forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims forall {t :: * -> *} {f :: * -> *}.
(Foldable t, MonadState (Map VName SizeSubst) f) =>
t VName -> Size -> Size -> f Size
f TypeBase Size a
t1 TypeBase Size a
t2) forall a. Monoid a => a
mempty
where
f :: t VName -> Size -> Size -> f Size
f t VName
bound Size
d1 (NamedSize QualName VName
d2)
| forall vn. QualName vn -> vn
qualLeaf QualName VName
d2 forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
bound = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d1
f t VName
_ (NamedSize QualName VName
d1) (NamedSize QualName VName
d2) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall vn. QualName vn -> vn
qualLeaf QualName VName
d1) forall a b. (a -> b) -> a -> b
$ QualName VName -> SizeSubst
SubstNamed QualName VName
d2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize QualName VName
d1
f t VName
_ (NamedSize QualName VName
d1) (ConstSize Int
d2) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (forall vn. QualName vn -> vn
qualLeaf QualName VName
d1) forall a b. (a -> b) -> a -> b
$ Int -> SizeSubst
SubstConst Int
d2
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize QualName VName
d1
f t VName
_ Size
d Size
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
dimMapping' ::
Monoid a =>
TypeBase Size a ->
TypeBase Size a ->
M.Map VName VName
dimMapping' :: forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName VName
dimMapping' TypeBase Size a
t1 TypeBase Size a
t2 = forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe SizeSubst -> Maybe VName
f forall a b. (a -> b) -> a -> b
$ forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping TypeBase Size a
t1 TypeBase Size a
t2
where
f :: SizeSubst -> Maybe VName
f (SubstNamed QualName VName
d) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
d
f SizeSubst
_ = forall a. Maybe a
Nothing
sizesToRename :: StaticVal -> S.Set VName
sizesToRename :: StaticVal -> Set VName
sizesToRename (DynamicFun (Exp
_, StaticVal
sv1) StaticVal
sv2) =
StaticVal -> Set VName
sizesToRename StaticVal
sv1 forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv2
sizesToRename StaticVal
IntrinsicSV =
forall a. Monoid a => a
mempty
sizesToRename Dynamic {} =
forall a. Monoid a => a
mempty
sizesToRename (RecordSV [(Name, StaticVal)]
fs) =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StaticVal -> Set VName
sizesToRename forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fs
sizesToRename (SumSV Name
_ [StaticVal]
svs [(Name, [PatType])]
_) =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StaticVal -> Set VName
sizesToRename [StaticVal]
svs
sizesToRename (LambdaSV PatBase Info VName
param StructRetType
_ ExtExp
_ Env
_) =
PatBase Info VName -> Set VName
freeInPat PatBase Info VName
param
forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall a. (a -> Bool) -> Set a -> Set a
S.filter forall {vn}. IdentBase Info vn -> Bool
couldBeSize forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents PatBase Info VName
param)
where
couldBeSize :: IdentBase Info vn -> Bool
couldBeSize IdentBase Info vn
ident =
forall a. Info a -> a
unInfo (forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info vn
ident) forall a. Eq a => a -> a -> Bool
== forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (forall dim as. PrimType -> ScalarTypeBase dim as
Prim (IntType -> PrimType
Signed IntType
Int64))
instStaticVal ::
MonadFreshNames m =>
S.Set VName ->
[VName] ->
StructType ->
StructType ->
StaticVal ->
m StaticVal
instStaticVal :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [VName] -> StructType -> StructType -> StaticVal -> m StaticVal
instStaticVal Set VName
globals [VName]
dims StructType
t StructType
sv_t StaticVal
sv = do
Map VName SizeSubst
fresh_substs <- forall {f :: * -> *}.
MonadFreshNames f =>
[VName] -> f (Map VName SizeSubst)
mkSubsts forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv
let dims' :: [VName]
dims' = forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
fresh_substs) [VName]
dims
isDim :: VName -> p -> Bool
isDim VName
k p
_ = VName
k forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
dims'
dim_substs :: Map VName SizeSubst
dim_substs =
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey forall {p}. VName -> p -> Bool
isDim forall a b. (a -> b) -> a -> b
$ forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName SizeSubst
dimMapping (forall als.
Map VName SizeSubst -> TypeBase Size als -> TypeBase Size als
replaceTypeSizes Map VName SizeSubst
fresh_substs StructType
sv_t) StructType
t
replace :: SizeSubst -> SizeSubst
replace (SubstNamed QualName VName
k) = forall a. a -> Maybe a -> a
fromMaybe (QualName VName -> SizeSubst
SubstNamed QualName VName
k) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. QualName vn -> vn
qualLeaf QualName VName
k) Map VName SizeSubst
dim_substs
replace SizeSubst
k = SizeSubst
k
substs :: Map VName SizeSubst
substs = forall a b k. (a -> b) -> Map k a -> Map k b
M.map SizeSubst -> SizeSubst
replace Map VName SizeSubst
fresh_substs forall a. Semigroup a => a -> a -> a
<> Map VName SizeSubst
dim_substs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
sv
where
mkSubsts :: [VName] -> f (Map VName SizeSubst)
mkSubsts [VName]
names =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> SizeSubst
SubstNamed forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. v -> QualName v
qualName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName [VName]
names
onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
Just (SubstNamed QualName VName
v') -> forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
Maybe SizeSubst
_ -> VName
v
defuncFun ::
[VName] ->
[Pat] ->
Exp ->
StructRetType ->
SrcLoc ->
DefM (Exp, StaticVal)
defuncFun :: [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
tparams [PatBase Info VName]
pats Exp
e0 StructRetType
ret SrcLoc
loc = do
let (PatBase Info VName
pat, StructRetType
ret', ExtExp
e0') = case [PatBase Info VName]
pats of
[] -> forall a. HasCallStack => [Char] -> a
error [Char]
"Received a lambda with no parameters."
[PatBase Info VName
pat'] -> (PatBase Info VName
pat', StructRetType
ret, Exp -> ExtExp
ExtExp Exp
e0)
(PatBase Info VName
pat' : [PatBase Info VName]
pats') ->
( PatBase Info VName
pat',
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (forall a b. (a -> b) -> [a] -> [b]
map (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName -> PatType
patternType) [PatBase Info VName]
pats') StructRetType
ret,
[PatBase Info VName] -> Exp -> StructRetType -> SrcLoc -> ExtExp
ExtLambda [PatBase Info VName]
pats' Exp
e0 StructRetType
ret SrcLoc
loc
)
let used :: FV
used =
Exp -> FV
freeInExp (forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
e0 forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
ret)) SrcLoc
loc)
FV -> Set VName -> FV
`freeWithout` forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams
Env
used_env <- FV -> DefM Env
restrictEnvTo FV
used
let sizes_of_arrays :: Set VName
sizes_of_arrays =
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StructType -> Set VName
arraySizes forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> PatType
typeFromSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV) Env
used_env
forall a. Semigroup a => a -> a -> a
<> PatBase Info VName -> Set VName
patternArraySizes PatBase Info VName
pat
notSize :: VName -> Bool
notSize = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes_of_arrays)
([FieldBase Info VName]
fields, Env
env) =
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. [(a, b)] -> ([a], [b])
unzip
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Bool
notSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Env
used_env
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fields SrcLoc
loc,
PatBase Info VName -> StructRetType -> ExtExp -> Env -> StaticVal
LambdaSV PatBase Info VName
pat StructRetType
ret' ExtExp
e0' Env
env
)
where
closureFromDynamicFun :: (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ (DynamicFun (Exp
clsr_env, StaticVal
sv) StaticVal
_)) =
let name :: Name
name = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString VName
vn
in ( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name Exp
clsr_env forall a. Monoid a => a
mempty,
(VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
)
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
let name :: Name
name = [Char] -> Name
nameFromString forall a b. (a -> b) -> a -> b
$ forall a. Pretty a => a -> [Char]
prettyString VName
vn
tp' :: PatType
tp' = StaticVal -> PatType
typeFromSV StaticVal
sv
in ( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
Name
name
(forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
vn) (forall a. a -> Info a
Info PatType
tp') forall a. Monoid a => a
mempty)
forall a. Monoid a => a
mempty,
(VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
)
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp e :: Exp
e@Literal {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@IntLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@FloatLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp e :: Exp
e@StringLit {} =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
defuncExp (Parens Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp (TupLit [Exp]
es SrcLoc
loc) = do
([Exp]
es', [StaticVal]
svs) <- forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [StaticVal]
svs)
defuncExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
([FieldBase Info VName]
fs', [(Name, StaticVal)]
names_svs) <- forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField [FieldBase Info VName]
fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fs' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV [(Name, StaticVal)]
names_svs)
where
defuncField :: FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField (RecordFieldExplicit Name
vn Exp
e SrcLoc
loc') = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn Exp
e' SrcLoc
loc', (Name
vn, StaticVal
sv))
defuncField (RecordFieldImplicit VName
vn (Info PatType
t) SrcLoc
loc') = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) VName
vn
case StaticVal
sv of
DynamicFun (Exp
e, StaticVal
sv') StaticVal
_ ->
let vn' :: Name
vn' = VName -> Name
baseName VName
vn
in forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn' Exp
e SrcLoc
loc',
(Name
vn', StaticVal
sv')
)
StaticVal
_ ->
let tp :: Info PatType
tp = forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv
in forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
vn -> f PatType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit VName
vn Info PatType
tp SrcLoc
loc', (VName -> Name
baseName VName
vn, StaticVal
sv))
defuncExp (ArrayLit [Exp]
es t :: Info PatType
t@(Info PatType
t') SrcLoc
loc) = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
ArrayLit [Exp]
es' Info PatType
t SrcLoc
loc, PatType -> StaticVal
Dynamic PatType
t')
defuncExp (AppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) Info AppRes
res) = do
Exp
e1' <- Exp -> DefM Exp
defuncExp' Exp
e1
Maybe Exp
me' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' Maybe Exp
me
Inclusiveness Exp
incl' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp' Inclusiveness Exp
incl
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' SrcLoc
loc) Info AppRes
res,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ AppRes -> PatType
appResType forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo Info AppRes
res
)
defuncExp e :: Exp
e@(Var QualName VName
qn (Info PatType
t) SrcLoc
loc) = do
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
case StaticVal
sv of
DynamicFun (Exp, StaticVal)
closure StaticVal
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp, StaticVal)
closure
StaticVal
IntrinsicSV -> do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e) Exp
e
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
StaticVal
_ ->
let tp :: PatType
tp = StaticVal -> PatType
typeFromSV StaticVal
sv
in forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info PatType
tp) SrcLoc
loc, StaticVal
sv)
defuncExp (Hole (Info PatType
t) SrcLoc
loc) =
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
Hole (forall a. a -> Info a
Info PatType
t) SrcLoc
loc, StaticVal
IntrinsicSV)
defuncExp (Ascript Exp
e0 TypeExp VName
tydecl SrcLoc
loc)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e0) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> ExpBase f vn
Ascript Exp
e0' TypeExp VName
tydecl SrcLoc
loc, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (Coerce Exp
e0 TypeExp VName
tydecl SrcLoc
loc) Info AppRes
res)
| forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e0) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> AppExpBase f vn
Coerce Exp
e0' TypeExp VName
tydecl SrcLoc
loc) Info AppRes
res, StaticVal
sv)
| Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName
pat Exp
e1 Exp
e2 SrcLoc
loc) (Info (AppRes PatType
t [VName]
retext))) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
let env :: Env
env = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv1
pat' :: PatBase Info VName
pat' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv1
(Exp
e2', StaticVal
sv2) <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let mapping :: Map VName VName
mapping = forall a.
Monoid a =>
TypeBase Size a -> TypeBase Size a -> Map VName VName
dimMapping' (Exp -> PatType
typeOf Exp
e2) PatType
t
subst :: VName -> VName
subst VName
v = forall a. a -> Maybe a -> a
fromMaybe VName
v forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
mapping
mapper :: ASTMapper Identity
mapper = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnName :: VName -> Identity VName
mapOnName = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> VName
subst}
t' :: PatType
t' = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper Identity
mapper) forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e2'
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes PatBase Info VName
pat' Exp
e1' Exp
e2' SrcLoc
loc) (forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
t' [VName]
retext)), StaticVal
sv2)
defuncExp (AppExp (LetFun VName
vn ([TypeParamBase VName], [PatBase Info VName],
Maybe (TypeExp VName), Info StructRetType, Exp)
_ Exp
_ SrcLoc
_) Info AppRes
_) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: Unexpected LetFun: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show VName
vn
defuncExp (AppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) Info AppRes
res) = do
(Exp
e1', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
(Exp
e3', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If Exp
e1' Exp
e2' Exp
e3' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
defuncExp e :: Exp
e@(AppExp (Apply f :: Exp
f@(Var QualName VName
f' Info PatType
_ SrcLoc
_) Exp
arg Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
res)
| VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
f') forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
TupLit [Exp]
es SrcLoc
tuploc <- Exp
arg = do
[Exp]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncSoacExp [Exp]
es
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f (forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [Exp]
es' SrcLoc
tuploc) Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
res,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e
)
defuncExp e :: Exp
e@(AppExp Apply {} Info AppRes
_) = Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
0 Exp
e
defuncExp (Negate Exp
e0 SrcLoc
loc) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate Exp
e0' SrcLoc
loc, StaticVal
sv)
defuncExp (Not Exp
e0 SrcLoc
loc) = do
(Exp
e0', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not Exp
e0' SrcLoc
loc, StaticVal
sv)
defuncExp (Lambda [PatBase Info VName]
pats Exp
e0 Maybe (TypeExp VName)
_ (Info (Set Alias
_, StructRetType
ret)) SrcLoc
loc) =
[VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [PatBase Info VName]
pats Exp
e0 StructRetType
ret SrcLoc
loc
defuncExp OpSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionLeft {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionRight {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp ProjectSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected projection section."
defuncExp IndexSection {} = forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected projection section."
defuncExp (AppExp (DoLoop [VName]
sparams PatBase Info VName
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 SrcLoc
loc) Info AppRes
res) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
let env1 :: Env
env1 = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv1
(LoopFormBase Info VName
form', Env
env2) <- case LoopFormBase Info VName
form of
For IdentBase Info VName
v Exp
e2 -> do
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
v Exp
e2', forall {k}. IdentBase Info k -> Map k Binding
envFromIdent IdentBase Info VName
v)
ForIn PatBase Info VName
pat2 Exp
e2 -> do
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn PatBase Info VName
pat2 Exp
e2', PatBase Info VName -> Env
envFromPat PatBase Info VName
pat2)
While Exp
e2 -> do
Exp
e2' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env1 forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While Exp
e2', forall a. Monoid a => a
mempty)
(Exp
e3', StaticVal
sv) <- forall a. Env -> DefM a -> DefM a
localEnv (Env
env1 forall a. Semigroup a => a -> a -> a
<> Env
env2) forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e3
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
[VName]
-> PatBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop [VName]
sparams PatBase Info VName
pat Exp
e1' LoopFormBase Info VName
form' Exp
e3' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
where
envFromIdent :: IdentBase Info k -> Map k Binding
envFromIdent (Ident k
vn (Info PatType
tp) SrcLoc
_) =
forall k a. k -> a -> Map k a
M.singleton k
vn forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic PatType
tp
defuncExp e :: Exp
e@(AppExp BinOp {} Info AppRes
_) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: unexpected binary operator: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e
defuncExp (Project Name
vn Exp
e0 tp :: Info PatType
tp@(Info PatType
tp') SrcLoc
loc) = do
(Exp
e0', StaticVal
sv0) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
case StaticVal
sv0 of
RecordSV [(Name, StaticVal)]
svs -> case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
vn [(Name, StaticVal)]
svs of
Just StaticVal
sv -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) SrcLoc
loc, StaticVal
sv)
Maybe StaticVal
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
Dynamic PatType
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info PatType
tp SrcLoc
loc, PatType -> StaticVal
Dynamic PatType
tp')
StaticVal
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Projection of an expression with static value " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv0
defuncExp (AppExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs Exp
e1 Exp
body SrcLoc
loc) Info AppRes
res) = do
Exp
e1' <- Exp -> DefM Exp
defuncExp' Exp
e1
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
let id1_binding :: Binding
id1_binding = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall a. Info a -> a
unInfo forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType IdentBase Info VName
id1
(Exp
body', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localEnv (forall k a. k -> a -> Map k a
M.singleton (forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
id1) Binding
id1_binding) forall a b. (a -> b) -> a -> b
$
Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
IdentBase f vn
-> IdentBase f vn
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs' Exp
e1' Exp
body' SrcLoc
loc) Info AppRes
res, StaticVal
sv)
defuncExp expr :: Exp
expr@(AppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) Info AppRes
res) = do
Exp
e0' <- Exp -> DefM Exp
defuncExp' Exp
e0
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index Exp
e0' SliceBase Info VName
idxs' SrcLoc
loc) Info AppRes
res,
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
expr
)
defuncExp (Update Exp
e1 SliceBase Info VName
idxs Exp
e2 SrcLoc
loc) = do
(Exp
e1', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
SliceBase Info VName
idxs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs
Exp
e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update Exp
e1' SliceBase Info VName
idxs' Exp
e2' SrcLoc
loc, StaticVal
sv)
defuncExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info PatType
_ SrcLoc
loc) = do
(Exp
e1', StaticVal
sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv2) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let sv :: StaticVal
sv = StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv1 StaticVal
sv2 [Name]
fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
RecordUpdate Exp
e1' [Name]
fs Exp
e2' (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv1) SrcLoc
loc,
StaticVal
sv
)
where
staticField :: StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (RecordSV [(Name, StaticVal)]
svs) StaticVal
sv2 (Name
f : [Name]
fs') =
case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
f [(Name, StaticVal)]
svs of
Just StaticVal
sv ->
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$
(Name
f, StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv StaticVal
sv2 [Name]
fs') forall a. a -> [a] -> [a]
: forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
/= Name
f) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Name, StaticVal)]
svs
Maybe StaticVal
Nothing -> forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
staticField (Dynamic t :: PatType
t@(Scalar Record {})) StaticVal
sv2 fs' :: [Name]
fs'@(Name
_ : [Name]
_) =
StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (PatType -> StaticVal
svFromType PatType
t) StaticVal
sv2 [Name]
fs'
staticField StaticVal
_ StaticVal
sv2 [Name]
_ = StaticVal
sv2
defuncExp (Assert Exp
e1 Exp
e2 Info Text
desc SrcLoc
loc) = do
(Exp
e1', StaticVal
_) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
(Exp
e2', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert Exp
e1' Exp
e2' Info Text
desc SrcLoc
loc, StaticVal
sv)
defuncExp (Constr Name
name [Exp]
es (Info (Scalar (Sum Map Name [PatType]
all_fs))) SrcLoc
loc) = do
([Exp]
es', [StaticVal]
svs) <- forall a b. [(a, b)] -> ([a], [b])
unzip forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
let sv :: StaticVal
sv =
Name -> [StaticVal] -> [(Name, [PatType])] -> StaticVal
SumSV Name
name [StaticVal]
svs forall a b. (a -> b) -> a -> b
$
forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$
Name
name forall k a. Ord k => k -> Map k a -> Map k a
`M.delete` forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType) Map Name [PatType]
all_fs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
Constr Name
name [Exp]
es' (forall a. a -> Info a
Info (StaticVal -> PatType
typeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
where
defuncType ::
Monoid als =>
TypeBase Size als ->
TypeBase Size als
defuncType :: forall als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType (Array als
as Uniqueness
u Shape Size
shape ScalarTypeBase Size ()
t) = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array als
as Uniqueness
u Shape Size
shape (forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar ScalarTypeBase Size ()
t)
defuncType (Scalar ScalarTypeBase Size als
t) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar ScalarTypeBase Size als
t
defuncScalar ::
Monoid als =>
ScalarTypeBase Size als ->
ScalarTypeBase Size als
defuncScalar :: forall als.
Monoid als =>
ScalarTypeBase Size als -> ScalarTypeBase Size als
defuncScalar (Record Map Name (TypeBase Size als)
fs) = forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType Map Name (TypeBase Size als)
fs
defuncScalar Arrow {} = forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a. Monoid a => a
mempty
defuncScalar (Sum Map Name [TypeBase Size als]
fs) = forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map (forall a b. (a -> b) -> [a] -> [b]
map forall als. Monoid als => TypeBase Size als -> TypeBase Size als
defuncType) Map Name [TypeBase Size als]
fs
defuncScalar (Prim PrimType
t) = forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
defuncScalar (TypeVar als
as Uniqueness
u QualName VName
tn [TypeArg Size]
targs) = forall dim as.
as
-> Uniqueness
-> QualName VName
-> [TypeArg dim]
-> ScalarTypeBase dim as
TypeVar als
as Uniqueness
u QualName VName
tn [TypeArg Size]
targs
defuncExp (Constr Name
name [Exp]
_ (Info PatType
t) SrcLoc
loc) =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Constructor "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
name
forall a. [a] -> [a] -> [a]
++ [Char]
" given type "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatType
t
forall a. [a] -> [a] -> [a]
++ [Char]
" at "
forall a. [a] -> [a] -> [a]
++ forall a. Located a => a -> [Char]
locStr SrcLoc
loc
defuncExp (AppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
res) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
NonEmpty (CaseBase Info VName, StaticVal)
csPairs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv) NonEmpty (CaseBase Info VName)
cs
let cs' :: NonEmpty (CaseBase Info VName)
cs' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst NonEmpty (CaseBase Info VName, StaticVal)
csPairs
sv' :: StaticVal
sv' = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> a
NE.head NonEmpty (CaseBase Info VName, StaticVal)
csPairs
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Exp
e' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) Info AppRes
res, StaticVal
sv')
defuncExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) = do
(Exp
e', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
info Exp
e' SrcLoc
loc, StaticVal
sv)
defuncExp' :: Exp -> DefM Exp
defuncExp' :: Exp -> DefM Exp
defuncExp' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> DefM (Exp, StaticVal)
defuncExp
defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)
defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)
defuncExtExp (ExtExp Exp
e) = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
defuncExtExp (ExtLambda [PatBase Info VName]
pats Exp
e0 StructRetType
ret SrcLoc
loc) =
[VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [PatBase Info VName]
pats Exp
e0 StructRetType
ret SrcLoc
loc
defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal)
defuncCase :: StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv (CasePat PatBase Info VName
p Exp
e SrcLoc
loc) = do
let p' :: PatBase Info VName
p' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
p StaticVal
sv
env :: Env
env = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv
(Exp
e', StaticVal
sv') <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat PatBase Info VName
p' Exp
e' SrcLoc
loc, StaticVal
sv')
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp e :: Exp
e@OpSection {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionLeft {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionRight {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@ProjectSection {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp (Parens Exp
e SrcLoc
loc) =
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM Exp
defuncSoacExp Exp
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
defuncSoacExp (Lambda [PatBase Info VName]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Set Alias, StructRetType)
tp SrcLoc
loc) = do
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
params
Exp
e0' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncSoacExp Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Set Alias, StructRetType)
tp SrcLoc
loc
defuncSoacExp Exp
e
| Scalar Arrow {} <- Exp -> PatType
typeOf Exp
e = do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e) Exp
e
let env :: Env
env = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
pats
Exp
body' <- forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncExp' Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body' forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
| Bool
otherwise = Exp -> DefM Exp
defuncExp' Exp
e
etaExpand :: PatType -> Exp -> DefM ([Pat], Exp, StructRetType)
etaExpand :: PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand PatType
e_t Exp
e = do
let ([(PName, StructType)]
ps, RetTypeBase Size (Set Alias)
ret) = forall {dim} {as}.
RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
e_t
([PatBase Info VName]
pats, [Exp]
vars) <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(PName, StructType)]
ps forall a b. (a -> b) -> a -> b
$ \(PName
p, StructType
t) -> do
let t' :: PatType
t' = forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
t
VName
x <- case PName
p of
Named VName
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
x
PName
Unnamed -> forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (forall a. a -> Info a
Info PatType
t') forall a. Monoid a => a
mempty,
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (forall v. v -> QualName v
qualName VName
x) (forall a. a -> Info a
Info PatType
t') forall a. Monoid a => a
mempty
)
let e' :: Exp
e' =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
( \Exp
e1 (Exp
e2, StructType
t2, [StructType]
argtypes) ->
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1 Exp
e2 (forall a. a -> Info a
Info (forall shape as. TypeBase shape as -> Diet
diet StructType
t2, forall a. Maybe a
Nothing)) forall a. Monoid a => a
mempty)
(forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [StructType]
argtypes RetTypeBase Size (Set Alias)
ret) []))
)
Exp
e
forall a b. (a -> b) -> a -> b
$ forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Exp]
vars (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(PName, StructType)]
ps) (forall a. Int -> [a] -> [a]
drop Int
1 forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [[a]]
tails forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(PName, StructType)]
ps)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PatBase Info VName]
pats, Exp
e', forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a b. a -> b -> a
const ()) RetTypeBase Size (Set Alias)
ret)
where
getType :: RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType (RetType [VName]
_ (Scalar (Arrow as
_ PName
p TypeBase dim ()
t1 RetTypeBase dim as
t2))) =
let ([(PName, TypeBase dim ())]
ps, RetTypeBase dim as
r) = RetTypeBase dim as
-> ([(PName, TypeBase dim ())], RetTypeBase dim as)
getType RetTypeBase dim as
t2 in ((PName
p, TypeBase dim ()
t1) forall a. a -> [a] -> [a]
: [(PName, TypeBase dim ())]
ps, RetTypeBase dim as
r)
getType RetTypeBase dim as
t = ([], RetTypeBase dim as
t)
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex (DimFix Exp
e1) = forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
defuncDimIndex (DimSlice Maybe Exp
me1 Maybe Exp
me2 Maybe Exp
me3) =
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me1 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me2 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me3
where
defunc' :: Maybe Exp -> DefM (Maybe Exp)
defunc' = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> DefM Exp
defuncExp'
defuncLet ::
[VName] ->
[Pat] ->
Exp ->
StructRetType ->
DefM ([VName], [Pat], Exp, StaticVal)
defuncLet :: [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet [VName]
dims ps :: [PatBase Info VName]
ps@(PatBase Info VName
pat : [PatBase Info VName]
pats) Exp
body (RetType [VName]
ret_dims StructType
rettype)
| forall vn. PatBase Info vn -> Bool
patternOrderZero PatBase Info VName
pat = do
let bound_by_pat :: VName -> Bool
bound_by_pat = (forall a. Ord a => a -> Set a -> Bool
`S.member` PatBase Info VName -> Set VName
freeInPat PatBase Info VName
pat)
([VName]
pat_dims, [VName]
rest_dims) = forall a. (a -> Bool) -> [a] -> ([a], [a])
partition VName -> Bool
bound_by_pat [VName]
dims
env :: Env
env = PatBase Info VName -> Env
envFromPat PatBase Info VName
pat forall a. Semigroup a => a -> a -> a
<> [VName] -> Env
envFromDimNames [VName]
pat_dims
([VName]
rest_dims', [PatBase Info VName]
pats', Exp
body', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet [VName]
rest_dims [PatBase Info VName]
pats Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype
(Exp, StaticVal)
closure <- [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [PatBase Info VName]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( [VName]
pat_dims forall a. [a] -> [a] -> [a]
++ [VName]
rest_dims',
PatBase Info VName
pat forall a. a -> [a] -> [a]
: [PatBase Info VName]
pats',
Exp
body',
(Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
closure StaticVal
sv
)
| Bool
otherwise = do
(Exp
e, StaticVal
sv) <- [VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [PatBase Info VName]
ps Exp
body (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], Exp
e, StaticVal
sv)
defuncLet [VName]
_ [] Exp
body (RetType [VName]
_ StructType
rettype) = do
(Exp
body', StaticVal
sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [], Exp
body', forall {as}. StaticVal -> TypeBase Size as -> StaticVal
imposeType StaticVal
sv StructType
rettype)
where
imposeType :: StaticVal -> TypeBase Size as -> StaticVal
imposeType Dynamic {} TypeBase Size as
t =
PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase Size as
t
imposeType (RecordSV [(Name, StaticVal)]
fs1) (Scalar (Record Map Name (TypeBase Size as)
fs2)) =
[(Name, StaticVal)] -> StaticVal
RecordSV forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith StaticVal -> TypeBase Size as -> StaticVal
imposeType (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, StaticVal)]
fs1) Map Name (TypeBase Size as)
fs2
imposeType StaticVal
sv TypeBase Size as
_ = StaticVal
sv
sizesForAll :: MonadFreshNames m => S.Set VName -> [Pat] -> m ([VName], [Pat])
sizesForAll :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params = do
([PatBase Info VName]
params', Set VName
sizes) <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT (Set VName) m)
tv) [PatBase Info VName]
params) forall a. Monoid a => a
mempty
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Set a -> [a]
S.toList Set VName
sizes, [PatBase Info VName]
params')
where
bound :: Set VName
bound = Set VName
bound_sizes forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
params
tv :: ASTMapper (StateT (Set VName) m)
tv = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatType :: PatType -> StateT (Set VName) m PatType
mapOnPatType = forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadState (Set VName) (t m), MonadTrans t, MonadFreshNames m) =>
Size -> t m Size
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure}
onDim :: Size -> t m Size
onDim (AnySize (Just VName
v)) = do
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
S.insert VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
v
onDim (AnySize Maybe VName
Nothing) = do
VName
v <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"size"
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall a. Ord a => a -> Set a -> Set a
S.insert VName
v
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
v
onDim (NamedSize QualName VName
d) = do
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall vn. QualName vn -> vn
qualLeaf QualName VName
d forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) forall a b. (a -> b) -> a -> b
$
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$
forall a. Ord a => a -> Set a -> Set a
S.insert forall a b. (a -> b) -> a -> b
$
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize QualName VName
d
onDim Size
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d
unRetType :: StructRetType -> StructType
unRetType :: StructRetType -> StructType
unRetType (RetType [] StructType
t) = StructType
t
unRetType (RetType [VName]
ext StructType
t) = forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Size -> Size
onDim StructType
t
where
onDim :: Size -> Size
onDim (NamedSize QualName VName
d) | forall vn. QualName vn -> vn
qualLeaf QualName VName
d forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
ext = Maybe VName -> Size
AnySize forall a. Maybe a
Nothing
onDim Size
d = Size
d
defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)
defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
depth e :: Exp
e@(AppExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d SrcLoc
loc) t :: Info AppRes
t@(Info (AppRes PatType
ret [VName]
ext))) = do
let ([StructType]
argtypes, StructType
_) = forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType PatType
ret
(Exp
e1', StaticVal
sv1) <- Int -> Exp -> DefM (Exp, StaticVal)
defuncApply (Int
depth forall a. Num a => a -> a -> a
+ Int
1) Exp
e1
(Exp
e2', StaticVal
sv2) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e2
let e' :: Exp
e' = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d SrcLoc
loc) Info AppRes
t
case StaticVal
sv1 of
LambdaSV PatBase Info VName
pat StructRetType
e0_t ExtExp
e0 Env
closure_env -> do
let env' :: Env
env' = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv2
dims :: [VName]
dims = forall a. Monoid a => a
mempty
(Exp
e0', StaticVal
sv) <-
forall a. Env -> DefM a -> DefM a
localNewEnv (Env
env' forall a. Semigroup a => a -> a -> a
<> Env
closure_env) forall a b. (a -> b) -> a -> b
$
ExtExp -> DefM (Exp, StaticVal)
defuncExtExp ExtExp
e0
let closure_pat :: PatBase Info VName
closure_pat = [VName] -> Env -> PatBase Info VName
buildEnvPat [VName]
dims Env
closure_env
pat' :: PatBase Info VName
pat' = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv2
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let params :: [PatBase Info VName]
params = [PatBase Info VName
closure_pat, PatBase Info VName
pat']
params_for_rettype :: [PatBase Info VName]
params_for_rettype = [PatBase Info VName]
params forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatBase Info VName]
svParams StaticVal
sv1 forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatBase Info VName]
svParams StaticVal
sv2
svParams :: StaticVal -> [PatBase Info VName]
svParams (LambdaSV PatBase Info VName
sv_pat StructRetType
_ ExtExp
_ Env
_) = [PatBase Info VName
sv_pat]
svParams StaticVal
_ = []
rettype :: PatType
rettype = Env -> [PatBase Info VName] -> StructType -> PatType -> PatType
buildRetType Env
closure_env [PatBase Info VName]
params_for_rettype (StructRetType -> StructType
unRetType StructRetType
e0_t) forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e0'
already_bound :: Set VName
already_bound =
Set VName
globals
forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims
forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
params)
more_dims :: [VName]
more_dims =
forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$
forall a. (a -> Bool) -> Set a -> Set a
S.filter (forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
already_bound) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Set VName
patternArraySizes [PatBase Info VName]
params
liftedName :: t -> ExpBase f VName -> [Char]
liftedName t
i (Var QualName VName
f f PatType
_ SrcLoc
_) =
[Char]
"defunc_" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show t
i forall a. [a] -> [a] -> [a]
++ [Char]
"_" forall a. [a] -> [a] -> [a]
++ VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
f)
liftedName t
i (AppExp (Apply ExpBase f VName
f ExpBase f VName
_ f (Diet, Maybe VName)
_ SrcLoc
_) f AppRes
_) =
t -> ExpBase f VName -> [Char]
liftedName (t
i forall a. Num a => a -> a -> a
+ t
1) ExpBase f VName
f
liftedName t
_ ExpBase f VName
_ = [Char]
"defunc"
let bound_sizes :: Set VName
bound_sizes = forall a. Ord a => [a] -> Set a
S.fromList ([VName]
dims forall a. Semigroup a => a -> a -> a
<> [VName]
more_dims) forall a. Semigroup a => a -> a -> a
<> Set VName
globals
([VName]
missing_dims, [PatBase Info VName]
params') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params
VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString forall a b. (a -> b) -> a -> b
$ forall {t} {f :: * -> *}.
(Show t, Num t) =>
t -> ExpBase f VName -> [Char]
liftedName (Int
0 :: Int) Exp
e1
VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec
VName
fname
(forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
rettype)
([VName]
dims forall a. [a] -> [a] -> [a]
++ [VName]
more_dims forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims)
[PatBase Info VName]
params'
Exp
e0'
let t1 :: StructType
t1 = forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e1'
t2 :: StructType
t2 = forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e2'
fname' :: QualName VName
fname' = forall v. v -> QualName v
qualName VName
fname
fname'' :: Exp
fname'' =
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
QualName VName
fname'
( forall a. a -> Info a
Info
( forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t2 forall a b. (a -> b) -> a -> b
$
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype
)
)
SrcLoc
loc
callret :: AppRes
callret
| forall dim as. TypeBase dim as -> Bool
orderZero PatType
ret = PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
ext
| Bool
otherwise = PatType -> [VName] -> AppRes
AppRes PatType
rettype [VName]
ext
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
( forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
fname'' Exp
e1' (forall a. a -> Info a
Info (Diet
Observe, forall a. Maybe a
Nothing)) SrcLoc
loc)
( forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$
PatType -> [VName] -> AppRes
AppRes
( forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
Unnamed StructType
t2 forall a b. (a -> b) -> a -> b
$
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype
)
[]
)
)
Exp
e2'
Info (Diet, Maybe VName)
d
SrcLoc
loc
)
(forall a. a -> Info a
Info AppRes
callret)
)
forall a. Monoid a => a
mempty,
StaticVal
sv
)
DynamicFun (Exp, StaticVal)
_ StaticVal
sv -> do
let ([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
argtypes
restype :: PatType
restype = forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype) forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall as shape. Monoid as => TypeBase shape as -> as
aliases PatType
ret
callret :: AppRes
callret = PatType -> [VName] -> AppRes
AppRes (forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes PatType
ret PatType
restype) [VName]
ext
apply_e :: Exp
apply_e = forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d SrcLoc
loc) (forall a. a -> Info a
Info AppRes
callret)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
apply_e, StaticVal
sv)
StaticVal
IntrinsicSV
| Int
depth forall a. Eq a => a -> a -> Bool
== Int
0 ->
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StructType]
argtypes
then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
e)
else do
([PatBase Info VName]
pats, Exp
body, StructRetType
tp) <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (Exp -> PatType
typeOf Exp
e') Exp
e'
Exp -> DefM (Exp, StaticVal)
defuncExp forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatBase Info VName]
pats Exp
body forall a. Maybe a
Nothing (forall a. a -> Info a
Info (forall a. Monoid a => a
mempty, StructRetType
tp)) forall a. Monoid a => a
mempty
| Bool
otherwise -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', StaticVal
IntrinsicSV)
StaticVal
_ ->
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Application of an expression\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Exp
e1
forall a. [a] -> [a] -> [a]
++ [Char]
"\nthat is neither a static lambda "
forall a. [a] -> [a] -> [a]
++ [Char]
"nor a dynamic function, but has static value:\n"
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv1
defuncApply Int
depth e :: Exp
e@(Var QualName VName
qn (Info PatType
t) SrcLoc
loc) = do
let ([StructType]
argtypes, StructType
_) = forall dim as.
TypeBase dim as -> ([TypeBase dim ()], TypeBase dim ())
unfoldFunType PatType
t
StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t) (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
case StaticVal
sv of
DynamicFun (Exp, StaticVal)
_ StaticVal
_
| StaticVal -> Int -> Bool
fullyApplied StaticVal
sv Int
depth -> do
let ([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
argtypes
forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
rettype)) SrcLoc
loc, StaticVal
sv)
| Bool
otherwise -> do
VName
fname <- forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName forall a b. (a -> b) -> a -> b
$ [Char]
"dyn_" forall a. Semigroup a => a -> a -> a
<> VName -> [Char]
baseString (forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
let ([PatBase Info VName]
pats, Exp
e0, StaticVal
sv') = [Char]
-> StaticVal -> Int -> ([PatBase Info VName], Exp, StaticVal)
liftDynFun (forall a. Pretty a => a -> [Char]
prettyString QualName VName
qn) StaticVal
sv Int
depth
([PatType]
argtypes', PatType
rettype) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv' [StructType]
argtypes
dims' :: [VName]
dims' = forall a. Monoid a => a
mempty
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let bound_sizes :: Set VName
bound_sizes = forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims' forall a. Semigroup a => a -> a -> a
<> Set VName
globals
([VName]
missing_dims, [PatBase Info VName]
pats') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
pats
VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec VName
fname (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
rettype) ([VName]
dims' forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims) [PatBase Info VName]
pats' Exp
e0
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
(forall v. v -> QualName v
qualName VName
fname)
(forall a. a -> Info a
Info (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType [PatType]
argtypes' 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 dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct PatType
rettype))
SrcLoc
loc,
StaticVal
sv'
)
StaticVal
IntrinsicSV -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, StaticVal
IntrinsicSV)
StaticVal
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (forall a. a -> Info a
Info (StaticVal -> PatType
typeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncApply Int
depth (Parens Exp
e SrcLoc
_) = Int -> Exp -> DefM (Exp, StaticVal)
defuncApply Int
depth Exp
e
defuncApply Int
_ Exp
expr = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
expr
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) Int
depth
| Int
depth forall a. Eq a => a -> a -> Bool
== Int
0 = Bool
False
| Int
depth forall a. Ord a => a -> a -> Bool
> Int
0 = StaticVal -> Int -> Bool
fullyApplied StaticVal
sv (Int
depth forall a. Num a => a -> a -> a
- Int
1)
fullyApplied StaticVal
_ Int
_ = Bool
True
liftDynFun :: String -> StaticVal -> Int -> ([Pat], Exp, StaticVal)
liftDynFun :: [Char]
-> StaticVal -> Int -> ([PatBase Info VName], Exp, StaticVal)
liftDynFun [Char]
_ (DynamicFun (Exp
e, StaticVal
sv) StaticVal
_) Int
0 = ([], Exp
e, StaticVal
sv)
liftDynFun [Char]
s (DynamicFun clsr :: (Exp, StaticVal)
clsr@(Exp
_, LambdaSV PatBase Info VName
pat StructRetType
_ ExtExp
_ Env
_) StaticVal
sv) Int
d
| Int
d forall a. Ord a => a -> a -> Bool
> Int
0 =
let ([PatBase Info VName]
pats, Exp
e', StaticVal
sv') = [Char]
-> StaticVal -> Int -> ([PatBase Info VName], Exp, StaticVal)
liftDynFun [Char]
s StaticVal
sv (Int
d forall a. Num a => a -> a -> a
- Int
1)
in (PatBase Info VName
pat forall a. a -> [a] -> [a]
: [PatBase Info VName]
pats, Exp
e', (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
clsr StaticVal
sv')
liftDynFun [Char]
s StaticVal
sv Int
d =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
s
forall a. [a] -> [a] -> [a]
++ [Char]
" Tried to lift a StaticVal "
forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take Int
100 (forall a. Show a => a -> [Char]
show StaticVal
sv)
forall a. [a] -> [a] -> [a]
++ [Char]
", but expected a dynamic function.\n"
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Int
d
envFromPat :: Pat -> Env
envFromPat :: PatBase Info VName -> Env
envFromPat PatBase Info VName
pat = case PatBase Info VName
pat of
TuplePat [PatBase Info VName]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
ps
RecordPat [(Name, PatBase Info VName)]
fs SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatBase Info VName -> Env
envFromPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, PatBase Info VName)]
fs
PatParens PatBase Info VName
p SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
PatAttr AttrInfo VName
_ PatBase Info VName
p SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
Id VName
vn (Info PatType
t) SrcLoc
_ -> forall k a. k -> a -> Map k a
M.singleton VName
vn forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic PatType
t
Wildcard Info PatType
_ SrcLoc
_ -> forall a. Monoid a => a
mempty
PatAscription PatBase Info VName
p TypeExp VName
_ SrcLoc
_ -> PatBase Info VName -> Env
envFromPat PatBase Info VName
p
PatLit {} -> forall a. Monoid a => a
mempty
PatConstr Name
_ Info PatType
_ [PatBase Info VName]
ps SrcLoc
_ -> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatBase Info VName -> Env
envFromPat [PatBase Info VName]
ps
envFromDimNames :: [VName] -> Env
envFromDimNames :: [VName] -> Env
envFromDimNames = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. [a] -> [b] -> [(a, b)]
zip (forall a. a -> [a]
repeat Binding
d)
where
d :: Binding
d = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
liftValDec :: VName -> StructRetType -> [VName] -> [Pat] -> Exp -> DefM ()
liftValDec :: VName
-> StructRetType
-> [VName]
-> [PatBase Info VName]
-> Exp
-> DefM ()
liftValDec VName
fname (RetType [VName]
ret_dims StructType
ret) [VName]
dims [PatBase Info VName]
pats Exp
body = ValBind -> DefM ()
addValBind ValBind
dec
where
dims' :: [TypeParamBase VName]
dims' = forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) [VName]
dims
bound_here :: Set VName
bound_here = forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
pats)
mkExt :: VName -> Maybe VName
mkExt VName
v
| Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ VName
v forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = forall a. a -> Maybe a
Just VName
v
mkExt VName
_ = forall a. Maybe a
Nothing
rettype_st :: StructRetType
rettype_st = forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe VName
mkExt (forall a. Set a -> [a]
S.toList (forall as. TypeBase Size as -> Set VName
freeInType StructType
ret)) forall a. [a] -> [a] -> [a]
++ [VName]
ret_dims) StructType
ret
dec :: ValBind
dec =
ValBind
{ valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = forall a. Maybe a
Nothing,
valBindName :: VName
valBindName = VName
fname,
valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = forall a. Maybe a
Nothing,
valBindRetType :: Info StructRetType
valBindRetType = forall a. a -> Info a
Info StructRetType
rettype_st,
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
dims',
valBindParams :: [PatBase Info VName]
valBindParams = [PatBase Info VName]
pats,
valBindBody :: Exp
valBindBody = Exp
body,
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
}
buildEnvPat :: [VName] -> Env -> Pat
buildEnvPat :: [VName] -> Env -> PatBase Info VName
buildEnvPat [VName]
sizes Env
env = forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat (forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (Name, PatBase Info VName)
buildField forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Env
env) forall a. Monoid a => a
mempty
where
buildField :: (VName, Binding) -> (Name, PatBase Info VName)
buildField (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
( [Char] -> Name
nameFromString (forall a. Pretty a => a -> [Char]
prettyString VName
vn),
if VName
vn forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
sizes
then forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
else forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) forall a. Monoid a => a
mempty
)
buildRetType :: Env -> [Pat] -> StructType -> PatType -> PatType
buildRetType :: Env -> [PatBase Info VName] -> StructType -> PatType -> PatType
buildRetType Env
env [PatBase Info VName]
pats = forall {t :: * -> *} {shape} {as}.
(Foldable t, Monoid (t Alias)) =>
TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb
where
bound :: Set VName
bound =
forall a. Ord a => [a] -> Set a
S.fromList (forall k a. Map k a -> [k]
M.keys Env
env) forall a. Semigroup a => a -> a -> a
<> forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map forall (f :: * -> *) vn. IdentBase f vn -> vn
identName (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
pats)
boundAsUnique :: VName -> Bool
boundAsUnique VName
v =
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall dim as. TypeBase dim as -> Bool
unique forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Info a -> a
unInfo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. IdentBase f vn -> f PatType
identType) forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== VName
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) vn. IdentBase f vn -> vn
identName) forall a b. (a -> b) -> a -> b
$
forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set (IdentBase f vn)
patIdents [PatBase Info VName]
pats
problematic :: VName -> Bool
problematic VName
v = (VName
v forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) Bool -> Bool -> Bool
&& Bool -> Bool
not (VName -> Bool
boundAsUnique VName
v)
comb :: TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb (Scalar (Record Map Name (TypeBase shape as)
fs_annot)) (Scalar (Record Map Name (TypeBase shape (t Alias))
fs_got)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb Map Name (TypeBase shape as)
fs_annot Map Name (TypeBase shape (t Alias))
fs_got
comb (Scalar (Sum Map Name [TypeBase shape as]
cs_annot)) (Scalar (Sum Map Name [TypeBase shape (t Alias)]
cs_got)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb) Map Name [TypeBase shape as]
cs_annot Map Name [TypeBase shape (t Alias)]
cs_got
comb (Scalar Arrow {}) TypeBase shape (t Alias)
t =
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend TypeBase shape (t Alias)
t
comb TypeBase shape as
got TypeBase shape (t Alias)
et =
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase shape as
got forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase shape (t Alias)
et
descend :: TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend t :: TypeBase dim (t Alias)
t@Array {}
| forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Bool
problematic forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim (t Alias)
t) = TypeBase dim (t Alias)
t forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
descend (Scalar (Record Map Name (TypeBase dim (t Alias))
t)) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend Map Name (TypeBase dim (t Alias))
t
descend TypeBase dim (t Alias)
t = TypeBase dim (t Alias)
t
typeFromSV :: StaticVal -> PatType
typeFromSV :: StaticVal -> PatType
typeFromSV (Dynamic PatType
tp) =
PatType
tp
typeFromSV (LambdaSV PatBase Info VName
_ StructRetType
_ ExtExp
_ Env
env) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ([Char] -> Name
nameFromString forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Pretty a => a -> [Char]
prettyString) (StaticVal -> PatType
typeFromSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV)) forall a b. (a -> b) -> a -> b
$
forall k a. Map k a -> [(k, a)]
M.toList Env
env
typeFromSV (RecordSV [(Name, StaticVal)]
ls) =
let ts :: [(Name, PatType)]
ts = forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> PatType
typeFromSV) [(Name, StaticVal)]
ls
in forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatType)]
ts
typeFromSV (DynamicFun (Exp
_, StaticVal
sv) StaticVal
_) =
StaticVal -> PatType
typeFromSV StaticVal
sv
typeFromSV (SumSV Name
name [StaticVal]
svs [(Name, [PatType])]
fields) =
let svs' :: [PatType]
svs' = forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> PatType
typeFromSV [StaticVal]
svs
in forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name [PatType]
svs' forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [PatType])]
fields
typeFromSV StaticVal
IntrinsicSV =
forall a. HasCallStack => [Char] -> a
error [Char]
"Tried to get the type from the static value of an intrinsic."
dynamicFunType :: StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType :: StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) (StructType
p : [StructType]
ps) =
let ([PatType]
ps', PatType
ret) = StaticVal -> [StructType] -> ([PatType], PatType)
dynamicFunType StaticVal
sv [StructType]
ps in (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
p forall a. a -> [a] -> [a]
: [PatType]
ps', PatType
ret)
dynamicFunType StaticVal
sv [StructType]
_ = ([], StaticVal -> PatType
typeFromSV StaticVal
sv)
matchPatSV :: PatBase Info VName -> StaticVal -> Env
matchPatSV :: PatBase Info VName -> StaticVal -> Env
matchPatSV (TuplePat [PatBase Info VName]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls) =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PatBase Info VName
p (Name
_, StaticVal
sv) -> PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv) [PatBase Info VName]
ps [(Name, StaticVal)]
ls
matchPatSV (RecordPat [(Name, PatBase Info VName)]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls)
| [(Name, PatBase Info VName)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, PatBase Info VName)]
ps,
[(Name, StaticVal)]
ls' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls,
forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Name, PatBase Info VName)]
ps' forall a. Eq a => a -> a -> Bool
== forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls' =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
_, PatBase Info VName
p) (Name
_, StaticVal
sv) -> PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
p StaticVal
sv) [(Name, PatBase Info VName)]
ps' [(Name, StaticVal)]
ls'
matchPatSV (PatParens PatBase Info VName
pat SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV (PatAttr AttrInfo VName
_ PatBase Info VName
pat SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV (Id VName
vn (Info PatType
t) SrcLoc
_) StaticVal
sv =
if StaticVal -> Bool
orderZeroSV StaticVal
sv
then Env
dim_env forall a. Semigroup a => a -> a -> a
<> forall k a. k -> a -> Map k a
M.singleton VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic PatType
t)
else Env
dim_env forall a. Semigroup a => a -> a -> a
<> forall k a. k -> a -> Map k a
M.singleton VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing StaticVal
sv)
where
dim_env :: Env
dim_env =
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (,Binding
i64) forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
S.toList forall a b. (a -> b) -> a -> b
$ forall as. TypeBase Size as -> Set VName
freeInType PatType
t
i64 :: Binding
i64 = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
Dynamic forall a b. (a -> b) -> a -> b
$ forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
matchPatSV (Wildcard Info PatType
_ SrcLoc
_) StaticVal
_ = forall a. Monoid a => a
mempty
matchPatSV (PatAscription PatBase Info VName
pat TypeExp VName
_ SrcLoc
_) StaticVal
sv = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat StaticVal
sv
matchPatSV PatLit {} StaticVal
_ = forall a. Monoid a => a
mempty
matchPatSV (PatConstr Name
c1 Info PatType
_ [PatBase Info VName]
ps SrcLoc
_) (SumSV Name
c2 [StaticVal]
ls [(Name, [PatType])]
fs)
| Name
c1 forall a. Eq a => a -> a -> Bool
== Name
c2 =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps [StaticVal]
ls
| Just [PatType]
ts <- forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
c1 [(Name, [PatType])]
fs =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PatType -> StaticVal
svFromType [PatType]
ts
| Bool
otherwise =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV (PatConstr Name
c1 Info PatType
_ [PatBase Info VName]
ps SrcLoc
_) (Dynamic (Scalar (Sum Map Name [PatType]
fs)))
| Just [PatType]
ts <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c1 Map Name [PatType]
fs =
forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> Env
matchPatSV [PatBase Info VName]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map PatType -> StaticVal
svFromType [PatType]
ts
| Bool
otherwise =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV PatBase Info VName
pat (Dynamic PatType
t) = PatBase Info VName -> StaticVal -> Env
matchPatSV PatBase Info VName
pat forall a b. (a -> b) -> a -> b
$ PatType -> StaticVal
svFromType PatType
t
matchPatSV PatBase Info VName
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to match pattern "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info VName
pat
forall a. [a] -> [a] -> [a]
++ [Char]
" with static value "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
forall a. [a] -> [a] -> [a]
++ [Char]
"."
orderZeroSV :: StaticVal -> Bool
orderZeroSV :: StaticVal -> Bool
orderZeroSV Dynamic {} = Bool
True
orderZeroSV (RecordSV [(Name, StaticVal)]
fields) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (StaticVal -> Bool
orderZeroSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fields
orderZeroSV StaticVal
_ = Bool
False
updatePat :: Pat -> StaticVal -> Pat
updatePat :: PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat (TuplePat [PatBase Info VName]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs) =
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat [PatBase Info VName]
ps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Name, StaticVal)]
svs) SrcLoc
loc
updatePat (RecordPat [(Name, PatBase Info VName)]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs)
| [(Name, PatBase Info VName)]
ps' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, PatBase Info VName)]
ps,
[(Name, StaticVal)]
svs' <- forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn forall a b. (a, b) -> a
fst [(Name, StaticVal)]
svs =
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat
(forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
n, PatBase Info VName
p) (Name
_, StaticVal
sv) -> (Name
n, PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
p StaticVal
sv)) [(Name, PatBase Info VName)]
ps' [(Name, StaticVal)]
svs')
SrcLoc
loc
updatePat (PatParens PatBase Info VName
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens (PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv) SrcLoc
loc
updatePat (PatAttr AttrInfo VName
attr PatBase Info VName
pat SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr (PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv) SrcLoc
loc
updatePat (Id VName
vn (Info PatType
tp) SrcLoc
loc) StaticVal
sv =
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ forall {dim} {as}.
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb PatType
tp (StaticVal -> PatType
typeFromSV StaticVal
sv forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique)) SrcLoc
loc
where
comb :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb (Scalar Arrow {}) TypeBase dim as
t2 = TypeBase dim as
t2
comb (Scalar (Record Map Name (TypeBase dim as)
m1)) (Scalar (Record Map Name (TypeBase dim as)
m2)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb Map Name (TypeBase dim as)
m1 Map Name (TypeBase dim as)
m2
comb (Scalar (Sum Map Name [TypeBase dim as]
m1)) (Scalar (Sum Map Name [TypeBase dim as]
m2)) =
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb) Map Name [TypeBase dim as]
m1 Map Name [TypeBase dim as]
m2
comb TypeBase dim as
t1 TypeBase dim as
_ = TypeBase dim as
t1
updatePat pat :: PatBase Info VName
pat@(Wildcard (Info PatType
tp) SrcLoc
loc) StaticVal
sv
| forall dim as. TypeBase dim as -> Bool
orderZero PatType
tp = PatBase Info VName
pat
| Bool
otherwise = forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ StaticVal -> PatType
typeFromSV StaticVal
sv) SrcLoc
loc
updatePat (PatAscription PatBase Info VName
pat TypeExp VName
_ SrcLoc
_) StaticVal
sv =
PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat StaticVal
sv
updatePat p :: PatBase Info VName
p@PatLit {} StaticVal
_ = PatBase Info VName
p
updatePat pat :: PatBase Info VName
pat@(PatConstr Name
c1 (Info PatType
t) [PatBase Info VName]
ps SrcLoc
loc) sv :: StaticVal
sv@(SumSV Name
_ [StaticVal]
svs [(Name, [PatType])]
_)
| forall dim as. TypeBase dim as -> Bool
orderZero PatType
t = PatBase Info VName
pat
| Bool
otherwise = forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
c1 (forall a. a -> Info a
Info PatType
t') [PatBase Info VName]
ps' SrcLoc
loc
where
t' :: PatType
t' = StaticVal -> PatType
typeFromSV StaticVal
sv forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
ps' :: [PatBase Info VName]
ps' = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat [PatBase Info VName]
ps [StaticVal]
svs
updatePat (PatConstr Name
c1 Info PatType
_ [PatBase Info VName]
ps SrcLoc
loc) (Dynamic PatType
t) =
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
c1 (forall a. a -> Info a
Info PatType
t) [PatBase Info VName]
ps SrcLoc
loc
updatePat PatBase Info VName
pat (Dynamic PatType
t) = PatBase Info VName -> StaticVal -> PatBase Info VName
updatePat PatBase Info VName
pat (PatType -> StaticVal
svFromType PatType
t)
updatePat PatBase Info VName
pat StaticVal
sv =
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
[Char]
"Tried to update pattern "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString PatBase Info VName
pat
forall a. [a] -> [a] -> [a]
++ [Char]
"to reflect the static value "
forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show StaticVal
sv
svFromType :: PatType -> StaticVal
svFromType :: PatType -> StaticVal
svFromType (Scalar (Record Map Name PatType
fs)) = [(Name, StaticVal)] -> StaticVal
RecordSV forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
M.toList forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map PatType -> StaticVal
svFromType Map Name PatType
fs
svFromType PatType
t = PatType -> StaticVal
Dynamic PatType
t
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp VName)
_ (Info (RetType [VName]
_ StructType
rettype)) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc)
| Scalar Arrow {} <- StructType
rettype = do
([PatBase Info VName]
body_pats, Exp
body', StructRetType
rettype') <- PatType -> Exp -> DefM ([PatBase Info VName], Exp, StructRetType)
etaExpand (forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
rettype) Exp
body
ValBind -> DefM (ValBind, Env)
defuncValBind forall a b. (a -> b) -> a -> b
$
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f StructRetType
-> [TypeParamBase vn]
-> [PatBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo vn]
-> SrcLoc
-> ValBindBase f vn
ValBind
Maybe (Info EntryPoint)
entry
VName
name
forall a. Maybe a
Nothing
(forall a. a -> Info a
Info StructRetType
rettype')
[TypeParamBase VName]
tparams
([PatBase Info VName]
params forall a. Semigroup a => a -> a -> a
<> [PatBase Info VName]
body_pats)
Exp
body'
forall a. Maybe a
Nothing
[AttrInfo VName]
attrs
SrcLoc
loc
defuncValBind valbind :: ValBind
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
retdecl (Info (RetType [VName]
ret_dims StructType
rettype)) [TypeParamBase VName]
tparams [PatBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_) = do
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any forall vn. TypeParamBase vn -> Bool
isTypeParam [TypeParamBase VName]
tparams) forall a b. (a -> b) -> a -> b
$
forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
forall a. Show a => a -> [Char]
show VName
name
forall a. [a] -> [a] -> [a]
++ [Char]
" has type parameters, "
forall a. [a] -> [a] -> [a]
++ [Char]
"but the defunctionaliser expects a monomorphic input program."
([VName]
tparams', [PatBase Info VName]
params', Exp
body', StaticVal
sv) <-
[VName]
-> [PatBase Info VName]
-> Exp
-> StructRetType
-> DefM ([VName], [PatBase Info VName], Exp, StaticVal)
defuncLet (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams) [PatBase Info VName]
params Exp
body forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims StructType
rettype
Set VName
globals <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a, b) -> a
fst
let bound_sizes :: Set VName
bound_sizes = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatBase f vn -> Set vn
patNames [PatBase Info VName]
params' forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams' forall a. Semigroup a => a -> a -> a
<> Set VName
globals
rettype' :: StructType
rettype' =
forall as.
Monoid as =>
TypeBase Size as -> TypeBase Size as -> TypeBase Size as
combineTypeShapes StructType
rettype forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Set VName -> Size -> Size
anyDimIfNotBound Set VName
bound_sizes) forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct forall a b. (a -> b) -> a -> b
$ Exp -> PatType
typeOf Exp
body'
ret_dims' :: [VName]
ret_dims' = forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Ord a => a -> Set a -> Bool
`S.member` forall as. TypeBase Size as -> Set VName
freeInType StructType
rettype') [VName]
ret_dims
([VName]
missing_dims, [PatBase Info VName]
params'') <- forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatBase Info VName] -> m ([VName], [PatBase Info VName])
sizesForAll Set VName
bound_sizes [PatBase Info VName]
params'
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( ValBind
valbind
{ valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
retdecl,
valBindRetType :: Info StructRetType
valBindRetType =
forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$
if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatBase Info VName]
params'
then forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' forall a b. (a -> b) -> a -> b
$ StructType
rettype' forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
else forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ret_dims' StructType
rettype',
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams =
forall a b. (a -> b) -> [a] -> [b]
map (forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` forall a. Monoid a => a
mempty) forall a b. (a -> b) -> a -> b
$ [VName]
tparams' forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims,
valBindParams :: [PatBase Info VName]
valBindParams = [PatBase Info VName]
params'',
valBindBody :: Exp
valBindBody = Exp
body'
},
forall k a. k -> a -> Map k a
M.singleton VName
name forall a b. (a -> b) -> a -> b
$
Maybe ([VName], StructType) -> StaticVal -> Binding
Binding
(forall a. a -> Maybe a
Just (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a b. (a -> b) -> [a] -> [b]
map forall vn. TypeParamBase vn -> vn
typeParamName) (ValBind -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBind
valbind)))
StaticVal
sv
)
where
anyDimIfNotBound :: Set VName -> Size -> Size
anyDimIfNotBound Set VName
bound_sizes (NamedSize QualName VName
v)
| forall vn. QualName vn -> vn
qualLeaf QualName VName
v forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
bound_sizes = Maybe VName -> Size
AnySize forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
v
anyDimIfNotBound Set VName
_ Size
d = Size
d
defuncVals :: [ValBind] -> DefM ()
defuncVals :: [ValBind] -> DefM ()
defuncVals [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
defuncVals (ValBind
valbind : [ValBind]
ds) = do
(ValBind
valbind', Env
env) <- ValBind -> DefM (ValBind, Env)
defuncValBind ValBind
valbind
ValBind -> DefM ()
addValBind ValBind
valbind'
let globals :: [VName]
globals = ValBind -> [VName]
valBindBound ValBind
valbind'
forall a. Env -> DefM a -> DefM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ forall a. [VName] -> DefM a -> DefM a
areGlobal [VName]
globals forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
defuncVals [ValBind]
ds
{-# NOINLINE transformProg #-}
transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
decs = forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
let ((), VNameSource
namesrc', [ValBind]
decs') = forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
namesrc forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
defuncVals [ValBind]
decs
in ([ValBind]
decs', VNameSource
namesrc')