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

-- | This monomorphization module converts a well-typed, polymorphic,
-- module-free Futhark program into an equivalent monomorphic program.
--
-- This pass also does a few other simplifications to make the job of
-- subsequent passes easier.  Specifically, it does the following:
--
-- * Turn operator sections into explicit lambdas.
--
-- * Converts identifiers of record type into record patterns (and
--   similarly for tuples).
--
-- * Converts applications of intrinsic SOACs into SOAC AST nodes
--   (Map, Reduce, etc).
--
-- * Elide functions that are not reachable from an entry point (this
--   is a side effect of the monomorphisation algorithm, which uses
--   the entry points as roots).
--
-- * Turns implicit record fields into explicit record fields.
--
-- * Rewrite BinOp nodes to Apply nodes.
--
-- Note that these changes are unfortunately not visible in the AST
-- representation.
module Futhark.Internalise.Monomorphise (transformProg) where

import Control.Monad.Identity
import Control.Monad.RWS hiding (Sum)
import Control.Monad.State
import Control.Monad.Writer hiding (Sum)
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition)
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Sequence as Seq
import qualified Data.Set as S
import Futhark.MonadFreshNames
import Futhark.Util.Pretty
import Language.Futhark
import Language.Futhark.Semantic (TypeBinding (..))
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types

i64 :: TypeBase dim als
i64 :: TypeBase dim als
i64 = ScalarTypeBase dim als -> TypeBase dim als
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim als -> TypeBase dim als)
-> ScalarTypeBase dim als -> TypeBase dim als
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim als
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase dim als)
-> PrimType -> ScalarTypeBase dim als
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- The monomorphization monad reads 'PolyBinding's and writes
-- 'ValBind's.  The 'TypeParam's in the 'ValBind's can only be size
-- parameters.
--
-- Each 'Polybinding' is also connected with the 'RecordReplacements'
-- that were active when the binding was defined.  This is used only
-- in local functions.
data PolyBinding
  = PolyBinding
      RecordReplacements
      ( VName,
        [TypeParam],
        [Pat],
        StructRetType,
        Exp,
        [AttrInfo VName],
        SrcLoc
      )

-- Mapping from record names to the variable names that contain the
-- fields.  This is used because the monomorphiser also expands all
-- record patterns.
type RecordReplacements = M.Map VName RecordReplacement

type RecordReplacement = M.Map Name (VName, PatType)

-- Monomorphization environment mapping names of polymorphic functions
-- to a representation of their corresponding function bindings.
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}

-- The monomorphization monad.
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 (f :: * -> *) a. Applicative f => a -> f a
pure (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 (f :: * -> *) a. Applicative f => a -> f a
pure 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

-- Given instantiated type of function, produce size arguments.
type InferSizeArgs = StructType -> [Exp]

data MonoSize
  = -- | The integer encodes an equivalence class, so we can keep
    -- track of sizes that are statically identical.
    MonoKnown Int
  | MonoAnon VName
  deriving (Int -> MonoSize -> ShowS
[MonoSize] -> ShowS
MonoSize -> String
(Int -> MonoSize -> ShowS)
-> (MonoSize -> String) -> ([MonoSize] -> ShowS) -> Show MonoSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MonoSize] -> ShowS
$cshowList :: [MonoSize] -> ShowS
show :: MonoSize -> String
$cshow :: MonoSize -> String
showsPrec :: Int -> MonoSize -> ShowS
$cshowsPrec :: Int -> MonoSize -> ShowS
Show)

-- We treat all MonoAnon as identical.
instance Eq MonoSize where
  MonoKnown Int
x == :: MonoSize -> MonoSize -> Bool
== MonoKnown Int
y = Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y
  MonoAnon VName
_ == MonoAnon VName
_ = Bool
True
  MonoSize
_ == MonoSize
_ = Bool
False

instance Pretty MonoSize where
  ppr :: MonoSize -> Doc
ppr (MonoKnown Int
i) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
i
  ppr (MonoAnon VName
v) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
v

instance Pretty (ShapeDecl MonoSize) where
  ppr :: ShapeDecl MonoSize -> Doc
ppr (ShapeDecl [MonoSize]
ds) = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((MonoSize -> Doc) -> [MonoSize] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (MonoSize -> Doc) -> MonoSize -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MonoSize -> Doc
forall a. Pretty a => a -> Doc
ppr) [MonoSize]
ds)

-- The kind of type relative to which we monomorphise.  What is most
-- important to us is not the specific dimensions, but merely whether
-- they are known or anonymous/local.
type MonoType = TypeBase MonoSize ()

monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType = (State (Int, Map (DimDecl VName) Int) MonoType
-> (Int, Map (DimDecl VName) Int) -> MonoType
forall s a. State s a -> s -> a
`evalState` (Int
0, Map (DimDecl VName) Int
forall a. Monoid a => a
mempty)) (State (Int, Map (DimDecl VName) Int) MonoType -> MonoType)
-> (TypeBase (DimDecl VName) als
    -> State (Int, Map (DimDecl VName) Int) MonoType)
-> TypeBase (DimDecl VName) als
-> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
 -> DimPos
 -> DimDecl VName
 -> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize)
-> TypeBase (DimDecl VName) ()
-> State (Int, Map (DimDecl VName) Int) MonoType
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos
-> DimDecl VName
-> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize
forall (f :: * -> *) p.
MonadState (Int, Map (DimDecl VName) Int) f =>
Set VName -> p -> DimDecl VName -> f MonoSize
onDim (TypeBase (DimDecl VName) ()
 -> State (Int, Map (DimDecl VName) Int) MonoType)
-> (TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) als
-> State (Int, Map (DimDecl VName) Int) MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct
  where
    onDim :: Set VName -> p -> DimDecl VName -> f MonoSize
onDim Set VName
bound p
_ (NamedDim QualName VName
d)
      -- A locally bound size.
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ VName -> MonoSize
MonoAnon (VName -> MonoSize) -> VName -> MonoSize
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    onDim Set VName
_ p
_ DimDecl VName
d = do
      (Int
i, Map (DimDecl VName) Int
m) <- f (Int, Map (DimDecl VName) Int)
forall s (m :: * -> *). MonadState s m => m s
get
      case DimDecl VName -> Map (DimDecl VName) Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup DimDecl VName
d Map (DimDecl VName) Int
m of
        Just Int
prev ->
          MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
prev
        Maybe Int
Nothing -> do
          (Int, Map (DimDecl VName) Int) -> f ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, DimDecl VName
-> Int -> Map (DimDecl VName) Int -> Map (DimDecl VName) Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert DimDecl VName
d Int
i Map (DimDecl VName) Int
m)
          MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
i

-- Mapping from function name and instance list to a new function name in case
-- the function has already been instantiated with those concrete types.
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 (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Exp
forall vn. QualName vn -> ExpBase Info vn
var QualName VName
fname
  | Bool
otherwise = do
      TypeBase (DimDecl VName) ()
t' <- TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t
      let mono_t :: MonoType
mono_t = TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t'
      Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t
      Maybe PolyBinding
maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
lookupFun (VName -> MonoM (Maybe PolyBinding))
-> VName -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
      case (Maybe (VName, InferSizeArgs)
maybe_fname, Maybe PolyBinding
maybe_funbind) of
        -- The function has already been monomorphised.
        (Just (VName
fname', InferSizeArgs
infer), Maybe PolyBinding
_) ->
          Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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'
        -- An intrinsic function.
        (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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
        -- A polymorphic function.
        (Maybe (VName, InferSizeArgs)
Nothing, Just PolyBinding
funbind) -> do
          (VName
fname', InferSizeArgs
infer, ValBind
funbind') <- Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
False PolyBinding
funbind MonoType
mono_t
          Seq (VName, ValBind) -> MonoM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (VName, ValBind) -> MonoM ())
-> Seq (VName, ValBind) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (VName, ValBind) -> Seq (VName, ValBind)
forall a. a -> Seq a
Seq.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname, ValBind
funbind')
          VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t (VName
fname', InferSizeArgs
infer)
          Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase (DimDecl VName) () -> [Exp] -> Exp
forall vn as.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
  where
    var :: QualName vn -> ExpBase Info vn
var QualName vn
fname' = QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (PatType -> Info PatType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t)) SrcLoc
loc

    applySizeArg :: (Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg (Int
i, ExpBase Info vn
f) ExpBase Info vn
size_arg =
      ( Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1,
        AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          (ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
f ExpBase Info vn
size_arg ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing)) SrcLoc
loc)
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes ([TypeBase (DimDecl VName) Any]
-> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (Int
-> TypeBase (DimDecl VName) Any -> [TypeBase (DimDecl VName) Any]
forall a. Int -> a -> [a]
replicate Int
i TypeBase (DimDecl VName) Any
forall dim als. TypeBase dim als
i64) ([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t))) [])
      )

    applySizeArgs :: vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs vn
fname' TypeBase (DimDecl VName) as
t' [ExpBase Info vn]
size_args =
      (Int, ExpBase Info vn) -> ExpBase Info vn
forall a b. (a, b) -> b
snd ((Int, ExpBase Info vn) -> ExpBase Info vn)
-> (Int, ExpBase Info vn) -> ExpBase Info vn
forall a b. (a -> b) -> a -> b
$
        ((Int, ExpBase Info vn)
 -> ExpBase Info vn -> (Int, ExpBase Info vn))
-> (Int, ExpBase Info vn)
-> [ExpBase Info vn]
-> (Int, ExpBase Info vn)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          (Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
forall vn.
(Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg
          ( [ExpBase Info vn] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info vn]
size_args Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1,
            QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
              (vn -> QualName vn
forall v. v -> QualName v
qualName vn
fname')
              ( PatType -> Info PatType
forall a. a -> Info a
Info
                  ( [TypeBase (DimDecl VName) Any]
-> RetTypeBase (DimDecl VName) Aliasing -> PatType
forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType
                      ((ExpBase Info vn -> TypeBase (DimDecl VName) Any)
-> [ExpBase Info vn] -> [TypeBase (DimDecl VName) Any]
forall a b. (a -> b) -> [a] -> [b]
map (TypeBase (DimDecl VName) Any
-> ExpBase Info vn -> TypeBase (DimDecl VName) Any
forall a b. a -> b -> a
const TypeBase (DimDecl VName) Any
forall dim als. TypeBase dim als
i64) [ExpBase Info vn]
size_args)
                      ([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (PatType -> RetTypeBase (DimDecl VName) Aliasing)
-> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) as -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) as
t')
                  )
              )
              SrcLoc
loc
          )
          [ExpBase Info vn]
size_args

-- This carries out record replacements in the alias information of a type.
transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType TypeBase dim Aliasing
t = do
  RecordReplacements
rrs <- (Env -> RecordReplacements) -> MonoM RecordReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> RecordReplacements
envRecordReplacements
  let replace :: Alias -> Aliasing
replace (AliasBound VName
v)
        | Just RecordReplacement
d <- VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v RecordReplacements
rrs =
            [Alias] -> Aliasing
forall a. Ord a => [a] -> Set a
S.fromList ([Alias] -> Aliasing) -> [Alias] -> Aliasing
forall a b. (a -> b) -> a -> b
$ ((VName, PatType) -> Alias) -> [(VName, PatType)] -> [Alias]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound (VName -> Alias)
-> ((VName, PatType) -> VName) -> (VName, PatType) -> Alias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, PatType) -> VName
forall a b. (a, b) -> a
fst) ([(VName, PatType)] -> [Alias]) -> [(VName, PatType)] -> [Alias]
forall a b. (a -> b) -> a -> b
$ RecordReplacement -> [(VName, PatType)]
forall k a. Map k a -> [a]
M.elems RecordReplacement
d
      replace Alias
x = Alias -> Aliasing
forall a. a -> Set a
S.singleton Alias
x
  -- As an attempt at an optimisation, only transform the aliases if
  -- they refer to a variable we have record-replaced.
  TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing))
-> TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall a b. (a -> b) -> a -> b
$
    if (Alias -> Bool) -> Aliasing -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> RecordReplacements -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` RecordReplacements
rrs) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (Aliasing -> Bool) -> Aliasing -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t
      then (Aliasing -> Aliasing)
-> TypeBase dim Aliasing -> TypeBase dim Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([Aliasing] -> Aliasing
forall a. Monoid a => [a] -> a
mconcat ([Aliasing] -> Aliasing)
-> (Aliasing -> [Aliasing]) -> Aliasing -> Aliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Aliasing) -> [Alias] -> [Aliasing]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> Aliasing
replace ([Alias] -> [Aliasing])
-> (Aliasing -> [Alias]) -> Aliasing -> [Aliasing]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliasing -> [Alias]
forall a. Set a -> [a]
S.toList) TypeBase dim Aliasing
t
      else TypeBase dim Aliasing
t

sizesForPat :: MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat :: Pat -> m ([VName], Pat)
sizesForPat Pat
pat = do
  (Pat
params', [VName]
sizes) <- StateT [VName] m Pat -> [VName] -> m (Pat, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ASTMapper (StateT [VName] m) -> Pat -> StateT [VName] m Pat
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT [VName] m)
tv Pat
pat) []
  ([VName], Pat) -> m ([VName], Pat)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName]
sizes, Pat
params')
  where
    tv :: ASTMapper (StateT [VName] m)
tv = ASTMapper (StateT [VName] m)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatType :: PatType -> StateT [VName] m PatType
mapOnPatType = (DimDecl VName -> StateT [VName] m (DimDecl VName))
-> (Aliasing -> StateT [VName] m Aliasing)
-> PatType
-> StateT [VName] m PatType
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse DimDecl VName -> StateT [VName] m (DimDecl VName)
forall (t :: (* -> *) -> * -> *) (m :: * -> *).
(MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
DimDecl VName -> t m (DimDecl VName)
onDim Aliasing -> StateT [VName] m Aliasing
forall (f :: * -> *) a. Applicative f => a -> f a
pure}
    onDim :: DimDecl VName -> t m (DimDecl VName)
onDim (AnyDim Maybe VName
_) = do
      VName
v <- m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"size"
      ([VName] -> [VName]) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:)
      DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
    onDim DimDecl VName
d = DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d

transformAppRes :: AppRes -> MonoM AppRes
transformAppRes :: AppRes -> MonoM AppRes
transformAppRes (AppRes PatType
t [VName]
ext) =
  PatType -> [VName] -> AppRes
AppRes (PatType -> [VName] -> AppRes)
-> MonoM PatType -> MonoM ([VName] -> AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t MonoM ([VName] -> AppRes) -> MonoM [VName] -> MonoM AppRes
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [VName] -> MonoM [VName]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VName]
ext

transformAppExp :: AppExp -> AppRes -> MonoM Exp
transformAppExp :: AppExp -> AppRes -> MonoM Exp
transformAppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) AppRes
res = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Maybe Exp
me' <- (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Maybe Exp
me
  Inclusiveness Exp
incl' <- (Exp -> MonoM Exp)
-> Inclusiveness Exp -> MonoM (Inclusiveness Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Inclusiveness Exp
incl
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> Maybe Exp -> Inclusiveness Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Coerce Exp
e TypeExp VName
tp SrcLoc
loc) AppRes
res =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> TypeExp VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> AppExpBase f vn
Coerce (Exp -> TypeExp VName -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (TypeExp VName -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeExp VName -> SrcLoc -> AppExp)
-> MonoM (TypeExp VName) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp VName -> MonoM (TypeExp VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp VName
tp MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (LetPat [SizeBinder VName]
sizes Pat
pat Exp
e1 Exp
e2 SrcLoc
loc) AppRes
res = do
  (Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
    (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes Pat
pat' (Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
            MonoM (Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (Exp -> MonoM Exp
transformExp Exp
e2)
            MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
        )
    MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat]
params, Maybe (TypeExp VName)
retdecl, Info StructRetType
ret, Exp
body) Exp
e SrcLoc
loc) AppRes
res
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParamBase VName]
tparams = do
      -- Retrieve the lifted monomorphic function bindings that are produced,
      -- filter those that are monomorphic versions of the current let-bound
      -- function and insert them at this point, and propagate the rest.
      RecordReplacements
rr <- (Env -> RecordReplacements) -> MonoM RecordReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> RecordReplacements
envRecordReplacements
      let funbind :: PolyBinding
funbind = RecordReplacements
-> (VName, [TypeParamBase VName], [Pat], StructRetType, Exp,
    [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
ret, Exp
body, [AttrInfo VName]
forall a. Monoid a => a
mempty, SrcLoc
loc)
      MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
 -> MonoM Exp)
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall a b. (a -> b) -> a -> b
$ do
        (Exp
e', Seq (VName, ValBind)
bs) <- MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (MonoM Exp -> MonoM (Exp, Seq (VName, ValBind)))
-> MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ VName -> PolyBinding -> MonoM Exp -> MonoM Exp
forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
fname PolyBinding
funbind (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
e
        -- Do not remember this one for next time we monomorphise this
        -- function.
        (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 (f :: * -> *) a. Applicative f => a -> f a
pure ([ValBind] -> Exp -> Exp
unfoldLetFuns (((VName, ValBind) -> ValBind) -> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> [a] -> [b]
map (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd ([(VName, ValBind)] -> [ValBind])
-> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> a -> b
$ Seq (VName, ValBind) -> [(VName, ValBind)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq (VName, ValBind)
bs_local) Exp
e', Seq (VName, ValBind)
-> Seq (VName, ValBind) -> Seq (VName, ValBind)
forall a b. a -> b -> a
const Seq (VName, ValBind)
bs_prop)
  | Bool
otherwise = do
      Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
      AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName
-> ([TypeParamBase VName], [Pat], Maybe (TypeExp VName),
    Info StructRetType, Exp)
-> Exp
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn], Maybe (TypeExp vn),
    f StructRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat]
params, Maybe (TypeExp VName)
retdecl, StructRetType -> Info StructRetType
forall a. a -> Info a
Info StructRetType
ret, Exp
body') (Exp -> SrcLoc -> AppExp) -> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
        MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) AppRes
res =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If (Exp -> Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Exp -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e3 MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d SrcLoc
loc) AppRes
res =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply (Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info (Diet, Maybe VName) -> SrcLoc -> AppExp)
-> MonoM (Info (Diet, Maybe VName)) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info (Diet, Maybe VName) -> MonoM (Info (Diet, Maybe VName))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info (Diet, Maybe VName)
d MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (DoLoop [VName]
sparams Pat
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 SrcLoc
loc) AppRes
res = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  LoopFormBase Info VName
form' <- case LoopFormBase Info VName
form of
    For IdentBase Info VName
ident Exp
e2 -> IdentBase Info VName -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
ident (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
    ForIn Pat
pat2 Exp
e2 -> Pat -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn Pat
pat2 (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
    While Exp
e2 -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
  Exp
e3' <- Exp -> MonoM Exp
transformExp Exp
e3
  -- Maybe monomorphisation introduced new arrays to the loop, and
  -- maybe they have AnyDim sizes.  This is not allowed.  Invent some
  -- sizes for them.
  ([VName]
pat_sizes, Pat
pat') <- Pat -> MonoM ([VName], Pat)
forall (m :: * -> *). MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat Pat
pat
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([VName]
-> Pat -> Exp -> LoopFormBase Info VName -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[VName]
-> PatBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
DoLoop ([VName]
sparams [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
pat_sizes) Pat
pat' Exp
e1' LoopFormBase Info VName
form' Exp
e3' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (BinOp (QualName VName
fname, SrcLoc
_) (Info PatType
t) (Exp
e1, Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2, Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) SrcLoc
loc) (AppRes PatType
ret [VName]
ext) = do
  Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
  if PatType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e1') Bool -> Bool -> Bool
&& PatType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatType
typeOf Exp
e2')
    then Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
e1' Exp
e2'
    else do
      -- We have to flip the arguments to the function, because
      -- operator application is left-to-right, while function
      -- application is outside-in.  This matters when the arguments
      -- produce existential sizes.  There are later places in the
      -- compiler where we transform BinOp to Apply, but anything that
      -- involves existential sizes will necessarily go through here.
      (Exp
x_param_e, Pat
x_param) <- Exp -> MonoM (Exp, Pat)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pat)
makeVarParam Exp
e1'
      (Exp
y_param_e, Pat
y_param) <- Exp -> MonoM (Exp, Pat)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pat)
makeVarParam Exp
e2'
      -- XXX: the type annotations here are wrong, but hopefully it
      -- doesn't matter as there will be an outer AppExp to handle
      -- them.
      Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
        AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( [SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
              []
              Pat
x_param
              Exp
e1'
              ( AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
                  ([SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] Pat
y_param Exp
e2' (Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
x_param_e Exp
y_param_e) SrcLoc
loc)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty)
              )
              SrcLoc
forall a. Monoid a => a
mempty
          )
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty))
  where
    applyOp :: ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp ExpBase Info vn
fname' ExpBase Info vn
x ExpBase Info vn
y =
      AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        ( ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
            ( AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
                (ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
fname' ExpBase Info vn
x ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d1))) SrcLoc
loc)
                (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
forall a. Monoid a => a
mempty)
            )
            ExpBase Info vn
y
            ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d2)))
            SrcLoc
loc
        )
        (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
ret [VName]
ext))

    makeVarParam :: Exp -> m (Exp, Pat)
makeVarParam Exp
arg = do
      let argtype :: PatType
argtype = Exp -> PatType
typeOf Exp
arg
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"binop_p"
      (Exp, Pat) -> m (Exp, Pat)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty
        )
transformAppExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs Exp
e1 Exp
body SrcLoc
loc) AppRes
res = do
  SliceBase Info VName
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (IdentBase Info VName
-> IdentBase Info VName
-> SliceBase Info VName
-> Exp
-> Exp
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
IdentBase f vn
-> IdentBase f vn
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 SliceBase Info VName
idxs' Exp
e1' Exp
body' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) AppRes
res =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
    (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SliceBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index (Exp -> SliceBase Info VName -> SrcLoc -> AppExp)
-> MonoM Exp -> MonoM (SliceBase Info VName -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e0 MonoM (SliceBase Info VName -> SrcLoc -> AppExp)
-> MonoM (SliceBase Info VName) -> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
    MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
transformAppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) AppRes
res =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
    (AppExp -> Info AppRes -> Exp)
-> MonoM AppExp -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match (Exp -> NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
-> MonoM Exp
-> MonoM (NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (NonEmpty (CaseBase Info VName) -> SrcLoc -> AppExp)
-> MonoM (NonEmpty (CaseBase Info VName))
-> MonoM (SrcLoc -> AppExp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (CaseBase Info VName -> MonoM (CaseBase Info VName))
-> NonEmpty (CaseBase Info VName)
-> MonoM (NonEmpty (CaseBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase NonEmpty (CaseBase Info VName)
cs MonoM (SrcLoc -> AppExp) -> MonoM SrcLoc -> MonoM AppExp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
    MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)

-- Monomorphization of expressions.
transformExp :: Exp -> MonoM Exp
transformExp :: Exp -> MonoM Exp
transformExp e :: Exp
e@Literal {} = Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@IntLit {} = Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@FloatLit {} = Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@StringLit {} = Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp (Parens Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) =
  (QualName VName, SrcLoc) -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (TupLit [Exp]
es SrcLoc
loc) =
  [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp) -> MonoM [Exp] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> [FieldBase Info VName] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField [FieldBase Info VName]
fs MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  where
    transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (RecordFieldExplicit Name
name Exp
e SrcLoc
loc') =
      Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name (Exp -> SrcLoc -> FieldBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> FieldBase Info VName)
-> MonoM SrcLoc -> MonoM (FieldBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc'
    transformField (RecordFieldImplicit VName
v Info PatType
t SrcLoc
_) = do
      Info PatType
t' <- (PatType -> MonoM PatType) -> Info PatType -> MonoM (Info PatType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatType
t
      FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> FieldBase Info VName -> MonoM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$
        Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
          (VName -> Name
baseName VName
v)
          (QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatType
t' SrcLoc
loc)
          SrcLoc
loc
transformExp (ArrayLit [Exp]
es Info PatType
t SrcLoc
loc) =
  [Exp] -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
ArrayLit ([Exp] -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatType -> MonoM PatType) -> Info PatType -> MonoM (Info PatType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (AppExp AppExp
e Info AppRes
res) = do
  PatType -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims (PatType -> MonoM ()) -> PatType -> MonoM ()
forall a b. (a -> b) -> a -> b
$ AppRes -> PatType
appResType (AppRes -> PatType) -> AppRes -> PatType
forall a b. (a -> b) -> a -> b
$ Info AppRes -> AppRes
forall a. Info a -> a
unInfo Info AppRes
res
  AppExp -> AppRes -> MonoM Exp
transformAppExp AppExp
e (AppRes -> MonoM Exp) -> MonoM AppRes -> MonoM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< AppRes -> MonoM AppRes
transformAppRes (Info AppRes -> AppRes
forall a. Info a -> a
unInfo Info AppRes
res)
transformExp (Var QualName VName
fname (Info PatType
t) SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (VName -> MonoM (Maybe RecordReplacement))
-> VName -> MonoM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
  case Maybe RecordReplacement
maybe_fs of
    Just RecordReplacement
fs -> do
      let toField :: (Name, (vn, PatType)) -> MonoM (FieldBase Info vn)
toField (Name
f, (vn
f_v, PatType
f_t)) = do
            PatType
f_t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
f_t
            let f_v' :: ExpBase Info vn
f_v' = QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
f_v) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
f_t') SrcLoc
loc
            FieldBase Info vn -> MonoM (FieldBase Info vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FieldBase Info vn -> MonoM (FieldBase Info vn))
-> FieldBase Info vn -> MonoM (FieldBase Info vn)
forall a b. (a -> b) -> a -> b
$ Name -> ExpBase Info vn -> SrcLoc -> FieldBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
f ExpBase Info vn
f_v' SrcLoc
loc
      [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, (VName, PatType)) -> MonoM (FieldBase Info VName))
-> [(Name, (VName, PatType))] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, (VName, PatType)) -> MonoM (FieldBase Info VName)
forall vn. (Name, (vn, PatType)) -> MonoM (FieldBase Info vn)
toField (RecordReplacement -> [(Name, (VName, PatType))]
forall k a. Map k a -> [(k, a)]
M.toList RecordReplacement
fs) MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    Maybe RecordReplacement
Nothing -> do
      PatType
t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t
      SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t')
transformExp (Hole Info PatType
t SrcLoc
loc) =
  Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f PatType -> SrcLoc -> ExpBase f vn
Hole (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatType -> MonoM PatType) -> Info PatType -> MonoM (Info PatType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Ascript Exp
e TypeExp VName
tp SrcLoc
loc) =
  Exp -> TypeExp VName -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp vn -> SrcLoc -> ExpBase f vn
Ascript (Exp -> TypeExp VName -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (TypeExp VName -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeExp VName -> SrcLoc -> Exp)
-> MonoM (TypeExp VName) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp VName -> MonoM (TypeExp VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp VName
tp MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Negate Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Not Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Lambda [Pat]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Aliasing, StructRetType)
tp SrcLoc
loc) = do
  Exp
e0' <- Exp -> MonoM Exp
transformExp Exp
e0
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [Pat]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Aliasing, StructRetType)
tp SrcLoc
loc
transformExp (OpSection QualName VName
qn Info PatType
t SrcLoc
loc) =
  Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info PatType
t SrcLoc
loc
transformExp (OpSectionLeft QualName VName
fname (Info PatType
t) Exp
e (Info (PName, TypeBase (DimDecl VName) (), Maybe VName),
 Info (PName, TypeBase (DimDecl VName) ()))
arg (Info RetTypeBase (DimDecl VName) Aliasing
rettype, Info [VName]
retext) SrcLoc
loc) = do
  let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype)) = (Info (PName, TypeBase (DimDecl VName) (), Maybe VName),
 Info (PName, TypeBase (DimDecl VName) ()))
arg
  Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
    Exp
fname'
    (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
    Maybe Exp
forall a. Maybe a
Nothing
    PatType
t
    (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext)
    (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
forall a. Maybe a
Nothing)
    (RetTypeBase (DimDecl VName) Aliasing
rettype, [VName]
retext)
    SrcLoc
loc
transformExp (OpSectionRight QualName VName
fname (Info PatType
t) Exp
e (Info (PName, TypeBase (DimDecl VName) ()),
 Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg (Info RetTypeBase (DimDecl VName) Aliasing
rettype) SrcLoc
loc) = do
  let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)) = (Info (PName, TypeBase (DimDecl VName) ()),
 Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg
  Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
    Exp
fname'
    Maybe Exp
forall a. Maybe a
Nothing
    (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
    PatType
t
    (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
forall a. Maybe a
Nothing)
    (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)
    (RetTypeBase (DimDecl VName) Aliasing
rettype, [])
    SrcLoc
loc
transformExp (ProjectSection [Name]
fields (Info PatType
t) SrcLoc
loc) =
  [Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields PatType
t SrcLoc
loc
transformExp (IndexSection SliceBase Info VName
idxs (Info PatType
t) SrcLoc
loc) = do
  SliceBase Info VName
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
  SliceBase Info VName -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection SliceBase Info VName
idxs' PatType
t SrcLoc
loc
transformExp (Project Name
n Exp
e Info PatType
tp SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- case Exp
e of
    Var QualName VName
qn Info PatType
_ SrcLoc
_ -> VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
    Exp
_ -> Maybe RecordReplacement -> MonoM (Maybe RecordReplacement)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe RecordReplacement
forall a. Maybe a
Nothing
  case Maybe RecordReplacement
maybe_fs of
    Just RecordReplacement
m
      | Just (VName
v, PatType
_) <- Name -> RecordReplacement -> Maybe (VName, PatType)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n RecordReplacement
m ->
          Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatType
tp SrcLoc
loc
    Maybe RecordReplacement
_ -> do
      Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
      Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info PatType
tp SrcLoc
loc
transformExp (Update Exp
e1 SliceBase Info VName
idxs Exp
e2 SrcLoc
loc) =
  Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM (SliceBase Info VName) -> MonoM (Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> SliceBase Info VName -> MonoM (SliceBase Info VName)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
    MonoM (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info PatType
t SrcLoc
loc) =
  Exp -> [Name] -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
RecordUpdate (Exp -> [Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([Name] -> Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Name] -> MonoM (Exp -> Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Name] -> MonoM [Name]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
fs
    MonoM (Exp -> Info PatType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatType -> MonoM (Info PatType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatType
t
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Assert Exp
e1 Exp
e2 Info String
desc SrcLoc
loc) =
  Exp -> Exp -> Info String -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f String -> SrcLoc -> ExpBase f vn
Assert (Exp -> Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Exp -> Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info String -> SrcLoc -> Exp)
-> MonoM (Info String) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info String -> MonoM (Info String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info String
desc MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Constr Name
name [Exp]
all_es Info PatType
t SrcLoc
loc) =
  Name -> [Exp] -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatType -> SrcLoc -> ExpBase f vn
Constr Name
name ([Exp] -> Info PatType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
all_es MonoM (Info PatType -> SrcLoc -> Exp)
-> MonoM (Info PatType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatType -> MonoM (Info PatType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) =
  AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
info (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

transformCase :: Case -> MonoM Case
transformCase :: CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase (CasePat Pat
p Exp
e SrcLoc
loc) = do
  (Pat
p', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
p
  Pat -> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pat
p' (Exp -> SrcLoc -> CaseBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> CaseBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (Exp -> MonoM Exp
transformExp Exp
e) MonoM (SrcLoc -> CaseBase Info VName)
-> MonoM SrcLoc -> MonoM (CaseBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex (DimFix Exp
e) = Exp -> DimIndexBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Exp -> DimIndexBase Info VName)
-> MonoM Exp -> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e
transformDimIndex (DimSlice Maybe Exp
me1 Maybe Exp
me2 Maybe Exp
me3) =
  Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe Exp -> Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp)
-> MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me1 MonoM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp)
-> MonoM (Maybe Exp -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me2 MonoM (Maybe Exp -> DimIndexBase Info VName)
-> MonoM (Maybe Exp) -> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> MonoM (Maybe Exp)
trans Maybe Exp
me3
  where
    trans :: Maybe Exp -> MonoM (Maybe Exp)
trans = (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp

-- Transform an operator section into a lambda.
desugarBinOpSection ::
  Exp ->
  Maybe Exp ->
  Maybe Exp ->
  PatType ->
  (PName, StructType, Maybe VName) ->
  (PName, StructType, Maybe VName) ->
  (PatRetType, [VName]) ->
  SrcLoc ->
  MonoM Exp
desugarBinOpSection :: Exp
-> Maybe Exp
-> Maybe Exp
-> PatType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (RetTypeBase (DimDecl VName) Aliasing, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection Exp
op Maybe Exp
e_left Maybe Exp
e_right PatType
t (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xext) (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yext) (RetType [VName]
dims PatType
rettype, [VName]
retext) SrcLoc
loc = do
  (VName
v1, Exp -> Exp
wrap_left, Exp
e1, [Pat]
p1) <- Maybe Exp -> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam Maybe Exp
e_left (PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat]))
-> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
xtype
  (VName
v2, Exp -> Exp
wrap_right, Exp
e2, [Pat]
p2) <- Maybe Exp -> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam Maybe Exp
e_right (PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat]))
-> PatType -> MonoM (VName, Exp -> Exp, Exp, [Pat])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype
  let apply_left :: Exp
apply_left =
        AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
              Exp
op
              Exp
e1
              ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
xext))
              SrcLoc
loc
          )
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> TypeBase (DimDecl VName) ()
-> RetTypeBase (DimDecl VName) Aliasing
-> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
yp TypeBase (DimDecl VName) ()
ytype ([VName] -> PatType -> RetTypeBase (DimDecl VName) Aliasing
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] PatType
t)) [])
      rettype' :: PatType
rettype' =
        let onDim :: DimDecl VName -> DimDecl VName
onDim (NamedDim QualName VName
d)
              | Named VName
p <- PName
xp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v1
              | Named VName
p <- PName
yp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v2
            onDim DimDecl VName
d = DimDecl VName
d
         in (DimDecl VName -> DimDecl VName) -> PatType -> PatType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
onDim PatType
rettype
      body :: Exp
body =
        AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( Exp -> Exp -> Info (Diet, Maybe VName) -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply
              Exp
apply_left
              Exp
e2
              ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
yext))
              SrcLoc
loc
          )
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
rettype' [VName]
retext)
      rettype'' :: TypeBase (DimDecl VName) ()
rettype'' = PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
rettype'
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
    Exp -> Exp
wrap_left (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
      Exp -> Exp
wrap_right (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$
        [Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda ([Pat]
p1 [Pat] -> [Pat] -> [Pat]
forall a. [a] -> [a] -> [a]
++ [Pat]
p2) Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype'')) SrcLoc
loc
  where
    patAndVar :: PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype = do
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
      (VName, Pat, Exp) -> m (VName, Pat, Exp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( VName
x,
          VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
x (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
argtype) SrcLoc
forall a. Monoid a => a
mempty
        )

    makeVarParam :: Maybe Exp -> PatType -> m (VName, Exp -> Exp, Exp, [Pat])
makeVarParam (Just Exp
e) PatType
argtype = do
      (VName
v, Pat
pat, Exp
var_e) <- PatType -> m (VName, Pat, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype
      let wrap :: Exp -> Exp
wrap Exp
body =
            AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([SizeBinder VName] -> Pat -> Exp -> Exp -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] Pat
pat Exp
e Exp
body SrcLoc
forall a. Monoid a => a
mempty) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (Exp -> PatType
typeOf Exp
body) [VName]
forall a. Monoid a => a
mempty)
      (VName, Exp -> Exp, Exp, [Pat])
-> m (VName, Exp -> Exp, Exp, [Pat])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, Exp -> Exp
wrap, Exp
var_e, [])
    makeVarParam Maybe Exp
Nothing PatType
argtype = do
      (VName
v, Pat
pat, Exp
var_e) <- PatType -> m (VName, Pat, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatType -> m (VName, Pat, Exp)
patAndVar PatType
argtype
      (VName, Exp -> Exp, Exp, [Pat])
-> m (VName, Exp -> Exp, Exp, [Pat])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, Exp -> Exp
forall a. a -> a
id, Exp
var_e, [Pat
pat])

desugarProjectSection :: [Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection :: [Name] -> PatType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields (Scalar (Arrow Aliasing
_ PName
_ TypeBase (DimDecl VName) ()
t1 (RetType [VName]
dims PatType
t2))) SrcLoc
loc = do
  VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"project_p"
  let body :: Exp
body = (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Name -> Exp
project (QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1') SrcLoc
forall a. Monoid a => a
mempty) [Name]
fields
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
    [Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda
      [VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
p (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1') SrcLoc
forall a. Monoid a => a
mempty]
      Exp
body
      Maybe (TypeExp VName)
forall a. Maybe a
Nothing
      ((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t2))
      SrcLoc
loc
  where
    t1' :: PatType
t1' = TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t1
    project :: Exp -> Name -> Exp
project Exp
e Name
field =
      case Exp -> PatType
typeOf Exp
e of
        Scalar (Record Map Name PatType
fs)
          | Just PatType
t <- Name -> Map Name PatType -> Maybe PatType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
field Map Name PatType
fs ->
              Name -> Exp -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatType -> SrcLoc -> ExpBase f vn
Project Name
field Exp
e (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t) SrcLoc
forall a. Monoid a => a
mempty
        PatType
t ->
          String -> Exp
forall a. HasCallStack => String -> a
error (String -> Exp) -> String -> Exp
forall a b. (a -> b) -> a -> b
$
            String
"desugarOpSection: type " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" does not have field "
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
field
desugarProjectSection [Name]
_ PatType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarOpSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t

desugarIndexSection :: [DimIndex] -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection :: SliceBase Info VName -> PatType -> SrcLoc -> MonoM Exp
desugarIndexSection SliceBase Info VName
idxs (Scalar (Arrow Aliasing
_ PName
_ TypeBase (DimDecl VName) ()
t1 (RetType [VName]
dims PatType
t2))) SrcLoc
loc = do
  VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"index_i"
  let body :: Exp
body = AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp -> SliceBase Info VName -> SrcLoc -> AppExp
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index (QualName VName -> Info PatType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t1') SrcLoc
loc) SliceBase Info VName
idxs SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (PatType -> [VName] -> AppRes
AppRes PatType
t2 []))
  Exp -> MonoM Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
    [Pat]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, StructRetType)
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda
      [VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
p (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatType
t1')) SrcLoc
forall a. Monoid a => a
mempty]
      Exp
body
      Maybe (TypeExp VName)
forall a. Maybe a
Nothing
      ((Aliasing, StructRetType) -> Info (Aliasing, StructRetType)
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ PatType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatType
t2))
      SrcLoc
loc
  where
    t1' :: PatType
t1' = TypeBase (DimDecl VName) () -> PatType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t1
desugarIndexSection SliceBase Info VName
_ PatType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarIndexSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatType -> String
forall a. Pretty a => a -> String
pretty PatType
t

noticeDims :: TypeBase (DimDecl VName) as -> MonoM ()
noticeDims :: TypeBase (DimDecl VName) as -> MonoM ()
noticeDims = (DimDecl VName -> MonoM ()) -> [DimDecl VName] -> MonoM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ DimDecl VName -> MonoM ()
notice ([DimDecl VName] -> MonoM ())
-> (TypeBase (DimDecl VName) as -> [DimDecl VName])
-> TypeBase (DimDecl VName) as
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) as -> [DimDecl VName]
forall as. TypeBase (DimDecl VName) as -> [DimDecl VName]
nestedDims
  where
    notice :: DimDecl VName -> MonoM ()
notice (NamedDim QualName VName
v) = MonoM Exp -> MonoM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MonoM Exp -> MonoM ()) -> MonoM Exp -> MonoM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
forall a. Monoid a => a
mempty QualName VName
v TypeBase (DimDecl VName) ()
forall dim als. TypeBase dim als
i64
    notice DimDecl VName
_ = () -> MonoM ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- Convert a collection of 'ValBind's to a nested sequence of let-bound,
-- monomorphic functions with the given expression at the bottom.
unfoldLetFuns :: [ValBind] -> Exp -> Exp
unfoldLetFuns :: [ValBind] -> Exp -> Exp
unfoldLetFuns [] Exp
e = Exp
e
unfoldLetFuns (ValBind Maybe (Info EntryPoint)
_ VName
fname Maybe (TypeExp VName)
_ (Info StructRetType
rettype) [TypeParamBase VName]
dim_params [Pat]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
loc : [ValBind]
rest) Exp
e =
  AppExp -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (VName
-> ([TypeParamBase VName], [Pat], Maybe (TypeExp VName),
    Info StructRetType, Exp)
-> Exp
-> SrcLoc
-> AppExp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn], Maybe (TypeExp vn),
    f StructRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
dim_params, [Pat]
params, Maybe (TypeExp VName)
forall a. Maybe a
Nothing, StructRetType -> Info StructRetType
forall a. a -> Info a
Info StructRetType
rettype, Exp
body) Exp
e' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes PatType
e_t [VName]
forall a. Monoid a => a
mempty)
  where
    e' :: Exp
e' = [ValBind] -> Exp -> Exp
unfoldLetFuns [ValBind]
rest Exp
e
    e_t :: PatType
e_t = Exp -> PatType
typeOf Exp
e'

transformPat :: Pat -> MonoM (Pat, RecordReplacements)
transformPat :: Pat -> MonoM (Pat, RecordReplacements)
transformPat (Id VName
v (Info (Scalar (Record Map Name PatType
fs))) SrcLoc
loc) = do
  let fs' :: [(Name, PatType)]
fs' = Map Name PatType -> [(Name, PatType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name PatType
fs
  ([VName]
fs_ks, [PatType]
fs_ts) <- ([(VName, PatType)] -> ([VName], [PatType]))
-> MonoM [(VName, PatType)] -> MonoM ([VName], [PatType])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, PatType)] -> ([VName], [PatType])
forall a b. [(a, b)] -> ([a], [b])
unzip (MonoM [(VName, PatType)] -> MonoM ([VName], [PatType]))
-> MonoM [(VName, PatType)] -> MonoM ([VName], [PatType])
forall a b. (a -> b) -> a -> b
$
    [(Name, PatType)]
-> ((Name, PatType) -> MonoM (VName, PatType))
-> MonoM [(VName, PatType)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Name, PatType)]
fs' (((Name, PatType) -> MonoM (VName, PatType))
 -> MonoM [(VName, PatType)])
-> ((Name, PatType) -> MonoM (VName, PatType))
-> MonoM [(VName, PatType)]
forall a b. (a -> b) -> a -> b
$ \(Name
f, PatType
ft) ->
      (,) (VName -> PatType -> (VName, PatType))
-> MonoM VName -> MonoM (PatType -> (VName, PatType))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (Name -> String
nameToString Name
f) MonoM (PatType -> (VName, PatType))
-> MonoM PatType -> MonoM (VName, PatType)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
ft
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ( [(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat
        ([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatType) -> Name) -> [(Name, PatType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatType)]
fs') ((VName -> Info PatType -> SrcLoc -> Pat)
-> [VName] -> [Info PatType] -> [SrcLoc] -> [Pat]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id [VName]
fs_ks ((PatType -> Info PatType) -> [PatType] -> [Info PatType]
forall a b. (a -> b) -> [a] -> [b]
map PatType -> Info PatType
forall a. a -> Info a
Info [PatType]
fs_ts) ([SrcLoc] -> [Pat]) -> [SrcLoc] -> [Pat]
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [SrcLoc]
forall a. a -> [a]
repeat SrcLoc
loc))
        SrcLoc
loc,
      VName -> RecordReplacement -> RecordReplacements
forall k a. k -> a -> Map k a
M.singleton VName
v (RecordReplacement -> RecordReplacements)
-> RecordReplacement -> RecordReplacements
forall a b. (a -> b) -> a -> b
$ [(Name, (VName, PatType))] -> RecordReplacement
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, PatType))] -> RecordReplacement)
-> [(Name, (VName, PatType))] -> RecordReplacement
forall a b. (a -> b) -> a -> b
$ [Name] -> [(VName, PatType)] -> [(Name, (VName, PatType))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatType) -> Name) -> [(Name, PatType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatType)]
fs') ([(VName, PatType)] -> [(Name, (VName, PatType))])
-> [(VName, PatType)] -> [(Name, (VName, PatType))]
forall a b. (a -> b) -> a -> b
$ [VName] -> [PatType] -> [(VName, PatType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
fs_ks [PatType]
fs_ts
    )
transformPat (Id VName
v Info PatType
t SrcLoc
loc) = (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
v Info PatType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (TuplePat [Pat]
pats SrcLoc
loc) = do
  ([Pat]
pats', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
pats
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat [Pat]
pats' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat (RecordPat [(Name, Pat)]
fields SrcLoc
loc) = do
  let ([Name]
field_names, [Pat]
field_pats) = [(Name, Pat)] -> ([Name], [Pat])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, Pat)]
fields
  ([Pat]
field_pats', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
field_pats
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat ([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
field_names [Pat]
field_pats') SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPat (PatParens Pat
pat SrcLoc
loc) = do
  (Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens Pat
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat (PatAttr AttrInfo VName
attr Pat
pat SrcLoc
loc) = do
  (Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttrInfo VName -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr Pat
pat' SrcLoc
loc, RecordReplacements
rr)
transformPat (Wildcard (Info PatType
t) SrcLoc
loc) = do
  PatType
t' <- PatType -> MonoM PatType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatType
t
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> SrcLoc -> Pat
wildcard PatType
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (PatAscription Pat
pat TypeExp VName
td SrcLoc
loc) = do
  (Pat
pat', RecordReplacements
rr) <- Pat -> MonoM (Pat, RecordReplacements)
transformPat Pat
pat
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat -> TypeExp VName -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatBase f vn -> TypeExp vn -> SrcLoc -> PatBase f vn
PatAscription Pat
pat' TypeExp VName
td SrcLoc
loc, RecordReplacements
rr)
transformPat (PatLit PatLit
e Info PatType
t SrcLoc
loc) = (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatLit -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatLit -> f PatType -> SrcLoc -> PatBase f vn
PatLit PatLit
e Info PatType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPat (PatConstr Name
name Info PatType
t [Pat]
all_ps SrcLoc
loc) = do
  ([Pat]
all_ps', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
all_ps
  (Pat, RecordReplacements) -> MonoM (Pat, RecordReplacements)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Info PatType -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
name Info PatType
t [Pat]
all_ps' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)

wildcard :: PatType -> SrcLoc -> Pat
wildcard :: PatType -> SrcLoc -> Pat
wildcard (Scalar (Record Map Name PatType
fs)) SrcLoc
loc =
  [(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat ([Name] -> [Pat] -> [(Name, Pat)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Map Name PatType -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name PatType
fs) ([Pat] -> [(Name, Pat)]) -> [Pat] -> [(Name, Pat)]
forall a b. (a -> b) -> a -> b
$ (PatType -> Pat) -> [PatType] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map ((Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
`Wildcard` SrcLoc
loc) (Info PatType -> Pat)
-> (PatType -> Info PatType) -> PatType -> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatType -> Info PatType
forall a. a -> Info a
Info) ([PatType] -> [Pat]) -> [PatType] -> [Pat]
forall a b. (a -> b) -> a -> b
$ Map Name PatType -> [PatType]
forall k a. Map k a -> [a]
M.elems Map Name PatType
fs) SrcLoc
loc
wildcard PatType
t SrcLoc
loc =
  Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info PatType
t) SrcLoc
loc

type DimInst = M.Map VName (DimDecl VName)

dimMapping ::
  Monoid a =>
  TypeBase (DimDecl VName) a ->
  TypeBase (DimDecl VName) a ->
  DimInst
dimMapping :: TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2 = State DimInst (TypeBase (DimDecl VName) a) -> DimInst -> DimInst
forall s a. State s a -> s -> s
execState (([VName]
 -> DimDecl VName
 -> DimDecl VName
 -> StateT DimInst Identity (DimDecl VName))
-> TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a
-> State DimInst (TypeBase (DimDecl VName) a)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims [VName]
-> DimDecl VName
-> DimDecl VName
-> StateT DimInst Identity (DimDecl VName)
forall (t :: * -> *) (f :: * -> *) vn a.
(Foldable t, MonadState (Map vn (DimDecl a)) f, Ord vn, Eq a) =>
t a -> DimDecl vn -> DimDecl a -> f (DimDecl vn)
f TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2) DimInst
forall a. Monoid a => a
mempty
  where
    f :: t a -> DimDecl vn -> DimDecl a -> f (DimDecl vn)
f t a
bound DimDecl vn
d1 (NamedDim QualName a
d2)
      | QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
d2 a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t a
bound = DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl vn
d1
    f t a
_ (NamedDim QualName vn
d1) DimDecl a
d2 = do
      (Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ())
-> (Map vn (DimDecl a) -> Map vn (DimDecl a)) -> f ()
forall a b. (a -> b) -> a -> b
$ vn -> DimDecl a -> Map vn (DimDecl a) -> Map vn (DimDecl a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName vn -> vn
forall vn. QualName vn -> vn
qualLeaf QualName vn
d1) DimDecl a
d2
      DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl vn -> f (DimDecl vn)) -> DimDecl vn -> f (DimDecl vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> DimDecl vn
forall vn. QualName vn -> DimDecl vn
NamedDim QualName vn
d1
    f t a
_ DimDecl vn
d DimDecl a
_ = DimDecl vn -> f (DimDecl vn)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl vn
d

inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp]
inferSizeArgs :: [TypeParamBase VName]
-> TypeBase (DimDecl VName) () -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t =
  (TypeParamBase VName -> Maybe Exp)
-> [TypeParamBase VName] -> [Exp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DimInst -> TypeParamBase VName -> Maybe Exp
forall k vn.
Ord k =>
Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg (TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> DimInst
forall a.
Monoid a =>
TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t)) [TypeParamBase VName]
tparams
  where
    tparamArg :: Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg Map k (DimDecl vn)
dinst TypeParamBase k
tp =
      case k -> Map k (DimDecl vn) -> Maybe (DimDecl vn)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (TypeParamBase k -> k
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase k
tp) Map k (DimDecl vn)
dinst of
        Just (NamedDim QualName vn
d) ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> Info PatType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName vn
d (PatType -> Info PatType
forall a. a -> Info a
Info PatType
forall dim als. TypeBase dim als
i64) SrcLoc
forall a. Monoid a => a
mempty
        Just (ConstDim Int
x) ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> Int64 -> IntValue
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) SrcLoc
forall a. Monoid a => a
mempty
        Maybe (DimDecl vn)
_ ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
0) SrcLoc
forall a. Monoid a => a
mempty

-- Monomorphising higher-order functions can result in function types
-- where the same named parameter occurs in multiple spots.  When
-- monomorphising we don't really need those parameter names anymore,
-- and the defunctionaliser can be confused if there are duplicates
-- (it doesn't handle shadowing), so let's just remove all parameter
-- names here.  This is safe because a MonoType does not contain sizes
-- anyway.
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 ShapeDecl dim
shape ScalarTypeBase dim ()
t) = ()
-> Uniqueness
-> ShapeDecl dim
-> ScalarTypeBase dim ()
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> ShapeDecl dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
u ShapeDecl dim
shape (ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t)
    f (Scalar ScalarTypeBase dim ()
t) = ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t
    f' :: ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' (Arrow () PName
_ TypeBase dim ()
t1 (RetType [VName]
dims TypeBase dim ()
t2)) =
      ()
-> PName
-> TypeBase dim ()
-> RetTypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) ([VName] -> TypeBase dim () -> RetTypeBase dim ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t2))
    f' (Record Map Name (TypeBase dim ())
fs) =
      Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim ()) -> ScalarTypeBase dim ())
-> Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ (TypeBase dim () -> TypeBase dim ())
-> Map Name (TypeBase dim ()) -> Map Name (TypeBase dim ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim () -> TypeBase dim ()
f Map Name (TypeBase dim ())
fs
    f' (Sum Map Name [TypeBase dim ()]
cs) =
      Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim ()] -> ScalarTypeBase dim ())
-> Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim ()] -> [TypeBase dim ()])
-> Map Name [TypeBase dim ()] -> Map Name [TypeBase dim ()]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim () -> TypeBase dim ())
-> [TypeBase dim ()] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim () -> TypeBase dim ()
f) Map Name [TypeBase dim ()]
cs
    f' ScalarTypeBase dim ()
t = ScalarTypeBase dim ()
t

-- Monomorphise a polymorphic function at the types given in the instance
-- list. Monomorphises the body of the function as well. Returns the fresh name
-- of the generated monomorphic function and its 'ValBind' representation.
monomorphiseBinding ::
  Bool ->
  PolyBinding ->
  MonoType ->
  MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding :: Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
entry (PolyBinding RecordReplacements
rr (VName
name, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) MonoType
inst_t =
  RecordReplacements
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr (MonoM (VName, InferSizeArgs, ValBind)
 -> MonoM (VName, InferSizeArgs, ValBind))
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ do
    let bind_t :: TypeBase (DimDecl VName) ()
bind_t = [TypeBase (DimDecl VName) ()]
-> StructRetType -> TypeBase (DimDecl VName) ()
forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType ((Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType [Pat]
params) StructRetType
rettype
    (Map VName StructRetType
substs, [TypeParamBase VName]
t_shape_params) <- SrcLoc
-> TypeBase () ()
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName StructRetType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc (TypeBase (DimDecl VName) () -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes TypeBase (DimDecl VName) ()
bind_t) (MonoType
 -> MonoM (Map VName StructRetType, [TypeParamBase VName]))
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
inst_t
    let substs' :: Map VName (Subst StructRetType)
substs' = (StructRetType -> Subst StructRetType)
-> Map VName StructRetType -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ([TypeParamBase VName] -> StructRetType -> Subst StructRetType
forall t. [TypeParamBase VName] -> t -> Subst t
Subst []) Map VName StructRetType
substs
        rettype' :: StructRetType
rettype' = TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') StructRetType
rettype
        substPatType :: PatType -> PatType
substPatType =
          (VName -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> PatType -> PatType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((Subst StructRetType
 -> Subst (RetTypeBase (DimDecl VName) Aliasing))
-> Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructRetType -> RetTypeBase (DimDecl VName) Aliasing)
-> Subst StructRetType
-> Subst (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((() -> Aliasing)
-> StructRetType -> RetTypeBase (DimDecl VName) Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Aliasing -> () -> Aliasing
forall a b. a -> b -> a
const Aliasing
forall a. Monoid a => a
mempty))) (Maybe (Subst StructRetType)
 -> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing)))
-> TypeSubs
-> VName
-> Maybe (Subst (RetTypeBase (DimDecl VName) Aliasing))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs'))
        params' :: [Pat]
params' = (Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
substPatType) [Pat]
params
        bind_t' :: TypeBase (DimDecl VName) ()
bind_t' = TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') TypeBase (DimDecl VName) ()
bind_t
        ([TypeParamBase VName]
shape_params_explicit, [TypeParamBase VName]
shape_params_implicit) =
          (TypeParamBase VName -> Bool)
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TypeBase (DimDecl VName) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
bind_t') (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) ([TypeParamBase VName]
 -> ([TypeParamBase VName], [TypeParamBase VName]))
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$
            [TypeParamBase VName]
shape_params [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
t_shape_params

    ([Pat]
params'', [RecordReplacements]
rrs) <- [(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pat, RecordReplacements)] -> ([Pat], [RecordReplacements]))
-> MonoM [(Pat, RecordReplacements)]
-> MonoM ([Pat], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pat -> MonoM (Pat, RecordReplacements))
-> [Pat] -> MonoM [(Pat, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pat -> MonoM (Pat, RecordReplacements)
transformPat [Pat]
params'

    (TypeBase (DimDecl VName) () -> MonoM ())
-> [TypeBase (DimDecl VName) ()] -> MonoM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ TypeBase (DimDecl VName) () -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims ([TypeBase (DimDecl VName) ()] -> MonoM ())
-> [TypeBase (DimDecl VName) ()] -> MonoM ()
forall a b. (a -> b) -> a -> b
$ StructRetType -> TypeBase (DimDecl VName) ()
forall dim as. RetTypeBase dim as -> TypeBase dim as
retType StructRetType
rettype TypeBase (DimDecl VName) ()
-> [TypeBase (DimDecl VName) ()] -> [TypeBase (DimDecl VName) ()]
forall a. a -> [a] -> [a]
: (Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType [Pat]
params''

    Exp
body' <- TypeSubs -> Exp -> MonoM Exp
forall (m :: * -> *). Monad m => TypeSubs -> Exp -> m Exp
updateExpTypes (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') Exp
body
    Exp
body'' <- RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body'
    Bool
seen_before <- VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem VName
name ([VName] -> Bool) -> (Lifts -> [VName]) -> Lifts -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
-> Lifts -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ((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) (Lifts -> Bool) -> MonoM Lifts -> MonoM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM Lifts
getLifts
    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 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
seen_before
        then VName -> MonoM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
name
        else VName -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
name

    (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( VName
name',
        [TypeParamBase VName]
-> TypeBase (DimDecl VName) () -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
shape_params_explicit TypeBase (DimDecl VName) ()
bind_t',
        if Bool
entry
          then
            VName
-> [TypeParamBase VName]
-> [Pat]
-> StructRetType
-> Exp
-> ValBind
toValBinding
              VName
name'
              ([TypeParamBase VName]
shape_params_explicit [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
shape_params_implicit)
              [Pat]
params''
              StructRetType
rettype'
              Exp
body''
          else
            VName
-> [TypeParamBase VName]
-> [Pat]
-> StructRetType
-> Exp
-> ValBind
toValBinding
              VName
name'
              [TypeParamBase VName]
shape_params_implicit
              ((TypeParamBase VName -> Pat) -> [TypeParamBase VName] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> Pat
forall vn. TypeParamBase vn -> PatBase Info vn
shapeParam [TypeParamBase VName]
shape_params_explicit [Pat] -> [Pat] -> [Pat]
forall a. [a] -> [a] -> [a]
++ [Pat]
params'')
              StructRetType
rettype'
              Exp
body''
      )
  where
    shape_params :: [TypeParamBase VName]
shape_params = (TypeParamBase VName -> Bool)
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> (TypeParamBase VName -> Bool) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam) [TypeParamBase VName]
tparams

    updateExpTypes :: TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (TypeSubs -> ASTMapper m
mapper TypeSubs
substs)

    mapper :: TypeSubs -> ASTMapper m
mapper TypeSubs
substs =
      ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatType -> m PatType)
-> (StructRetType -> m StructRetType)
-> (RetTypeBase (DimDecl VName) Aliasing
    -> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
        { mapOnExp :: Exp -> m Exp
mapOnExp = TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs,
          mapOnName :: VName -> m VName
mapOnName = VName -> m VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnQualName :: QualName VName -> m (QualName VName)
mapOnQualName = QualName VName -> m (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnStructType :: TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> m (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
          mapOnPatType :: PatType -> m PatType
mapOnPatType = PatType -> m PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> m PatType)
-> (PatType -> PatType) -> PatType -> m PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
          mapOnStructRetType :: StructRetType -> m StructRetType
mapOnStructRetType = StructRetType -> m StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> m StructRetType)
-> (StructRetType -> StructRetType)
-> StructRetType
-> m StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
          mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) Aliasing
 -> m (RetTypeBase (DimDecl VName) Aliasing))
-> (RetTypeBase (DimDecl VName) Aliasing
    -> RetTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> m (RetTypeBase (DimDecl VName) Aliasing)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs
        }

    shapeParam :: TypeParamBase vn -> PatBase Info vn
shapeParam TypeParamBase vn
tp = vn -> Info PatType -> SrcLoc -> PatBase Info vn
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id (TypeParamBase vn -> vn
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase vn
tp) (PatType -> Info PatType
forall a. a -> Info a
Info PatType
forall dim als. TypeBase dim als
i64) (SrcLoc -> PatBase Info vn) -> SrcLoc -> PatBase Info vn
forall a b. (a -> b) -> a -> b
$ TypeParamBase vn -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf TypeParamBase vn
tp

    toValBinding :: VName
-> [TypeParamBase VName]
-> [Pat]
-> StructRetType
-> Exp
-> ValBind
toValBinding VName
name' [TypeParamBase VName]
tparams' [Pat]
params'' StructRetType
rettype' Exp
body'' =
      ValBind :: forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f StructRetType
-> [TypeParamBase vn]
-> [PatBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo vn]
-> SrcLoc
-> ValBindBase f vn
ValBind
        { valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing,
          valBindName :: VName
valBindName = VName
name',
          valBindRetType :: Info StructRetType
valBindRetType = StructRetType -> Info StructRetType
forall a. a -> Info a
Info StructRetType
rettype',
          valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
forall a. Maybe a
Nothing,
          valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
tparams',
          valBindParams :: [Pat]
valBindParams = [Pat]
params'',
          valBindBody :: Exp
valBindBody = Exp
body'',
          valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing,
          valBindAttrs :: [AttrInfo VName]
valBindAttrs = [AttrInfo VName]
attrs,
          valBindLocation :: SrcLoc
valBindLocation = SrcLoc
loc
        }

typeSubstsM ::
  MonadFreshNames m =>
  SrcLoc ->
  TypeBase () () ->
  MonoType ->
  m (M.Map VName StructRetType, [TypeParam])
typeSubstsM :: SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName StructRetType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc TypeBase () ()
orig_t1 MonoType
orig_t2 =
  WriterT [TypeParamBase VName] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParamBase VName])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT [TypeParamBase VName] m (Map VName StructRetType)
 -> m (Map VName StructRetType, [TypeParamBase VName]))
-> WriterT [TypeParamBase VName] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ (Map VName StructRetType, Map Int VName) -> Map VName StructRetType
forall a b. (a, b) -> a
fst ((Map VName StructRetType, Map Int VName)
 -> Map VName StructRetType)
-> WriterT
     [TypeParamBase VName] m (Map VName StructRetType, Map Int VName)
-> WriterT [TypeParamBase VName] m (Map VName StructRetType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT
  (Map VName StructRetType, Map Int VName)
  (WriterT [TypeParamBase VName] m)
  ()
-> (Map VName StructRetType, Map Int VName)
-> WriterT
     [TypeParamBase VName] m (Map VName StructRetType, Map Int VName)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (TypeBase () ()
-> MonoType
-> StateT
     (Map VName StructRetType, Map Int VName)
     (WriterT [TypeParamBase VName] m)
     ()
forall (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
       (m :: * -> *) dim.
(MonadState (Map VName StructRetType, Map Int VName) (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)),
 Pretty (ShapeDecl dim), Monad (t m)) =>
TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase () ()
orig_t1 MonoType
orig_t2) (Map VName StructRetType
forall a. Monoid a => a
mempty, Map Int VName
forall a. Monoid a => a
mempty)
  where
    subRet :: TypeBase dim () -> RetTypeBase MonoSize () -> t (t m) ()
subRet (Scalar (TypeVar ()
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) RetTypeBase MonoSize ()
rt =
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Int
baseTag (TypeName -> VName
typeLeaf TypeName
v) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
        TypeName -> RetTypeBase MonoSize () -> t (t m) ()
forall as (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
       (m :: * -> *).
(MonadState
   (Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
   (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst TypeName
v RetTypeBase MonoSize ()
rt
    subRet TypeBase dim ()
t1 (RetType [VName]
_ MonoType
t2) =
      TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1 MonoType
t2

    sub :: TypeBase dim () -> MonoType -> t (t m) ()
sub t1 :: TypeBase dim ()
t1@Array {} t2 :: MonoType
t2@Array {}
      | Just TypeBase dim ()
t1' <- Int -> TypeBase dim () -> Maybe (TypeBase dim ())
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim () -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim ()
t1) TypeBase dim ()
t1,
        Just MonoType
t2' <- Int -> MonoType -> Maybe MonoType
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim () -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim ()
t1) MonoType
t2 =
          TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1' MonoType
t2'
    sub (Scalar (TypeVar ()
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) MonoType
t =
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Int
baseTag (TypeName -> VName
typeLeaf TypeName
v) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
        TypeName -> RetTypeBase MonoSize () -> t (t m) ()
forall as (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
       (m :: * -> *).
(MonadState
   (Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
   (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst TypeName
v (RetTypeBase MonoSize () -> t (t m) ())
-> RetTypeBase MonoSize () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ [VName] -> MonoType -> RetTypeBase MonoSize ()
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] MonoType
t
    sub (Scalar (Record Map Name (TypeBase dim ())
fields1)) (Scalar (Record Map Name MonoType
fields2)) =
      (TypeBase dim () -> MonoType -> t (t m) ())
-> [TypeBase dim ()] -> [MonoType] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
        TypeBase dim () -> MonoType -> t (t m) ()
sub
        (((Name, TypeBase dim ()) -> TypeBase dim ())
-> [(Name, TypeBase dim ())] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase dim ()) -> TypeBase dim ()
forall a b. (a, b) -> b
snd ([(Name, TypeBase dim ())] -> [TypeBase dim ()])
-> [(Name, TypeBase dim ())] -> [TypeBase dim ()]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim ()) -> [(Name, TypeBase dim ())]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase dim ())
fields1)
        (((Name, MonoType) -> MonoType) -> [(Name, MonoType)] -> [MonoType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, MonoType) -> MonoType
forall a b. (a, b) -> b
snd ([(Name, MonoType)] -> [MonoType])
-> [(Name, MonoType)] -> [MonoType]
forall a b. (a -> b) -> a -> b
$ Map Name MonoType -> [(Name, MonoType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name MonoType
fields2)
    sub (Scalar Prim {}) (Scalar Prim {}) = () -> t (t m) ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    sub (Scalar (Arrow ()
_ PName
_ TypeBase dim ()
t1a (RetType [VName]
_ TypeBase dim ()
t1b))) (Scalar (Arrow ()
_ PName
_ MonoType
t2a RetTypeBase MonoSize ()
t2b)) = do
      TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1a MonoType
t2a
      TypeBase dim () -> RetTypeBase MonoSize () -> t (t m) ()
subRet TypeBase dim ()
t1b RetTypeBase MonoSize ()
t2b
    sub (Scalar (Sum Map Name [TypeBase dim ()]
cs1)) (Scalar (Sum Map Name [MonoType]
cs2)) =
      ((Name, [TypeBase dim ()]) -> (Name, [MonoType]) -> t (t m) [()])
-> [(Name, [TypeBase dim ()])]
-> [(Name, [MonoType])]
-> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Name, [TypeBase dim ()]) -> (Name, [MonoType]) -> t (t m) [()]
forall a a.
(a, [TypeBase dim ()]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (Map Name [TypeBase dim ()] -> [(Name, [TypeBase dim ()])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase dim ()]
cs1) (Map Name [MonoType] -> [(Name, [MonoType])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [MonoType]
cs2)
      where
        typeSubstClause :: (a, [TypeBase dim ()]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (a
_, [TypeBase dim ()]
ts1) (a
_, [MonoType]
ts2) = (TypeBase dim () -> MonoType -> t (t m) ())
-> [TypeBase dim ()] -> [MonoType] -> t (t m) [()]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase dim () -> MonoType -> t (t m) ()
sub [TypeBase dim ()]
ts1 [MonoType]
ts2
    sub t1 :: TypeBase dim ()
t1@(Scalar Sum {}) MonoType
t2 = TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1 MonoType
t2
    sub TypeBase dim ()
t1 t2 :: MonoType
t2@(Scalar Sum {}) = TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1 MonoType
t2
    sub TypeBase dim ()
t1 MonoType
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 () -> String
forall a. Pretty a => a -> String
pretty TypeBase dim ()
t1, MonoType -> String
forall a. Pretty a => a -> String
pretty MonoType
t2]

    addSubst :: TypeName -> RetTypeBase MonoSize as -> t (t m) ()
addSubst (TypeName [VName]
_ VName
v) (RetType [VName]
ext TypeBase MonoSize as
t) = do
      (Map VName (RetTypeBase (DimDecl VName) as)
ts, Map Int VName
sizes) <- t (t m) (Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
forall s (m :: * -> *). MonadState s m => m s
get
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName
v VName -> Map VName (RetTypeBase (DimDecl VName) as) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (RetTypeBase (DimDecl VName) as)
ts) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ do
        TypeBase (DimDecl VName) as
t' <- (MonoSize -> t (t m) (DimDecl VName))
-> (as -> t (t m) as)
-> TypeBase MonoSize as
-> t (t m) (TypeBase (DimDecl VName) as)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse MonoSize -> t (t m) (DimDecl VName)
forall a (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
       (m :: * -> *).
(MonadState (a, Map Int VName) (t (t m)), MonadTrans t,
 MonadTrans t, Monad (t m), MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m))) =>
MonoSize -> t (t m) (DimDecl VName)
onDim as -> t (t m) as
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase MonoSize as
t
        (Map VName (RetTypeBase (DimDecl VName) as), Map Int VName)
-> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VName
-> RetTypeBase (DimDecl VName) as
-> Map VName (RetTypeBase (DimDecl VName) as)
-> Map VName (RetTypeBase (DimDecl VName) as)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v ([VName]
-> TypeBase (DimDecl VName) as -> RetTypeBase (DimDecl VName) as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase (DimDecl VName) as
t') Map VName (RetTypeBase (DimDecl VName) as)
ts, Map Int VName
sizes)

    onDim :: MonoSize -> t (t m) (DimDecl VName)
onDim (MonoKnown Int
i) = do
      (a
ts, Map Int VName
sizes) <- t (t m) (a, Map Int VName)
forall s (m :: * -> *). MonadState s m => m s
get
      case Int -> Map Int VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Int
i Map Int VName
sizes of
        Maybe VName
Nothing -> do
          VName
d <- t m VName -> t (t m) VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (t m VName -> t (t m) VName) -> t m VName -> t (t m) VName
forall a b. (a -> b) -> a -> b
$ m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"d"
          [TypeParamBase VName] -> t (t m) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
TypeParamDim VName
d SrcLoc
loc]
          (a, Map Int VName) -> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (a
ts, Int -> VName -> Map Int VName -> Map Int VName
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i VName
d Map Int VName
sizes)
          DimDecl VName -> t (t m) (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d
        Just VName
d ->
          DimDecl VName -> t (t m) (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
d
    onDim (MonoAnon VName
v) = DimDecl VName -> t (t m) (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t (t m) (DimDecl VName))
-> DimDecl VName -> t (t m) (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ Maybe VName -> DimDecl VName
forall vn. Maybe vn -> DimDecl vn
AnyDim (Maybe VName -> DimDecl VName) -> Maybe VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v

-- Perform a given substitution on the types in a pattern.
substPat :: Bool -> (PatType -> PatType) -> Pat -> Pat
substPat :: Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
pat = case Pat
pat of
  TuplePat [Pat]
pats SrcLoc
loc -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn. [PatBase f vn] -> SrcLoc -> PatBase f vn
TuplePat ((Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f) [Pat]
pats) SrcLoc
loc
  RecordPat [(Name, Pat)]
fs SrcLoc
loc -> [(Name, Pat)] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
[(Name, PatBase f vn)] -> SrcLoc -> PatBase f vn
RecordPat (((Name, Pat) -> (Name, Pat)) -> [(Name, Pat)] -> [(Name, Pat)]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Pat) -> (Name, Pat)
forall a. (a, Pat) -> (a, Pat)
substField [(Name, Pat)]
fs) SrcLoc
loc
    where
      substField :: (a, Pat) -> (a, Pat)
substField (a
n, Pat
p) = (a
n, Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p)
  PatParens Pat
p SrcLoc
loc -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn. PatBase f vn -> SrcLoc -> PatBase f vn
PatParens (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p) SrcLoc
loc
  PatAttr AttrInfo VName
attr Pat
p SrcLoc
loc -> AttrInfo VName -> Pat -> SrcLoc -> Pat
forall (f :: * -> *) vn.
AttrInfo vn -> PatBase f vn -> SrcLoc -> PatBase f vn
PatAttr AttrInfo VName
attr (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry PatType -> PatType
f Pat
p) SrcLoc
loc
  Id VName
vn (Info PatType
tp) SrcLoc
loc -> VName -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. vn -> f PatType -> SrcLoc -> PatBase f vn
Id VName
vn (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
  Wildcard (Info PatType
tp) SrcLoc
loc -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn. f PatType -> SrcLoc -> PatBase f vn
Wildcard (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
  PatAscription Pat
p TypeExp VName
td SrcLoc
loc
    | Bool
entry -> Pat -> TypeExp VName -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatBase f vn -> TypeExp vn -> SrcLoc -> PatBase f vn
PatAscription (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
False PatType -> PatType
f Pat
p) TypeExp VName
td SrcLoc
loc
    | Bool
otherwise -> Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
False PatType -> PatType
f Pat
p
  PatLit PatLit
e (Info PatType
tp) SrcLoc
loc -> PatLit -> Info PatType -> SrcLoc -> Pat
forall (f :: * -> *) vn.
PatLit -> f PatType -> SrcLoc -> PatBase f vn
PatLit PatLit
e (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) SrcLoc
loc
  PatConstr Name
n (Info PatType
tp) [Pat]
ps SrcLoc
loc -> Name -> Info PatType -> [Pat] -> SrcLoc -> Pat
forall (f :: * -> *) vn.
Name -> f PatType -> [PatBase f vn] -> SrcLoc -> PatBase f vn
PatConstr Name
n (PatType -> Info PatType
forall a. a -> Info a
Info (PatType -> Info PatType) -> PatType -> Info PatType
forall a b. (a -> b) -> a -> b
$ PatType -> PatType
f PatType
tp) [Pat]
ps SrcLoc
loc

toPolyBinding :: ValBind -> PolyBinding
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding (ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
_ (Info StructRetType
rettype) [TypeParamBase VName]
tparams [Pat]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) =
  RecordReplacements
-> (VName, [TypeParamBase VName], [Pat], StructRetType, Exp,
    [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [Pat]
params, StructRetType
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)

-- Remove all type variables and type abbreviations from a value binding.
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables Bool
entry ValBind
valbind = do
  let (ValBind Maybe (Info EntryPoint)
_ VName
_ Maybe (TypeExp VName)
_ (Info (RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype)) [TypeParamBase VName]
_ [Pat]
pats Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_) = ValBind
valbind
  Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
 -> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  let mapper :: ASTMapper MonoM
mapper =
        ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatType -> m PatType)
-> (StructRetType -> m StructRetType)
-> (RetTypeBase (DimDecl VName) Aliasing
    -> m (RetTypeBase (DimDecl VName) Aliasing))
-> ASTMapper m
ASTMapper
          { mapOnExp :: Exp -> MonoM Exp
mapOnExp = Exp -> MonoM Exp
onExp,
            mapOnName :: VName -> MonoM VName
mapOnName = VName -> MonoM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnQualName :: QualName VName -> MonoM (QualName VName)
mapOnQualName = QualName VName -> MonoM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnStructType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) ()
 -> MonoM (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
            mapOnPatType :: PatType -> MonoM PatType
mapOnPatType = PatType -> MonoM PatType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatType -> MonoM PatType)
-> (PatType -> PatType) -> PatType -> MonoM PatType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
            mapOnStructRetType :: StructRetType -> MonoM StructRetType
mapOnStructRetType = StructRetType -> MonoM StructRetType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructRetType -> MonoM StructRetType)
-> (StructRetType -> StructRetType)
-> StructRetType
-> MonoM StructRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs),
            mapOnPatRetType :: RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
mapOnPatRetType = RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase (DimDecl VName) Aliasing
 -> MonoM (RetTypeBase (DimDecl VName) Aliasing))
-> (RetTypeBase (DimDecl VName) Aliasing
    -> RetTypeBase (DimDecl VName) Aliasing)
-> RetTypeBase (DimDecl VName) Aliasing
-> MonoM (RetTypeBase (DimDecl VName) Aliasing)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase (DimDecl VName) Aliasing
-> RetTypeBase (DimDecl VName) Aliasing
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs)
          }

      onExp :: Exp -> MonoM Exp
onExp = ASTMapper MonoM -> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper

  Exp
body' <- Exp -> MonoM Exp
onExp Exp
body

  ValBind -> MonoM ValBind
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    ValBind
valbind
      { valBindRetType :: Info StructRetType
valBindRetType = StructRetType -> Info StructRetType
forall a. a -> Info a
Info (TypeSubs -> StructRetType -> StructRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) (StructRetType -> StructRetType) -> StructRetType -> StructRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase (DimDecl VName) ()
rettype),
        valBindParams :: [Pat]
valBindParams = (Pat -> Pat) -> [Pat] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
entry ((PatType -> PatType) -> Pat -> Pat)
-> (PatType -> PatType) -> Pat -> Pat
forall a b. (a -> b) -> a -> b
$ TypeSubs -> PatType -> PatType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs)) [Pat]
pats,
        valBindBody :: Exp
valBindBody = Exp
body'
      }

removeTypeVariablesInType :: StructType -> MonoM StructType
removeTypeVariablesInType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t = do
  Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
 -> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) ()
 -> MonoM (TypeBase (DimDecl VName) ()))
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$ TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) TypeBase (DimDecl VName) ()
t

transformValBind :: ValBind -> MonoM Env
transformValBind :: ValBind -> MonoM Env
transformValBind ValBind
valbind = do
  PolyBinding
valbind' <-
    ValBind -> PolyBinding
toPolyBinding
      (ValBind -> PolyBinding) -> MonoM ValBind -> MonoM PolyBinding
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> ValBind -> MonoM ValBind
removeTypeVariables (Maybe (Info EntryPoint) -> Bool
forall a. Maybe a -> Bool
isJust (ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind)) ValBind
valbind

  Bool -> MonoM () -> MonoM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe (Info EntryPoint) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Info EntryPoint) -> Bool)
-> Maybe (Info EntryPoint) -> Bool
forall a b. (a -> b) -> a -> b
$ ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind) (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ do
    TypeBase (DimDecl VName) ()
t <-
      TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType (TypeBase (DimDecl VName) ()
 -> MonoM (TypeBase (DimDecl VName) ()))
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$
        [TypeBase (DimDecl VName) ()]
-> StructRetType -> TypeBase (DimDecl VName) ()
forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType
          ((Pat -> TypeBase (DimDecl VName) ())
-> [Pat] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pat -> TypeBase (DimDecl VName) ()
patternStructType (ValBind -> [Pat]
forall (f :: * -> *) vn. ValBindBase f vn -> [PatBase f vn]
valBindParams ValBind
valbind))
          (StructRetType -> TypeBase (DimDecl VName) ())
-> StructRetType -> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ Info StructRetType -> StructRetType
forall a. Info a -> a
unInfo (Info StructRetType -> StructRetType)
-> Info StructRetType -> StructRetType
forall a b. (a -> b) -> a -> b
$ ValBind -> Info StructRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f StructRetType
valBindRetType ValBind
valbind
    (VName
name, InferSizeArgs
infer, ValBind
valbind'') <- Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
True PolyBinding
valbind' (MonoType -> MonoM (VName, InferSizeArgs, ValBind))
-> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t
    Seq (VName, ValBind) -> MonoM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (VName, ValBind) -> MonoM ())
-> Seq (VName, ValBind) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (VName, ValBind) -> Seq (VName, ValBind)
forall a. a -> Seq a
Seq.singleton (VName
name, ValBind
valbind'' {valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind})
    VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) (TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t) (VName
name, InferSizeArgs
infer)

  Env -> MonoM Env
forall (f :: * -> *) a. Applicative f => a -> f a
pure Env
forall a. Monoid a => a
mempty {envPolyBindings :: Map VName PolyBinding
envPolyBindings = VName -> PolyBinding -> Map VName PolyBinding
forall k a. k -> a -> Map k a
M.singleton (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) PolyBinding
valbind'}

transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeExp VName
_ (Info (RetType [VName]
dims TypeBase (DimDecl VName) ()
t)) Maybe DocComment
_ SrcLoc
_) = do
  Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
 -> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  TypeBase (DimDecl VName) () -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims TypeBase (DimDecl VName) ()
t
  let tbinding :: TypeBinding
tbinding = Liftedness -> [TypeParamBase VName] -> StructRetType -> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase (DimDecl VName) () -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase (DimDecl VName) () -> StructRetType)
-> TypeBase (DimDecl VName) () -> StructRetType
forall a b. (a -> b) -> a -> b
$ TypeSubs
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) TypeBase (DimDecl VName) ()
t
  Env -> MonoM Env
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 (f :: * -> *) a. Applicative f => a -> f a
pure ()
transformDecs (ValDec ValBind
valbind : [Dec]
ds) = do
  Env
env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (TypeDec TypeBind
typebind : [Dec]
ds) = do
  Env
env <- TypeBind -> MonoM Env
transformTypeBind TypeBind
typebind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (Dec
dec : [Dec]
_) =
  String -> MonoM ()
forall a. HasCallStack => String -> a
error (String -> MonoM ()) -> String -> MonoM ()
forall a b. (a -> b) -> a -> b
$
    String
"The monomorphization module expects a module-free "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"input program, but received: "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Dec -> String
forall a. Pretty a => a -> String
pretty Dec
dec

-- | Monomorphise a list of top-level declarations. A module-free input program
-- is expected, so only value declarations and type declaration are accepted.
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