-- | 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.List.NonEmpty qualified as NE
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 vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
          ExpBase Info vn
f
          [(Diet
Observe, forall a. Maybe a
Nothing, ExpBase Info vn
size_arg)]
          (PatType -> [VName] -> AppRes
AppRes (forall as dim pas.
Monoid as =>
[(Diet, TypeBase dim pas)] -> RetTypeBase dim as -> TypeBase dim as
foldFunType (forall a. Int -> a -> [a]
replicate Int
i (Diet
Observe, 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 =>
[(Diet, 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 (Diet
Observe, 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 Info 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 f 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 Info 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 Info 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 f vn),
    f StructRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat]
params, Maybe (TypeExp Info 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
fe NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
_) AppRes
res =
  forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
fe
    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 forall {t} {t}. (Info (t, t), Exp) -> MonoM (t, t, Exp)
onArg (forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Diet, Maybe VName), Exp)
args)
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure AppRes
res
  where
    onArg :: (Info (t, t), Exp) -> MonoM (t, t, Exp)
onArg (Info (t
d, t
ext), Exp
e) = (t
d,t
ext,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e
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 vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
        (forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply ExpBase Info vn
fname' [(Diet
Observe, forall a b. (a, b) -> b
snd (forall a. Info a -> a
unInfo Info (StructType, Maybe VName)
d1), ExpBase Info vn
x)] (PatType -> [VName] -> AppRes
AppRes PatType
ret forall a. Monoid a => a
mempty))
        [(Diet
Observe, forall a b. (a, b) -> b
snd (forall a. Info a -> a
unInfo Info (StructType, Maybe VName)
d2), ExpBase Info vn
y)]
        (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 Info VName
tp SrcLoc
loc) =
  forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp f 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 Info 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 Info 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 f vn)
-> f (Aliasing, StructRetType)
-> SrcLoc
-> ExpBase f vn
Lambda [Pat]
params Exp
e0' Maybe (TypeExp Info 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 vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
          Exp
op
          [(Diet
Observe, Maybe VName
xext, Exp
e1)]
          (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
-> Diet
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow forall a. Monoid a => a
mempty PName
yp Diet
Observe 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 vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
apply_left [(Diet
Observe, Maybe VName
yext, Exp
e2)] (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 f 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
_ Diet
_ 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 f 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
_ Diet
_ 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 f 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 Info 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 f 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 (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM 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 (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM 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 Info 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 f vn -> SrcLoc -> PatBase f vn
PatAscription Pat
pat' TypeExp Info 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 (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM 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 Int64
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 Int64
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
_ Diet
d1 TypeBase dim ()
t1 (RetType [VName]
dims TypeBase dim ()
t2)) =
      forall dim as.
as
-> PName
-> Diet
-> TypeBase dim ()
-> RetTypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed Diet
d1 (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 = [Pat] -> StructRetType -> StructType
funType [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 (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM 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 Info 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
_ Diet
_ TypeBase dim ()
t1a (RetType [VName]
_ TypeBase dim ()
t1b))) (Scalar (Arrow ()
_ PName
_ Diet
_ 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 Info VName
td SrcLoc
loc
    | Bool
entry -> forall (f :: * -> *) vn.
PatBase f vn -> TypeExp f vn -> SrcLoc -> PatBase f vn
PatAscription (Bool -> (PatType -> PatType) -> Pat -> Pat
substPat Bool
False PatType -> PatType
f Pat
p) TypeExp Info 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 Info 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 Info 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 Info VName)
te) =
      StructType -> Maybe (TypeExp Info 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 Info 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
$
          [Pat] -> StructRetType -> StructType
funType (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 Info 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