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

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

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

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

-- The monomorphization monad reads 'PolyBinding's and writes
-- 'ValBind's.  The 'TypeParam's in the 'ValBind's can only be size
-- parameters.
--
-- Each 'Polybinding' is also connected with the 'RecordReplacements'
-- that were active when the binding was defined.  This is used only
-- in local functions.
data PolyBinding
  = PolyBinding
      RecordReplacements
      ( VName,
        [TypeParam],
        [Pattern],
        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],
    TypeBase (DimDecl VName) (), [VName], Exp, [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [Pattern]
params, 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 (PName, TypeBase (DimDecl VName) (), Maybe VName),
 Info (PName, TypeBase (DimDecl VName) ()))
arg (Info PatternType, Info [VName])
ret SrcLoc
loc) = do
  let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype)) = (Info (PName, TypeBase (DimDecl VName) (), Maybe VName),
 Info (PName, TypeBase (DimDecl VName) ()))
arg
      (Info PatternType
rettype, Info [VName]
retext) = (Info PatternType, Info [VName])
ret
  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
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, 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
    (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xargext)
    (PName
yp, 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 (PName, TypeBase (DimDecl VName) ()),
 Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg (Info PatternType
rettype) SrcLoc
loc) = do
  let (Info (PName
xp, TypeBase (DimDecl VName) ()
xtype), Info (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yargext)) = (Info (PName, TypeBase (DimDecl VName) ()),
 Info (PName, TypeBase (DimDecl VName) (), Maybe VName))
arg
  Exp
fname' <- SrcLoc
-> QualName VName -> TypeBase (DimDecl VName) () -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM Exp)
-> TypeBase (DimDecl VName) () -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ 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
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, 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
    (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
forall a. Maybe a
Nothing)
    (PName
yp, 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
_) (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
  if PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf Exp
e1') Bool -> Bool -> Bool
&& PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> PatternType
typeOf Exp
e2')
    then Exp -> MonoM Exp
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp -> Exp
forall vn.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp Exp
fname' Exp
e1' Exp
e2'
    else do
      -- 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 ->
  (PName, StructType, Maybe VName) ->
  (PName, StructType, Maybe VName) ->
  (PatternType, [VName]) ->
  SrcLoc ->
  MonoM Exp
desugarBinOpSection :: Exp
-> Maybe Exp
-> Maybe Exp
-> PatternType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection Exp
op Maybe Exp
e_left Maybe Exp
e_right PatternType
t (PName
xp, TypeBase (DimDecl VName) ()
xtype, Maybe VName
xext) (PName
yp, TypeBase (DimDecl VName) ()
ytype, Maybe VName
yext) (PatternType
rettype, [VName]
retext) SrcLoc
loc = do
  (VName
v1, Exp -> Exp
wrap_left, Exp
e1, [Pattern]
p1) <- Maybe Exp
-> PatternType -> MonoM (VName, Exp -> Exp, Exp, [Pattern])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (VName, Exp -> Exp, Exp, [Pattern])
makeVarParam Maybe Exp
e_left (PatternType -> MonoM (VName, Exp -> Exp, Exp, [Pattern]))
-> PatternType -> MonoM (VName, Exp -> Exp, 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
  (VName
v2, Exp -> Exp
wrap_right, Exp
e2, [Pattern]
p2) <- Maybe Exp
-> PatternType -> MonoM (VName, Exp -> Exp, Exp, [Pattern])
forall (m :: * -> *).
MonadFreshNames m =>
Maybe Exp -> PatternType -> m (VName, Exp -> Exp, Exp, [Pattern])
makeVarParam Maybe Exp
e_right (PatternType -> MonoM (VName, Exp -> Exp, Exp, [Pattern]))
-> PatternType -> MonoM (VName, Exp -> Exp, 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
$ ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) Aliasing -> PatternType)
-> ScalarTypeBase (DimDecl VName) Aliasing -> PatternType
forall a b. (a -> b) -> a -> b
$ Aliasing
-> PName
-> PatternType
-> PatternType
-> ScalarTypeBase (DimDecl VName) Aliasing
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow Aliasing
forall a. Monoid a => a
mempty PName
yp (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
      rettype' :: PatternType
rettype' =
        let onDim :: DimDecl VName -> DimDecl VName
onDim (NamedDim QualName VName
d)
              | Named VName
p <- PName
xp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v1
              | Named VName
p <- PName
yp, QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
p = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v2
            onDim DimDecl VName
d = DimDecl VName
d
         in (DimDecl VName -> DimDecl VName) -> PatternType -> PatternType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
onDim PatternType
rettype
      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
$ Exp -> Exp
wrap_left (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp
wrap_right (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ [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
    patAndVar :: PatternType -> m (VName, Pattern, Exp)
patAndVar PatternType
argtype = do
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
      (VName, Pattern, Exp) -> m (VName, Pattern, Exp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( VName
x,
          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
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          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
        )

    makeVarParam :: Maybe Exp -> PatternType -> m (VName, Exp -> Exp, Exp, [Pattern])
makeVarParam (Just Exp
e) PatternType
argtype = do
      (VName
v, Pattern
pat, Exp
var_e) <- PatternType -> m (VName, Pattern, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m (VName, Pattern, Exp)
patAndVar PatternType
argtype
      let wrap :: Exp -> Exp
wrap Exp
body =
            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
e Exp
body (PatternType -> Info PatternType
forall a. a -> Info a
Info (Exp -> PatternType
typeOf Exp
body), [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
forall a. Monoid a => a
mempty) SrcLoc
forall a. Monoid a => a
mempty
      (VName, Exp -> Exp, Exp, [Pattern])
-> m (VName, Exp -> Exp, Exp, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, Exp -> Exp
wrap, Exp
var_e, [])
    makeVarParam Maybe Exp
Nothing PatternType
argtype = do
      (VName
v, Pattern
pat, Exp
var_e) <- PatternType -> m (VName, Pattern, Exp)
forall (m :: * -> *).
MonadFreshNames m =>
PatternType -> m (VName, Pattern, Exp)
patAndVar PatternType
argtype
      (VName, Exp -> Exp, Exp, [Pattern])
-> m (VName, Exp -> Exp, Exp, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, Exp -> Exp
forall a. a -> a
id, Exp
var_e, [Pattern
pat])

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)
_ ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
0) SrcLoc
forall a. Monoid a => a
mempty

-- Monomorphising higher-order functions can result in function types
-- where the same named parameter occurs in multiple spots.  When
-- monomorphising we don't really need those parameter names anymore,
-- and the defunctionaliser can be confused if there are duplicates
-- (it doesn't handle shadowing), so let's just remove all parameter
-- names here.  This is safe because a MonoType does not contain sizes
-- anyway.
noNamedParams :: MonoType -> MonoType
noNamedParams :: MonoType -> MonoType
noNamedParams = MonoType -> MonoType
forall dim. TypeBase dim () -> TypeBase dim ()
f
  where
    f :: TypeBase dim () -> TypeBase dim ()
f (Array () Uniqueness
u 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, 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) () -> Set VName
mustBeExplicit TypeBase (DimDecl VName) ()
bind_t') (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) ([TypeParamBase VName]
 -> ([TypeParamBase VName], [TypeParamBase VName]))
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$
            [TypeParamBase VName]
shape_params [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
t_shape_params

    ([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
forall vn.
vn
-> [TypeParamBase vn]
-> [PatternBase Info vn]
-> (TypeBase (DimDecl VName) (), [VName])
-> ExpBase Info vn
-> ValBindBase Info vn
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
forall vn.
vn
-> [TypeParamBase vn]
-> [PatternBase Info vn]
-> (TypeBase (DimDecl VName) (), [VName])
-> ExpBase Info vn
-> ValBindBase Info vn
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 :: vn
-> [TypeParamBase vn]
-> [PatternBase Info vn]
-> (TypeBase (DimDecl VName) (), [VName])
-> ExpBase Info vn
-> ValBindBase Info vn
toValBinding vn
name' [TypeParamBase vn]
tparams' [PatternBase Info vn]
params'' (TypeBase (DimDecl VName) (), [VName])
rettype' ExpBase Info vn
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 :: vn
valBindName = vn
name',
          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',
          valBindRetDecl :: Maybe (TypeExp vn)
valBindRetDecl = Maybe (TypeExp vn)
forall a. Maybe a
Nothing,
          valBindTypeParams :: [TypeParamBase vn]
valBindTypeParams = [TypeParamBase vn]
tparams',
          valBindParams :: [PatternBase Info vn]
valBindParams = [PatternBase Info vn]
params'',
          valBindBody :: ExpBase Info vn
valBindBody = ExpBase Info vn
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)
_ (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],
    TypeBase (DimDecl VName) (), [VName], Exp, [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [Pattern]
params, 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