-- | 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).
--
-- * 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 (isJust, isNothing)
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.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
      ( Maybe EntryPoint,
        VName,
        [TypeParam],
        [Pat ParamType],
        ResRetType,
        Exp,
        [AttrInfo VName],
        SrcLoc
      )

-- | To deduplicate size expressions, we want a looser notion 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 = pure . expReplace 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] -> RetTypeBase Exp Uniqueness -> StructType
foldFunType [ParamType
forall dim als. TypeBase dim als
bool, ParamType
forall dim als. TypeBase dim als
bool] (RetTypeBase Exp Uniqueness -> StructType)
-> RetTypeBase Exp Uniqueness -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
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 -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply Exp
andop [(Maybe VName
forall a. Maybe a
Nothing, Exp
x'), (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 -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply Exp
eqop [(Maybe VName
forall a. Maybe a
Nothing, Exp
x'), (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 -> 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
pb1 Set VName
sc1 Set VName
gs1 ExpReplacements
pr1 <> :: Env -> Env -> Env
<> Env Map VName PolyBinding
pb2 Set VName
sc2 Set VName
gs2 ExpReplacements
pr2 = Map VName PolyBinding
-> Set VName -> Set VName -> ExpReplacements -> Env
Env (Map VName PolyBinding
pb1 Map VName PolyBinding
-> Map VName PolyBinding -> Map VName PolyBinding
forall a. Semigroup a => a -> a -> a
<> Map VName PolyBinding
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
-> Set VName -> Set VName -> ExpReplacements -> Env
Env Map VName PolyBinding
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 Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>)

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

-- | These now have monomorphic types in the given action. This is
-- used to handle shadowing.
withMono :: [VName] -> MonoM a -> MonoM a
withMono :: forall a. [VName] -> MonoM a -> MonoM a
withMono [] = MonoM a -> MonoM a
forall a. a -> a
id
withMono [VName]
vs = (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) -> MonoM a -> MonoM a)
-> (Env -> Env) -> MonoM a -> MonoM a
forall a b. (a -> b) -> a -> b
$ \Env
env ->
  Env
env {envPolyBindings = M.filterWithKey keep (envPolyBindings env)}
  where
    keep :: VName -> p -> Bool
keep VName
v p
_ = VName
v VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
vs

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 = 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 = 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
  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 M.lookup vn 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
  scope <- (Env -> Set VName) -> MonoM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Set VName
envScope
  scope' <- asks $ S.union scope . envGlobalScope
  scope'' <- asks $ S.union scope' . M.keysSet . envPolyBindings
  S.union scope'' . S.fromList . map (fst . snd) <$> 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 Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`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
  intros <- Set VName -> MonoM (Set VName)
askIntros Set VName
argset
  let usesIntros = Bool -> Bool
not (Bool -> Bool) -> (Exp -> Bool) -> Exp -> 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) -> (Exp -> Set VName) -> Exp -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
  (params, nxtBind) <- gets $ partition (usesIntros . unReplaced . fst)
  put nxtBind
  pure 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
    depends :: (ReplacedExp, b) -> (ReplacedExp, b) -> Bool
depends (ReplacedExp
a, b
_) (ReplacedExp
b, b
_) = ReplacedExp -> Exp
unReplaced ReplacedExp
b Exp -> [Exp] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Exp -> [Exp]
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
      reName <- VName -> f VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName VName
vn
      let 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
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 = 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]
      foldCalc repls $
        AppExp
          ( LetPat
              []
              (Id vn (Info i64) (srclocOf expr))
              expr
              body'
              mempty
          )
          appRes

unscoping :: S.Set VName -> Exp -> MonoM Exp
unscoping :: Set VName -> Exp -> MonoM Exp
unscoping Set VName
argset Exp
body = do
  localDims <- Set VName -> MonoM ExpReplacements
parametrizing Set VName
argset
  scope <- S.union argset <$> askScope
  calculateDims body $ canCalculate scope 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]

-- | The integer encodes an equivalence class, so we can keep
-- track of sizes that are statically identical.
data MonoSize
  = MonoKnown Int
  | MonoAnon Int
  deriving (MonoSize -> MonoSize -> Bool
(MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool) -> Eq MonoSize
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MonoSize -> MonoSize -> Bool
== :: MonoSize -> MonoSize -> Bool
$c/= :: MonoSize -> MonoSize -> Bool
/= :: MonoSize -> MonoSize -> Bool
Eq, 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)

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 (MonoAnon 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

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 {m :: * -> *} {p}.
MonadState (Int, Map Exp Int) m =>
Set VName -> p -> Exp -> m 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 -> m MonoSize
onDim Set VName
bound p
_ Exp
d
      -- 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
d = do
          (i, m) <- m (Int, Map Exp Int)
forall s (m :: * -> *). MonadState s m => m s
get
          case M.lookup d m of
            Just Int
prev ->
              MonoSize -> m MonoSize
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoAnon Int
prev
            Maybe Int
Nothing -> do
              (Int, Map Exp Int) -> m ()
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 -> m MonoSize
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoAnon Int
i
    onDim Set VName
_ p
_ Exp
d = do
      (i, m) <- m (Int, Map Exp Int)
forall s (m :: * -> *). MonadState s m => m s
get
      case M.lookup d m of
        Just Int
prev ->
          MonoSize -> m MonoSize
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> m MonoSize) -> MonoSize -> m MonoSize
forall a b. (a -> b) -> a -> b
$ Int -> MonoSize
MonoKnown Int
prev
        Maybe Int
Nothing -> do
          -- Ensure that each instance of anySize is treated distinctly.
          (Int, Map Exp Int) -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, if Exp
d Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
anySize then Map Exp Int
m else 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 -> m MonoSize
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MonoSize -> m MonoSize) -> MonoSize -> m 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) ((VName, MonoType), (VName, InferSizeArgs)) -> Lifts -> Lifts
forall a. a -> [a] -> [a]
:)

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

sizeVarName :: Exp -> String
sizeVarName :: Exp -> [Char]
sizeVarName Exp
e = [Char]
"d<{" [Char] -> ShowS
forall a. Semigroup 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. Semigroup a => a -> a -> a
<> [Char]
"}>"

-- | 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
      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'
      prev_param <- asks $ lookup e' . envParametrized
      case (prev_param, 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
          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
$ Exp -> [Char]
sizeVarName Exp
e
          modify ((e', vn) :)
          pure $ sizeFromName (qualName vn) (srclocOf 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
ft = do
  t' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
ft
  let mono_t = StructType -> MonoType
forall als. TypeBase Exp als -> MonoType
monoType StructType
ft
  if baseTag (qualLeaf fname) <= maxIntrinsicTag
    then pure $ var fname t'
    else do
      maybe_fname <- lookupLifted (qualLeaf fname) mono_t
      maybe_funbind <- lookupFun $ qualLeaf fname
      case (maybe_fname, 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
          (fname', infer, funbind') <- PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding PolyBinding
funbind MonoType
mono_t
          tell $ Seq.singleton (qualLeaf fname, funbind')
          addLifted (qualLeaf fname) mono_t (fname', infer)
          applySizeArgs fname' (toRes Nonunique t') <$> infer 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
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply
          ExpBase Info vn
f
          [(Maybe VName
forall a. Maybe a
Nothing, ExpBase Info vn
size_arg)]
          (StructType -> [VName] -> AppRes
AppRes ([ParamType] -> RetTypeBase Exp Uniqueness -> 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 -> RetTypeBase Exp Uniqueness
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] -> RetTypeBase Exp Uniqueness -> 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 -> RetTypeBase Exp Uniqueness
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 RetTypeBase Exp Uniqueness
retT) =
      u
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Uniqueness
-> 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 -> RetTypeBase Exp Uniqueness -> ScalarTypeBase Exp u)
-> MonoM StructType
-> MonoM (RetTypeBase Exp Uniqueness -> 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 (RetTypeBase Exp Uniqueness -> ScalarTypeBase Exp u)
-> MonoM (RetTypeBase Exp Uniqueness)
-> 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
<*> Set VName
-> RetTypeBase Exp Uniqueness -> MonoM (RetTypeBase Exp Uniqueness)
forall as.
Set VName -> RetTypeBase Exp as -> MonoM (RetTypeBase Exp as)
transformRetTypeSizes Set VName
argset RetTypeBase Exp Uniqueness
retT
      where
        argset :: Set VName
argset =
          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
  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
$ [VName] -> MonoM (TypeBase Exp as) -> MonoM (TypeBase Exp as)
forall a. [VName] -> MonoM a -> MonoM a
withMono [VName]
dims (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
  rl <- parametrizing argset
  let 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
  pure $ RetType dims' 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
  (params', 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) []
  pure (sizes, 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
          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"
          modify (v :)
          pure $ sizeFromName (qualName v) 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
  e1' <- Exp -> MonoM Exp
transformExp Exp
e1
  me' <- mapM transformExp me
  incl' <- mapM transformExp incl
  res' <- transformAppRes res
  pure $ AppExp (Range e1' me' incl' loc) (Info res')
transformAppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e Exp
body SrcLoc
loc) AppRes
res = do
  e' <- Exp -> MonoM Exp
transformExp Exp
e
  let 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)
  implicitDims <- withArgs dimArgs $ askIntros $ fvVars $ freeInPat pat
  let dimArgs' = Set VName
dimArgs Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
implicitDims
      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
dimArgs' Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
letArgs
  pat' <- withArgs dimArgs' $ transformPat pat
  params <- parametrizing dimArgs'
  let 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)
  body' <- withParams params $ scoping argset $ transformExp body
  res' <- transformAppRes res
  pure $ AppExp (LetPat sizes' pat' e' body' loc) (Info res')
transformAppExp LetFun {} AppRes
_ =
  [Char] -> MonoM Exp
forall a. HasCallStack => [Char] -> a
error [Char]
"transformAppExp: LetFun is not supposed to occur"
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 (Maybe VName), Exp)
args SrcLoc
_) AppRes
res =
  Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply
    (Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp)
-> MonoM Exp -> MonoM ([(Maybe VName, Exp)] -> AppRes -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
fe
    MonoM ([(Maybe VName, Exp)] -> AppRes -> Exp)
-> MonoM [(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 (Maybe VName), Exp) -> MonoM (Maybe VName, Exp))
-> [(Info (Maybe VName), Exp)] -> MonoM [(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 (Maybe VName), Exp) -> MonoM (Maybe VName, Exp)
forall {t}. (Info t, Exp) -> MonoM (t, Exp)
onArg (NonEmpty (Info (Maybe VName), Exp) -> [(Info (Maybe VName), Exp)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Info (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, Exp) -> MonoM (t, Exp)
onArg (Info t
ext, Exp
e) = (t
ext,) (Exp -> (t, Exp)) -> MonoM Exp -> MonoM (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 LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
body SrcLoc
loc) AppRes
res = do
  e1' <- Exp -> MonoM Exp
transformExp (Exp -> MonoM Exp) -> Exp -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit

  let dimArgs = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
sparams
  pat' <- withArgs dimArgs $ transformPat pat
  params <- parametrizing dimArgs
  let 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
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)

  (form', formArgs) <- case 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
      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
      (,S.fromList (patNames pat2)) . ForIn pat2' <$> transformExp 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
mergeArgs Set VName -> Set VName -> Set VName
forall a. Ord a => Set a -> Set a -> Set a
`S.union` Set VName
formArgs

  body' <- withParams params $ scoping argset $ transformExp body
  -- Maybe monomorphisation introduced new arrays to the loop, and
  -- maybe they have AnySize sizes.  This is not allowed.  Invent some
  -- sizes for them.
  (pat_sizes, pat'') <- sizesForPat pat'
  res' <- transformAppRes res
  pure $ AppExp (Loop (sparams' ++ pat_sizes) pat'' (LoopInitExplicit e1') form' body' loc) (Info 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 ret ext) <- AppRes -> MonoM AppRes
transformAppRes AppRes
res
  fname' <- transformFName loc fname (toStruct t)
  e1' <- transformExp e1
  e2' <- transformExp e2
  if orderZero (typeOf e1') && orderZero (typeOf e2')
    then pure $ applyOp ret ext fname' e1' 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.
      (x_param_e, x_param) <- makeVarParam e1'
      (y_param_e, y_param) <- makeVarParam e2'
      -- XXX: the type annotations here are wrong, but hopefully it
      -- doesn't matter as there will be an outer AppExp to handle
      -- them.
      pure $
        AppExp
          ( LetPat
              []
              x_param
              e1'
              ( AppExp
                  (LetPat [] y_param e2' (applyOp ret ext fname' x_param_e y_param_e) loc)
                  (Info $ AppRes ret mempty)
              )
              mempty
          )
          (Info (AppRes ret 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
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply
        (ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply ExpBase Info vn
fname' [(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))
        [(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
      x <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"binop_p"
      pure
        ( Var (qualName x) (Info argtype) mempty,
          Id x (Info argtype) mempty
        )
transformAppExp LetWith {} AppRes
_ =
  [Char] -> MonoM Exp
forall a. HasCallStack => [Char] -> a
error [Char]
"transformAppExp: LetWith is not supposed to occur"
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
  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
  e' <- transformExp e
  cs' <- mapM (transformCase implicitDims) cs
  res' <- transformAppRes res
  if S.null implicitDims
    then pure $ AppExp (Match e' cs' loc) (Info res')
    else do
      tmpVar <- newNameFromString "matched_variable"
      pure $
        AppExp
          ( LetPat
              (map (`SizeBinder` mempty) $ S.toList implicitDims)
              (Id tmpVar (Info $ typeOf e') mempty)
              e'
              ( AppExp
                  (Match (Var (qualName tmpVar) (Info $ typeOf e') mempty) cs' loc)
                  (Info res)
              )
              mempty
          )
          (Info 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 L Name
name Exp
e SrcLoc
loc') =
      L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit L 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 (L Loc
vloc VName
v) Info StructType
t SrcLoc
_) = do
      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
      transformField $
        RecordFieldExplicit
          (L vloc (baseName v))
          (Var (qualName v) t' loc)
          loc
transformExp (ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  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
$ [PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t 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 Exp VName
tp SrcLoc
loc) =
  Exp -> TypeExp Exp VName -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn -> TypeExp (ExpBase f vn) vn -> SrcLoc -> ExpBase f vn
Ascript (Exp -> TypeExp Exp VName -> SrcLoc -> Exp)
-> MonoM Exp -> MonoM (TypeExp Exp VName -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e MonoM (TypeExp Exp VName -> SrcLoc -> Exp)
-> MonoM (TypeExp Exp 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 Exp VName -> MonoM (TypeExp Exp VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Exp 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 Exp VName
te Info StructType
t SrcLoc
loc) =
  Exp -> TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeExp (ExpBase f vn) vn
-> f StructType
-> SrcLoc
-> ExpBase f vn
Coerce (Exp -> TypeExp Exp VName -> Info StructType -> SrcLoc -> Exp)
-> MonoM Exp
-> MonoM (TypeExp Exp 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 Exp VName -> Info StructType -> SrcLoc -> Exp)
-> MonoM (TypeExp Exp 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 Exp VName -> MonoM (TypeExp Exp VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TypeExp Exp 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 {}) =
  [Char] -> MonoM Exp
forall a. HasCallStack => [Char] -> a
error [Char]
"transformExp: Lambda is not supposed to occur"
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 RetTypeBase Exp Uniqueness
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
  e' <- Exp -> MonoM Exp
transformExp Exp
e
  desugarBinOpSection
    fname
    (Just e')
    Nothing
    t
    (xp, xtype, xargext)
    (yp, ytype, Nothing)
    (rettype, retext)
    loc
transformExp (OpSectionRight QualName VName
fname (Info StructType
t) Exp
e (Info (PName, ParamType), Info (PName, ParamType, Maybe VName))
arg (Info RetTypeBase Exp Uniqueness
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
  e' <- Exp -> MonoM Exp
transformExp Exp
e
  desugarBinOpSection
    fname
    Nothing
    (Just e')
    t
    (xp, xtype, Nothing)
    (yp, ytype, yargext)
    (rettype, [])
    loc
transformExp (ProjectSection [Name]
fields (Info StructType
t) SrcLoc
loc) = do
  t' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
t
  desugarProjectSection fields t' loc
transformExp (IndexSection SliceBase Info VName
idxs (Info StructType
t) SrcLoc
loc) = do
  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
  desugarIndexSection idxs' t loc
transformExp (Project Name
n Exp
e Info StructType
tp SrcLoc
loc) = do
  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
  e' <- transformExp e
  pure $ Project n e' tp' 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
  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
  CasePat p' <$> scoping (S.fromList (patNames p) `S.union` implicitDims) (transformExp e) <*> pure 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)
-> (RetTypeBase Exp Uniqueness, [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
  t' <- StructType -> MonoM StructType
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType StructType
t
  op <- transformFName loc fname $ toStruct t
  (v1, wrap_left, e1, p1) <- makeVarParam e_left =<< transformType xtype
  (v2, wrap_right, e2, p2) <- makeVarParam e_right =<< transformType ytype
  let apply_left =
        Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply
          Exp
op
          [(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
-> RetTypeBase Exp Uniqueness
-> 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 -> RetTypeBase Exp Uniqueness
forall dim as. [VName] -> TypeBase dim as -> RetTypeBase dim as
RetType [] (TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness)
-> TypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
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 (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' = (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
  body <-
    scoping (S.fromList [v1, v2]) $
      mkApply apply_left [(yext, e2)]
        <$> transformAppRes (AppRes (toStruct rettype') retext)
  rettype'' <- transformRetTypeSizes (S.fromList [v1, v2]) $ RetType dims rettype'
  pure . wrap_left . wrap_right $
    Lambda (p1 ++ p2) body Nothing (Info rettype'') loc
  where
    patAndVar :: TypeBase Exp u
-> m (VName, PatBase Info VName (TypeBase Exp u), Exp)
patAndVar TypeBase Exp u
argtype = do
      x <- [Char] -> m VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newNameFromString [Char]
"x"
      pure
        ( x,
          Id x (Info argtype) mempty,
          Var (qualName x) (Info (toStruct argtype)) 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
      (v, pat, 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
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)
      pure (v, wrap, var_e, [])
    makeVarParam Maybe Exp
Nothing TypeBase Exp u
argtype = do
      (v, pat, 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
      pure (v, id, var_e, [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
  p <- [Char] -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"project_p"
  let 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
  pure $
    Lambda
      [Id p (Info $ toParam Observe t1) mempty]
      body
      Nothing
      (Info (RetType dims t2))
      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
  p <- [Char] -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => [Char] -> m VName
newVName [Char]
"index_i"
  t1' <- transformType t1
  t2' <- transformType t2
  let 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') []))
  pure $
    Lambda
      [Id p (Info $ toParam Observe t1') mempty]
      body
      Nothing
      (Info (RetType dims t2'))
      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

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, Monoid (t VName)) =>
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, Monoid (t VName)) =>
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

    -- XXX: It is intentional that we throw away the 'bound'
    -- information after looking up ExpReplacements, as there are
    -- cases (particularly including the function types that occur
    -- after lambda lifting) where some troublesome shadowing occurs,
    -- and the names are not actually locally bound. We cannot just
    -- ignore the bound information entirely, and expect that any
    -- instantiation uses only names in scope at the outer level, due
    -- to truly exotic cases like entry-lifted.fut, where the real
    -- sizes were not visible to the type checker. Arguably that is
    -- the thing that should be fixed, but it requires fiddling with
    -- the defunctorisation of size-lifted types.

    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
forall a. Monoid a => a
mempty (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
forall a. Monoid a => a
mempty 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 :: [TypeParam] -> StructType -> ExpReplacements -> InferSizeArgs
inferSizeArgs [TypeParam]
tparams StructType
bind_t ExpReplacements
bind_r StructType
t = do
  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 = 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
  mapM (tparamArg dinst) 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
          -- In some cases we infer anySizes for size arguments. This
          -- only occurs when those sizes don't actually matter (knock
          -- on wood...), but we should never actually insert anySize
          -- as a concrete argument.
          | Exp
e Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
/= Exp
anySize ->
              Exp -> MonoM Exp
replaceExp Exp
e
        Maybe Exp
_ ->
          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

-- | 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
      (ty', (_, 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
      pure (RetType (filter (`S.member` canExt) dims') ty', first (`S.difference` 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 RetTypeBase Exp Uniqueness
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'
        retT' <- (Set VName, [VName])
-> Set VName
-> RetTypeBase Exp Uniqueness
-> Writer (Set VName, Set VName) (RetTypeBase Exp Uniqueness)
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 RetTypeBase Exp Uniqueness
retT
        pure (Arrow as argName d argT retT', bimap (intros `S.union`) (const 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 RetTypeBase Exp Uniqueness
retT) =
      as
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Uniqueness
-> 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
-> RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
forall as. Set VName -> RetTypeBase Exp as -> RetTypeBase Exp as
arrowCleanRetType Set VName
paramed RetTypeBase Exp Uniqueness
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 ::
  PolyBinding ->
  MonoType ->
  MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding :: PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding (PolyBinding (Maybe EntryPoint
entry, VName
name, [TypeParam]
tparams, [Pat ParamType]
params, RetTypeBase Exp Uniqueness
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) MonoType
inst_t = 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] -> RetTypeBase Exp Uniqueness -> StructType
funType [Pat ParamType]
params RetTypeBase Exp Uniqueness
rettype
  (substs, t_shape_params) <-
    SrcLoc
-> StructType
-> MonoType
-> MonoM (Map VName StructRetType, [TypeParam])
forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> StructType
-> MonoType
-> m (Map VName StructRetType, [TypeParam])
typeSubstsM SrcLoc
loc StructType
bind_t (MonoType -> MonoM (Map VName StructRetType, [TypeParam]))
-> MonoType -> MonoM (Map VName StructRetType, [TypeParam])
forall a b. (a -> b) -> a -> b
$ MonoType -> MonoType
noNamedParams MonoType
inst_t
  let 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
$ (TypeParam -> VName) -> [TypeParam] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParam -> VName
forall vn. TypeParamBase vn -> vn
typeParamName ([TypeParam] -> [VName]) -> [TypeParam] -> [VName]
forall a b. (a -> b) -> a -> b
$ [TypeParam]
shape_params [TypeParam] -> [TypeParam] -> [TypeParam]
forall a. [a] -> [a] -> [a]
++ [TypeParam]
t_shape_params
      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 ([TypeParam] -> StructRetType -> Subst StructRetType
forall t. [TypeParam] -> t -> Subst t
Subst []) Map VName StructRetType
substs
      substStructType =
        (VName -> Maybe (Subst (RetTypeBase Exp Diet)))
-> ParamType -> ParamType
forall u.
Monoid u =>
(VName -> Maybe (Subst (RetTypeBase Exp u)))
-> TypeBase Exp u -> TypeBase Exp u
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 -> Pat ParamType)
-> [Pat ParamType] -> [Pat ParamType]
forall a b. (a -> b) -> [a] -> [b]
map ((ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
forall t. (t -> t) -> Pat t -> Pat t
substPat ParamType -> ParamType
substStructType) [Pat ParamType]
params
  params'' <- withArgs shape_names $ mapM transformPat params'
  exp_naming <- get <* put mempty

  let 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 = ((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

  rettype' <-
    withParams exp_naming $
      withArgs (args <> shape_names) $
        hardTransformRetType (applySubst (`M.lookup` substs') rettype)
  extNaming <- get <* put mempty
  scope <- S.union shape_names <$> askScope'
  let (rettype'', new_params) = arrowArg scope args arg_params rettype'
      bind_t' = TypeSubs -> StructType -> StructType
forall u.
Monoid u =>
(VName -> Maybe (Subst (RetTypeBase Exp u)))
-> TypeBase Exp u -> TypeBase Exp u
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
      mkExplicit =
        (VName -> Set VName -> Bool) -> Set VName -> VName -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip
          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. Semigroup a => a -> a -> a
<> StructType -> Set VName
mustBeExplicitInBinding StructType
bind_t')
      (shape_params_explicit, shape_params_implicit) =
        partition (mkExplicit . typeParamName) $
          shape_params ++ t_shape_params ++ map (`TypeParamDim` mempty) (S.toList new_params)
      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'' = [Pat ParamType] -> RetTypeBase Exp Uniqueness -> StructType
funType [Pat ParamType]
params'' RetTypeBase Exp Uniqueness
rettype''
      bind_r = ExpReplacements
exp_naming ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
<> ExpReplacements
extNaming
  body' <- updateExpTypes (`M.lookup` substs') body
  body'' <- withParams exp_naming' $ withArgs (shape_names <> args) $ transformExp body'
  scope' <- S.union (shape_names <> args) <$> askScope'
  body''' <-
    expReplace exp_naming' <$> (calculateDims body'' . canCalculate scope' =<< get)

  seen_before <- elem name . map (fst . fst) <$> getLifts
  name' <-
    if null tparams && isNothing entry && not seen_before
      then pure name
      else newName name

  pure
    ( name',
      -- If the function is an entry point, then it cannot possibly
      -- need any explicit size arguments (checked by type checker).
      if isJust entry
        then const $ pure []
        else inferSizeArgs shape_params_explicit bind_t'' bind_r,
      if isJust entry
        then
          toValBinding
            name'
            (shape_params_explicit ++ shape_params_implicit)
            params''
            rettype''
            (entryAssert exp_naming body''')
        else
          toValBinding
            name'
            shape_params_implicit
            (map shapeParam shape_params_explicit ++ params'')
            rettype''
            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` RetTypeBase Exp Uniqueness -> [VName]
forall dim as. RetTypeBase dim as -> [VName]
retDims RetTypeBase Exp Uniqueness
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 :: [TypeParam]
shape_params = (TypeParam -> Bool) -> [TypeParam] -> [TypeParam]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (TypeParam -> Bool) -> TypeParam -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeParam -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam) [TypeParam]
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
      ty' <- TypeBase Exp as -> MonoM (TypeBase Exp as)
forall u. TypeBase Exp u -> MonoM (TypeBase Exp u)
transformType TypeBase Exp as
ty
      unbounded <- askIntros $ fvVars $ freeInType ty'
      let dims' = Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
unbounded
      pure $ RetType (dims' <> dims) 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 :: RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness)
mapOnResRetType = RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Exp Uniqueness -> m (RetTypeBase Exp Uniqueness))
-> (RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness)
-> RetTypeBase Exp Uniqueness
-> m (RetTypeBase Exp Uniqueness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs
-> RetTypeBase Exp Uniqueness -> RetTypeBase Exp Uniqueness
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
-> [TypeParam]
-> [Pat ParamType]
-> RetTypeBase Exp Uniqueness
-> Exp
-> ValBind
toValBinding VName
name' [TypeParam]
tparams' [Pat ParamType]
params'' RetTypeBase Exp Uniqueness
rettype' Exp
body'' =
      ValBind
        { valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = EntryPoint -> Info EntryPoint
forall a. a -> Info a
Info (EntryPoint -> Info EntryPoint)
-> Maybe EntryPoint -> Maybe (Info EntryPoint)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe EntryPoint
entry,
          valBindName :: VName
valBindName = VName
name',
          valBindRetType :: Info (RetTypeBase Exp Uniqueness)
valBindRetType = RetTypeBase Exp Uniqueness -> Info (RetTypeBase Exp Uniqueness)
forall a. a -> Info a
Info RetTypeBase Exp Uniqueness
rettype',
          valBindRetDecl :: Maybe (TypeExp Exp VName)
valBindRetDecl = Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing,
          valBindTypeParams :: [TypeParam]
valBindTypeParams = [TypeParam]
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 ->
  StructType ->
  MonoType ->
  m (M.Map VName StructRetType, [TypeParam])
typeSubstsM :: forall (m :: * -> *).
MonadFreshNames m =>
SrcLoc
-> StructType
-> MonoType
-> m (Map VName StructRetType, [TypeParam])
typeSubstsM SrcLoc
loc StructType
orig_t1 MonoType
orig_t2 =
  WriterT [TypeParam] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParam])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT [TypeParam] m (Map VName StructRetType)
 -> m (Map VName StructRetType, [TypeParam]))
-> WriterT [TypeParam] m (Map VName StructRetType)
-> m (Map VName StructRetType, [TypeParam])
forall a b. (a -> b) -> a -> b
$ (Map VName StructRetType, Map Int Exp) -> Map VName StructRetType
forall a b. (a, b) -> a
fst ((Map VName StructRetType, Map Int Exp) -> Map VName StructRetType)
-> WriterT [TypeParam] m (Map VName StructRetType, Map Int Exp)
-> WriterT [TypeParam] m (Map VName StructRetType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT
  (Map VName StructRetType, Map Int Exp) (WriterT [TypeParam] m) ()
-> (Map VName StructRetType, Map Int Exp)
-> WriterT [TypeParam] m (Map VName StructRetType, Map Int Exp)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT (StructType
-> MonoType
-> StateT
     (Map VName StructRetType, Map Int Exp) (WriterT [TypeParam] m) ()
forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *}.
(MonadState (Map VName StructRetType, Map Int Exp) (t (t m)),
 MonadWriter [TypeParam] (t (t m)), MonadFreshNames m, MonadTrans t,
 MonadTrans t, Monad (t m)) =>
StructType -> MonoType -> t (t m) ()
sub StructType
orig_t1 MonoType
orig_t2) (Map VName StructRetType
forall a. Monoid a => a
mempty, Map Int Exp
forall a. Monoid a => a
mempty)
  where
    subRet :: StructType -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
subRet (Scalar (TypeVar NoUniqueness
_ QualName VName
v [TypeArg Exp]
_)) 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 {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {k} {as}.
(MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParam] (t (t m)), Ord k,
 MonadState (Map k (RetTypeBase Exp as), Map Int Exp) (t (t m)),
 Monad (t m)) =>
QualName k -> RetTypeBase MonoSize as -> t (t m) ()
addSubst QualName VName
v RetTypeBase MonoSize NoUniqueness
rt
    subRet StructType
t1 (RetType [VName]
_ MonoType
t2) =
      StructType -> MonoType -> t (t m) ()
sub StructType
t1 MonoType
t2

    sub :: StructType -> MonoType -> t (t m) ()
sub t1 :: StructType
t1@(Array NoUniqueness
_ (Shape (Exp
d1 : [Exp]
_)) ScalarTypeBase Exp NoUniqueness
_) t2 :: MonoType
t2@(Array NoUniqueness
_ (Shape (MonoSize
d2 : [MonoSize]
_)) ScalarTypeBase MonoSize NoUniqueness
_) = do
      case MonoSize
d2 of
        MonoAnon Int
i -> do
          (ts, sizes) <- t (t m) (Map VName StructRetType, Map Int Exp)
forall s (m :: * -> *). MonadState s m => m s
get
          put (ts, M.insert i d1 sizes)
        MonoSize
_ -> () -> t (t m) ()
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      StructType -> MonoType -> t (t m) ()
sub (Int -> StructType -> StructType
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 StructType
t1) (Int -> MonoType -> MonoType
forall dim as. Int -> TypeBase dim as -> TypeBase dim as
stripArray Int
1 MonoType
t2)
    sub (Scalar (TypeVar NoUniqueness
_ QualName VName
v [TypeArg Exp]
_)) 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 {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {k} {as}.
(MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParam] (t (t m)), Ord k,
 MonadState (Map k (RetTypeBase Exp as), Map Int Exp) (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 StructType
fields1)) (Scalar (Record Map Name MonoType
fields2)) =
      (StructType -> MonoType -> t (t m) ())
-> [StructType] -> [MonoType] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
        StructType -> MonoType -> t (t m) ()
sub
        (((Name, StructType) -> StructType)
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StructType) -> StructType
forall a b. (a, b) -> b
snd ([(Name, StructType)] -> [StructType])
-> [(Name, StructType)] -> [StructType]
forall a b. (a -> b) -> a -> b
$ Map Name StructType -> [(Name, StructType)]
forall a. Map Name a -> [(Name, a)]
sortFields Map Name StructType
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
_ StructType
t1a (RetType [VName]
_ TypeBase Exp Uniqueness
t1b))) (Scalar (Arrow NoUniqueness
_ PName
_ Diet
_ MonoType
t2a RetTypeBase MonoSize Uniqueness
t2b)) = do
      StructType -> MonoType -> t (t m) ()
sub StructType
t1a MonoType
t2a
      StructType -> RetTypeBase MonoSize NoUniqueness -> t (t m) ()
subRet (TypeBase Exp Uniqueness -> StructType
forall dim u. TypeBase dim u -> TypeBase dim NoUniqueness
toStruct TypeBase Exp 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 [StructType]
cs1)) (Scalar (Sum Map Name [MonoType]
cs2)) =
      ((Name, [StructType]) -> (Name, [MonoType]) -> t (t m) [()])
-> [(Name, [StructType])] -> [(Name, [MonoType])] -> t (t m) ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (Name, [StructType]) -> (Name, [MonoType]) -> t (t m) [()]
forall {a} {a}.
(a, [StructType]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (Map Name [StructType] -> [(Name, [StructType])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [StructType]
cs1) (Map Name [MonoType] -> [(Name, [MonoType])]
forall a. Map Name a -> [(Name, a)]
sortConstrs Map Name [MonoType]
cs2)
      where
        typeSubstClause :: (a, [StructType]) -> (a, [MonoType]) -> t (t m) [()]
typeSubstClause (a
_, [StructType]
ts1) (a
_, [MonoType]
ts2) = (StructType -> MonoType -> t (t m) ())
-> [StructType] -> [MonoType] -> t (t m) [()]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM StructType -> MonoType -> t (t m) ()
sub [StructType]
ts1 [MonoType]
ts2
    sub t1 :: StructType
t1@(Scalar Sum {}) MonoType
t2 = StructType -> MonoType -> t (t m) ()
sub StructType
t1 MonoType
t2
    sub StructType
t1 t2 :: MonoType
t2@(Scalar Sum {}) = StructType -> MonoType -> t (t m) ()
sub StructType
t1 MonoType
t2
    sub StructType
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:", StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
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
      (ts, sizes) <- t (t m) (Map k (RetTypeBase Exp as), Map Int Exp)
forall s (m :: * -> *). MonadState s m => m s
get
      unless (v `M.member` ts) $ do
        t' <- bitraverse onDim pure t
        put (M.insert v (RetType ext t') ts, sizes)

    onDim :: MonoSize -> t (t m) Exp
onDim (MonoKnown Int
i) = do
      (ts, sizes) <- t (t m) (a, Map Int Exp)
forall s (m :: * -> *). MonadState s m => m s
get
      case M.lookup i sizes of
        Maybe Exp
Nothing -> do
          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"
          tell [TypeParamDim d loc]
          put (ts, M.insert i (sizeFromName (qualName d) mempty) sizes)
          pure $ sizeFromName (qualName d) mempty
        Just Exp
d ->
          Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d
    onDim (MonoAnon Int
i) = do
      (_, sizes) <- t (t m) (a, Map Int Exp)
forall s (m :: * -> *). MonadState s m => m s
get
      case M.lookup i sizes of
        Maybe Exp
Nothing -> Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
anySize
        Just Exp
d -> Exp -> t (t m) Exp
forall a. a -> t (t m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d

-- Perform a given substitution on the types in a pattern.
substPat :: (t -> t) -> Pat t -> Pat t
substPat :: forall t. (t -> t) -> Pat t -> Pat t
substPat 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 ((t -> t) -> Pat t -> Pat t
forall t. (t -> t) -> Pat t -> Pat t
substPat t -> t
f) [Pat t]
pats) SrcLoc
loc
  RecordPat [(L Name, Pat t)]
fs SrcLoc
loc -> [(L Name, Pat t)] -> SrcLoc -> Pat t
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat (((L Name, Pat t) -> (L Name, Pat t))
-> [(L Name, Pat t)] -> [(L Name, Pat t)]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, Pat t) -> (L Name, Pat t)
forall {a}. (a, Pat t) -> (a, Pat t)
substField [(L Name, Pat t)]
fs) SrcLoc
loc
    where
      substField :: (a, Pat t) -> (a, Pat t)
substField (a
n, Pat t
p) = (a
n, (t -> t) -> Pat t -> Pat t
forall t. (t -> t) -> Pat t -> Pat t
substPat 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 ((t -> t) -> Pat t -> Pat t
forall t. (t -> t) -> Pat t -> Pat t
substPat 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 ((t -> t) -> Pat t -> Pat t
forall t. (t -> t) -> Pat t -> Pat t
substPat 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 Exp VName
_ SrcLoc
_ -> (t -> t) -> Pat t -> Pat t
forall t. (t -> t) -> Pat t -> Pat t
substPat 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)
entry VName
name Maybe (TypeExp Exp VName)
_ (Info RetTypeBase Exp Uniqueness
rettype) [TypeParam]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) =
  (Maybe EntryPoint, VName, [TypeParam], [Pat ParamType],
 RetTypeBase Exp Uniqueness, Exp, [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding (Info EntryPoint -> EntryPoint
forall a. Info a -> a
unInfo (Info EntryPoint -> EntryPoint)
-> Maybe (Info EntryPoint) -> Maybe EntryPoint
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Info EntryPoint)
entry, VName
name, [TypeParam]
tparams, [Pat ParamType]
params, RetTypeBase Exp Uniqueness
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)

transformValBind :: ValBind -> MonoM Env
transformValBind :: ValBind -> MonoM Env
transformValBind ValBind
valbind = do
  let valbind' :: PolyBinding
valbind' = ValBind -> PolyBinding
toPolyBinding ValBind
valbind

  Bool -> MonoM () -> MonoM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe (Info EntryPoint) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Info EntryPoint) -> Bool)
-> Maybe (Info EntryPoint) -> Bool
forall a b. (a -> b) -> a -> b
$ ValBind -> Maybe (Info EntryPoint)
forall (f :: * -> *) vn. ValBindBase f vn -> Maybe (f EntryPoint)
valBindEntryPoint ValBind
valbind) (MonoM () -> MonoM ()) -> MonoM () -> MonoM ()
forall a b. (a -> b) -> a -> b
$ do
    let t :: StructType
t =
          [Pat ParamType] -> RetTypeBase Exp Uniqueness -> StructType
funType (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind) (RetTypeBase Exp Uniqueness -> StructType)
-> RetTypeBase Exp Uniqueness -> StructType
forall a b. (a -> b) -> a -> b
$
            Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness
forall a. Info a -> a
unInfo (Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness)
-> Info (RetTypeBase Exp Uniqueness) -> RetTypeBase Exp Uniqueness
forall a b. (a -> b) -> a -> b
$
              ValBind -> Info (RetTypeBase Exp Uniqueness)
forall (f :: * -> *) vn.
ValBindBase f vn -> f (RetTypeBase Exp Uniqueness)
valBindRetType ValBind
valbind
    (name, infer, valbind'') <- PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs, ValBind)
monomorphiseBinding 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
    tell $ Seq.singleton (name, valbind'')
    addLifted (valBindName valbind) (monoType t) (name, 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 = M.singleton (valBindName valbind) valbind',
        envGlobalScope =
          if null (valBindParams valbind)
            then S.fromList $ retDims $ unInfo $ valBindRetType valbind
            else mempty
      }

transformValBinds :: [ValBind] -> MonoM ()
transformValBinds :: [ValBind] -> MonoM ()
transformValBinds [] = () -> MonoM ()
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
transformValBinds (ValBind
valbind : [ValBind]
ds) = do
  env <- ValBind -> MonoM Env
transformValBind ValBind
valbind
  localEnv env $ transformValBinds ds

-- | Monomorphise a list of top-level value bindings.
transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
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
$ [ValBind] -> MonoM ()
transformValBinds [ValBind]
decs