{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE Trustworthy #-}
module Futhark.Internalise.Monomorphise (transformProg) where
import Control.Monad.Identity
import Control.Monad.RWS hiding (Sum)
import Control.Monad.State
import Control.Monad.Writer hiding (Sum)
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Sequence as Seq
import qualified Data.Set as S
import Futhark.MonadFreshNames
import Futhark.Util.Pretty
import Language.Futhark
import Language.Futhark.Semantic (TypeBinding (..))
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types
i64 :: TypeBase dim als
i64 :: TypeBase dim als
i64 = ScalarTypeBase dim als -> TypeBase dim als
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim als -> TypeBase dim als)
-> ScalarTypeBase dim als -> TypeBase dim als
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim als
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase dim als)
-> PrimType -> ScalarTypeBase dim als
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
data PolyBinding
= PolyBinding
RecordReplacements
( VName,
[TypeParam],
[Pat],
StructRetType,
[VName],
Exp,
[AttrInfo VName],
SrcLoc
)
type RecordReplacements = M.Map VName RecordReplacement
type RecordReplacement = M.Map Name (VName, PatType)
data Env = Env
{ Env -> Map VName PolyBinding
envPolyBindings :: M.Map VName PolyBinding,
Env -> Map VName TypeBinding
envTypeBindings :: M.Map VName TypeBinding,
Env -> RecordReplacements
envRecordReplacements :: RecordReplacements
}
instance Semigroup Env where
Env Map VName PolyBinding
tb1 Map VName TypeBinding
pb1 RecordReplacements
rr1 <> :: Env -> Env -> Env
<> Env Map VName PolyBinding
tb2 Map VName TypeBinding
pb2 RecordReplacements
rr2 = Map VName PolyBinding
-> Map VName TypeBinding -> RecordReplacements -> Env
Env (Map VName PolyBinding
tb1 Map VName PolyBinding
-> Map VName PolyBinding -> Map VName PolyBinding
forall a. Semigroup a => a -> a -> a
<> Map VName PolyBinding
tb2) (Map VName TypeBinding
pb1 Map VName TypeBinding
-> Map VName TypeBinding -> Map VName TypeBinding
forall a. Semigroup a => a -> a -> a
<> Map VName TypeBinding
pb2) (RecordReplacements
rr1 RecordReplacements -> RecordReplacements -> RecordReplacements
forall a. Semigroup a => a -> a -> a
<> RecordReplacements
rr2)
instance Monoid Env where
mempty :: Env
mempty = Map VName PolyBinding
-> Map VName TypeBinding -> RecordReplacements -> Env
Env Map VName PolyBinding
forall a. Monoid a => a
mempty Map VName TypeBinding
forall a. Monoid a => a
mempty RecordReplacements
forall a. Monoid a => a
mempty
localEnv :: Env -> MonoM a -> MonoM a
localEnv :: Env -> MonoM a -> MonoM a
localEnv Env
env = (Env -> Env) -> MonoM a -> MonoM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>)
extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
vn PolyBinding
binding =
Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv
Env
forall a. Monoid a => a
mempty {envPolyBindings :: Map VName PolyBinding
envPolyBindings = VName -> PolyBinding -> Map VName PolyBinding
forall k a. k -> a -> Map k a
M.singleton VName
vn PolyBinding
binding}
withRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr = Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv Env
forall a. Monoid a => a
mempty {envRecordReplacements :: RecordReplacements
envRecordReplacements = RecordReplacements
rr}
replaceRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr = (Env -> Env) -> MonoM a -> MonoM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> MonoM a -> MonoM a)
-> (Env -> Env) -> MonoM a -> MonoM a
forall a b. (a -> b) -> a -> b
$ \Env
env -> Env
env {envRecordReplacements :: RecordReplacements
envRecordReplacements = RecordReplacements
rr}
newtype MonoM a
= MonoM
( RWST
Env
(Seq.Seq (VName, ValBind))
VNameSource
(State Lifts)
a
)
deriving
( a -> MonoM b -> MonoM a
(a -> b) -> MonoM a -> MonoM b
(forall a b. (a -> b) -> MonoM a -> MonoM b)
-> (forall a b. a -> MonoM b -> MonoM a) -> Functor MonoM
forall a b. a -> MonoM b -> MonoM a
forall a b. (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> MonoM b -> MonoM a
$c<$ :: forall a b. a -> MonoM b -> MonoM a
fmap :: (a -> b) -> MonoM a -> MonoM b
$cfmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
Functor,
Functor MonoM
a -> MonoM a
Functor MonoM
-> (forall a. a -> MonoM a)
-> (forall a b. MonoM (a -> b) -> MonoM a -> MonoM b)
-> (forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c)
-> (forall a b. MonoM a -> MonoM b -> MonoM b)
-> (forall a b. MonoM a -> MonoM b -> MonoM a)
-> Applicative MonoM
MonoM a -> MonoM b -> MonoM b
MonoM a -> MonoM b -> MonoM a
MonoM (a -> b) -> MonoM a -> MonoM b
(a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
forall a. a -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM b
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM 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
<* :: MonoM a -> MonoM b -> MonoM a
$c<* :: forall a b. MonoM a -> MonoM b -> MonoM a
*> :: MonoM a -> MonoM b -> MonoM b
$c*> :: forall a b. MonoM a -> MonoM b -> MonoM b
liftA2 :: (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
$cliftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
<*> :: MonoM (a -> b) -> MonoM a -> MonoM b
$c<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
pure :: a -> MonoM a
$cpure :: forall a. a -> MonoM a
$cp1Applicative :: Functor MonoM
Applicative,
Applicative MonoM
a -> MonoM a
Applicative MonoM
-> (forall a b. MonoM a -> (a -> MonoM b) -> MonoM b)
-> (forall a b. MonoM a -> MonoM b -> MonoM b)
-> (forall a. a -> MonoM a)
-> Monad MonoM
MonoM a -> (a -> MonoM b) -> MonoM b
MonoM a -> MonoM b -> MonoM b
forall a. a -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM b
forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> MonoM a
$creturn :: forall a. a -> MonoM a
>> :: MonoM a -> MonoM b -> MonoM b
$c>> :: forall a b. MonoM a -> MonoM b -> MonoM b
>>= :: MonoM a -> (a -> MonoM b) -> MonoM b
$c>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
$cp1Monad :: Applicative MonoM
Monad,
MonadReader Env,
MonadWriter (Seq.Seq (VName, ValBind)),
Monad MonoM
Applicative MonoM
MonoM VNameSource
Applicative MonoM
-> Monad MonoM
-> MonoM VNameSource
-> (VNameSource -> MonoM ())
-> MonadFreshNames MonoM
VNameSource -> MonoM ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> MonoM ()
$cputNameSource :: VNameSource -> MonoM ()
getNameSource :: MonoM VNameSource
$cgetNameSource :: MonoM VNameSource
$cp2MonadFreshNames :: Monad MonoM
$cp1MonadFreshNames :: Applicative MonoM
MonadFreshNames
)
runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource)
runMonoM :: VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
runMonoM VNameSource
src (MonoM RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
m) = ((a
a, Seq (VName, ValBind)
defs), VNameSource
src')
where
(a
a, VNameSource
src', Seq (VName, ValBind)
defs) = State Lifts (a, VNameSource, Seq (VName, ValBind))
-> Lifts -> (a, VNameSource, Seq (VName, ValBind))
forall s a. State s a -> s -> a
evalState (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> Env
-> VNameSource
-> State Lifts (a, VNameSource, Seq (VName, ValBind))
forall r w s (m :: * -> *) a.
RWST r w s m a -> r -> s -> m (a, s, w)
runRWST RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
m Env
forall a. Monoid a => a
mempty VNameSource
src) Lifts
forall a. Monoid a => a
mempty
lookupFun :: VName -> MonoM (Maybe PolyBinding)
lookupFun :: VName -> MonoM (Maybe PolyBinding)
lookupFun VName
vn = do
Map VName PolyBinding
env <- (Env -> Map VName PolyBinding) -> MonoM (Map VName PolyBinding)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Map VName PolyBinding
envPolyBindings
case VName -> Map VName PolyBinding -> Maybe PolyBinding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
vn Map VName PolyBinding
env of
Just PolyBinding
valbind -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe PolyBinding -> MonoM (Maybe PolyBinding))
-> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ PolyBinding -> Maybe PolyBinding
forall a. a -> Maybe a
Just PolyBinding
valbind
Maybe PolyBinding
Nothing -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe PolyBinding
forall a. Maybe a
Nothing
lookupRecordReplacement :: VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement :: VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement VName
v = (Env -> Maybe RecordReplacement) -> MonoM (Maybe RecordReplacement)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Maybe RecordReplacement)
-> MonoM (Maybe RecordReplacement))
-> (Env -> Maybe RecordReplacement)
-> MonoM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (RecordReplacements -> Maybe RecordReplacement)
-> (Env -> RecordReplacements) -> Env -> Maybe RecordReplacement
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> RecordReplacements
envRecordReplacements
type InferSizeArgs = StructType -> [Exp]
data MonoSize
=
MonoKnown Int
| MonoAnon VName
deriving (Int -> MonoSize -> ShowS
[MonoSize] -> ShowS
MonoSize -> String
(Int -> MonoSize -> ShowS)
-> (MonoSize -> String) -> ([MonoSize] -> ShowS) -> Show MonoSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MonoSize] -> ShowS
$cshowList :: [MonoSize] -> ShowS
show :: MonoSize -> String
$cshow :: MonoSize -> String
showsPrec :: Int -> MonoSize -> ShowS
$cshowsPrec :: Int -> MonoSize -> ShowS
Show)
instance Eq MonoSize where
MonoKnown Int
x == :: MonoSize -> MonoSize -> Bool
== MonoKnown Int
y = Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y
MonoAnon VName
_ == MonoAnon VName
_ = Bool
True
MonoSize
_ == MonoSize
_ = Bool
False
instance Pretty MonoSize where
ppr :: MonoSize -> Doc
ppr (MonoKnown Int
i) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
i
ppr (MonoAnon VName
v) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
v
instance Pretty (ShapeDecl MonoSize) where
ppr :: ShapeDecl MonoSize -> Doc
ppr (ShapeDecl [MonoSize]
ds) = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((MonoSize -> Doc) -> [MonoSize] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (MonoSize -> Doc) -> MonoSize -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MonoSize -> Doc
forall a. Pretty a => a -> Doc
ppr) [MonoSize]
ds)
type MonoType = TypeBase MonoSize ()
monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType = (State (Int, Map (DimDecl VName) Int) MonoType
-> (Int, Map (DimDecl VName) Int) -> MonoType
forall s a. State s a -> s -> a
`evalState` (Int
0, Map (DimDecl VName) Int
forall a. Monoid a => a
mempty)) (State (Int, Map (DimDecl VName) Int) MonoType -> MonoType)
-> (TypeBase (DimDecl VName) als
-> State (Int, Map (DimDecl VName) Int) MonoType)
-> TypeBase (DimDecl VName) als
-> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
-> DimPos
-> DimDecl VName
-> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize)
-> TypeBase (DimDecl VName) ()
-> State (Int, Map (DimDecl VName) Int) MonoType
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos
-> DimDecl VName
-> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize
forall (f :: * -> *) p.
MonadState (Int, Map (DimDecl VName) Int) f =>
Set VName -> p -> DimDecl VName -> f MonoSize
onDim (TypeBase (DimDecl VName) ()
-> State (Int, Map (DimDecl VName) Int) MonoType)
-> (TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) als
-> State (Int, Map (DimDecl VName) Int) MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct
where
onDim :: Set VName -> p -> DimDecl VName -> f MonoSize
onDim Set VName
bound p
_ (NamedDim QualName VName
d)
| QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ VName -> MonoSize
MonoAnon (VName -> MonoSize) -> VName -> MonoSize
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
onDim Set VName
_ p
_ DimDecl VName
d = do
(Int
i, Map (DimDecl VName) Int
m) <- f (Int, Map (DimDecl VName) Int)
forall s (m :: * -> *). MonadState s m => m s
get
case DimDecl VName -> Map (DimDecl VName) Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup DimDecl VName
d Map (DimDecl VName) Int
m of
Just Int
prev ->
MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
prev
Maybe Int
Nothing -> do
(Int, Map (DimDecl VName) Int) -> f ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, DimDecl VName
-> Int -> Map (DimDecl VName) Int -> Map (DimDecl VName) Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert DimDecl VName
d Int
i Map (DimDecl VName) Int
m)
MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
i
type Lifts = [((VName, MonoType), (VName, InferSizeArgs))]
getLifts :: MonoM Lifts
getLifts :: MonoM Lifts
getLifts = RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
-> MonoM Lifts
forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
-> MonoM Lifts)
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
-> MonoM Lifts
forall a b. (a -> b) -> a -> b
$ State Lifts Lifts
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift State Lifts Lifts
forall s (m :: * -> *). MonadState s m => m s
get
modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts = RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
-> MonoM ()
forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
-> MonoM ())
-> ((Lifts -> Lifts)
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ())
-> (Lifts -> Lifts)
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. State Lifts ()
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (State Lifts ()
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ())
-> ((Lifts -> Lifts) -> State Lifts ())
-> (Lifts -> Lifts)
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Lifts -> Lifts) -> State Lifts ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify
addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted VName
fname MonoType
il (VName, InferSizeArgs)
liftf =
(Lifts -> Lifts) -> MonoM ()
modifyLifts (((VName
fname, MonoType
il), (VName, InferSizeArgs)
liftf) ((VName, MonoType), (VName, InferSizeArgs)) -> Lifts -> Lifts
forall a. a -> [a] -> [a]
:)
lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted VName
fname MonoType
t = (VName, MonoType) -> Lifts -> Maybe (VName, InferSizeArgs)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (VName
fname, MonoType
t) (Lifts -> Maybe (VName, InferSizeArgs))
-> MonoM Lifts -> MonoM (Maybe (VName, InferSizeArgs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM Lifts
getLifts
transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName :: SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname TypeBase (DimDecl VName) ()
t
| VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Exp
forall vn. QualName vn -> ExpBase Info vn
var QualName VName
fname
| Bool
otherwise = do
TypeBase (DimDecl VName) ()
t' <- TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t
let mono_t :: MonoType
mono_t = TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t'
Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t
Maybe PolyBinding
maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
lookupFun (VName -> MonoM (Maybe PolyBinding))
-> VName -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
case (Maybe (VName, InferSizeArgs)
maybe_fname, Maybe PolyBinding
maybe_funbind) of
(Just (VName
fname', InferSizeArgs
infer), Maybe PolyBinding
_) ->
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase (DimDecl VName) () -> [Exp] -> Exp
forall vn as.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
(Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Exp
forall vn. QualName vn -> ExpBase Info vn
var QualName VName
fname
(Maybe (VName, InferSizeArgs)
Nothing, Just PolyBinding
funbind) -> do
(VName
fname', InferSizeArgs
infer, ValBind
funbind') <- Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
False PolyBinding
funbind MonoType
mono_t
Seq (VName, ValBind) -> MonoM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (VName, ValBind) -> MonoM ())
-> Seq (VName, ValBind) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (VName, ValBind) -> Seq (VName, ValBind)
forall a. a -> Seq a
Seq.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname, ValBind
funbind')
VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t (VName
fname', InferSizeArgs
infer)
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase (DimDecl VName) () -> [Exp] -> Exp
forall vn as.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
where
var :: QualName vn -> ExpBase Info vn
var QualName vn
fname' = QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (PatType -> Info PatType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t)) SrcLoc
loc
applySizeArg :: (Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg (Int
i, ExpBase Info vn
f) ExpBase Info vn
size_arg =
( Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,
AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
f ExpBase Info vn
size_arg ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing)) SrcLoc
loc)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes ([PatType] -> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall as dim.
Monoid as =>
[TypeBase dim as] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (Int -> PatType -> [PatType]
forall a. Int -> a -> [a]
replicate Int
i PatType
forall dim als. TypeBase dim als
i64) ([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t))) [])
)
applySizeArgs :: vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs vn
fname' TypeBase (DimDecl VName) as
t' [ExpBase Info vn]
size_args =
(Int, ExpBase Info vn) -> ExpBase Info vn
forall a b. (a, b) -> b
snd ((Int, ExpBase Info vn) -> ExpBase Info vn)
-> (Int, ExpBase Info vn) -> ExpBase Info vn
forall a b. (a -> b) -> a -> b
$
((Int, ExpBase Info vn)
-> ExpBase Info vn -> (Int, ExpBase Info vn))
-> (Int, ExpBase Info vn)
-> [ExpBase Info vn]
-> (Int, ExpBase Info vn)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
(Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
forall vn.
(Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg
( [ExpBase Info vn] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info vn]
size_args Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1,
QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
(vn -> QualName vn
forall v. v -> QualName v
qualName vn
fname')
( PatType -> Info PatType
forall a. a -> Info a
Info
( [PatType] -> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall as dim.
Monoid as =>
[TypeBase dim as] -> RetTypeBase dim as -> TypeBase dim as
foldFunType
((ExpBase Info vn -> PatType) -> [ExpBase Info vn] -> [PatType]
forall a b. (a -> b) -> [a] -> [b]
map (PatType -> ExpBase Info vn -> PatType
forall a b. a -> b -> a
const PatType
forall dim als. TypeBase dim als
i64) [ExpBase Info vn]
size_args)
([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (PatType -> RetTypeBase (DimDecl VName) Aliasing)
-> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) as -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) as
t')
)
)
SrcLoc
loc
)
[ExpBase Info vn]
size_args
transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType TypeBase dim Aliasing
t = do
RecordReplacements
rrs <- (Env -> RecordReplacements) -> MonoM RecordReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> RecordReplacements
envRecordReplacements
let replace :: Alias -> Aliasing
replace (AliasBound VName
v)
| Just RecordReplacement
d <- VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v RecordReplacements
rrs =
[Alias] -> Aliasing
forall a. Ord a => [a] -> Set a
S.fromList ([Alias] -> Aliasing) -> [Alias] -> Aliasing
forall a b. (a -> b) -> a -> b
$ ((VName, PatType) -> Alias) -> [(VName, PatType)] -> [Alias]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound (VName -> Alias)
-> ((VName, PatType) -> VName) -> (VName, PatType) -> Alias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, PatType) -> VName
forall a b. (a, b) -> a
fst) ([(VName, PatType)] -> [Alias]) -> [(VName, PatType)] -> [Alias]
forall a b. (a -> b) -> a -> b
$ RecordReplacement -> [(VName, PatType)]
forall k a. Map k a -> [a]
M.elems RecordReplacement
d
replace Alias
x = Alias -> Aliasing
forall a. a -> Set a
S.singleton Alias
x
TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing))
-> TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall a b. (a -> b) -> a -> b
$
if (Alias -> Bool) -> Aliasing -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> RecordReplacements -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` RecordReplacements
rrs) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (Aliasing -> Bool) -> Aliasing -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t
then (Aliasing -> Aliasing)
-> TypeBase dim Aliasing -> TypeBase dim Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([Aliasing] -> Aliasing
forall a. Monoid a => [a] -> a
mconcat ([Aliasing] -> Aliasing)
-> (Aliasing -> [Aliasing]) -> Aliasing -> Aliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Aliasing) -> [Alias] -> [Aliasing]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> Aliasing
replace ([Alias] -> [Aliasing])
-> (Aliasing -> [Alias]) -> Aliasing -> [Aliasing]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliasing -> [Alias]
forall a. Set a -> [a]
S.toList) TypeBase dim Aliasing
t
else TypeBase dim Aliasing
t
sizesForPat :: MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat :: Pat -> m ([VName], Pat)
sizesForPat Pat
pat = do
(Pat
params', [VName]
sizes) <- StateT [VName] m Pat -> [VName] -> m (Pat, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ASTMapper (StateT [VName] m) -> Pat -> StateT [VName] m Pat
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT [VName] m)
tv Pat
pat) []
([VName], Pat) -> m ([VName], Pat)
forall (m :: * -> *) a. Monad m => a -> m a
return ([VName]
sizes, Pat
params')
where
tv :: ASTMapper (StateT [VName] m)
tv = ASTMapper (StateT [VName] m)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatType :: PatType -> StateT [VName] m PatType
mapOnPatType = (DimDecl VName -> StateT [VName] m (DimDecl VName))
-> (Aliasing -> StateT [VName] m Aliasing)
-> PatType
-> StateT [VName] m PatType
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 DimDecl VName -> StateT [VName] m (DimDecl VName)
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
DimDecl VName -> t m (DimDecl VName)
onDim Aliasing -> StateT [VName] m Aliasing
forall (f :: * -> *) a. Applicative f => a -> f a
pure}
onDim :: DimDecl VName -> t m (DimDecl VName)
onDim (AnyDim Maybe VName
_) = do
VName
v <- m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"size"
([VName] -> [VName]) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:)
DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
onDim DimDecl VName
d = DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d
transformAppRes :: AppRes -> MonoM AppRes
transformAppRes :: AppRes -> MonoM AppRes
transformAppRes (AppRes PatType
t [VName]
ext) =
PatType -> [VName] -> AppRes
AppRes (PatType -> [VName] -> AppRes)
-> MonoM PatType -> MonoM ([VName] -> AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t MonoM ([VName] -> AppRes) -> MonoM [VName] -> MonoM AppRes
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [VName] -> MonoM [VName]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName]
ext
transformAppExp :: AppExp -> AppRes -> MonoM Exp
transformAppExp :: AppExp -> AppRes -> MonoM Exp
transformAppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) AppRes
res = do
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Maybe Exp
me' <- (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Maybe Exp
me
Inclusiveness Exp
incl' <- (Exp -> MonoM Exp)
-> Inclusiveness Exp -> MonoM (Inclusiveness Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Inclusiveness Exp
incl
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> Maybe Exp -> Inclusiveness Exp -> SrcLoc -> 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) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Coerce Exp
e TypeDeclBase Info VName
tp SrcLoc
loc) AppRes
res =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> TypeDeclBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeDeclBase f vn -> SrcLoc -> AppExpBase f vn
Coerce (Exp -> TypeDeclBase Info VName -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (TypeDeclBase Info VName -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeDeclBase Info VName -> SrcLoc -> AppExp)
-> MonoM (TypeDeclBase Info VName) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeDeclBase Info VName -> MonoM (TypeDeclBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeDeclBase Info VName
tp MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (LetPat [SizeBinder VName]
sizes Pat
pat Exp
e1 Exp
e2 SrcLoc
loc) AppRes
res = do
(Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes Pat
pat' (Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
MonoM (Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (Exp -> MonoM Exp
transformExp Exp
e2)
MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
)
MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat]
params, Maybe (TypeExp VName)
retdecl, Info StructRetType
ret, Exp
body) Exp
e SrcLoc
loc) AppRes
res
| Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParamBase VName]
tparams = do
RecordReplacements
rr <- (Env -> RecordReplacements) -> MonoM RecordReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> RecordReplacements
envRecordReplacements
let funbind :: PolyBinding
funbind = RecordReplacements
-> (VName, [TypeParamBase VName], [Pat], StructRetType, [VName],
Exp, [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
ret, [], Exp
body, [AttrInfo VName]
forall a. Monoid a => a
mempty, SrcLoc
loc)
MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp)
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall a b. (a -> b) -> a -> b
$ do
(Exp
e', Seq (VName, ValBind)
bs) <- MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (MonoM Exp -> MonoM (Exp, Seq (VName, ValBind)))
-> MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ VName -> PolyBinding -> MonoM Exp -> MonoM Exp
forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
fname PolyBinding
funbind (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
e
(Lifts -> Lifts) -> MonoM ()
modifyLifts ((Lifts -> Lifts) -> MonoM ()) -> (Lifts -> Lifts) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (((VName, MonoType), (VName, InferSizeArgs)) -> Bool)
-> Lifts -> Lifts
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
fname) (VName -> Bool)
-> (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
-> ((VName, MonoType), (VName, InferSizeArgs))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, MonoType) -> VName
forall a b. (a, b) -> a
fst ((VName, MonoType) -> VName)
-> (((VName, MonoType), (VName, InferSizeArgs))
-> (VName, MonoType))
-> ((VName, MonoType), (VName, InferSizeArgs))
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, MonoType), (VName, InferSizeArgs)) -> (VName, MonoType)
forall a b. (a, b) -> a
fst)
let (Seq (VName, ValBind)
bs_local, Seq (VName, ValBind)
bs_prop) = ((VName, ValBind) -> Bool)
-> Seq (VName, ValBind)
-> (Seq (VName, ValBind), Seq (VName, ValBind))
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
fname) (VName -> Bool)
-> ((VName, ValBind) -> VName) -> (VName, ValBind) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, ValBind) -> VName
forall a b. (a, b) -> a
fst) Seq (VName, ValBind)
bs
(Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
forall (m :: * -> *) a. Monad m => a -> m a
return ([ValBind] -> Exp -> Exp
unfoldLetFuns (((VName, ValBind) -> ValBind) -> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> [a] -> [b]
map (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd ([(VName, ValBind)] -> [ValBind])
-> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> a -> b
$ Seq (VName, ValBind) -> [(VName, ValBind)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq (VName, ValBind)
bs_local) Exp
e', Seq (VName, ValBind)
-> Seq (VName, ValBind) -> Seq (VName, ValBind)
forall a b. a -> b -> a
const Seq (VName, ValBind)
bs_prop)
| Bool
otherwise = do
Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName
-> ([TypeParamBase VName], [Pat], Maybe (TypeExp VName),
Info StructRetType, Exp)
-> Exp
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn], Maybe (TypeExp vn),
f StructRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat]
params, Maybe (TypeExp VName)
retdecl, StructRetType -> Info StructRetType
forall a. a -> Info a
Info StructRetType
ret, Exp
body') (Exp -> SrcLoc -> AppExp) -> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) AppRes
res =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If (Exp -> Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e3 MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d SrcLoc
loc) AppRes
res =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply (Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM (Info (Diet, Maybe VName)) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info (Diet, Maybe VName) -> MonoM (Info (Diet, Maybe VName))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info (Diet, Maybe VName)
d MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (DoLoop [VName]
sparams Pat
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 SrcLoc
loc) AppRes
res = do
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
LoopFormBase Info VName
form' <- case LoopFormBase Info VName
form of
For IdentBase Info VName
ident Exp
e2 -> IdentBase Info VName -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
ident (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
ForIn Pat
pat2 Exp
e2 -> Pat -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn Pat
pat2 (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
While Exp
e2 -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
Exp
e3' <- Exp -> MonoM Exp
transformExp Exp
e3
([VName]
pat_sizes, Pat
pat') <- Pat -> MonoM ([VName], Pat)
forall (m :: * -> *). MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat Pat
pat
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([VName]
-> Pat -> Exp -> LoopFormBase Info VName -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[VName]
-> PatBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop ([VName]
sparams [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
pat_sizes) Pat
pat' Exp
e1' LoopFormBase Info VName
form' Exp
e3' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (BinOp (QualName VName
fname, SrcLoc
_) (Info PatType
t) (Exp
e1, Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2, Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) SrcLoc
loc) (AppRes PatType
ret [VName]
ext) = do
Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
if PatType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e1') Bool -> Bool -> Bool
&& PatType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e2')
then Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
e1' Exp
e2'
else do
(Exp
x_param_e, Pat
x_param) <- Exp -> MonoM (Exp, Pat)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pat)
makeVarParam Exp
e1'
(Exp
y_param_e, Pat
y_param) <- Exp -> MonoM (Exp, Pat)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pat)
makeVarParam Exp
e2'
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( [SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
[]
Pat
x_param
Exp
e1'
( AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
([SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] Pat
y_param Exp
e2' (Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
x_param_e Exp
y_param_e) SrcLoc
loc)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty)
)
SrcLoc
forall a. Monoid a => a
mempty
)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty))
where
applyOp :: ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp ExpBase Info vn
fname' ExpBase Info vn
x ExpBase Info vn
y =
AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
( AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
fname' ExpBase Info vn
x ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d1))) SrcLoc
loc)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty)
)
ExpBase Info vn
y
((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d2)))
SrcLoc
loc
)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
ext))
makeVarParam :: Exp -> m (Exp, Pat)
makeVarParam Exp
arg = do
let argtype :: PatType
argtype = Exp -> PatType
typeOf Exp
arg
VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"binop_p"
(Exp, Pat) -> m (Exp, Pat)
forall (m :: * -> *) a. Monad m => a -> m a
return
( QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty
)
transformAppExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs Exp
e1 Exp
body SrcLoc
loc) AppRes
res = do
SliceBase Info VName
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (IdentBase Info VName
-> IdentBase Info VName
-> SliceBase Info VName
-> Exp
-> Exp
-> SrcLoc
-> 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) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) AppRes
res =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SliceBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index (Exp -> SliceBase Info VName -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SliceBase Info VName -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e0 MonoM (SliceBase Info VName -> SrcLoc -> AppExp)
-> MonoM (SliceBase Info VName) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) AppRes
res =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
(AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match (Exp -> NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
-> MonoM (NonEmpty (CaseBase Info VName))
-> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (CaseBase Info VName -> MonoM (CaseBase Info VName))
-> NonEmpty (CaseBase Info VName)
-> MonoM (NonEmpty (CaseBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase NonEmpty (CaseBase Info VName)
cs MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformExp :: Exp -> MonoM Exp
transformExp :: Exp -> MonoM Exp
transformExp e :: Exp
e@Literal {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@IntLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@FloatLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@StringLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp (Parens Exp
e SrcLoc
loc) =
Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) =
(QualName VName, SrcLoc) -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (TupLit [Exp]
es SrcLoc
loc) =
[Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp) -> MonoM [Exp] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
[FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> [FieldBase Info VName] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField [FieldBase Info VName]
fs MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
where
transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (RecordFieldExplicit Name
name Exp
e SrcLoc
loc') =
Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name (Exp -> SrcLoc -> FieldBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> FieldBase Info VName)
-> MonoM SrcLoc -> MonoM (FieldBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc'
transformField (RecordFieldImplicit VName
v Info PatType
t SrcLoc
_) = do
Info PatType
t' <- (PatType -> MonoM PatType) -> Info PatType -> MonoM (Info PatType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatType
t
FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> FieldBase Info VName -> MonoM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$
Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
(VName -> Name
baseName VName
v)
(QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatType
t' SrcLoc
loc)
SrcLoc
loc
transformExp (ArrayLit [Exp]
es Info PatType
t SrcLoc
loc) =
[Exp] -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
ArrayLit ([Exp] -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatType -> MonoM PatType) -> Info PatType -> MonoM (Info PatType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (AppExp AppExp
e Info AppRes
res) = do
PatType -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims (PatType -> MonoM ()) -> PatType -> MonoM ()
forall a b. (a -> b) -> a -> b
$ AppRes -> PatType
appResType (AppRes -> PatType) -> AppRes -> PatType
forall a b. (a -> b) -> a -> b
$ Info AppRes -> AppRes
forall a. Info a -> a
unInfo Info AppRes
res
AppExp -> AppRes -> MonoM Exp
transformAppExp AppExp
e (AppRes -> MonoM Exp) -> MonoM AppRes -> MonoM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< AppRes -> MonoM AppRes
transformAppRes (Info AppRes -> AppRes
forall a. Info a -> a
unInfo Info AppRes
res)
transformExp (Var QualName VName
fname (Info PatType
t) SrcLoc
loc) = do
Maybe RecordReplacement
maybe_fs <- VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (VName -> MonoM (Maybe RecordReplacement))
-> VName -> MonoM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
case Maybe RecordReplacement
maybe_fs of
Just RecordReplacement
fs -> do
let toField :: (Name, (vn, PatType)) -> MonoM (FieldBase Info vn)
toField (Name
f, (vn
f_v, PatType
f_t)) = do
PatType
f_t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
f_t
let f_v' :: ExpBase Info vn
f_v' = QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
f_v) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
f_t') SrcLoc
loc
FieldBase Info vn -> MonoM (FieldBase Info vn)
forall (m :: * -> *) a. Monad m => a -> m a
return (FieldBase Info vn -> MonoM (FieldBase Info vn))
-> FieldBase Info vn -> MonoM (FieldBase Info vn)
forall a b. (a -> b) -> a -> b
$ Name -> ExpBase Info vn -> SrcLoc -> FieldBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
f ExpBase Info vn
f_v' SrcLoc
loc
[FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, (VName, PatType)) -> MonoM (FieldBase Info VName))
-> [(Name, (VName, PatType))] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, (VName, PatType)) -> MonoM (FieldBase Info VName)
forall vn. (Name, (vn, PatType)) -> MonoM (FieldBase Info vn)
toField (RecordReplacement -> [(Name, (VName, PatType))]
forall k a. Map k a -> [(k, a)]
M.toList RecordReplacement
fs) MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
Maybe RecordReplacement
Nothing -> do
PatType
t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t
SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t')
transformExp (Ascript Exp
e TypeDeclBase Info VName
tp SrcLoc
loc) =
Exp -> TypeDeclBase Info VName -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeDeclBase f vn -> SrcLoc -> ExpBase f vn
Ascript (Exp -> TypeDeclBase Info VName -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (TypeDeclBase Info VName -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeDeclBase Info VName -> SrcLoc -> Exp)
-> MonoM (TypeDeclBase Info VName) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeDeclBase Info VName -> MonoM (TypeDeclBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeDeclBase Info VName
tp MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Negate Exp
e SrcLoc
loc) =
Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Not Exp
e SrcLoc
loc) =
Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Lambda [Pat]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Aliasing, StructRetType)
tp SrcLoc
loc) = do
Exp
e0' <- Exp -> MonoM Exp
transformExp Exp
e0
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [Pat]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Aliasing, StructRetType)
tp SrcLoc
loc
transformExp (OpSection QualName VName
qn Info PatType
t SrcLoc
loc) =
Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info PatType
t SrcLoc
loc
transformExp (OpSectionLeft QualName VName
fname (Info PatType
t) Exp
e (Info (PName, TypeBase (DimDecl VName) (), Maybe VName),
Info (PName, TypeBase (DimDecl VName) ()))
arg (Info RetTypeBase (DimDecl VName) Aliasing
rettype, Info [VName]
retext) SrcLoc
loc) = do
let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype)) = (Info (PName, TypeBase (DimDecl VName) (), Maybe VName),
Info (PName, TypeBase (DimDecl VName) ()))
arg
Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
Exp
fname'
(Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
Maybe Exp
forall a. Maybe a
Nothing
PatType
t
(PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext)
(PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
forall a. Maybe a
Nothing)
(RetTypeBase (DimDecl VName) Aliasing
rettype, [VName]
retext)
SrcLoc
loc
transformExp (OpSectionRight QualName VName
fname (Info PatType
t) Exp
e (Info (PName, TypeBase (DimDecl VName) ()),
Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg (Info RetTypeBase (DimDecl VName) Aliasing
rettype) SrcLoc
loc) = do
let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)) = (Info (PName, TypeBase (DimDecl VName) ()),
Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg
Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
Exp
fname'
Maybe Exp
forall a. Maybe a
Nothing
(Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
PatType
t
(PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
forall a. Maybe a
Nothing)
(PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)
(RetTypeBase (DimDecl VName) Aliasing
rettype, [])
SrcLoc
loc
transformExp (ProjectSection [Name]
fields (Info PatType
t) SrcLoc
loc) =
[Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields PatType
t SrcLoc
loc
transformExp (IndexSection SliceBase Info VName
idxs (Info PatType
t) SrcLoc
loc) = do
SliceBase Info VName
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
SliceBase Info VName -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection SliceBase Info VName
idxs' PatType
t SrcLoc
loc
transformExp (Project Name
n Exp
e Info PatType
tp SrcLoc
loc) = do
Maybe RecordReplacement
maybe_fs <- case Exp
e of
Var QualName VName
qn Info PatType
_ SrcLoc
_ -> VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
Exp
_ -> Maybe RecordReplacement -> MonoM (Maybe RecordReplacement)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe RecordReplacement
forall a. Maybe a
Nothing
case Maybe RecordReplacement
maybe_fs of
Just RecordReplacement
m
| Just (VName
v, PatType
_) <- Name -> RecordReplacement -> Maybe (VName, PatType)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n RecordReplacement
m ->
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatType
tp SrcLoc
loc
Maybe RecordReplacement
_ -> do
Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info PatType
tp SrcLoc
loc
transformExp (Update Exp
e1 SliceBase Info VName
idxs Exp
e2 SrcLoc
loc) =
Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM (SliceBase Info VName) -> MonoM (Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
MonoM (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info PatType
t SrcLoc
loc) =
Exp -> [Name] -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
RecordUpdate (Exp -> [Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Name] -> MonoM (Exp -> Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Name] -> MonoM [Name]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
fs
MonoM (Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatType -> MonoM (Info PatType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatType
t
MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Assert Exp
e1 Exp
e2 Info String
desc SrcLoc
loc) =
Exp -> Exp -> Info String -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f String -> SrcLoc -> ExpBase f vn
Assert (Exp -> Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Exp -> Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info String -> SrcLoc -> Exp)
-> MonoM (Info String) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info String -> MonoM (Info String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info String
desc MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Constr Name
name [Exp]
all_es Info PatType
t SrcLoc
loc) =
Name -> [Exp] -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
Constr Name
name ([Exp] -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
all_es MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatType -> MonoM (Info PatType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) =
AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
info (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformCase :: Case -> MonoM Case
transformCase :: CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase (CasePat Pat
p Exp
e SrcLoc
loc) = do
(Pat
p', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
p
Pat -> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pat
p' (Exp -> SrcLoc -> CaseBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> CaseBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (Exp -> MonoM Exp
transformExp Exp
e) MonoM (SrcLoc -> CaseBase Info VName)
-> MonoM SrcLoc -> MonoM (CaseBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex (DimFix Exp
e) = Exp -> DimIndexBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Exp -> DimIndexBase Info VName)
-> MonoM Exp -> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e
transformDimIndex (DimSlice Maybe Exp
me1 Maybe Exp
me2 Maybe Exp
me3) =
Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp)
-> MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me1 MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp)
-> MonoM (Maybe Exp -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me2 MonoM (Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp) -> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me3
where
trans :: Maybe Exp -> MonoM (Maybe Exp)
trans = (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp
desugarBinOpSection ::
Exp ->
Maybe Exp ->
Maybe Exp ->
PatType ->
(PName, StructType, Maybe VName) ->
(PName, StructType, Maybe VName) ->
(PatRetType, [VName]) ->
SrcLoc ->
MonoM Exp
desugarBinOpSection :: Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection Exp
op Maybe Exp
e_left Maybe Exp
e_right PatType
t (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xext) (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yext) (RetType [VName]
dims PatType
rettype, [VName]
retext) SrcLoc
loc = do
(VName
v1, Exp -> Exp
wrap_left, Exp
e1, [Pat]
p1) <- Maybe Exp -> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam Maybe Exp
e_left (PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat]))
-> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
xtype
(VName
v2, Exp -> Exp
wrap_right, Exp
e2, [Pat]
p2) <- Maybe Exp -> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam Maybe Exp
e_right (PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat]))
-> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype
let apply_left :: Exp
apply_left =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
Exp
op
Exp
e1
((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
xext))
SrcLoc
loc
)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> PatType
-> RetTypeBase (DimDecl VName) Aliasing
-> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
yp (TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype) ([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
t)) [])
rettype' :: PatType
rettype' =
let onDim :: DimDecl VName -> DimDecl VName
onDim (NamedDim QualName VName
d)
| Named VName
p <- PName
xp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v1
| Named VName
p <- PName
yp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v2
onDim DimDecl VName
d = DimDecl VName
d
in (DimDecl VName -> DimDecl VName) -> PatType -> PatType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
onDim PatType
rettype
body :: Exp
body =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
( Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
Exp
apply_left
Exp
e2
((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
yext))
SrcLoc
loc
)
(AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
rettype' [VName]
retext)
rettype'' :: TypeBase (DimDecl VName) ()
rettype'' = PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
rettype'
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
Exp -> Exp
wrap_left (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
Exp -> Exp
wrap_right (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
[Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda ([Pat]
p1 [Pat] -> [Pat] -> [Pat]
forall a. [a] -> [a] -> [a]
++ [Pat]
p2) Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype'')) SrcLoc
loc
where
patAndVar :: PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype = do
VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
(VName, Pat, Exp) -> m (VName, Pat, Exp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
( VName
x,
VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty
)
makeVarParam :: Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam (Just Exp
e) PatType
argtype = do
(VName
v, Pat
pat, Exp
var_e) <- PatType -> m (VName, Pat, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype
let wrap :: Exp -> Exp
wrap Exp
body =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] Pat
pat Exp
e Exp
body SrcLoc
forall a. Monoid a => a
mempty) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (Exp -> PatType
typeOf Exp
body) [VName]
forall a. Monoid a => a
mempty)
(VName, Exp -> Exp, Exp, [Pat])
-> m (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, Exp -> Exp
wrap, Exp
var_e, [])
makeVarParam Maybe Exp
Nothing PatType
argtype = do
(VName
v, Pat
pat, Exp
var_e) <- PatType -> m (VName, Pat, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype
(VName, Exp -> Exp, Exp, [Pat])
-> m (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, Exp -> Exp
forall a. a -> a
id, Exp
var_e, [Pat
pat])
desugarProjectSection :: [Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection :: [Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields (Scalar (Arrow Aliasing
_ PName
_ PatType
t1 (RetType [VName]
dims PatType
t2))) SrcLoc
loc = do
VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"project_p"
let body :: Exp
body = (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Name -> Exp
project (QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1) SrcLoc
forall a. Monoid a => a
mempty) [Name]
fields
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
[Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda
[VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
p (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1) SrcLoc
forall a. Monoid a => a
mempty]
Exp
body
Maybe (TypeExp VName)
forall a. Maybe a
Nothing
((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t2))
SrcLoc
loc
where
project :: Exp -> Name -> Exp
project Exp
e Name
field =
case Exp -> PatType
typeOf Exp
e of
Scalar (Record Map Name PatType
fs)
| Just PatType
t <- Name -> Map Name PatType -> Maybe PatType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
field Map Name PatType
fs ->
Name -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
field Exp
e (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t) SrcLoc
forall a. Monoid a => a
mempty
PatType
t ->
String -> Exp
forall a. HasCallStack => String -> a
error (String -> Exp) -> String -> Exp
forall a b. (a -> b) -> a -> b
$
String
"desugarOpSection: type " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" does not have field "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
field
desugarProjectSection [Name]
_ PatType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarOpSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t
desugarIndexSection :: [DimIndex] -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection :: SliceBase Info VName -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection SliceBase Info VName
idxs (Scalar (Arrow Aliasing
_ PName
_ PatType
t1 (RetType [VName]
dims PatType
t2))) SrcLoc
loc = do
VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"index_i"
let body :: Exp
body = AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> SliceBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index (QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1) SrcLoc
loc) SliceBase Info VName
idxs SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
t2 []))
Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
[Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda
[VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
p (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1) SrcLoc
forall a. Monoid a => a
mempty]
Exp
body
Maybe (TypeExp VName)
forall a. Maybe a
Nothing
((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t2))
SrcLoc
loc
desugarIndexSection SliceBase Info VName
_ PatType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarIndexSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t
noticeDims :: TypeBase (DimDecl VName) as -> MonoM ()
noticeDims :: TypeBase (DimDecl VName) as -> MonoM ()
noticeDims = (DimDecl VName -> MonoM ()) -> [DimDecl VName] -> MonoM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ DimDecl VName -> MonoM ()
notice ([DimDecl VName] -> MonoM ())
-> (TypeBase (DimDecl VName) as -> [DimDecl VName])
-> TypeBase (DimDecl VName) as
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims
where
notice :: DimDecl VName -> MonoM ()
notice (NamedDim QualName VName
v) = MonoM Exp -> MonoM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MonoM Exp -> MonoM ()) -> MonoM Exp -> MonoM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
forall a. Monoid a => a
mempty QualName VName
v TypeBase (DimDecl VName) ()
forall dim als. TypeBase dim als
i64
notice DimDecl VName
_ = () -> MonoM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
unfoldLetFuns :: [ValBind] -> Exp -> Exp
unfoldLetFuns :: [ValBind] -> Exp -> Exp
unfoldLetFuns [] Exp
e = Exp
e
unfoldLetFuns (ValBind Maybe (Info EntryPoint)
_ VName
fname Maybe (TypeExp VName)
_ (Info (StructRetType
rettype, [VName]
_)) [TypeParamBase VName]
dim_params [Pat]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
loc : [ValBind]
rest) Exp
e =
AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (VName
-> ([TypeParamBase VName], [Pat], Maybe (TypeExp VName),
Info StructRetType, Exp)
-> Exp
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn], Maybe (TypeExp vn),
f StructRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
dim_params, [Pat]
params, Maybe (TypeExp VName)
forall a. Maybe a
Nothing, StructRetType -> Info StructRetType
forall a. a -> Info a
Info StructRetType
rettype, Exp
body) Exp
e' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
e_t [VName]
forall a. Monoid a => a
mempty)
where
e' :: Exp
e' = [ValBind] -> Exp -> Exp
unfoldLetFuns [ValBind]
rest Exp
e
e_t :: PatType
e_t = Exp -> PatType
typeOf Exp
e'
transformPat :: Pat -> MonoM (Pat, RecordReplacements)
transformPat :: Pat -> MonoM (Pat, RecordReplacements)
transformPat (Id VName
v (Info (Scalar (Record Map Name PatType
fs))) SrcLoc
loc) = do
let fs' :: [(Name, PatType)]
fs' = Map Name PatType -> [(Name, PatType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name PatType
fs
([VName]
fs_ks, [PatType]
fs_ts) <- ([(VName, PatType)] -> ([VName], [PatType]))
-> MonoM [(VName, PatType)] -> MonoM ([VName], [PatType])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, PatType)] -> ([VName], [PatType])
forall a b. [(a, b)] -> ([a], [b])
unzip (MonoM [(VName, PatType)] -> MonoM ([VName], [PatType]))
-> MonoM [(VName, PatType)] -> MonoM ([VName], [PatType])
forall a b. (a -> b) -> a -> b
$
[(Name, PatType)]
-> ((Name, PatType) -> MonoM (VName, PatType))
-> MonoM [(VName, PatType)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Name, PatType)]
fs' (((Name, PatType) -> MonoM (VName, PatType))
-> MonoM [(VName, PatType)])
-> ((Name, PatType) -> MonoM (VName, PatType))
-> MonoM [(VName, PatType)]
forall a b. (a -> b) -> a -> b
$ \(Name
f, PatType
ft) ->
(,) (VName -> PatType -> (VName, PatType))
-> MonoM VName -> MonoM (PatType -> (VName, PatType))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (Name -> String
nameToString Name
f) MonoM (PatType -> (VName, PatType))
-> MonoM PatType -> MonoM (VName, PatType)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
ft
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return
( [(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat
([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatType) -> Name) -> [(Name, PatType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatType)]
fs') ((VName -> Info PatType -> SrcLoc -> Pat)
-> [VName] -> [Info PatType] -> [SrcLoc] -> [Pat]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id [VName]
fs_ks ((PatType -> Info PatType) -> [PatType] -> [Info PatType]
forall a b. (a -> b) -> [a] -> [b]
map PatType -> Info PatType
forall a. a -> Info a
Info [PatType]
fs_ts) ([SrcLoc] -> [Pat]) -> [SrcLoc] -> [Pat]
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [SrcLoc]
forall a. a -> [a]
repeat SrcLoc
loc))
SrcLoc
loc,
VName -> RecordReplacement -> RecordReplacements
forall k a. k -> a -> Map k a
M.singleton VName
v (RecordReplacement -> RecordReplacements)
-> RecordReplacement -> RecordReplacements
forall a b. (a -> b) -> a -> b
$ [(Name, (VName, PatType))] -> RecordReplacement
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, PatType))] -> RecordReplacement)
-> [(Name, (VName, PatType))] -> RecordReplacement
forall a b. (a -> b) -> a -> b
$ [Name] -> [(VName, PatType)] -> [(Name, (VName, PatType))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatType) -> Name) -> [(Name, PatType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatType)]
fs') ([(VName, PatType)] -> [(Name, (VName, PatType))])
-> [(VName, PatType)] -> [(Name, (VName, PatType))]
forall a b. (a -> b) -> a -> b
$ [VName] -> [PatType] -> [(VName, PatType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
fs_ks [PatType]
fs_ts
)
transformPat (Id VName
v Info PatType
t SrcLoc
loc) = (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
v Info PatType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (TuplePat [Pat]
pats SrcLoc
loc) = do
([Pat]
pats', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
pats
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat [Pat]
pats' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat (RecordPat [(Name, Pat)]
fields SrcLoc
loc) = do
let ([Name]
field_names, [Pat]
field_pats) = [(Name, Pat)] -> ([Name], [Pat])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, Pat)]
fields
([Pat]
field_pats', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
field_pats
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat ([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
field_names [Pat]
field_pats') SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat (PatParens Pat
pat SrcLoc
loc) = do
(Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens Pat
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat (PatAttr AttrInfo VName
attr Pat
pat SrcLoc
loc) = do
(Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (AttrInfo VName -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr Pat
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat (Wildcard (Info PatType
t) SrcLoc
loc) = do
PatType
t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatType -> SrcLoc -> Pat
wildcard PatType
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (PatAscription Pat
pat TypeDeclBase Info VName
td SrcLoc
loc) = do
(Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Pat -> TypeDeclBase Info VName -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatBase f vn
PatAscription Pat
pat' TypeDeclBase Info VName
td SrcLoc
loc, RecordReplacements
rr)
transformPat (PatLit PatLit
e Info PatType
t SrcLoc
loc) = (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatLit -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatLit -> f PatType -> SrcLoc -> PatBase f vn
PatLit PatLit
e Info PatType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (PatConstr Name
name Info PatType
t [Pat]
all_ps SrcLoc
loc) = do
([Pat]
all_ps', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
all_ps
(Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Info PatType -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
name Info PatType
t [Pat]
all_ps' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
wildcard :: PatType -> SrcLoc -> Pat
wildcard :: PatType -> SrcLoc -> Pat
wildcard (Scalar (Record Map Name PatType
fs)) SrcLoc
loc =
[(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat ([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Map Name PatType -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name PatType
fs) ([Pat] -> [(Name, Pat)]) -> [Pat] -> [(Name, Pat)]
forall a b. (a -> b) -> a -> b
$ (PatType -> Pat) -> [PatType] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map ((Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
`Wildcard` SrcLoc
loc) (Info PatType -> Pat)
-> (PatType -> Info PatType) -> PatType -> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatType -> Info PatType
forall a. a -> Info a
Info) ([PatType] -> [Pat]) -> [PatType] -> [Pat]
forall a b. (a -> b) -> a -> b
$ Map Name PatType -> [PatType]
forall k a. Map k a -> [a]
M.elems Map Name PatType
fs) SrcLoc
loc
wildcard PatType
t SrcLoc
loc =
Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t) SrcLoc
loc
type DimInst = M.Map VName (DimDecl VName)
dimMapping ::
Monoid a =>
TypeBase (DimDecl VName) a ->
TypeBase (DimDecl VName) a ->
DimInst
dimMapping :: TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2 = State DimInst (TypeBase (DimDecl VName) a) -> DimInst -> DimInst
forall s a. State s a -> s -> s
execState (([VName]
-> DimDecl VName
-> DimDecl VName
-> StateT DimInst Identity (DimDecl VName))
-> TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a
-> State DimInst (TypeBase (DimDecl VName) a)
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 [VName]
-> DimDecl VName
-> DimDecl VName
-> StateT DimInst Identity (DimDecl VName)
forall (t :: * -> *) (f :: * -> *) vn a.
(Foldable t, MonadState (Map vn (DimDecl a)) f, Ord vn, Eq a) =>
t a -> DimDecl vn -> DimDecl a -> f (DimDecl vn)
f TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2) DimInst
forall a. Monoid a => a
mempty
where
f :: t a -> DimDecl vn -> DimDecl a -> f (DimDecl vn)
f t a
bound DimDecl vn
d1 (NamedDim QualName a
d2)
| QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
d2 a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t a
bound = DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl vn
d1
f t a
_ (NamedDim QualName vn
d1) DimDecl a
d2 = do
(Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ())
-> (Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ()
forall a b. (a -> b) -> a -> b
$ vn -> DimDecl a -> Map vn (DimDecl a) -> Map vn (DimDecl a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName vn -> vn
forall vn. QualName vn -> vn
qualLeaf QualName vn
d1) DimDecl a
d2
DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl vn -> f (DimDecl vn)) -> DimDecl vn -> f (DimDecl vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> DimDecl vn
forall vn. QualName vn -> DimDecl vn
NamedDim QualName vn
d1
f t a
_ DimDecl vn
d DimDecl a
_ = DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl vn
d
inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp]
inferSizeArgs :: [TypeParamBase VName]
-> TypeBase (DimDecl VName) () -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t =
(TypeParamBase VName -> Maybe Exp)
-> [TypeParamBase VName] -> [Exp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DimInst -> TypeParamBase VName -> Maybe Exp
forall k vn.
Ord k =>
Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg (TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> DimInst
forall a.
Monoid a =>
TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t)) [TypeParamBase VName]
tparams
where
tparamArg :: Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg Map k (DimDecl vn)
dinst TypeParamBase k
tp =
case k -> Map k (DimDecl vn) -> Maybe (DimDecl vn)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (TypeParamBase k -> k
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase k
tp) Map k (DimDecl vn)
dinst of
Just (NamedDim QualName vn
d) ->
ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName vn
d (PatType -> Info PatType
forall a. a -> Info a
Info PatType
forall dim als. TypeBase dim als
i64) SrcLoc
forall a. Monoid a => a
mempty
Just (ConstDim Int
x) ->
ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> Int64 -> IntValue
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) SrcLoc
forall a. Monoid a => a
mempty
Maybe (DimDecl vn)
_ ->
ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
0) SrcLoc
forall a. Monoid a => a
mempty
noNamedParams :: MonoType -> MonoType
noNamedParams :: MonoType -> MonoType
noNamedParams = MonoType -> MonoType
forall dim. TypeBase dim () -> TypeBase dim ()
f
where
f :: TypeBase dim () -> TypeBase dim ()
f (Array () Uniqueness
u ScalarTypeBase dim ()
t ShapeDecl dim
shape) = ()
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
u (ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t) ShapeDecl dim
shape
f (Scalar ScalarTypeBase dim ()
t) = ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t
f' :: ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' (Arrow () PName
_ TypeBase dim ()
t1 (RetType [VName]
dims TypeBase dim ()
t2)) =
()
-> PName
-> TypeBase dim ()
-> RetTypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) ([VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t2))
f' (Record Map Name (TypeBase dim ())
fs) =
Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim ()) -> ScalarTypeBase dim ())
-> Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ (TypeBase dim () -> TypeBase dim ())
-> Map Name (TypeBase dim ()) -> Map Name (TypeBase dim ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim () -> TypeBase dim ()
f Map Name (TypeBase dim ())
fs
f' (Sum Map Name [TypeBase dim ()]
cs) =
Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim ()] -> ScalarTypeBase dim ())
-> Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim ()] -> [TypeBase dim ()])
-> Map Name [TypeBase dim ()] -> Map Name [TypeBase dim ()]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim () -> TypeBase dim ())
-> [TypeBase dim ()] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim () -> TypeBase dim ()
f) Map Name [TypeBase dim ()]
cs
f' ScalarTypeBase dim ()
t = ScalarTypeBase dim ()
t
monomorphiseBinding ::
Bool ->
PolyBinding ->
MonoType ->
MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding :: Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
entry (PolyBinding RecordReplacements
rr (VName
name, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
rettype, [VName]
retext, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) MonoType
inst_t =
RecordReplacements
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr (MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind))
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ do
let bind_t :: TypeBase (DimDecl VName) ()
bind_t = [TypeBase (DimDecl VName) ()]
-> StructRetType -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
[TypeBase dim as] -> RetTypeBase dim as -> TypeBase dim as
foldFunType ((Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType [Pat]
params) StructRetType
rettype
(Map VName StructRetType
substs, [TypeParamBase VName]
t_shape_params) <- SrcLoc
-> TypeBase () ()
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName StructRetType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc (TypeBase (DimDecl VName) () -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes TypeBase (DimDecl VName) ()
bind_t) (MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName]))
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
inst_t
let substs' :: Map VName (Subst StructRetType)
substs' = (StructRetType -> Subst StructRetType)
-> Map VName StructRetType -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ([TypeParamBase VName] -> StructRetType -> Subst StructRetType
forall t. [TypeParamBase VName] -> t -> Subst t
Subst []) Map VName StructRetType
substs
rettype' :: StructRetType
rettype' = TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') StructRetType
rettype
substPatType :: PatType -> PatType
substPatType =
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> PatType -> PatType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((Subst StructRetType
-> Subst (RetTypeBase (DimDecl VName) Aliasing))
-> Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructRetType -> RetTypeBase (DimDecl VName) Aliasing)
-> Subst StructRetType
-> Subst (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((() -> Aliasing)
-> StructRetType -> RetTypeBase (DimDecl VName) Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliasing -> () -> Aliasing
forall a b. a -> b -> a
const Aliasing
forall a. Monoid a => a
mempty))) (Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> TypeSubs
-> VName
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs'))
params' :: [Pat]
params' = (Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
substPatType) [Pat]
params
bind_t' :: TypeBase (DimDecl VName) ()
bind_t' = TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') TypeBase (DimDecl VName) ()
bind_t
([TypeParamBase VName]
shape_params_explicit, [TypeParamBase VName]
shape_params_implicit) =
(TypeParamBase VName -> Bool)
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TypeBase (DimDecl VName) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
bind_t') (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) ([TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName]))
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$
[TypeParamBase VName]
shape_params [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
t_shape_params
([Pat]
params'', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
params'
(TypeBase (DimDecl VName) () -> MonoM ())
-> [TypeBase (DimDecl VName) ()] -> MonoM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ TypeBase (DimDecl VName) () -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims ([TypeBase (DimDecl VName) ()] -> MonoM ())
-> [TypeBase (DimDecl VName) ()] -> MonoM ()
forall a b. (a -> b) -> a -> b
$ StructRetType -> TypeBase (DimDecl VName) ()
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType StructRetType
rettype TypeBase (DimDecl VName) ()
-> [TypeBase (DimDecl VName) ()] -> [TypeBase (DimDecl VName) ()]
forall a. a -> [a] -> [a]
: (Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType [Pat]
params''
Exp
body' <- TypeSubs -> Exp -> MonoM Exp
forall (m :: * -> *). Monad m => TypeSubs -> Exp -> m Exp
updateExpTypes (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') Exp
body
Exp
body'' <- RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body'
VName
name' <- if [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParamBase VName]
tparams Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
entry then VName -> MonoM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
name else VName -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
name
(VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall (m :: * -> *) a. Monad m => a -> m a
return
( VName
name',
[TypeParamBase VName]
-> TypeBase (DimDecl VName) () -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
shape_params_explicit TypeBase (DimDecl VName) ()
bind_t',
if Bool
entry
then
VName
-> [TypeParamBase VName]
-> [Pat]
-> (StructRetType, [VName])
-> Exp
-> ValBind
toValBinding
VName
name'
([TypeParamBase VName]
shape_params_explicit [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
shape_params_implicit)
[Pat]
params''
(StructRetType
rettype', [VName]
retext)
Exp
body''
else
VName
-> [TypeParamBase VName]
-> [Pat]
-> (StructRetType, [VName])
-> Exp
-> ValBind
toValBinding
VName
name'
[TypeParamBase VName]
shape_params_implicit
((TypeParamBase VName -> Pat) -> [TypeParamBase VName] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> Pat
forall vn. TypeParamBase vn -> PatBase Info vn
shapeParam [TypeParamBase VName]
shape_params_explicit [Pat] -> [Pat] -> [Pat]
forall a. [a] -> [a] -> [a]
++ [Pat]
params'')
(StructRetType
rettype', [VName]
retext)
Exp
body''
)
where
shape_params :: [TypeParamBase VName]
shape_params = (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> (TypeParamBase VName -> Bool) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam) [TypeParamBase VName]
tparams
updateExpTypes :: TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs Exp
e = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (TypeSubs -> ASTMapper m
mapper TypeSubs
substs) Exp
e
mapper :: TypeSubs -> ASTMapper m
mapper TypeSubs
substs =
ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatType -> m PatType)
-> (StructRetType -> m StructRetType)
-> (RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
{ mapOnExp :: Exp -> m Exp
mapOnExp = TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs,
mapOnName :: VName -> m VName
mapOnName = VName -> m VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnQualName :: QualName VName -> m (QualName VName)
mapOnQualName = QualName VName -> m (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnStructType :: TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> m (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
mapOnPatType :: PatType -> m PatType
mapOnPatType = PatType -> m PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> m PatType)
-> (PatType -> PatType) -> PatType -> m PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
mapOnStructRetType :: StructRetType -> m StructRetType
mapOnStructRetType = StructRetType -> m StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> m StructRetType)
-> (StructRetType -> StructRetType)
-> StructRetType
-> m StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing))
-> (RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs
}
shapeParam :: TypeParamBase vn -> PatBase Info vn
shapeParam TypeParamBase vn
tp = vn -> Info PatType -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id (TypeParamBase vn -> vn
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase vn
tp) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
forall dim als. TypeBase dim als
i64) (SrcLoc -> PatBase Info vn) -> SrcLoc -> PatBase Info vn
forall a b. (a -> b) -> a -> b
$ TypeParamBase vn -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeParamBase vn
tp
toValBinding :: VName
-> [TypeParamBase VName]
-> [Pat]
-> (StructRetType, [VName])
-> Exp
-> ValBind
toValBinding VName
name' [TypeParamBase VName]
tparams' [Pat]
params'' (StructRetType, [VName])
rettype' Exp
body'' =
ValBind :: forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructRetType, [VName])
-> [TypeParamBase vn]
-> [PatBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo vn]
-> SrcLoc
-> ValBindBase f vn
ValBind
{ valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing,
valBindName :: VName
valBindName = VName
name',
valBindRetType :: Info (StructRetType, [VName])
valBindRetType = (StructRetType, [VName]) -> Info (StructRetType, [VName])
forall a. a -> Info a
Info (StructRetType, [VName])
rettype',
valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
forall a. Maybe a
Nothing,
valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
tparams',
valBindParams :: [Pat]
valBindParams = [Pat]
params'',
valBindBody :: Exp
valBindBody = Exp
body'',
valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing,
valBindAttrs :: [AttrInfo VName]
valBindAttrs = [AttrInfo VName]
attrs,
valBindLocation :: SrcLoc
valBindLocation = SrcLoc
loc
}
typeSubstsM ::
MonadFreshNames m =>
SrcLoc ->
TypeBase () () ->
MonoType ->
m (M.Map VName StructRetType, [TypeParam])
typeSubstsM :: SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName StructRetType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc TypeBase () ()
orig_t1 MonoType
orig_t2 =
WriterT [TypeParamBase VName] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParamBase VName])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT [TypeParamBase VName] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParamBase VName]))
-> WriterT [TypeParamBase VName] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ (Map VName StructRetType, Map Int VName) -> Map VName StructRetType
forall a b. (a, b) -> a
fst ((Map VName StructRetType, Map Int VName)
-> Map VName StructRetType)
-> WriterT
[TypeParamBase VName] m (Map VName StructRetType, Map Int VName)
-> WriterT [TypeParamBase VName] m (Map VName StructRetType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT
(Map VName StructRetType, Map Int VName)
(WriterT [TypeParamBase VName] m)
()
-> (Map VName StructRetType, Map Int VName)
-> WriterT
[TypeParamBase VName] m (Map VName StructRetType, Map Int VName)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (TypeBase () ()
-> MonoType
-> StateT
(Map VName StructRetType, Map Int VName)
(WriterT [TypeParamBase VName] m)
()
forall as (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *) dim as.
(MonadState
(Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
(t (t m)),
MonadTrans t, MonadTrans t, MonadFreshNames m,
MonadWriter [TypeParamBase VName] (t (t m)),
Pretty (ShapeDecl dim), Monad (t m)) =>
TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase () ()
orig_t1 MonoType
orig_t2) (Map VName StructRetType
forall a. Monoid a => a
mempty, Map Int VName
forall a. Monoid a => a
mempty)
where
subRet :: TypeBase dim as -> RetTypeBase MonoSize as -> t (t m) ()
subRet (Scalar (TypeVar as
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) RetTypeBase MonoSize as
rt =
Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Int
baseTag (TypeName -> VName
typeLeaf TypeName
v) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
forall as (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *).
(MonadState
(Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
(t (t m)),
MonadTrans t, MonadTrans t, MonadFreshNames m,
MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst TypeName
v RetTypeBase MonoSize as
rt
subRet TypeBase dim as
t1 (RetType [VName]
_ TypeBase MonoSize as
t2) =
TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase dim as
t1 TypeBase MonoSize as
t2
sub :: TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub t1 :: TypeBase dim as
t1@Array {} t2 :: TypeBase MonoSize as
t2@Array {}
| Just TypeBase dim as
t1' <- Int -> TypeBase dim as -> Maybe (TypeBase dim as)
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim as -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim as
t1) TypeBase dim as
t1,
Just TypeBase MonoSize as
t2' <- Int -> TypeBase MonoSize as -> Maybe (TypeBase MonoSize as)
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim as -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim as
t1) TypeBase MonoSize as
t2 =
TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase dim as
t1' TypeBase MonoSize as
t2'
sub (Scalar (TypeVar as
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) TypeBase MonoSize as
t =
Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Int
baseTag (TypeName -> VName
typeLeaf TypeName
v) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
forall as (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *).
(MonadState
(Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
(t (t m)),
MonadTrans t, MonadTrans t, MonadFreshNames m,
MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst TypeName
v (RetTypeBase MonoSize as -> t (t m) ())
-> RetTypeBase MonoSize as -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase MonoSize as -> RetTypeBase MonoSize as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase MonoSize as
t
sub (Scalar (Record Map Name (TypeBase dim as)
fields1)) (Scalar (Record Map Name (TypeBase MonoSize as)
fields2)) =
(TypeBase dim as -> TypeBase MonoSize as -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase MonoSize as] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub
(((Name, TypeBase dim as) -> TypeBase dim as)
-> [(Name, TypeBase dim as)] -> [TypeBase dim as]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase dim as) -> TypeBase dim as
forall a b. (a, b) -> b
snd ([(Name, TypeBase dim as)] -> [TypeBase dim as])
-> [(Name, TypeBase dim as)] -> [TypeBase dim as]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> [(Name, TypeBase dim as)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase dim as)
fields1)
(((Name, TypeBase MonoSize as) -> TypeBase MonoSize as)
-> [(Name, TypeBase MonoSize as)] -> [TypeBase MonoSize as]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase MonoSize as) -> TypeBase MonoSize as
forall a b. (a, b) -> b
snd ([(Name, TypeBase MonoSize as)] -> [TypeBase MonoSize as])
-> [(Name, TypeBase MonoSize as)] -> [TypeBase MonoSize as]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase MonoSize as) -> [(Name, TypeBase MonoSize as)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase MonoSize as)
fields2)
sub (Scalar Prim {}) (Scalar Prim {}) = () -> t (t m) ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
sub (Scalar (Arrow as
_ PName
_ TypeBase dim as
t1a (RetType [VName]
_ TypeBase dim as
t1b))) (Scalar (Arrow as
_ PName
_ TypeBase MonoSize as
t2a RetTypeBase MonoSize as
t2b)) = do
TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase dim as
t1a TypeBase MonoSize as
t2a
TypeBase dim as -> RetTypeBase MonoSize as -> t (t m) ()
subRet TypeBase dim as
t1b RetTypeBase MonoSize as
t2b
sub (Scalar (Sum Map Name [TypeBase dim as]
cs1)) (Scalar (Sum Map Name [TypeBase MonoSize as]
cs2)) =
((Name, [TypeBase dim as])
-> (Name, [TypeBase MonoSize as]) -> t (t m) [()])
-> [(Name, [TypeBase dim as])]
-> [(Name, [TypeBase MonoSize as])]
-> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Name, [TypeBase dim as])
-> (Name, [TypeBase MonoSize as]) -> t (t m) [()]
forall a a.
(a, [TypeBase dim as])
-> (a, [TypeBase MonoSize as]) -> t (t m) [()]
typeSubstClause (Map Name [TypeBase dim as] -> [(Name, [TypeBase dim as])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase dim as]
cs1) (Map Name [TypeBase MonoSize as] -> [(Name, [TypeBase MonoSize as])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase MonoSize as]
cs2)
where
typeSubstClause :: (a, [TypeBase dim as])
-> (a, [TypeBase MonoSize as]) -> t (t m) [()]
typeSubstClause (a
_, [TypeBase dim as]
ts1) (a
_, [TypeBase MonoSize as]
ts2) = (TypeBase dim as -> TypeBase MonoSize as -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase MonoSize as] -> t (t m) [()]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub [TypeBase dim as]
ts1 [TypeBase MonoSize as]
ts2
sub t1 :: TypeBase dim as
t1@(Scalar Sum {}) TypeBase MonoSize as
t2 = TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase dim as
t1 TypeBase MonoSize as
t2
sub TypeBase dim as
t1 t2 :: TypeBase MonoSize as
t2@(Scalar Sum {}) = TypeBase dim as -> TypeBase MonoSize as -> t (t m) ()
sub TypeBase dim as
t1 TypeBase MonoSize as
t2
sub TypeBase dim as
t1 TypeBase MonoSize as
t2 = String -> t (t m) ()
forall a. HasCallStack => String -> a
error (String -> t (t m) ()) -> String -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines [String
"typeSubstsM: mismatched types:", TypeBase dim as -> String
forall a. Pretty a => a -> String
pretty TypeBase dim as
t1, TypeBase MonoSize as -> String
forall a. Pretty a => a -> String
pretty TypeBase MonoSize as
t2]
addSubst :: TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst (TypeName [VName]
_ VName
v) (RetType [VName]
ext TypeBase MonoSize as
t) = do
(Map VName (RetTypeBase (DimDecl VName) as)
ts, Map Int VName
sizes) <- t (t m) (Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
forall s (m :: * -> *). MonadState s m => m s
get
Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName
v VName -> Map VName (RetTypeBase (DimDecl VName) as) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (RetTypeBase (DimDecl VName) as)
ts) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ do
TypeBase (DimDecl VName) as
t' <- (MonoSize -> t (t m) (DimDecl VName))
-> (as -> t (t m) as)
-> TypeBase MonoSize as
-> t (t m) (TypeBase (DimDecl VName) as)
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 MonoSize -> t (t m) (DimDecl VName)
forall a (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *).
(MonadState (a, Map Int VName) (t (t m)), MonadTrans t,
MonadTrans t, Monad (t m), MonadFreshNames m,
MonadWriter [TypeParamBase VName] (t (t m))) =>
MonoSize -> t (t m) (DimDecl VName)
onDim as -> t (t m) as
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase MonoSize as
t
(Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
-> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VName
-> RetTypeBase (DimDecl VName) as
-> Map VName (RetTypeBase (DimDecl VName) as)
-> Map VName (RetTypeBase (DimDecl VName) as)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v ([VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase (DimDecl VName) as
t') Map VName (RetTypeBase (DimDecl VName) as)
ts, Map Int VName
sizes)
onDim :: MonoSize -> t (t m) (DimDecl VName)
onDim (MonoKnown Int
i) = do
(a
ts, Map Int VName
sizes) <- t (t m) (a, Map Int VName)
forall s (m :: * -> *). MonadState s m => m s
get
case Int -> Map Int VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
i Map Int VName
sizes of
Maybe VName
Nothing -> do
VName
d <- t m VName -> t (t m) VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (t m VName -> t (t m) VName) -> t m VName -> t (t m) VName
forall a b. (a -> b) -> a -> b
$ m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"d"
[TypeParamBase VName] -> t (t m) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
TypeParamDim VName
d SrcLoc
loc]
(a, Map Int VName) -> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (a
ts, Int -> VName -> Map Int VName -> Map Int VName
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i VName
d Map Int VName
sizes)
DimDecl VName -> t (t m) (DimDecl VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d
Just VName
d ->
DimDecl VName -> t (t m) (DimDecl VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d
onDim (MonoAnon VName
v) = DimDecl VName -> t (t m) (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ Maybe VName -> DimDecl VName
forall vn. Maybe vn -> DimDecl vn
AnyDim (Maybe VName -> DimDecl VName) -> Maybe VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
substPat :: Bool -> (PatType -> PatType) -> Pat -> Pat
substPat :: Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
pat = case Pat
pat of
TuplePat [Pat]
pats SrcLoc
loc -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat ((Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f) [Pat]
pats) SrcLoc
loc
RecordPat [(Name, Pat)]
fs SrcLoc
loc -> [(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat (((Name, Pat) -> (Name, Pat)) -> [(Name, Pat)] -> [(Name, Pat)]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Pat) -> (Name, Pat)
forall a. (a, Pat) -> (a, Pat)
substField [(Name, Pat)]
fs) SrcLoc
loc
where
substField :: (a, Pat) -> (a, Pat)
substField (a
n, Pat
p) = (a
n, Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p)
PatParens Pat
p SrcLoc
loc -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p) SrcLoc
loc
PatAttr AttrInfo VName
attr Pat
p SrcLoc
loc -> AttrInfo VName -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p) SrcLoc
loc
Id VName
vn (Info PatType
tp) SrcLoc
loc -> VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
Wildcard (Info PatType
tp) SrcLoc
loc -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
PatAscription Pat
p TypeDeclBase Info VName
td SrcLoc
loc
| Bool
entry -> Pat -> TypeDeclBase Info VName -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatBase f vn
PatAscription (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
False PatType -> PatType
f Pat
p) TypeDeclBase Info VName
td SrcLoc
loc
| Bool
otherwise -> Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
False PatType -> PatType
f Pat
p
PatLit PatLit
e (Info PatType
tp) SrcLoc
loc -> PatLit -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatLit -> f PatType -> SrcLoc -> PatBase f vn
PatLit PatLit
e (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
PatConstr Name
n (Info PatType
tp) [Pat]
ps SrcLoc
loc -> Name -> Info PatType -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
n (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) [Pat]
ps SrcLoc
loc
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding (ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
_ (Info (StructRetType
rettype, [VName]
retext)) [TypeParamBase VName]
tparams [Pat]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) =
RecordReplacements
-> (VName, [TypeParamBase VName], [Pat], StructRetType, [VName],
Exp, [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
rettype, [VName]
retext, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables Bool
entry ValBind
valbind = do
let (ValBind Maybe (Info EntryPoint)
_ VName
_ Maybe (TypeExp VName)
_ (Info (RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype, [VName]
retext)) [TypeParamBase VName]
_ [Pat]
pats Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_) = ValBind
valbind
Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
let mapper :: ASTMapper MonoM
mapper =
ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatType -> m PatType)
-> (StructRetType -> m StructRetType)
-> (RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
{ mapOnExp :: Exp -> MonoM Exp
mapOnExp = Exp -> MonoM Exp
onExp,
mapOnName :: VName -> MonoM VName
mapOnName = VName -> MonoM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnQualName :: QualName VName -> MonoM (QualName VName)
mapOnQualName = QualName VName -> MonoM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
mapOnStructType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
mapOnPatType :: PatType -> MonoM PatType
mapOnPatType = PatType -> MonoM PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> MonoM PatType)
-> (PatType -> PatType) -> PatType -> MonoM PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
mapOnStructRetType :: StructRetType -> MonoM StructRetType
mapOnStructRetType = StructRetType -> MonoM StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> MonoM StructRetType)
-> (StructRetType -> StructRetType)
-> StructRetType
-> MonoM StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing))
-> (RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs)
}
onExp :: Exp -> MonoM Exp
onExp Exp
e = ASTMapper MonoM -> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper Exp
e
Exp
body' <- Exp -> MonoM Exp
onExp Exp
body
ValBind -> MonoM ValBind
forall (m :: * -> *) a. Monad m => a -> m a
return
ValBind
valbind
{ valBindRetType :: Info (StructRetType, [VName])
valBindRetType = (StructRetType, [VName]) -> Info (StructRetType, [VName])
forall a. a -> Info a
Info (TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) (StructRetType -> StructRetType) -> StructRetType -> StructRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype, [VName]
retext),
valBindParams :: [Pat]
valBindParams = (Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry ((PatType -> PatType) -> Pat -> Pat)
-> (PatType -> PatType) -> Pat -> Pat
forall a b. (a -> b) -> a -> b
$ TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs)) [Pat]
pats,
valBindBody :: Exp
valBindBody = Exp
body'
}
removeTypeVariablesInType :: StructType -> MonoM StructType
removeTypeVariablesInType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t = do
Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ()))
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$ TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) TypeBase (DimDecl VName) ()
t
transformValBind :: ValBind -> MonoM Env
transformValBind :: ValBind -> MonoM Env
transformValBind ValBind
valbind = do
PolyBinding
valbind' <-
ValBind -> PolyBinding
toPolyBinding
(ValBind -> PolyBinding) -> MonoM ValBind -> MonoM PolyBinding
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> ValBind -> MonoM ValBind
removeTypeVariables (Maybe (Info EntryPoint) -> Bool
forall a. Maybe a -> Bool
isJust (ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind)) ValBind
valbind
Bool -> MonoM () -> MonoM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe (Info EntryPoint) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Info EntryPoint) -> Bool)
-> Maybe (Info EntryPoint) -> Bool
forall a b. (a -> b) -> a -> b
$ ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind) (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ do
TypeBase (DimDecl VName) ()
t <-
TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType (TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ()))
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$
[TypeBase (DimDecl VName) ()]
-> StructRetType -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
[TypeBase dim as] -> RetTypeBase dim as -> TypeBase dim as
foldFunType
((Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType (ValBind -> [Pat]
forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBind
valbind))
(StructRetType -> TypeBase (DimDecl VName) ())
-> StructRetType -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ (StructRetType, [VName]) -> StructRetType
forall a b. (a, b) -> a
fst ((StructRetType, [VName]) -> StructRetType)
-> (StructRetType, [VName]) -> StructRetType
forall a b. (a -> b) -> a -> b
$ Info (StructRetType, [VName]) -> (StructRetType, [VName])
forall a. Info a -> a
unInfo (Info (StructRetType, [VName]) -> (StructRetType, [VName]))
-> Info (StructRetType, [VName]) -> (StructRetType, [VName])
forall a b. (a -> b) -> a -> b
$ ValBind -> Info (StructRetType, [VName])
forall (f :: * -> *) vn.
ValBindBase f vn -> f (StructRetType, [VName])
valBindRetType ValBind
valbind
(VName
name, InferSizeArgs
infer, ValBind
valbind'') <- Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
True PolyBinding
valbind' (MonoType -> MonoM (VName, InferSizeArgs, ValBind))
-> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t
Seq (VName, ValBind) -> MonoM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (VName, ValBind) -> MonoM ())
-> Seq (VName, ValBind) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (VName, ValBind) -> Seq (VName, ValBind)
forall a. a -> Seq a
Seq.singleton (VName
name, ValBind
valbind'' {valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind})
VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) (TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t) (VName
name, InferSizeArgs
infer)
Env -> MonoM Env
forall (m :: * -> *) a. Monad m => a -> m a
return Env
forall a. Monoid a => a
mempty {envPolyBindings :: Map VName PolyBinding
envPolyBindings = VName -> PolyBinding -> Map VName PolyBinding
forall k a. k -> a -> Map k a
M.singleton (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) PolyBinding
valbind'}
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeExp VName
_ (Info (RetType [VName]
dims TypeBase (DimDecl VName) ()
t)) Maybe DocComment
_ SrcLoc
_) = do
Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
TypeBase (DimDecl VName) () -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims TypeBase (DimDecl VName) ()
t
let tbinding :: TypeBinding
tbinding = Liftedness -> [TypeParamBase VName] -> StructRetType -> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) TypeBase (DimDecl VName) ()
t
Env -> MonoM Env
forall (m :: * -> *) a. Monad m => a -> m a
return Env
forall a. Monoid a => a
mempty {envTypeBindings :: Map VName TypeBinding
envTypeBindings = VName -> TypeBinding -> Map VName TypeBinding
forall k a. k -> a -> Map k a
M.singleton VName
name TypeBinding
tbinding}
transformDecs :: [Dec] -> MonoM ()
transformDecs :: [Dec] -> MonoM ()
transformDecs [] = () -> MonoM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
transformDecs (ValDec ValBind
valbind : [Dec]
ds) = do
Env
env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (TypeDec TypeBind
typebind : [Dec]
ds) = do
Env
env <- TypeBind -> MonoM Env
transformTypeBind TypeBind
typebind
Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (Dec
dec : [Dec]
_) =
String -> MonoM ()
forall a. HasCallStack => String -> a
error (String -> MonoM ()) -> String -> MonoM ()
forall a b. (a -> b) -> a -> b
$
String
"The monomorphization module expects a module-free "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"input program, but received: "
String -> ShowS
forall a. [a] -> [a] -> [a]
++ Dec -> String
forall a. Pretty a => a -> String
pretty Dec
dec
transformProg :: MonadFreshNames m => [Dec] -> m [ValBind]
transformProg :: [Dec] -> m [ValBind]
transformProg [Dec]
decs =
(((), Seq (VName, ValBind)) -> [ValBind])
-> m ((), Seq (VName, ValBind)) -> m [ValBind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Seq ValBind -> [ValBind]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq ValBind -> [ValBind])
-> (((), Seq (VName, ValBind)) -> Seq ValBind)
-> ((), Seq (VName, ValBind))
-> [ValBind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, ValBind) -> ValBind)
-> Seq (VName, ValBind) -> Seq ValBind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd (Seq (VName, ValBind) -> Seq ValBind)
-> (((), Seq (VName, ValBind)) -> Seq (VName, ValBind))
-> ((), Seq (VName, ValBind))
-> Seq ValBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((), Seq (VName, ValBind)) -> Seq (VName, ValBind)
forall a b. (a, b) -> b
snd) (m ((), Seq (VName, ValBind)) -> m [ValBind])
-> m ((), Seq (VName, ValBind)) -> m [ValBind]
forall a b. (a -> b) -> a -> b
$
(VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
-> m ((), Seq (VName, ValBind))
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
-> m ((), Seq (VName, ValBind)))
-> (VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
-> m ((), Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
VNameSource
-> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
forall a.
VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
runMonoM VNameSource
namesrc (MonoM () -> (((), Seq (VName, ValBind)), VNameSource))
-> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
decs