-- | 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 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.
--
-- * Replace all size expressions by constants or variables,
--   complex expressions replaced by variables are calculated in
--   let binding or replaced by size parameters if in argument.
--
-- Note that these changes are unfortunately not visible in the AST
-- representation.
module Futhark.Internalise.Monomorphise (transformProg) where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.RWS (MonadReader (..), MonadWriter (..), RWST, asks, runRWST)
import Control.Monad.State
import Control.Monad.Writer (Writer, runWriter, runWriterT)
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Sequence qualified as Seq
import Data.Set qualified as S
import Futhark.MonadFreshNames
import Futhark.Util (nubOrd, topologicalSort)
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 u. ScalarTypeBase dim u -> TypeBase dim u
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 u. PrimType -> ScalarTypeBase dim u
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.
newtype PolyBinding
  = PolyBinding
      ( VName,
        [TypeParam],
        [Pat ParamType],
        ResRetType,
        Exp,
        [AttrInfo VName],
        SrcLoc
      )

-- | To deduplicate size expressions, we want a looser notation of
-- equality than the strict syntactical equality provided by the Eq
-- instance on Exp.  This newtype wrapper provides such a looser
-- notion of equality.
newtype ReplacedExp = ReplacedExp {ReplacedExp -> Exp
unReplaced :: Exp}
  deriving (Int -> ReplacedExp -> ShowS
[ReplacedExp] -> ShowS
ReplacedExp -> [Char]
(Int -> ReplacedExp -> ShowS)
-> (ReplacedExp -> [Char])
-> ([ReplacedExp] -> ShowS)
-> Show ReplacedExp
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ReplacedExp -> ShowS
showsPrec :: Int -> ReplacedExp -> ShowS
$cshow :: ReplacedExp -> [Char]
show :: ReplacedExp -> [Char]
$cshowList :: [ReplacedExp] -> ShowS
showList :: [ReplacedExp] -> ShowS
Show)

instance Pretty ReplacedExp where
  pretty :: forall ann. ReplacedExp -> Doc ann
pretty (ReplacedExp Exp
e) = Exp -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Exp -> Doc ann
pretty Exp
e

instance Eq ReplacedExp where
  ReplacedExp Exp
e1 == :: ReplacedExp -> ReplacedExp -> Bool
== ReplacedExp Exp
e2
    | Just [(Exp, Exp)]
es <- Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1 Exp
e2 =
        ((Exp, Exp) -> Bool) -> [(Exp, Exp)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((ReplacedExp -> ReplacedExp -> Bool)
-> (ReplacedExp, ReplacedExp) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ReplacedExp -> ReplacedExp -> Bool
forall a. Eq a => a -> a -> Bool
(==) ((ReplacedExp, ReplacedExp) -> Bool)
-> ((Exp, Exp) -> (ReplacedExp, ReplacedExp)) -> (Exp, Exp) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp -> ReplacedExp)
-> (Exp -> ReplacedExp) -> (Exp, Exp) -> (ReplacedExp, ReplacedExp)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Exp -> ReplacedExp
ReplacedExp Exp -> ReplacedExp
ReplacedExp) [(Exp, Exp)]
es
  ReplacedExp
_ == ReplacedExp
_ = Bool
False

type ExpReplacements = [(ReplacedExp, VName)]

canCalculate :: S.Set VName -> ExpReplacements -> ExpReplacements
canCalculate :: Set VName -> ExpReplacements -> ExpReplacements
canCalculate Set VName
scope ExpReplacements
mapping = do
  ((ReplacedExp, VName) -> Bool)
-> ExpReplacements -> ExpReplacements
forall a. (a -> Bool) -> [a] -> [a]
filter
    ( (Set VName -> Set VName -> Bool
forall a. Ord a => Set a -> Set a -> Bool
`S.isSubsetOf` Set VName
scope)
        (Set VName -> Bool)
-> ((ReplacedExp, VName) -> Set VName)
-> (ReplacedExp, VName)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter VName -> Bool
notIntrisic
        (Set VName -> Set VName)
-> ((ReplacedExp, VName) -> Set VName)
-> (ReplacedExp, VName)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars
        (FV -> Set VName)
-> ((ReplacedExp, VName) -> FV)
-> (ReplacedExp, VName)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> FV
freeInExp
        (Exp -> FV)
-> ((ReplacedExp, VName) -> Exp) -> (ReplacedExp, VName) -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReplacedExp -> Exp
unReplaced
        (ReplacedExp -> Exp)
-> ((ReplacedExp, VName) -> ReplacedExp)
-> (ReplacedExp, VName)
-> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ReplacedExp, VName) -> ReplacedExp
forall a b. (a, b) -> a
fst
    )
    ExpReplacements
mapping
  where
    notIntrisic :: VName -> Bool
notIntrisic VName
vn = VName -> Int
baseTag VName
vn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIntrinsicTag

-- Replace some expressions by a parameter.
expReplace :: ExpReplacements -> Exp -> Exp
expReplace :: ExpReplacements -> Exp -> Exp
expReplace ExpReplacements
mapping Exp
e
  | Just VName
vn <- ReplacedExp -> ExpReplacements -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Exp -> ReplacedExp
ReplacedExp Exp
e) ExpReplacements
mapping =
      QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
vn) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e) (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
expReplace ExpReplacements
mapping Exp
e = Identity Exp -> Exp
forall a. Identity a -> a
runIdentity (Identity Exp -> Exp) -> Identity Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ASTMapper Identity -> Exp -> Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper Identity
mapper Exp
e
  where
    mapper :: ASTMapper Identity
mapper = ASTMapper Identity
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp :: Exp -> Identity Exp
mapOnExp = Exp -> Identity Exp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Identity Exp) -> (Exp -> Exp) -> Exp -> Identity Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExpReplacements -> Exp -> Exp
expReplace ExpReplacements
mapping}

-- Construct an Assert expression that checks that the names (values)
-- in the mapping have the same value as the expression they
-- represent.  This is injected into entry points, where we cannot
-- otherwise trust the input.  XXX: the error message generated from
-- this is not great; we should rework it eventually.
entryAssert :: ExpReplacements -> Exp -> Exp
entryAssert :: ExpReplacements -> Exp -> Exp
entryAssert [] Exp
body = Exp
body
entryAssert ((ReplacedExp, VName)
x : ExpReplacements
xs) Exp
body =
  Exp -> Exp -> Info Text -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert ((Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
logAnd ((ReplacedExp, VName) -> Exp
cmpExp (ReplacedExp, VName)
x) ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ ((ReplacedExp, VName) -> Exp) -> ExpReplacements -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (ReplacedExp, VName) -> Exp
cmpExp ExpReplacements
xs) Exp
body Info Text
errmsg (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
body)
  where
    errmsg :: Info Text
errmsg = Text -> Info Text
forall a. a -> Info a
Info Text
"entry point arguments have invalid sizes."
    bool :: TypeBase dim u
bool = ScalarTypeBase dim u -> TypeBase dim u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase dim u -> TypeBase dim u)
-> ScalarTypeBase dim u -> TypeBase dim u
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim u
forall dim u. PrimType -> ScalarTypeBase dim u
Prim PrimType
Bool
    opt :: StructType
opt = [ParamType] -> ResRetType -> StructType
foldFunType [ParamType
forall dim als. TypeBase dim als
bool, ParamType
forall dim als. TypeBase dim als
bool] (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
forall dim als. TypeBase dim als
bool
    andop :: Exp
andop = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName (Name -> VName
intrinsicVar Name
"&&")) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
opt) SrcLoc
forall a. Monoid a => a
mempty
    eqop :: Exp
eqop = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName (Name -> VName
intrinsicVar Name
"==")) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
opt) SrcLoc
forall a. Monoid a => a
mempty
    logAnd :: Exp -> Exp -> Exp
logAnd Exp
x' Exp
y =
      Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
andop [(Diet
Observe, Maybe VName
forall a. Maybe a
Nothing, Exp
x'), (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing, Exp
y)] (AppRes -> Exp) -> AppRes -> Exp
forall a b. (a -> b) -> a -> b
$
        StructType -> [VName] -> AppRes
AppRes StructType
forall dim als. TypeBase dim als
bool []
    cmpExp :: (ReplacedExp, VName) -> Exp
cmpExp (ReplacedExp Exp
x', VName
y) =
      Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
eqop [(Diet
Observe, Maybe VName
forall a. Maybe a
Nothing, Exp
x'), (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing, Exp
y')] (AppRes -> Exp) -> AppRes -> Exp
forall a b. (a -> b) -> a -> b
$
        StructType -> [VName] -> AppRes
AppRes StructType
forall dim als. TypeBase dim als
bool []
      where
        y' :: Exp
y' = QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
y) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
forall dim als. TypeBase dim als
i64) SrcLoc
forall a. Monoid a => a
mempty

-- 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 -> Set VName
envScope :: S.Set VName,
    Env -> Set VName
envGlobalScope :: S.Set VName,
    Env -> ExpReplacements
envParametrized :: ExpReplacements
  }

instance Semigroup Env where
  Env Map VName PolyBinding
tb1 Map VName TypeBinding
pb1 Set VName
sc1 Set VName
gs1 ExpReplacements
pr1 <> :: Env -> Env -> Env
<> Env Map VName PolyBinding
tb2 Map VName TypeBinding
pb2 Set VName
sc2 Set VName
gs2 ExpReplacements
pr2 = Map VName PolyBinding
-> Map VName TypeBinding
-> Set VName
-> Set VName
-> ExpReplacements
-> 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) (Set VName
sc1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
sc2) (Set VName
gs1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
gs2) (ExpReplacements
pr1 ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
<> ExpReplacements
pr2)

instance Monoid Env where
  mempty :: Env
mempty = Map VName PolyBinding
-> Map VName TypeBinding
-> Set VName
-> Set VName
-> ExpReplacements
-> Env
Env Map VName PolyBinding
forall a. Monoid a => a
mempty Map VName TypeBinding
forall a. Monoid a => a
mempty Set VName
forall a. Monoid a => a
mempty Set VName
forall a. Monoid a => a
mempty ExpReplacements
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 a. (Env -> Env) -> MonoM a -> MonoM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (Env
env <>)

extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv :: forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
vn PolyBinding
binding =
  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}

isolateNormalisation :: MonoM a -> MonoM a
isolateNormalisation :: forall a. MonoM a -> MonoM a
isolateNormalisation MonoM a
m = do
  ExpReplacements
prevRepl <- MonoM ExpReplacements
forall s (m :: * -> *). MonadState s m => m s
get
  ExpReplacements -> MonoM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ExpReplacements
forall a. Monoid a => a
mempty
  a
ret <- (Env -> Env) -> MonoM a -> MonoM a
forall a. (Env -> Env) -> MonoM a -> MonoM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\Env
env -> Env
env {envScope :: Set VName
envScope = Set VName
forall a. Monoid a => a
mempty, envParametrized :: ExpReplacements
envParametrized = ExpReplacements
forall a. Monoid a => a
mempty}) MonoM a
m
  ExpReplacements -> MonoM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ExpReplacements
prevRepl
  a -> MonoM a
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
ret

withArgs :: S.Set VName -> MonoM a -> MonoM a
withArgs :: forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
args = Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv (Env -> MonoM a -> MonoM a) -> Env -> MonoM a -> MonoM a
forall a b. (a -> b) -> a -> b
$ Env
forall a. Monoid a => a
mempty {envScope :: Set VName
envScope = Set VName
args}

withParams :: ExpReplacements -> MonoM a -> MonoM a
withParams :: forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
params = Env -> MonoM a -> MonoM a
forall a. Env -> MonoM a -> MonoM a
localEnv (Env -> MonoM a -> MonoM a) -> Env -> MonoM a -> MonoM a
forall a b. (a -> b) -> a -> b
$ Env
forall a. Monoid a => a
mempty {envParametrized :: ExpReplacements
envParametrized = ExpReplacements
params}

-- The monomorphization monad.
newtype MonoM a
  = MonoM
      ( RWST
          Env
          (Seq.Seq (VName, ValBind))
          (ExpReplacements, 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
$cfmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
fmap :: forall a b. (a -> b) -> MonoM a -> MonoM b
$c<$ :: forall a b. a -> MonoM b -> MonoM a
<$ :: forall a b. a -> MonoM b -> MonoM a
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
$cpure :: forall a. a -> MonoM a
pure :: forall a. a -> MonoM a
$c<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
<*> :: forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
$cliftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
liftA2 :: forall a b c. (a -> b -> c) -> MonoM a -> MonoM b -> MonoM c
$c*> :: forall a b. MonoM a -> MonoM b -> MonoM b
*> :: forall a b. MonoM a -> MonoM b -> MonoM b
$c<* :: forall a b. MonoM a -> MonoM b -> MonoM a
<* :: forall a b. MonoM a -> MonoM b -> 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
$c>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
>>= :: forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
$c>> :: forall a b. MonoM a -> MonoM b -> MonoM b
>> :: forall a b. MonoM a -> MonoM b -> MonoM b
$creturn :: forall a. a -> MonoM a
return :: forall a. a -> MonoM a
Monad,
      MonadReader Env,
      MonadWriter (Seq.Seq (VName, ValBind))
    )

instance MonadFreshNames MonoM where
  getNameSource :: MonoM VNameSource
getNameSource = RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  VNameSource
-> MonoM VNameSource
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   VNameSource
 -> MonoM VNameSource)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     VNameSource
-> MonoM VNameSource
forall a b. (a -> b) -> a -> b
$ ((ExpReplacements, VNameSource) -> VNameSource)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (ExpReplacements, VNameSource) -> VNameSource
forall a b. (a, b) -> b
snd
  putNameSource :: VNameSource -> MonoM ()
putNameSource = RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  ()
-> MonoM ()
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   ()
 -> MonoM ())
-> (VNameSource
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ())
-> VNameSource
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
 -> RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ())
-> (VNameSource
    -> (ExpReplacements, VNameSource)
    -> (ExpReplacements, VNameSource))
-> VNameSource
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VNameSource -> VNameSource)
-> (ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((VNameSource -> VNameSource)
 -> (ExpReplacements, VNameSource)
 -> (ExpReplacements, VNameSource))
-> (VNameSource -> VNameSource -> VNameSource)
-> VNameSource
-> (ExpReplacements, VNameSource)
-> (ExpReplacements, VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VNameSource -> VNameSource -> VNameSource
forall a b. a -> b -> a
const

instance MonadState ExpReplacements MonoM where
  get :: MonoM ExpReplacements
get = RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  ExpReplacements
-> MonoM ExpReplacements
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   ExpReplacements
 -> MonoM ExpReplacements)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ExpReplacements
-> MonoM ExpReplacements
forall a b. (a -> b) -> a -> b
$ ((ExpReplacements, VNameSource) -> ExpReplacements)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ExpReplacements
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (ExpReplacements, VNameSource) -> ExpReplacements
forall a b. (a, b) -> a
fst
  put :: ExpReplacements -> MonoM ()
put = RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  ()
-> MonoM ()
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   ()
 -> MonoM ())
-> (ExpReplacements
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ())
-> ExpReplacements
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource))
 -> RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ())
-> (ExpReplacements
    -> (ExpReplacements, VNameSource)
    -> (ExpReplacements, VNameSource))
-> ExpReplacements
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExpReplacements -> ExpReplacements)
-> (ExpReplacements, VNameSource) -> (ExpReplacements, VNameSource)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((ExpReplacements -> ExpReplacements)
 -> (ExpReplacements, VNameSource)
 -> (ExpReplacements, VNameSource))
-> (ExpReplacements -> ExpReplacements -> ExpReplacements)
-> ExpReplacements
-> (ExpReplacements, VNameSource)
-> (ExpReplacements, VNameSource)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExpReplacements -> ExpReplacements -> ExpReplacements
forall a b. a -> b -> a
const

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))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
m) = ((a
a, Seq (VName, ValBind)
defs), VNameSource
src')
  where
    (a
a, (ExpReplacements
_, VNameSource
src'), Seq (VName, ValBind)
defs) = State
  Lifts (a, (ExpReplacements, VNameSource), Seq (VName, ValBind))
-> Lifts
-> (a, (ExpReplacements, VNameSource), Seq (VName, ValBind))
forall s a. State s a -> s -> a
evalState (RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> Env
-> (ExpReplacements, VNameSource)
-> State
     Lifts (a, (ExpReplacements, 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))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
m Env
forall a. Monoid a => a
mempty (ExpReplacements
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 a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (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 a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe PolyBinding
forall a. Maybe a
Nothing

askScope :: MonoM (S.Set VName)
askScope :: MonoM (Set VName)
askScope = do
  Set VName
scope <- (Env -> Set VName) -> MonoM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Set VName
envScope
  Set VName
scope' <- (Env -> Set VName) -> MonoM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Set VName) -> MonoM (Set VName))
-> (Env -> Set VName) -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VName
scope (Set VName -> Set VName) -> (Env -> Set VName) -> Env -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Set VName
envGlobalScope
  Set VName
scope'' <- (Env -> Set VName) -> MonoM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Set VName) -> MonoM (Set VName))
-> (Env -> Set VName) -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VName
scope' (Set VName -> Set VName) -> (Env -> Set VName) -> Env -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName PolyBinding -> Set VName
forall k a. Map k a -> Set k
M.keysSet (Map VName PolyBinding -> Set VName)
-> (Env -> Map VName PolyBinding) -> Env -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName PolyBinding
envPolyBindings
  Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VName
scope'' (Set VName -> Set VName)
-> (Lifts -> Set VName) -> Lifts -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> (Lifts -> [VName]) -> Lifts -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
-> Lifts -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ((VName, InferSizeArgs) -> VName
forall a b. (a, b) -> a
fst ((VName, InferSizeArgs) -> VName)
-> (((VName, MonoType), (VName, InferSizeArgs))
    -> (VName, InferSizeArgs))
-> ((VName, MonoType), (VName, InferSizeArgs))
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, MonoType), (VName, InferSizeArgs))
-> (VName, InferSizeArgs)
forall a b. (a, b) -> b
snd) (Lifts -> Set VName) -> MonoM Lifts -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM Lifts
getLifts

-- | Asks the introduced variables in a set of argument,
-- that is arguments not currently in scope.
askIntros :: S.Set VName -> MonoM (S.Set VName)
askIntros :: Set VName -> MonoM (Set VName)
askIntros Set VName
argset =
  ((VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter VName -> Bool
notIntrisic Set VName
argset `S.difference`) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM (Set VName)
askScope
  where
    notIntrisic :: VName -> Bool
notIntrisic VName
vn = VName -> Int
baseTag VName
vn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIntrinsicTag

-- | Gets and removes expressions that could not be calculated when
-- the arguments set will be unscoped.
-- This should be called without argset in scope, for good detection of intros.
parametrizing :: S.Set VName -> MonoM ExpReplacements
parametrizing :: Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset = do
  Set VName
intros <- Set VName -> MonoM (Set VName)
askIntros Set VName
argset
  (ExpReplacements
params, ExpReplacements
nxtBind) <- (ExpReplacements -> (ExpReplacements, ExpReplacements))
-> MonoM (ExpReplacements, ExpReplacements)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((ExpReplacements -> (ExpReplacements, ExpReplacements))
 -> MonoM (ExpReplacements, ExpReplacements))
-> (ExpReplacements -> (ExpReplacements, ExpReplacements))
-> MonoM (ExpReplacements, ExpReplacements)
forall a b. (a -> b) -> a -> b
$ ((ReplacedExp, VName) -> Bool)
-> ExpReplacements -> (ExpReplacements, ExpReplacements)
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (Bool -> Bool
not (Bool -> Bool)
-> ((ReplacedExp, VName) -> Bool) -> (ReplacedExp, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set VName -> Set VName -> Bool
forall a. Ord a => Set a -> Set a -> Bool
S.disjoint Set VName
intros (Set VName -> Bool)
-> ((ReplacedExp, VName) -> Set VName)
-> (ReplacedExp, VName)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FV -> Set VName
fvVars (FV -> Set VName)
-> ((ReplacedExp, VName) -> FV)
-> (ReplacedExp, VName)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> FV
freeInExp (Exp -> FV)
-> ((ReplacedExp, VName) -> Exp) -> (ReplacedExp, VName) -> FV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReplacedExp -> Exp
unReplaced (ReplacedExp -> Exp)
-> ((ReplacedExp, VName) -> ReplacedExp)
-> (ReplacedExp, VName)
-> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ReplacedExp, VName) -> ReplacedExp
forall a b. (a, b) -> a
fst)
  ExpReplacements -> MonoM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ExpReplacements
nxtBind
  ExpReplacements -> MonoM ExpReplacements
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpReplacements
params

calculateDims :: Exp -> ExpReplacements -> MonoM Exp
calculateDims :: Exp -> ExpReplacements -> MonoM Exp
calculateDims Exp
body ExpReplacements
repl =
  ExpReplacements -> Exp -> MonoM Exp
forall {f :: * -> *}.
MonadFreshNames f =>
ExpReplacements -> Exp -> f Exp
foldCalc ExpReplacements
top_repl (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ ExpReplacements -> Exp -> Exp
expReplace ExpReplacements
top_repl Exp
body
  where
    -- list of strict sub-expressions of e
    subExps :: Exp -> [ReplacedExp]
subExps Exp
e
      | Just Exp
e' <- Exp -> Maybe Exp
stripExp Exp
e = Exp -> [ReplacedExp]
subExps Exp
e'
      | Bool
otherwise = ASTMapper (StateT [ReplacedExp] Identity)
-> Exp -> StateT [ReplacedExp] Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper (StateT [ReplacedExp] Identity)
mapper Exp
e StateT [ReplacedExp] Identity Exp -> [ReplacedExp] -> [ReplacedExp]
forall s a. State s a -> s -> s
`execState` [ReplacedExp]
forall a. Monoid a => a
mempty
      where
        mapOnExp :: Exp -> StateT [ReplacedExp] Identity Exp
mapOnExp Exp
e'
          | Just Exp
e'' <- Exp -> Maybe Exp
stripExp Exp
e' = Exp -> StateT [ReplacedExp] Identity Exp
mapOnExp Exp
e''
          | Bool
otherwise = do
              ([ReplacedExp] -> [ReplacedExp])
-> StateT [ReplacedExp] Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Exp -> ReplacedExp
ReplacedExp Exp
e' :)
              ASTMapper (StateT [ReplacedExp] Identity)
-> Exp -> StateT [ReplacedExp] Identity Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap ASTMapper (StateT [ReplacedExp] Identity)
mapper Exp
e'
        mapper :: ASTMapper (StateT [ReplacedExp] Identity)
mapper = ASTMapper (StateT [ReplacedExp] Identity)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {Exp -> StateT [ReplacedExp] Identity Exp
mapOnExp :: Exp -> StateT [ReplacedExp] Identity Exp
mapOnExp :: Exp -> StateT [ReplacedExp] Identity Exp
mapOnExp}
    depends :: (ReplacedExp, b) -> (ReplacedExp, b) -> Bool
depends (ReplacedExp
a, b
_) (ReplacedExp
b, b
_) = ReplacedExp
b ReplacedExp -> [ReplacedExp] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Exp -> [ReplacedExp]
subExps (ReplacedExp -> Exp
unReplaced ReplacedExp
a)
    top_repl :: ExpReplacements
top_repl =
      ((ReplacedExp, VName) -> (ReplacedExp, VName) -> Bool)
-> ExpReplacements -> ExpReplacements
forall a. (a -> a -> Bool) -> [a] -> [a]
topologicalSort (ReplacedExp, VName) -> (ReplacedExp, VName) -> Bool
forall {b} {b}. (ReplacedExp, b) -> (ReplacedExp, b) -> Bool
depends ExpReplacements
repl

    ---- Calculus insertion
    foldCalc :: ExpReplacements -> Exp -> f Exp
foldCalc [] Exp
body' = Exp -> f Exp
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
body'
    foldCalc ((ReplacedExp
dim, VName
vn) : ExpReplacements
repls) Exp
body' = do
      VName
reName <- VName -> f VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
vn
      let expr :: Exp
expr = ExpReplacements -> Exp -> Exp
expReplace ExpReplacements
repls (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ ReplacedExp -> Exp
unReplaced ReplacedExp
dim
          subst :: VName -> Maybe (Subst t)
subst VName
vn' =
            if VName
vn' VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
vn
              then Subst t -> Maybe (Subst t)
forall a. a -> Maybe a
Just (Subst t -> Maybe (Subst t)) -> Subst t -> Maybe (Subst t)
forall a b. (a -> b) -> a -> b
$ Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> Exp -> Subst t
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
reName) SrcLoc
forall a. Monoid a => a
mempty
              else Maybe (Subst t)
forall a. Maybe a
Nothing
          appRes :: Info AppRes
appRes = case Exp
body' of
            (AppExp AppExpBase Info VName
_ (Info (AppRes StructType
ty [VName]
ext))) -> AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
subst StructType
ty) (VName
reName VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
: [VName]
ext)
            Exp
e -> AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
forall {t}. VName -> Maybe (Subst t)
subst (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e) [VName
reName]
      ExpReplacements -> Exp -> f Exp
foldCalc ExpReplacements
repls (Exp -> f Exp) -> Exp -> f Exp
forall a b. (a -> b) -> a -> b
$
        AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
              []
              (VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
vn (StructType -> Info StructType
forall a. a -> Info a
Info StructType
forall dim als. TypeBase dim als
i64) (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
expr))
              Exp
expr
              Exp
body'
              SrcLoc
forall a. Monoid a => a
mempty
          )
          Info AppRes
appRes

unscoping :: S.Set VName -> Exp -> MonoM Exp
unscoping :: Set VName -> Exp -> MonoM Exp
unscoping Set VName
argset Exp
body = do
  ExpReplacements
localDims <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset
  Set VName
scope <- Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VName
argset (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM (Set VName)
askScope
  Exp -> ExpReplacements -> MonoM Exp
calculateDims Exp
body (ExpReplacements -> MonoM Exp) -> ExpReplacements -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> ExpReplacements -> ExpReplacements
canCalculate Set VName
scope ExpReplacements
localDims

scoping :: S.Set VName -> MonoM Exp -> MonoM Exp
scoping :: Set VName -> MonoM Exp -> MonoM Exp
scoping Set VName
argset MonoM Exp
m =
  Set VName -> MonoM Exp -> MonoM Exp
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
argset MonoM Exp
m MonoM Exp -> (Exp -> MonoM Exp) -> MonoM Exp
forall a b. MonoM a -> (a -> MonoM b) -> MonoM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Set VName -> Exp -> MonoM Exp
unscoping Set VName
argset

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

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

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

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

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

monoType :: TypeBase Size als -> MonoType
monoType :: forall als. TypeBase Exp als -> MonoType
monoType = MonoType -> MonoType
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts (MonoType -> MonoType)
-> (TypeBase Exp als -> MonoType) -> TypeBase Exp als -> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (State (Int, Map Exp Int) MonoType -> (Int, Map Exp Int) -> MonoType
forall s a. State s a -> s -> a
`evalState` (Int
0, Map Exp Int
forall a. Monoid a => a
mempty)) (State (Int, Map Exp Int) MonoType -> MonoType)
-> (TypeBase Exp als -> State (Int, Map Exp Int) MonoType)
-> TypeBase Exp als
-> MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Set VName
 -> DimPos -> Exp -> StateT (Int, Map Exp Int) Identity MonoSize)
-> StructType -> State (Int, Map Exp 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 -> Exp -> StateT (Int, Map Exp Int) Identity MonoSize
forall {f :: * -> *} {p}.
MonadState (Int, Map Exp Int) f =>
Set VName -> p -> Exp -> f MonoSize
onDim (StructType -> State (Int, Map Exp Int) MonoType)
-> (TypeBase Exp als -> StructType)
-> TypeBase Exp als
-> State (Int, Map Exp Int) MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Exp als -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct
  where
    -- Remove exts from return types because we don't use them anymore.
    noExts :: TypeBase MonoSize u -> TypeBase MonoSize u
    noExts :: forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts (Array u
u Shape MonoSize
shape ScalarTypeBase MonoSize NoUniqueness
t) = u
-> Shape MonoSize
-> ScalarTypeBase MonoSize NoUniqueness
-> TypeBase MonoSize u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u Shape MonoSize
shape (ScalarTypeBase MonoSize NoUniqueness -> TypeBase MonoSize u)
-> ScalarTypeBase MonoSize NoUniqueness -> TypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize NoUniqueness
-> ScalarTypeBase MonoSize NoUniqueness
forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
noExtsScalar ScalarTypeBase MonoSize NoUniqueness
t
    noExts (Scalar ScalarTypeBase MonoSize u
t) = ScalarTypeBase MonoSize u -> TypeBase MonoSize u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase MonoSize u -> TypeBase MonoSize u)
-> ScalarTypeBase MonoSize u -> TypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
noExtsScalar ScalarTypeBase MonoSize u
t
    noExtsScalar :: ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
noExtsScalar (Record Map Name (TypeBase MonoSize u)
fs) = Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u)
-> Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ (TypeBase MonoSize u -> TypeBase MonoSize u)
-> Map Name (TypeBase MonoSize u) -> Map Name (TypeBase MonoSize u)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase MonoSize u -> TypeBase MonoSize u
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts Map Name (TypeBase MonoSize u)
fs
    noExtsScalar (Sum Map Name [TypeBase MonoSize u]
fs) = Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u)
-> Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ ([TypeBase MonoSize u] -> [TypeBase MonoSize u])
-> Map Name [TypeBase MonoSize u] -> Map Name [TypeBase MonoSize u]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase MonoSize u -> TypeBase MonoSize u)
-> [TypeBase MonoSize u] -> [TypeBase MonoSize u]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase MonoSize u -> TypeBase MonoSize u
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts) Map Name [TypeBase MonoSize u]
fs
    noExtsScalar (Arrow u
as PName
p Diet
d MonoType
t1 (RetType [VName]
_ TypeBase MonoSize Uniqueness
t2)) =
      u
-> PName
-> Diet
-> MonoType
-> RetTypeBase MonoSize Uniqueness
-> ScalarTypeBase MonoSize u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
as PName
p Diet
d (MonoType -> MonoType
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts MonoType
t1) ([VName]
-> TypeBase MonoSize Uniqueness -> RetTypeBase MonoSize Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase MonoSize Uniqueness -> TypeBase MonoSize Uniqueness
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
noExts TypeBase MonoSize Uniqueness
t2))
    noExtsScalar ScalarTypeBase MonoSize u
t = ScalarTypeBase MonoSize u
t
    onDim :: Set VName -> p -> Exp -> f MonoSize
onDim Set VName
bound p
_ Exp
e
      -- A locally bound size.
      | (VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) (Set VName -> Bool) -> Set VName -> Bool
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ Exp -> FV
freeInExp Exp
e =
          MonoSize -> f MonoSize
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MonoSize
MonoAnon
    onDim Set VName
_ p
_ Exp
d = do
      (Int
i, Map Exp Int
m) <- f (Int, Map Exp Int)
forall s (m :: * -> *). MonadState s m => m s
get
      case Exp -> Map Exp Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Exp
d Map Exp Int
m of
        Just Int
prev ->
          MonoSize -> f MonoSize
forall a. a -> f a
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 Exp 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, Exp -> Int -> Map Exp Int -> Map Exp Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Exp
d Int
i Map Exp Int
m)
          MonoSize -> f MonoSize
forall a. a -> f a
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))
  (ExpReplacements, VNameSource)
  (State Lifts)
  Lifts
-> MonoM Lifts
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   Lifts
 -> MonoM Lifts)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     Lifts
-> MonoM Lifts
forall a b. (a -> b) -> a -> b
$ State Lifts Lifts
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     Lifts
forall (m :: * -> *) a.
Monad m =>
m a
-> RWST
     Env (Seq (VName, ValBind)) (ExpReplacements, VNameSource) m a
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))
  (ExpReplacements, VNameSource)
  (State Lifts)
  ()
-> MonoM ()
forall a.
RWST
  Env
  (Seq (VName, ValBind))
  (ExpReplacements, VNameSource)
  (State Lifts)
  a
-> MonoM a
MonoM (RWST
   Env
   (Seq (VName, ValBind))
   (ExpReplacements, VNameSource)
   (State Lifts)
   ()
 -> MonoM ())
-> ((Lifts -> Lifts)
    -> RWST
         Env
         (Seq (VName, ValBind))
         (ExpReplacements, VNameSource)
         (State Lifts)
         ())
-> (Lifts -> Lifts)
-> MonoM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. State Lifts ()
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, VNameSource)
     (State Lifts)
     ()
forall (m :: * -> *) a.
Monad m =>
m a
-> RWST
     Env (Seq (VName, ValBind)) (ExpReplacements, VNameSource) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (State Lifts ()
 -> RWST
      Env
      (Seq (VName, ValBind))
      (ExpReplacements, VNameSource)
      (State Lifts)
      ())
-> ((Lifts -> Lifts) -> State Lifts ())
-> (Lifts -> Lifts)
-> RWST
     Env
     (Seq (VName, ValBind))
     (ExpReplacements, 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) :)

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

-- | Creates a new expression replacement if needed, this always produces normalised sizes.
-- (e.g. single variable or constant)
replaceExp :: Exp -> MonoM Exp
replaceExp :: Exp -> MonoM Exp
replaceExp Exp
e =
  case Exp -> Maybe Exp
maybeNormalisedSize Exp
e of
    Just Exp
e' -> Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e'
    Maybe Exp
Nothing -> do
      let e' :: ReplacedExp
e' = Exp -> ReplacedExp
ReplacedExp Exp
e
      Maybe VName
prev <- (ExpReplacements -> Maybe VName) -> MonoM (Maybe VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((ExpReplacements -> Maybe VName) -> MonoM (Maybe VName))
-> (ExpReplacements -> Maybe VName) -> MonoM (Maybe VName)
forall a b. (a -> b) -> a -> b
$ ReplacedExp -> ExpReplacements -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ReplacedExp
e'
      Maybe VName
prev_param <- (Env -> Maybe VName) -> MonoM (Maybe VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Maybe VName) -> MonoM (Maybe VName))
-> (Env -> Maybe VName) -> MonoM (Maybe VName)
forall a b. (a -> b) -> a -> b
$ ReplacedExp -> ExpReplacements -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ReplacedExp
e' (ExpReplacements -> Maybe VName)
-> (Env -> ExpReplacements) -> Env -> Maybe VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> ExpReplacements
envParametrized
      case (Maybe VName
prev_param, Maybe VName
prev) of
        (Just VName
vn, Maybe VName
_) -> Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
vn) (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
        (Maybe VName
Nothing, Just VName
vn) -> Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
vn) (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
        (Maybe VName
Nothing, Maybe VName
Nothing) -> do
          VName
vn <- [Char] -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString ([Char] -> MonoM VName) -> [Char] -> MonoM VName
forall a b. (a -> b) -> a -> b
$ [Char]
"d<{" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ExpBase NoInfo VName -> [Char]
forall a. Pretty a => a -> [Char]
prettyString (Exp -> ExpBase NoInfo VName
bareExp Exp
e) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"}>"
          (ExpReplacements -> ExpReplacements) -> MonoM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((ReplacedExp
e', VName
vn) :)
          Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
vn) (Exp -> SrcLoc
forall a. Located a => a -> SrcLoc
srclocOf Exp
e)
  where
    -- Avoid replacing of some 'already normalised' sizes that are just surounded by some parentheses.
    maybeNormalisedSize :: Exp -> Maybe Exp
maybeNormalisedSize Exp
e'
      | Just Exp
e'' <- Exp -> Maybe Exp
stripExp Exp
e' = Exp -> Maybe Exp
maybeNormalisedSize Exp
e''
    maybeNormalisedSize (Var QualName VName
qn Info StructType
_ SrcLoc
loc) = Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName QualName VName
qn SrcLoc
loc
    maybeNormalisedSize (IntLit Integer
v Info StructType
_ SrcLoc
loc) = Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Maybe Exp) -> Exp -> Maybe Exp
forall a b. (a -> b) -> a -> b
$ Integer -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Integer -> f StructType -> SrcLoc -> ExpBase f vn
IntLit Integer
v (StructType -> Info StructType
forall a. a -> Info a
Info StructType
forall dim als. TypeBase dim als
i64) SrcLoc
loc
    maybeNormalisedSize Exp
_ = Maybe Exp
forall a. Maybe a
Nothing

transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname StructType
t = do
  StructType
t' <- StructType -> MonoM StructType
removeTypeVariablesInType StructType
t
  StructType
t'' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
t'
  let mono_t :: MonoType
mono_t = StructType -> MonoType
forall als. TypeBase Exp als -> MonoType
monoType StructType
t'
  if 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
    then Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> StructType -> Exp
forall {vn}. QualName vn -> StructType -> ExpBase Info vn
var QualName VName
fname StructType
t''
    else do
      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
_) ->
          VName -> TypeBase Exp Uniqueness -> [Exp] -> Exp
forall {vn}.
vn
-> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs VName
fname' (Uniqueness -> StructType -> TypeBase Exp Uniqueness
forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t'') ([Exp] -> Exp) -> MonoM [Exp] -> MonoM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InferSizeArgs
infer StructType
t''
        -- An intrinsic function.
        (Maybe (VName, InferSizeArgs)
Nothing, Maybe PolyBinding
Nothing) -> Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> StructType -> Exp
forall {vn}. QualName vn -> StructType -> ExpBase Info vn
var QualName VName
fname StructType
t''
        -- 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)
          VName -> TypeBase Exp Uniqueness -> [Exp] -> Exp
forall {vn}.
vn
-> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs VName
fname' (Uniqueness -> StructType -> TypeBase Exp Uniqueness
forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t'') ([Exp] -> Exp) -> MonoM [Exp] -> MonoM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> InferSizeArgs
infer StructType
t''
  where
    var :: QualName vn -> StructType -> ExpBase Info vn
var QualName vn
fname' StructType
t'' = QualName vn -> Info StructType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName vn
fname' (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t'') SrcLoc
loc

    applySizeArg :: TypeBase Exp Uniqueness
-> (Int, ExpBase Info vn)
-> ExpBase Info vn
-> (Int, ExpBase Info vn)
applySizeArg TypeBase Exp Uniqueness
t' (Int
i, ExpBase Info vn
f) ExpBase Info vn
size_arg =
      ( Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1,
        ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
          ExpBase Info vn
f
          [(Diet
Observe, Maybe VName
forall a. Maybe a
Nothing, ExpBase Info vn
size_arg)]
          (StructType -> [VName] -> AppRes
AppRes ([ParamType] -> ResRetType -> StructType
foldFunType (Int -> ParamType -> [ParamType]
forall a. Int -> a -> [a]
replicate Int
i ParamType
forall dim als. TypeBase dim als
i64) ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
t')) [])
      )

    applySizeArgs :: vn
-> TypeBase Exp Uniqueness -> [ExpBase Info vn] -> ExpBase Info vn
applySizeArgs vn
fname' TypeBase Exp Uniqueness
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 b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          (TypeBase Exp Uniqueness
-> (Int, ExpBase Info vn)
-> ExpBase Info vn
-> (Int, ExpBase Info vn)
forall {vn}.
TypeBase Exp Uniqueness
-> (Int, ExpBase Info vn)
-> ExpBase Info vn
-> (Int, ExpBase Info vn)
applySizeArg TypeBase Exp Uniqueness
t')
          ( [ExpBase Info vn] -> Int
forall a. [a] -> 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 StructType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var
              (vn -> QualName vn
forall v. v -> QualName v
qualName vn
fname')
              (StructType -> Info StructType
forall a. a -> Info a
Info ([ParamType] -> ResRetType -> StructType
foldFunType ((ExpBase Info vn -> ParamType) -> [ExpBase Info vn] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (ParamType -> ExpBase Info vn -> ParamType
forall a b. a -> b -> a
const ParamType
forall dim als. TypeBase dim als
i64) [ExpBase Info vn]
size_args) ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] TypeBase Exp Uniqueness
t')))
              SrcLoc
loc
          )
          [ExpBase Info vn]
size_args

transformType :: TypeBase Size u -> MonoM (TypeBase Size u)
transformType :: forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType TypeBase Exp u
typ =
  case TypeBase Exp u
typ of
    Scalar ScalarTypeBase Exp u
scalar -> ScalarTypeBase Exp u -> TypeBase Exp u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp u -> TypeBase Exp u)
-> MonoM (ScalarTypeBase Exp u) -> MonoM (TypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
transformScalarSizes ScalarTypeBase Exp u
scalar
    Array u
u Shape Exp
shape ScalarTypeBase Exp NoUniqueness
scalar -> u -> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u (Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
-> MonoM (Shape Exp)
-> MonoM (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> Shape Exp -> MonoM (Shape Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Shape a -> m (Shape b)
mapM Exp -> MonoM Exp
onDim Shape Exp
shape MonoM (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp u)
-> MonoM (ScalarTypeBase Exp NoUniqueness)
-> MonoM (TypeBase Exp u)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ScalarTypeBase Exp NoUniqueness
-> MonoM (ScalarTypeBase Exp NoUniqueness)
forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
transformScalarSizes ScalarTypeBase Exp NoUniqueness
scalar
  where
    transformScalarSizes :: ScalarTypeBase Size u -> MonoM (ScalarTypeBase Size u)
    transformScalarSizes :: forall u. ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
transformScalarSizes (Record Map Name (TypeBase Exp u)
fs) =
      Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u)
-> MonoM (Map Name (TypeBase Exp u))
-> MonoM (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp u -> MonoM (TypeBase Exp u))
-> Map Name (TypeBase Exp u) -> MonoM (Map Name (TypeBase Exp u))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse TypeBase Exp u -> MonoM (TypeBase Exp u)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Map Name (TypeBase Exp u)
fs
    transformScalarSizes (Sum Map Name [TypeBase Exp u]
cs) =
      Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u)
-> MonoM (Map Name [TypeBase Exp u])
-> MonoM (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase Exp u] -> MonoM [TypeBase Exp u])
-> Map Name [TypeBase Exp u] -> MonoM (Map Name [TypeBase Exp u])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (([TypeBase Exp u] -> MonoM [TypeBase Exp u])
 -> Map Name [TypeBase Exp u] -> MonoM (Map Name [TypeBase Exp u]))
-> ((TypeBase Exp u -> MonoM (TypeBase Exp u))
    -> [TypeBase Exp u] -> MonoM [TypeBase Exp u])
-> (TypeBase Exp u -> MonoM (TypeBase Exp u))
-> Map Name [TypeBase Exp u]
-> MonoM (Map Name [TypeBase Exp u])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp u -> MonoM (TypeBase Exp u))
-> [TypeBase Exp u] -> MonoM [TypeBase Exp u]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse) TypeBase Exp u -> MonoM (TypeBase Exp u)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Map Name [TypeBase Exp u]
cs
    transformScalarSizes (Arrow u
as PName
argName Diet
d StructType
argT ResRetType
retT) = do
      ResRetType
retT' <- Set VName -> ResRetType -> MonoM ResRetType
forall as.
Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
transformRetTypeSizes Set VName
argset ResRetType
retT
      u
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
as PName
argName Diet
d (StructType -> ResRetType -> ScalarTypeBase Exp u)
-> MonoM StructType -> MonoM (ResRetType -> ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
argT MonoM (ResRetType -> ScalarTypeBase Exp u)
-> MonoM ResRetType -> MonoM (ScalarTypeBase Exp u)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ResRetType -> MonoM ResRetType
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ResRetType
retT'
      where
        argset :: Set VName
argset =
          FV -> Set VName
fvVars (StructType -> FV
forall u. TypeBase Exp u -> FV
freeInType StructType
argT)
            Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> case PName
argName of
              PName
Unnamed -> Set VName
forall a. Monoid a => a
mempty
              Named VName
vn -> VName -> Set VName
forall a. a -> Set a
S.singleton VName
vn
    transformScalarSizes (TypeVar u
u QualName VName
qn [TypeArg Exp]
args) =
      u -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp u)
-> MonoM [TypeArg Exp] -> MonoM (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp -> MonoM (TypeArg Exp))
-> [TypeArg Exp] -> MonoM [TypeArg Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeArg Exp -> MonoM (TypeArg Exp)
onArg [TypeArg Exp]
args
      where
        onArg :: TypeArg Exp -> MonoM (TypeArg Exp)
onArg (TypeArgDim Exp
dim) = Exp -> TypeArg Exp
forall dim. dim -> TypeArg dim
TypeArgDim (Exp -> TypeArg Exp) -> MonoM Exp -> MonoM (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
onDim Exp
dim
        onArg (TypeArgType StructType
ty) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp)
-> MonoM StructType -> MonoM (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
ty
    transformScalarSizes ty :: ScalarTypeBase Exp u
ty@Prim {} = ScalarTypeBase Exp u -> MonoM (ScalarTypeBase Exp u)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase Exp u
ty

    onDim :: Exp -> MonoM Exp
onDim Exp
e
      | Exp
e Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
anySize = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      | Bool
otherwise = Exp -> MonoM Exp
replaceExp (Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Exp -> MonoM Exp
transformExp Exp
e

transformRetTypeSizes :: S.Set VName -> RetTypeBase Size as -> MonoM (RetTypeBase Size as)
transformRetTypeSizes :: forall as.
Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
transformRetTypeSizes Set VName
argset (RetType [VName]
dims TypeBase Exp as
ty) = do
  TypeBase Exp as
ty' <- Set VName -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
argset (MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as))
-> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
forall a b. (a -> b) -> a -> b
$ TypeBase Exp as -> MonoM (TypeBase Exp as)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType TypeBase Exp as
ty
  ExpReplacements
rl <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset
  let dims' :: [VName]
dims' = [VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (ReplacedExp, VName) -> VName
forall a b. (a, b) -> b
snd ExpReplacements
rl
  RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Exp as -> MonoM (RetTypeBase Exp as))
-> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp as -> RetTypeBase Exp as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims' TypeBase Exp as
ty'

sizesForPat :: (MonadFreshNames m) => Pat ParamType -> m ([VName], Pat ParamType)
sizesForPat :: forall (m :: * -> *).
MonadFreshNames m =>
Pat ParamType -> m ([VName], Pat ParamType)
sizesForPat Pat ParamType
pat = do
  (Pat ParamType
params', [VName]
sizes) <- StateT [VName] m (Pat ParamType)
-> [VName] -> m (Pat ParamType, [VName])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((ParamType -> StateT [VName] m ParamType)
-> Pat ParamType -> StateT [VName] m (Pat ParamType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PatBase Info VName a -> f (PatBase Info VName b)
traverse ((Exp -> StateT [VName] m Exp)
-> (Diet -> StateT [VName] m Diet)
-> ParamType
-> StateT [VName] m ParamType
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TypeBase a b -> f (TypeBase c 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 Exp -> StateT [VName] m Exp
forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadTrans t, MonadFreshNames m, MonadState [VName] (t m)) =>
Exp -> t m Exp
onDim Diet -> StateT [VName] m Diet
forall a. a -> StateT [VName] m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) Pat ParamType
pat) []
  ([VName], Pat ParamType) -> m ([VName], Pat ParamType)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName]
sizes, Pat ParamType
params')
  where
    onDim :: Exp -> t m Exp
onDim Exp
d
      | Exp
d Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
anySize = do
          VName
v <- m VName -> t m VName
forall (m :: * -> *) a. Monad m => m a -> t m a
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
$ [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"size"
          ([VName] -> [VName]) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName
v :)
          Exp -> t m Exp
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> t m Exp) -> Exp -> t m Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) SrcLoc
forall a. Monoid a => a
mempty
      | Bool
otherwise = Exp -> t m Exp
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d

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

transformAppExp :: AppExp -> AppRes -> MonoM Exp
transformAppExp :: AppExpBase Info VName -> AppRes -> MonoM Exp
transformAppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) AppRes
res = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Maybe Exp
me' <- (Exp -> MonoM Exp) -> Maybe Exp -> MonoM (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Maybe a -> m (Maybe b)
mapM Exp -> MonoM Exp
transformExp Maybe Exp
me
  Inclusiveness Exp
incl' <- (Exp -> MonoM Exp)
-> Inclusiveness Exp -> MonoM (Inclusiveness Exp)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Inclusiveness a -> m (Inclusiveness b)
mapM Exp -> MonoM Exp
transformExp Inclusiveness Exp
incl
  AppRes
res' <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> Maybe Exp
-> Inclusiveness Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Range Exp
e1' Maybe Exp
me' Inclusiveness Exp
incl' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')
transformAppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
loc) AppRes
res = do
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  let dimArgs :: Set VName
dimArgs = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((SizeBinder VName -> VName) -> [SizeBinder VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map SizeBinder VName -> VName
forall vn. SizeBinder vn -> vn
sizeName [SizeBinder VName]
sizes)
  Set VName
implicitDims <- Set VName -> MonoM (Set VName) -> MonoM (Set VName)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
dimArgs (MonoM (Set VName) -> MonoM (Set VName))
-> MonoM (Set VName) -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM (Set VName)
askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ PatBase Info VName StructType -> FV
forall u. Pat (TypeBase Exp u) -> FV
freeInPat PatBase Info VName StructType
pat
  let dimArgs' :: Set VName
dimArgs' = Set VName
dimArgs Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
implicitDims
      letArgs :: Set VName
letArgs = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
pat
      argset :: Set VName
argset = Set VName
dimArgs' Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
letArgs
  PatBase Info VName StructType
pat' <- Set VName
-> MonoM (PatBase Info VName StructType)
-> MonoM (PatBase Info VName StructType)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
dimArgs' (MonoM (PatBase Info VName StructType)
 -> MonoM (PatBase Info VName StructType))
-> MonoM (PatBase Info VName StructType)
-> MonoM (PatBase Info VName StructType)
forall a b. (a -> b) -> a -> b
$ PatBase Info VName StructType
-> MonoM (PatBase Info VName StructType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat PatBase Info VName StructType
pat
  ExpReplacements
params <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
dimArgs'
  let sizes' :: [SizeBinder VName]
sizes' = [SizeBinder VName]
sizes [SizeBinder VName] -> [SizeBinder VName] -> [SizeBinder VName]
forall a. Semigroup a => a -> a -> a
<> (VName -> SizeBinder VName) -> [VName] -> [SizeBinder VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> SizeBinder VName
forall vn. vn -> SrcLoc -> SizeBinder vn
`SizeBinder` SrcLoc
forall a. Monoid a => a
mempty) (((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (ReplacedExp, VName) -> VName
forall a b. (a, b) -> b
snd ExpReplacements
params [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
implicitDims)
  Exp
body' <- ExpReplacements -> MonoM Exp -> MonoM Exp
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM Exp -> MonoM Exp
scoping Set VName
argset (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body
  AppRes
res' <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [SizeBinder VName]
sizes' PatBase Info VName StructType
pat' Exp
e' Exp
body' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')
transformAppExp (LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat ParamType]
params, Maybe (TypeExp Info VName)
retdecl, Info ResRetType
ret, Exp
body) Exp
e SrcLoc
loc) AppRes
res
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [TypeParamBase VName] -> Bool
forall a. [a] -> 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.
      let funbind :: PolyBinding
funbind = (VName, [TypeParamBase VName], [Pat ParamType], ResRetType, Exp,
 [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding (VName
fname, [TypeParamBase VName]
tparams, [Pat ParamType]
params, ResRetType
ret, Exp
body, [AttrInfo VName]
forall a. Monoid a => a
mempty, SrcLoc
loc)
      MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall a.
MonoM (a, Seq (VName, ValBind) -> Seq (VName, ValBind)) -> MonoM a
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
 -> MonoM Exp)
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM Exp
forall a b. (a -> b) -> a -> b
$ do
        (Exp
e', Seq (VName, ValBind)
bs) <- MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall a. MonoM a -> MonoM (a, Seq (VName, ValBind))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (MonoM Exp -> MonoM (Exp, Seq (VName, ValBind)))
-> MonoM Exp -> MonoM (Exp, Seq (VName, ValBind))
forall a b. (a -> b) -> a -> b
$ VName -> PolyBinding -> MonoM Exp -> MonoM Exp
forall a. VName -> PolyBinding -> MonoM a -> MonoM a
extendEnv VName
fname PolyBinding
funbind (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM Exp -> MonoM Exp
scoping (VName -> Set VName
forall a. a -> Set a
S.singleton VName
fname) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
e
        -- Do not remember this one for next time we monomorphise this
        -- function.
        (Lifts -> Lifts) -> MonoM ()
modifyLifts ((Lifts -> Lifts) -> MonoM ()) -> (Lifts -> Lifts) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ (((VName, MonoType), (VName, InferSizeArgs)) -> Bool)
-> Lifts -> Lifts
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
/= VName
fname) (VName -> Bool)
-> (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
-> ((VName, MonoType), (VName, InferSizeArgs))
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, MonoType) -> VName
forall a b. (a, b) -> a
fst ((VName, MonoType) -> VName)
-> (((VName, MonoType), (VName, InferSizeArgs))
    -> (VName, MonoType))
-> ((VName, MonoType), (VName, InferSizeArgs))
-> VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, MonoType), (VName, InferSizeArgs)) -> (VName, MonoType)
forall a b. (a, b) -> a
fst)
        let (Seq (VName, ValBind)
bs_local, Seq (VName, ValBind)
bs_prop) = ((VName, ValBind) -> Bool)
-> Seq (VName, ValBind)
-> (Seq (VName, ValBind), Seq (VName, ValBind))
forall a. (a -> Bool) -> Seq a -> (Seq a, Seq a)
Seq.partition ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
fname) (VName -> Bool)
-> ((VName, ValBind) -> VName) -> (VName, ValBind) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, ValBind) -> VName
forall a b. (a, b) -> a
fst) Seq (VName, ValBind)
bs
        (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
-> MonoM (Exp, Seq (VName, ValBind) -> Seq (VName, ValBind))
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ValBind] -> Exp -> Exp
unfoldLetFuns (((VName, ValBind) -> ValBind) -> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> [a] -> [b]
map (VName, ValBind) -> ValBind
forall a b. (a, b) -> b
snd ([(VName, ValBind)] -> [ValBind])
-> [(VName, ValBind)] -> [ValBind]
forall a b. (a -> b) -> a -> b
$ Seq (VName, ValBind) -> [(VName, ValBind)]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq (VName, ValBind)
bs_local) Exp
e', Seq (VName, ValBind)
-> Seq (VName, ValBind) -> Seq (VName, ValBind)
forall a b. a -> b -> a
const Seq (VName, ValBind)
bs_prop)
  | Bool
otherwise = do
      Exp
body' <- Set VName -> MonoM Exp -> MonoM Exp
scoping ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body
      ResRetType
ret' <- Set VName -> ResRetType -> MonoM ResRetType
forall as.
Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
transformRetTypeSizes ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)) ResRetType
ret
      AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
        (AppExpBase Info VName -> Info AppRes -> Exp)
-> MonoM (AppExpBase Info VName) -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( VName
-> ([TypeParamBase VName], [Pat ParamType],
    Maybe (TypeExp Info VName), Info ResRetType, Exp)
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
vn
-> ([TypeParamBase vn], [PatBase f vn ParamType],
    Maybe (TypeExp f vn), f ResRetType, ExpBase f vn)
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetFun VName
fname ([TypeParamBase VName]
tparams, [Pat ParamType]
params, Maybe (TypeExp Info VName)
retdecl, ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
ret', Exp
body')
                (Exp -> SrcLoc -> AppExpBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> AppExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> MonoM Exp -> MonoM Exp
scoping (VName -> Set VName
forall a. a -> Set a
S.singleton VName
fname) (Exp -> MonoM Exp
transformExp Exp
e)
                MonoM (SrcLoc -> AppExpBase Info VName)
-> MonoM SrcLoc -> MonoM (AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
            )
        MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> MonoM AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AppRes -> MonoM AppRes
transformAppRes AppRes
res)
transformAppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) AppRes
res =
  AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (AppExpBase Info VName -> Info AppRes -> Exp)
-> MonoM (AppExpBase Info VName) -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn -> ExpBase f vn -> SrcLoc -> AppExpBase f vn
If (Exp -> Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
-> MonoM Exp
-> MonoM (Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Exp -> SrcLoc -> AppExpBase Info VName)
-> MonoM Exp -> MonoM (Exp -> SrcLoc -> AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Exp -> SrcLoc -> AppExpBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e3 MonoM (SrcLoc -> AppExpBase Info VName)
-> MonoM SrcLoc -> MonoM (AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc) MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> MonoM AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AppRes -> MonoM AppRes
transformAppRes AppRes
res)
transformAppExp (Apply Exp
fe NonEmpty (Info (Diet, Maybe VName), Exp)
args SrcLoc
_) AppRes
res =
  Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
    (Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp)
-> MonoM Exp -> MonoM ([(Diet, Maybe VName, Exp)] -> AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
fe
    MonoM ([(Diet, Maybe VName, Exp)] -> AppRes -> Exp)
-> MonoM [(Diet, Maybe VName, Exp)] -> MonoM (AppRes -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((Info (Diet, Maybe VName), Exp) -> MonoM (Diet, Maybe VName, Exp))
-> [(Info (Diet, Maybe VName), Exp)]
-> MonoM [(Diet, Maybe VName, Exp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Info (Diet, Maybe VName), Exp) -> MonoM (Diet, Maybe VName, Exp)
forall {t} {t}. (Info (t, t), Exp) -> MonoM (t, t, Exp)
onArg (NonEmpty (Info (Diet, Maybe VName), Exp)
-> [(Info (Diet, Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (Diet, Maybe VName), Exp)
args)
    MonoM (AppRes -> Exp) -> MonoM AppRes -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> AppRes -> MonoM AppRes
transformAppRes AppRes
res
  where
    onArg :: (Info (t, t), Exp) -> MonoM (t, t, Exp)
onArg (Info (t
d, t
ext), Exp
e) = (t
d,t
ext,) (Exp -> (t, t, Exp)) -> MonoM Exp -> MonoM (t, t, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e
transformAppExp (Loop [VName]
sparams Pat ParamType
pat Exp
e1 LoopFormBase Info VName
form Exp
body SrcLoc
loc) AppRes
res = do
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1

  let dimArgs :: Set VName
dimArgs = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams
  Pat ParamType
pat' <- Set VName -> MonoM (Pat ParamType) -> MonoM (Pat ParamType)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
dimArgs (MonoM (Pat ParamType) -> MonoM (Pat ParamType))
-> MonoM (Pat ParamType) -> MonoM (Pat ParamType)
forall a b. (a -> b) -> a -> b
$ Pat ParamType -> MonoM (Pat ParamType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat Pat ParamType
pat
  ExpReplacements
params <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
dimArgs
  let sparams' :: [VName]
sparams' = [VName]
sparams [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (ReplacedExp, VName) -> VName
forall a b. (a, b) -> b
snd ExpReplacements
params
      mergeArgs :: Set VName
mergeArgs = Set VName
dimArgs Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames Pat ParamType
pat)

  (LoopFormBase Info VName
form', Set VName
formArgs) <- case LoopFormBase Info VName
form of
    For IdentBase Info VName StructType
ident Exp
e2 -> (,VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
ident) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
-> (Exp -> LoopFormBase Info VName)
-> Exp
-> (LoopFormBase Info VName, Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase Info VName StructType -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName StructType
ident (Exp -> (LoopFormBase Info VName, Set VName))
-> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
    ForIn PatBase Info VName StructType
pat2 Exp
e2 -> do
      PatBase Info VName StructType
pat2' <- PatBase Info VName StructType
-> MonoM (PatBase Info VName StructType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat PatBase Info VName StructType
pat2
      (,[VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
pat2)) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
-> (Exp -> LoopFormBase Info VName)
-> Exp
-> (LoopFormBase Info VName, Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatBase Info VName StructType -> Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> LoopFormBase f vn
ForIn PatBase Info VName StructType
pat2' (Exp -> (LoopFormBase Info VName, Set VName))
-> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e2
    While Exp
e2 ->
      (Exp -> (LoopFormBase Info VName, Set VName))
-> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
forall a b. (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,Set VName
forall a. Monoid a => a
mempty) (LoopFormBase Info VName -> (LoopFormBase Info VName, Set VName))
-> (Exp -> LoopFormBase Info VName)
-> Exp
-> (LoopFormBase Info VName, Set VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While) (MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName))
-> MonoM Exp -> MonoM (LoopFormBase Info VName, Set VName)
forall a b. (a -> b) -> a -> b
$
        ExpReplacements -> MonoM Exp -> MonoM Exp
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
          Set VName -> MonoM Exp -> MonoM Exp
scoping Set VName
mergeArgs (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
            Exp -> MonoM Exp
transformExp Exp
e2
  let argset :: Set VName
argset = Set VName
mergeArgs Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
formArgs

  Exp
body' <- ExpReplacements -> MonoM Exp -> MonoM Exp
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
params (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM Exp -> MonoM Exp
scoping Set VName
argset (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body
  -- Maybe monomorphisation introduced new arrays to the loop, and
  -- maybe they have AnySize sizes.  This is not allowed.  Invent some
  -- sizes for them.
  ([VName]
pat_sizes, Pat ParamType
pat'') <- Pat ParamType -> MonoM ([VName], Pat ParamType)
forall (m :: * -> *).
MonadFreshNames m =>
Pat ParamType -> m ([VName], Pat ParamType)
sizesForPat Pat ParamType
pat'
  AppRes
res' <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([VName]
-> Pat ParamType
-> Exp
-> LoopFormBase Info VName
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[VName]
-> PatBase f vn ParamType
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
Loop ([VName]
sparams' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
pat_sizes) Pat ParamType
pat'' Exp
e1' LoopFormBase Info VName
form' Exp
body' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')
transformAppExp (BinOp (QualName VName
fname, SrcLoc
_) (Info StructType
t) (Exp
e1, Info (Maybe VName)
d1) (Exp
e2, Info (Maybe VName)
d2) SrcLoc
loc) AppRes
res = do
  (AppRes StructType
ret [VName]
ext) <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  Exp
fname' <- SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t)
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
e2' <- Exp -> MonoM Exp
transformExp Exp
e2
  if StructType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> StructType
typeOf Exp
e1') Bool -> Bool -> Bool
&& StructType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (Exp -> StructType
typeOf Exp
e2')
    then Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> Exp -> Exp -> Exp -> Exp
forall {vn}.
StructType
-> [VName]
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
applyOp StructType
ret [VName]
ext Exp
fname' Exp
e1' Exp
e2'
    else do
      -- We have to flip the arguments to the function, because
      -- operator application is left-to-right, while function
      -- application is outside-in.  This matters when the arguments
      -- produce existential sizes.  There are later places in the
      -- compiler where we transform BinOp to Apply, but anything that
      -- involves existential sizes will necessarily go through here.
      (Exp
x_param_e, PatBase Info VName StructType
x_param) <- Exp -> MonoM (Exp, PatBase Info VName StructType)
forall {m :: * -> *}.
MonadFreshNames m =>
Exp -> m (Exp, PatBase Info VName StructType)
makeVarParam Exp
e1'
      (Exp
y_param_e, PatBase Info VName StructType
y_param) <- Exp -> MonoM (Exp, PatBase Info VName StructType)
forall {m :: * -> *}.
MonadFreshNames m =>
Exp -> m (Exp, PatBase Info VName StructType)
makeVarParam Exp
e2'
      -- XXX: the type annotations here are wrong, but hopefully it
      -- doesn't matter as there will be an outer AppExp to handle
      -- them.
      Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
        AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
              []
              PatBase Info VName StructType
x_param
              Exp
e1'
              ( AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
                  ([SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] PatBase Info VName StructType
y_param Exp
e2' (StructType -> [VName] -> Exp -> Exp -> Exp -> Exp
forall {vn}.
StructType
-> [VName]
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
applyOp StructType
ret [VName]
ext Exp
fname' Exp
x_param_e Exp
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
$ StructType -> [VName] -> AppRes
AppRes StructType
ret [VName]
forall a. Monoid a => a
mempty)
              )
              SrcLoc
forall a. Monoid a => a
mempty
          )
          (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes StructType
ret [VName]
forall a. Monoid a => a
mempty))
  where
    applyOp :: StructType
-> [VName]
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
-> ExpBase Info vn
applyOp StructType
ret [VName]
ext ExpBase Info vn
fname' ExpBase Info vn
x ExpBase Info vn
y =
      ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
        (ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply ExpBase Info vn
fname' [(Diet
Observe, Info (Maybe VName) -> Maybe VName
forall a. Info a -> a
unInfo Info (Maybe VName)
d1, ExpBase Info vn
x)] (StructType -> [VName] -> AppRes
AppRes StructType
ret [VName]
forall a. Monoid a => a
mempty))
        [(Diet
Observe, Info (Maybe VName) -> Maybe VName
forall a. Info a -> a
unInfo Info (Maybe VName)
d2, ExpBase Info vn
y)]
        (StructType -> [VName] -> AppRes
AppRes StructType
ret [VName]
ext)

    makeVarParam :: Exp -> m (Exp, PatBase Info VName StructType)
makeVarParam Exp
arg = do
      let argtype :: StructType
argtype = Exp -> StructType
typeOf Exp
arg
      VName
x <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"binop_p"
      (Exp, PatBase Info VName StructType)
-> m (Exp, PatBase Info VName StructType)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
x (StructType -> Info StructType
forall a. a -> Info a
Info StructType
argtype) SrcLoc
forall a. Monoid a => a
mempty
        )
transformAppExp (LetWith IdentBase Info VName StructType
id1 IdentBase Info VName StructType
id2 SliceBase Info VName
idxs Exp
e1 Exp
body SrcLoc
loc) AppRes
res = do
  IdentBase Info VName StructType
id1' <- IdentBase Info VName StructType
-> MonoM (IdentBase Info VName StructType)
forall {f :: * -> *} {vn} {u}.
Traversable f =>
IdentBase f vn (TypeBase Exp u)
-> MonoM (IdentBase f vn (TypeBase Exp u))
transformIdent IdentBase Info VName StructType
id1
  IdentBase Info VName StructType
id2' <- IdentBase Info VName StructType
-> MonoM (IdentBase Info VName StructType)
forall {f :: * -> *} {vn} {u}.
Traversable f =>
IdentBase f vn (TypeBase Exp u)
-> MonoM (IdentBase f vn (TypeBase Exp u))
transformIdent IdentBase Info VName StructType
id2
  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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
  Exp
e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  Exp
body' <- Set VName -> MonoM Exp -> MonoM Exp
scoping (VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName StructType -> VName
forall {k} (f :: k -> *) vn (t :: k). IdentBase f vn t -> vn
identName IdentBase Info VName StructType
id1') (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body
  AppRes
res' <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (IdentBase Info VName StructType
-> IdentBase Info VName StructType
-> SliceBase Info VName
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn StructType
-> IdentBase f vn StructType
-> SliceBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetWith IdentBase Info VName StructType
id1' IdentBase Info VName StructType
id2' SliceBase Info VName
idxs' Exp
e1' Exp
body' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')
  where
    transformIdent :: IdentBase f vn (TypeBase Exp u)
-> MonoM (IdentBase f vn (TypeBase Exp u))
transformIdent (Ident vn
v f (TypeBase Exp u)
t SrcLoc
vloc) =
      vn
-> f (TypeBase Exp u) -> SrcLoc -> IdentBase f vn (TypeBase Exp u)
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident vn
v (f (TypeBase Exp u) -> SrcLoc -> IdentBase f vn (TypeBase Exp u))
-> MonoM (f (TypeBase Exp u))
-> MonoM (SrcLoc -> IdentBase f vn (TypeBase Exp u))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp u -> MonoM (TypeBase Exp u))
-> f (TypeBase Exp u) -> MonoM (f (TypeBase Exp u))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> f a -> f (f b)
traverse TypeBase Exp u -> MonoM (TypeBase Exp u)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType f (TypeBase Exp u)
t MonoM (SrcLoc -> IdentBase f vn (TypeBase Exp u))
-> MonoM SrcLoc -> MonoM (IdentBase f vn (TypeBase Exp u))
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
vloc
transformAppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) AppRes
res =
  AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
    (AppExpBase Info VName -> Info AppRes -> Exp)
-> MonoM (AppExpBase Info VName) -> MonoM (Info AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index (Exp -> SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
-> MonoM Exp
-> MonoM (SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e0 MonoM (SliceBase Info VName -> SrcLoc -> AppExpBase Info VName)
-> MonoM (SliceBase Info VName)
-> MonoM (SrcLoc -> AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs MonoM (SrcLoc -> AppExpBase Info VName)
-> MonoM SrcLoc -> MonoM (AppExpBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc)
    MonoM (Info AppRes -> Exp) -> MonoM (Info AppRes) -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> MonoM AppRes -> MonoM (Info AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AppRes -> MonoM AppRes
transformAppRes AppRes
res)
transformAppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) AppRes
res = do
  Set VName
implicitDims <- Set VName -> MonoM (Set VName)
askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ StructType -> FV
forall u. TypeBase Exp u -> FV
freeInType (StructType -> FV) -> StructType -> FV
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  NonEmpty (CaseBase Info VName)
cs' <- (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)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM (Set VName -> CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase Set VName
implicitDims) NonEmpty (CaseBase Info VName)
cs
  AppRes
res' <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  if Set VName -> Bool
forall a. Set a -> Bool
S.null Set VName
implicitDims
    then Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (CaseBase Info VName)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match Exp
e' NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')
    else do
      VName
tmpVar <- [Char] -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"matched_variable"
      Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
        AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
          ( [SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat
              ((VName -> SizeBinder VName) -> [VName] -> [SizeBinder VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> SizeBinder VName
forall vn. vn -> SrcLoc -> SizeBinder vn
`SizeBinder` SrcLoc
forall a. Monoid a => a
mempty) ([VName] -> [SizeBinder VName]) -> [VName] -> [SizeBinder VName]
forall a b. (a -> b) -> a -> b
$ Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
implicitDims)
              (VName -> Info StructType -> SrcLoc -> PatBase Info VName StructType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
tmpVar (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e') SrcLoc
forall a. Monoid a => a
mempty)
              Exp
e'
              ( AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp
                  (Exp
-> NonEmpty (CaseBase Info VName)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn) -> SrcLoc -> AppExpBase f vn
Match (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
tmpVar) (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e') SrcLoc
forall a. Monoid a => a
mempty) NonEmpty (CaseBase Info VName)
cs' SrcLoc
loc)
                  (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res)
              )
              SrcLoc
forall a. Monoid a => a
mempty
          )
          (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
res')

-- Monomorphization of expressions.
transformExp :: Exp -> MonoM Exp
transformExp :: Exp -> MonoM Exp
transformExp e :: Exp
e@Literal {} = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@IntLit {} = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@FloatLit {} = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp e :: Exp
e@StringLit {} = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
transformExp (Parens Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) =
  (QualName VName, SrcLoc) -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (TupLit [Exp]
es SrcLoc
loc) =
  [Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit ([Exp] -> SrcLoc -> Exp) -> MonoM [Exp] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) =
  [FieldBase Info VName] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit ([FieldBase Info VName] -> SrcLoc -> Exp)
-> MonoM [FieldBase Info VName] -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> [FieldBase Info VName] -> MonoM [FieldBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField [FieldBase Info VName]
fs MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
  where
    transformField :: FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (RecordFieldExplicit Name
name Exp
e SrcLoc
loc') =
      Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name (Exp -> SrcLoc -> FieldBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> FieldBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> FieldBase Info VName)
-> MonoM SrcLoc -> MonoM (FieldBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc'
    transformField (RecordFieldImplicit VName
v Info StructType
t SrcLoc
_) = do
      Info StructType
t' <- (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t
      FieldBase Info VName -> MonoM (FieldBase Info VName)
transformField (FieldBase Info VName -> MonoM (FieldBase Info VName))
-> FieldBase Info VName -> MonoM (FieldBase Info VName)
forall a b. (a -> b) -> a -> b
$
        Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
          (VName -> Name
baseName VName
v)
          (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) Info StructType
t' SrcLoc
loc)
          SrcLoc
loc
transformExp (ArrayLit [Exp]
es Info StructType
t SrcLoc
loc) =
  [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
ArrayLit ([Exp] -> Info StructType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info StructType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> MonoM Exp
transformExp [Exp]
es MonoM (Info StructType -> SrcLoc -> Exp)
-> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (AppExp AppExpBase Info VName
e Info AppRes
res) =
  AppExpBase Info VName -> AppRes -> MonoM Exp
transformAppExp AppExpBase Info VName
e (Info AppRes -> AppRes
forall a. Info a -> a
unInfo Info AppRes
res)
transformExp (Var QualName VName
fname (Info StructType
t) SrcLoc
loc) =
  SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t)
transformExp (Hole Info StructType
t SrcLoc
loc) =
  Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (Info StructType -> SrcLoc -> Exp)
-> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Ascript Exp
e TypeExp Info VName
tp SrcLoc
loc) =
  Exp -> TypeExp Info VName -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp f vn -> SrcLoc -> ExpBase f vn
Ascript (Exp -> TypeExp Info VName -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (TypeExp Info VName -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeExp Info VName -> SrcLoc -> Exp)
-> MonoM (TypeExp Info VName) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp Info VName -> MonoM (TypeExp Info VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Info VName
tp MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Coerce Exp
e TypeExp Info VName
te Info StructType
t SrcLoc
loc) =
  Exp -> TypeExp Info VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp f vn -> f StructType -> SrcLoc -> ExpBase f vn
Coerce (Exp -> TypeExp Info VName -> Info StructType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (TypeExp Info VName -> Info StructType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeExp Info VName -> Info StructType -> SrcLoc -> Exp)
-> MonoM (TypeExp Info VName)
-> MonoM (Info StructType -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TypeExp Info VName -> MonoM (TypeExp Info VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Info VName
te MonoM (Info StructType -> SrcLoc -> Exp)
-> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Negate Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Not Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Not (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Lambda [Pat ParamType]
params Exp
e0 Maybe (TypeExp Info VName)
decl Info ResRetType
tp SrcLoc
loc) = do
  let patArgs :: Set VName
patArgs = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params
  Set VName
dimArgs <- Set VName -> MonoM (Set VName) -> MonoM (Set VName)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
patArgs (MonoM (Set VName) -> MonoM (Set VName))
-> MonoM (Set VName) -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM (Set VName)
askIntros ((Pat ParamType -> Set VName) -> [Pat ParamType] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (FV -> Set VName
fvVars (FV -> Set VName)
-> (Pat ParamType -> FV) -> Pat ParamType -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> FV
forall u. Pat (TypeBase Exp u) -> FV
freeInPat) [Pat ParamType]
params)
  let argset :: Set VName
argset = Set VName
dimArgs Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
patArgs
  [Pat ParamType]
params' <- (Pat ParamType -> MonoM (Pat ParamType))
-> [Pat ParamType] -> MonoM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> MonoM (Pat ParamType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat [Pat ParamType]
params
  ExpReplacements
paramed <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset
  [Pat ParamType]
-> Exp
-> Maybe (TypeExp Info VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
params'
    (Exp
 -> Maybe (TypeExp Info VName) -> Info ResRetType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM
     (Maybe (TypeExp Info VName) -> Info ResRetType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpReplacements -> MonoM Exp -> MonoM Exp
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
paramed (Set VName -> MonoM Exp -> MonoM Exp
scoping Set VName
argset (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
e0)
    MonoM
  (Maybe (TypeExp Info VName) -> Info ResRetType -> SrcLoc -> Exp)
-> MonoM (Maybe (TypeExp Info VName))
-> MonoM (Info ResRetType -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (TypeExp Info VName) -> MonoM (Maybe (TypeExp Info VName))
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (TypeExp Info VName)
decl
    MonoM (Info ResRetType -> SrcLoc -> Exp)
-> MonoM (Info ResRetType) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (ResRetType -> MonoM ResRetType)
-> Info ResRetType -> MonoM (Info ResRetType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse ResRetType -> MonoM ResRetType
forall u. RetTypeBase Exp u -> MonoM (RetTypeBase Exp u)
transformRetType Info ResRetType
tp
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (OpSection QualName VName
qn Info StructType
t SrcLoc
loc) =
  Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn Info StructType
t SrcLoc
loc
transformExp (OpSectionLeft QualName VName
fname (Info StructType
t) Exp
e (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
arg (Info ResRetType
rettype, Info [VName]
retext) SrcLoc
loc) = do
  let (Info (PName
xp, ParamType
xtype, Maybe VName
xargext), Info (PName
yp, ParamType
ytype)) = (Info (PName, ParamType, Maybe VName), Info (PName, ParamType))
arg
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  QualName VName
-> Maybe Exp
-> Maybe Exp
-> StructType
-> (PName, ParamType, Maybe VName)
-> (PName, ParamType, Maybe VName)
-> (ResRetType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
    QualName VName
fname
    (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
    Maybe Exp
forall a. Maybe a
Nothing
    StructType
t
    (PName
xp, ParamType
xtype, Maybe VName
xargext)
    (PName
yp, ParamType
ytype, Maybe VName
forall a. Maybe a
Nothing)
    (ResRetType
rettype, [VName]
retext)
    SrcLoc
loc
transformExp (OpSectionRight QualName VName
fname (Info StructType
t) Exp
e (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
arg (Info ResRetType
rettype) SrcLoc
loc) = do
  let (Info (PName
xp, ParamType
xtype), Info (PName
yp, ParamType
ytype, Maybe VName
yargext)) = (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
arg
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  QualName VName
-> Maybe Exp
-> Maybe Exp
-> StructType
-> (PName, ParamType, Maybe VName)
-> (PName, ParamType, Maybe VName)
-> (ResRetType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection
    QualName VName
fname
    Maybe Exp
forall a. Maybe a
Nothing
    (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
e')
    StructType
t
    (PName
xp, ParamType
xtype, Maybe VName
forall a. Maybe a
Nothing)
    (PName
yp, ParamType
ytype, Maybe VName
yargext)
    (ResRetType
rettype, [])
    SrcLoc
loc
transformExp (ProjectSection [Name]
fields (Info StructType
t) SrcLoc
loc) = do
  StructType
t' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
t
  [Name] -> StructType -> SrcLoc -> MonoM Exp
desugarProjectSection [Name]
fields StructType
t' SrcLoc
loc
transformExp (IndexSection SliceBase Info VName
idxs (Info StructType
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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
  SliceBase Info VName -> StructType -> SrcLoc -> MonoM Exp
desugarIndexSection SliceBase Info VName
idxs' StructType
t SrcLoc
loc
transformExp (Project Name
n Exp
e Info StructType
tp SrcLoc
loc) = do
  Info StructType
tp' <- (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
tp
  Exp
e' <- Exp -> MonoM Exp
transformExp Exp
e
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
n Exp
e' Info StructType
tp' SrcLoc
loc
transformExp (Update Exp
e1 SliceBase Info VName
idxs Exp
e2 SrcLoc
loc) =
  Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> SliceBase f vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update
    (Exp -> SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
    MonoM (SliceBase Info VName -> Exp -> SrcLoc -> Exp)
-> MonoM (SliceBase Info VName) -> MonoM (Exp -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM DimIndexBase Info VName -> MonoM (DimIndexBase Info VName)
transformDimIndex SliceBase Info VName
idxs
    MonoM (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (RecordUpdate Exp
e1 [Name]
fs Exp
e2 Info StructType
t SrcLoc
loc) =
  Exp -> [Name] -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name] -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
RecordUpdate
    (Exp -> [Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM ([Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
    MonoM ([Name] -> Exp -> Info StructType -> SrcLoc -> Exp)
-> MonoM [Name] -> MonoM (Exp -> Info StructType -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Name] -> MonoM [Name]
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
fs
    MonoM (Exp -> Info StructType -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info StructType -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2
    MonoM (Info StructType -> SrcLoc -> Exp)
-> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t
    MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Assert Exp
e1 Exp
e2 Info Text
desc SrcLoc
loc) =
  Exp -> Exp -> Info Text -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f Text -> SrcLoc -> ExpBase f vn
Assert (Exp -> Exp -> Info Text -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Exp -> Info Text -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1 MonoM (Exp -> Info Text -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (Info Text -> SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> MonoM Exp
transformExp Exp
e2 MonoM (Info Text -> SrcLoc -> Exp)
-> MonoM (Info Text) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Info Text -> MonoM (Info Text)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Info Text
desc MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Constr Name
name [Exp]
all_es Info StructType
t SrcLoc
loc) =
  Name -> [Exp] -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f StructType -> SrcLoc -> ExpBase f vn
Constr Name
name ([Exp] -> Info StructType -> SrcLoc -> Exp)
-> MonoM [Exp] -> MonoM (Info StructType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> MonoM Exp) -> [Exp] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Exp -> MonoM Exp
transformExp [Exp]
all_es MonoM (Info StructType -> SrcLoc -> Exp)
-> MonoM (Info StructType) -> MonoM (SrcLoc -> Exp)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (StructType -> MonoM StructType)
-> Info StructType -> MonoM (Info StructType)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Info a -> f (Info b)
traverse StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType Info StructType
t MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
transformExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) =
  AttrInfo VName -> Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn.
AttrInfo vn -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo VName
info (Exp -> SrcLoc -> Exp) -> MonoM Exp -> MonoM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (SrcLoc -> Exp) -> MonoM SrcLoc -> MonoM Exp
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

transformCase :: S.Set VName -> Case -> MonoM Case
transformCase :: Set VName -> CaseBase Info VName -> MonoM (CaseBase Info VName)
transformCase Set VName
implicitDims (CasePat PatBase Info VName StructType
p Exp
e SrcLoc
loc) = do
  PatBase Info VName StructType
p' <- PatBase Info VName StructType
-> MonoM (PatBase Info VName StructType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat PatBase Info VName StructType
p
  PatBase Info VName StructType
-> Exp -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatBase f vn StructType -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat PatBase Info VName StructType
p' (Exp -> SrcLoc -> CaseBase Info VName)
-> MonoM Exp -> MonoM (SrcLoc -> CaseBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set VName -> MonoM Exp -> MonoM Exp
scoping ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (PatBase Info VName StructType -> [VName]
forall t. Pat t -> [VName]
patNames PatBase Info VName StructType
p) Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
implicitDims) (Exp -> MonoM Exp
transformExp Exp
e) MonoM (SrcLoc -> CaseBase Info VName)
-> MonoM SrcLoc -> MonoM (CaseBase Info VName)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> MonoM SrcLoc
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc

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

-- Transform an operator section into a lambda.
desugarBinOpSection ::
  QualName VName ->
  Maybe Exp ->
  Maybe Exp ->
  StructType ->
  (PName, ParamType, Maybe VName) ->
  (PName, ParamType, Maybe VName) ->
  (ResRetType, [VName]) ->
  SrcLoc ->
  MonoM Exp
desugarBinOpSection :: QualName VName
-> Maybe Exp
-> Maybe Exp
-> StructType
-> (PName, ParamType, Maybe VName)
-> (PName, ParamType, Maybe VName)
-> (ResRetType, [VName])
-> SrcLoc
-> MonoM Exp
desugarBinOpSection QualName VName
fname Maybe Exp
e_left Maybe Exp
e_right StructType
t (PName
xp, ParamType
xtype, Maybe VName
xext) (PName
yp, ParamType
ytype, Maybe VName
yext) (RetType [VName]
dims TypeBase Exp Uniqueness
rettype, [VName]
retext) SrcLoc
loc = do
  StructType
t' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
t
  Exp
op <- SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname (StructType -> MonoM Exp) -> StructType -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct StructType
t
  (VName
v1, Exp -> Exp
wrap_left, Exp
e1, [Pat ParamType]
p1) <- Maybe Exp
-> ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
forall {m :: * -> *} {u}.
MonadFreshNames m =>
Maybe Exp
-> TypeBase Exp u
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp u)])
makeVarParam Maybe Exp
e_left (ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType]))
-> MonoM ParamType
-> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ParamType -> MonoM ParamType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType ParamType
xtype
  (VName
v2, Exp -> Exp
wrap_right, Exp
e2, [Pat ParamType]
p2) <- Maybe Exp
-> ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
forall {m :: * -> *} {u}.
MonadFreshNames m =>
Maybe Exp
-> TypeBase Exp u
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp u)])
makeVarParam Maybe Exp
e_right (ParamType -> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType]))
-> MonoM ParamType
-> MonoM (VName, Exp -> Exp, Exp, [Pat ParamType])
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ParamType -> MonoM ParamType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType ParamType
ytype
  let apply_left :: Exp
apply_left =
        Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply
          Exp
op
          [(Diet
Observe, Maybe VName
xext, Exp
e1)]
          (StructType -> [VName] -> AppRes
AppRes (ScalarTypeBase Exp NoUniqueness -> StructType
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp NoUniqueness -> StructType)
-> ScalarTypeBase Exp NoUniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ NoUniqueness
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp NoUniqueness
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow NoUniqueness
forall a. Monoid a => a
mempty PName
yp (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
ytype) (ParamType -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct ParamType
ytype) ([VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase Exp Uniqueness -> ResRetType)
-> TypeBase Exp Uniqueness -> ResRetType
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StructType -> TypeBase Exp Uniqueness
forall u. Uniqueness -> TypeBase Exp u -> TypeBase Exp Uniqueness
toRes Uniqueness
Nonunique StructType
t')) [])
      onDim :: ExpBase f VName -> ExpBase f VName
onDim (Var QualName VName
d f StructType
typ SrcLoc
_)
        | 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 -> f StructType -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v1) f StructType
typ SrcLoc
loc
        | 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 -> f StructType -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v2) f StructType
typ SrcLoc
loc
      onDim ExpBase f VName
d = ExpBase f VName
d
      rettype' :: TypeBase Exp Uniqueness
rettype' = (Exp -> Exp) -> TypeBase Exp Uniqueness -> TypeBase Exp Uniqueness
forall a b c. (a -> b) -> TypeBase a c -> TypeBase b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Exp -> Exp
forall {f :: * -> *}. ExpBase f VName -> ExpBase f VName
onDim TypeBase Exp Uniqueness
rettype
  Exp
body <-
    Set VName -> MonoM Exp -> MonoM Exp
scoping ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName
v1, VName
v2]) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
      Exp -> [(Diet, Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Diet, Maybe VName, ExpBase Info vn)]
-> AppRes
-> ExpBase Info vn
mkApply Exp
apply_left [(Diet
Observe, Maybe VName
yext, Exp
e2)]
        (AppRes -> Exp) -> MonoM AppRes -> MonoM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AppRes -> MonoM AppRes
transformAppRes (StructType -> [VName] -> AppRes
AppRes (TypeBase Exp Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp Uniqueness
rettype') [VName]
retext)
  ResRetType
rettype'' <- Set VName -> ResRetType -> MonoM ResRetType
forall as.
Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
transformRetTypeSizes ([VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName
v1, VName
v2]) (ResRetType -> MonoM ResRetType) -> ResRetType -> MonoM ResRetType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp Uniqueness -> ResRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims TypeBase Exp Uniqueness
rettype'
  Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> (Exp -> Exp) -> Exp -> MonoM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp
wrap_left (Exp -> Exp) -> (Exp -> Exp) -> Exp -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp
wrap_right (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$
    [Pat ParamType]
-> Exp
-> Maybe (TypeExp Info VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp f vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda ([Pat ParamType]
p1 [Pat ParamType] -> [Pat ParamType] -> [Pat ParamType]
forall a. [a] -> [a] -> [a]
++ [Pat ParamType]
p2) Exp
body Maybe (TypeExp Info VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
rettype'') SrcLoc
loc
  where
    patAndVar :: TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
patAndVar TypeBase Exp u
argtype = do
      VName
x <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
      (VName, PatBase Info VName (TypeBase Exp u), Exp)
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( VName
x,
          VName
-> Info (TypeBase Exp u)
-> SrcLoc
-> PatBase Info VName (TypeBase Exp u)
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
x (TypeBase Exp u -> Info (TypeBase Exp u)
forall a. a -> Info a
Info TypeBase Exp u
argtype) SrcLoc
forall a. Monoid a => a
mempty,
          QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (StructType -> Info StructType
forall a. a -> Info a
Info (TypeBase Exp u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp u
argtype)) SrcLoc
forall a. Monoid a => a
mempty
        )

    makeVarParam :: Maybe Exp
-> TypeBase Exp u
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp u)])
makeVarParam (Just Exp
e) TypeBase Exp u
argtype = do
      (VName
v, PatBase Info VName (TypeBase Exp u)
pat, Exp
var_e) <- TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
forall {m :: * -> *} {u}.
MonadFreshNames m =>
TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
patAndVar TypeBase Exp u
argtype
      let wrap :: Exp -> Exp
wrap Exp
body =
            AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp ([SizeBinder VName]
-> PatBase Info VName StructType
-> Exp
-> Exp
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
[SizeBinder vn]
-> PatBase f vn StructType
-> ExpBase f vn
-> ExpBase f vn
-> SrcLoc
-> AppExpBase f vn
LetPat [] ((TypeBase Exp u -> StructType)
-> PatBase Info VName (TypeBase Exp u)
-> PatBase Info VName StructType
forall a b.
(a -> b) -> PatBase Info VName a -> PatBase Info VName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase Exp u -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct PatBase Info VName (TypeBase Exp u)
pat) Exp
e Exp
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
$ StructType -> [VName] -> AppRes
AppRes (Exp -> StructType
typeOf Exp
body) [VName]
forall a. Monoid a => a
mempty)
      (VName, Exp -> Exp, Exp, [PatBase Info VName (TypeBase Exp u)])
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp u)])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, Exp -> Exp
wrap, Exp
var_e, [])
    makeVarParam Maybe Exp
Nothing TypeBase Exp u
argtype = do
      (VName
v, PatBase Info VName (TypeBase Exp u)
pat, Exp
var_e) <- TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
forall {m :: * -> *} {u}.
MonadFreshNames m =>
TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
patAndVar TypeBase Exp u
argtype
      (VName, Exp -> Exp, Exp, [PatBase Info VName (TypeBase Exp u)])
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp u)])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName
v, Exp -> Exp
forall a. a -> a
id, Exp
var_e, [PatBase Info VName (TypeBase Exp u)
pat])

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

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

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

transformPat :: Pat (TypeBase Size u) -> MonoM (Pat (TypeBase Size u))
transformPat :: forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat = (TypeBase Exp u -> MonoM (TypeBase Exp u))
-> PatBase Info VName (TypeBase Exp u)
-> MonoM (PatBase Info VName (TypeBase Exp u))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PatBase Info VName a -> f (PatBase Info VName b)
traverse TypeBase Exp u -> MonoM (TypeBase Exp u)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType

type DimInst = M.Map VName Size

dimMapping ::
  (Monoid a) =>
  TypeBase Size a ->
  TypeBase Size a ->
  ExpReplacements ->
  ExpReplacements ->
  DimInst
dimMapping :: forall a.
Monoid a =>
TypeBase Exp a
-> TypeBase Exp a -> ExpReplacements -> ExpReplacements -> DimInst
dimMapping TypeBase Exp a
t1 TypeBase Exp a
t2 ExpReplacements
r1 ExpReplacements
r2 = State DimInst (TypeBase Exp a) -> DimInst -> DimInst
forall s a. State s a -> s -> s
execState (([VName] -> Exp -> Exp -> StateT DimInst Identity Exp)
-> TypeBase Exp a
-> TypeBase Exp a
-> State DimInst (TypeBase Exp a)
forall as (m :: * -> *) d1 d2.
(Monoid as, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 as -> TypeBase d2 as -> m (TypeBase d1 as)
matchDims [VName] -> Exp -> Exp -> StateT DimInst Identity Exp
forall {m :: * -> *} {t :: * -> *}.
(Foldable t, MonadState DimInst m) =>
t VName -> Exp -> Exp -> m Exp
onDims TypeBase Exp a
t1 TypeBase Exp a
t2) DimInst
forall a. Monoid a => a
mempty
  where
    revMap :: [(b, a)] -> [(a, b)]
revMap = ((b, a) -> (a, b)) -> [(b, a)] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map (\(b
k, a
v) -> (a
v, b
k))
    named1 :: [(VName, ReplacedExp)]
named1 = ExpReplacements -> [(VName, ReplacedExp)]
forall {b} {a}. [(b, a)] -> [(a, b)]
revMap ExpReplacements
r1
    named2 :: [(VName, ReplacedExp)]
named2 = ExpReplacements -> [(VName, ReplacedExp)]
forall {b} {a}. [(b, a)] -> [(a, b)]
revMap ExpReplacements
r2

    onDims :: t VName -> Exp -> Exp -> m Exp
onDims t VName
bound Exp
e1 Exp
e2 = do
      t VName -> Exp -> Exp -> m ()
forall {m :: * -> *} {t :: * -> *}.
(Foldable t, MonadState DimInst m) =>
t VName -> Exp -> Exp -> m ()
onExps t VName
bound Exp
e1 Exp
e2
      Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e1

    onExps :: t VName -> Exp -> Exp -> m ()
onExps t VName
bound (Var QualName VName
v Info StructType
_ SrcLoc
_) Exp
e = do
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ((VName -> Bool) -> Set VName -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> t VName -> Bool
forall a. Eq a => a -> t a -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` t VName
bound) (Set VName -> Bool) -> Set VName -> Bool
forall a b. (a -> b) -> a -> b
$ Exp -> Set VName
freeVarsInExp Exp
e) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
        (DimInst -> DimInst) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (VName -> Exp -> DimInst -> DimInst
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Exp
e)
      case VName -> [(VName, ReplacedExp)] -> Maybe ReplacedExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) [(VName, ReplacedExp)]
named1 of
        Just ReplacedExp
rexp -> t VName -> Exp -> Exp -> m ()
onExps t VName
bound (ReplacedExp -> Exp
unReplaced ReplacedExp
rexp) Exp
e
        Maybe ReplacedExp
Nothing -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    onExps t VName
bound Exp
e (Var QualName VName
v Info StructType
_ SrcLoc
_)
      | Just ReplacedExp
rexp <- VName -> [(VName, ReplacedExp)] -> Maybe ReplacedExp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) [(VName, ReplacedExp)]
named2 =
          t VName -> Exp -> Exp -> m ()
onExps t VName
bound Exp
e (ReplacedExp -> Exp
unReplaced ReplacedExp
rexp)
    onExps t VName
bound Exp
e1 Exp
e2
      | Just [(Exp, Exp)]
es <- Exp -> Exp -> Maybe [(Exp, Exp)]
similarExps Exp
e1 Exp
e2 =
          ((Exp, Exp) -> m ()) -> [(Exp, Exp)] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((Exp -> Exp -> m ()) -> (Exp, Exp) -> m ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Exp -> Exp -> m ()) -> (Exp, Exp) -> m ())
-> (Exp -> Exp -> m ()) -> (Exp, Exp) -> m ()
forall a b. (a -> b) -> a -> b
$ t VName -> Exp -> Exp -> m ()
onExps t VName
bound) [(Exp, Exp)]
es
    onExps t VName
_ Exp
_ Exp
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
forall a. Monoid a => a
mempty

    freeVarsInExp :: Exp -> Set VName
freeVarsInExp = FV -> Set VName
fvVars (FV -> Set VName) -> (Exp -> FV) -> Exp -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> FV
freeInExp

inferSizeArgs :: [TypeParam] -> StructType -> ExpReplacements -> StructType -> MonoM [Exp]
inferSizeArgs :: [TypeParamBase VName]
-> StructType -> ExpReplacements -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
tparams StructType
bind_t ExpReplacements
bind_r StructType
t = do
  ExpReplacements
r <- (ExpReplacements -> ExpReplacements -> ExpReplacements)
-> MonoM (ExpReplacements -> ExpReplacements)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
(<>) MonoM (ExpReplacements -> ExpReplacements)
-> MonoM ExpReplacements -> MonoM ExpReplacements
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Env -> ExpReplacements) -> MonoM ExpReplacements
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> ExpReplacements
envParametrized
  let dinst :: DimInst
dinst = StructType
-> StructType -> ExpReplacements -> ExpReplacements -> DimInst
forall a.
Monoid a =>
TypeBase Exp a
-> TypeBase Exp a -> ExpReplacements -> ExpReplacements -> DimInst
dimMapping StructType
bind_t StructType
t ExpReplacements
bind_r ExpReplacements
r
  (TypeParamBase VName -> MonoM Exp)
-> [TypeParamBase VName] -> MonoM [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (DimInst -> TypeParamBase VName -> MonoM Exp
forall {k}. Ord k => Map k Exp -> TypeParamBase k -> MonoM Exp
tparamArg DimInst
dinst) [TypeParamBase VName]
tparams
  where
    tparamArg :: Map k Exp -> TypeParamBase k -> MonoM Exp
tparamArg Map k Exp
dinst TypeParamBase k
tp =
      case k -> Map k Exp -> Maybe Exp
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 Exp
dinst of
        Just Exp
e ->
          Exp -> MonoM Exp
replaceExp Exp
e
        Maybe Exp
Nothing ->
          Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Integer -> SrcLoc -> Exp
sizeFromInteger Integer
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 u. TypeBase MonoSize u -> TypeBase MonoSize u
f
  where
    f :: TypeBase MonoSize u -> TypeBase MonoSize u
    f :: forall u. TypeBase MonoSize u -> TypeBase MonoSize u
f (Array u
u Shape MonoSize
shape ScalarTypeBase MonoSize NoUniqueness
t) = u
-> Shape MonoSize
-> ScalarTypeBase MonoSize NoUniqueness
-> TypeBase MonoSize u
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array u
u Shape MonoSize
shape (ScalarTypeBase MonoSize NoUniqueness
-> ScalarTypeBase MonoSize NoUniqueness
forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
f' ScalarTypeBase MonoSize NoUniqueness
t)
    f (Scalar ScalarTypeBase MonoSize u
t) = ScalarTypeBase MonoSize u -> TypeBase MonoSize u
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase MonoSize u -> TypeBase MonoSize u)
-> ScalarTypeBase MonoSize u -> TypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
f' ScalarTypeBase MonoSize u
t
    f' :: ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
    f' :: forall {u}. ScalarTypeBase MonoSize u -> ScalarTypeBase MonoSize u
f' (Record Map Name (TypeBase MonoSize u)
fs) = Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u)
-> Map Name (TypeBase MonoSize u) -> ScalarTypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ (TypeBase MonoSize u -> TypeBase MonoSize u)
-> Map Name (TypeBase MonoSize u) -> Map Name (TypeBase MonoSize u)
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase MonoSize u -> TypeBase MonoSize u
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
f Map Name (TypeBase MonoSize u)
fs
    f' (Sum Map Name [TypeBase MonoSize u]
cs) = Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u)
-> Map Name [TypeBase MonoSize u] -> ScalarTypeBase MonoSize u
forall a b. (a -> b) -> a -> b
$ ([TypeBase MonoSize u] -> [TypeBase MonoSize u])
-> Map Name [TypeBase MonoSize u] -> Map Name [TypeBase MonoSize u]
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeBase MonoSize u -> TypeBase MonoSize u)
-> [TypeBase MonoSize u] -> [TypeBase MonoSize u]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase MonoSize u -> TypeBase MonoSize u
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
f) Map Name [TypeBase MonoSize u]
cs
    f' (Arrow u
u PName
_ Diet
d1 MonoType
t1 (RetType [VName]
dims TypeBase MonoSize Uniqueness
t2)) =
      u
-> PName
-> Diet
-> MonoType
-> RetTypeBase MonoSize Uniqueness
-> ScalarTypeBase MonoSize u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
u PName
Unnamed Diet
d1 (MonoType -> MonoType
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
f MonoType
t1) ([VName]
-> TypeBase MonoSize Uniqueness -> RetTypeBase MonoSize Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (TypeBase MonoSize Uniqueness -> TypeBase MonoSize Uniqueness
forall u. TypeBase MonoSize u -> TypeBase MonoSize u
f TypeBase MonoSize Uniqueness
t2))
    f' ScalarTypeBase MonoSize u
t = ScalarTypeBase MonoSize u
t

transformRetType :: RetTypeBase Size u -> MonoM (RetTypeBase Size u)
transformRetType :: forall u. RetTypeBase Exp u -> MonoM (RetTypeBase Exp u)
transformRetType (RetType [VName]
ext TypeBase Exp u
t) = [VName] -> TypeBase Exp u -> RetTypeBase Exp u
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
ext (TypeBase Exp u -> RetTypeBase Exp u)
-> MonoM (TypeBase Exp u) -> MonoM (RetTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeBase Exp u -> MonoM (TypeBase Exp u)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType TypeBase Exp u
t

-- | arrowArg takes a return type and returns it
-- with the existentials bound moved at the right of arrows.
-- It also gives the new set of parameters to consider.
arrowArg ::
  S.Set VName -> -- scope
  S.Set VName -> -- set of argument
  [VName] -> -- size parameters
  RetTypeBase Size as ->
  (RetTypeBase Size as, S.Set VName)
arrowArg :: forall as.
Set VName
-> Set VName
-> [VName]
-> RetTypeBase Exp as
-> (RetTypeBase Exp as, Set VName)
arrowArg Set VName
scope Set VName
argset [VName]
args_params RetTypeBase Exp as
rety =
  let (RetTypeBase Exp as
rety', (Set VName
funArgs, Set VName
_)) = Writer (Set VName, Set VName) (RetTypeBase Exp as)
-> (RetTypeBase Exp as, (Set VName, Set VName))
forall w a. Writer w a -> (a, w)
runWriter ((Set VName, [VName])
-> Set VName
-> RetTypeBase Exp as
-> Writer (Set VName, Set VName) (RetTypeBase Exp as)
forall as'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp as'
-> Writer (Set VName, Set VName) (RetTypeBase Exp as')
arrowArgRetType (Set VName
scope, [VName]
forall a. Monoid a => a
mempty) Set VName
argset RetTypeBase Exp as
rety)
      new_params :: Set VName
new_params = Set VName
funArgs Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
args_params
   in (Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
arrowCleanRetType Set VName
new_params RetTypeBase Exp as
rety', Set VName
new_params)
  where
    -- \| takes a type (or return type) and returns it
    -- with the existentials bound moved at the right of arrows.
    -- It also gives (through writer monad) size variables used in arrow arguments
    -- and variables that are constructively used.
    -- The returned type should be cleanned, as too many existentials are introduced.
    arrowArgRetType ::
      (S.Set VName, [VName]) ->
      S.Set VName ->
      RetTypeBase Size as' ->
      Writer (S.Set VName, S.Set VName) (RetTypeBase Size as')
    arrowArgRetType :: forall as'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp as'
-> Writer (Set VName, Set VName) (RetTypeBase Exp as')
arrowArgRetType (Set VName
scope', [VName]
dimsToPush) Set VName
argset' (RetType [VName]
dims TypeBase Exp as'
ty) = WriterT
  (Set VName, Set VName)
  Identity
  (RetTypeBase Exp as',
   (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as')
forall a.
WriterT
  (Set VName, Set VName)
  Identity
  (a, (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity a
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (WriterT
   (Set VName, Set VName)
   Identity
   (RetTypeBase Exp as',
    (Set VName, Set VName) -> (Set VName, Set VName))
 -> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as'))
-> WriterT
     (Set VName, Set VName)
     Identity
     (RetTypeBase Exp as',
      (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp as')
forall a b. (a -> b) -> a -> b
$ do
      let dims' :: [VName]
dims' = [VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dimsToPush
      (TypeBase Exp as'
ty', (Set VName
_, Set VName
canExt)) <- WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
-> WriterT
     (Set VName, Set VName)
     Identity
     (TypeBase Exp as', (Set VName, Set VName))
forall a.
WriterT (Set VName, Set VName) Identity a
-> WriterT
     (Set VName, Set VName) Identity (a, (Set VName, Set VName))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen (WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
 -> WriterT
      (Set VName, Set VName)
      Identity
      (TypeBase Exp as', (Set VName, Set VName)))
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
-> WriterT
     (Set VName, Set VName)
     Identity
     (TypeBase Exp as', (Set VName, Set VName))
forall a b. (a -> b) -> a -> b
$ (Set VName, [VName])
-> TypeBase Exp as'
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
forall as'.
(Set VName, [VName])
-> TypeBase Exp as'
-> Writer (Set VName, Set VName) (TypeBase Exp as')
arrowArgType (Set VName
argset' Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
scope', [VName]
dims') TypeBase Exp as'
ty
      (RetTypeBase Exp as',
 (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT
     (Set VName, Set VName)
     Identity
     (RetTypeBase Exp as',
      (Set VName, Set VName) -> (Set VName, Set VName))
forall a. a -> WriterT (Set VName, Set VName) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([VName] -> TypeBase Exp as' -> RetTypeBase Exp as'
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ((VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
canExt) [VName]
dims') TypeBase Exp as'
ty', (Set VName -> Set VName)
-> (Set VName, Set VName) -> (Set VName, Set VName)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set VName
canExt))

    arrowArgScalar :: (Set VName, [VName])
-> ScalarTypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
arrowArgScalar (Set VName, [VName])
env (Record Map Name (TypeBase Exp u)
fs) =
      Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase Exp u) -> ScalarTypeBase Exp u)
-> WriterT
     (Set VName, Set VName) Identity (Map Name (TypeBase Exp u))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp u
 -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
-> Map Name (TypeBase Exp u)
-> WriterT
     (Set VName, Set VName) Identity (Map Name (TypeBase Exp u))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse ((Set VName, [VName])
-> TypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp u)
forall as'.
(Set VName, [VName])
-> TypeBase Exp as'
-> Writer (Set VName, Set VName) (TypeBase Exp as')
arrowArgType (Set VName, [VName])
env) Map Name (TypeBase Exp u)
fs
    arrowArgScalar (Set VName, [VName])
env (Sum Map Name [TypeBase Exp u]
cs) =
      Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase Exp u] -> ScalarTypeBase Exp u)
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase Exp u]
 -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
-> Map Name [TypeBase Exp u]
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Name a -> f (Map Name b)
traverse (([TypeBase Exp u]
  -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
 -> Map Name [TypeBase Exp u]
 -> WriterT
      (Set VName, Set VName) Identity (Map Name [TypeBase Exp u]))
-> ((TypeBase Exp u
     -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
    -> [TypeBase Exp u]
    -> WriterT (Set VName, Set VName) Identity [TypeBase Exp u])
-> (TypeBase Exp u
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
-> Map Name [TypeBase Exp u]
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp u])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp u
 -> WriterT (Set VName, Set VName) Identity (TypeBase Exp u))
-> [TypeBase Exp u]
-> WriterT (Set VName, Set VName) Identity [TypeBase Exp u]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse) ((Set VName, [VName])
-> TypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp u)
forall as'.
(Set VName, [VName])
-> TypeBase Exp as'
-> Writer (Set VName, Set VName) (TypeBase Exp as')
arrowArgType (Set VName, [VName])
env) Map Name [TypeBase Exp u]
cs
    arrowArgScalar (Set VName
scope', [VName]
dimsToPush) (Arrow u
as PName
argName Diet
d StructType
argT ResRetType
retT) =
      WriterT
  (Set VName, Set VName)
  Identity
  (ScalarTypeBase Exp u,
   (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall a.
WriterT
  (Set VName, Set VName)
  Identity
  (a, (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity a
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (WriterT
   (Set VName, Set VName)
   Identity
   (ScalarTypeBase Exp u,
    (Set VName, Set VName) -> (Set VName, Set VName))
 -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u))
-> WriterT
     (Set VName, Set VName)
     Identity
     (ScalarTypeBase Exp u,
      (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall a b. (a -> b) -> a -> b
$ do
        let intros :: Set VName
intros = (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter VName -> Bool
notIntrisic Set VName
argset' Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.difference` Set VName
scope'
        ResRetType
retT' <- (Set VName, [VName])
-> Set VName
-> ResRetType
-> Writer (Set VName, Set VName) ResRetType
forall as'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp as'
-> Writer (Set VName, Set VName) (RetTypeBase Exp as')
arrowArgRetType (Set VName
scope', (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
intros) [VName]
dimsToPush) Set VName
fullArgset ResRetType
retT
        (ScalarTypeBase Exp u,
 (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT
     (Set VName, Set VName)
     Identity
     (ScalarTypeBase Exp u,
      (Set VName, Set VName) -> (Set VName, Set VName))
forall a. a -> WriterT (Set VName, Set VName) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (u
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp u
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow u
as PName
argName Diet
d StructType
argT ResRetType
retT', (Set VName -> Set VName)
-> (Set VName -> Set VName)
-> (Set VName, Set VName)
-> (Set VName, Set VName)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Set VName
intros `S.union`) (Set VName -> Set VName -> Set VName
forall a b. a -> b -> a
const Set VName
forall a. Monoid a => a
mempty))
      where
        notIntrisic :: VName -> Bool
notIntrisic VName
vn = VName -> Int
baseTag VName
vn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxIntrinsicTag
        argset' :: Set VName
argset' = FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ StructType -> FV
forall u. TypeBase Exp u -> FV
freeInType StructType
argT
        fullArgset :: Set VName
fullArgset =
          Set VName
argset'
            Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> case PName
argName of
              PName
Unnamed -> Set VName
forall a. Monoid a => a
mempty
              Named VName
vn -> VName -> Set VName
forall a. a -> Set a
S.singleton VName
vn
    arrowArgScalar (Set VName, [VName])
env (TypeVar u
u QualName VName
qn [TypeArg Exp]
args) =
      u -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp u
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar u
u QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp u)
-> WriterT (Set VName, Set VName) Identity [TypeArg Exp]
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeArg Exp
 -> WriterT (Set VName, Set VName) Identity (TypeArg Exp))
-> [TypeArg Exp]
-> WriterT (Set VName, Set VName) Identity [TypeArg Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM TypeArg Exp
-> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
arrowArgArg [TypeArg Exp]
args
      where
        arrowArgArg :: TypeArg Exp
-> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
arrowArgArg (TypeArgDim Exp
dim) = Exp -> TypeArg Exp
forall dim. dim -> TypeArg dim
TypeArgDim (Exp -> TypeArg Exp)
-> WriterT (Set VName, Set VName) Identity Exp
-> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> WriterT (Set VName, Set VName) Identity Exp
forall {a} {a} {m :: * -> *} {f :: * -> *}.
(MonadWriter (a, Set a) m, Monoid a, Ord a) =>
ExpBase f a -> m (ExpBase f a)
arrowArgSize Exp
dim
        arrowArgArg (TypeArgType StructType
ty) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp)
-> WriterT (Set VName, Set VName) Identity StructType
-> WriterT (Set VName, Set VName) Identity (TypeArg Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Set VName, [VName])
-> StructType -> WriterT (Set VName, Set VName) Identity StructType
forall as'.
(Set VName, [VName])
-> TypeBase Exp as'
-> Writer (Set VName, Set VName) (TypeBase Exp as')
arrowArgType (Set VName, [VName])
env StructType
ty
    arrowArgScalar (Set VName, [VName])
_ ScalarTypeBase Exp u
ty = ScalarTypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
forall a. a -> WriterT (Set VName, Set VName) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase Exp u
ty

    arrowArgType ::
      (S.Set VName, [VName]) ->
      TypeBase Size as' ->
      Writer (S.Set VName, S.Set VName) (TypeBase Size as')
    arrowArgType :: forall as'.
(Set VName, [VName])
-> TypeBase Exp as'
-> Writer (Set VName, Set VName) (TypeBase Exp as')
arrowArgType (Set VName, [VName])
env (Array as'
u Shape Exp
shape ScalarTypeBase Exp NoUniqueness
scalar) =
      as'
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as'
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array as'
u (Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
-> WriterT (Set VName, Set VName) Identity (Shape Exp)
-> WriterT
     (Set VName, Set VName)
     Identity
     (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> WriterT (Set VName, Set VName) Identity Exp)
-> Shape Exp -> WriterT (Set VName, Set VName) Identity (Shape Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Shape a -> f (Shape b)
traverse Exp -> WriterT (Set VName, Set VName) Identity Exp
forall {a} {a} {m :: * -> *} {f :: * -> *}.
(MonadWriter (a, Set a) m, Monoid a, Ord a) =>
ExpBase f a -> m (ExpBase f a)
arrowArgSize Shape Exp
shape WriterT
  (Set VName, Set VName)
  Identity
  (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as')
-> WriterT
     (Set VName, Set VName) Identity (ScalarTypeBase Exp NoUniqueness)
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
forall a b.
WriterT (Set VName, Set VName) Identity (a -> b)
-> WriterT (Set VName, Set VName) Identity a
-> WriterT (Set VName, Set VName) Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Set VName, [VName])
-> ScalarTypeBase Exp NoUniqueness
-> WriterT
     (Set VName, Set VName) Identity (ScalarTypeBase Exp NoUniqueness)
forall {u}.
(Set VName, [VName])
-> ScalarTypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
arrowArgScalar (Set VName, [VName])
env ScalarTypeBase Exp NoUniqueness
scalar
    arrowArgType (Set VName, [VName])
env (Scalar ScalarTypeBase Exp as'
ty) =
      ScalarTypeBase Exp as' -> TypeBase Exp as'
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp as' -> TypeBase Exp as')
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp as')
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp as')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Set VName, [VName])
-> ScalarTypeBase Exp as'
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp as')
forall {u}.
(Set VName, [VName])
-> ScalarTypeBase Exp u
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp u)
arrowArgScalar (Set VName, [VName])
env ScalarTypeBase Exp as'
ty

    arrowArgSize :: ExpBase f a -> m (ExpBase f a)
arrowArgSize s :: ExpBase f a
s@(Var QualName a
qn f StructType
_ SrcLoc
_) = (ExpBase f a, (a, Set a)) -> m (ExpBase f a)
forall a. (a, (a, Set a)) -> m a
forall w (m :: * -> *) a. MonadWriter w m => (a, w) -> m a
writer (ExpBase f a
s, (a
forall a. Monoid a => a
mempty, a -> Set a
forall a. a -> Set a
S.singleton (a -> Set a) -> a -> Set a
forall a b. (a -> b) -> a -> b
$ QualName a -> a
forall vn. QualName vn -> vn
qualLeaf QualName a
qn))
    arrowArgSize ExpBase f a
s = ExpBase f a -> m (ExpBase f a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f a
s

    -- \| arrowClean cleans the mess in the type
    arrowCleanRetType :: S.Set VName -> RetTypeBase Size as -> RetTypeBase Size as
    arrowCleanRetType :: forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
arrowCleanRetType Set VName
paramed (RetType [VName]
dims TypeBase Exp as
ty) =
      [VName] -> TypeBase Exp as -> RetTypeBase Exp as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName] -> [VName]
forall a. Ord a => [a] -> [a]
nubOrd ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
paramed) [VName]
dims) (Set VName -> TypeBase Exp as -> TypeBase Exp as
forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
arrowCleanType (Set VName
paramed Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims) TypeBase Exp as
ty)

    arrowCleanScalar :: S.Set VName -> ScalarTypeBase Size as -> ScalarTypeBase Size as
    arrowCleanScalar :: forall as.
Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
arrowCleanScalar Set VName
paramed (Record Map Name (TypeBase Exp as)
fs) =
      Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as
forall dim u. Map Name (TypeBase dim u) -> ScalarTypeBase dim u
Record (Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as)
-> Map Name (TypeBase Exp as) -> ScalarTypeBase Exp as
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp as -> TypeBase Exp as)
-> Map Name (TypeBase Exp as) -> Map Name (TypeBase Exp as)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Set VName -> TypeBase Exp as -> TypeBase Exp as
forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
arrowCleanType Set VName
paramed) Map Name (TypeBase Exp as)
fs
    arrowCleanScalar Set VName
paramed (Sum Map Name [TypeBase Exp as]
cs) =
      Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as
forall dim u. Map Name [TypeBase dim u] -> ScalarTypeBase dim u
Sum (Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as)
-> Map Name [TypeBase Exp as] -> ScalarTypeBase Exp as
forall a b. (a -> b) -> a -> b
$ (([TypeBase Exp as] -> [TypeBase Exp as])
-> Map Name [TypeBase Exp as] -> Map Name [TypeBase Exp as]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (([TypeBase Exp as] -> [TypeBase Exp as])
 -> Map Name [TypeBase Exp as] -> Map Name [TypeBase Exp as])
-> ((TypeBase Exp as -> TypeBase Exp as)
    -> [TypeBase Exp as] -> [TypeBase Exp as])
-> (TypeBase Exp as -> TypeBase Exp as)
-> Map Name [TypeBase Exp as]
-> Map Name [TypeBase Exp as]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp as -> TypeBase Exp as)
-> [TypeBase Exp as] -> [TypeBase Exp as]
forall a b. (a -> b) -> [a] -> [b]
map) (Set VName -> TypeBase Exp as -> TypeBase Exp as
forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
arrowCleanType Set VName
paramed) Map Name [TypeBase Exp as]
cs
    arrowCleanScalar Set VName
paramed (Arrow as
as PName
argName Diet
d StructType
argT ResRetType
retT) =
      as
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp as
forall dim u.
u
-> PName
-> Diet
-> TypeBase dim NoUniqueness
-> RetTypeBase dim Uniqueness
-> ScalarTypeBase dim u
Arrow as
as PName
argName Diet
d StructType
argT (Set VName -> ResRetType -> ResRetType
forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
arrowCleanRetType Set VName
paramed ResRetType
retT)
    arrowCleanScalar Set VName
paramed (TypeVar as
u QualName VName
qn [TypeArg Exp]
args) =
      as -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp as
forall dim u.
u -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim u
TypeVar as
u QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp as)
-> [TypeArg Exp] -> ScalarTypeBase Exp as
forall a b. (a -> b) -> a -> b
$ (TypeArg Exp -> TypeArg Exp) -> [TypeArg Exp] -> [TypeArg Exp]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Exp -> TypeArg Exp
arrowCleanArg [TypeArg Exp]
args
      where
        arrowCleanArg :: TypeArg Exp -> TypeArg Exp
arrowCleanArg (TypeArgDim Exp
dim) = Exp -> TypeArg Exp
forall dim. dim -> TypeArg dim
TypeArgDim Exp
dim
        arrowCleanArg (TypeArgType StructType
ty) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoUniqueness -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp) -> StructType -> TypeArg Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> StructType -> StructType
forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
arrowCleanType Set VName
paramed StructType
ty
    arrowCleanScalar Set VName
_ ScalarTypeBase Exp as
ty = ScalarTypeBase Exp as
ty

    arrowCleanType :: S.Set VName -> TypeBase Size as -> TypeBase Size as
    arrowCleanType :: forall as. Set VName -> TypeBase Exp as -> TypeBase Exp as
arrowCleanType Set VName
paramed (Array as
u Shape Exp
shape ScalarTypeBase Exp NoUniqueness
scalar) =
      as
-> Shape Exp -> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as
forall dim u.
u -> Shape dim -> ScalarTypeBase dim NoUniqueness -> TypeBase dim u
Array as
u Shape Exp
shape (ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as)
-> ScalarTypeBase Exp NoUniqueness -> TypeBase Exp as
forall a b. (a -> b) -> a -> b
$ Set VName
-> ScalarTypeBase Exp NoUniqueness
-> ScalarTypeBase Exp NoUniqueness
forall as.
Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
arrowCleanScalar Set VName
paramed ScalarTypeBase Exp NoUniqueness
scalar
    arrowCleanType Set VName
paramed (Scalar ScalarTypeBase Exp as
ty) =
      ScalarTypeBase Exp as -> TypeBase Exp as
forall dim u. ScalarTypeBase dim u -> TypeBase dim u
Scalar (ScalarTypeBase Exp as -> TypeBase Exp as)
-> ScalarTypeBase Exp as -> TypeBase Exp as
forall a b. (a -> b) -> a -> b
$ Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
forall as.
Set VName -> ScalarTypeBase Exp as -> ScalarTypeBase Exp as
arrowCleanScalar Set VName
paramed ScalarTypeBase Exp as
ty

-- 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 (VName
name, [TypeParamBase VName]
tparams, [Pat ParamType]
params, ResRetType
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) MonoType
inst_t = do
  Bool
letFun <- (Env -> Bool) -> MonoM Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Bool) -> MonoM Bool) -> (Env -> Bool) -> MonoM Bool
forall a b. (a -> b) -> a -> b
$ VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member VName
name (Set VName -> Bool) -> (Env -> Set VName) -> Env -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Set VName
envScope
  let paramGetClean :: Set VName -> MonoM ExpReplacements
paramGetClean Set VName
argset =
        if Bool
letFun
          then Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset
          else do
            ExpReplacements
ret <- MonoM ExpReplacements
forall s (m :: * -> *). MonadState s m => m s
get
            ExpReplacements -> MonoM ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ExpReplacements
forall a. Monoid a => a
mempty
            ExpReplacements -> MonoM ExpReplacements
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpReplacements
ret
  (if Bool
letFun then MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. a -> a
id else MonoM (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. MonoM a -> MonoM a
isolateNormalisation) (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 = [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
params ResRetType
rettype
    (Map VName StructRetType
substs, [TypeParamBase VName]
t_shape_params) <-
      SrcLoc
-> TypeBase () NoUniqueness
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> TypeBase () NoUniqueness
-> MonoType
-> m (Map VName StructRetType, [TypeParamBase VName])
typeSubstsM SrcLoc
loc (StructType -> TypeBase () NoUniqueness
forall as. TypeBase Exp as -> TypeBase () as
noSizes StructType
bind_t) (MonoType
 -> MonoM (Map VName StructRetType, [TypeParamBase VName]))
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParamBase VName])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
inst_t
    let shape_names :: Set VName
shape_names = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName ([TypeParamBase VName] -> [VName])
-> [TypeParamBase VName] -> [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
        substs' :: Map VName (Subst StructRetType)
substs' = (StructRetType -> Subst StructRetType)
-> Map VName StructRetType -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ([TypeParamBase VName] -> StructRetType -> Subst StructRetType
forall t. [TypeParamBase VName] -> t -> Subst t
Subst []) Map VName StructRetType
substs
        substStructType :: ParamType -> ParamType
substStructType =
          (VName -> Maybe (Subst (RetTypeBase Exp Diet)))
-> ParamType -> ParamType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase Exp as)))
-> TypeBase Exp as -> TypeBase Exp as
substTypesAny ((Subst StructRetType -> Subst (RetTypeBase Exp Diet))
-> Maybe (Subst StructRetType)
-> Maybe (Subst (RetTypeBase Exp Diet))
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StructRetType -> RetTypeBase Exp Diet)
-> Subst StructRetType -> Subst (RetTypeBase Exp Diet)
forall a b. (a -> b) -> Subst a -> Subst b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((NoUniqueness -> Diet) -> StructRetType -> RetTypeBase Exp Diet
forall b c a. (b -> c) -> RetTypeBase a b -> RetTypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Diet -> NoUniqueness -> Diet
forall a b. a -> b -> a
const Diet
forall a. Monoid a => a
mempty))) (Maybe (Subst StructRetType)
 -> Maybe (Subst (RetTypeBase Exp Diet)))
-> TypeSubs -> VName -> Maybe (Subst (RetTypeBase Exp Diet))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs'))
        params' :: [Pat ParamType]
params' = (Pat ParamType -> Pat ParamType)
-> [Pat ParamType] -> [Pat ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> (ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
forall t. Bool -> (t -> t) -> Pat t -> Pat t
substPat Bool
entry ParamType -> ParamType
substStructType) [Pat ParamType]
params
    [Pat ParamType]
params'' <- Set VName -> MonoM [Pat ParamType] -> MonoM [Pat ParamType]
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
shape_names (MonoM [Pat ParamType] -> MonoM [Pat ParamType])
-> MonoM [Pat ParamType] -> MonoM [Pat ParamType]
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> MonoM (Pat ParamType))
-> [Pat ParamType] -> MonoM [Pat ParamType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Pat ParamType -> MonoM (Pat ParamType)
forall u. Pat (TypeBase Exp u) -> MonoM (Pat (TypeBase Exp u))
transformPat [Pat ParamType]
params'
    ExpReplacements
exp_naming <- Set VName -> MonoM ExpReplacements
paramGetClean Set VName
shape_names

    let args :: Set VName
args = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params
        arg_params :: [VName]
arg_params = ((ReplacedExp, VName) -> VName) -> ExpReplacements -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (ReplacedExp, VName) -> VName
forall a b. (a, b) -> b
snd ExpReplacements
exp_naming

    ResRetType
rettype' <-
      ExpReplacements -> MonoM ResRetType -> MonoM ResRetType
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
exp_naming (MonoM ResRetType -> MonoM ResRetType)
-> MonoM ResRetType -> MonoM ResRetType
forall a b. (a -> b) -> a -> b
$
        Set VName -> MonoM ResRetType -> MonoM ResRetType
forall a. Set VName -> MonoM a -> MonoM a
withArgs (Set VName
args Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
shape_names) (MonoM ResRetType -> MonoM ResRetType)
-> MonoM ResRetType -> MonoM ResRetType
forall a b. (a -> b) -> a -> b
$
          ResRetType -> MonoM ResRetType
forall u. RetTypeBase Exp u -> MonoM (RetTypeBase Exp u)
hardTransformRetType (TypeSubs -> ResRetType -> ResRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') ResRetType
rettype)
    ExpReplacements
extNaming <- Set VName -> MonoM ExpReplacements
paramGetClean (Set VName
args Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
shape_names)
    Set VName
scope <- Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union Set VName
shape_names (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM (Set VName)
askScope'
    let (ResRetType
rettype'', Set VName
new_params) = Set VName
-> Set VName -> [VName] -> ResRetType -> (ResRetType, Set VName)
forall as.
Set VName
-> Set VName
-> [VName]
-> RetTypeBase Exp as
-> (RetTypeBase Exp as, Set VName)
arrowArg Set VName
scope Set VName
args [VName]
arg_params ResRetType
rettype'
        bind_t' :: StructType
bind_t' = TypeSubs -> StructType -> StructType
forall as.
Monoid as =>
(VName -> Maybe (Subst (RetTypeBase Exp as)))
-> TypeBase Exp as -> TypeBase Exp as
substTypesAny (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') StructType
bind_t
        ([TypeParamBase VName]
shape_params_explicit, [TypeParamBase VName]
shape_params_implicit) =
          (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
mustBeExplicitInBinding StructType
bind_t'' Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` StructType -> Set VName
mustBeExplicitInBinding 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 [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ (VName -> TypeParamBase VName) -> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) (Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
new_params)
        exp_naming' :: ExpReplacements
exp_naming' = ((ReplacedExp, VName) -> Bool)
-> ExpReplacements -> ExpReplacements
forall a. (a -> Bool) -> [a] -> [a]
filter ((VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
new_params) (VName -> Bool)
-> ((ReplacedExp, VName) -> VName) -> (ReplacedExp, VName) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ReplacedExp, VName) -> VName
forall a b. (a, b) -> b
snd) (ExpReplacements
extNaming ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
<> ExpReplacements
exp_naming)

        bind_t'' :: StructType
bind_t'' = [Pat ParamType] -> ResRetType -> StructType
funType [Pat ParamType]
params'' ResRetType
rettype''
        bind_r :: ExpReplacements
bind_r = ExpReplacements
exp_naming ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
<> ExpReplacements
extNaming
    Exp
body' <- TypeSubs -> Exp -> MonoM Exp
forall {m :: * -> *}. Monad m => TypeSubs -> Exp -> m Exp
updateExpTypes (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') Exp
body
    Exp
body'' <- ExpReplacements -> MonoM Exp -> MonoM Exp
forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
exp_naming' (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> MonoM Exp -> MonoM Exp
forall a. Set VName -> MonoM a -> MonoM a
withArgs (Set VName
shape_names Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
args) (MonoM Exp -> MonoM Exp) -> MonoM Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> MonoM Exp
transformExp Exp
body'
    Set VName
scope' <- Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
S.union (Set VName
shape_names Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
args) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM (Set VName)
askScope'
    Exp
body''' <-
      if Bool
letFun
        then Set VName -> Exp -> MonoM Exp
unscoping (Set VName
shape_names Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
args) Exp
body''
        else ExpReplacements -> Exp -> Exp
expReplace ExpReplacements
exp_naming' (Exp -> Exp) -> MonoM Exp -> MonoM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> ExpReplacements -> MonoM Exp
calculateDims Exp
body'' (ExpReplacements -> MonoM Exp)
-> (ExpReplacements -> ExpReplacements)
-> ExpReplacements
-> MonoM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set VName -> ExpReplacements -> ExpReplacements
canCalculate Set VName
scope' (ExpReplacements -> MonoM Exp)
-> MonoM ExpReplacements -> MonoM Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MonoM ExpReplacements
forall s (m :: * -> *). MonadState s m => m s
get)

    Bool
seen_before <- VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem VName
name ([VName] -> Bool) -> (Lifts -> [VName]) -> Lifts -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (((VName, MonoType), (VName, InferSizeArgs)) -> VName)
-> Lifts -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map ((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) (Lifts -> Bool) -> MonoM Lifts -> MonoM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM Lifts
getLifts
    VName
name' <-
      if [TypeParamBase VName] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeParamBase VName]
tparams Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
entry Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
seen_before
        then VName -> MonoM VName
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
name
        else VName -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
name

    (VName, InferSizeArgs, ValBind)
-> MonoM (VName, InferSizeArgs, ValBind)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      ( VName
name',
        [TypeParamBase VName]
-> StructType -> ExpReplacements -> InferSizeArgs
inferSizeArgs [TypeParamBase VName]
shape_params_explicit StructType
bind_t'' ExpReplacements
bind_r,
        if Bool
entry
          then
            VName
-> [TypeParamBase VName]
-> [Pat ParamType]
-> ResRetType
-> Exp
-> ValBind
toValBinding
              VName
name'
              ([TypeParamBase VName]
shape_params_explicit [TypeParamBase VName]
-> [TypeParamBase VName] -> [TypeParamBase VName]
forall a. [a] -> [a] -> [a]
++ [TypeParamBase VName]
shape_params_implicit)
              [Pat ParamType]
params''
              ResRetType
rettype''
              (ExpReplacements -> Exp -> Exp
entryAssert ExpReplacements
exp_naming Exp
body''')
          else
            VName
-> [TypeParamBase VName]
-> [Pat ParamType]
-> ResRetType
-> Exp
-> ValBind
toValBinding
              VName
name'
              [TypeParamBase VName]
shape_params_implicit
              ((TypeParamBase VName -> Pat ParamType)
-> [TypeParamBase VName] -> [Pat ParamType]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> Pat ParamType
forall {vn} {dim} {als}.
TypeParamBase vn -> PatBase Info vn (TypeBase dim als)
shapeParam [TypeParamBase VName]
shape_params_explicit [Pat ParamType] -> [Pat ParamType] -> [Pat ParamType]
forall a. [a] -> [a] -> [a]
++ [Pat ParamType]
params'')
              ResRetType
rettype''
              Exp
body'''
      )
  where
    askScope' :: MonoM (Set VName)
askScope' = (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter (VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` ResRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims ResRetType
rettype) (Set VName -> Set VName) -> MonoM (Set VName) -> MonoM (Set VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM (Set VName)
askScope

    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 :: TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs = ASTMapper m -> Exp -> m Exp
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
forall (m :: * -> *). Monad m => ASTMapper m -> Exp -> m Exp
astMap (TypeSubs -> ASTMapper m
mapper TypeSubs
substs)

    hardTransformRetType :: RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
hardTransformRetType (RetType [VName]
dims TypeBase Exp as
ty) = do
      TypeBase Exp as
ty' <- TypeBase Exp as -> MonoM (TypeBase Exp as)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType TypeBase Exp as
ty
      Set VName
unbounded <- Set VName -> MonoM (Set VName)
askIntros (Set VName -> MonoM (Set VName)) -> Set VName -> MonoM (Set VName)
forall a b. (a -> b) -> a -> b
$ FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ TypeBase Exp as -> FV
forall u. TypeBase Exp u -> FV
freeInType TypeBase Exp as
ty'
      let dims' :: [VName]
dims' = Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
unbounded
      RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Exp as -> MonoM (RetTypeBase Exp as))
-> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp as -> RetTypeBase Exp as
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType ([VName]
dims' [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
dims) TypeBase Exp as
ty'

    mapper :: TypeSubs -> ASTMapper m
mapper TypeSubs
substs =
      ASTMapper
        { mapOnExp :: Exp -> m Exp
mapOnExp = TypeSubs -> Exp -> m Exp
updateExpTypes TypeSubs
substs,
          mapOnName :: QualName VName -> m (QualName VName)
mapOnName = QualName VName -> m (QualName VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnStructType :: StructType -> m StructType
mapOnStructType = StructType -> m StructType
forall a. a -> m a
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
. TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
          mapOnParamType :: ParamType -> m ParamType
mapOnParamType = ParamType -> m ParamType
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ParamType -> m ParamType)
-> (ParamType -> ParamType) -> ParamType -> m ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> ParamType -> ParamType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs,
          mapOnResRetType :: ResRetType -> m ResRetType
mapOnResRetType = ResRetType -> m ResRetType
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ResRetType -> m ResRetType)
-> (ResRetType -> ResRetType) -> ResRetType -> m ResRetType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> ResRetType -> ResRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs
        }

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

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

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

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

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

    onDim :: MonoSize -> t (t m) Exp
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 (m :: * -> *) a. Monad m => m a -> t m a
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 (m :: * -> *) a. Monad m => m a -> t m a
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
$ [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"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)
          Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> t (t m) Exp) -> Exp -> t (t m) Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) SrcLoc
forall a. Monoid a => a
mempty
        Just VName
d ->
          Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> t (t m) Exp) -> Exp -> t (t m) Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
d) SrcLoc
forall a. Monoid a => a
mempty
    onDim MonoSize
MonoAnon = Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
anySize

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

toPolyBinding :: ValBind -> PolyBinding
toPolyBinding :: ValBind -> PolyBinding
toPolyBinding (ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp Info VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) =
  (VName, [TypeParamBase VName], [Pat ParamType], ResRetType, Exp,
 [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding (VName
name, [TypeParamBase VName]
tparams, [Pat ParamType]
params, ResRetType
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)

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

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

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

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

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

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

transformValBind :: ValBind -> MonoM Env
transformValBind :: ValBind -> MonoM Env
transformValBind ValBind
valbind = do
  PolyBinding
valbind' <-
    ValBind -> PolyBinding
toPolyBinding
      (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

  case ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind of
    Maybe (Info EntryPoint)
Nothing -> () -> MonoM ()
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just (Info EntryPoint
entry) -> do
      StructType
t <-
        StructType -> MonoM StructType
removeTypeVariablesInType (StructType -> MonoM StructType) -> StructType -> MonoM StructType
forall a b. (a -> b) -> a -> b
$
          [Pat ParamType] -> ResRetType -> StructType
funType (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind) (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$
            Info ResRetType -> ResRetType
forall a. Info a -> a
unInfo (Info ResRetType -> ResRetType) -> Info ResRetType -> ResRetType
forall a b. (a -> b) -> a -> b
$
              ValBind -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
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 Exp als -> MonoType
monoType StructType
t
      EntryPoint
entry' <- EntryPoint -> MonoM EntryPoint
transformEntryPoint EntryPoint
entry
      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 = Info EntryPoint -> Maybe (Info EntryPoint)
forall a. a -> Maybe a
Just (Info EntryPoint -> Maybe (Info EntryPoint))
-> Info EntryPoint -> Maybe (Info EntryPoint)
forall a b. (a -> b) -> a -> b
$ EntryPoint -> Info EntryPoint
forall a. a -> Info a
Info EntryPoint
entry'})
      VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) (StructType -> MonoType
forall als. TypeBase Exp als -> MonoType
monoType StructType
t) (VName
name, InferSizeArgs
infer)

  Env -> MonoM Env
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    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',
        envGlobalScope :: Set VName
envGlobalScope =
          if [Pat ParamType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind)
            then [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ ResRetType -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims (ResRetType -> [VName]) -> ResRetType -> [VName]
forall a b. (a -> b) -> a -> b
$ Info ResRetType -> ResRetType
forall a. Info a -> a
unInfo (Info ResRetType -> ResRetType) -> Info ResRetType -> ResRetType
forall a b. (a -> b) -> a -> b
$ ValBind -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
valBindRetType ValBind
valbind
            else Set VName
forall a. Monoid a => a
mempty
      }

transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind :: TypeBind -> MonoM Env
transformTypeBind (TypeBind VName
name Liftedness
l [TypeParamBase VName]
tparams TypeExp Info VName
_ (Info (RetType [VName]
dims StructType
t)) Maybe DocComment
_ SrcLoc
_) = do
  Map VName (Subst StructRetType)
subs <- (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Env -> Map VName (Subst StructRetType))
 -> MonoM (Map VName (Subst StructRetType)))
-> (Env -> Map VName (Subst StructRetType))
-> MonoM (Map VName (Subst StructRetType))
forall a b. (a -> b) -> a -> b
$ (TypeBinding -> Subst StructRetType)
-> Map VName TypeBinding -> Map VName (Subst StructRetType)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBinding -> Subst StructRetType
substFromAbbr (Map VName TypeBinding -> Map VName (Subst StructRetType))
-> (Env -> Map VName TypeBinding)
-> Env
-> Map VName (Subst StructRetType)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env -> Map VName TypeBinding
envTypeBindings
  let tbinding :: TypeBinding
tbinding = Liftedness -> [TypeParamBase VName] -> StructRetType -> TypeBinding
TypeAbbr Liftedness
l [TypeParamBase VName]
tparams (StructRetType -> TypeBinding) -> StructRetType -> TypeBinding
forall a b. (a -> b) -> a -> b
$ [VName] -> StructType -> StructRetType
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [VName]
dims (StructType -> StructRetType) -> StructType -> StructRetType
forall a b. (a -> b) -> a -> b
$ TypeSubs -> StructType -> StructType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
subs) StructType
t
  Env -> MonoM Env
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
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]
_) =
  [Char] -> MonoM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> MonoM ()) -> [Char] -> MonoM ()
forall a b. (a -> b) -> a -> b
$
    [Char]
"The monomorphization module expects a module-free "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"input program, but received: "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Dec -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Dec
dec

-- | Monomorphise a list of top-level declarations. A module-free input program
-- is expected, so only value declarations and type declaration are accepted.
transformProg :: (MonadFreshNames m) => [Dec] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [Dec] -> m [ValBind]
transformProg [Dec]
decs =
  (((), Seq (VName, ValBind)) -> [ValBind])
-> m ((), Seq (VName, ValBind)) -> m [ValBind]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Seq ValBind -> [ValBind]
forall a. Seq a -> [a]
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 a b. (a -> b) -> Seq a -> Seq b
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