{-# 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 :: forall dim als. 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 :: forall a. 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 :: forall a. 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 :: forall a. 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 :: forall a. 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
    ( (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
<$ :: forall a b. a -> MonoM b -> MonoM a
$c<$ :: forall a b. a -> MonoM b -> MonoM a
fmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
$cfmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
Functor,
      Functor MonoM
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
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
<* :: forall a b. MonoM a -> MonoM b -> MonoM a
$c<* :: forall a b. MonoM a -> MonoM b -> MonoM a
*> :: forall a b. MonoM a -> MonoM b -> MonoM b
$c*> :: forall a b. MonoM a -> MonoM b -> MonoM b
liftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
$cliftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
$c<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
pure :: forall a. a -> MonoM a
$cpure :: forall a. a -> MonoM a
Applicative,
      Applicative MonoM
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
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 :: forall a. a -> MonoM a
$creturn :: forall a. a -> MonoM a
>> :: forall a b. MonoM a -> MonoM b -> MonoM b
$c>> :: forall a b. MonoM a -> MonoM b -> MonoM b
>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
$c>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
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
MonadFreshNames
    )

runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource)
runMonoM :: forall a.
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
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 :: forall als. 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 (ExpBase Info VName)
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 = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> ExpBase Info VName
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
_) ->
        ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ VName
-> TypeBase (DimDecl VName) ()
-> [ExpBase Info VName]
-> ExpBase Info VName
forall {vn} {as}.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([ExpBase Info VName] -> ExpBase Info VName)
-> [ExpBase Info VName] -> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer TypeBase (DimDecl VName) ()
t'
      -- An intrinsic function.
      (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> ExpBase Info VName
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)
        ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ VName
-> TypeBase (DimDecl VName) ()
-> [ExpBase Info VName]
-> ExpBase Info VName
forall {vn} {as}.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' TypeBase (DimDecl VName) ()
t' ([ExpBase Info VName] -> ExpBase Info VName)
-> [ExpBase Info VName] -> ExpBase Info VName
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 :: forall dim. 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 :: forall (m :: * -> *).
MonadFreshNames m =>
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 :: ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp e :: ExpBase Info VName
e@Literal {} = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
transformExp e :: ExpBase Info VName
e@IntLit {} = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
transformExp e :: ExpBase Info VName
e@FloatLit {} = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
transformExp e :: ExpBase Info VName
e@StringLit {} = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
transformExp (Parens ExpBase Info VName
e SrcLoc
loc) =
  ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (QualParens (QualName VName, SrcLoc)
qn ExpBase Info VName
e SrcLoc
loc) =
  (QualName VName, SrcLoc)
-> ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (TupLit [ExpBase Info VName]
es SrcLoc
loc) =
  [ExpBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([ExpBase Info VName] -> SrcLoc -> ExpBase Info VName)
-> MonoM [ExpBase Info VName]
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> [ExpBase Info VName] -> MonoM [ExpBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp [ExpBase Info VName]
es MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> ExpBase Info VName)
-> MonoM [FieldBase Info VName]
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
  where
    transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (RecordFieldExplicit Name
name ExpBase Info VName
e SrcLoc
loc') =
      Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name (ExpBase Info VName -> SrcLoc -> FieldBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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 -> ExpBase Info VName -> 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 -> ExpBase Info VName
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 [ExpBase Info VName]
es Info PatternType
t SrcLoc
loc) =
  [ExpBase Info VName]
-> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
ArrayLit ([ExpBase Info VName]
 -> Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM [ExpBase Info VName]
-> MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> [ExpBase Info VName] -> MonoM [ExpBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp [ExpBase Info VName]
es MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Range ExpBase Info VName
e1 Maybe (ExpBase Info VName)
me Inclusiveness (ExpBase Info VName)
incl (Info PatternType, Info [VName])
tp SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  Maybe (ExpBase Info VName)
me' <- (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp Maybe (ExpBase Info VName)
me
  Inclusiveness (ExpBase Info VName)
incl' <- (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> Inclusiveness (ExpBase Info VName)
-> MonoM (Inclusiveness (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp Inclusiveness (ExpBase Info VName)
incl
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> Maybe (ExpBase Info VName)
-> Inclusiveness (ExpBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Range ExpBase Info VName
e1' Maybe (ExpBase Info VName)
me' Inclusiveness (ExpBase Info VName)
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 -> ExpBase Info VName
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> ExpBase Info VName)
-> MonoM [FieldBase Info VName]
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
    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 (ExpBase Info VName)
transformFName SrcLoc
loc QualName VName
fname (PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t')
transformExp (Ascript ExpBase Info VName
e TypeDeclBase Info VName
tp SrcLoc
loc) =
  ExpBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> TypeDeclBase f vn -> SrcLoc -> ExpBase f vn
Ascript (ExpBase Info VName
 -> TypeDeclBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (TypeDeclBase Info VName -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (TypeDeclBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (TypeDeclBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Coerce ExpBase Info VName
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
  ExpBase Info VName
-> TypeDeclBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeDeclBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Coerce (ExpBase Info VName
 -> TypeDeclBase Info VName
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     (TypeDeclBase Info VName
      -> (Info PatternType, Info [VName])
      -> SrcLoc
      -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM
  (TypeDeclBase Info VName
   -> (Info PatternType, Info [VName])
   -> SrcLoc
   -> ExpBase Info VName)
-> MonoM (TypeDeclBase Info VName)
-> MonoM
     ((Info PatternType, Info [VName]) -> SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM (Info PatternType, Info [VName])
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (LetPat Pattern
pat ExpBase Info VName
e1 ExpBase Info VName
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
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat Pattern
pat' (ExpBase Info VName
 -> ExpBase Info VName
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     (ExpBase Info VName
      -> (Info PatternType, Info [VName])
      -> SrcLoc
      -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
    MonoM
  (ExpBase Info VName
   -> (Info PatternType, Info [VName])
   -> SrcLoc
   -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     ((Info PatternType, Info [VName]) -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> RecordReplacements
-> MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName)
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2)
    MonoM
  ((Info PatternType, Info [VName]) -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info PatternType, Info [VName])
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [Pattern]
params, Maybe (TypeExp VName)
retdecl, Info TypeBase (DimDecl VName) ()
ret, ExpBase Info VName
body) ExpBase Info VName
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], ExpBase Info VName,
    [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
rr (VName
fname, [TypeParamBase VName]
tparams, [Pattern]
params, TypeBase (DimDecl VName) ()
ret, [], ExpBase Info VName
body, [AttrInfo]
forall a. Monoid a => a
mempty, SrcLoc
loc)
    MonoM
  (ExpBase Info VName, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM (ExpBase Info VName)
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (MonoM
   (ExpBase Info VName, Seq (VName, ValBind) -> Seq (VName, ValBind))
 -> MonoM (ExpBase Info VName))
-> MonoM
     (ExpBase Info VName, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ do
      (ExpBase Info VName
e', Seq (VName, ValBind)
bs) <- MonoM (ExpBase Info VName)
-> MonoM (ExpBase Info VName, Seq (VName, ValBind))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (MonoM (ExpBase Info VName)
 -> MonoM (ExpBase Info VName, Seq (VName, ValBind)))
-> MonoM (ExpBase Info VName)
-> MonoM (ExpBase Info VName, Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ VName
-> PolyBinding
-> MonoM (ExpBase Info VName)
-> MonoM (ExpBase Info VName)
forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
fname PolyBinding
funbind (MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName))
-> MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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
      (ExpBase Info VName, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM
     (ExpBase Info VName, Seq (VName, ValBind) -> Seq (VName, ValBind))
forall (m :: * -> *) a. Monad m => a -> m a
return ([ValBind] -> ExpBase Info VName -> ExpBase Info VName
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) ExpBase Info VName
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
    ExpBase Info VName
body' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
body
    VName
-> ([TypeParamBase VName], [Pattern], Maybe (TypeExp VName),
    Info (TypeBase (DimDecl VName) ()), ExpBase Info VName)
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
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, ExpBase Info VName
body')
      (ExpBase Info VName
 -> Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (If ExpBase Info VName
e1 ExpBase Info VName
e2 ExpBase Info VName
e3 (Info PatternType
tp, Info [VName]
retext) SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  ExpBase Info VName
e2' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
  ExpBase Info VName
e3' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
If ExpBase Info VName
e1' ExpBase Info VName
e2' ExpBase Info VName
e3' (Info PatternType
tp', Info [VName]
retext) SrcLoc
loc
transformExp (Apply ExpBase Info VName
e1 ExpBase Info VName
e2 Info (Diet, Maybe VName)
d (Info PatternType
ret, Info [VName]
ext) SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  ExpBase Info VName
e2' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply ExpBase Info VName
e1' ExpBase Info VName
e2' Info (Diet, Maybe VName)
d (Info PatternType
ret', Info [VName]
ext) SrcLoc
loc
transformExp (Negate ExpBase Info VName
e SrcLoc
loc) =
  ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Lambda [Pattern]
params ExpBase Info VName
e0 Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
tp SrcLoc
loc) = do
  ExpBase Info VName
e0' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e0
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase f vn
Lambda [Pattern]
params ExpBase Info VName
e0' Maybe (TypeExp VName)
decl Info (Aliasing, TypeBase (DimDecl VName) ())
tp SrcLoc
loc
transformExp (OpSection QualName VName
qn Info PatternType
t SrcLoc
loc) =
  ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
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) ExpBase Info VName
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
  ExpBase Info VName
fname' <- SrcLoc
-> QualName VName
-> TypeBase (DimDecl VName) ()
-> MonoM (ExpBase Info VName)
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName))
-> TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
  ExpBase Info VName
e' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e
  ExpBase Info VName
-> Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> PatternType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM (ExpBase Info VName)
desugarBinOpSection
    ExpBase Info VName
fname'
    (ExpBase Info VName -> Maybe (ExpBase Info VName)
forall a. a -> Maybe a
Just ExpBase Info VName
e')
    Maybe (ExpBase Info VName)
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) ExpBase Info VName
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
  ExpBase Info VName
fname' <- SrcLoc
-> QualName VName
-> TypeBase (DimDecl VName) ()
-> MonoM (ExpBase Info VName)
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName))
-> TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
  ExpBase Info VName
e' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e
  ExpBase Info VName
-> Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> PatternType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM (ExpBase Info VName)
desugarBinOpSection
    ExpBase Info VName
fname'
    Maybe (ExpBase Info VName)
forall a. Maybe a
Nothing
    (ExpBase Info VName -> Maybe (ExpBase Info VName)
forall a. a -> Maybe a
Just ExpBase Info VName
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 (ExpBase Info VName)
desugarProjectSection [Name]
fields PatternType
t SrcLoc
loc
transformExp (IndexSection [DimIndexBase Info VName]
idxs (Info PatternType
t) SrcLoc
loc) =
  [DimIndexBase Info VName]
-> PatternType -> SrcLoc -> MonoM (ExpBase Info VName)
desugarIndexSection [DimIndexBase Info VName]
idxs PatternType
t SrcLoc
loc
transformExp (DoLoop [VName]
sparams Pattern
pat ExpBase Info VName
e1 LoopFormBase Info VName
form ExpBase Info VName
e3 Info (PatternType, [VName])
ret SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  LoopFormBase Info VName
form' <- case LoopFormBase Info VName
form of
    For IdentBase Info VName
ident ExpBase Info VName
e2 -> IdentBase Info VName
-> ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
ident (ExpBase Info VName -> LoopFormBase Info VName)
-> MonoM (ExpBase Info VName) -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
    ForIn Pattern
pat2 ExpBase Info VName
e2 -> Pattern -> ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn Pattern
pat2 (ExpBase Info VName -> LoopFormBase Info VName)
-> MonoM (ExpBase Info VName) -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
    While ExpBase Info VName
e2 -> ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While (ExpBase Info VName -> LoopFormBase Info VName)
-> MonoM (ExpBase Info VName) -> MonoM (LoopFormBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
  ExpBase Info VName
e3' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [VName]
-> Pattern
-> ExpBase Info VName
-> LoopFormBase Info VName
-> ExpBase Info VName
-> Info (PatternType, [VName])
-> SrcLoc
-> ExpBase Info VName
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' ExpBase Info VName
e1' LoopFormBase Info VName
form' ExpBase Info VName
e3' Info (PatternType, [VName])
ret SrcLoc
loc
transformExp (BinOp (QualName VName
fname, SrcLoc
_) (Info PatternType
t) (ExpBase Info VName
e1, Info (TypeBase (DimDecl VName) (), Maybe VName)
d1) (ExpBase Info VName
e2, Info (TypeBase (DimDecl VName) (), Maybe VName)
d2) Info PatternType
tp Info [VName]
ext SrcLoc
loc) = do
  ExpBase Info VName
fname' <- SrcLoc
-> QualName VName
-> TypeBase (DimDecl VName) ()
-> MonoM (ExpBase Info VName)
transformFName SrcLoc
loc QualName VName
fname (TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName))
-> TypeBase (DimDecl VName) () -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ PatternType -> TypeBase (DimDecl VName) ()
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  ExpBase Info VName
e2' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
  if PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e1') Bool -> Bool -> Bool
&& PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e2')
    then ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName
-> ExpBase Info VName -> ExpBase Info VName -> ExpBase Info VName
forall {vn}.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp ExpBase Info VName
fname' ExpBase Info VName
e1' ExpBase Info VName
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.
      (ExpBase Info VName
x_param_e, Pattern
x_param) <- ExpBase Info VName -> MonoM (ExpBase Info VName, Pattern)
forall {m :: * -> *}.
MonadFreshNames m =>
ExpBase Info VName -> m (ExpBase Info VName, Pattern)
makeVarParam ExpBase Info VName
e1'
      (ExpBase Info VName
y_param_e, Pattern
y_param) <- ExpBase Info VName -> MonoM (ExpBase Info VName, Pattern)
forall {m :: * -> *}.
MonadFreshNames m =>
ExpBase Info VName -> m (ExpBase Info VName, Pattern)
makeVarParam ExpBase Info VName
e2'
      ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$
        Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat
          Pattern
x_param
          ExpBase Info VName
e1'
          ( Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat
              Pattern
y_param
              ExpBase Info VName
e2'
              (ExpBase Info VName
-> ExpBase Info VName -> ExpBase Info VName -> ExpBase Info VName
forall {vn}.
ExpBase Info vn
-> ExpBase Info vn -> ExpBase Info vn -> ExpBase Info vn
applyOp ExpBase Info VName
fname' ExpBase Info VName
x_param_e ExpBase Info VName
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 :: ExpBase Info VName -> m (ExpBase Info VName, Pattern)
makeVarParam ExpBase Info VName
arg = do
      let argtype :: PatternType
argtype = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
arg
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"binop_p"
      (ExpBase Info VName, Pattern) -> m (ExpBase Info VName, Pattern)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
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 ExpBase Info VName
e Info PatternType
tp SrcLoc
loc) = do
  Maybe RecordReplacement
maybe_fs <- case ExpBase Info VName
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)
    ExpBase Info VName
_ -> 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 ->
        ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
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
      ExpBase Info VName
e' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e
      ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ Name
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
n ExpBase Info VName
e' Info PatternType
tp SrcLoc
loc
transformExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs ExpBase Info VName
e1 ExpBase Info VName
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
  ExpBase Info VName
e1' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1
  ExpBase Info VName
body' <- ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
body
  PatternType
t' <- PatternType -> MonoM PatternType
forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType PatternType
t
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName
-> IdentBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
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' ExpBase Info VName
e1' ExpBase Info VName
body' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t') SrcLoc
loc
transformExp (Index ExpBase Info VName
e0 [DimIndexBase Info VName]
idxs (Info PatternType, Info [VName])
info SrcLoc
loc) =
  ExpBase Info VName
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (ExpBase Info VName
 -> [DimIndexBase Info VName]
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     ([DimIndexBase Info VName]
      -> (Info PatternType, Info [VName])
      -> SrcLoc
      -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e0 MonoM
  ([DimIndexBase Info VName]
   -> (Info PatternType, Info [VName])
   -> SrcLoc
   -> ExpBase Info VName)
-> MonoM [DimIndexBase Info VName]
-> MonoM
     ((Info PatternType, Info [VName]) -> SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM (Info PatternType, Info [VName])
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Update ExpBase Info VName
e1 [DimIndexBase Info VName]
idxs ExpBase Info VName
e2 SrcLoc
loc) =
  ExpBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update (ExpBase Info VName
 -> [DimIndexBase Info VName]
 -> ExpBase Info VName
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     ([DimIndexBase Info VName]
      -> ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1 MonoM
  ([DimIndexBase Info VName]
   -> ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM [DimIndexBase Info VName]
-> MonoM (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
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 (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
    MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (RecordUpdate ExpBase Info VName
e1 [Name]
fs ExpBase Info VName
e2 Info PatternType
t SrcLoc
loc) =
  ExpBase Info VName
-> [Name]
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name]
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
RecordUpdate (ExpBase Info VName
 -> [Name]
 -> ExpBase Info VName
 -> Info PatternType
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     ([Name]
      -> ExpBase Info VName
      -> Info PatternType
      -> SrcLoc
      -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1 MonoM
  ([Name]
   -> ExpBase Info VName
   -> Info PatternType
   -> SrcLoc
   -> ExpBase Info VName)
-> MonoM [Name]
-> MonoM
     (ExpBase Info VName
      -> Info PatternType -> SrcLoc -> ExpBase Info VName)
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
  (ExpBase Info VName
   -> Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2
    MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Assert ExpBase Info VName
e1 ExpBase Info VName
e2 Info String
desc SrcLoc
loc) =
  ExpBase Info VName
-> ExpBase Info VName
-> Info String
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f String -> SrcLoc -> ExpBase f vn
Assert (ExpBase Info VName
 -> ExpBase Info VName
 -> Info String
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     (ExpBase Info VName -> Info String -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e1 MonoM
  (ExpBase Info VName -> Info String -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (Info String -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e2 MonoM (Info String -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info String) -> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Constr Name
name [ExpBase Info VName]
all_es Info PatternType
t SrcLoc
loc) =
  Name
-> [ExpBase Info VName]
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
Constr Name
name ([ExpBase Info VName]
 -> Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM [ExpBase Info VName]
-> MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> [ExpBase Info VName] -> MonoM [ExpBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp [ExpBase Info VName]
all_es MonoM (Info PatternType -> SrcLoc -> ExpBase Info VName)
-> MonoM (Info PatternType) -> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) =
  ExpBase Info VName
-> NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Match (ExpBase Info VName
 -> NonEmpty (CaseBase Info VName)
 -> (Info PatternType, Info [VName])
 -> SrcLoc
 -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM
     (NonEmpty (CaseBase Info VName)
      -> (Info PatternType, Info [VName])
      -> SrcLoc
      -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM
  (NonEmpty (CaseBase Info VName)
   -> (Info PatternType, Info [VName])
   -> SrcLoc
   -> ExpBase Info VName)
-> MonoM (NonEmpty (CaseBase Info VName))
-> MonoM
     ((Info PatternType, Info [VName]) -> SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM (Info PatternType, Info [VName])
-> MonoM (SrcLoc -> ExpBase Info VName)
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 -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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
transformExp (Attr AttrInfo
info ExpBase Info VName
e SrcLoc
loc) =
  AttrInfo -> ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
AttrInfo -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo
info (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e MonoM (SrcLoc -> ExpBase Info VName)
-> MonoM SrcLoc -> MonoM (ExpBase 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

transformCase :: Case -> MonoM Case
transformCase :: CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase (CasePat Pattern
p ExpBase Info VName
e SrcLoc
loc) = do
  (Pattern
p', RecordReplacements
rr) <- Pattern -> MonoM (Pattern, RecordReplacements)
transformPattern Pattern
p
  Pattern -> ExpBase Info VName -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat Pattern
p' (ExpBase Info VName -> SrcLoc -> CaseBase Info VName)
-> MonoM (ExpBase Info VName)
-> MonoM (SrcLoc -> CaseBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RecordReplacements
-> MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName)
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr (ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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 ExpBase Info VName
e) = ExpBase Info VName -> DimIndexBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (ExpBase Info VName -> DimIndexBase Info VName)
-> MonoM (ExpBase Info VName) -> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
e
transformDimIndex (DimSlice Maybe (ExpBase Info VName)
me1 Maybe (ExpBase Info VName)
me2 Maybe (ExpBase Info VName)
me3) =
  Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> DimIndexBase Info VName
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe (ExpBase Info VName)
 -> Maybe (ExpBase Info VName)
 -> Maybe (ExpBase Info VName)
 -> DimIndexBase Info VName)
-> MonoM (Maybe (ExpBase Info VName))
-> MonoM
     (Maybe (ExpBase Info VName)
      -> Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
trans Maybe (ExpBase Info VName)
me1 MonoM
  (Maybe (ExpBase Info VName)
   -> Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
-> MonoM (Maybe (ExpBase Info VName))
-> MonoM (Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
trans Maybe (ExpBase Info VName)
me2 MonoM (Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
-> MonoM (Maybe (ExpBase Info VName))
-> MonoM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
trans Maybe (ExpBase Info VName)
me3
  where
    trans :: Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
trans = (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> Maybe (ExpBase Info VName) -> MonoM (Maybe (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> MonoM (ExpBase Info VName)
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 :: ExpBase Info VName
-> Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> PatternType
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PName, TypeBase (DimDecl VName) (), Maybe VName)
-> (PatternType, [VName])
-> SrcLoc
-> MonoM (ExpBase Info VName)
desugarBinOpSection ExpBase Info VName
op Maybe (ExpBase Info VName)
e_left Maybe (ExpBase Info VName)
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, ExpBase Info VName -> ExpBase Info VName
wrap_left, ExpBase Info VName
e1, [Pattern]
p1) <- Maybe (ExpBase Info VName)
-> PatternType
-> MonoM
     (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
forall {m :: * -> *}.
MonadFreshNames m =>
Maybe (ExpBase Info VName)
-> PatternType
-> m (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
makeVarParam Maybe (ExpBase Info VName)
e_left (PatternType
 -> MonoM
      (VName, ExpBase Info VName -> ExpBase Info VName,
       ExpBase Info VName, [Pattern]))
-> PatternType
-> MonoM
     (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [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, ExpBase Info VName -> ExpBase Info VName
wrap_right, ExpBase Info VName
e2, [Pattern]
p2) <- Maybe (ExpBase Info VName)
-> PatternType
-> MonoM
     (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
forall {m :: * -> *}.
MonadFreshNames m =>
Maybe (ExpBase Info VName)
-> PatternType
-> m (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
makeVarParam Maybe (ExpBase Info VName)
e_right (PatternType
 -> MonoM
      (VName, ExpBase Info VName -> ExpBase Info VName,
       ExpBase Info VName, [Pattern]))
-> PatternType
-> MonoM
     (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [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 :: ExpBase Info VName
apply_left =
        ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
          ExpBase Info VName
op
          ExpBase Info VName
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 :: ExpBase Info VName
body =
        ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
          ExpBase Info VName
apply_left
          ExpBase Info VName
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'
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> ExpBase Info VName
wrap_left (ExpBase Info VName -> ExpBase Info VName)
-> ExpBase Info VName -> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> ExpBase Info VName
wrap_right (ExpBase Info VName -> ExpBase Info VName)
-> ExpBase Info VName -> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase Info VName
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) ExpBase Info VName
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, ExpBase Info VName)
patAndVar PatternType
argtype = do
      VName
x <- String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
      (VName, Pattern, ExpBase Info VName)
-> m (VName, Pattern, ExpBase Info VName)
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 -> ExpBase Info VName
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 (ExpBase Info VName)
-> PatternType
-> m (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
makeVarParam (Just ExpBase Info VName
e) PatternType
argtype = do
      (VName
v, Pattern
pat, ExpBase Info VName
var_e) <- PatternType -> m (VName, Pattern, ExpBase Info VName)
forall {m :: * -> *}.
MonadFreshNames m =>
PatternType -> m (VName, Pattern, ExpBase Info VName)
patAndVar PatternType
argtype
      let wrap :: ExpBase Info VName -> ExpBase Info VName
wrap ExpBase Info VName
body =
            Pattern
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat Pattern
pat ExpBase Info VName
e ExpBase Info VName
body (PatternType -> Info PatternType
forall a. a -> Info a
Info (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
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, ExpBase Info VName -> ExpBase Info VName,
 ExpBase Info VName, [Pattern])
-> m (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, ExpBase Info VName -> ExpBase Info VName
wrap, ExpBase Info VName
var_e, [])
    makeVarParam Maybe (ExpBase Info VName)
Nothing PatternType
argtype = do
      (VName
v, Pattern
pat, ExpBase Info VName
var_e) <- PatternType -> m (VName, Pattern, ExpBase Info VName)
forall {m :: * -> *}.
MonadFreshNames m =>
PatternType -> m (VName, Pattern, ExpBase Info VName)
patAndVar PatternType
argtype
      (VName, ExpBase Info VName -> ExpBase Info VName,
 ExpBase Info VName, [Pattern])
-> m (VName, ExpBase Info VName -> ExpBase Info VName,
      ExpBase Info VName, [Pattern])
forall (m :: * -> *) a. Monad m => a -> m a
return (VName
v, ExpBase Info VName -> ExpBase Info VName
forall a. a -> a
id, ExpBase Info VName
var_e, [Pattern
pat])

desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp
desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM (ExpBase Info VName)
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 :: ExpBase Info VName
body = (ExpBase Info VName -> Name -> ExpBase Info VName)
-> ExpBase Info VName -> [Name] -> ExpBase Info VName
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ExpBase Info VName -> Name -> ExpBase Info VName
project (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
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
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase Info VName
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] ExpBase Info VName
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 :: ExpBase Info VName -> Name -> ExpBase Info VName
project ExpBase Info VName
e Name
field =
      case ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
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
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
field ExpBase Info VName
e (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
forall a. Monoid a => a
mempty
        PatternType
t ->
          String -> ExpBase Info VName
forall a. HasCallStack => String -> a
error (String -> ExpBase Info VName) -> String -> ExpBase Info VName
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 (ExpBase Info VName)
forall a. HasCallStack => String -> a
error (String -> MonoM (ExpBase Info VName))
-> String -> MonoM (ExpBase Info VName)
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 (ExpBase Info VName)
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 :: ExpBase Info VName
body = ExpBase Info VName
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
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
  ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [Pattern]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Aliasing, TypeBase (DimDecl VName) ())
-> SrcLoc
-> ExpBase Info VName
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] ExpBase Info VName
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 (ExpBase Info VName)
forall a. HasCallStack => String -> a
error (String -> MonoM (ExpBase Info VName))
-> String -> MonoM (ExpBase Info VName)
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 :: forall as. 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 (ExpBase Info VName) -> MonoM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MonoM (ExpBase Info VName) -> MonoM ())
-> MonoM (ExpBase Info VName) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ SrcLoc
-> QualName VName
-> TypeBase (DimDecl VName) ()
-> MonoM (ExpBase Info VName)
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] -> ExpBase Info VName -> ExpBase Info VName
unfoldLetFuns [] ExpBase Info VName
e = ExpBase Info VName
e
unfoldLetFuns (ValBind Maybe (Info EntryPoint)
_ VName
fname Maybe (TypeExp VName)
_ (Info (TypeBase (DimDecl VName) ()
rettype, [VName]
_)) [TypeParamBase VName]
dim_params [Pattern]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
loc : [ValBind]
rest) ExpBase Info VName
e =
  VName
-> ([TypeParamBase VName], [Pattern], Maybe (TypeExp VName),
    Info (TypeBase (DimDecl VName) ()), ExpBase Info VName)
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
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, ExpBase Info VName
body) ExpBase Info VName
e' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
e_t) SrcLoc
loc
  where
    e' :: ExpBase Info VName
e' = [ValBind] -> ExpBase Info VName -> ExpBase Info VName
unfoldLetFuns [ValBind]
rest ExpBase Info VName
e
    e_t :: PatternType
e_t = ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
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 :: forall a.
Monoid a =>
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 (ExpBase Info VName))
-> [TypeParamBase VName] -> [ExpBase Info VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DimInst -> TypeParamBase VName -> Maybe (ExpBase Info VName)
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, ExpBase Info VName
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''

    ExpBase Info VName
body' <- (VName -> Maybe (Subst (TypeBase (DimDecl VName) ())))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
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') ExpBase Info VName
body
    ExpBase Info VName
body'' <- RecordReplacements
-> MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName)
forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements ([RecordReplacements] -> RecordReplacements
forall a. Monoid a => [a] -> a
mconcat [RecordReplacements]
rrs) (MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName))
-> MonoM (ExpBase Info VName) -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> MonoM (ExpBase Info VName)
transformExp ExpBase Info VName
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])
-> ExpBase Info VName
-> 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)
              ExpBase Info VName
body''
          else
            VName
-> [TypeParamBase VName]
-> [Pattern]
-> (TypeBase (DimDecl VName) (), [VName])
-> ExpBase Info VName
-> 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)
              ExpBase Info VName
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 :: * -> *).
(ExpBase Info VName -> m (ExpBase Info VName))
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
        { mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName)
mapOnExp = ASTMapper m -> ExpBase Info VName -> m (ExpBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (ASTMapper m -> ExpBase Info VName -> m (ExpBase Info VName))
-> ASTMapper m -> ExpBase Info VName -> m (ExpBase Info VName)
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 :: forall (m :: * -> *).
MonadFreshNames m =>
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 ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc) =
  RecordReplacements
-> (VName, [TypeParamBase VName], [Pattern],
    TypeBase (DimDecl VName) (), [VName], ExpBase Info VName,
    [AttrInfo], SrcLoc)
-> PolyBinding
PolyBinding RecordReplacements
forall a. Monoid a => a
mempty (VName
name, [TypeParamBase VName]
tparams, [Pattern]
params, TypeBase (DimDecl VName) ()
rettype, [VName]
retext, ExpBase Info VName
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 ExpBase Info VName
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 :: * -> *).
(ExpBase Info VName -> m (ExpBase Info VName))
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (TypeBase (DimDecl VName) () -> m (TypeBase (DimDecl VName) ()))
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
          { mapOnExp :: ExpBase Info VName -> MonoM (ExpBase Info VName)
mapOnExp = ASTMapper MonoM -> ExpBase Info VName -> MonoM (ExpBase Info VName)
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
          }

  ExpBase Info VName
body' <- ASTMapper MonoM -> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper ExpBase Info VName
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 :: ExpBase Info VName
valBindBody = ExpBase Info VName
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 :: forall (m :: * -> *). MonadFreshNames m => [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