-- | 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 Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence qualified as Seq
import Data.Set qualified 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 = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ forall dim as. PrimType -> ScalarTypeBase dim as
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

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

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

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

-- Monomorphization environment mapping names of polymorphic functions
-- to a representation of their corresponding function bindings.
data Env = Env
  { Env -> Map VName PolyBinding
envPolyBindings :: M.Map VName PolyBinding,
    Env -> Map VName TypeBinding
envTypeBindings :: M.Map VName TypeBinding,
    Env -> RecordReplacements
envRecordReplacements :: RecordReplacements
  }

instance Semigroup Env where
  Env Map VName PolyBinding
tb1 Map VName TypeBinding
pb1 RecordReplacements
rr1 <> :: Env -> Env -> Env
<> Env Map VName PolyBinding
tb2 Map VName TypeBinding
pb2 RecordReplacements
rr2 = Map VName PolyBinding
-> Map VName TypeBinding -> RecordReplacements -> Env
Env (Map VName PolyBinding
tb1 forall a. Semigroup a => a -> a -> a
<> Map VName PolyBinding
tb2) (Map VName TypeBinding
pb1 forall a. Semigroup a => a -> a -> a
<> Map VName TypeBinding
pb2) (RecordReplacements
rr1 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 forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty

localEnv :: Env -> MonoM a -> MonoM a
localEnv :: forall a. Env -> MonoM a -> MonoM a
localEnv Env
env = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Env
env <>)

extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv :: forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
vn PolyBinding
binding =
  forall a. Env -> MonoM a -> MonoM a
localEnv
    forall a. Monoid a => a
mempty {envPolyBindings :: Map VName PolyBinding
envPolyBindings = 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 = forall a. Env -> MonoM a -> MonoM a
localEnv 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 = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local 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 -> 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
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
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
MonoM VNameSource
VNameSource -> MonoM ()
forall (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) = forall s a. State s a -> s -> a
evalState (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 forall a. Monoid a => a
mempty VNameSource
src) forall a. Monoid a => a
mempty

lookupFun :: VName -> MonoM (Maybe PolyBinding)
lookupFun :: VName -> MonoM (Maybe PolyBinding)
lookupFun VName
vn = do
  Map VName PolyBinding
env <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Map VName PolyBinding
envPolyBindings
  case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
vn Map VName PolyBinding
env of
    Just PolyBinding
valbind -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just PolyBinding
valbind
    Maybe PolyBinding
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

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

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

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

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

instance Pretty MonoSize where
  pretty :: forall ann. MonoSize -> Doc ann
pretty (MonoKnown Int
i) = Doc ann
"?" forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Int
i
  pretty (MonoAnon VName
v) = Doc ann
"?" forall a. Semigroup a => a -> a -> a
<> forall v a. IsName v => v -> Doc a
prettyName VName
v

instance Pretty (Shape MonoSize) where
  pretty :: forall ann. Shape MonoSize -> Doc ann
pretty (Shape [MonoSize]
ds) = forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map (forall ann. Doc ann -> Doc ann
brackets forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Pretty a => a -> Doc ann
pretty) [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 Size als -> MonoType
monoType :: forall als. TypeBase Size als -> MonoType
monoType = (forall s a. State s a -> s -> a
`evalState` (Int
0, forall a. Monoid a => a
mempty)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims forall {f :: * -> *} {p}.
MonadState (Int, Map Size Int) f =>
Set VName -> p -> Size -> f MonoSize
onDim forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct
  where
    onDim :: Set VName -> p -> Size -> f MonoSize
onDim Set VName
bound p
_ (NamedSize QualName VName
d)
      -- A locally bound size.
      | forall vn. QualName vn -> vn
qualLeaf QualName VName
d forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ VName -> MonoSize
MonoAnon forall a b. (a -> b) -> a -> b
$ forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    onDim Set VName
_ p
_ Size
d = do
      (Int
i, Map Size Int
m) <- forall s (m :: * -> *). MonadState s m => m s
get
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Size
d Map Size Int
m of
        Just Int
prev ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
prev
        Maybe Int
Nothing -> do
          forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i forall a. Num a => a -> a -> a
+ Int
1, forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Size
d Int
i Map Size Int
m)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 = forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall s (m :: * -> *). MonadState s m => m s
get

modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts = forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) :)

lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted VName
fname MonoType
t = forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (VName
fname, MonoType
t) 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 -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname StructType
t
  | VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {vn}. QualName vn -> ExpBase Info vn
var QualName VName
fname
  | Bool
otherwise = do
      StructType
t' <- StructType -> MonoM StructType
removeTypeVariablesInType StructType
t
      let mono_t :: MonoType
mono_t = forall als. TypeBase Size als -> MonoType
monoType StructType
t'
      Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t
      Maybe PolyBinding
maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
lookupFun forall a b. (a -> b) -> a -> b
$ 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
_) ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {vn} {as}.
vn -> TypeBase Size as -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs VName
fname' StructType
t' forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer StructType
t'
        -- An intrinsic function.
        (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ 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
          forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell forall a b. (a -> b) -> a -> b
$ forall a. a -> Seq a
Seq.singleton (forall vn. QualName vn -> vn
qualLeaf QualName VName
fname, ValBind
funbind')
          VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t (VName
fname', InferSizeArgs
infer)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {vn} {as}.
vn -> TypeBase Size as -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs VName
fname' StructType
t' forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer StructType
t'
  where
    var :: QualName vn -> ExpBase Info vn
var QualName vn
fname' = forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (forall a. a -> Info a
Info (forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct StructType
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 forall a. Num a => a -> a -> a
- Int
1,
        forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          (forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
f ExpBase Info vn
size_arg (forall a. a -> Info a
Info (Diet
Observe, forall a. Maybe a
Nothing)) SrcLoc
loc)
          (forall a. a -> Info a
Info forall a b. (a -> b) -> a -> b
$ PatType -> [VName] -> AppRes
AppRes (forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (forall a. Int -> a -> [a]
replicate Int
i forall dim als. TypeBase dim als
i64) (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct StructType
t))) [])
      )

    applySizeArgs :: vn -> TypeBase Size as -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs vn
fname' TypeBase Size as
t' [ExpBase Info vn]
size_args =
      forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          forall {vn}.
(Int, ExpBase Info vn) -> ExpBase Info vn -> (Int, ExpBase Info vn)
applySizeArg
          ( forall (t :: * -> *) a. Foldable t => t a -> Int
length [ExpBase Info vn]
size_args forall a. Num a => a -> a -> a
- Int
1,
            forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var
              (forall v. v -> QualName v
qualName vn
fname')
              ( forall a. a -> Info a
Info
                  ( forall as dim pas.
Monoid as =>
[TypeBase dim pas] -> RetTypeBase dim as -> TypeBase dim as
foldFunType
                      (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const forall dim als. TypeBase dim als
i64) [ExpBase Info vn]
size_args)
                      (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] forall a b. (a -> b) -> a -> b
$ forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct TypeBase Size 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 <- 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 <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v RecordReplacements
rrs =
            forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [a]
M.elems RecordReplacement
d
      replace Alias
x = 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.
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((forall k a. Ord k => k -> Map k a -> Bool
`M.member` RecordReplacements
rrs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) forall a b. (a -> b) -> a -> b
$ forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t
      then forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Alias -> Aliasing
replace forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
S.toList) TypeBase dim Aliasing
t
      else TypeBase dim Aliasing
t

sizesForPat :: MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat :: forall (m :: * -> *). MonadFreshNames m => Pat -> m ([VName], Pat)
sizesForPat Pat
pat = do
  (Pat
params', [VName]
sizes) <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT [VName] m)
tv Pat
pat) []
  forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName]
sizes, Pat
params')
  where
    tv :: ASTMapper (StateT [VName] m)
tv = forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatType :: PatType -> StateT [VName] m PatType
mapOnPatType = 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 forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
Size -> t m Size
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure}
    onDim :: Size -> t m Size
onDim (AnySize Maybe VName
_) = do
      VName
v <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"size"
      forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v :)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
v
    onDim Size
d = forall (f :: * -> *) a. Applicative f => a -> f a
pure Size
d

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

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

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

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

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

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

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

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

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

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

noticeDims :: TypeBase Size as -> MonoM ()
noticeDims :: forall as. TypeBase Size as -> MonoM ()
noticeDims = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ VName -> MonoM ()
notice forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall as. TypeBase Size as -> Set VName
freeInType
  where
    notice :: VName -> MonoM ()
notice VName
v = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName forall a. Monoid a => a
mempty (forall v. v -> QualName v
qualName VName
v) forall dim als. TypeBase dim als
i64

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

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

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

type DimInst = M.Map VName Size

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

inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp]
inferSizeArgs :: [TypeParamBase VName] -> StructType -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams StructType
bind_t StructType
t =
  forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall {k}. Ord k => Map k Size -> TypeParamBase k -> Maybe Exp
tparamArg (forall a. Monoid a => TypeBase Size a -> TypeBase Size a -> DimInst
dimMapping StructType
bind_t StructType
t)) [TypeParamBase VName]
tparams
  where
    tparamArg :: Map k Size -> TypeParamBase k -> Maybe Exp
tparamArg Map k Size
dinst TypeParamBase k
tp =
      case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase k
tp) Map k Size
dinst of
        Just (NamedSize QualName VName
d) ->
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn.
QualName vn -> f PatType -> SrcLoc -> ExpBase f vn
Var QualName VName
d (forall a. a -> Info a
Info forall dim als. TypeBase dim als
i64) forall a. Monoid a => a
mempty
        Just (ConstSize Int
x) ->
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) forall a. Monoid a => a
mempty
        Maybe Size
_ ->
          forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
0) 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 = forall {dim}. TypeBase dim () -> TypeBase dim ()
f
  where
    f :: TypeBase dim () -> TypeBase dim ()
f (Array () Uniqueness
u Shape dim
shape ScalarTypeBase dim ()
t) = forall dim as.
as
-> Uniqueness
-> Shape dim
-> ScalarTypeBase dim ()
-> TypeBase dim as
Array () Uniqueness
u Shape dim
shape (ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t)
    f (Scalar ScalarTypeBase dim ()
t) = forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t
    f' :: ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' (Arrow () PName
_ TypeBase dim ()
t1 (RetType [VName]
dims TypeBase dim ()
t2)) =
      forall dim as.
as
-> PName
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t2))
    f' (Record Map Name (TypeBase dim ())
fs) =
      forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record forall a b. (a -> b) -> a -> b
$ 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) =
      forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim () -> TypeBase dim ()
f) Map Name [TypeBase dim ()]
cs
    f' ScalarTypeBase dim ()
t = ScalarTypeBase dim ()
t

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

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

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

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

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

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

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

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

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

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

    sub :: TypeBase dim () -> MonoType -> t (t m) ()
sub t1 :: TypeBase dim ()
t1@Array {} t2 :: MonoType
t2@Array {}
      | Just TypeBase dim ()
t1' <- forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim ()
t1) TypeBase dim ()
t1,
        Just MonoType
t2' <- forall dim as. Int -> TypeBase dim as -> Maybe (TypeBase dim as)
peelArray (forall dim as. TypeBase dim as -> Int
arrayRank TypeBase dim ()
t1) MonoType
t2 =
          TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1' MonoType
t2'
    sub (Scalar (TypeVar ()
_ Uniqueness
_ QualName VName
v [TypeArg dim]
_)) MonoType
t =
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Int
baseTag (forall vn. QualName vn -> vn
qualLeaf QualName VName
v) forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag) forall a b. (a -> b) -> a -> b
$
        forall {k} {as} {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *}.
(Ord k,
 MonadState (Map k (RetTypeBase Size as), Map Int VName) (t (t m)),
 MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m)), Monad (t m)) =>
QualName k -> RetTypeBase MonoSize as -> t (t m) ()
addSubst QualName VName
v forall a b. (a -> b) -> a -> b
$
          forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] MonoType
t
    sub (Scalar (Record Map Name (TypeBase dim ())
fields1)) (Scalar (Record Map Name MonoType
fields2)) =
      forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
        TypeBase dim () -> MonoType -> t (t m) ()
sub
        (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
sortFields Map Name (TypeBase dim ())
fields1)
        (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall a. Map Name a -> [(Name, a)]
sortFields Map Name MonoType
fields2)
    sub (Scalar Prim {}) (Scalar Prim {}) = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    sub (Scalar (Arrow ()
_ PName
_ TypeBase dim ()
t1a (RetType [VName]
_ TypeBase dim ()
t1b))) (Scalar (Arrow ()
_ PName
_ MonoType
t2a RetTypeBase MonoSize ()
t2b)) = do
      TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1a MonoType
t2a
      TypeBase dim () -> RetTypeBase MonoSize () -> t (t m) ()
subRet TypeBase dim ()
t1b RetTypeBase MonoSize ()
t2b
    sub (Scalar (Sum Map Name [TypeBase dim ()]
cs1)) (Scalar (Sum Map Name [MonoType]
cs2)) =
      forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ forall {a} {a}.
(a, [TypeBase dim ()]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [TypeBase dim ()]
cs1) (forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [MonoType]
cs2)
      where
        typeSubstClause :: (a, [TypeBase dim ()]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (a
_, [TypeBase dim ()]
ts1) (a
_, [MonoType]
ts2) = forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM TypeBase dim () -> MonoType -> t (t m) ()
sub [TypeBase dim ()]
ts1 [MonoType]
ts2
    sub t1 :: TypeBase dim ()
t1@(Scalar Sum {}) MonoType
t2 = TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1 MonoType
t2
    sub TypeBase dim ()
t1 t2 :: MonoType
t2@(Scalar Sum {}) = TypeBase dim () -> MonoType -> t (t m) ()
sub TypeBase dim ()
t1 MonoType
t2
    sub TypeBase dim ()
t1 MonoType
t2 = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines [[Char]
"typeSubstsM: mismatched types:", forall a. Pretty a => a -> [Char]
prettyString TypeBase dim ()
t1, forall a. Pretty a => a -> [Char]
prettyString MonoType
t2]

    addSubst :: QualName k -> RetTypeBase MonoSize as -> t (t m) ()
addSubst (QualName [k]
_ k
v) (RetType [VName]
ext TypeBase MonoSize as
t) = do
      (Map k (RetTypeBase Size as)
ts, Map Int VName
sizes) <- forall s (m :: * -> *). MonadState s m => m s
get
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (k
v forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map k (RetTypeBase Size as)
ts) forall a b. (a -> b) -> a -> b
$ do
        TypeBase Size as
t' <- 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 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) Size
onDim forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeBase MonoSize as
t
        forall s (m :: * -> *). MonadState s m => s -> m ()
put (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
v (forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext TypeBase Size as
t') Map k (RetTypeBase Size as)
ts, Map Int VName
sizes)

    onDim :: MonoSize -> t (t m) Size
onDim (MonoKnown Int
i) = do
      (a
ts, Map Int VName
sizes) <- forall s (m :: * -> *). MonadState s m => m s
get
      case 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 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"d"
          forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell [forall vn. vn -> SrcLoc -> TypeParamBase vn
TypeParamDim VName
d SrcLoc
loc]
          forall s (m :: * -> *). MonadState s m => s -> m ()
put (a
ts, forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Int
i VName
d Map Int VName
sizes)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
d
        Just VName
d ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ QualName VName -> Size
NamedSize forall a b. (a -> b) -> a -> b
$ forall v. v -> QualName v
qualName VName
d
    onDim (MonoAnon VName
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Maybe VName -> Size
AnySize forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just VName
v

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

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

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

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

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

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

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

transformEntryPoint :: EntryPoint -> MonoM EntryPoint
transformEntryPoint :: EntryPoint -> MonoM EntryPoint
transformEntryPoint (EntryPoint [EntryParam]
params EntryType
ret) =
  [EntryParam] -> EntryType -> EntryPoint
EntryPoint forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM EntryParam -> MonoM EntryParam
onEntryParam [EntryParam]
params forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> EntryType -> MonoM EntryType
onEntryType EntryType
ret
  where
    onEntryParam :: EntryParam -> MonoM EntryParam
onEntryParam (EntryParam Name
v EntryType
t) =
      Name -> EntryType -> EntryParam
EntryParam Name
v forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EntryType -> MonoM EntryType
onEntryType EntryType
t
    onEntryType :: EntryType -> MonoM EntryType
onEntryType (EntryType StructType
t Maybe (TypeExp VName)
te) =
      StructType -> Maybe (TypeExp VName) -> EntryType
EntryType forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> MonoM StructType
removeTypeVariablesInType StructType
t forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (TypeExp VName)
te

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

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

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

transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeExp VName
_ (Info (RetType [VName]
dims StructType
t)) Maybe DocComment
_ SrcLoc
_) = do
  Map VName (Subst StructRetType)
subs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  forall as. TypeBase Size as -> MonoM ()
noticeDims StructType
t
  let tbinding :: TypeBinding
tbinding = Liftedness -> [TypeParamBase VName] -> StructRetType -> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams forall a b. (a -> b) -> a -> b
$ forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims forall a b. (a -> b) -> a -> b
$ forall a. Substitutable a => TypeSubs -> a -> a
applySubst (forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) StructType
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty {envTypeBindings :: Map VName TypeBinding
envTypeBindings = forall k a. k -> a -> Map k a
M.singleton VName
name TypeBinding
tbinding}

transformDecs :: [Dec] -> MonoM ()
transformDecs :: [Dec] -> MonoM ()
transformDecs [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
transformDecs (ValDec ValBind
valbind : [Dec]
ds) = do
  Env
env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
  forall a. Env -> MonoM a -> MonoM a
localEnv Env
env 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
  forall a. Env -> MonoM a -> MonoM a
localEnv Env
env forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (Dec
dec : [Dec]
_) =
  forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$
    [Char]
"The monomorphization module expects a module-free "
      forall a. [a] -> [a] -> [a]
++ [Char]
"input program, but received: "
      forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> [Char]
prettyString 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 =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
      forall a.
VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
runMonoM VNameSource
namesrc forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
decs