{-# LANGUAGE FlexibleContexts #-}
{-# 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.Bitraversable
import Data.Bifunctor
import Data.List (partition)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Set as S
import qualified Data.Sequence as Seq
import Data.Foldable
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.Semantic (TypeBinding(..))
import Language.Futhark.TypeChecker.Types
i32 :: TypeBase dim als
i32 :: TypeBase dim als
i32 = 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
Int32
data PolyBinding = PolyBinding RecordReplacements
(VName, [TypeParam], [Pattern],
Maybe (TypeExp VName), StructType, [VName], Exp,
[AttrInfo], SrcLoc)
type RecordReplacements = M.Map VName RecordReplacement
type RecordReplacement = M.Map Name (VName, PatternType)
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]
type MonoType = TypeBase Bool ()
monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType = Identity MonoType -> MonoType
forall a. Identity a -> a
runIdentity (Identity MonoType -> MonoType)
-> (TypeBase (DimDecl VName) als -> Identity MonoType)
-> TypeBase (DimDecl VName) als
-> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName -> DimPos -> DimDecl VName -> Identity Bool)
-> TypeBase (DimDecl VName) () -> Identity 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 -> Identity Bool
forall a (f :: * -> *) p.
(Ord a, Applicative f) =>
Set a -> p -> DimDecl a -> f Bool
onDim (TypeBase (DimDecl VName) () -> Identity MonoType)
-> (TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) als
-> Identity 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 a -> p -> DimDecl a -> f Bool
onDim Set a
bound p
_ (NamedDim QualName a
d)
| QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
d a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
bound = Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
onDim Set a
_ p
_ DimDecl a
AnyDim = Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
onDim Set a
_ p
_ DimDecl a
_ = Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
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
Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) (TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
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
TypeBase (DimDecl VName) ()
t' <- TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t
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 (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 (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) (TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
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 PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (PatternType -> Info PatternType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) () -> PatternType
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
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,
ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase 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))
(PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType (Int -> PatternType -> [PatternType]
forall a. Int -> a -> [a]
replicate Int
i PatternType
forall dim als. TypeBase dim als
i32) (TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t)), [VName] -> Info [VName]
forall a. a -> Info a
Info [])
SrcLoc
loc)
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 PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
fname')
(PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ((ExpBase Info vn -> PatternType)
-> [ExpBase Info vn] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map (PatternType -> ExpBase Info vn -> PatternType
forall a b. a -> b -> a
const PatternType
forall dim als. TypeBase dim als
i32) [ExpBase Info vn]
size_args)
(TypeBase (DimDecl VName) as -> PatternType
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, PatternType) -> Alias)
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound (VName -> Alias)
-> ((VName, PatternType) -> VName) -> (VName, PatternType) -> Alias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, PatternType) -> VName
forall a b. (a, b) -> a
fst) ([(VName, PatternType)] -> [Alias])
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> a -> b
$ RecordReplacement -> [(VName, PatternType)]
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
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 PatternType
t SrcLoc
_) = do
Info PatternType
t' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
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 PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatternType
t' SrcLoc
loc) SrcLoc
loc
transformExp (ArrayLit [Exp]
es Info PatternType
t SrcLoc
loc) =
[Exp] -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
ArrayLit ([Exp] -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatternType -> 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 PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
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 (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl (Info PatternType, Info [VName])
tp SrcLoc
loc) = 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
$ Exp
-> Maybe Exp
-> Inclusiveness Exp
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' (Info PatternType, Info [VName])
tp SrcLoc
loc
transformExp (Var QualName VName
fname (Info PatternType
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, PatternType)) -> MonoM (FieldBase Info vn)
toField (Name
f, (vn
f_v, PatternType
f_t)) = do
PatternType
f_t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
f_t
let f_v' :: ExpBase Info vn
f_v' = QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
f_v) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
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, PatternType)) -> MonoM (FieldBase Info VName))
-> [(Name, (VName, PatternType))] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, (VName, PatternType)) -> MonoM (FieldBase Info VName)
forall vn. (Name, (vn, PatternType)) -> MonoM (FieldBase Info vn)
toField (RecordReplacement -> [(Name, (VName, PatternType))]
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
PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
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 (Coerce Exp
e TypeDeclBase Info VName
tp (Info PatternType
t, Info [VName]
ext) SrcLoc
loc) = do
PatternType -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims PatternType
t
Exp
-> TypeDeclBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeDeclBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Coerce (Exp
-> TypeDeclBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp)
-> MonoM Exp
-> MonoM
(TypeDeclBase Info VName
-> (Info PatternType, 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
-> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (TypeDeclBase Info VName)
-> MonoM ((Info PatternType, Info [VName]) -> 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 ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
((,) (Info PatternType
-> Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info PatternType)
-> MonoM (Info [VName] -> (Info PatternType, Info [VName]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> MonoM PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t) MonoM (Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info [VName]) -> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info [VName] -> MonoM (Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info [VName]
ext) 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 (LetPat PatternBase Info VName
pat Exp
e1 Exp
e2 (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) = do
(PatternBase Info VName
pat', RecordReplacements
rr) <- PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern PatternBase Info VName
pat
PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
PatternBase Info VName
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat PatternBase Info VName
pat' (Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
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 ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
(Info PatternType, Info [VName])
-> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t', Info [VName]
retext) 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 (LetFun VName
fname ([TypeParamBase VName]
tparams, [PatternBase Info VName]
params, Maybe (TypeExp VName)
retdecl, Info TypeBase (DimDecl VName) ()
ret, Exp
body) Exp
e Info PatternType
e_t SrcLoc
loc)
| (TypeParamBase VName -> Bool) -> [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam [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], [PatternBase Info VName],
Maybe (TypeExp VName), TypeBase (DimDecl VName) (), [VName], Exp,
[AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [PatternBase Info VName]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
ret, [], Exp
body, [AttrInfo]
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
VName
-> ([TypeParamBase VName], [PatternBase Info VName],
Maybe (TypeExp VName), Info (TypeBase (DimDecl VName) ()), Exp)
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn),
f (TypeBase (DimDecl VName) ()), ExpBase f vn)
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [PatternBase Info VName]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) () -> Info (TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info TypeBase (DimDecl VName) ()
ret, Exp
body') (Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
Exp -> MonoM Exp
transformExp Exp
e MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
e_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 (If Exp
e1 Exp
e2 Exp
e3 (Info PatternType
tp, Info [VName]
retext) SrcLoc
loc) = do
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
Exp
e3' <- Exp -> MonoM Exp
transformExp Exp
e3
Info PatternType
tp' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
tp
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 -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
If Exp
e1' Exp
e2' Exp
e3' (Info PatternType
tp', Info [VName]
retext) SrcLoc
loc
transformExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d (Info PatternType
ret, Info [VName]
ext) SrcLoc
loc) = do
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
Info PatternType
ret' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
ret
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
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d (Info PatternType
ret', Info [VName]
ext) 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 (Lambda [PatternBase Info VName]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
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
$ [PatternBase Info VName]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
tp SrcLoc
loc
transformExp (OpSection QualName VName
qn Info PatternType
t SrcLoc
loc) =
Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info PatternType
t SrcLoc
loc
transformExp (OpSectionLeft QualName VName
fname (Info PatternType
t) Exp
e
(Info (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext), Info TypeBase (DimDecl VName) ()
ytype) (Info PatternType
rettype, Info [VName]
retext) SrcLoc
loc) = 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
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [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
PatternType
t (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext) (TypeBase (DimDecl VName) ()
ytype, Maybe VName
forall a. Maybe a
Nothing) (PatternType
rettype, [VName]
retext) SrcLoc
loc
transformExp (OpSectionRight QualName VName
fname (Info PatternType
t) Exp
e
(Info TypeBase (DimDecl VName) ()
xtype, Info (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)) (Info PatternType
rettype) SrcLoc
loc) = 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
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [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')
PatternType
t (TypeBase (DimDecl VName) ()
xtype, Maybe VName
forall a. Maybe a
Nothing) (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext) (PatternType
rettype, []) SrcLoc
loc
transformExp (ProjectSection [Name]
fields (Info PatternType
t) SrcLoc
loc) =
[Name] -> PatternType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields PatternType
t SrcLoc
loc
transformExp (IndexSection [DimIndexBase Info VName]
idxs (Info PatternType
t) SrcLoc
loc) =
[DimIndexBase Info VName] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection [DimIndexBase Info VName]
idxs PatternType
t SrcLoc
loc
transformExp (DoLoop [VName]
sparams PatternBase Info VName
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 Info (PatternType, [VName])
ret SrcLoc
loc) = 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 PatternBase Info VName
pat2 Exp
e2 -> PatternBase Info VName -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn PatternBase Info VName
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
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]
-> PatternBase Info VName
-> Exp
-> LoopFormBase Info VName
-> Exp
-> Info (PatternType, [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[VName]
-> PatternBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> f (PatternType, [VName])
-> SrcLoc
-> ExpBase f vn
DoLoop [VName]
sparams PatternBase Info VName
pat Exp
e1' LoopFormBase Info VName
form' Exp
e3' Info (PatternType, [VName])
ret SrcLoc
loc
transformExp (BinOp (QualName VName
fname, SrcLoc
oploc) (Info PatternType
t) (Exp
e1, Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2, Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) Info PatternType
tp Info [VName]
ext SrcLoc
loc) = 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
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
case Exp
fname' of
Var QualName VName
fname'' Info PatternType
_ SrcLoc
_ | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf Exp
e1'), PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf 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
$ (QualName VName, SrcLoc)
-> Info PatternType
-> (Exp, Info (TypeBase (DimDecl VName) (), Maybe VName))
-> (Exp, Info (TypeBase (DimDecl VName) (), Maybe VName))
-> Info PatternType
-> Info [VName]
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc)
-> f PatternType
-> (ExpBase f vn, f (TypeBase (DimDecl VName) (), Maybe VName))
-> (ExpBase f vn, f (TypeBase (DimDecl VName) (), Maybe VName))
-> f PatternType
-> f [VName]
-> SrcLoc
-> ExpBase f vn
BinOp (QualName VName
fname'', SrcLoc
oploc) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) (Exp
e1', Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2', Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) Info PatternType
tp Info [VName]
ext SrcLoc
loc
Exp
_ -> do
(Exp
x_param_e, PatternBase Info VName
x_param) <- Exp -> MonoM (Exp, PatternBase Info VName)
forall (m :: * -> *).
MonadFreshNames m =>
Exp -> m (Exp, PatternBase Info VName)
makeVarParam Exp
e1'
(Exp
y_param_e, PatternBase Info VName
y_param) <- Exp -> MonoM (Exp, PatternBase Info VName)
forall (m :: * -> *).
MonadFreshNames m =>
Exp -> m (Exp, PatternBase Info VName)
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
$ PatternBase Info VName
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat PatternBase Info VName
x_param Exp
e1'
(PatternBase Info VName
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat PatternBase Info VName
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) (Info PatternType
tp, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty)
(Info PatternType
tp, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty) SrcLoc
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 =
ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply (ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase 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)))
(PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (TypeBase (DimDecl VName) () -> PatternType)
-> TypeBase (DimDecl VName) () -> PatternType
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) (), Maybe VName)
-> TypeBase (DimDecl VName) ()
forall a b. (a, b) -> a
fst (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d2)] (Info PatternType -> PatternType
forall a. Info a -> a
unInfo Info PatternType
tp)),
[VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty) SrcLoc
loc)
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))) (Info PatternType
tp, Info [VName]
ext) SrcLoc
loc
makeVarParam :: Exp -> m (Exp, PatternBase Info VName)
makeVarParam Exp
arg = do
let argtype :: PatternType
argtype = Exp -> PatternType
typeOf Exp
arg
VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"binop_p"
(Exp, PatternBase Info VName) -> m (Exp, PatternBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
x (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty)
transformExp (Project Name
n Exp
e Info PatternType
tp SrcLoc
loc) = do
Maybe RecordReplacement
maybe_fs <- case Exp
e of
Var QualName VName
qn Info PatternType
_ 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, PatternType
_) <- Name -> RecordReplacement -> Maybe (VName, PatternType)
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 PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatternType
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 PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info PatternType
tp SrcLoc
loc
transformExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs Exp
e1 Exp
body (Info PatternType
t) SrcLoc
loc) = do
[DimIndexBase Info VName]
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase 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 [DimIndexBase Info VName]
idxs
Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
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
$ IdentBase Info VName
-> IdentBase Info VName
-> [DimIndexBase Info VName]
-> Exp
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
IdentBase f vn
-> IdentBase f vn
-> [DimIndexBase f vn]
-> ExpBase f vn
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs' Exp
e1' Exp
body' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t') SrcLoc
loc
transformExp (Index Exp
e0 [DimIndexBase Info VName]
idxs (Info PatternType, Info [VName])
info SrcLoc
loc) =
Exp
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (Exp
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp)
-> MonoM Exp
-> MonoM
([DimIndexBase Info VName]
-> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e0 MonoM
([DimIndexBase Info VName]
-> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM [DimIndexBase Info VName]
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase 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 [DimIndexBase Info VName]
idxs MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Info PatternType, Info [VName])
-> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Info PatternType, Info [VName])
info 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 (Update Exp
e1 [DimIndexBase Info VName]
idxs Exp
e2 SrcLoc
loc) =
Exp -> [DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (Exp -> [DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
-> MonoM [DimIndexBase 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))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase 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 [DimIndexBase 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 PatternType
t SrcLoc
loc) =
Exp -> [Name] -> Exp -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name]
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
RecordUpdate (Exp -> [Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Name] -> MonoM (Exp -> Info PatternType -> 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 PatternType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatternType
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 PatternType
t SrcLoc
loc) =
Name -> [Exp] -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
Constr Name
name ([Exp] -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatternType -> 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 PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatternType
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 (Match Exp
e NonEmpty (CaseBase Info VName)
cs (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) =
Exp
-> NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Match (Exp
-> NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp)
-> MonoM Exp
-> MonoM
(NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM
(NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (NonEmpty (CaseBase Info VName))
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
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 ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
((,) (Info PatternType
-> Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info PatternType)
-> MonoM (Info [VName] -> (Info PatternType, Info [VName]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
t MonoM (Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info [VName]) -> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info [VName] -> MonoM (Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info [VName]
retext) 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
info Exp
e SrcLoc
loc) =
AttrInfo -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo
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 PatternBase Info VName
p Exp
e SrcLoc
loc) = do
(PatternBase Info VName
p', RecordReplacements
rr) <- PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern PatternBase Info VName
p
PatternBase Info VName -> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat PatternBase Info VName
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
-> PatternType
-> (StructType, Maybe VName) -> (StructType, Maybe VName)
-> (PatternType, [VName]) -> SrcLoc -> MonoM Exp
desugarBinOpSection :: Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection Exp
op Maybe Exp
e_left Maybe Exp
e_right PatternType
t (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xext) (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yext) (PatternType
rettype, [VName]
retext) SrcLoc
loc = do
(Exp
e1, [PatternBase Info VName]
p1) <- Maybe Exp -> PatternType -> MonoM (Exp, [PatternBase Info VName])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (Exp, [PatternBase Info VName])
makeVarParam Maybe Exp
e_left (PatternType -> MonoM (Exp, [PatternBase Info VName]))
-> PatternType -> MonoM (Exp, [PatternBase Info VName])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
xtype
(Exp
e2, [PatternBase Info VName]
p2) <- Maybe Exp -> PatternType -> MonoM (Exp, [PatternBase Info VName])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (Exp, [PatternBase Info VName])
makeVarParam Maybe Exp
e_right (PatternType -> MonoM (Exp, [PatternBase Info VName]))
-> PatternType -> MonoM (Exp, [PatternBase Info VName])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype
let apply_left :: Exp
apply_left = Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply Exp
op Exp
e1 ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
xext))
(PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype] PatternType
t, [VName] -> Info [VName]
forall a. a -> Info a
Info []) SrcLoc
loc
body :: Exp
body = Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply Exp
apply_left Exp
e2 ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
yext))
(PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
rettype, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
retext) SrcLoc
loc
rettype' :: TypeBase (DimDecl VName) ()
rettype' = PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
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
$ [PatternBase Info VName]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda ([PatternBase Info VName]
p1 [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. [a] -> [a] -> [a]
++ [PatternBase Info VName]
p2) Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, TypeBase (DimDecl VName) ())
-> Info (Aliasing, TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, TypeBase (DimDecl VName) ()
rettype')) SrcLoc
loc
where makeVarParam :: Maybe Exp -> PatternType -> m (Exp, [PatternBase Info VName])
makeVarParam (Just Exp
e) PatternType
_ = (Exp, [PatternBase Info VName])
-> m (Exp, [PatternBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp
e, [])
makeVarParam Maybe Exp
Nothing PatternType
argtype = do
VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
(Exp, [PatternBase Info VName])
-> m (Exp, [PatternBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
[VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
x (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty])
desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp
desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields (Scalar (Arrow Aliasing
_ PName
_ PatternType
t1 PatternType
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 PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
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
$ [PatternBase Info VName]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
p (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t1) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, TypeBase (DimDecl VName) ())
-> Info (Aliasing, TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t2)) SrcLoc
loc
where project :: Exp -> Name -> Exp
project Exp
e Name
field =
case Exp -> PatternType
typeOf Exp
e of
Scalar (Record Map Name PatternType
fs)
| Just PatternType
t <- Name -> Map Name PatternType -> Maybe PatternType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
field Map Name PatternType
fs ->
Name -> Exp -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
field Exp
e (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
forall a. Monoid a => a
mempty
PatternType
t -> String -> Exp
forall a. HasCallStack => String -> a
error (String -> Exp) -> String -> Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarOpSection: type " String -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t String -> String -> String
forall a. [a] -> [a] -> [a]
++
String
" does not have field " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
field
desugarProjectSection [Name]
_ PatternType
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 -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t
desugarIndexSection :: [DimIndex] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection :: [DimIndexBase Info VName] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection [DimIndexBase Info VName]
idxs (Scalar (Arrow Aliasing
_ PName
_ PatternType
t1 PatternType
t2)) SrcLoc
loc = do
VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"index_i"
let body :: Exp
body = Exp
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t1) SrcLoc
loc) [DimIndexBase Info VName]
idxs (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t2, [VName] -> Info [VName]
forall a. a -> Info a
Info []) SrcLoc
loc
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
$ [PatternBase Info VName]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
p (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t1) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, TypeBase (DimDecl VName) ())
-> Info (Aliasing, TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t2)) SrcLoc
loc
desugarIndexSection [DimIndexBase Info VName]
_ PatternType
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 -> String -> String
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
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
i32
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 (TypeBase (DimDecl VName) ()
rettype, [VName]
_)) [TypeParamBase VName]
dim_params [PatternBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
loc : [ValBind]
rest) Exp
e =
VName
-> ([TypeParamBase VName], [PatternBase Info VName],
Maybe (TypeExp VName), Info (TypeBase (DimDecl VName) ()), Exp)
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn),
f (TypeBase (DimDecl VName) ()), ExpBase f vn)
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
dim_params, [PatternBase Info VName]
params, Maybe (TypeExp VName)
forall a. Maybe a
Nothing, TypeBase (DimDecl VName) () -> Info (TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info TypeBase (DimDecl VName) ()
rettype, Exp
body) Exp
e' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
e_t) SrcLoc
loc
where e' :: Exp
e' = [ValBind] -> Exp -> Exp
unfoldLetFuns [ValBind]
rest Exp
e
e_t :: PatternType
e_t = Exp -> PatternType
typeOf Exp
e'
transformPattern :: Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern :: PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern (Id VName
v (Info (Scalar (Record Map Name PatternType
fs))) SrcLoc
loc) = do
let fs' :: [(Name, PatternType)]
fs' = Map Name PatternType -> [(Name, PatternType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name PatternType
fs
([VName]
fs_ks, [PatternType]
fs_ts) <- ([(VName, PatternType)] -> ([VName], [PatternType]))
-> MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, PatternType)] -> ([VName], [PatternType])
forall a b. [(a, b)] -> ([a], [b])
unzip (MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType]))
-> MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType])
forall a b. (a -> b) -> a -> b
$ [(Name, PatternType)]
-> ((Name, PatternType) -> MonoM (VName, PatternType))
-> MonoM [(VName, PatternType)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Name, PatternType)]
fs' (((Name, PatternType) -> MonoM (VName, PatternType))
-> MonoM [(VName, PatternType)])
-> ((Name, PatternType) -> MonoM (VName, PatternType))
-> MonoM [(VName, PatternType)]
forall a b. (a -> b) -> a -> b
$ \(Name
f, PatternType
ft) ->
(,) (VName -> PatternType -> (VName, PatternType))
-> MonoM VName -> MonoM (PatternType -> (VName, PatternType))
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 (PatternType -> (VName, PatternType))
-> MonoM PatternType -> MonoM (VName, PatternType)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
ft
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern ([Name]
-> [PatternBase Info VName] -> [(Name, PatternBase Info VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatternType) -> Name) -> [(Name, PatternType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternType)]
fs')
((VName -> Info PatternType -> SrcLoc -> PatternBase Info VName)
-> [VName]
-> [Info PatternType]
-> [SrcLoc]
-> [PatternBase Info VName]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id [VName]
fs_ks ((PatternType -> Info PatternType)
-> [PatternType] -> [Info PatternType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> Info PatternType
forall a. a -> Info a
Info [PatternType]
fs_ts) ([SrcLoc] -> [PatternBase Info VName])
-> [SrcLoc] -> [PatternBase Info VName]
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, PatternType))] -> RecordReplacement
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, PatternType))] -> RecordReplacement)
-> [(Name, (VName, PatternType))] -> RecordReplacement
forall a b. (a -> b) -> a -> b
$ [Name] -> [(VName, PatternType)] -> [(Name, (VName, PatternType))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatternType) -> Name) -> [(Name, PatternType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternType)]
fs') ([(VName, PatternType)] -> [(Name, (VName, PatternType))])
-> [(VName, PatternType)] -> [(Name, (VName, PatternType))]
forall a b. (a -> b) -> a -> b
$ [VName] -> [PatternType] -> [(VName, PatternType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
fs_ks [PatternType]
fs_ts)
transformPattern (Id VName
v Info PatternType
t SrcLoc
loc) = (PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
v Info PatternType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (TuplePattern [PatternBase Info VName]
pats SrcLoc
loc) = do
([PatternBase Info VName]
pats', [RecordReplacements]
rrs) <- [(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements]))
-> MonoM [(PatternBase Info VName, RecordReplacements)]
-> MonoM ([PatternBase Info VName], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements))
-> [PatternBase Info VName]
-> MonoM [(PatternBase Info VName, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern [PatternBase Info VName]
pats
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([PatternBase Info VName] -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern [PatternBase Info VName]
pats' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPattern (RecordPattern [(Name, PatternBase Info VName)]
fields SrcLoc
loc) = do
let ([Name]
field_names, [PatternBase Info VName]
field_pats) = [(Name, PatternBase Info VName)]
-> ([Name], [PatternBase Info VName])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, PatternBase Info VName)]
fields
([PatternBase Info VName]
field_pats', [RecordReplacements]
rrs) <- [(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements]))
-> MonoM [(PatternBase Info VName, RecordReplacements)]
-> MonoM ([PatternBase Info VName], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements))
-> [PatternBase Info VName]
-> MonoM [(PatternBase Info VName, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern [PatternBase Info VName]
field_pats
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern ([Name]
-> [PatternBase Info VName] -> [(Name, PatternBase Info VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
field_names [PatternBase Info VName]
field_pats') SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPattern (PatternParens PatternBase Info VName
pat SrcLoc
loc) = do
(PatternBase Info VName
pat', RecordReplacements
rr) <- PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern PatternBase Info VName
pat
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> SrcLoc -> PatternBase f vn
PatternParens PatternBase Info VName
pat' SrcLoc
loc, RecordReplacements
rr)
transformPattern (Wildcard (Info PatternType
t) SrcLoc
loc) = do
PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternType -> SrcLoc -> PatternBase Info VName
wildcard PatternType
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (PatternAscription PatternBase Info VName
pat TypeDeclBase Info VName
td SrcLoc
loc) = do
(PatternBase Info VName
pat', RecordReplacements
rr) <- PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern PatternBase Info VName
pat
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatternBase f vn
PatternAscription PatternBase Info VName
pat' TypeDeclBase Info VName
td SrcLoc
loc, RecordReplacements
rr)
transformPattern (PatternLit Exp
e Info PatternType
t SrcLoc
loc) = (PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> f PatternType -> SrcLoc -> PatternBase f vn
PatternLit Exp
e Info PatternType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (PatternConstr Name
name Info PatternType
t [PatternBase Info VName]
all_ps SrcLoc
loc) = do
([PatternBase Info VName]
all_ps', [RecordReplacements]
rrs) <- [(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements]))
-> MonoM [(PatternBase Info VName, RecordReplacements)]
-> MonoM ([PatternBase Info VName], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements))
-> [PatternBase Info VName]
-> MonoM [(PatternBase Info VName, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern [PatternBase Info VName]
all_ps
(PatternBase Info VName, RecordReplacements)
-> MonoM (PatternBase Info VName, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
-> Info PatternType
-> [PatternBase Info VName]
-> SrcLoc
-> PatternBase Info VName
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
name Info PatternType
t [PatternBase Info VName]
all_ps' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
wildcard :: PatternType -> SrcLoc -> Pattern
wildcard :: PatternType -> SrcLoc -> PatternBase Info VName
wildcard (Scalar (Record Map Name PatternType
fs)) SrcLoc
loc =
[(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern ([Name]
-> [PatternBase Info VName] -> [(Name, PatternBase Info VName)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Map Name PatternType -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name PatternType
fs) ([PatternBase Info VName] -> [(Name, PatternBase Info VName)])
-> [PatternBase Info VName] -> [(Name, PatternBase Info VName)]
forall a b. (a -> b) -> a -> b
$ (PatternType -> PatternBase Info VName)
-> [PatternType] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map ((Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
`Wildcard` SrcLoc
loc) (Info PatternType -> PatternBase Info VName)
-> (PatternType -> Info PatternType)
-> PatternType
-> PatternBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternType -> Info PatternType
forall a. a -> Info a
Info) ([PatternType] -> [PatternBase Info VName])
-> [PatternType] -> [PatternBase Info VName]
forall a b. (a -> b) -> a -> b
$ Map Name PatternType -> [PatternType]
forall k a. Map k a -> [a]
M.elems Map Name PatternType
fs) SrcLoc
loc
wildcard PatternType
t SrcLoc
loc =
Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
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 ((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) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims DimDecl VName
-> DimDecl VName -> StateT DimInst Identity (DimDecl VName)
forall (m :: * -> *) vn a.
(MonadState (Map vn a) m, Ord vn) =>
DimDecl vn -> a -> m (DimDecl vn)
f TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2) DimInst
forall a. Monoid a => a
mempty
where f :: DimDecl vn -> a -> m (DimDecl vn)
f (NamedDim QualName vn
d1) a
d2 = do
(Map vn a -> Map vn a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map vn a -> Map vn a) -> m ()) -> (Map vn a -> Map vn a) -> m ()
forall a b. (a -> b) -> a -> b
$ vn -> a -> Map vn a -> Map vn 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) a
d2
DimDecl vn -> m (DimDecl vn)
forall (m :: * -> *) a. Monad m => a -> m a
return (DimDecl vn -> m (DimDecl vn)) -> DimDecl vn -> m (DimDecl vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> DimDecl vn
forall vn. QualName vn -> DimDecl vn
NamedDim QualName vn
d1
f DimDecl vn
d a
_ = DimDecl vn -> m (DimDecl vn)
forall (m :: * -> *) a. Monad m => a -> m a
return 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 PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
d (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
forall dim als. TypeBase dim als
i32) 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
$ Int32 -> IntValue
Int32Value (Int32 -> IntValue) -> Int32 -> IntValue
forall a b. (a -> b) -> a -> b
$ Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) SrcLoc
forall a. Monoid a => a
mempty
Maybe (DimDecl vn)
_ ->
Maybe (ExpBase Info vn)
forall a. Maybe a
Nothing
explicitSizes :: StructType -> MonoType -> S.Set VName
explicitSizes :: TypeBase (DimDecl VName) () -> MonoType -> Set VName
explicitSizes TypeBase (DimDecl VName) ()
t1 MonoType
t2 =
State (Set VName) (TypeBase (DimDecl VName) ())
-> Set VName -> Set VName
forall s a. State s a -> s -> s
execState ((DimDecl VName
-> Bool -> StateT (Set VName) Identity (DimDecl VName))
-> TypeBase (DimDecl VName) ()
-> MonoType
-> State (Set VName) (TypeBase (DimDecl VName) ())
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims DimDecl VName
-> Bool -> StateT (Set VName) Identity (DimDecl VName)
forall (m :: * -> *) a.
(MonadState (Set a) m, Ord a) =>
DimDecl a -> Bool -> m (DimDecl a)
onDims TypeBase (DimDecl VName) ()
t1 MonoType
t2) Set VName
forall a. Monoid a => a
mempty Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` TypeBase (DimDecl VName) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
t1
where onDims :: DimDecl a -> Bool -> m (DimDecl a)
onDims DimDecl a
d1 Bool
d2 = do
case (DimDecl a
d1, Bool
d2) of
(NamedDim QualName a
v, Bool
True) -> (Set a -> Set a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set a -> Set a) -> m ()) -> (Set a -> Set a) -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert (a -> Set a -> Set a) -> a -> Set a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
(DimDecl a, Bool)
_ -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
DimDecl a -> m (DimDecl a)
forall (m :: * -> *) a. Monad m => a -> m a
return DimDecl a
d1
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 TypeBase dim ()
t2) =
()
-> PName
-> TypeBase dim ()
-> TypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) (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, [PatternBase Info VName]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
rettype, [VName]
retext, Exp
body, [AttrInfo]
attrs, SrcLoc
loc)) MonoType
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) ()]
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ((PatternBase Info VName -> TypeBase (DimDecl VName) ())
-> [PatternBase Info VName] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> TypeBase (DimDecl VName) ()
patternStructType [PatternBase Info VName]
params) TypeBase (DimDecl VName) ()
rettype
(Map VName (TypeBase (DimDecl VName) ())
substs, [TypeParamBase VName]
t_shape_params) <- SrcLoc
-> TypeBase () ()
-> MonoType
-> MonoM
(Map VName (TypeBase (DimDecl VName) ()), [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName (TypeBase (DimDecl VName) ()),
[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 (TypeBase (DimDecl VName) ()), [TypeParamBase VName]))
-> MonoType
-> MonoM
(Map VName (TypeBase (DimDecl VName) ()), [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
t
let substs' :: Map VName (Subst (TypeBase (DimDecl VName) ()))
substs' = (TypeBase (DimDecl VName) ()
-> Subst (TypeBase (DimDecl VName) ()))
-> Map VName (TypeBase (DimDecl VName) ())
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase (DimDecl VName) () -> Subst (TypeBase (DimDecl VName) ())
forall t. t -> Subst t
Subst Map VName (TypeBase (DimDecl VName) ())
substs
rettype' :: TypeBase (DimDecl VName) ()
rettype' = (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs') TypeBase (DimDecl VName) ()
rettype
substPatternType :: PatternType -> PatternType
substPatternType =
(VName -> Maybe (Subst PatternType)) -> PatternType -> PatternType
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((Subst (TypeBase (DimDecl VName) ()) -> Subst PatternType)
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase (DimDecl VName) () -> PatternType)
-> Subst (TypeBase (DimDecl VName) ()) -> Subst PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct) (Maybe (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst PatternType))
-> (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> VName
-> Maybe (Subst PatternType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs'))
params' :: [PatternBase Info VName]
params' = (PatternBase Info VName -> PatternBase Info VName)
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry PatternType -> PatternType
substPatternType) [PatternBase Info VName]
params
bind_t' :: TypeBase (DimDecl VName) ()
bind_t' = (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
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) () -> MonoType -> Set VName
explicitSizes TypeBase (DimDecl VName) ()
bind_t' MonoType
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
([PatternBase Info VName]
params'', [RecordReplacements]
rrs) <- [(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(PatternBase Info VName, RecordReplacements)]
-> ([PatternBase Info VName], [RecordReplacements]))
-> MonoM [(PatternBase Info VName, RecordReplacements)]
-> MonoM ([PatternBase Info VName], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements))
-> [PatternBase Info VName]
-> MonoM [(PatternBase Info VName, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM PatternBase Info VName
-> MonoM (PatternBase Info VName, RecordReplacements)
transformPattern [PatternBase Info VName]
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
$ TypeBase (DimDecl VName) ()
rettype TypeBase (DimDecl VName) ()
-> [TypeBase (DimDecl VName) ()] -> [TypeBase (DimDecl VName) ()]
forall a. a -> [a] -> [a]
: (PatternBase Info VName -> TypeBase (DimDecl VName) ())
-> [PatternBase Info VName] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> TypeBase (DimDecl VName) ()
patternStructType [PatternBase Info VName]
params''
Exp
body' <- (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> x -> m x
updateExpTypes (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
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]
-> [PatternBase Info VName]
-> (TypeBase (DimDecl VName) (), [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) [PatternBase Info VName]
params''
(TypeBase (DimDecl VName) ()
rettype', [VName]
retext) Exp
body''
else VName
-> [TypeParamBase VName]
-> [PatternBase Info VName]
-> (TypeBase (DimDecl VName) (), [VName])
-> Exp
-> ValBind
toValBinding VName
name' [TypeParamBase VName]
shape_params_implicit
((TypeParamBase VName -> PatternBase Info VName)
-> [TypeParamBase VName] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> PatternBase Info VName
forall vn. TypeParamBase vn -> PatternBase Info vn
shapeParam [TypeParamBase VName]
shape_params_explicit [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. [a] -> [a] -> [a]
++ [PatternBase Info VName]
params'')
(TypeBase (DimDecl VName) ()
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 :: (VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> x -> m x
updateExpTypes VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs = ASTMapper m -> x -> m x
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m -> x -> m x) -> ASTMapper m -> x -> m x
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
forall (m :: * -> *).
Monad m =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs
mapper :: (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs = ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper { mapOnExp :: Exp -> m Exp
mapOnExp = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m -> Exp -> m Exp) -> ASTMapper m -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
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
. (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a.
Substitutable a =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> a -> a
applySubst VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs
, mapOnPatternType :: PatternType -> m PatternType
mapOnPatternType = PatternType -> m PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> m PatternType)
-> (PatternType -> PatternType) -> PatternType -> m PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> PatternType -> PatternType
forall a.
Substitutable a =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> a -> a
applySubst VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs
}
shapeParam :: TypeParamBase vn -> PatternBase Info vn
shapeParam TypeParamBase vn
tp = vn -> Info PatternType -> SrcLoc -> PatternBase Info vn
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id (TypeParamBase vn -> vn
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase vn
tp) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
forall dim als. TypeBase dim als
i32) (SrcLoc -> PatternBase Info vn) -> SrcLoc -> PatternBase 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]
-> [PatternBase Info VName]
-> (TypeBase (DimDecl VName) (), [VName])
-> Exp
-> ValBind
toValBinding VName
name' [TypeParamBase VName]
tparams' [PatternBase Info VName]
params'' (TypeBase (DimDecl VName) (), [VName])
rettype' Exp
body'' =
ValBind :: forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (TypeBase (DimDecl VName) (), [VName])
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
ValBind { valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing
, valBindName :: VName
valBindName = VName
name'
, valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
retdecl
, valBindRetType :: Info (TypeBase (DimDecl VName) (), [VName])
valBindRetType = (TypeBase (DimDecl VName) (), [VName])
-> Info (TypeBase (DimDecl VName) (), [VName])
forall a. a -> Info a
Info (TypeBase (DimDecl VName) (), [VName])
rettype'
, valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
tparams'
, valBindParams :: [PatternBase Info VName]
valBindParams = [PatternBase Info VName]
params''
, valBindBody :: Exp
valBindBody = Exp
body''
, valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing
, valBindAttrs :: [AttrInfo]
valBindAttrs = [AttrInfo]
attrs
, valBindLocation :: SrcLoc
valBindLocation = SrcLoc
loc
}
typeSubstsM :: MonadFreshNames m =>
SrcLoc -> TypeBase () () -> MonoType
-> m (M.Map VName StructType, [TypeParam])
typeSubstsM :: SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName (TypeBase (DimDecl VName) ()),
[TypeParamBase VName])
typeSubstsM SrcLoc
loc TypeBase () ()
orig_t1 MonoType
orig_t2 =
let m :: StateT
(Map VName (TypeBase (DimDecl VName) ()))
(WriterT [TypeParamBase VName] m)
()
m = TypeBase () ()
-> MonoType
-> StateT
(Map VName (TypeBase (DimDecl VName) ()))
(WriterT [TypeParamBase VName] m)
()
forall d (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *) dim as.
(MonadState (Map VName (TypeBase (DimDecl VName) d)) (t (t m)),
MonadWriter [TypeParamBase VName] (t (t m)), MonadFreshNames m,
MonadTrans t, MonadTrans t, Pretty (ShapeDecl dim), Monad (t m)) =>
TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase () ()
orig_t1 MonoType
orig_t2
in WriterT
[TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
-> m (Map VName (TypeBase (DimDecl VName) ()),
[TypeParamBase VName])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT
[TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
-> m (Map VName (TypeBase (DimDecl VName) ()),
[TypeParamBase VName]))
-> WriterT
[TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
-> m (Map VName (TypeBase (DimDecl VName) ()),
[TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ StateT
(Map VName (TypeBase (DimDecl VName) ()))
(WriterT [TypeParamBase VName] m)
()
-> Map VName (TypeBase (DimDecl VName) ())
-> WriterT
[TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT StateT
(Map VName (TypeBase (DimDecl VName) ()))
(WriterT [TypeParamBase VName] m)
()
m Map VName (TypeBase (DimDecl VName) ())
forall a. Monoid a => a
mempty
where sub :: TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub t1 :: TypeBase dim as
t1@Array{} t2 :: TypeBase Bool d
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 Bool d
t2' <- Int -> TypeBase Bool d -> Maybe (TypeBase Bool d)
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 Bool d
t2 =
TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase dim as
t1' TypeBase Bool d
t2'
sub (Scalar (TypeVar as
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) TypeBase Bool d
t = TypeName -> TypeBase Bool d -> t (t m) ()
forall (t :: * -> * -> *) d (t :: (* -> *) -> * -> *)
(t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadState (Map VName (t (DimDecl VName) d)) (t (t m)),
MonadWriter [TypeParamBase VName] (t (t m)), MonadFreshNames m,
MonadTrans t, MonadTrans t, Bitraversable t, Monad (t m)) =>
TypeName -> t Bool d -> t (t m) ()
addSubst TypeName
v TypeBase Bool d
t
sub (Scalar (Record Map Name (TypeBase dim as)
fields1)) (Scalar (Record Map Name (TypeBase Bool d)
fields2)) =
(TypeBase dim as -> TypeBase Bool d -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase Bool d] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ TypeBase dim as -> TypeBase Bool d -> 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 Bool d) -> TypeBase Bool d)
-> [(Name, TypeBase Bool d)] -> [TypeBase Bool d]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase Bool d) -> TypeBase Bool d
forall a b. (a, b) -> b
snd ([(Name, TypeBase Bool d)] -> [TypeBase Bool d])
-> [(Name, TypeBase Bool d)] -> [TypeBase Bool d]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Bool d) -> [(Name, TypeBase Bool d)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase Bool d)
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 TypeBase dim as
t1b)) (Scalar (Arrow d
_ PName
_ TypeBase Bool d
t2a TypeBase Bool d
t2b)) = do
TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase dim as
t1a TypeBase Bool d
t2a
TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase dim as
t1b TypeBase Bool d
t2b
sub (Scalar (Sum Map Name [TypeBase dim as]
cs1)) (Scalar (Sum Map Name [TypeBase Bool d]
cs2)) =
((Name, [TypeBase dim as])
-> (Name, [TypeBase Bool d]) -> t (t m) [()])
-> [(Name, [TypeBase dim as])]
-> [(Name, [TypeBase Bool d])]
-> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Name, [TypeBase dim as])
-> (Name, [TypeBase Bool d]) -> t (t m) [()]
forall a a.
(a, [TypeBase dim as]) -> (a, [TypeBase Bool d]) -> 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 Bool d] -> [(Name, [TypeBase Bool d])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase Bool d]
cs2)
where typeSubstClause :: (a, [TypeBase dim as]) -> (a, [TypeBase Bool d]) -> t (t m) [()]
typeSubstClause (a
_, [TypeBase dim as]
ts1) (a
_, [TypeBase Bool d]
ts2) = (TypeBase dim as -> TypeBase Bool d -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase Bool d] -> t (t m) [()]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub [TypeBase dim as]
ts1 [TypeBase Bool d]
ts2
sub t1 :: TypeBase dim as
t1@(Scalar Sum{}) TypeBase Bool d
t2 = TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase dim as
t1 TypeBase Bool d
t2
sub TypeBase dim as
t1 t2 :: TypeBase Bool d
t2@(Scalar Sum{}) = TypeBase dim as -> TypeBase Bool d -> t (t m) ()
sub TypeBase dim as
t1 TypeBase Bool d
t2
sub TypeBase dim as
t1 TypeBase Bool d
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 Bool d -> String
forall a. Pretty a => a -> String
pretty TypeBase Bool d
t2]
addSubst :: TypeName -> t Bool d -> t (t m) ()
addSubst (TypeName [VName]
_ VName
v) t Bool d
t = do
Bool
exists <- (Map VName (t (DimDecl VName) d) -> Bool) -> t (t m) Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((Map VName (t (DimDecl VName) d) -> Bool) -> t (t m) Bool)
-> (Map VName (t (DimDecl VName) d) -> Bool) -> t (t m) Bool
forall a b. (a -> b) -> a -> b
$ VName -> Map VName (t (DimDecl VName) d) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member VName
v
Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
exists (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ do
t (DimDecl VName) d
t' <- (Bool -> t (t m) (DimDecl VName))
-> (d -> t (t m) d) -> t Bool d -> t (t m) (t (DimDecl VName) d)
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 Bool -> t (t m) (DimDecl VName)
forall (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
(m :: * -> *).
(MonadTrans t, MonadTrans t, Monad (t m), MonadFreshNames m,
MonadWriter [TypeParamBase VName] (t (t m))) =>
Bool -> t (t m) (DimDecl VName)
onDim d -> t (t m) d
forall (f :: * -> *) a. Applicative f => a -> f a
pure t Bool d
t
(Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d))
-> t (t m) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d))
-> t (t m) ())
-> (Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d))
-> t (t m) ()
forall a b. (a -> b) -> a -> b
$ VName
-> t (DimDecl VName) d
-> Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v t (DimDecl VName) d
t'
onDim :: Bool -> t (t m) (DimDecl VName)
onDim Bool
True = 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]
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 Bool
False = DimDecl VName -> t (t m) (DimDecl VName)
forall (m :: * -> *) a. Monad m => a -> m a
return DimDecl VName
forall vn. DimDecl vn
AnyDim
substPattern :: Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern :: Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry PatternType -> PatternType
f PatternBase Info VName
pat = case PatternBase Info VName
pat of
TuplePattern [PatternBase Info VName]
pats SrcLoc
loc -> [PatternBase Info VName] -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern ((PatternBase Info VName -> PatternBase Info VName)
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry PatternType -> PatternType
f) [PatternBase Info VName]
pats) SrcLoc
loc
RecordPattern [(Name, PatternBase Info VName)]
fs SrcLoc
loc -> [(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern (((Name, PatternBase Info VName) -> (Name, PatternBase Info VName))
-> [(Name, PatternBase Info VName)]
-> [(Name, PatternBase Info VName)]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternBase Info VName) -> (Name, PatternBase Info VName)
forall a.
(a, PatternBase Info VName) -> (a, PatternBase Info VName)
substField [(Name, PatternBase Info VName)]
fs) SrcLoc
loc
where substField :: (a, PatternBase Info VName) -> (a, PatternBase Info VName)
substField (a
n, PatternBase Info VName
p) = (a
n, Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry PatternType -> PatternType
f PatternBase Info VName
p)
PatternParens PatternBase Info VName
p SrcLoc
loc -> PatternBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> SrcLoc -> PatternBase f vn
PatternParens (Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry PatternType -> PatternType
f PatternBase Info VName
p) SrcLoc
loc
Id VName
vn (Info PatternType
tp) SrcLoc
loc -> VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
vn (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
Wildcard (Info PatternType
tp) SrcLoc
loc -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
PatternAscription PatternBase Info VName
p TypeDeclBase Info VName
td SrcLoc
loc | Bool
entry -> PatternBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatternBase f vn
PatternAscription (Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
False PatternType -> PatternType
f PatternBase Info VName
p) TypeDeclBase Info VName
td SrcLoc
loc
| Bool
otherwise -> Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
False PatternType -> PatternType
f PatternBase Info VName
p
PatternLit Exp
e (Info PatternType
tp) SrcLoc
loc -> Exp -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> f PatternType -> SrcLoc -> PatternBase f vn
PatternLit Exp
e (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
PatternConstr Name
n (Info PatternType
tp) [PatternBase Info VName]
ps SrcLoc
loc -> Name
-> Info PatternType
-> [PatternBase Info VName]
-> SrcLoc
-> PatternBase Info VName
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
n (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) [PatternBase Info VName]
ps SrcLoc
loc
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding (ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
retdecl (Info (TypeBase (DimDecl VName) ()
rettype, [VName]
retext)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params Exp
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) =
RecordReplacements
-> (VName, [TypeParamBase VName], [PatternBase Info VName],
Maybe (TypeExp VName), TypeBase (DimDecl VName) (), [VName], Exp,
[AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [PatternBase Info VName]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
rettype, [VName]
retext, Exp
body, [AttrInfo]
attrs, SrcLoc
loc)
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables Bool
entry valbind :: ValBind
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
_ Maybe (TypeExp VName)
_ (Info (TypeBase (DimDecl VName) ()
rettype, [VName]
retext)) [TypeParamBase VName]
_ [PatternBase Info VName]
pats Exp
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
_) = do
Map VName TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
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) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper {
mapOnExp :: Exp -> MonoM Exp
mapOnExp = ASTMapper MonoM -> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper
, 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
. Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs
, mapOnPatternType :: PatternType -> MonoM PatternType
mapOnPatternType = PatternType -> MonoM PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> MonoM PatternType)
-> (PatternType -> PatternType) -> PatternType -> MonoM PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName TypeSub -> PatternType -> PatternType
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs
}
Exp
body' <- ASTMapper MonoM -> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper Exp
body
ValBind -> MonoM ValBind
forall (m :: * -> *) a. Monad m => a -> m a
return ValBind
valbind { valBindRetType :: Info (TypeBase (DimDecl VName) (), [VName])
valBindRetType = (TypeBase (DimDecl VName) (), [VName])
-> Info (TypeBase (DimDecl VName) (), [VName])
forall a. a -> Info a
Info (Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs TypeBase (DimDecl VName) ()
rettype, [VName]
retext)
, valBindParams :: [PatternBase Info VName]
valBindParams = (PatternBase Info VName -> PatternBase Info VName)
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (Bool
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
substPattern Bool
entry ((PatternType -> PatternType)
-> PatternBase Info VName -> PatternBase Info VName)
-> (PatternType -> PatternType)
-> PatternBase Info VName
-> PatternBase Info VName
forall a b. (a -> b) -> a -> b
$ Map VName TypeSub -> PatternType -> PatternType
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs) [PatternBase Info VName]
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 TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
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
$ Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
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) ()]
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType
((PatternBase Info VName -> TypeBase (DimDecl VName) ())
-> [PatternBase Info VName] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map PatternBase Info VName -> TypeBase (DimDecl VName) ()
patternStructType (ValBind -> [PatternBase Info VName]
forall (f :: * -> *) vn. ValBindBase f vn -> [PatternBase f vn]
valBindParams ValBind
valbind)) (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$
(TypeBase (DimDecl VName) (), [VName])
-> TypeBase (DimDecl VName) ()
forall a b. (a, b) -> a
fst ((TypeBase (DimDecl VName) (), [VName])
-> TypeBase (DimDecl VName) ())
-> (TypeBase (DimDecl VName) (), [VName])
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) (), [VName])
-> (TypeBase (DimDecl VName) (), [VName])
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) (), [VName])
-> (TypeBase (DimDecl VName) (), [VName]))
-> Info (TypeBase (DimDecl VName) (), [VName])
-> (TypeBase (DimDecl VName) (), [VName])
forall a b. (a -> b) -> a -> b
$ ValBind -> Info (TypeBase (DimDecl VName) (), [VName])
forall (f :: * -> *) vn.
ValBindBase f vn -> f (TypeBase (DimDecl VName) (), [VName])
valBindRetType ValBind
valbind
(VName
name, InferSizeArgs
_, 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})
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 TypeDeclBase Info VName
tydecl Maybe DocComment
_ SrcLoc
_) = do
Map VName TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
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) () -> MonoM ())
-> TypeBase (DimDecl VName) () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info (TypeBase (DimDecl VName) ())
forall (f :: * -> *) vn.
TypeDeclBase f vn -> f (TypeBase (DimDecl VName) ())
expandedType TypeDeclBase Info VName
tydecl
let tp :: TypeBase (DimDecl VName) ()
tp = Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> (Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info (TypeBase (DimDecl VName) ())
forall (f :: * -> *) vn.
TypeDeclBase f vn -> f (TypeBase (DimDecl VName) ())
expandedType TypeDeclBase Info VName
tydecl
tbinding :: TypeBinding
tbinding = Liftedness
-> [TypeParamBase VName]
-> TypeBase (DimDecl VName) ()
-> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams TypeBase (DimDecl VName) ()
tp
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 -> String -> String
forall a. [a] -> [a] -> [a]
++
String
"input program, but received: " String -> String -> String
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