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

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

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

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

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

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

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

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

instance Semigroup Env where
  Env Map VName PolyBinding
tb1 Map VName TypeBinding
pb1 RecordReplacements
rr1 <> :: Env -> Env -> Env
<> Env Map VName PolyBinding
tb2 Map VName TypeBinding
pb2 RecordReplacements
rr2 = Map VName PolyBinding
-> Map VName TypeBinding -> RecordReplacements -> Env
Env (Map VName PolyBinding
tb1 Map VName PolyBinding
-> Map VName PolyBinding -> Map VName PolyBinding
forall a. Semigroup a => a -> a -> a
<> Map VName PolyBinding
tb2) (Map VName TypeBinding
pb1 Map VName TypeBinding
-> Map VName TypeBinding -> Map VName TypeBinding
forall a. Semigroup a => a -> a -> a
<> Map VName TypeBinding
pb2) (RecordReplacements
rr1 RecordReplacements -> RecordReplacements -> RecordReplacements
forall a. Semigroup a => a -> a -> a
<> RecordReplacements
rr2)

instance Monoid Env where
  mempty :: Env
mempty = Map VName PolyBinding
-> Map VName TypeBinding -> RecordReplacements -> Env
Env Map VName PolyBinding
forall a. Monoid a => a
mempty Map VName TypeBinding
forall a. Monoid a => a
mempty RecordReplacements
forall a. Monoid a => a
mempty

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

extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv :: forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
vn PolyBinding
binding =
  Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv
    Env
forall a. Monoid a => a
mempty {envPolyBindings :: Map VName PolyBinding
envPolyBindings = VName -> PolyBinding -> Map VName PolyBinding
forall k a. k -> a -> Map k a
M.singleton VName
vn PolyBinding
binding}

withRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements :: forall a. RecordReplacements -> MonoM a -> MonoM a
withRecordReplacements RecordReplacements
rr = Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv Env
forall a. Monoid a => a
mempty {envRecordReplacements :: RecordReplacements
envRecordReplacements = RecordReplacements
rr}

replaceRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements :: forall a. RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr = (Env -> Env) -> MonoM a -> MonoM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> MonoM a -> MonoM a)
-> (Env -> Env) -> MonoM a -> MonoM a
forall a b. (a -> b) -> a -> b
$ \Env
env -> Env
env {envRecordReplacements :: RecordReplacements
envRecordReplacements = RecordReplacements
rr}

-- The monomorphization monad.
newtype MonoM a
  = MonoM
      ( RWST
          Env
          (Seq.Seq (VName, ValBind))
          VNameSource
          (State Lifts)
          a
      )
  deriving
    ( (forall a b. (a -> b) -> MonoM a -> MonoM b)
-> (forall a b. a -> MonoM b -> MonoM a) -> Functor MonoM
forall a b. a -> MonoM b -> MonoM a
forall a b. (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> MonoM b -> MonoM a
$c<$ :: forall a b. a -> MonoM b -> MonoM a
fmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
$cfmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
Functor,
      Functor MonoM
Functor MonoM
-> (forall a. a -> MonoM a)
-> (forall a b. MonoM (a -> b) -> MonoM a -> MonoM b)
-> (forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c)
-> (forall a b. MonoM a -> MonoM b -> MonoM b)
-> (forall a b. MonoM a -> MonoM b -> MonoM a)
-> Applicative MonoM
forall a. a -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM b
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. MonoM a -> MonoM b -> MonoM a
$c<* :: forall a b. MonoM a -> MonoM b -> MonoM a
*> :: forall a b. MonoM a -> MonoM b -> MonoM b
$c*> :: forall a b. MonoM a -> MonoM b -> MonoM b
liftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
$cliftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
$c<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
pure :: forall a. a -> MonoM a
$cpure :: forall a. a -> MonoM a
Applicative,
      Applicative MonoM
Applicative MonoM
-> (forall a b. MonoM a -> (a -> MonoM b) -> MonoM b)
-> (forall a b. MonoM a -> MonoM b -> MonoM b)
-> (forall a. a -> MonoM a)
-> Monad MonoM
forall a. a -> MonoM a
forall a b. MonoM a -> MonoM b -> MonoM b
forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> MonoM a
$creturn :: forall a. a -> MonoM a
>> :: forall a b. MonoM a -> MonoM b -> MonoM b
$c>> :: forall a b. MonoM a -> MonoM b -> MonoM b
>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
$c>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
Monad,
      MonadReader Env,
      MonadWriter (Seq.Seq (VName, ValBind)),
      Monad MonoM
Applicative MonoM
MonoM VNameSource
Applicative MonoM
-> Monad MonoM
-> MonoM VNameSource
-> (VNameSource -> MonoM ())
-> MonadFreshNames MonoM
VNameSource -> MonoM ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> MonoM ()
$cputNameSource :: VNameSource -> MonoM ()
getNameSource :: MonoM VNameSource
$cgetNameSource :: MonoM VNameSource
MonadFreshNames
    )

runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource)
runMonoM :: forall a.
VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
runMonoM VNameSource
src (MonoM RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
m) = ((a
a, Seq (VName, ValBind)
defs), VNameSource
src')
  where
    (a
a, VNameSource
src', Seq (VName, ValBind)
defs) = State Lifts (a, VNameSource, Seq (VName, ValBind))
-> Lifts -> (a, VNameSource, Seq (VName, ValBind))
forall s a. State s a -> s -> a
evalState (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> Env
-> VNameSource
-> State Lifts (a, VNameSource, Seq (VName, ValBind))
forall r w s (m :: * -> *) a.
RWST r w s m a -> r -> s -> m (a, s, w)
runRWST RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
m Env
forall a. Monoid a => a
mempty VNameSource
src) Lifts
forall a. Monoid a => a
mempty

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

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

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

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

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

instance Pretty MonoSize where
  ppr :: MonoSize -> Doc
ppr (MonoKnown Int
i) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Int -> Doc
forall a. Pretty a => a -> Doc
ppr Int
i
  ppr (MonoAnon Maybe VName
Nothing) = Doc
forall a. Monoid a => a
mempty
  ppr (MonoAnon (Just VName
v)) = String -> Doc
text String
"?" Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> VName -> Doc
forall v. IsName v => v -> Doc
pprName VName
v

instance Pretty (ShapeDecl MonoSize) where
  ppr :: ShapeDecl MonoSize -> Doc
ppr (ShapeDecl [MonoSize]
ds) = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((MonoSize -> Doc) -> [MonoSize] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (MonoSize -> Doc) -> MonoSize -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MonoSize -> Doc
forall a. Pretty a => a -> Doc
ppr) [MonoSize]
ds)

-- The kind of type relative to which we monomorphise.  What is most
-- important to us is not the specific dimensions, but merely whether
-- they are known or anonymous/local.
type MonoType = TypeBase MonoSize ()

monoType :: TypeBase (DimDecl VName) als -> MonoType
monoType :: forall als. TypeBase (DimDecl VName) als -> MonoType
monoType = (State (Int, Map (DimDecl VName) Int) MonoType
-> (Int, Map (DimDecl VName) Int) -> MonoType
forall s a. State s a -> s -> a
`evalState` (Int
0, Map (DimDecl VName) Int
forall a. Monoid a => a
mempty)) (State (Int, Map (DimDecl VName) Int) MonoType -> MonoType)
-> (TypeBase (DimDecl VName) als
    -> State (Int, Map (DimDecl VName) Int) MonoType)
-> TypeBase (DimDecl VName) als
-> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
 -> DimPos
 -> DimDecl VName
 -> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize)
-> StructType -> State (Int, Map (DimDecl VName) Int) MonoType
forall (f :: * -> *) fdim tdim als.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim als -> f (TypeBase tdim als)
traverseDims Set VName
-> DimPos
-> DimDecl VName
-> StateT (Int, Map (DimDecl VName) Int) Identity MonoSize
forall {f :: * -> *} {p}.
MonadState (Int, Map (DimDecl VName) Int) f =>
Set VName -> p -> DimDecl VName -> f MonoSize
onDim (StructType -> State (Int, Map (DimDecl VName) Int) MonoType)
-> (TypeBase (DimDecl VName) als -> StructType)
-> TypeBase (DimDecl VName) als
-> State (Int, Map (DimDecl VName) Int) MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase (DimDecl VName) als -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct
  where
    onDim :: Set VName -> p -> DimDecl VName -> f MonoSize
onDim Set VName
bound p
_ (NamedDim QualName VName
d)
      -- A locally bound size.
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Maybe VName -> MonoSize
MonoAnon (Maybe VName -> MonoSize) -> Maybe VName -> MonoSize
forall a b. (a -> b) -> a -> b
$ VName -> Maybe VName
forall a. a -> Maybe a
Just (VName -> Maybe VName) -> VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    onDim Set VName
_ p
_ (AnyDim Maybe VName
v) = MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Maybe VName -> MonoSize
MonoAnon Maybe VName
v
    onDim Set VName
_ p
_ DimDecl VName
d = do
      (Int
i, Map (DimDecl VName) Int
m) <- f (Int, Map (DimDecl VName) Int)
forall s (m :: * -> *). MonadState s m => m s
get
      case DimDecl VName -> Map (DimDecl VName) Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup DimDecl VName
d Map (DimDecl VName) Int
m of
        Just Int
prev ->
          MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
prev
        Maybe Int
Nothing -> do
          (Int, Map (DimDecl VName) Int) -> f ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, DimDecl VName
-> Int -> Map (DimDecl VName) Int -> Map (DimDecl VName) Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert DimDecl VName
d Int
i Map (DimDecl VName) Int
m)
          MonoSize -> f MonoSize
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> f MonoSize) -> MonoSize -> f MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
i

-- Mapping from function name and instance list to a new function name in case
-- the function has already been instantiated with those concrete types.
type Lifts = [((VName, MonoType), (VName, InferSizeArgs))]

getLifts :: MonoM Lifts
getLifts :: MonoM Lifts
getLifts = RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
-> MonoM Lifts
forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
 -> MonoM Lifts)
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
-> MonoM Lifts
forall a b. (a -> b) -> a -> b
$ State Lifts Lifts
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) Lifts
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift State Lifts Lifts
forall s (m :: * -> *). MonadState s m => m s
get

modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts :: (Lifts -> Lifts) -> MonoM ()
modifyLifts = RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
-> MonoM ()
forall a.
RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) a
-> MonoM a
MonoM (RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
 -> MonoM ())
-> ((Lifts -> Lifts)
    -> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ())
-> (Lifts -> Lifts)
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. State Lifts ()
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (State Lifts ()
 -> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ())
-> ((Lifts -> Lifts) -> State Lifts ())
-> (Lifts -> Lifts)
-> RWST Env (Seq (VName, ValBind)) VNameSource (State Lifts) ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Lifts -> Lifts) -> State Lifts ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify

addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted VName
fname MonoType
il (VName, InferSizeArgs)
liftf =
  (Lifts -> Lifts) -> MonoM ()
modifyLifts (((VName
fname, MonoType
il), (VName, InferSizeArgs)
liftf) ((VName, MonoType), (VName, InferSizeArgs)) -> Lifts -> Lifts
forall a. a -> [a] -> [a]
:)

lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted VName
fname MonoType
t = (VName, MonoType) -> Lifts -> Maybe (VName, InferSizeArgs)
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (VName
fname, MonoType
t) (Lifts -> Maybe (VName, InferSizeArgs))
-> MonoM Lifts -> MonoM (Maybe (VName, InferSizeArgs))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM Lifts
getLifts

transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName :: SrcLoc
-> QualName VName -> StructType -> MonoM (ExpBase Info VName)
transformFName SrcLoc
loc QualName VName
fname StructType
t
  | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag = ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> ExpBase Info VName
forall {vn}. QualName vn -> ExpBase Info vn
var QualName VName
fname
  | Bool
otherwise = do
    StructType
t' <- StructType -> MonoM StructType
removeTypeVariablesInType StructType
t
    let mono_t :: MonoType
mono_t = StructType -> MonoType
forall als. TypeBase (DimDecl VName) als -> MonoType
monoType StructType
t'
    Maybe (VName, InferSizeArgs)
maybe_fname <- VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs))
lookupLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t
    Maybe PolyBinding
maybe_funbind <- VName -> MonoM (Maybe PolyBinding)
lookupFun (VName -> MonoM (Maybe PolyBinding))
-> VName -> MonoM (Maybe PolyBinding)
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname
    case (Maybe (VName, InferSizeArgs)
maybe_fname, Maybe PolyBinding
maybe_funbind) of
      -- The function has already been monomorphised.
      (Just (VName
fname', InferSizeArgs
infer), Maybe PolyBinding
_) ->
        ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ VName -> StructType -> [ExpBase Info VName] -> ExpBase Info VName
forall {vn} {as}.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' StructType
t' ([ExpBase Info VName] -> ExpBase Info VName)
-> [ExpBase Info VName] -> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer StructType
t'
      -- An intrinsic function.
      (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> ExpBase Info VName
forall {vn}. QualName vn -> ExpBase Info vn
var QualName VName
fname
      -- A polymorphic function.
      (Maybe (VName, InferSizeArgs)
Nothing, Just PolyBinding
funbind) -> do
        (VName
fname', InferSizeArgs
infer, ValBind
funbind') <- Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
False PolyBinding
funbind MonoType
mono_t
        Seq (VName, ValBind) -> MonoM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (VName, ValBind) -> MonoM ())
-> Seq (VName, ValBind) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (VName, ValBind) -> Seq (VName, ValBind)
forall a. a -> Seq a
Seq.singleton (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname, ValBind
funbind')
        VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
fname) MonoType
mono_t (VName
fname', InferSizeArgs
infer)
        ExpBase Info VName -> MonoM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> MonoM (ExpBase Info VName))
-> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ VName -> StructType -> [ExpBase Info VName] -> ExpBase Info VName
forall {vn} {as}.
vn
-> TypeBase (DimDecl VName) as
-> [ExpBase Info vn]
-> ExpBase Info vn
applySizeArgs VName
fname' StructType
t' ([ExpBase Info VName] -> ExpBase Info VName)
-> [ExpBase Info VName] -> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ InferSizeArgs
infer StructType
t'
  where
    var :: QualName vn -> ExpBase Info vn
var QualName vn
fname' = QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (PatternType -> Info PatternType
forall a. a -> Info a
Info (StructType -> PatternType
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,
        AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          (ExpBase Info vn
-> ExpBase Info vn
-> Info (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> SrcLoc
-> AppExpBase f vn
Apply ExpBase Info vn
f ExpBase Info vn
size_arg ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing)) SrcLoc
loc)
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ PatternType -> [VName] -> AppRes
AppRes ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType (Int -> PatternType -> [PatternType]
forall a. Int -> a -> [a]
replicate Int
i PatternType
forall dim als. TypeBase dim als
i64) (StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct StructType
t)) [])
      )

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

-- This carries out record replacements in the alias information of a type.
transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType :: forall dim. TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
transformType TypeBase dim Aliasing
t = do
  RecordReplacements
rrs <- (Env -> RecordReplacements) -> MonoM RecordReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> RecordReplacements
envRecordReplacements
  let replace :: Alias -> Aliasing
replace (AliasBound VName
v)
        | Just RecordReplacement
d <- VName -> RecordReplacements -> Maybe RecordReplacement
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v RecordReplacements
rrs =
          [Alias] -> Aliasing
forall a. Ord a => [a] -> Set a
S.fromList ([Alias] -> Aliasing) -> [Alias] -> Aliasing
forall a b. (a -> b) -> a -> b
$ ((VName, PatternType) -> Alias)
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> Alias
AliasBound (VName -> Alias)
-> ((VName, PatternType) -> VName) -> (VName, PatternType) -> Alias
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, PatternType) -> VName
forall a b. (a, b) -> a
fst) ([(VName, PatternType)] -> [Alias])
-> [(VName, PatternType)] -> [Alias]
forall a b. (a -> b) -> a -> b
$ RecordReplacement -> [(VName, PatternType)]
forall k a. Map k a -> [a]
M.elems RecordReplacement
d
      replace Alias
x = Alias -> Aliasing
forall a. a -> Set a
S.singleton Alias
x
  -- As an attempt at an optimisation, only transform the aliases if
  -- they refer to a variable we have record-replaced.
  TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing))
-> TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing)
forall a b. (a -> b) -> a -> b
$
    if (Alias -> Bool) -> Aliasing -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((VName -> RecordReplacements -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` RecordReplacements
rrs) (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (Aliasing -> Bool) -> Aliasing -> Bool
forall a b. (a -> b) -> a -> b
$ TypeBase dim Aliasing -> Aliasing
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim Aliasing
t
      then (Aliasing -> Aliasing)
-> TypeBase dim Aliasing -> TypeBase dim Aliasing
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([Aliasing] -> Aliasing
forall a. Monoid a => [a] -> a
mconcat ([Aliasing] -> Aliasing)
-> (Aliasing -> [Aliasing]) -> Aliasing -> Aliasing
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Alias -> Aliasing) -> [Alias] -> [Aliasing]
forall a b. (a -> b) -> [a] -> [b]
map Alias -> Aliasing
replace ([Alias] -> [Aliasing])
-> (Aliasing -> [Alias]) -> Aliasing -> [Aliasing]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Aliasing -> [Alias]
forall a. Set a -> [a]
S.toList) TypeBase dim Aliasing
t
      else TypeBase dim Aliasing
t

sizesForPat :: MonadFreshNames m => Pattern -> m ([VName], Pattern)
sizesForPat :: forall (m :: * -> *).
MonadFreshNames m =>
Pattern -> m ([VName], Pattern)
sizesForPat Pattern
pat = do
  (Pattern
params', [VName]
sizes) <- StateT [VName] m Pattern -> [VName] -> m (Pattern, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (ASTMapper (StateT [VName] m) -> Pattern -> StateT [VName] m Pattern
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT [VName] m)
tv Pattern
pat) []
  ([VName], Pattern) -> m ([VName], Pattern)
forall (m :: * -> *) a. Monad m => a -> m a
return ([VName]
sizes, Pattern
params')
  where
    tv :: ASTMapper (StateT [VName] m)
tv = ASTMapper (StateT [VName] m)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatternType :: PatternType -> StateT [VName] m PatternType
mapOnPatternType = (DimDecl VName -> StateT [VName] m (DimDecl VName))
-> (Aliasing -> StateT [VName] m Aliasing)
-> PatternType
-> StateT [VName] m PatternType
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse DimDecl VName -> StateT [VName] m (DimDecl VName)
forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
DimDecl VName -> t m (DimDecl VName)
onDim Aliasing -> StateT [VName] m Aliasing
forall (f :: * -> *) a. Applicative f => a -> f a
pure}
    onDim :: DimDecl VName -> t m (DimDecl VName)
onDim (AnyDim Maybe VName
_) = do
      VName
v <- m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"size"
      ([VName] -> [VName]) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:)
      DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
    onDim DimDecl VName
d = DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d

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

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

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

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

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

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

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

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

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

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

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

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

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

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

type DimInst = M.Map VName (DimDecl VName)

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

inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp]
inferSizeArgs :: [TypeParamBase VName] -> StructType -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams StructType
bind_t StructType
t =
  (TypeParamBase VName -> Maybe (ExpBase Info VName))
-> [TypeParamBase VName] -> [ExpBase Info VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DimInst -> TypeParamBase VName -> Maybe (ExpBase Info VName)
forall {k} {vn}.
Ord k =>
Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg (StructType -> StructType -> DimInst
forall a.
Monoid a =>
TypeBase (DimDecl VName) a -> TypeBase (DimDecl VName) a -> DimInst
dimMapping StructType
bind_t StructType
t)) [TypeParamBase VName]
tparams
  where
    tparamArg :: Map k (DimDecl vn) -> TypeParamBase k -> Maybe (ExpBase Info vn)
tparamArg Map k (DimDecl vn)
dinst TypeParamBase k
tp =
      case k -> Map k (DimDecl vn) -> Maybe (DimDecl vn)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (TypeParamBase k -> k
forall vn. TypeParamBase vn -> vn
typeParamName TypeParamBase k
tp) Map k (DimDecl vn)
dinst of
        Just (NamedDim QualName vn
d) ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ QualName vn -> Info PatternType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName vn
d (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
forall dim als. TypeBase dim als
i64) SrcLoc
forall a. Monoid a => a
mempty
        Just (ConstDim Int
x) ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value (Int64 -> IntValue) -> Int64 -> IntValue
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) SrcLoc
forall a. Monoid a => a
mempty
        Maybe (DimDecl vn)
_ ->
          ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a. a -> Maybe a
Just (ExpBase Info vn -> Maybe (ExpBase Info vn))
-> ExpBase Info vn -> Maybe (ExpBase Info vn)
forall a b. (a -> b) -> a -> b
$ PrimValue -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (IntValue -> PrimValue) -> IntValue -> PrimValue
forall a b. (a -> b) -> a -> b
$ Int64 -> IntValue
Int64Value Int64
0) SrcLoc
forall a. Monoid a => a
mempty

-- Monomorphising higher-order functions can result in function types
-- where the same named parameter occurs in multiple spots.  When
-- monomorphising we don't really need those parameter names anymore,
-- and the defunctionaliser can be confused if there are duplicates
-- (it doesn't handle shadowing), so let's just remove all parameter
-- names here.  This is safe because a MonoType does not contain sizes
-- anyway.
noNamedParams :: MonoType -> MonoType
noNamedParams :: MonoType -> MonoType
noNamedParams = MonoType -> MonoType
forall {dim}. TypeBase dim () -> TypeBase dim ()
f
  where
    f :: TypeBase dim () -> TypeBase dim ()
f (Array () Uniqueness
u ScalarTypeBase dim ()
t ShapeDecl dim
shape) = ()
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim ()
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array () Uniqueness
u (ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t) ShapeDecl dim
shape
    f (Scalar ScalarTypeBase dim ()
t) = ScalarTypeBase dim () -> TypeBase dim ()
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim () -> TypeBase dim ())
-> ScalarTypeBase dim () -> TypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' ScalarTypeBase dim ()
t
    f' :: ScalarTypeBase dim () -> ScalarTypeBase dim ()
f' (Arrow () PName
_ TypeBase dim ()
t1 TypeBase dim ()
t2) =
      ()
-> PName
-> TypeBase dim ()
-> TypeBase dim ()
-> ScalarTypeBase dim ()
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow () PName
Unnamed (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t1) (TypeBase dim () -> TypeBase dim ()
f TypeBase dim ()
t2)
    f' (Record Map Name (TypeBase dim ())
fs) =
      Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim ()) -> ScalarTypeBase dim ())
-> Map Name (TypeBase dim ()) -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ (TypeBase dim () -> TypeBase dim ())
-> Map Name (TypeBase dim ()) -> Map Name (TypeBase dim ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim () -> TypeBase dim ()
f Map Name (TypeBase dim ())
fs
    f' (Sum Map Name [TypeBase dim ()]
cs) =
      Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim ()] -> ScalarTypeBase dim ())
-> Map Name [TypeBase dim ()] -> ScalarTypeBase dim ()
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim ()] -> [TypeBase dim ()])
-> Map Name [TypeBase dim ()] -> Map Name [TypeBase dim ()]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase dim () -> TypeBase dim ())
-> [TypeBase dim ()] -> [TypeBase dim ()]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase dim () -> TypeBase dim ()
f) Map Name [TypeBase dim ()]
cs
    f' ScalarTypeBase dim ()
t = ScalarTypeBase dim ()
t

-- Monomorphise a polymorphic function at the types given in the instance
-- list. Monomorphises the body of the function as well. Returns the fresh name
-- of the generated monomorphic function and its 'ValBind' representation.
monomorphiseBinding ::
  Bool ->
  PolyBinding ->
  MonoType ->
  MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding :: Bool
-> PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding Bool
entry (PolyBinding RecordReplacements
rr (VName
name, [TypeParamBase VName]
tparams, [Pattern]
params, StructType
rettype, [VName]
retext, ExpBase Info VName
body, [AttrInfo]
attrs, SrcLoc
loc)) MonoType
t =
  RecordReplacements
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. RecordReplacements -> MonoM a -> MonoM a
replaceRecordReplacements RecordReplacements
rr (MonoM (VName, InferSizeArgs, ValBind)
 -> MonoM (VName, InferSizeArgs, ValBind))
-> MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a b. (a -> b) -> a -> b
$ do
    let bind_t :: StructType
bind_t = [StructType] -> StructType -> StructType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ((Pattern -> StructType) -> [Pattern] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map Pattern -> StructType
patternStructType [Pattern]
params) StructType
rettype
    (Map VName StructType
substs, [TypeParamBase VName]
t_shape_params) <- SrcLoc
-> TypeBase () ()
-> MonoType
-> MonoM (Map VName StructType, [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () ()
-> MonoType
-> m (Map VName StructType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc (StructType -> TypeBase () ()
forall vn as. TypeBase (DimDecl vn) as -> TypeBase () as
noSizes StructType
bind_t) (MonoType -> MonoM (Map VName StructType, [TypeParamBase VName]))
-> MonoType -> MonoM (Map VName StructType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
t
    let substs' :: Map VName (Subst StructType)
substs' = (StructType -> Subst StructType)
-> Map VName StructType -> Map VName (Subst StructType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ([TypeParamBase VName] -> StructType -> Subst StructType
forall t. [TypeParamBase VName] -> t -> Subst t
Subst []) Map VName StructType
substs
        rettype' :: StructType
rettype' = (VName -> Maybe (Subst StructType)) -> StructType -> StructType
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
substs') StructType
rettype
        substPatternType :: PatternType -> PatternType
substPatternType =
          (VName -> Maybe (Subst PatternType)) -> PatternType -> PatternType
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny ((Subst StructType -> Subst PatternType)
-> Maybe (Subst StructType) -> Maybe (Subst PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructType -> PatternType)
-> Subst StructType -> Subst PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim Aliasing
fromStruct) (Maybe (Subst StructType) -> Maybe (Subst PatternType))
-> (VName -> Maybe (Subst StructType))
-> VName
-> Maybe (Subst PatternType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
substs'))
        params' :: [Pattern]
params' = (Pattern -> Pattern) -> [Pattern] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry PatternType -> PatternType
substPatternType) [Pattern]
params
        bind_t' :: StructType
bind_t' = (VName -> Maybe (Subst StructType)) -> StructType -> StructType
forall as.
Monoid as =>
(VName -> Maybe (Subst (TypeBase (DimDecl VName) as)))
-> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as
substTypesAny (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
substs') StructType
bind_t
        ([TypeParamBase VName]
shape_params_explicit, [TypeParamBase VName]
shape_params_implicit) =
          (TypeParamBase VName -> Bool)
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` StructType -> Set VName
mustBeExplicit StructType
bind_t') (VName -> Bool)
-> (TypeParamBase VName -> VName) -> TypeParamBase VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) ([TypeParamBase VName]
 -> ([TypeParamBase VName], [TypeParamBase VName]))
-> [TypeParamBase VName]
-> ([TypeParamBase VName], [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$
            [TypeParamBase VName]
shape_params [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
t_shape_params

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

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

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

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

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

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

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

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

    addSubst :: TypeName -> t MonoSize d -> t (t m) ()
addSubst (TypeName [VName]
_ VName
v) t MonoSize d
t = do
      (Map VName (t (DimDecl VName) d)
ts, Map Int VName
sizes) <- t (t m) (Map VName (t (DimDecl VName) d), Map Int VName)
forall s (m :: * -> *). MonadState s m => m s
get
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName
v VName -> Map VName (t (DimDecl VName) d) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName (t (DimDecl VName) d)
ts) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ do
        t (DimDecl VName) d
t' <- (MonoSize -> t (t m) (DimDecl VName))
-> (d -> t (t m) d)
-> t MonoSize d
-> t (t m) (t (DimDecl VName) d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse MonoSize -> t (t m) (DimDecl VName)
forall {a} {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *}.
(MonadState (a, Map Int VName) (t (t m)), MonadTrans t,
 MonadTrans t, Monad (t m), MonadFreshNames m,
 MonadWriter [TypeParamBase VName] (t (t m))) =>
MonoSize -> t (t m) (DimDecl VName)
onDim d -> t (t m) d
forall (f :: * -> *) a. Applicative f => a -> f a
pure t MonoSize d
t
        (Map VName (t (DimDecl VName) d), Map Int VName) -> t (t m) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VName
-> t (DimDecl VName) d
-> Map VName (t (DimDecl VName) d)
-> Map VName (t (DimDecl VName) d)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
v t (DimDecl VName) d
t' Map VName (t (DimDecl VName) d)
ts, Map Int VName
sizes)

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

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

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

-- Remove all type variables and type abbreviations from a value binding.
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables :: Bool -> ValBind -> MonoM ValBind
removeTypeVariables Bool
entry valbind :: ValBind
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
_ Maybe (TypeExp VName)
_ (Info (StructType
rettype, [VName]
retext)) [TypeParamBase VName]
_ [Pattern]
pats ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
_ SrcLoc
_) = do
  Map VName (Subst StructType)
subs <- (Env -> Map VName (Subst StructType))
-> MonoM (Map VName (Subst StructType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructType))
 -> MonoM (Map VName (Subst StructType)))
-> (Env -> Map VName (Subst StructType))
-> MonoM (Map VName (Subst StructType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructType)
-> Map VName TypeBinding -> Map VName (Subst StructType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  let mapper :: ASTMapper MonoM
mapper =
        ASTMapper :: forall (m :: * -> *).
(ExpBase Info VName -> m (ExpBase Info VName))
-> (VName -> m VName)
-> (QualName VName -> m (QualName VName))
-> (StructType -> m StructType)
-> (PatternType -> m PatternType)
-> ASTMapper m
ASTMapper
          { mapOnExp :: ExpBase Info VName -> MonoM (ExpBase Info VName)
mapOnExp = ASTMapper MonoM -> ExpBase Info VName -> MonoM (ExpBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper MonoM
mapper,
            mapOnName :: VName -> MonoM VName
mapOnName = VName -> MonoM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnQualName :: QualName VName -> MonoM (QualName VName)
mapOnQualName = QualName VName -> MonoM (QualName VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
            mapOnStructType :: StructType -> MonoM StructType
mapOnStructType = StructType -> MonoM StructType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType -> MonoM StructType)
-> (StructType -> StructType) -> StructType -> MonoM StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe (Subst StructType)) -> StructType -> StructType
forall a.
Substitutable a =>
(VName -> Maybe (Subst StructType)) -> a -> a
applySubst (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
subs),
            mapOnPatternType :: PatternType -> MonoM PatternType
mapOnPatternType = PatternType -> MonoM PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> MonoM PatternType)
-> (PatternType -> PatternType) -> PatternType -> MonoM PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Maybe (Subst StructType)) -> PatternType -> PatternType
forall a.
Substitutable a =>
(VName -> Maybe (Subst StructType)) -> a -> a
applySubst (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
subs)
          }

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

  ValBind -> MonoM ValBind
forall (m :: * -> *) a. Monad m => a -> m a
return
    ValBind
valbind
      { valBindRetType :: Info (StructType, [VName])
valBindRetType = (StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info ((VName -> Maybe (Subst StructType)) -> StructType -> StructType
forall a.
Substitutable a =>
(VName -> Maybe (Subst StructType)) -> a -> a
applySubst (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
subs) StructType
rettype, [VName]
retext),
        valBindParams :: [Pattern]
valBindParams = (Pattern -> Pattern) -> [Pattern] -> [Pattern]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (PatternType -> PatternType) -> Pattern -> Pattern
substPattern Bool
entry ((PatternType -> PatternType) -> Pattern -> Pattern)
-> (PatternType -> PatternType) -> Pattern -> Pattern
forall a b. (a -> b) -> a -> b
$ (VName -> Maybe (Subst StructType)) -> PatternType -> PatternType
forall a.
Substitutable a =>
(VName -> Maybe (Subst StructType)) -> a -> a
applySubst (VName -> Map VName (Subst StructType) -> Maybe (Subst StructType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructType)
subs)) [Pattern]
pats,
        valBindBody :: ExpBase Info VName
valBindBody = ExpBase Info VName
body'
      }

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

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

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

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

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

transformDecs :: [Dec] -> MonoM ()
transformDecs :: [Dec] -> MonoM ()
transformDecs [] = () -> MonoM ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
transformDecs (ValDec ValBind
valbind : [Dec]
ds) = do
  Env
env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (TypeDec TypeBind
typebind : [Dec]
ds) = do
  Env
env <- TypeBind -> MonoM Env
transformTypeBind TypeBind
typebind
  Env -> MonoM () -> MonoM ()
forall a. Env -> MonoM a -> MonoM a
localEnv Env
env (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
ds
transformDecs (Dec
dec : [Dec]
_) =
  String -> MonoM ()
forall a. HasCallStack => String -> a
error (String -> MonoM ()) -> String -> MonoM ()
forall a b. (a -> b) -> a -> b
$
    String
"The monomorphization module expects a module-free "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"input program, but received: "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Dec -> String
forall a. Pretty a => a -> String
pretty Dec
dec

-- | Monomorphise a list of top-level declarations. A module-free input program
-- is expected, so only value declarations and type declaration are accepted.
transformProg :: MonadFreshNames m => [Dec] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [Dec] -> m [ValBind]
transformProg [Dec]
decs =
  (((), Seq (VName, ValBind)) -> [ValBind])
-> m ((), Seq (VName, ValBind)) -> m [ValBind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Seq ValBind -> [ValBind]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq ValBind -> [ValBind])
-> (((), Seq (VName, ValBind)) -> Seq ValBind)
-> ((), Seq (VName, ValBind))
-> [ValBind]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, ValBind) -> ValBind)
-> Seq (VName, ValBind) -> Seq ValBind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd (Seq (VName, ValBind) -> Seq ValBind)
-> (((), Seq (VName, ValBind)) -> Seq (VName, ValBind))
-> ((), Seq (VName, ValBind))
-> Seq ValBind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((), Seq (VName, ValBind)) -> Seq (VName, ValBind)
forall a b. (a, b) -> b
snd) (m ((), Seq (VName, ValBind)) -> m [ValBind])
-> m ((), Seq (VName, ValBind)) -> m [ValBind]
forall a b. (a -> b) -> a -> b
$
    (VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
-> m ((), Seq (VName, ValBind))
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
 -> m ((), Seq (VName, ValBind)))
-> (VNameSource -> (((), Seq (VName, ValBind)), VNameSource))
-> m ((), Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
      VNameSource
-> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
forall a.
VNameSource -> MonoM a -> ((a, Seq (VName, ValBind)), VNameSource)
runMonoM VNameSource
namesrc (MonoM () -> (((), Seq (VName, ValBind)), VNameSource))
-> MonoM () -> (((), Seq (VName, ValBind)), VNameSource)
forall a b. (a -> b) -> a -> b
$ [Dec] -> MonoM ()
transformDecs [Dec]
decs