{-# 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.
--
-- 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],
        [Pattern],
        Maybe (TypeExp VName),
        StructType,
        [VName],
        Exp,
        [AttrInfo],
        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, PatternType)

-- 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 (m :: * -> *) a. Monad m => a -> m a
return (Maybe PolyBinding -> MonoM (Maybe PolyBinding))
-> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ PolyBinding -> Maybe PolyBinding
forall a. a -> Maybe a
Just PolyBinding
valbind
    Maybe PolyBinding
Nothing -> Maybe PolyBinding -> MonoM (Maybe PolyBinding)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe PolyBinding
forall a. Maybe a
Nothing

lookupRecordReplacement :: VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement :: VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement VName
v = (Env -> Maybe RecordReplacement) -> MonoM (Maybe RecordReplacement)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Maybe RecordReplacement)
 -> MonoM (Maybe RecordReplacement))
-> (Env -> Maybe RecordReplacement)
-> MonoM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (RecordReplacements -> Maybe RecordReplacement)
-> (Env -> RecordReplacements) -> Env -> Maybe RecordReplacement
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> RecordReplacements
envRecordReplacements

-- 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
  deriving (MonoSize -> MonoSize -> Bool
(MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool) -> Eq MonoSize
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: MonoSize -> MonoSize -> Bool
$c/= :: MonoSize -> MonoSize -> Bool
== :: MonoSize -> MonoSize -> Bool
$c== :: MonoSize -> MonoSize -> Bool
Eq, Eq MonoSize
Eq MonoSize
-> (MonoSize -> MonoSize -> Ordering)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> MonoSize)
-> (MonoSize -> MonoSize -> MonoSize)
-> Ord MonoSize
MonoSize -> MonoSize -> Bool
MonoSize -> MonoSize -> Ordering
MonoSize -> MonoSize -> MonoSize
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: MonoSize -> MonoSize -> MonoSize
$cmin :: MonoSize -> MonoSize -> MonoSize
max :: MonoSize -> MonoSize -> MonoSize
$cmax :: MonoSize -> MonoSize -> MonoSize
>= :: MonoSize -> MonoSize -> Bool
$c>= :: MonoSize -> MonoSize -> Bool
> :: MonoSize -> MonoSize -> Bool
$c> :: MonoSize -> MonoSize -> Bool
<= :: MonoSize -> MonoSize -> Bool
$c<= :: MonoSize -> MonoSize -> Bool
< :: MonoSize -> MonoSize -> Bool
$c< :: MonoSize -> MonoSize -> Bool
compare :: MonoSize -> MonoSize -> Ordering
$ccompare :: MonoSize -> MonoSize -> Ordering
$cp1Ord :: Eq MonoSize
Ord, Int -> MonoSize -> ShowS
[MonoSize] -> ShowS
MonoSize -> String
(Int -> MonoSize -> ShowS)
-> (MonoSize -> String) -> ([MonoSize] -> ShowS) -> Show MonoSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MonoSize] -> ShowS
$cshowList :: [MonoSize] -> ShowS
show :: MonoSize -> String
$cshow :: MonoSize -> String
showsPrec :: Int -> MonoSize -> ShowS
$cshowsPrec :: Int -> MonoSize -> ShowS
Show)

instance 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 MonoSize
MonoAnon = String -> Doc
text String
"?"

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 :: * -> *) a p.
(MonadState (Int, Map (DimDecl a) Int) f, Ord a,
 Ord (DimDecl a)) =>
Set a -> p -> DimDecl a -> 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 a -> p -> DimDecl a -> f MonoSize
onDim Set a
bound p
_ (NamedDim QualName a
d)
      -- A locally bound size.
      | QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
d a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set a
bound = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure MonoSize
MonoAnon
    onDim Set a
_ p
_ DimDecl a
AnyDim = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure MonoSize
MonoAnon
    onDim Set a
_ p
_ DimDecl a
d = do
      (Int
i, Map (DimDecl a) Int
m) <- f (Int, Map (DimDecl a) Int)
forall s (m :: * -> *). MonadState s m => m s
get
      case DimDecl a -> Map (DimDecl a) Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup DimDecl a
d Map (DimDecl a) 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 a) 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 a -> Int -> Map (DimDecl a) Int -> Map (DimDecl a) Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert DimDecl a
d Int
i Map (DimDecl a) 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 (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Exp
forall vn. QualName vn -> ExpBase Info vn
var QualName VName
fname
  | Bool
otherwise = do
    TypeBase (DimDecl VName) ()
t' <- TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t
    let mono_t :: MonoType
mono_t = TypeBase (DimDecl VName) () -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType TypeBase (DimDecl VName) ()
t'
    Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t
    Maybe PolyBinding
maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
lookupFun (VName -> MonoM (Maybe PolyBinding))
-> VName -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
    case (Maybe (VName, InferSizeArgs)
maybe_fname, Maybe PolyBinding
maybe_funbind) of
      -- The function has already been monomorphised.
      (Just (VName
fname', InferSizeArgs
infer), Maybe PolyBinding
_) ->
        Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase (DimDecl VName) () -> [Exp] -> Exp
forall vn as.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
      -- An intrinsic function.
      (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Exp
forall vn. QualName vn -> ExpBase Info vn
var QualName VName
fname
      -- 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 (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ VName -> TypeBase (DimDecl VName) () -> [Exp] -> Exp
forall vn as.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
  where
    var :: QualName vn -> ExpBase Info vn
var QualName vn
fname' = QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (PatternType -> Info PatternType
forall a. a -> Info a
Info (TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t)) SrcLoc
loc

    applySizeArg :: (Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg (Int
i, ExpBase Info vn
f) ExpBase Info vn
size_arg =
      ( Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,
        ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
          ExpBase Info vn
f
          ExpBase Info vn
size_arg
          ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing))
          (PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType (Int -> PatternType -> [PatternType]
forall a. Int -> a -> [a]
replicate Int
i PatternType
forall dim als. TypeBase dim als
i64) (TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
t)), [VName] -> Info [VName]
forall a. a -> Info a
Info [])
          SrcLoc
loc
      )

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

-- 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, PatternType) -> Alias)
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound (VName -> Alias)
-> ((VName, PatternType) -> VName) -> (VName, PatternType) -> Alias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, PatternType) -> VName
forall a b. (a, b) -> a
fst) ([(VName, PatternType)] -> [Alias])
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> a -> b
$ RecordReplacement -> [(VName, PatternType)]
forall k a. Map k a -> [a]
M.elems RecordReplacement
d
      replace Alias
x = Alias -> Aliasing
forall a. a -> Set a
S.singleton Alias
x
  -- 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 (m :: * -> *) a. Monad m => a -> m a
return (TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing))
-> TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall a b. (a -> b) -> a -> b
$
    if (Alias -> Bool) -> Aliasing -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> RecordReplacements -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` RecordReplacements
rrs) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (Aliasing -> Bool) -> Aliasing -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t
      then (Aliasing -> Aliasing)
-> TypeBase dim Aliasing -> TypeBase dim Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([Aliasing] -> Aliasing
forall a. Monoid a => [a] -> a
mconcat ([Aliasing] -> Aliasing)
-> (Aliasing -> [Aliasing]) -> Aliasing -> Aliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Aliasing) -> [Alias] -> [Aliasing]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> Aliasing
replace ([Alias] -> [Aliasing])
-> (Aliasing -> [Alias]) -> Aliasing -> [Aliasing]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliasing -> [Alias]
forall a. Set a -> [a]
S.toList) TypeBase dim Aliasing
t
      else TypeBase dim Aliasing
t

sizesForPat :: MonadFreshNames m => Pattern -> m ([VName], Pattern)
sizesForPat :: Pattern -> m ([VName], Pattern)
sizesForPat Pattern
pat = do
  (Pattern
params', [VName]
sizes) <- StateT [VName] m Pattern -> [VName] -> m (Pattern, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ASTMapper (StateT [VName] m) -> Pattern -> StateT [VName] m Pattern
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT [VName] m)
tv Pattern
pat) []
  ([VName], Pattern) -> m ([VName], Pattern)
forall (m :: * -> *) a. Monad m => a -> m a
return ([VName]
sizes, Pattern
params')
  where
    tv :: ASTMapper (StateT [VName] m)
tv = ASTMapper (StateT [VName] m)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatternType :: PatternType -> StateT [VName] m PatternType
mapOnPatternType = (DimDecl VName -> StateT [VName] m (DimDecl VName))
-> (Aliasing -> StateT [VName] m Aliasing)
-> PatternType
-> StateT [VName] m PatternType
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 DimDecl VName
AnyDim = 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

-- Monomorphization of expressions.
transformExp :: Exp -> MonoM Exp
transformExp :: Exp -> MonoM Exp
transformExp e :: Exp
e@Literal {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@IntLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@FloatLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp e :: Exp
e@StringLit {} = Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
transformExp (Parens Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) =
  (QualName VName, SrcLoc) -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (TupLit [Exp]
es SrcLoc
loc) =
  [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp) -> MonoM [Exp] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> [FieldBase Info VName] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField [FieldBase Info VName]
fs MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  where
    transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (RecordFieldExplicit Name
name Exp
e SrcLoc
loc') =
      Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name (Exp -> SrcLoc -> FieldBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> FieldBase Info VName)
-> MonoM SrcLoc -> MonoM (FieldBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc'
    transformField (RecordFieldImplicit VName
v Info PatternType
t SrcLoc
_) = do
      Info PatternType
t' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
t
      FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> FieldBase Info VName -> MonoM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$
        Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
          (VName -> Name
baseName VName
v)
          (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatternType
t' SrcLoc
loc)
          SrcLoc
loc
transformExp (ArrayLit [Exp]
es Info PatternType
t SrcLoc
loc) =
  [Exp] -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
ArrayLit ([Exp] -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl (Info PatternType, Info [VName])
tp SrcLoc
loc) = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Maybe Exp
me' <- (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Maybe Exp
me
  Inclusiveness Exp
incl' <- (Exp -> MonoM Exp)
-> Inclusiveness Exp -> MonoM (Inclusiveness Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp Inclusiveness Exp
incl
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp
-> Maybe Exp
-> Inclusiveness Exp
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' (Info PatternType, Info [VName])
tp SrcLoc
loc
transformExp (Var QualName VName
fname (Info PatternType
t) SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (VName -> MonoM (Maybe RecordReplacement))
-> VName -> MonoM (Maybe RecordReplacement)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
  case Maybe RecordReplacement
maybe_fs of
    Just RecordReplacement
fs -> do
      let toField :: (Name, (vn, PatternType)) -> MonoM (FieldBase Info vn)
toField (Name
f, (vn
f_v, PatternType
f_t)) = do
            PatternType
f_t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
f_t
            let f_v' :: ExpBase Info vn
f_v' = QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (vn -> QualName vn
forall v. v -> QualName v
qualName vn
f_v) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
f_t') SrcLoc
loc
            FieldBase Info vn -> MonoM (FieldBase Info vn)
forall (m :: * -> *) a. Monad m => a -> m a
return (FieldBase Info vn -> MonoM (FieldBase Info vn))
-> FieldBase Info vn -> MonoM (FieldBase Info vn)
forall a b. (a -> b) -> a -> b
$ Name -> ExpBase Info vn -> SrcLoc -> FieldBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
f ExpBase Info vn
f_v' SrcLoc
loc
      [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, (VName, PatternType)) -> MonoM (FieldBase Info VName))
-> [(Name, (VName, PatternType))] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name, (VName, PatternType)) -> MonoM (FieldBase Info VName)
forall vn. (Name, (vn, PatternType)) -> MonoM (FieldBase Info vn)
toField (RecordReplacement -> [(Name, (VName, PatternType))]
forall k a. Map k a -> [(k, a)]
M.toList RecordReplacement
fs) MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
    Maybe RecordReplacement
Nothing -> do
      PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
      SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t')
transformExp (Ascript Exp
e TypeDeclBase Info VName
tp SrcLoc
loc) =
  Exp -> TypeDeclBase Info VName -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeDeclBase f vn -> SrcLoc -> ExpBase f vn
Ascript (Exp -> TypeDeclBase Info VName -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (TypeDeclBase Info VName -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeDeclBase Info VName -> SrcLoc -> Exp)
-> MonoM (TypeDeclBase Info VName) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeDeclBase Info VName -> MonoM (TypeDeclBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeDeclBase Info VName
tp MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Coerce Exp
e TypeDeclBase Info VName
tp (Info PatternType
t, Info [VName]
ext) SrcLoc
loc) = do
  PatternType -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims PatternType
t
  Exp
-> TypeDeclBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeDeclBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Coerce (Exp
 -> TypeDeclBase Info VName
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> Exp)
-> MonoM Exp
-> MonoM
     (TypeDeclBase Info VName
      -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM
  (TypeDeclBase Info VName
   -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (TypeDeclBase Info VName)
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeDeclBase Info VName -> MonoM (TypeDeclBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeDeclBase Info VName
tp
    MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((,) (Info PatternType
 -> Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info PatternType)
-> MonoM (Info [VName] -> (Info PatternType, Info [VName]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> MonoM PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t) MonoM (Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info [VName]) -> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info [VName] -> MonoM (Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info [VName]
ext)
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (LetPat Pattern
pat Exp
e1 Exp
e2 (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) = do
  (Pattern
pat', RecordReplacements
rr) <- Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern Pattern
pat
  PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
  Pattern
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat Pattern
pat' (Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
    MonoM (Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (Exp -> MonoM Exp
transformExp Exp
e2)
    MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Info PatternType, Info [VName])
-> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t', Info [VName]
retext)
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, Info TypeBase (DimDecl VName) ()
ret, Exp
body) Exp
e Info PatternType
e_t SrcLoc
loc)
  | 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], [Pattern], Maybe (TypeExp VName),
    TypeBase (DimDecl VName) (), [VName], Exp, [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
ret, [], Exp
body, [AttrInfo]
forall a. Monoid a => a
mempty, SrcLoc
loc)
    MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
 -> MonoM Exp)
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall a b. (a -> b) -> a -> b
$ do
      (Exp
e', Seq (VName, ValBind)
bs) <- MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (MonoM Exp -> MonoM (Exp, Seq (VName, ValBind)))
-> MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ VName -> PolyBinding -> MonoM Exp -> MonoM Exp
forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
fname PolyBinding
funbind (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
e
      -- 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 (m :: * -> *) a. Monad m => a -> m a
return ([ValBind] -> Exp -> Exp
unfoldLetFuns (((VName, ValBind) -> ValBind) -> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> [a] -> [b]
map (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd ([(VName, ValBind)] -> [ValBind])
-> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> a -> b
$ Seq (VName, ValBind) -> [(VName, ValBind)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq (VName, ValBind)
bs_local) Exp
e', Seq (VName, ValBind)
-> Seq (VName, ValBind) -> Seq (VName, ValBind)
forall a b. a -> b -> a
const Seq (VName, ValBind)
bs_prop)
  | Bool
otherwise = do
    Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
    VName
-> ([TypeParamBase VName], [Pattern], Maybe (TypeExp VName),
    Info (TypeBase (DimDecl VName) ()), Exp)
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn),
    f (TypeBase (DimDecl VName) ()), ExpBase f vn)
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) () -> Info (TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info TypeBase (DimDecl VName) ()
ret, Exp
body')
      (Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
e_t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (If Exp
e1 Exp
e2 Exp
e3 (Info PatternType
tp, Info [VName]
retext) SrcLoc
loc) = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
  Exp
e3' <- Exp -> MonoM Exp
transformExp Exp
e3
  Info PatternType
tp' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
tp
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
If Exp
e1' Exp
e2' Exp
e3' (Info PatternType
tp', Info [VName]
retext) SrcLoc
loc
transformExp (Apply Exp
e1 Exp
e2 Info (Diet, Maybe VName)
d (Info PatternType
ret, Info [VName]
ext) SrcLoc
loc) = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
  Info PatternType
ret' <- (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
ret
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply Exp
e1' Exp
e2' Info (Diet, Maybe VName)
d (Info PatternType
ret', Info [VName]
ext) SrcLoc
loc
transformExp (Negate Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Lambda [Pattern]
params Exp
e0 Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
tp SrcLoc
loc) = do
  Exp
e0' <- Exp -> MonoM Exp
transformExp Exp
e0
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [Pattern]
params Exp
e0' Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
tp SrcLoc
loc
transformExp (OpSection QualName VName
qn Info PatternType
t SrcLoc
loc) =
  Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info PatternType
t SrcLoc
loc
transformExp
  ( OpSectionLeft
      QualName VName
fname
      (Info PatternType
t)
      Exp
e
      (Info (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext), Info TypeBase (DimDecl VName) ()
ytype)
      (Info PatternType
rettype, Info [VName]
retext)
      SrcLoc
loc
    ) = do
    Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
    Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
    Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
      Exp
fname'
      (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
      Maybe Exp
forall a. Maybe a
Nothing
      PatternType
t
      (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext)
      (TypeBase (DimDecl VName) ()
ytype, Maybe VName
forall a. Maybe a
Nothing)
      (PatternType
rettype, [VName]
retext)
      SrcLoc
loc
transformExp
  ( OpSectionRight
      QualName VName
fname
      (Info PatternType
t)
      Exp
e
      (Info TypeBase (DimDecl VName) ()
xtype, Info (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext))
      (Info PatternType
rettype)
      SrcLoc
loc
    ) = do
    Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
    Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
    Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
      Exp
fname'
      Maybe Exp
forall a. Maybe a
Nothing
      (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
      PatternType
t
      (TypeBase (DimDecl VName) ()
xtype, Maybe VName
forall a. Maybe a
Nothing)
      (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)
      (PatternType
rettype, [])
      SrcLoc
loc
transformExp (ProjectSection [Name]
fields (Info PatternType
t) SrcLoc
loc) =
  [Name] -> PatternType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields PatternType
t SrcLoc
loc
transformExp (IndexSection [DimIndexBase Info VName]
idxs (Info PatternType
t) SrcLoc
loc) =
  [DimIndexBase Info VName] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection [DimIndexBase Info VName]
idxs PatternType
t SrcLoc
loc
transformExp (DoLoop [VName]
sparams Pattern
pat Exp
e1 LoopFormBase Info VName
form Exp
e3 Info (PatternType, [VName])
ret SrcLoc
loc) = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  LoopFormBase Info VName
form' <- case LoopFormBase Info VName
form of
    For IdentBase Info VName
ident Exp
e2 -> IdentBase Info VName -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
ident (Exp -> LoopFormBase Info VName)
-> MonoM Exp -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
    ForIn Pattern
pat2 Exp
e2 -> Pattern -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn Pattern
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, Pattern
pat') <- Pattern -> MonoM ([VName], Pattern)
forall (m :: * -> *).
MonadFreshNames m =>
Pattern -> m ([VName], Pattern)
sizesForPat Pattern
pat
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [VName]
-> Pattern
-> Exp
-> LoopFormBase Info VName
-> Exp
-> Info (PatternType, [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[VName]
-> PatternBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> f (PatternType, [VName])
-> SrcLoc
-> ExpBase f vn
DoLoop ([VName]
sparams [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
pat_sizes) Pattern
pat' Exp
e1' LoopFormBase Info VName
form' Exp
e3' Info (PatternType, [VName])
ret SrcLoc
loc
transformExp (BinOp (QualName VName
fname, SrcLoc
oploc) (Info PatternType
t) (Exp
e1, Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2, Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) Info PatternType
tp Info [VName]
ext SrcLoc
loc) = do
  Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
  case Exp
fname' of
    Var QualName VName
fname'' Info PatternType
_ SrcLoc
_
      | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf Exp
e1'),
        PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf Exp
e2') ->
        Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ (QualName VName, SrcLoc)
-> Info PatternType
-> (Exp, Info (TypeBase (DimDecl VName) (), Maybe VName))
-> (Exp, Info (TypeBase (DimDecl VName) (), Maybe VName))
-> Info PatternType
-> Info [VName]
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc)
-> f PatternType
-> (ExpBase f vn, f (TypeBase (DimDecl VName) (), Maybe VName))
-> (ExpBase f vn, f (TypeBase (DimDecl VName) (), Maybe VName))
-> f PatternType
-> f [VName]
-> SrcLoc
-> ExpBase f vn
BinOp (QualName VName
fname'', SrcLoc
oploc) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) (Exp
e1', Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (Exp
e2', Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) Info PatternType
tp Info [VName]
ext SrcLoc
loc
    Exp
_ -> do
      -- 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, Pattern
x_param) <- Exp -> MonoM (Exp, Pattern)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pattern)
makeVarParam Exp
e1'
      (Exp
y_param_e, Pattern
y_param) <- Exp -> MonoM (Exp, Pattern)
forall (m :: * -> *). MonadFreshNames m => Exp -> m (Exp, Pattern)
makeVarParam Exp
e2'
      Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
        Pattern
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat
          Pattern
x_param
          Exp
e1'
          ( Pattern
-> Exp -> Exp -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat
              Pattern
y_param
              Exp
e2'
              (Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
x_param_e Exp
y_param_e)
              (Info PatternType
tp, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty)
              SrcLoc
forall a. Monoid a => a
mempty
          )
          (Info PatternType
tp, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty)
          SrcLoc
forall a. Monoid a => a
mempty
  where
    applyOp :: ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp ExpBase Info vn
fname' ExpBase Info vn
x ExpBase Info vn
y =
      ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
        ( ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
            ExpBase Info vn
fname'
            ExpBase Info vn
x
            ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d1)))
            ( PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct (TypeBase (DimDecl VName) () -> PatternType)
-> TypeBase (DimDecl VName) () -> PatternType
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) (), Maybe VName)
-> TypeBase (DimDecl VName) ()
forall a b. (a, b) -> a
fst (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d2)] (Info PatternType -> PatternType
forall a. Info a -> a
unInfo Info PatternType
tp)),
              [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty
            )
            SrcLoc
loc
        )
        ExpBase Info vn
y
        ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, (TypeBase (DimDecl VName) (), Maybe VName) -> Maybe VName
forall a b. (a, b) -> b
snd (Info (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
forall a. Info a -> a
unInfo Info (TypeBase (DimDecl VName) (), Maybe VName)
d2)))
        (Info PatternType
tp, Info [VName]
ext)
        SrcLoc
loc

    makeVarParam :: Exp -> m (Exp, Pattern)
makeVarParam Exp
arg = do
      let argtype :: PatternType
argtype = Exp -> PatternType
typeOf Exp
arg
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"binop_p"
      (Exp, Pattern) -> m (Exp, Pattern)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
x (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty
        )
transformExp (Project Name
n Exp
e Info PatternType
tp SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- case Exp
e of
    Var QualName VName
qn Info PatternType
_ SrcLoc
_ -> VName -> MonoM (Maybe RecordReplacement)
lookupRecordReplacement (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
    Exp
_ -> Maybe RecordReplacement -> MonoM (Maybe RecordReplacement)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe RecordReplacement
forall a. Maybe a
Nothing
  case Maybe RecordReplacement
maybe_fs of
    Just RecordReplacement
m
      | Just (VName
v, PatternType
_) <- Name -> RecordReplacement -> Maybe (VName, PatternType)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n RecordReplacement
m ->
        Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info PatternType
tp SrcLoc
loc
    Maybe RecordReplacement
_ -> do
      Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
      Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info PatternType
tp SrcLoc
loc
transformExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs Exp
e1 Exp
body (Info PatternType
t) SrcLoc
loc) = do
  [DimIndexBase Info VName]
idxs' <- (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex [DimIndexBase Info VName]
idxs
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
body' <- Exp -> MonoM Exp
transformExp Exp
body
  PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName
-> IdentBase Info VName
-> [DimIndexBase Info VName]
-> Exp
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
IdentBase f vn
-> IdentBase f vn
-> [DimIndexBase f vn]
-> ExpBase f vn
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs' Exp
e1' Exp
body' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t') SrcLoc
loc
transformExp (Index Exp
e0 [DimIndexBase Info VName]
idxs (Info PatternType, Info [VName])
info SrcLoc
loc) =
  Exp
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (Exp
 -> [DimIndexBase Info VName]
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> Exp)
-> MonoM Exp
-> MonoM
     ([DimIndexBase Info VName]
      -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e0 MonoM
  ([DimIndexBase Info VName]
   -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM [DimIndexBase Info VName]
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex [DimIndexBase Info VName]
idxs MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Info PatternType, Info [VName])
-> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Info PatternType, Info [VName])
info MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Update Exp
e1 [DimIndexBase Info VName]
idxs Exp
e2 SrcLoc
loc) =
  Exp -> [DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (Exp -> [DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([DimIndexBase Info VName] -> Exp -> SrcLoc -> Exp)
-> MonoM [DimIndexBase Info VName] -> MonoM (Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (DimIndexBase Info VName -> MonoM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> MonoM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex [DimIndexBase Info VName]
idxs
    MonoM (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info PatternType
t SrcLoc
loc) =
  Exp -> [Name] -> Exp -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name]
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
RecordUpdate (Exp -> [Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM ([Name] -> Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Name] -> MonoM (Exp -> Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Name] -> MonoM [Name]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
fs
    MonoM (Exp -> Info PatternType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatternType
t
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Assert Exp
e1 Exp
e2 Info String
desc SrcLoc
loc) =
  Exp -> Exp -> Info String -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f String -> SrcLoc -> ExpBase f vn
Assert (Exp -> Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Exp -> Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info String -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info String -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info String -> SrcLoc -> Exp)
-> MonoM (Info String) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info String -> MonoM (Info String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info String
desc MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Constr Name
name [Exp]
all_es Info PatternType
t SrcLoc
loc) =
  Name -> [Exp] -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
Constr Name
name ([Exp] -> Info PatternType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info PatternType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Exp -> MonoM Exp
transformExp [Exp]
all_es MonoM (Info PatternType -> SrcLoc -> Exp)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info PatternType -> MonoM (Info PatternType)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info PatternType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) =
  Exp
-> NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Match (Exp
 -> NonEmpty (CaseBase Info VName)
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> Exp)
-> MonoM Exp
-> MonoM
     (NonEmpty (CaseBase Info VName)
      -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM
  (NonEmpty (CaseBase Info VName)
   -> (Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (NonEmpty (CaseBase Info VName))
-> MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (CaseBase Info VName -> MonoM (CaseBase Info VName))
-> NonEmpty (CaseBase Info VName)
-> MonoM (NonEmpty (CaseBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase NonEmpty (CaseBase Info VName)
cs
    MonoM ((Info PatternType, Info [VName]) -> SrcLoc -> Exp)
-> MonoM (Info PatternType, Info [VName]) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((,) (Info PatternType
 -> Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info PatternType)
-> MonoM (Info [VName] -> (Info PatternType, Info [VName]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PatternType -> MonoM PatternType)
-> Info PatternType -> MonoM (Info PatternType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType Info PatternType
t MonoM (Info [VName] -> (Info PatternType, Info [VName]))
-> MonoM (Info [VName]) -> MonoM (Info PatternType, Info [VName])
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info [VName] -> MonoM (Info [VName])
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info [VName]
retext)
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Attr AttrInfo
info Exp
e SrcLoc
loc) =
  AttrInfo -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo
info (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

transformCase :: Case -> MonoM Case
transformCase :: CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase (CasePat Pattern
p Exp
e SrcLoc
loc) = do
  (Pattern
p', RecordReplacements
rr) <- Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern Pattern
p
  Pattern -> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pattern
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 ->
  PatternType ->
  (StructType, Maybe VName) ->
  (StructType, Maybe VName) ->
  (PatternType, [VName]) ->
  SrcLoc ->
  MonoM Exp
desugarBinOpSection :: Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection Exp
op Maybe Exp
e_left Maybe Exp
e_right PatternType
t (TypeBase (DimDecl VName) ()
xtype, Maybe VName
xext) (TypeBase (DimDecl VName) ()
ytype, Maybe VName
yext) (PatternType
rettype, [VName]
retext) SrcLoc
loc = do
  (Exp
e1, [Pattern]
p1) <- Maybe Exp -> PatternType -> MonoM (Exp, [Pattern])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (Exp, [Pattern])
makeVarParam Maybe Exp
e_left (PatternType -> MonoM (Exp, [Pattern]))
-> PatternType -> MonoM (Exp, [Pattern])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
xtype
  (Exp
e2, [Pattern]
p2) <- Maybe Exp -> PatternType -> MonoM (Exp, [Pattern])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (Exp, [Pattern])
makeVarParam Maybe Exp
e_right (PatternType -> MonoM (Exp, [Pattern]))
-> PatternType -> MonoM (Exp, [Pattern])
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype
  let apply_left :: Exp
apply_left =
        Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
          Exp
op
          Exp
e1
          ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
xext))
          (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase (DimDecl VName) ()
ytype] PatternType
t, [VName] -> Info [VName]
forall a. a -> Info a
Info [])
          SrcLoc
loc
      body :: Exp
body =
        Exp
-> Exp
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
          Exp
apply_left
          Exp
e2
          ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
yext))
          (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
rettype, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
retext)
          SrcLoc
loc
      rettype' :: TypeBase (DimDecl VName) ()
rettype' = PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
rettype
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda ([Pattern]
p1 [Pattern] -> [Pattern] -> [Pattern]
forall a. [a] -> [a] -> [a]
++ [Pattern]
p2) Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, TypeBase (DimDecl VName) ())
-> Info (Aliasing, TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, TypeBase (DimDecl VName) ()
rettype')) SrcLoc
loc
  where
    makeVarParam :: Maybe Exp -> PatternType -> m (Exp, [Pattern])
makeVarParam (Just Exp
e) PatternType
_ = (Exp, [Pattern]) -> m (Exp, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp
e, [])
    makeVarParam Maybe Exp
Nothing PatternType
argtype = do
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
      (Exp, [Pattern]) -> m (Exp, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          [VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
x (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct PatternType
argtype) SrcLoc
forall a. Monoid a => a
mempty]
        )

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

desugarIndexSection :: [DimIndex] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection :: [DimIndexBase Info VName] -> PatternType -> SrcLoc -> MonoM Exp
desugarIndexSection [DimIndexBase Info VName]
idxs (Scalar (Arrow Aliasing
_ PName
_ PatternType
t1 PatternType
t2)) SrcLoc
loc = do
  VName
p <- String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"index_i"
  let body :: Exp
body = Exp
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (QualName VName -> Info PatternType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
p) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t1) SrcLoc
loc) [DimIndexBase Info VName]
idxs (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t2, [VName] -> Info [VName]
forall a. a -> Info a
Info []) SrcLoc
loc
  Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> Exp
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
p (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t1) SrcLoc
forall a. Monoid a => a
mempty] Exp
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Aliasing, TypeBase (DimDecl VName) ())
-> Info (Aliasing, TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info (Aliasing
forall a. Monoid a => a
mempty, PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t2)) SrcLoc
loc
desugarIndexSection [DimIndexBase Info VName]
_ PatternType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarIndexSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t

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

-- 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 (TypeBase (DimDecl VName) ()
rettype, [VName]
_)) [TypeParamBase VName]
dim_params [Pattern]
params Exp
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
loc : [ValBind]
rest) Exp
e =
  VName
-> ([TypeParamBase VName], [Pattern], Maybe (TypeExp VName),
    Info (TypeBase (DimDecl VName) ()), Exp)
-> Exp
-> Info PatternType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn),
    f (TypeBase (DimDecl VName) ()), ExpBase f vn)
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
dim_params, [Pattern]
params, Maybe (TypeExp VName)
forall a. Maybe a
Nothing, TypeBase (DimDecl VName) () -> Info (TypeBase (DimDecl VName) ())
forall a. a -> Info a
Info TypeBase (DimDecl VName) ()
rettype, Exp
body) Exp
e' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
e_t) SrcLoc
loc
  where
    e' :: Exp
e' = [ValBind] -> Exp -> Exp
unfoldLetFuns [ValBind]
rest Exp
e
    e_t :: PatternType
e_t = Exp -> PatternType
typeOf Exp
e'

transformPattern :: Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern :: Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern (Id VName
v (Info (Scalar (Record Map Name PatternType
fs))) SrcLoc
loc) = do
  let fs' :: [(Name, PatternType)]
fs' = Map Name PatternType -> [(Name, PatternType)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name PatternType
fs
  ([VName]
fs_ks, [PatternType]
fs_ts) <- ([(VName, PatternType)] -> ([VName], [PatternType]))
-> MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(VName, PatternType)] -> ([VName], [PatternType])
forall a b. [(a, b)] -> ([a], [b])
unzip (MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType]))
-> MonoM [(VName, PatternType)] -> MonoM ([VName], [PatternType])
forall a b. (a -> b) -> a -> b
$
    [(Name, PatternType)]
-> ((Name, PatternType) -> MonoM (VName, PatternType))
-> MonoM [(VName, PatternType)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Name, PatternType)]
fs' (((Name, PatternType) -> MonoM (VName, PatternType))
 -> MonoM [(VName, PatternType)])
-> ((Name, PatternType) -> MonoM (VName, PatternType))
-> MonoM [(VName, PatternType)]
forall a b. (a -> b) -> a -> b
$ \(Name
f, PatternType
ft) ->
      (,) (VName -> PatternType -> (VName, PatternType))
-> MonoM VName -> MonoM (PatternType -> (VName, PatternType))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName (Name -> String
nameToString Name
f) MonoM (PatternType -> (VName, PatternType))
-> MonoM PatternType -> MonoM (VName, PatternType)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
ft
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( [(Name, Pattern)] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern
        ( [Name] -> [Pattern] -> [(Name, Pattern)]
forall a b. [a] -> [b] -> [(a, b)]
zip
            (((Name, PatternType) -> Name) -> [(Name, PatternType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternType)]
fs')
            ((VName -> Info PatternType -> SrcLoc -> Pattern)
-> [VName] -> [Info PatternType] -> [SrcLoc] -> [Pattern]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id [VName]
fs_ks ((PatternType -> Info PatternType)
-> [PatternType] -> [Info PatternType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> Info PatternType
forall a. a -> Info a
Info [PatternType]
fs_ts) ([SrcLoc] -> [Pattern]) -> [SrcLoc] -> [Pattern]
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [SrcLoc]
forall a. a -> [a]
repeat SrcLoc
loc)
        )
        SrcLoc
loc,
      VName -> RecordReplacement -> RecordReplacements
forall k a. k -> a -> Map k a
M.singleton VName
v (RecordReplacement -> RecordReplacements)
-> RecordReplacement -> RecordReplacements
forall a b. (a -> b) -> a -> b
$ [(Name, (VName, PatternType))] -> RecordReplacement
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, (VName, PatternType))] -> RecordReplacement)
-> [(Name, (VName, PatternType))] -> RecordReplacement
forall a b. (a -> b) -> a -> b
$ [Name] -> [(VName, PatternType)] -> [(Name, (VName, PatternType))]
forall a b. [a] -> [b] -> [(a, b)]
zip (((Name, PatternType) -> Name) -> [(Name, PatternType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternType) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternType)]
fs') ([(VName, PatternType)] -> [(Name, (VName, PatternType))])
-> [(VName, PatternType)] -> [(Name, (VName, PatternType))]
forall a b. (a -> b) -> a -> b
$ [VName] -> [PatternType] -> [(VName, PatternType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
fs_ks [PatternType]
fs_ts
    )
transformPattern (Id VName
v Info PatternType
t SrcLoc
loc) = (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
v Info PatternType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (TuplePattern [Pattern]
pats SrcLoc
loc) = do
  ([Pattern]
pats', [RecordReplacements]
rrs) <- [(Pattern, RecordReplacements)]
-> ([Pattern], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pattern, RecordReplacements)]
 -> ([Pattern], [RecordReplacements]))
-> MonoM [(Pattern, RecordReplacements)]
-> MonoM ([Pattern], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pattern -> MonoM (Pattern, RecordReplacements))
-> [Pattern] -> MonoM [(Pattern, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern [Pattern]
pats
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern [Pattern]
pats' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPattern (RecordPattern [(Name, Pattern)]
fields SrcLoc
loc) = do
  let ([Name]
field_names, [Pattern]
field_pats) = [(Name, Pattern)] -> ([Name], [Pattern])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, Pattern)]
fields
  ([Pattern]
field_pats', [RecordReplacements]
rrs) <- [(Pattern, RecordReplacements)]
-> ([Pattern], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pattern, RecordReplacements)]
 -> ([Pattern], [RecordReplacements]))
-> MonoM [(Pattern, RecordReplacements)]
-> MonoM ([Pattern], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pattern -> MonoM (Pattern, RecordReplacements))
-> [Pattern] -> MonoM [(Pattern, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern [Pattern]
field_pats
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, Pattern)] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern ([Name] -> [Pattern] -> [(Name, Pattern)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
field_names [Pattern]
field_pats') SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)
transformPattern (PatternParens Pattern
pat SrcLoc
loc) = do
  (Pattern
pat', RecordReplacements
rr) <- Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern Pattern
pat
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Pattern -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatternBase f vn -> SrcLoc -> PatternBase f vn
PatternParens Pattern
pat' SrcLoc
loc, RecordReplacements
rr)
transformPattern (Wildcard (Info PatternType
t) SrcLoc
loc) = do
  PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternType -> SrcLoc -> Pattern
wildcard PatternType
t' SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (PatternAscription Pattern
pat TypeDeclBase Info VName
td SrcLoc
loc) = do
  (Pattern
pat', RecordReplacements
rr) <- Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern Pattern
pat
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Pattern -> TypeDeclBase Info VName -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatternBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatternBase f vn
PatternAscription Pattern
pat' TypeDeclBase Info VName
td SrcLoc
loc, RecordReplacements
rr)
transformPattern (PatternLit PatLit
e Info PatternType
t SrcLoc
loc) = (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatLit -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatLit -> f PatternType -> SrcLoc -> PatternBase f vn
PatternLit PatLit
e Info PatternType
t SrcLoc
loc, RecordReplacements
forall a. Monoid a => a
mempty)
transformPattern (PatternConstr Name
name Info PatternType
t [Pattern]
all_ps SrcLoc
loc) = do
  ([Pattern]
all_ps', [RecordReplacements]
rrs) <- [(Pattern, RecordReplacements)]
-> ([Pattern], [RecordReplacements])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Pattern, RecordReplacements)]
 -> ([Pattern], [RecordReplacements]))
-> MonoM [(Pattern, RecordReplacements)]
-> MonoM ([Pattern], [RecordReplacements])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Pattern -> MonoM (Pattern, RecordReplacements))
-> [Pattern] -> MonoM [(Pattern, RecordReplacements)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern [Pattern]
all_ps
  (Pattern, RecordReplacements)
-> MonoM (Pattern, RecordReplacements)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Info PatternType -> [Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
name Info PatternType
t [Pattern]
all_ps' SrcLoc
loc, [RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs)

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

type DimInst = M.Map VName (DimDecl VName)

dimMapping ::
  Monoid a =>
  TypeBase (DimDecl VName) a ->
  TypeBase (DimDecl VName) a ->
  DimInst
dimMapping :: TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2 = State DimInst (TypeBase (DimDecl VName) a) -> DimInst -> DimInst
forall s a. State s a -> s -> s
execState ((DimDecl VName
 -> DimDecl VName -> StateT DimInst Identity (DimDecl VName))
-> TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a
-> State DimInst (TypeBase (DimDecl VName) a)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims DimDecl VName
-> DimDecl VName -> StateT DimInst Identity (DimDecl VName)
forall (m :: * -> *) vn a.
(MonadState (Map vn a) m, Ord vn) =>
DimDecl vn -> a -> m (DimDecl vn)
f TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
t2) DimInst
forall a. Monoid a => a
mempty
  where
    f :: DimDecl vn -> a -> m (DimDecl vn)
f (NamedDim QualName vn
d1) a
d2 = do
      (Map vn a -> Map vn a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map vn a -> Map vn a) -> m ()) -> (Map vn a -> Map vn a) -> m ()
forall a b. (a -> b) -> a -> b
$ vn -> a -> Map vn a -> Map vn a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName vn -> vn
forall vn. QualName vn -> vn
qualLeaf QualName vn
d1) a
d2
      DimDecl vn -> m (DimDecl vn)
forall (m :: * -> *) a. Monad m => a -> m a
return (DimDecl vn -> m (DimDecl vn)) -> DimDecl vn -> m (DimDecl vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> DimDecl vn
forall vn. QualName vn -> DimDecl vn
NamedDim QualName vn
d1
    f DimDecl vn
d a
_ = DimDecl vn -> m (DimDecl vn)
forall (m :: * -> *) a. Monad m => a -> m a
return DimDecl vn
d

inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp]
inferSizeArgs :: [TypeParamBase VName]
-> TypeBase (DimDecl VName) () -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t =
  (TypeParamBase VName -> Maybe Exp)
-> [TypeParamBase VName] -> [Exp]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DimInst -> TypeParamBase VName -> Maybe Exp
forall k vn.
Ord k =>
Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg (TypeBase (DimDecl VName) ()
-> TypeBase (DimDecl VName) () -> DimInst
forall a.
Monoid a =>
TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping TypeBase (DimDecl VName) ()
bind_t TypeBase (DimDecl VName) ()
t)) [TypeParamBase VName]
tparams
  where
    tparamArg :: Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg Map k (DimDecl vn)
dinst TypeParamBase k
tp =
      case k -> Map k (DimDecl vn) -> Maybe (DimDecl vn)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (TypeParamBase k -> k
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase k
tp) Map k (DimDecl vn)
dinst of
        Just (NamedDim QualName vn
d) ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
d (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
forall dim als. TypeBase dim als
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)
_ ->
          Maybe (ExpBase Info vn)
forall a. Maybe a
Nothing

explicitSizes :: StructType -> MonoType -> S.Set VName
explicitSizes :: TypeBase (DimDecl VName) () -> MonoType -> Set VName
explicitSizes TypeBase (DimDecl VName) ()
t1 MonoType
t2 =
  State (Set VName) (TypeBase (DimDecl VName) ())
-> Set VName -> Set VName
forall s a. State s a -> s -> s
execState ((DimDecl VName
 -> MonoSize -> StateT (Set VName) Identity (DimDecl VName))
-> TypeBase (DimDecl VName) ()
-> MonoType
-> State (Set VName) (TypeBase (DimDecl VName) ())
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
(d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims DimDecl VName
-> MonoSize -> StateT (Set VName) Identity (DimDecl VName)
forall (m :: * -> *) a.
(MonadState (Set a) m, Ord a) =>
DimDecl a -> MonoSize -> m (DimDecl a)
onDims TypeBase (DimDecl VName) ()
t1 MonoType
t2) Set VName
forall a. Monoid a => a
mempty Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` TypeBase (DimDecl VName) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
t1
  where
    onDims :: DimDecl a -> MonoSize -> m (DimDecl a)
onDims DimDecl a
d1 MonoSize
d2 = do
      case (DimDecl a
d1, MonoSize
d2) of
        (NamedDim QualName a
v, MonoKnown Int
_) -> (Set a -> Set a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set a -> Set a) -> m ()) -> (Set a -> Set a) -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert (a -> Set a -> Set a) -> a -> Set a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
v
        (DimDecl a, MonoSize)
_ -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      DimDecl a -> m (DimDecl a)
forall (m :: * -> *) a. Monad m => a -> m a
return DimDecl a
d1

-- 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 ScalarTypeBase dim ()
t ShapeDecl dim
shape) = ()
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
u (ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t) ShapeDecl dim
shape
    f (Scalar ScalarTypeBase dim ()
t) = ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t
    f' :: ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' (Arrow () PName
_ TypeBase dim ()
t1 TypeBase dim ()
t2) =
      ()
-> PName
-> TypeBase dim ()
-> TypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t2)
    f' (Record Map Name (TypeBase dim ())
fs) =
      Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim ()) -> ScalarTypeBase dim ())
-> Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ (TypeBase dim () -> TypeBase dim ())
-> Map Name (TypeBase dim ()) -> Map Name (TypeBase dim ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim () -> TypeBase dim ()
f Map Name (TypeBase dim ())
fs
    f' (Sum Map Name [TypeBase dim ()]
cs) =
      Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim ()] -> ScalarTypeBase dim ())
-> Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim ()] -> [TypeBase dim ()])
-> Map Name [TypeBase dim ()] -> Map Name [TypeBase dim ()]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim () -> TypeBase dim ())
-> [TypeBase dim ()] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim () -> TypeBase dim ()
f) Map Name [TypeBase dim ()]
cs
    f' ScalarTypeBase dim ()
t = ScalarTypeBase dim ()
t

-- 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, [Pattern]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
rettype, [VName]
retext, Exp
body, [AttrInfo]
attrs, SrcLoc
loc)) MonoType
t =
  RecordReplacements
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr (MonoM (VName, InferSizeArgs, ValBind)
 -> MonoM (VName, InferSizeArgs, ValBind))
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ do
    let bind_t :: TypeBase (DimDecl VName) ()
bind_t = [TypeBase (DimDecl VName) ()]
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ((Pattern -> TypeBase (DimDecl VName) ())
-> [Pattern] -> [TypeBase (DimDecl VName) ()]
forall a b. (a -> b) -> [a] -> [b]
map Pattern -> TypeBase (DimDecl VName) ()
patternStructType [Pattern]
params) TypeBase (DimDecl VName) ()
rettype
    (Map VName (TypeBase (DimDecl VName) ())
substs, [TypeParamBase VName]
t_shape_params) <- SrcLoc
-> TypeBase () ()
-> MonoType
-> MonoM
     (Map VName (TypeBase (DimDecl VName) ()), [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName (TypeBase (DimDecl VName) ()),
      [TypeParamBase VName])
typeSubstsM SrcLoc
loc (TypeBase (DimDecl VName) () -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes TypeBase (DimDecl VName) ()
bind_t) (MonoType
 -> MonoM
      (Map VName (TypeBase (DimDecl VName) ()), [TypeParamBase VName]))
-> MonoType
-> MonoM
     (Map VName (TypeBase (DimDecl VName) ()), [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
t
    let substs' :: Map VName (Subst (TypeBase (DimDecl VName) ()))
substs' = (TypeBase (DimDecl VName) ()
 -> Subst (TypeBase (DimDecl VName) ()))
-> Map VName (TypeBase (DimDecl VName) ())
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase (DimDecl VName) () -> Subst (TypeBase (DimDecl VName) ())
forall t. t -> Subst t
Subst Map VName (TypeBase (DimDecl VName) ())
substs
        rettype' :: TypeBase (DimDecl VName) ()
rettype' = (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs') TypeBase (DimDecl VName) ()
rettype
        substPatternType :: PatternType -> PatternType
substPatternType =
          (VName -> Maybe (Subst PatternType)) -> PatternType -> PatternType
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((Subst (TypeBase (DimDecl VName) ()) -> Subst PatternType)
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase (DimDecl VName) () -> PatternType)
-> Subst (TypeBase (DimDecl VName) ()) -> Subst PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase (DimDecl VName) () -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct) (Maybe (Subst (TypeBase (DimDecl VName) ()))
 -> Maybe (Subst PatternType))
-> (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> VName
-> Maybe (Subst PatternType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs'))
        params' :: [Pattern]
params' = (Pattern -> Pattern) -> [Pattern] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
substPatternType) [Pattern]
params
        bind_t' :: TypeBase (DimDecl VName) ()
bind_t' = (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs') TypeBase (DimDecl VName) ()
bind_t
        ([TypeParamBase VName]
shape_params_explicit, [TypeParamBase VName]
shape_params_implicit) =
          (TypeParamBase VName -> Bool)
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` TypeBase (DimDecl VName) () -> MonoType -> Set VName
explicitSizes TypeBase (DimDecl VName) ()
bind_t' MonoType
t) (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) ([TypeParamBase VName]
 -> ([TypeParamBase VName], [TypeParamBase VName]))
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$
            [TypeParamBase VName]
shape_params [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
t_shape_params

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

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

    Exp
body' <- (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> x -> m x
updateExpTypes (VName
-> Map VName (Subst (TypeBase (DimDecl VName) ()))
-> Maybe (Subst (TypeBase (DimDecl VName) ()))
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst (TypeBase (DimDecl VName) ()))
substs') Exp
body
    Exp
body'' <- RecordReplacements -> MonoM Exp -> MonoM Exp
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body'
    VName
name' <- if [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParamBase VName]
tparams Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
entry then VName -> MonoM VName
forall (m :: * -> *) a. Monad m => a -> m a
return VName
name else VName -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
name

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

    updateExpTypes :: (VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> x -> m x
updateExpTypes VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs = ASTMapper m -> x -> m x
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m -> x -> m x) -> ASTMapper m -> x -> m x
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
forall (m :: * -> *).
Monad m =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs
    mapper :: (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs =
      ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
        { mapOnExp :: Exp -> m Exp
mapOnExp = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m -> Exp -> m Exp) -> ASTMapper m -> Exp -> m Exp
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ASTMapper m
mapper VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs,
          mapOnName :: VName -> m VName
mapOnName = VName -> m VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnQualName :: QualName VName -> m (QualName VName)
mapOnQualName = QualName VName -> m (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnStructType :: TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> m (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall a.
Substitutable a =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> a -> a
applySubst VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs,
          mapOnPatternType :: PatternType -> m PatternType
mapOnPatternType = PatternType -> m PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> m PatternType)
-> (PatternType -> PatternType) -> PatternType -> m PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> PatternType -> PatternType
forall a.
Substitutable a =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))) -> a -> a
applySubst VName -> Maybe (Subst (TypeBase (DimDecl VName) ()))
substs
        }

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

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

typeSubstsM ::
  MonadFreshNames m =>
  SrcLoc ->
  TypeBase () () ->
  MonoType ->
  m (M.Map VName StructType, [TypeParam])
typeSubstsM :: SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName (TypeBase (DimDecl VName) ()),
      [TypeParamBase VName])
typeSubstsM SrcLoc
loc TypeBase () ()
orig_t1 MonoType
orig_t2 =
  let m :: StateT
  (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
  (WriterT [TypeParamBase VName] m)
  ()
m = TypeBase () ()
-> MonoType
-> StateT
     (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
     (WriterT [TypeParamBase VName] m)
     ()
forall d (t :: (* -> *) -> * -> *) (t :: (* -> *) -> * -> *)
       (m :: * -> *) dim as.
(MonadState
   (Map VName (TypeBase (DimDecl VName) d), Map Int VName) (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)),
 Pretty (ShapeDecl dim), Monad (t m)) =>
TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase () ()
orig_t1 MonoType
orig_t2
   in WriterT
  [TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
-> m (Map VName (TypeBase (DimDecl VName) ()),
      [TypeParamBase VName])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT
   [TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
 -> m (Map VName (TypeBase (DimDecl VName) ()),
       [TypeParamBase VName]))
-> WriterT
     [TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
-> m (Map VName (TypeBase (DimDecl VName) ()),
      [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
-> Map VName (TypeBase (DimDecl VName) ())
forall a b. (a, b) -> a
fst ((Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
 -> Map VName (TypeBase (DimDecl VName) ()))
-> WriterT
     [TypeParamBase VName]
     m
     (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
-> WriterT
     [TypeParamBase VName] m (Map VName (TypeBase (DimDecl VName) ()))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT
  (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
  (WriterT [TypeParamBase VName] m)
  ()
-> (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
-> WriterT
     [TypeParamBase VName]
     m
     (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT StateT
  (Map VName (TypeBase (DimDecl VName) ()), Map Int VName)
  (WriterT [TypeParamBase VName] m)
  ()
m (Map VName (TypeBase (DimDecl VName) ())
forall a. Monoid a => a
mempty, Map Int VName
forall a. Monoid a => a
mempty)
  where
    sub :: TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub t1 :: TypeBase dim as
t1@Array {} t2 :: TypeBase MonoSize d
t2@Array {}
      | Just TypeBase dim as
t1' <- Int -> TypeBase dim as -> Maybe (TypeBase dim as)
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim as -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim as
t1) TypeBase dim as
t1,
        Just TypeBase MonoSize d
t2' <- Int -> TypeBase MonoSize d -> Maybe (TypeBase MonoSize d)
forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (TypeBase dim as -> Int
forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim as
t1) TypeBase MonoSize d
t2 =
        TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase dim as
t1' TypeBase MonoSize d
t2'
    sub (Scalar (TypeVar as
_ Uniqueness
_ TypeName
v [TypeArg dim]
_)) TypeBase MonoSize d
t = TypeName -> TypeBase MonoSize d -> t (t m) ()
forall (t :: * -> * -> *) d (t :: (* -> *) -> * -> *)
       (t :: (* -> *) -> * -> *) (m :: * -> *).
(Bitraversable t,
 MonadState
   (Map VName (t (DimDecl VName) d), Map Int VName) (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
TypeName -> t MonoSize d -> t (t m) ()
addSubst TypeName
v TypeBase MonoSize d
t
    sub (Scalar (Record Map Name (TypeBase dim as)
fields1)) (Scalar (Record Map Name (TypeBase MonoSize d)
fields2)) =
      (TypeBase dim as -> TypeBase MonoSize d -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase MonoSize d] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
        TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub
        (((Name, TypeBase dim as) -> TypeBase dim as)
-> [(Name, TypeBase dim as)] -> [TypeBase dim as]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase dim as) -> TypeBase dim as
forall a b. (a, b) -> b
snd ([(Name, TypeBase dim as)] -> [TypeBase dim as])
-> [(Name, TypeBase dim as)] -> [TypeBase dim as]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> [(Name, TypeBase dim as)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase dim as)
fields1)
        (((Name, TypeBase MonoSize d) -> TypeBase MonoSize d)
-> [(Name, TypeBase MonoSize d)] -> [TypeBase MonoSize d]
forall a b. (a -> b) -> [a] -> [b]
map (Name, TypeBase MonoSize d) -> TypeBase MonoSize d
forall a b. (a, b) -> b
snd ([(Name, TypeBase MonoSize d)] -> [TypeBase MonoSize d])
-> [(Name, TypeBase MonoSize d)] -> [TypeBase MonoSize d]
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase MonoSize d) -> [(Name, TypeBase MonoSize d)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase MonoSize d)
fields2)
    sub (Scalar Prim {}) (Scalar Prim {}) = () -> t (t m) ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    sub (Scalar (Arrow as
_ PName
_ TypeBase dim as
t1a TypeBase dim as
t1b)) (Scalar (Arrow d
_ PName
_ TypeBase MonoSize d
t2a TypeBase MonoSize d
t2b)) = do
      TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase dim as
t1a TypeBase MonoSize d
t2a
      TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase dim as
t1b TypeBase MonoSize d
t2b
    sub (Scalar (Sum Map Name [TypeBase dim as]
cs1)) (Scalar (Sum Map Name [TypeBase MonoSize d]
cs2)) =
      ((Name, [TypeBase dim as])
 -> (Name, [TypeBase MonoSize d]) -> t (t m) [()])
-> [(Name, [TypeBase dim as])]
-> [(Name, [TypeBase MonoSize d])]
-> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Name, [TypeBase dim as])
-> (Name, [TypeBase MonoSize d]) -> t (t m) [()]
forall a a.
(a, [TypeBase dim as])
-> (a, [TypeBase MonoSize d]) -> t (t m) [()]
typeSubstClause (Map Name [TypeBase dim as] -> [(Name, [TypeBase dim as])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase dim as]
cs1) (Map Name [TypeBase MonoSize d] -> [(Name, [TypeBase MonoSize d])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase MonoSize d]
cs2)
      where
        typeSubstClause :: (a, [TypeBase dim as])
-> (a, [TypeBase MonoSize d]) -> t (t m) [()]
typeSubstClause (a
_, [TypeBase dim as]
ts1) (a
_, [TypeBase MonoSize d]
ts2) = (TypeBase dim as -> TypeBase MonoSize d -> t (t m) ())
-> [TypeBase dim as] -> [TypeBase MonoSize d] -> t (t m) [()]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub [TypeBase dim as]
ts1 [TypeBase MonoSize d]
ts2
    sub t1 :: TypeBase dim as
t1@(Scalar Sum {}) TypeBase MonoSize d
t2 = TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase dim as
t1 TypeBase MonoSize d
t2
    sub TypeBase dim as
t1 t2 :: TypeBase MonoSize d
t2@(Scalar Sum {}) = TypeBase dim as -> TypeBase MonoSize d -> t (t m) ()
sub TypeBase dim as
t1 TypeBase MonoSize d
t2
    sub TypeBase dim as
t1 TypeBase MonoSize d
t2 = String -> t (t m) ()
forall a. HasCallStack => String -> a
error (String -> t (t m) ()) -> String -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines [String
"typeSubstsM: mismatched types:", TypeBase dim as -> String
forall a. Pretty a => a -> String
pretty TypeBase dim as
t1, TypeBase MonoSize d -> String
forall a. Pretty a => a -> String
pretty TypeBase MonoSize d
t2]

    addSubst :: TypeName -> t MonoSize d -> t (t m) ()
addSubst (TypeName [VName]
_ VName
v) t MonoSize d
t = do
      (Map VName (t (DimDecl VName) d)
ts, Map Int VName
sizes) <- t (t m) (Map VName (t (DimDecl VName) d), 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 (t (DimDecl VName) d) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (t (DimDecl VName) d)
ts) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ do
        t (DimDecl VName) d
t' <- (MonoSize -> t (t m) (DimDecl VName))
-> (d -> t (t m) d)
-> t MonoSize d
-> t (t m) (t (DimDecl VName) d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse 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 d -> t (t m) d
forall (f :: * -> *) a. Applicative f => a -> f a
pure t MonoSize d
t
        (Map VName (t (DimDecl VName) d), Map Int VName) -> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VName
-> t (DimDecl VName) d
-> Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v t (DimDecl VName) d
t' Map VName (t (DimDecl VName) d)
ts, Map Int VName
sizes)

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

-- Perform a given substitution on the types in a pattern.
substPattern :: Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern :: Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
f Pattern
pat = case Pattern
pat of
  TuplePattern [Pattern]
pats SrcLoc
loc -> [Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern ((Pattern -> Pattern) -> [Pattern] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
f) [Pattern]
pats) SrcLoc
loc
  RecordPattern [(Name, Pattern)]
fs SrcLoc
loc -> [(Name, Pattern)] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern (((Name, Pattern) -> (Name, Pattern))
-> [(Name, Pattern)] -> [(Name, Pattern)]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Pattern) -> (Name, Pattern)
forall a. (a, Pattern) -> (a, Pattern)
substField [(Name, Pattern)]
fs) SrcLoc
loc
    where
      substField :: (a, Pattern) -> (a, Pattern)
substField (a
n, Pattern
p) = (a
n, Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
f Pattern
p)
  PatternParens Pattern
p SrcLoc
loc -> Pattern -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatternBase f vn -> SrcLoc -> PatternBase f vn
PatternParens (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
f Pattern
p) SrcLoc
loc
  Id VName
vn (Info PatternType
tp) SrcLoc
loc -> VName -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
vn (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
  Wildcard (Info PatternType
tp) SrcLoc
loc -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
  PatternAscription Pattern
p TypeDeclBase Info VName
td SrcLoc
loc
    | Bool
entry -> Pattern -> TypeDeclBase Info VName -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatternBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatternBase f vn
PatternAscription (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
False PatternType -> PatternType
f Pattern
p) TypeDeclBase Info VName
td SrcLoc
loc
    | Bool
otherwise -> Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
False PatternType -> PatternType
f Pattern
p
  PatternLit PatLit
e (Info PatternType
tp) SrcLoc
loc -> PatLit -> Info PatternType -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
PatLit -> f PatternType -> SrcLoc -> PatternBase f vn
PatternLit PatLit
e (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) SrcLoc
loc
  PatternConstr Name
n (Info PatternType
tp) [Pattern]
ps SrcLoc
loc -> Name -> Info PatternType -> [Pattern] -> SrcLoc -> Pattern
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
n (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
f PatternType
tp) [Pattern]
ps SrcLoc
loc

toPolyBinding :: ValBind -> PolyBinding
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding (ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
retdecl (Info (TypeBase (DimDecl VName) ()
rettype, [VName]
retext)) [TypeParamBase VName]
tparams [Pattern]
params Exp
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) =
  RecordReplacements
-> (VName, [TypeParamBase VName], [Pattern], Maybe (TypeExp VName),
    TypeBase (DimDecl VName) (), [VName], Exp, [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, TypeBase (DimDecl VName) ()
rettype, [VName]
retext, Exp
body, [AttrInfo]
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
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
_ Maybe (TypeExp VName)
_ (Info (TypeBase (DimDecl VName) ()
rettype, [VName]
retext)) [TypeParamBase VName]
_ [Pattern]
pats Exp
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
_) = do
  Map VName TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  let mapper :: ASTMapper MonoM
mapper =
        ASTMapper :: forall (m :: * -> *).
(Exp -> m Exp)
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
          { mapOnExp :: Exp -> MonoM Exp
mapOnExp = ASTMapper MonoM -> Exp -> MonoM Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper,
            mapOnName :: VName -> MonoM VName
mapOnName = VName -> MonoM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnQualName :: QualName VName -> MonoM (QualName VName)
mapOnQualName = QualName VName -> MonoM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnStructType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
mapOnStructType = TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeBase (DimDecl VName) ()
 -> MonoM (TypeBase (DimDecl VName) ()))
-> (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs,
            mapOnPatternType :: PatternType -> MonoM PatternType
mapOnPatternType = PatternType -> MonoM PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> MonoM PatternType)
-> (PatternType -> PatternType) -> PatternType -> MonoM PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName TypeSub -> PatternType -> PatternType
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs
          }

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

  ValBind -> MonoM ValBind
forall (m :: * -> *) a. Monad m => a -> m a
return
    ValBind
valbind
      { valBindRetType :: Info (TypeBase (DimDecl VName) (), [VName])
valBindRetType = (TypeBase (DimDecl VName) (), [VName])
-> Info (TypeBase (DimDecl VName) (), [VName])
forall a. a -> Info a
Info (Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs TypeBase (DimDecl VName) ()
rettype, [VName]
retext),
        valBindParams :: [Pattern]
valBindParams = (Pattern -> Pattern) -> [Pattern] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry ((PatternType -> PatternType) -> Pattern -> Pattern)
-> (PatternType -> PatternType) -> Pattern -> Pattern
forall a b. (a -> b) -> a -> b
$ Map VName TypeSub -> PatternType -> PatternType
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs) [Pattern]
pats,
        valBindBody :: Exp
valBindBody = Exp
body'
      }

removeTypeVariablesInType :: StructType -> MonoM StructType
removeTypeVariablesInType :: TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
removeTypeVariablesInType TypeBase (DimDecl VName) ()
t = do
  Map VName TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  TypeBase (DimDecl VName) () -> MonoM (TypeBase (DimDecl VName) ())
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeBase (DimDecl VName) ()
 -> MonoM (TypeBase (DimDecl VName) ()))
-> TypeBase (DimDecl VName) ()
-> MonoM (TypeBase (DimDecl VName) ())
forall a b. (a -> b) -> a -> b
$ Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs TypeBase (DimDecl VName) ()
t

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

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

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

transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeDeclBase Info VName
tydecl Maybe DocComment
_ SrcLoc
_) = do
  Map VName TypeSub
subs <- (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub))
-> (Env -> Map VName TypeSub) -> MonoM (Map VName TypeSub)
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> TypeSub)
-> Map VName TypeBinding -> Map VName TypeSub
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> TypeSub
TypeSub (Map VName TypeBinding -> Map VName TypeSub)
-> (Env -> Map VName TypeBinding) -> Env -> Map VName TypeSub
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  TypeBase (DimDecl VName) () -> MonoM ()
forall as. TypeBase (DimDecl VName) as -> MonoM ()
noticeDims (TypeBase (DimDecl VName) () -> MonoM ())
-> TypeBase (DimDecl VName) () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info (TypeBase (DimDecl VName) ())
forall (f :: * -> *) vn.
TypeDeclBase f vn -> f (TypeBase (DimDecl VName) ())
expandedType TypeDeclBase Info VName
tydecl
  let tp :: TypeBase (DimDecl VName) ()
tp = Map VName TypeSub
-> TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ()
forall als.
Monoid als =>
Map VName TypeSub
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
substituteTypes Map VName TypeSub
subs (TypeBase (DimDecl VName) () -> TypeBase (DimDecl VName) ())
-> (Info (TypeBase (DimDecl VName) ())
    -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ()
forall a. Info a -> a
unInfo (Info (TypeBase (DimDecl VName) ()) -> TypeBase (DimDecl VName) ())
-> Info (TypeBase (DimDecl VName) ())
-> TypeBase (DimDecl VName) ()
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info (TypeBase (DimDecl VName) ())
forall (f :: * -> *) vn.
TypeDeclBase f vn -> f (TypeBase (DimDecl VName) ())
expandedType TypeDeclBase Info VName
tydecl
      tbinding :: TypeBinding
tbinding = Liftedness
-> [TypeParamBase VName]
-> TypeBase (DimDecl VName) ()
-> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams TypeBase (DimDecl VName) ()
tp
  Env -> MonoM Env
forall (m :: * -> *) a. Monad m => a -> m a
return Env
forall a. Monoid a => a
mempty {envTypeBindings :: Map VName TypeBinding
envTypeBindings = VName -> TypeBinding -> Map VName TypeBinding
forall k a. k -> a -> Map k a
M.singleton VName
name TypeBinding
tbinding}

transformDecs :: [Dec] -> MonoM ()
transformDecs :: [Dec] -> MonoM ()
transformDecs [] = () -> MonoM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
transformDecs (ValDec ValBind
valbind : [Dec]
ds) = do
  Env
env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (TypeDec TypeBind
typebind : [Dec]
ds) = do
  Env
env <- TypeBind -> MonoM Env
transformTypeBind TypeBind
typebind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (Dec
dec : [Dec]
_) =
  String -> MonoM ()
forall a. HasCallStack => String -> a
error (String -> MonoM ()) -> String -> MonoM ()
forall a b. (a -> b) -> a -> b
$
    String
"The monomorphization module expects a module-free "
      String -> 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