-- | 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,
    MonoType,
    MonoStats,
  )
where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.Function
import Data.List (intersperse, partition, sortBy)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe (isJust, isNothing)
import Data.Ord (Down (..), comparing)
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 o
i64 :: forall dim o. TypeBase dim o
i64 = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase dim o)
-> PrimType -> ScalarTypeBase dim o
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 -> String
(Int -> ReplacedExp -> ShowS)
-> (ReplacedExp -> String)
-> ([ReplacedExp] -> ShowS)
-> Show ReplacedExp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ReplacedExp -> ShowS
showsPrec :: Int -> ReplacedExp -> ShowS
$cshow :: ReplacedExp -> String
show :: ReplacedExp -> String
$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 (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isIntrinsic)
        (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

-- 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 o
bool = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
Bool
    opt :: TypeBase d NoMode
opt = [TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [TypeBase d Diet
forall dim o. TypeBase dim o
bool, TypeBase d Diet
forall dim o. TypeBase dim o
bool] (RetTypeBase d Freshness -> TypeBase d NoMode)
-> RetTypeBase d Freshness -> TypeBase d NoMode
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase d Freshness -> RetTypeBase d Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase d Freshness
forall dim o. TypeBase dim o
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
forall {d}. TypeBase d NoMode
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
forall {d}. TypeBase d NoMode
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 o. TypeBase dim o
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 o. TypeBase dim o
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 o. TypeBase dim o
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
  }

isolateNormalisation :: MonoM a -> MonoM a
isolateNormalisation :: forall a. MonoM a -> MonoM a
isolateNormalisation MonoM a
m = do
  prevRepl <- MonoM ExpReplacements
getExpReplacements
  putExpReplacements mempty
  ret <-
    local
      ( \Env
env ->
          Env
env
            { envScope = envGlobalScope env <> M.keysSet (envPolyBindings env),
              envParametrized = mempty
            }
      )
      m
  putExpReplacements 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 -> 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 {envScope = args <> envScope env}

withParams :: ExpReplacements -> MonoM a -> MonoM a
withParams :: forall a. ExpReplacements -> MonoM a -> MonoM a
withParams ExpReplacements
params = (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 {envParametrized = params <> envParametrized env}

-- 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 = M.Map (VName, MonoType) (VName, InferSizeArgs)

data MonoState = MonoState
  { MonoState -> VNameSource
sVNameSource :: !VNameSource,
    MonoState -> ExpReplacements
sExpReplacements :: ExpReplacements,
    MonoState -> [(ReplacedExp, Int)]
sAnySizes :: [(ReplacedExp, Int)],
    MonoState -> Lifts
sLifts :: Lifts,
    MonoState -> Set VName
sLiftedNames :: S.Set VName,
    MonoState -> [ValBind]
sValBinds :: [ValBind]
  }

-- The monomorphization monad.
newtype MonoM a
  = MonoM (ReaderT Env (State MonoState) 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,
      MonadState MonoState
    )

instance MonadFreshNames MonoM where
  getNameSource :: MonoM VNameSource
getNameSource = (MonoState -> VNameSource) -> MonoM VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MonoState -> VNameSource
sVNameSource
  putNameSource :: VNameSource -> MonoM ()
putNameSource VNameSource
src = (MonoState -> MonoState) -> MonoM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((MonoState -> MonoState) -> MonoM ())
-> (MonoState -> MonoState) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ \MonoState
s -> MonoState
s {sVNameSource = src}

runMonoM :: VNameSource -> MonoM () -> (([ValBind], Lifts), VNameSource)
runMonoM :: VNameSource -> MonoM () -> (([ValBind], Lifts), VNameSource)
runMonoM VNameSource
src (MonoM ReaderT Env (State MonoState) ()
m) =
  ( ([ValBind] -> [ValBind]
forall a. [a] -> [a]
reverse (MonoState -> [ValBind]
sValBinds MonoState
final_state), MonoState -> Lifts
sLifts MonoState
final_state),
    MonoState -> VNameSource
sVNameSource MonoState
final_state
  )
  where
    ((), MonoState
final_state) = State MonoState () -> MonoState -> ((), MonoState)
forall s a. State s a -> s -> (a, s)
runState (ReaderT Env (State MonoState) () -> Env -> State MonoState ()
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Env (State MonoState) ()
m Env
initial_env) MonoState
initial_state
    initial_state :: MonoState
initial_state = VNameSource
-> ExpReplacements
-> [(ReplacedExp, Int)]
-> Lifts
-> Set VName
-> [ValBind]
-> MonoState
MonoState VNameSource
src ExpReplacements
forall a. Monoid a => a
mempty [(ReplacedExp, Int)]
forall a. Monoid a => a
mempty Lifts
forall a. Monoid a => a
mempty Set VName
forall a. Monoid a => a
mempty [ValBind]
forall a. Monoid a => a
mempty
    initial_env :: Env
initial_env = 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

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

addValBind :: ValBind -> MonoM ()
addValBind :: ValBind -> MonoM ()
addValBind ValBind
funbind =
  (MonoState -> MonoState) -> MonoM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((MonoState -> MonoState) -> MonoM ())
-> (MonoState -> MonoState) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ \MonoState
s -> MonoState
s {sValBinds = funbind : sValBinds s}

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
  gets $ S.union scope . sLiftedNames

getExpReplacements :: MonoM ExpReplacements
getExpReplacements :: MonoM ExpReplacements
getExpReplacements = (MonoState -> ExpReplacements) -> MonoM ExpReplacements
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MonoState -> ExpReplacements
sExpReplacements

putExpReplacements :: ExpReplacements -> MonoM ()
putExpReplacements :: ExpReplacements -> MonoM ()
putExpReplacements ExpReplacements
x = (MonoState -> MonoState) -> MonoM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((MonoState -> MonoState) -> MonoM ())
-> (MonoState -> MonoState) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ \MonoState
s -> MonoState
s {sExpReplacements = x}

getLifts :: MonoM Lifts
getLifts :: MonoM Lifts
getLifts = (MonoState -> Lifts) -> MonoM Lifts
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MonoState -> Lifts
sLifts

addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted :: VName -> MonoType -> (VName, InferSizeArgs) -> MonoM ()
addLifted VName
fname MonoType
il (VName, InferSizeArgs)
liftf =
  (MonoState -> MonoState) -> MonoM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((MonoState -> MonoState) -> MonoM ())
-> (MonoState -> MonoState) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ \MonoState
s ->
    MonoState
s
      { sLifts = M.insert (fname, il) liftf (sLifts s),
        sLiftedNames = S.insert (fst liftf) $ sLiftedNames s
      }

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 k a. Ord k => k -> Map k a -> Maybe a
M.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

-- | 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 (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isIntrinsic) 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

-- | 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) <-
    partition (usesIntros . unReplaced . fst) <$> getExpReplacements
  putExpReplacements 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 the arguments it must be
-- applied to: its size arguments, and @()@ for a polymorphic value.
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, Eq MonoSize
Eq MonoSize =>
(MonoSize -> MonoSize -> Ordering)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> Bool)
-> (MonoSize -> MonoSize -> MonoSize)
-> (MonoSize -> MonoSize -> MonoSize)
-> Ord MonoSize
MonoSize -> MonoSize -> Bool
MonoSize -> MonoSize -> Ordering
MonoSize -> MonoSize -> MonoSize
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: MonoSize -> MonoSize -> Ordering
compare :: MonoSize -> MonoSize -> Ordering
$c< :: MonoSize -> MonoSize -> Bool
< :: MonoSize -> MonoSize -> Bool
$c<= :: MonoSize -> MonoSize -> Bool
<= :: MonoSize -> MonoSize -> Bool
$c> :: MonoSize -> MonoSize -> Bool
> :: MonoSize -> MonoSize -> Bool
$c>= :: MonoSize -> MonoSize -> Bool
>= :: MonoSize -> MonoSize -> Bool
$cmax :: MonoSize -> MonoSize -> MonoSize
max :: MonoSize -> MonoSize -> MonoSize
$cmin :: MonoSize -> MonoSize -> MonoSize
min :: MonoSize -> MonoSize -> MonoSize
Ord, Int -> MonoSize -> ShowS
[MonoSize] -> ShowS
MonoSize -> String
(Int -> MonoSize -> ShowS)
-> (MonoSize -> String) -> ([MonoSize] -> ShowS) -> Show MonoSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MonoSize -> ShowS
showsPrec :: Int -> MonoSize -> ShowS
$cshow :: MonoSize -> String
show :: MonoSize -> String
$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 NoMode

monoType :: TypeBase Size o -> MonoType
monoType :: forall o. TypeBase Exp o -> MonoType
monoType = MonoType -> MonoType
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts (MonoType -> MonoType)
-> (TypeBase Exp o -> MonoType) -> TypeBase Exp o -> 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 o -> State (Int, Map Exp Int) MonoType)
-> TypeBase Exp o
-> 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 o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
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 o -> StructType)
-> TypeBase Exp o
-> State (Int, Map Exp Int) MonoType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeBase Exp o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct
  where
    -- Remove exts from return types because we don't use them anymore.
    noExts :: TypeBase MonoSize o -> TypeBase MonoSize o
    noExts :: forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts (Array o
o Shape MonoSize
shape ScalarTypeBase MonoSize NoMode
t) = o
-> Shape MonoSize
-> ScalarTypeBase MonoSize NoMode
-> TypeBase MonoSize o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape MonoSize
shape (ScalarTypeBase MonoSize NoMode -> TypeBase MonoSize o)
-> ScalarTypeBase MonoSize NoMode -> TypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize NoMode -> ScalarTypeBase MonoSize NoMode
forall {o}. ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
noExtsScalar ScalarTypeBase MonoSize NoMode
t
    noExts (Scalar ScalarTypeBase MonoSize o
t) = ScalarTypeBase MonoSize o -> TypeBase MonoSize o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase MonoSize o -> TypeBase MonoSize o)
-> ScalarTypeBase MonoSize o -> TypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
forall {o}. ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
noExtsScalar ScalarTypeBase MonoSize o
t
    noExtsScalar :: ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
noExtsScalar (Record Map Name (TypeBase MonoSize o)
fs) = Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o)
-> Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ (TypeBase MonoSize o -> TypeBase MonoSize o)
-> Map Name (TypeBase MonoSize o) -> Map Name (TypeBase MonoSize o)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase MonoSize o -> TypeBase MonoSize o
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts Map Name (TypeBase MonoSize o)
fs
    noExtsScalar (Sum Map Name [TypeBase MonoSize o]
fs) = Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o)
-> Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ ([TypeBase MonoSize o] -> [TypeBase MonoSize o])
-> Map Name [TypeBase MonoSize o] -> Map Name [TypeBase MonoSize o]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase MonoSize o -> TypeBase MonoSize o)
-> [TypeBase MonoSize o] -> [TypeBase MonoSize o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase MonoSize o -> TypeBase MonoSize o
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts) Map Name [TypeBase MonoSize o]
fs
    noExtsScalar (Arrow o
as PName
p Diet
d MonoType
t1 (RetType [VName]
_ TypeBase MonoSize Freshness
t2)) =
      o
-> PName
-> Diet
-> MonoType
-> RetTypeBase MonoSize Freshness
-> ScalarTypeBase MonoSize o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
as PName
p Diet
d (MonoType -> MonoType
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts MonoType
t1) ([VName]
-> TypeBase MonoSize Freshness -> RetTypeBase MonoSize Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase MonoSize Freshness -> TypeBase MonoSize Freshness
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
noExts TypeBase MonoSize Freshness
t2))
    noExtsScalar ScalarTypeBase MonoSize o
t = ScalarTypeBase MonoSize o
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
          (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
MonoKnown Int
i

sizeVarName :: Exp -> Name
sizeVarName :: Exp -> Name
sizeVarName Exp
e = Name
"d<{" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Text -> Name
nameFromText (ExpBase NoInfo VName -> Text
forall a. Pretty a => a -> Text
prettyText (Exp -> ExpBase NoInfo VName
bareExp Exp
e)) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"}>"

-- | Is this size expression already normalised (a single variable or
-- constant), modulo being surrounded by some parentheses? If so, return the
-- normalised form.
maybeNormalisedSize :: Exp -> Maybe Exp
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 o. TypeBase dim o
i64) SrcLoc
loc
maybeNormalisedSize Exp
_ = Maybe Exp
forall a. Maybe a
Nothing

-- | 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 <- ReplacedExp -> ExpReplacements -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ReplacedExp
e' (ExpReplacements -> Maybe VName)
-> MonoM ExpReplacements -> MonoM (Maybe VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM ExpReplacements
getExpReplacements
      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 <- Name -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> MonoM VName) -> Name -> MonoM VName
forall a b. (a -> b) -> a -> b
$ Exp -> Name
sizeVarName Exp
e
          putExpReplacements . ((e', vn) :) =<< getExpReplacements
          pure $ sizeFromName (qualName vn) (srclocOf e)

-- | Turn a non-trivial size expression occurring in a *type* into an
-- 'anySize', rather than trying to hoist it into a named size parameter that
-- 'inferSizeArgs' would later have to (fallibly) reconstruct a concrete
-- argument for at every reference to the function. Some such expressions
-- (e.g. a product of several of the function's own parameters, as in #2230)
-- cannot in general be recovered at the point of reference, so instead of
-- pretending we can, we erase the information here and let internalisation
-- resolve it from the actual sizes of the values involved, exactly as is
-- already done for other 'anySize's (see Note [AnySize]).
--
-- Repeated occurrences of the same expression are mapped to the same
-- equivalence class, so that size equalities implied by the original
-- expression are not lost.
anySizeForType :: Exp -> MonoM Exp
anySizeForType :: Exp -> MonoM Exp
anySizeForType Exp
e = do
  let e' :: ReplacedExp
e' = Exp -> ReplacedExp
ReplacedExp Exp
e
  substs <- (MonoState -> [(ReplacedExp, Int)]) -> MonoM [(ReplacedExp, Int)]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets MonoState -> [(ReplacedExp, Int)]
sAnySizes
  case lookup e' substs of
    Just Int
i -> 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
$ Int -> Exp
anySize Int
i
    Maybe Int
Nothing -> do
      i <- VName -> Int
baseTag (VName -> Int) -> MonoM VName -> MonoM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"any_size"
      modify $ \MonoState
s -> MonoState
s {sAnySizes = (e', i) : substs}
      pure $ anySize i

transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp
transformFName SrcLoc
loc QualName VName
fname StructType
ft = do
  -- The type 't' is later matched structurally, syntactically, against the
  -- callee's own declared signature (in 'inferSizeArgs', via 'dimMapping'),
  -- so it must stay fully precise here, even if we are currently inside the
  -- return type of some outer higher-order parameter.
  t' <- StructType -> MonoM StructType
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType StructType
ft
  let mono_t = StructType -> MonoType
forall o. TypeBase Exp o -> MonoType
monoType StructType
ft
  if isIntrinsic (qualLeaf fname)
    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 Freshness -> [Exp] -> Exp
applyArgs VName
fname' (Freshness -> StructType -> TypeBase Exp Freshness
forall o. Freshness -> TypeBase Exp o -> TypeBase Exp Freshness
toRes Freshness
Nonfresh 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) <- PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs)
monomorphiseBinding PolyBinding
funbind MonoType
mono_t
          applyArgs fname' (toRes Nonfresh 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

    argType :: Exp -> ParamType
argType = Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe (StructType -> ParamType)
-> (Exp -> StructType) -> Exp -> ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> StructType
typeOf

    applyArg :: TypeBase Exp Freshness
-> ([ParamType], ExpBase Info vn)
-> ExpBase Info vn
-> ([ParamType], ExpBase Info vn)
applyArg TypeBase Exp Freshness
t ([ParamType]
arg_ts, ExpBase Info vn
f) ExpBase Info vn
arg =
      ( Int -> [ParamType] -> [ParamType]
forall a. Int -> [a] -> [a]
drop Int
1 [ParamType]
arg_ts,
        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
arg)]
          (StructType -> [VName] -> AppRes
AppRes ([ParamType] -> RetTypeBase Exp Freshness -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType (Int -> [ParamType] -> [ParamType]
forall a. Int -> [a] -> [a]
drop Int
1 [ParamType]
arg_ts) ([VName] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase Exp Freshness
t)) [])
      )

    applyArgs :: VName -> TypeBase Exp Freshness -> [Exp] -> Exp
applyArgs VName
fname' TypeBase Exp Freshness
t [Exp]
args =
      SrcLoc -> Exp -> Exp
setApplyLoc SrcLoc
loc (Exp -> Exp)
-> (([ParamType], Exp) -> Exp) -> ([ParamType], Exp) -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([ParamType], Exp) -> Exp
forall a b. (a, b) -> b
snd (([ParamType], Exp) -> Exp) -> ([ParamType], Exp) -> Exp
forall a b. (a -> b) -> a -> b
$
        (([ParamType], Exp) -> Exp -> ([ParamType], Exp))
-> ([ParamType], Exp) -> [Exp] -> ([ParamType], Exp)
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 Freshness
-> ([ParamType], Exp) -> Exp -> ([ParamType], Exp)
forall {vn}.
TypeBase Exp Freshness
-> ([ParamType], ExpBase Info vn)
-> ExpBase Info vn
-> ([ParamType], ExpBase Info vn)
applyArg TypeBase Exp Freshness
t)
          ( (Exp -> ParamType) -> [Exp] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> ParamType
argType [Exp]
args,
            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
fname')
              (StructType -> Info StructType
forall a. a -> Info a
Info ([ParamType] -> RetTypeBase Exp Freshness -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType ((Exp -> ParamType) -> [Exp] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> ParamType
argType [Exp]
args) ([VName] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] TypeBase Exp Freshness
t)))
              SrcLoc
loc
          )
          [Exp]
args

-- | General-purpose size transformation for a type, used everywhere except when
-- reconstructing the type of a function reference (see 'transformFNameType').
-- Sizes immediately in the type (before any 'Arrow' is crossed) are hoisted
-- into a reference to a fresh named size variable (see 'replaceExp'), while
-- sizes found within a higher-order (function-valued) parameter or return
-- type - i.e. anywhere 'traverseDims' reports a 'DimPos' other than
-- 'PosImmediate' - are instead erased into an 'anySize'. Whether such a
-- hoisted name ends up as an actual explicit parameter of the enclosing
-- function, an ordinary local @let@-binding, or something else again,
-- depends on what the caller of 'transformType' does with the resulting
-- 'ExpReplacements' afterwards; 'transformType' itself does not add any
-- parameters. See Note [Higher-Order Parameter Sizes].
transformType :: TypeBase Size o -> MonoM (TypeBase Size o)
transformType :: forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType = (Set VName -> DimPos -> Exp -> MonoM Exp)
-> TypeBase Exp o -> MonoM (TypeBase Exp o)
forall (f :: * -> *) fdim tdim o.
Applicative f =>
(Set VName -> DimPos -> fdim -> f tdim)
-> TypeBase fdim o -> f (TypeBase tdim o)
traverseDims Set VName -> DimPos -> Exp -> MonoM Exp
forall {p}. p -> DimPos -> Exp -> MonoM Exp
onDim
  where
    onDim :: p -> DimPos -> Exp -> MonoM Exp
onDim p
_ DimPos
pos Exp
e
      | Just Int
_ <- Exp -> Maybe Int
isAnySize Exp
e = Exp -> MonoM Exp
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
      | Bool
otherwise = do
          e' <- Exp -> MonoM Exp
transformExp Exp
e
          case maybeNormalisedSize 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
              | DimPos
PosImmediate <- DimPos
pos -> Exp -> MonoM Exp
replaceExp Exp
e'
              | Bool
otherwise -> Exp -> MonoM Exp
anySizeForType Exp
e'

-- | Like 'transformType', but replaces every non-trivial size, regardless of
-- where it occurs, with a reference to a fresh named size variable (see
-- 'replaceExp'), rather than erasing sizes found within higher-order
-- parameter/return types into an 'anySize'. Note that this does not turn any
-- of these sizes into an actual parameter of anything: the resulting type has
-- exactly the same 'Arrow' structure (and hence the same arity) as the input.
-- The named variables are merely placeholders, used only so that
-- 'dimMapping' can structurally line this reconstructed type up against the
-- callee's own declared signature - whose *own* analogous sizes did, when
-- the callee itself was monomorphised, become real explicit parameters - so
-- that 'inferSizeArgs' can read off a concrete argument expression for each
-- of those parameters. Used only to reconstruct the type of a plain function
-- *reference* (in 'transformFName'); that is why it must stay fully precise,
-- even through further nested 'Arrow's - which, for this particular type,
-- are simply the callee's own remaining curried parameters, not necessarily
-- higher-order values. This also means we cannot use 'traverseDims' here
-- (unlike in 'transformType'): a size hoisted from deep within this type may
-- refer to a curried parameter bound by some enclosing 'Arrow' in the very
-- same type, so each 'Arrow' we cross must, before returning, re-scope any
-- such size locally to its own return type (via 'transformRetTypeSizesWith').
-- See Note [Higher-Order Parameter Sizes].
transformFNameType :: TypeBase Size o -> MonoM (TypeBase Size o)
transformFNameType :: forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType TypeBase Exp o
typ =
  case TypeBase Exp o
typ of
    Scalar ScalarTypeBase Exp o
scalar ->
      ScalarTypeBase Exp o -> TypeBase Exp o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o -> TypeBase Exp o)
-> MonoM (ScalarTypeBase Exp o) -> MonoM (TypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ScalarTypeBase Exp o -> MonoM (ScalarTypeBase Exp o)
forall o. ScalarTypeBase Exp o -> MonoM (ScalarTypeBase Exp o)
transformScalarSizes ScalarTypeBase Exp o
scalar
    Array o
o Shape Exp
shape ScalarTypeBase Exp NoMode
scalar ->
      o -> Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o (Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o)
-> MonoM (Shape Exp)
-> MonoM (ScalarTypeBase Exp NoMode -> TypeBase Exp o)
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 NoMode -> TypeBase Exp o)
-> MonoM (ScalarTypeBase Exp NoMode) -> MonoM (TypeBase Exp o)
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 NoMode -> MonoM (ScalarTypeBase Exp NoMode)
forall o. ScalarTypeBase Exp o -> MonoM (ScalarTypeBase Exp o)
transformScalarSizes ScalarTypeBase Exp NoMode
scalar
  where
    transformScalarSizes :: ScalarTypeBase Size o -> MonoM (ScalarTypeBase Size o)
    transformScalarSizes :: forall o. ScalarTypeBase Exp o -> MonoM (ScalarTypeBase Exp o)
transformScalarSizes (Record Map Name (TypeBase Exp o)
fs) =
      Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> MonoM (Map Name (TypeBase Exp o))
-> MonoM (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Map Name (TypeBase Exp o) -> MonoM (Map Name (TypeBase Exp o))
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 o -> MonoM (TypeBase Exp o)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType Map Name (TypeBase Exp o)
fs
    transformScalarSizes (Sum Map Name [TypeBase Exp o]
cs) =
      Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> MonoM (Map Name [TypeBase Exp o])
-> MonoM (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase Exp o] -> MonoM [TypeBase Exp o])
-> Map Name [TypeBase Exp o] -> MonoM (Map Name [TypeBase Exp o])
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 o] -> MonoM [TypeBase Exp o])
 -> Map Name [TypeBase Exp o] -> MonoM (Map Name [TypeBase Exp o]))
-> ((TypeBase Exp o -> MonoM (TypeBase Exp o))
    -> [TypeBase Exp o] -> MonoM [TypeBase Exp o])
-> (TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Map Name [TypeBase Exp o]
-> MonoM (Map Name [TypeBase Exp o])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp o -> MonoM (TypeBase Exp o))
-> [TypeBase Exp o] -> MonoM [TypeBase Exp o]
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 o -> MonoM (TypeBase Exp o)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType Map Name [TypeBase Exp o]
cs
    transformScalarSizes (Arrow o
as PName
argName Diet
d StructType
argT RetTypeBase Exp Freshness
retT) =
      o
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Freshness
-> ScalarTypeBase Exp o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
as PName
argName Diet
d
        (StructType -> RetTypeBase Exp Freshness -> ScalarTypeBase Exp o)
-> MonoM StructType
-> MonoM (RetTypeBase Exp Freshness -> ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> MonoM StructType
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType StructType
argT
        MonoM (RetTypeBase Exp Freshness -> ScalarTypeBase Exp o)
-> MonoM (RetTypeBase Exp Freshness)
-> MonoM (ScalarTypeBase Exp o)
forall a b. MonoM (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TypeBase Exp Freshness -> MonoM (TypeBase Exp Freshness))
-> Set VName
-> RetTypeBase Exp Freshness
-> MonoM (RetTypeBase Exp Freshness)
forall o.
(TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Set VName -> RetTypeBase Exp o -> MonoM (RetTypeBase Exp o)
transformRetTypeSizesWith TypeBase Exp Freshness -> MonoM (TypeBase Exp Freshness)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType Set VName
argset RetTypeBase Exp Freshness
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 o
o QualName VName
qn [TypeArg Exp]
args) =
      o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp o)
-> MonoM [TypeArg Exp] -> MonoM (ScalarTypeBase Exp o)
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 NoMode -> 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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformFNameType StructType
ty
    transformScalarSizes ty :: ScalarTypeBase Exp o
ty@Prim {} = ScalarTypeBase Exp o -> MonoM (ScalarTypeBase Exp o)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase Exp o
ty

    onDim :: Exp -> MonoM Exp
onDim Exp
e
      | Just Int
_ <- Exp -> Maybe Int
isAnySize Exp
e = 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

-- | Transform the sizes of a return type using the given size transformer,
-- given the set of parameter names newly in scope for this return type (usually
-- the name of the 'Arrow' parameter it is the return type of). Any size hoisted
-- while transforming that turns out to depend on one of these parameters cannot
-- be a top-level named parameter (those names are not in scope outside), so it
-- is instead added to the return type's own existentially-bound sizes.
transformRetTypeSizesWith ::
  (TypeBase Size o -> MonoM (TypeBase Size o)) ->
  S.Set VName ->
  RetTypeBase Size o ->
  MonoM (RetTypeBase Size o)
transformRetTypeSizesWith :: forall o.
(TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Set VName -> RetTypeBase Exp o -> MonoM (RetTypeBase Exp o)
transformRetTypeSizesWith TypeBase Exp o -> MonoM (TypeBase Exp o)
f Set VName
argset (RetType [VName]
dims TypeBase Exp o
ty) = do
  ty' <- Set VName -> MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o)
forall a. Set VName -> MonoM a -> MonoM a
withArgs Set VName
argset (MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o))
-> MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ [VName] -> MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o)
forall a. [VName] -> MonoM a -> MonoM a
withMono [VName]
dims (MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o))
-> MonoM (TypeBase Exp o) -> MonoM (TypeBase Exp o)
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> MonoM (TypeBase Exp o)
f TypeBase Exp o
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'

transformRetTypeSizes :: S.Set VName -> RetTypeBase Size o -> MonoM (RetTypeBase Size o)
transformRetTypeSizes :: forall o.
Set VName -> RetTypeBase Exp o -> MonoM (RetTypeBase Exp o)
transformRetTypeSizes = (TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Set VName -> RetTypeBase Exp o -> MonoM (RetTypeBase Exp o)
forall o.
(TypeBase Exp o -> MonoM (TypeBase Exp o))
-> Set VName -> RetTypeBase Exp o -> MonoM (RetTypeBase Exp o)
transformRetTypeSizesWith TypeBase Exp o -> MonoM (TypeBase Exp o)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType

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 [(Int, VName)] m (Pat ParamType)
-> [(Int, VName)] -> m (Pat ParamType, [(Int, VName)])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((ParamType -> StateT [(Int, VName)] m ParamType)
-> Pat ParamType -> StateT [(Int, 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 [(Int, VName)] m Exp)
-> (Diet -> StateT [(Int, VName)] m Diet)
-> ParamType
-> StateT [(Int, 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 [(Int, VName)] m Exp
forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadState [(Int, VName)] (t m), MonadTrans t,
 MonadFreshNames m) =>
Exp -> t m Exp
onDim Diet -> StateT [(Int, VName)] m Diet
forall a. a -> StateT [(Int, VName)] m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) Pat ParamType
pat) []
  pure (map snd sizes, params')
  where
    onDim :: Exp -> t m Exp
onDim Exp
d
      | Just Int
k <- Exp -> Maybe Int
isAnySize Exp
d = do
          prev <- ([(Int, VName)] -> Maybe VName) -> t m (Maybe VName)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (([(Int, VName)] -> Maybe VName) -> t m (Maybe VName))
-> ([(Int, VName)] -> Maybe VName) -> t m (Maybe VName)
forall a b. (a -> b) -> a -> b
$ Int -> [(Int, VName)] -> Maybe VName
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Int
k
          case prev of
            Maybe VName
Nothing -> 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
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"size"
              modify ((k, v) :)
              pure $ sizeFromName (qualName v) mempty
            Just VName
v ->
              Exp -> t m Exp
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> t m Exp) -> Exp -> t m Exp
forall a b. (a -> b) -> a -> b
$ QualName VName -> SrcLoc -> Exp
sizeFromName (VName -> QualName VName
forall v. v -> QualName v
qualName VName
v) SrcLoc
forall a. Monoid a => a
mempty
      | Bool
otherwise = Exp -> t m Exp
forall a. a -> t m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d

transformAppRes :: AppRes -> MonoM AppRes
transformAppRes :: AppRes -> MonoM AppRes
transformAppRes (AppRes StructType
t [VName]
ext) =
  StructType -> [VName] -> AppRes
AppRes (StructType -> [VName] -> AppRes)
-> MonoM StructType -> MonoM ([VName] -> AppRes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructType -> MonoM StructType
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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
_ =
  String -> MonoM Exp
forall a. HasCallStack => String -> a
error String
"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
loc) AppRes
res =
  SrcLoc -> Exp -> Exp
setApplyLoc SrcLoc
loc
    (Exp -> Exp) -> MonoM Exp -> MonoM Exp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (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 o. Pat (TypeBase Exp o) -> MonoM (Pat (TypeBase Exp o))
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 <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"binop_p"
      pure
        ( Var (qualName x) (Info argtype) mempty,
          Id x (Info argtype) mempty
        )
transformAppExp LetWith {} AppRes
_ =
  String -> MonoM Exp
forall a. HasCallStack => String -> a
error String
"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 o. TypeBase Exp o -> 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 <- newVName "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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 o. TypeBase dim o -> TypeBase dim NoMode
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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 {}) =
  String -> MonoM Exp
forall a. HasCallStack => String -> a
error String
"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 Freshness
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 Freshness
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 (UpdateSection [UpdateStep Info VName]
steps (Info StructType
t) SrcLoc
loc) = do
  t' <- StructType -> MonoM StructType
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType StructType
t
  steps' <- mapM transformStep steps
  desugarUpdateSection steps' t' loc
  where
    transformStep :: UpdateStep Info VName -> MonoM (UpdateStep Info VName)
transformStep (UpdateStepSlice SliceBase Info VName
idxs) =
      SliceBase Info VName -> UpdateStep Info VName
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice (SliceBase Info VName -> UpdateStep Info VName)
-> MonoM (SliceBase Info VName) -> MonoM (UpdateStep Info VName)
forall (f :: * -> *) a b. Functor 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
    transformStep (UpdateStepField Name
f) =
      UpdateStep Info VName -> MonoM (UpdateStep Info VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep Info VName -> MonoM (UpdateStep Info VName))
-> UpdateStep Info VName -> MonoM (UpdateStep Info VName)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep Info VName
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType Info StructType
tp
  e' <- transformExp e
  pure $ Project n e' tp' loc
transformExp (Update Exp
e1 [UpdateStep Info VName]
steps Exp
e2 Info StructType
t SrcLoc
loc) =
  Exp
-> [UpdateStep Info VName]
-> Exp
-> Info StructType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
ExpBase f vn
-> [UpdateStep f vn]
-> ExpBase f vn
-> f StructType
-> SrcLoc
-> ExpBase f vn
Update
    (Exp
 -> [UpdateStep Info VName]
 -> Exp
 -> Info StructType
 -> SrcLoc
 -> Exp)
-> MonoM Exp
-> MonoM
     ([UpdateStep Info VName]
      -> Exp -> Info StructType -> SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> MonoM Exp
transformExp Exp
e1
    MonoM
  ([UpdateStep Info VName]
   -> Exp -> Info StructType -> SrcLoc -> Exp)
-> MonoM [UpdateStep Info VName]
-> 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
<*> (UpdateStep Info VName -> MonoM (UpdateStep Info VName))
-> [UpdateStep Info VName] -> MonoM [UpdateStep 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 UpdateStep Info VName -> MonoM (UpdateStep Info VName)
transformStep [UpdateStep Info VName]
steps
    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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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
  where
    transformStep :: UpdateStep Info VName -> MonoM (UpdateStep Info VName)
transformStep (UpdateStepSlice SliceBase Info VName
idxs) =
      SliceBase Info VName -> UpdateStep Info VName
forall (f :: * -> *) vn. SliceBase f vn -> UpdateStep f vn
UpdateStepSlice (SliceBase Info VName -> UpdateStep Info VName)
-> MonoM (SliceBase Info VName) -> MonoM (UpdateStep Info VName)
forall (f :: * -> *) a b. Functor 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
    transformStep (UpdateStepField Name
f) =
      UpdateStep Info VName -> MonoM (UpdateStep Info VName)
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep Info VName -> MonoM (UpdateStep Info VName))
-> UpdateStep Info VName -> MonoM (UpdateStep Info VName)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep Info VName
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
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 o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 o. Pat (TypeBase Exp o) -> MonoM (Pat (TypeBase Exp o))
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 Freshness, [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 Freshness
rettype, [VName]
retext) SrcLoc
loc = do
  t' <- StructType -> MonoM StructType
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
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 NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp NoMode -> StructType)
-> ScalarTypeBase Exp NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ NoMode
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Freshness
-> ScalarTypeBase Exp NoMode
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow NoMode
forall a. Monoid a => a
mempty PName
yp (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet ParamType
ytype) (ParamType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ParamType
ytype) ([VName] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase Exp Freshness -> RetTypeBase Exp Freshness)
-> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> TypeBase Exp Freshness
forall o. Freshness -> TypeBase Exp o -> TypeBase Exp Freshness
toRes Freshness
Nonfresh 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 Freshness -> TypeBase Exp Freshness
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 Freshness
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 o
-> m (VName, PatBase Info VName (TypeBase Exp o), Exp)
patAndVar TypeBase Exp o
argtype = do
      x <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"x"
      pure
        ( x,
          Id x (Info argtype) mempty,
          Var (qualName x) (Info (toStruct argtype)) mempty
        )

    makeVarParam :: Maybe Exp
-> TypeBase Exp o
-> m (VName, Exp -> Exp, Exp,
      [PatBase Info VName (TypeBase Exp o)])
makeVarParam (Just Exp
e) TypeBase Exp o
argtype = do
      (v, pat, var_e) <- TypeBase Exp o
-> m (VName, PatBase Info VName (TypeBase Exp o), Exp)
forall {m :: * -> *} {o}.
MonadFreshNames m =>
TypeBase Exp o
-> m (VName, PatBase Info VName (TypeBase Exp o), Exp)
patAndVar TypeBase Exp o
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 o -> StructType)
-> PatBase Info VName (TypeBase Exp o)
-> 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 o -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct PatBase Info VName (TypeBase Exp o)
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 o
argtype = do
      (v, pat, var_e) <- TypeBase Exp o
-> m (VName, PatBase Info VName (TypeBase Exp o), Exp)
forall {m :: * -> *} {o}.
MonadFreshNames m =>
TypeBase Exp o
-> m (VName, PatBase Info VName (TypeBase Exp o), Exp)
patAndVar TypeBase Exp o
argtype
      pure (v, id, var_e, [pat])

desugarUpdateSection :: [UpdateStep Info VName] -> StructType -> SrcLoc -> MonoM Exp
desugarUpdateSection :: [UpdateStep Info VName] -> StructType -> SrcLoc -> MonoM Exp
desugarUpdateSection [UpdateStep Info VName]
steps (Scalar (Arrow NoMode
_ PName
_ Diet
_ StructType
t1 (RetType [VName]
dims TypeBase Exp Freshness
t2))) SrcLoc
loc = do
  p <- Name -> MonoM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"section_p"
  let body = (Exp, StructType) -> Exp
forall a b. (a, b) -> a
fst ((Exp, StructType) -> Exp) -> (Exp, StructType) -> Exp
forall a b. (a -> b) -> a -> b
$ ((Exp, StructType) -> UpdateStep Info VName -> (Exp, StructType))
-> (Exp, StructType)
-> [UpdateStep Info VName]
-> (Exp, StructType)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Exp, StructType) -> UpdateStep Info VName -> (Exp, StructType)
forall {vn}.
(ExpBase Info vn, StructType)
-> UpdateStep Info vn -> (ExpBase Info vn, StructType)
applyStep (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, StructType
t1) [UpdateStep Info VName]
steps
  pure $
    Lambda
      [Id p (Info $ toParam Observe t1) mempty]
      body
      Nothing
      (Info (RetType dims t2))
      loc
  where
    applyStep :: (ExpBase Info vn, StructType)
-> UpdateStep Info vn -> (ExpBase Info vn, StructType)
applyStep (ExpBase Info vn
e, StructType
t) (UpdateStepField Name
field) =
      case StructType
t 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
-> ExpBase Info vn -> Info StructType -> SrcLoc -> ExpBase Info vn
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
field ExpBase Info vn
e (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t') SrcLoc
forall a. Monoid a => a
mempty, StructType
t')
        StructType
_ ->
          String -> (ExpBase Info vn, StructType)
forall a. HasCallStack => String -> a
error (String -> (ExpBase Info vn, StructType))
-> String -> (ExpBase Info vn, StructType)
forall a b. (a -> b) -> a -> b
$
            String
"desugarUpdateSection: type "
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
t
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" does not have field "
              String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
prettyString Name
field
    applyStep (ExpBase Info vn
e, StructType
t) (UpdateStepSlice SliceBase Info vn
idxs) =
      let t' :: StructType
t' = Int -> StructType -> StructType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray (SliceBase Info vn -> Int
forall {f :: * -> *} {vn}. [DimIndexBase f vn] -> Int
fixedDims SliceBase Info vn
idxs) StructType
t
          e' :: ExpBase Info vn
e' = AppExpBase Info vn -> Info AppRes -> ExpBase Info vn
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (ExpBase Info vn
-> SliceBase Info vn -> SrcLoc -> AppExpBase Info vn
forall (f :: * -> *) vn.
ExpBase f vn -> SliceBase f vn -> SrcLoc -> AppExpBase f vn
Index ExpBase Info vn
e SliceBase Info vn
idxs SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info (StructType -> [VName] -> AppRes
AppRes StructType
t' []))
       in (ExpBase Info vn
e', StructType
t')

    fixedDims :: [DimIndexBase f vn] -> Int
fixedDims = [DimIndexBase f vn] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([DimIndexBase f vn] -> Int)
-> ([DimIndexBase f vn] -> [DimIndexBase f vn])
-> [DimIndexBase f vn]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DimIndexBase f vn -> Bool)
-> [DimIndexBase f vn] -> [DimIndexBase f vn]
forall a. (a -> Bool) -> [a] -> [a]
filter DimIndexBase f vn -> Bool
forall {f :: * -> *} {vn}. DimIndexBase f vn -> Bool
isFix
    isFix :: DimIndexBase f vn -> Bool
isFix DimFix {} = Bool
True
    isFix DimIndexBase f vn
_ = Bool
False
desugarUpdateSection [UpdateStep Info VName]
_ StructType
t SrcLoc
_ = String -> MonoM Exp
forall a. HasCallStack => String -> a
error (String -> MonoM Exp) -> String -> MonoM Exp
forall a b. (a -> b) -> a -> b
$ String
"desugarUpdateSection: not a function type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
t

transformPat :: Pat (TypeBase Size o) -> MonoM (Pat (TypeBase Size o))
transformPat :: forall o. Pat (TypeBase Exp o) -> MonoM (Pat (TypeBase Exp o))
transformPat = (TypeBase Exp o -> MonoM (TypeBase Exp o))
-> PatBase Info VName (TypeBase Exp o)
-> MonoM (PatBase Info VName (TypeBase Exp o))
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 o -> MonoM (TypeBase Exp o)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType

type DimInst = M.Map VName Size

dimMapping ::
  (Monoid o) =>
  TypeBase Size o ->
  TypeBase Size o ->
  ExpReplacements ->
  ExpReplacements ->
  DimInst
dimMapping :: forall o.
Monoid o =>
TypeBase Exp o
-> TypeBase Exp o -> ExpReplacements -> ExpReplacements -> DimInst
dimMapping TypeBase Exp o
t1 TypeBase Exp o
t2 ExpReplacements
r1 ExpReplacements
r2 = State DimInst (TypeBase Exp o) -> DimInst -> DimInst
forall s a. State s a -> s -> s
execState (([VName] -> Exp -> Exp -> StateT DimInst Identity Exp)
-> TypeBase Exp o
-> TypeBase Exp o
-> State DimInst (TypeBase Exp o)
forall o (m :: * -> *) d1 d2.
(Monoid o, Monad m) =>
([VName] -> d1 -> d2 -> m d1)
-> TypeBase d1 o -> TypeBase d2 o -> m (TypeBase d1 o)
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 o
t1 TypeBase Exp o
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
forall a. Semigroup a => a -> a -> a
(<>) (ExpReplacements -> ExpReplacements -> ExpReplacements)
-> MonoM ExpReplacements
-> MonoM (ExpReplacements -> ExpReplacements)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MonoM ExpReplacements
getExpReplacements 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 o.
Monoid o =>
TypeBase Exp o
-> TypeBase Exp o -> 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.
          | Maybe Int
Nothing <- Exp -> Maybe Int
isAnySize Exp
e ->
              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 o. TypeBase MonoSize o -> TypeBase MonoSize o
f
  where
    f :: TypeBase MonoSize o -> TypeBase MonoSize o
    f :: forall o. TypeBase MonoSize o -> TypeBase MonoSize o
f (Array o
o Shape MonoSize
shape ScalarTypeBase MonoSize NoMode
t) = o
-> Shape MonoSize
-> ScalarTypeBase MonoSize NoMode
-> TypeBase MonoSize o
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o
o Shape MonoSize
shape (ScalarTypeBase MonoSize NoMode -> ScalarTypeBase MonoSize NoMode
forall {o}. ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
f' ScalarTypeBase MonoSize NoMode
t)
    f (Scalar ScalarTypeBase MonoSize o
t) = ScalarTypeBase MonoSize o -> TypeBase MonoSize o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase MonoSize o -> TypeBase MonoSize o)
-> ScalarTypeBase MonoSize o -> TypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
forall {o}. ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
f' ScalarTypeBase MonoSize o
t
    f' :: ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
    f' :: forall {o}. ScalarTypeBase MonoSize o -> ScalarTypeBase MonoSize o
f' (Record Map Name (TypeBase MonoSize o)
fs) = Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o)
-> Map Name (TypeBase MonoSize o) -> ScalarTypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ (TypeBase MonoSize o -> TypeBase MonoSize o)
-> Map Name (TypeBase MonoSize o) -> Map Name (TypeBase MonoSize o)
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 o -> TypeBase MonoSize o
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
f Map Name (TypeBase MonoSize o)
fs
    f' (Sum Map Name [TypeBase MonoSize o]
cs) = Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o)
-> Map Name [TypeBase MonoSize o] -> ScalarTypeBase MonoSize o
forall a b. (a -> b) -> a -> b
$ ([TypeBase MonoSize o] -> [TypeBase MonoSize o])
-> Map Name [TypeBase MonoSize o] -> Map Name [TypeBase MonoSize o]
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 o -> TypeBase MonoSize o)
-> [TypeBase MonoSize o] -> [TypeBase MonoSize o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase MonoSize o -> TypeBase MonoSize o
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
f) Map Name [TypeBase MonoSize o]
cs
    f' (Arrow o
o PName
_ Diet
d1 MonoType
t1 (RetType [VName]
dims TypeBase MonoSize Freshness
t2)) =
      o
-> PName
-> Diet
-> MonoType
-> RetTypeBase MonoSize Freshness
-> ScalarTypeBase MonoSize o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
o PName
Unnamed Diet
d1 (MonoType -> MonoType
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
f MonoType
t1) ([VName]
-> TypeBase MonoSize Freshness -> RetTypeBase MonoSize Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims (TypeBase MonoSize Freshness -> TypeBase MonoSize Freshness
forall o. TypeBase MonoSize o -> TypeBase MonoSize o
f TypeBase MonoSize Freshness
t2))
    f' ScalarTypeBase MonoSize o
t = ScalarTypeBase MonoSize o
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 o ->
  (RetTypeBase Size o, S.Set VName)
arrowArg :: forall o.
Set VName
-> Set VName
-> [VName]
-> RetTypeBase Exp o
-> (RetTypeBase Exp o, Set VName)
arrowArg Set VName
scope Set VName
argset [VName]
args_params RetTypeBase Exp o
rety =
  let (RetTypeBase Exp o
rety', (Set VName
funArgs, Set VName
_)) = Writer (Set VName, Set VName) (RetTypeBase Exp o)
-> (RetTypeBase Exp o, (Set VName, Set VName))
forall w a. Writer w a -> (a, w)
runWriter ((Set VName, [VName])
-> Set VName
-> RetTypeBase Exp o
-> Writer (Set VName, Set VName) (RetTypeBase Exp o)
forall o'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp o'
-> Writer (Set VName, Set VName) (RetTypeBase Exp o')
arrowArgRetType (Set VName
scope, [VName]
forall a. Monoid a => a
mempty) Set VName
argset RetTypeBase Exp o
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 o -> RetTypeBase Exp o
forall o. Set VName -> RetTypeBase Exp o -> RetTypeBase Exp o
arrowCleanRetType Set VName
new_params RetTypeBase Exp o
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 o' ->
      Writer (S.Set VName, S.Set VName) (RetTypeBase Size o')
    arrowArgRetType :: forall o'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp o'
-> Writer (Set VName, Set VName) (RetTypeBase Exp o')
arrowArgRetType (Set VName
scope', [VName]
dimsToPush) Set VName
argset' (RetType [VName]
dims TypeBase Exp o'
ty) = WriterT
  (Set VName, Set VName)
  Identity
  (RetTypeBase Exp o',
   (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp o')
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 o',
    (Set VName, Set VName) -> (Set VName, Set VName))
 -> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp o'))
-> WriterT
     (Set VName, Set VName)
     Identity
     (RetTypeBase Exp o',
      (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (RetTypeBase Exp o')
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 o')
-> WriterT
     (Set VName, Set VName)
     Identity
     (TypeBase Exp o', (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 o')
 -> WriterT
      (Set VName, Set VName)
      Identity
      (TypeBase Exp o', (Set VName, Set VName)))
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o')
-> WriterT
     (Set VName, Set VName)
     Identity
     (TypeBase Exp o', (Set VName, Set VName))
forall a b. (a -> b) -> a -> b
$ (Set VName, [VName])
-> TypeBase Exp o'
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o')
forall o'.
(Set VName, [VName])
-> TypeBase Exp o'
-> Writer (Set VName, Set VName) (TypeBase Exp o')
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 o'
ty
      pure (RetType (filter (`S.member` canExt) dims') ty', first (`S.difference` canExt))

    arrowArgScalar :: (Set VName, [VName])
-> ScalarTypeBase Exp o
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
arrowArgScalar (Set VName, [VName])
env (Record Map Name (TypeBase Exp o)
fs) =
      Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> WriterT
     (Set VName, Set VName) Identity (Map Name (TypeBase Exp o))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TypeBase Exp o
 -> WriterT (Set VName, Set VName) Identity (TypeBase Exp o))
-> Map Name (TypeBase Exp o)
-> WriterT
     (Set VName, Set VName) Identity (Map Name (TypeBase Exp o))
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 o
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o)
forall o'.
(Set VName, [VName])
-> TypeBase Exp o'
-> Writer (Set VName, Set VName) (TypeBase Exp o')
arrowArgType (Set VName, [VName])
env) Map Name (TypeBase Exp o)
fs
    arrowArgScalar (Set VName, [VName])
env (Sum Map Name [TypeBase Exp o]
cs) =
      Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp o])
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([TypeBase Exp o]
 -> WriterT (Set VName, Set VName) Identity [TypeBase Exp o])
-> Map Name [TypeBase Exp o]
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp o])
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 o]
  -> WriterT (Set VName, Set VName) Identity [TypeBase Exp o])
 -> Map Name [TypeBase Exp o]
 -> WriterT
      (Set VName, Set VName) Identity (Map Name [TypeBase Exp o]))
-> ((TypeBase Exp o
     -> WriterT (Set VName, Set VName) Identity (TypeBase Exp o))
    -> [TypeBase Exp o]
    -> WriterT (Set VName, Set VName) Identity [TypeBase Exp o])
-> (TypeBase Exp o
    -> WriterT (Set VName, Set VName) Identity (TypeBase Exp o))
-> Map Name [TypeBase Exp o]
-> WriterT
     (Set VName, Set VName) Identity (Map Name [TypeBase Exp o])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp o
 -> WriterT (Set VName, Set VName) Identity (TypeBase Exp o))
-> [TypeBase Exp o]
-> WriterT (Set VName, Set VName) Identity [TypeBase Exp o]
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 o
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o)
forall o'.
(Set VName, [VName])
-> TypeBase Exp o'
-> Writer (Set VName, Set VName) (TypeBase Exp o')
arrowArgType (Set VName, [VName])
env) Map Name [TypeBase Exp o]
cs
    arrowArgScalar (Set VName
scope', [VName]
dimsToPush) (Arrow o
as PName
argName Diet
d StructType
argT RetTypeBase Exp Freshness
retT) =
      WriterT
  (Set VName, Set VName)
  Identity
  (ScalarTypeBase Exp o,
   (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
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 o,
    (Set VName, Set VName) -> (Set VName, Set VName))
 -> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o))
-> WriterT
     (Set VName, Set VName)
     Identity
     (ScalarTypeBase Exp o,
      (Set VName, Set VName) -> (Set VName, Set VName))
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
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 (Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Bool
isIntrinsic) 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 Freshness
-> Writer (Set VName, Set VName) (RetTypeBase Exp Freshness)
forall o'.
(Set VName, [VName])
-> Set VName
-> RetTypeBase Exp o'
-> Writer (Set VName, Set VName) (RetTypeBase Exp o')
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 Freshness
retT
        pure (Arrow as argName d argT retT', bimap (intros `S.union`) (const mempty))
      where
        argset' :: Set VName
argset' = FV -> Set VName
fvVars (FV -> Set VName) -> FV -> Set VName
forall a b. (a -> b) -> a -> b
$ StructType -> FV
forall o. TypeBase Exp o -> FV
freeInType StructType
argT
        fullArgset :: Set VName
fullArgset =
          case PName
argName of
            PName
Unnamed -> Set VName
argset'
            Named VName
vn -> VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
vn Set VName
argset'
    arrowArgScalar (Set VName, [VName])
env (TypeVar o
o QualName VName
qn [TypeArg Exp]
args) =
      o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp o)
-> WriterT (Set VName, Set VName) Identity [TypeArg Exp]
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
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 NoMode -> 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 o'.
(Set VName, [VName])
-> TypeBase Exp o'
-> Writer (Set VName, Set VName) (TypeBase Exp o')
arrowArgType (Set VName, [VName])
env StructType
ty
    arrowArgScalar (Set VName, [VName])
_ ScalarTypeBase Exp o
ty = ScalarTypeBase Exp o
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
forall a. a -> WriterT (Set VName, Set VName) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ScalarTypeBase Exp o
ty

    arrowArgType ::
      (S.Set VName, [VName]) ->
      TypeBase Size o' ->
      Writer (S.Set VName, S.Set VName) (TypeBase Size o')
    arrowArgType :: forall o'.
(Set VName, [VName])
-> TypeBase Exp o'
-> Writer (Set VName, Set VName) (TypeBase Exp o')
arrowArgType (Set VName, [VName])
env (Array o'
o Shape Exp
shape ScalarTypeBase Exp NoMode
scalar) =
      o' -> Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o'
forall dim o.
o -> Shape dim -> ScalarTypeBase dim NoMode -> TypeBase dim o
Array o'
o (Shape Exp -> ScalarTypeBase Exp NoMode -> TypeBase Exp o')
-> WriterT (Set VName, Set VName) Identity (Shape Exp)
-> WriterT
     (Set VName, Set VName)
     Identity
     (ScalarTypeBase Exp NoMode -> TypeBase Exp o')
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 NoMode -> TypeBase Exp o')
-> WriterT
     (Set VName, Set VName) Identity (ScalarTypeBase Exp NoMode)
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o')
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 NoMode
-> WriterT
     (Set VName, Set VName) Identity (ScalarTypeBase Exp NoMode)
forall {o}.
(Set VName, [VName])
-> ScalarTypeBase Exp o
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
arrowArgScalar (Set VName, [VName])
env ScalarTypeBase Exp NoMode
scalar
    arrowArgType (Set VName, [VName])
env (Scalar ScalarTypeBase Exp o'
ty) =
      ScalarTypeBase Exp o' -> TypeBase Exp o'
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp o' -> TypeBase Exp o')
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o')
-> WriterT (Set VName, Set VName) Identity (TypeBase Exp o')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Set VName, [VName])
-> ScalarTypeBase Exp o'
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o')
forall {o}.
(Set VName, [VName])
-> ScalarTypeBase Exp o
-> WriterT (Set VName, Set VName) Identity (ScalarTypeBase Exp o)
arrowArgScalar (Set VName, [VName])
env ScalarTypeBase Exp o'
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 o -> RetTypeBase Size o
    arrowCleanRetType :: forall o. Set VName -> RetTypeBase Exp o -> RetTypeBase Exp o
arrowCleanRetType Set VName
paramed (RetType [VName]
dims TypeBase Exp o
ty) =
      [VName] -> TypeBase Exp o -> RetTypeBase Exp o
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
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 o -> TypeBase Exp o
forall o. Set VName -> TypeBase Exp o -> TypeBase Exp o
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 o
ty)

    arrowCleanScalar :: S.Set VName -> ScalarTypeBase Size o -> ScalarTypeBase Size o
    arrowCleanScalar :: forall o. Set VName -> ScalarTypeBase Exp o -> ScalarTypeBase Exp o
arrowCleanScalar Set VName
paramed (Record Map Name (TypeBase Exp o)
fs) =
      Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o)
-> Map Name (TypeBase Exp o) -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp o -> TypeBase Exp o)
-> Map Name (TypeBase Exp o) -> Map Name (TypeBase Exp o)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (Set VName -> TypeBase Exp o -> TypeBase Exp o
forall o. Set VName -> TypeBase Exp o -> TypeBase Exp o
arrowCleanType Set VName
paramed) Map Name (TypeBase Exp o)
fs
    arrowCleanScalar Set VName
paramed (Sum Map Name [TypeBase Exp o]
cs) =
      Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o)
-> Map Name [TypeBase Exp o] -> ScalarTypeBase Exp o
forall a b. (a -> b) -> a -> b
$ (([TypeBase Exp o] -> [TypeBase Exp o])
-> Map Name [TypeBase Exp o] -> Map Name [TypeBase Exp o]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (([TypeBase Exp o] -> [TypeBase Exp o])
 -> Map Name [TypeBase Exp o] -> Map Name [TypeBase Exp o])
-> ((TypeBase Exp o -> TypeBase Exp o)
    -> [TypeBase Exp o] -> [TypeBase Exp o])
-> (TypeBase Exp o -> TypeBase Exp o)
-> Map Name [TypeBase Exp o]
-> Map Name [TypeBase Exp o]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeBase Exp o -> TypeBase Exp o)
-> [TypeBase Exp o] -> [TypeBase Exp o]
forall a b. (a -> b) -> [a] -> [b]
map) (Set VName -> TypeBase Exp o -> TypeBase Exp o
forall o. Set VName -> TypeBase Exp o -> TypeBase Exp o
arrowCleanType Set VName
paramed) Map Name [TypeBase Exp o]
cs
    arrowCleanScalar Set VName
paramed (Arrow o
as PName
argName Diet
d StructType
argT RetTypeBase Exp Freshness
retT) =
      o
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Freshness
-> ScalarTypeBase Exp o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
as PName
argName Diet
d StructType
argT (Set VName -> RetTypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall o. Set VName -> RetTypeBase Exp o -> RetTypeBase Exp o
arrowCleanRetType Set VName
paramed RetTypeBase Exp Freshness
retT)
    arrowCleanScalar Set VName
paramed (TypeVar o
o QualName VName
qn [TypeArg Exp]
args) =
      o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
qn ([TypeArg Exp] -> ScalarTypeBase Exp o)
-> [TypeArg Exp] -> ScalarTypeBase Exp o
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 NoMode -> TypeArg dim
TypeArgType (StructType -> TypeArg Exp) -> StructType -> TypeArg Exp
forall a b. (a -> b) -> a -> b
$ Set VName -> StructType -> StructType
forall o. Set VName -> TypeBase Exp o -> TypeBase Exp o
arrowCleanType Set VName
paramed StructType
ty
    arrowCleanScalar Set VName
_ ScalarTypeBase Exp o
ty = ScalarTypeBase Exp o
ty

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

removeEntryPoint :: PolyBinding -> PolyBinding
removeEntryPoint :: PolyBinding -> PolyBinding
removeEntryPoint (PolyBinding (Maybe EntryPoint
_, VName
name, [TypeParam]
tparams, [Pat ParamType]
params, RetTypeBase Exp Freshness
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) =
  (Maybe EntryPoint, VName, [TypeParam], [Pat ParamType],
 RetTypeBase Exp Freshness, Exp, [AttrInfo VName], SrcLoc)
-> PolyBinding
PolyBinding (Maybe EntryPoint
forall a. Maybe a
Nothing, VName
name, [TypeParam]
tparams, [Pat ParamType]
params, RetTypeBase Exp Freshness
rettype, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)

-- | Copy freshness from the instantiated type into the result of the declared
-- one. See 'freshenFromInst'.
freshenAsType :: TypeBase d Freshness -> TypeBase Size u -> TypeBase Size u
freshenAsType :: forall d u.
TypeBase d Freshness -> TypeBase Exp u -> TypeBase Exp u
freshenAsType
  (Scalar (Arrow Freshness
_ PName
_ Diet
_ TypeBase d NoMode
_ (RetType [VName]
_ TypeBase d Freshness
ir)))
  (Scalar (Arrow u
u PName
pn Diet
d StructType
a (RetType [VName]
ext TypeBase Exp Freshness
r))) =
    ScalarTypeBase Exp u -> TypeBase Exp u
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp u -> TypeBase Exp u)
-> ScalarTypeBase Exp u -> TypeBase Exp u
forall a b. (a -> b) -> a -> b
$ u
-> PName
-> Diet
-> StructType
-> RetTypeBase Exp Freshness
-> ScalarTypeBase Exp u
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow u
u PName
pn Diet
d StructType
a (RetTypeBase Exp Freshness -> ScalarTypeBase Exp u)
-> RetTypeBase Exp Freshness -> ScalarTypeBase Exp u
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ext (TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
freshenAs TypeBase d Freshness
ir TypeBase Exp Freshness
r)
freshenAsType TypeBase d Freshness
_ TypeBase Exp u
t = TypeBase Exp u
t

freshenAs :: TypeBase d Freshness -> ResType -> ResType
freshenAs :: forall d.
TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
freshenAs it :: TypeBase d Freshness
it@(Scalar Arrow {}) TypeBase Exp Freshness
t = TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall d u.
TypeBase d Freshness -> TypeBase Exp u -> TypeBase Exp u
freshenAsType TypeBase d Freshness
it TypeBase Exp Freshness
t
freshenAs (Scalar (Record Map Name (TypeBase d Freshness)
ifs)) (Scalar (Record Map Name (TypeBase Exp Freshness)
fs))
  | Map Name (TypeBase d Freshness) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase d Freshness)
ifs [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name (TypeBase Exp Freshness) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp Freshness)
fs =
      ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness)
-> ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase Exp Freshness) -> ScalarTypeBase Exp Freshness
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name (TypeBase Exp Freshness) -> ScalarTypeBase Exp Freshness)
-> Map Name (TypeBase Exp Freshness)
-> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ (TypeBase d Freshness
 -> TypeBase Exp Freshness -> TypeBase Exp Freshness)
-> Map Name (TypeBase d Freshness)
-> Map Name (TypeBase Exp Freshness)
-> Map Name (TypeBase Exp Freshness)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
freshenAs Map Name (TypeBase d Freshness)
ifs Map Name (TypeBase Exp Freshness)
fs
freshenAs (Scalar (Sum Map Name [TypeBase d Freshness]
ics)) (Scalar (Sum Map Name [TypeBase Exp Freshness]
cs))
  | Map Name [TypeBase d Freshness] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase d Freshness]
ics [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name [TypeBase Exp Freshness] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Exp Freshness]
cs =
      ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness)
-> ScalarTypeBase Exp Freshness -> TypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase Exp Freshness] -> ScalarTypeBase Exp Freshness
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [TypeBase Exp Freshness] -> ScalarTypeBase Exp Freshness)
-> Map Name [TypeBase Exp Freshness]
-> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ ([TypeBase d Freshness]
 -> [TypeBase Exp Freshness] -> [TypeBase Exp Freshness])
-> Map Name [TypeBase d Freshness]
-> Map Name [TypeBase Exp Freshness]
-> Map Name [TypeBase Exp Freshness]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeBase d Freshness
 -> TypeBase Exp Freshness -> TypeBase Exp Freshness)
-> [TypeBase d Freshness]
-> [TypeBase Exp Freshness]
-> [TypeBase Exp Freshness]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
freshenAs) Map Name [TypeBase d Freshness]
ics Map Name [TypeBase Exp Freshness]
cs
freshenAs TypeBase d Freshness
it TypeBase Exp Freshness
t
  -- 'setMode' writes every node, and 'freshness' of a record is Fresh
  -- if *any* field is, so this must not be reached for compound types.
  | TypeBase d Freshness -> Bool
forall dim o. TypeBase dim o -> Bool
compound TypeBase d Freshness
it Bool -> Bool -> Bool
|| TypeBase Exp Freshness -> Bool
forall dim o. TypeBase dim o -> Bool
compound TypeBase Exp Freshness
t = TypeBase Exp Freshness
t
  | TypeBase d Freshness -> Freshness
forall shape. TypeBase shape Freshness -> Freshness
freshness TypeBase d Freshness
it Freshness -> Freshness -> Bool
forall a. Eq a => a -> a -> Bool
== Freshness
Fresh = TypeBase Exp Freshness
t TypeBase Exp Freshness -> Freshness -> TypeBase Exp Freshness
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
`setMode` Freshness
Fresh
  | Bool
otherwise = TypeBase Exp Freshness
t
  where
    compound :: TypeBase dim o -> Bool
compound (Scalar Record {}) = Bool
True
    compound (Scalar Sum {}) = Bool
True
    compound TypeBase dim o
_ = Bool
False

-- | Consumption checking may refine the instantiated type of a polymorphic name
-- to say that its result is *fresh*, which the declared type cannot express.
-- The instantiation can, so take the freshness of the result from there, past
-- the given parameters. The declared return type must already have the type
-- substitution applied: a type parameter has a single mode, so where it is
-- instantiated at a record or sum, freshness can only be copied component by
-- component once the parameter has been replaced. See Note [Parametric
-- results] in Language.Futhark.TypeChecker.Consumption.
freshenFromInst :: TypeBase d Freshness -> [Pat ParamType] -> ResRetType -> ResRetType
freshenFromInst :: forall d.
TypeBase d Freshness
-> [Pat ParamType]
-> RetTypeBase Exp Freshness
-> RetTypeBase Exp Freshness
freshenFromInst (Scalar (Arrow Freshness
_ PName
_ Diet
_ TypeBase d NoMode
_ (RetType [VName]
_ TypeBase d Freshness
ir))) (Pat ParamType
_ : [Pat ParamType]
ps) RetTypeBase Exp Freshness
rt =
  TypeBase d Freshness
-> [Pat ParamType]
-> RetTypeBase Exp Freshness
-> RetTypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> [Pat ParamType]
-> RetTypeBase Exp Freshness
-> RetTypeBase Exp Freshness
freshenFromInst TypeBase d Freshness
ir [Pat ParamType]
ps RetTypeBase Exp Freshness
rt
freshenFromInst TypeBase d Freshness
it [] (RetType [VName]
ext TypeBase Exp Freshness
t) = [VName] -> TypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ext (TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> TypeBase Exp Freshness -> TypeBase Exp Freshness
freshenAs TypeBase d Freshness
it TypeBase Exp Freshness
t)
freshenFromInst TypeBase d Freshness
_ [Pat ParamType]
_ RetTypeBase Exp Freshness
rt = RetTypeBase Exp Freshness
rt

-- | Recursion is monomorphic, so a recursive reference in the body of an
-- instance is to that same instance. Its type is the declared one, however,
-- without the freshness that 'freshenFromInst' copied into the instance, and
-- would otherwise give rise to a second instance that does not justify it.
freshenRecursive :: VName -> MonoType -> Exp -> Exp
freshenRecursive :: VName -> MonoType -> Exp -> Exp
freshenRecursive VName
name MonoType
inst_t = Exp -> Exp
onExp
  where
    onExp :: Exp -> Exp
onExp (Var QualName VName
v (Info StructType
t) SrcLoc
loc)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
name =
          QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ TypeBase MonoSize Freshness -> StructType -> StructType
forall d u.
TypeBase d Freshness -> TypeBase Exp u -> TypeBase Exp u
freshenAsType ((NoMode -> Freshness) -> MonoType -> TypeBase MonoSize Freshness
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Freshness -> NoMode -> Freshness
forall a b. a -> b -> a
const Freshness
Nonfresh) MonoType
inst_t) StructType
t) SrcLoc
loc
    onExp 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
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnExp = pure . onExp} Exp
e

-- Monomorphises the body of the function as well. Returns the fresh name of the
-- generated monomorphic function as well a function for constructing additional
-- size arguments.
monomorphiseBinding ::
  PolyBinding ->
  MonoType ->
  MonoM (VName, InferSizeArgs)
monomorphiseBinding :: PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs)
monomorphiseBinding (PolyBinding (Maybe EntryPoint
entry, VName
name, [TypeParam]
tparams, [Pat ParamType]
params0, RetTypeBase Exp Freshness
rettype0, Exp
body, [AttrInfo VName]
attrs, SrcLoc
loc)) MonoType
inst_t = MonoM (VName, InferSizeArgs) -> MonoM (VName, InferSizeArgs)
forall a. MonoM a -> MonoM a
isolateNormalisation (MonoM (VName, InferSizeArgs) -> MonoM (VName, InferSizeArgs))
-> MonoM (VName, InferSizeArgs) -> MonoM (VName, InferSizeArgs)
forall a b. (a -> b) -> a -> b
$ do
  let bind_t :: StructType
bind_t = [Pat ParamType] -> RetTypeBase Exp Freshness -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
params0 RetTypeBase Exp Freshness
rettype0
  (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 o.
Monoid o =>
(VName -> Maybe (Subst (RetTypeBase Exp o)))
-> TypeBase Exp o -> TypeBase Exp o
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 ((NoMode -> 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 -> NoMode -> 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]
params0
      rettype =
        TypeBase MonoSize Freshness
-> [Pat ParamType]
-> RetTypeBase Exp Freshness
-> RetTypeBase Exp Freshness
forall d.
TypeBase d Freshness
-> [Pat ParamType]
-> RetTypeBase Exp Freshness
-> RetTypeBase Exp Freshness
freshenFromInst
          ((NoMode -> Freshness) -> MonoType -> TypeBase MonoSize Freshness
forall b c a. (b -> c) -> TypeBase a b -> TypeBase a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Freshness -> NoMode -> Freshness
forall a b. a -> b -> a
const Freshness
Nonfresh) MonoType
inst_t)
          [Pat ParamType]
params'
          (TypeSubs -> RetTypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
substs') RetTypeBase Exp Freshness
rettype0)
  params'' <- withArgs shape_names $ mapM transformPat params'
  exp_naming <- getExpReplacements <* putExpReplacements 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]
params0
      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 rettype
  extNaming <- getExpReplacements <* putExpReplacements mempty
  scope <- S.union shape_names <$> askScope'
  let (rettype'', new_params) = arrowArg scope args arg_params rettype'
      bind_t' = TypeSubs -> StructType -> StructType
forall o.
Monoid o =>
(VName -> Maybe (Subst (RetTypeBase Exp o)))
-> TypeBase Exp o -> TypeBase Exp o
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 Freshness -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
params'' RetTypeBase Exp Freshness
rettype''
      bind_r = ExpReplacements
exp_naming ExpReplacements -> ExpReplacements -> ExpReplacements
forall a. Semigroup a => a -> a -> a
<> ExpReplacements
extNaming

  -- It is important that we record the lifted function before transforming the
  -- body, in order to handle recursion. Fortunately the "calling convention"
  -- does not depend on the body.
  seen_before <- elem name . map fst . M.keys <$> getLifts
  name' <-
    if not (any isTypeParam tparams) && isNothing entry && not seen_before
      then pure name
      else newName name

  let infer
        -- If the function is an entry point, then it cannot possibly
        -- need any explicit size arguments (checked by type checker).
        | Maybe EntryPoint -> Bool
forall a. Maybe a -> Bool
isJust Maybe EntryPoint
entry = MonoM [Exp] -> InferSizeArgs
forall a b. a -> b -> a
const (MonoM [Exp] -> InferSizeArgs) -> MonoM [Exp] -> InferSizeArgs
forall a b. (a -> b) -> a -> b
$ [Exp] -> MonoM [Exp]
forall a. a -> MonoM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        | Bool
poly_value =
            ([Exp] -> [Exp]) -> MonoM [Exp] -> MonoM [Exp]
forall a b. (a -> b) -> MonoM a -> MonoM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Exp] -> [Exp] -> [Exp]
forall a. Semigroup a => a -> a -> a
<> [[Exp] -> SrcLoc -> Exp
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [] SrcLoc
loc]) (MonoM [Exp] -> MonoM [Exp]) -> InferSizeArgs -> InferSizeArgs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [TypeParam] -> StructType -> ExpReplacements -> InferSizeArgs
inferSizeArgs [TypeParam]
shape_params_explicit StructType
bind_t'' ExpReplacements
bind_r
        | Bool
otherwise = [TypeParam] -> StructType -> ExpReplacements -> InferSizeArgs
inferSizeArgs [TypeParam]
shape_params_explicit StructType
bind_t'' ExpReplacements
bind_r

  addLifted name inst_t (name', infer)

  body' <- freshenRecursive name inst_t <$> 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' =<< getExpReplacements)

  addValBind $
    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'' ++ unit_param)
          rettype''
          body'''

  pure (name', infer)
  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 Freshness -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims RetTypeBase Exp Freshness
rettype0) (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

    -- A polymorphic value becomes a function of @()@, so that each use computes
    -- it anew, as a size-polymorphic value is a function of its sizes.
    -- Consumption checking relies on this; see Note [Parametric results].
    poly_value :: Bool
poly_value = [Pat ParamType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Pat ParamType]
params0 Bool -> Bool -> Bool
&& (TypeParam -> Bool) -> [TypeParam] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeParam -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam [TypeParam]
tparams Bool -> Bool -> Bool
&& Maybe EntryPoint -> Bool
forall a. Maybe a -> Bool
isNothing Maybe EntryPoint
entry
    unit_param :: [PatBase Info vn (TypeBase dim o)]
unit_param = [Info (TypeBase dim o) -> SrcLoc -> PatBase Info vn (TypeBase dim o)
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (TypeBase dim o -> Info (TypeBase dim o)
forall a. a -> Info a
Info (TypeBase dim o -> Info (TypeBase dim o))
-> TypeBase dim o -> Info (TypeBase dim o)
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim o) -> ScalarTypeBase dim o
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record Map Name (TypeBase dim o)
forall a. Monoid a => a
mempty) SrcLoc
loc | Bool
poly_value]

    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 o -> MonoM (RetTypeBase Exp o)
hardTransformRetType (RetType [VName]
dims TypeBase Exp o
ty) = do
      ty' <- TypeBase Exp o -> MonoM (TypeBase Exp o)
forall o. TypeBase Exp o -> MonoM (TypeBase Exp o)
transformType TypeBase Exp o
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 Freshness -> m (RetTypeBase Exp Freshness)
mapOnResRetType = RetTypeBase Exp Freshness -> m (RetTypeBase Exp Freshness)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RetTypeBase Exp Freshness -> m (RetTypeBase Exp Freshness))
-> (RetTypeBase Exp Freshness -> RetTypeBase Exp Freshness)
-> RetTypeBase Exp Freshness
-> m (RetTypeBase Exp Freshness)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeSubs -> RetTypeBase Exp Freshness -> RetTypeBase Exp Freshness
forall a. Substitutable a => TypeSubs -> a -> a
applySubst TypeSubs
substs
        }

    shapeParam :: TypeParamBase vn -> PatBase Info vn (TypeBase dim o)
shapeParam TypeParamBase vn
tp = vn
-> Info (TypeBase dim o)
-> SrcLoc
-> PatBase Info vn (TypeBase dim o)
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 o -> Info (TypeBase dim o)
forall a. a -> Info a
Info TypeBase dim o
forall dim o. TypeBase dim o
i64) (SrcLoc -> PatBase Info vn (TypeBase dim o))
-> SrcLoc -> PatBase Info vn (TypeBase dim o)
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 Freshness
-> Exp
-> ValBind
toValBinding VName
name' [TypeParam]
tparams' [Pat ParamType]
params'' RetTypeBase Exp Freshness
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',
          valBindNameLoc :: SrcLoc
valBindNameLoc = SrcLoc
forall a. Monoid a => a
mempty,
          valBindRetType :: Info (RetTypeBase Exp Freshness)
valBindRetType = RetTypeBase Exp Freshness -> Info (RetTypeBase Exp Freshness)
forall a. a -> Info a
Info RetTypeBase Exp Freshness
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 NoMode -> t (t m) ()
subRet (Scalar (TypeVar NoMode
_ QualName VName
v [TypeArg Exp]
_)) RetTypeBase MonoSize NoMode
rt =
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Bool
isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v)) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
        QualName VName -> RetTypeBase MonoSize NoMode -> t (t m) ()
forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {k} {o}.
(MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParam] (t (t m)), Ord k,
 MonadState (Map k (RetTypeBase Exp o), Map Int Exp) (t (t m)),
 Monad (t m)) =>
QualName k -> RetTypeBase MonoSize o -> t (t m) ()
addSubst QualName VName
v RetTypeBase MonoSize NoMode
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 NoMode
_ (Shape (Exp
d1 : [Exp]
_)) ScalarTypeBase Exp NoMode
_) t2 :: MonoType
t2@(Array NoMode
_ (Shape (MonoSize
d2 : [MonoSize]
_)) ScalarTypeBase MonoSize NoMode
_) = 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 o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 StructType
t1) (Int -> MonoType -> MonoType
forall dim o. Int -> TypeBase dim o -> TypeBase dim o
stripArray Int
1 MonoType
t2)
    sub (Scalar (TypeVar NoMode
_ QualName VName
v [TypeArg Exp]
_)) MonoType
t =
      Bool -> t (t m) () -> t (t m) ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (VName -> Bool
isIntrinsic (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v)) (t (t m) () -> t (t m) ()) -> t (t m) () -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
        QualName VName -> RetTypeBase MonoSize NoMode -> t (t m) ()
forall {t :: (* -> *) -> * -> *} {t :: (* -> *) -> * -> *}
       {m :: * -> *} {k} {o}.
(MonadTrans t, MonadTrans t, MonadFreshNames m,
 MonadWriter [TypeParam] (t (t m)), Ord k,
 MonadState (Map k (RetTypeBase Exp o), Map Int Exp) (t (t m)),
 Monad (t m)) =>
QualName k -> RetTypeBase MonoSize o -> t (t m) ()
addSubst QualName VName
v (RetTypeBase MonoSize NoMode -> t (t m) ())
-> RetTypeBase MonoSize NoMode -> t (t m) ()
forall a b. (a -> b) -> a -> b
$
          [VName] -> MonoType -> RetTypeBase MonoSize NoMode
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
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 NoMode
_ PName
_ Diet
_ StructType
t1a (RetType [VName]
_ TypeBase Exp Freshness
t1b))) (Scalar (Arrow NoMode
_ PName
_ Diet
_ MonoType
t2a RetTypeBase MonoSize Freshness
t2b)) = do
      StructType -> MonoType -> t (t m) ()
sub StructType
t1a MonoType
t2a
      StructType -> RetTypeBase MonoSize NoMode -> t (t m) ()
subRet (TypeBase Exp Freshness -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase Exp Freshness
t1b) ((Freshness -> NoMode)
-> RetTypeBase MonoSize Freshness -> RetTypeBase MonoSize NoMode
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 (NoMode -> Freshness -> NoMode
forall a b. a -> b -> a
const NoMode
NoMode) RetTypeBase MonoSize Freshness
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 = String -> t (t m) ()
forall a. HasCallStack => String -> a
error (String -> t (t m) ()) -> String -> t (t m) ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines [String
"typeSubstsM: mismatched types:", StructType -> String
forall a. Pretty a => a -> String
prettyString StructType
t1, MonoType -> String
forall a. Pretty a => a -> String
prettyString MonoType
t2]
    addSubst :: QualName k -> RetTypeBase MonoSize o -> t (t m) ()
addSubst (QualName [k]
_ k
v) (RetType [VName]
ext TypeBase MonoSize o
t) = do
      (ts, sizes) <- t (t m) (Map k (RetTypeBase Exp o), 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
$ Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"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 -> t (t m) Exp) -> Exp -> t (t m) Exp
forall a b. (a -> b) -> a -> b
$ Int -> Exp
anySize Int
i
        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 SrcLoc
_ Maybe (TypeExp Exp VName)
_ (Info RetTypeBase Exp Freshness
rettype) [TypeParam]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc) =
  (Maybe EntryPoint, VName, [TypeParam], [Pat ParamType],
 RetTypeBase Exp Freshness, 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 Freshness
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 Freshness -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind) (RetTypeBase Exp Freshness -> StructType)
-> RetTypeBase Exp Freshness -> StructType
forall a b. (a -> b) -> a -> b
$
            Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness
forall a. Info a -> a
unInfo (Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness)
-> Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$
              ValBind -> Info (RetTypeBase Exp Freshness)
forall (f :: * -> *) vn.
ValBindBase f vn -> f (RetTypeBase Exp Freshness)
valBindRetType ValBind
valbind
    MonoM (VName, InferSizeArgs) -> MonoM ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MonoM (VName, InferSizeArgs) -> MonoM ())
-> MonoM (VName, InferSizeArgs) -> MonoM ()
forall a b. (a -> b) -> a -> b
$ PolyBinding -> MonoType -> MonoM (VName, InferSizeArgs)
monomorphiseBinding PolyBinding
valbind' (MonoType -> MonoM (VName, InferSizeArgs))
-> MonoType -> MonoM (VName, InferSizeArgs)
forall a b. (a -> b) -> a -> b
$ StructType -> MonoType
forall o. TypeBase Exp o -> MonoType
monoType StructType
t

  let global :: Set VName
global =
        if [Pat ParamType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind)
          then [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName] -> Set VName) -> [VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ RetTypeBase Exp Freshness -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims (RetTypeBase Exp Freshness -> [VName])
-> RetTypeBase Exp Freshness -> [VName]
forall a b. (a -> b) -> a -> b
$ Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness
forall a. Info a -> a
unInfo (Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness)
-> Info (RetTypeBase Exp Freshness) -> RetTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ ValBind -> Info (RetTypeBase Exp Freshness)
forall (f :: * -> *) vn.
ValBindBase f vn -> f (RetTypeBase Exp Freshness)
valBindRetType ValBind
valbind
          else Set VName
forall a. Monoid a => a
mempty

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

  pure
    env
      { envPolyBindings =
          M.insert (valBindName valbind) (removeEntryPoint valbind') $
            envPolyBindings env,
        envGlobalScope = global <> envGlobalScope env,
        envScope = S.insert (valBindName valbind) global <> envScope env
      }

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
  local (const env) $ transformValBinds ds

-- | Statistics about which functions were monomorphised.
newtype MonoStats = MonoStats [(VName, [MonoType])]

instance Pretty MonoStats where
  pretty :: forall ann. MonoStats -> Doc ann
pretty (MonoStats [(VName, [MonoType])]
l) =
    [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
stack ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
intersperse Doc ann
"" ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
forall a b. (a -> b) -> a -> b
$ ((VName, [MonoType]) -> Doc ann)
-> [(VName, [MonoType])] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (VName, [MonoType]) -> Doc ann
forall {a}. Pretty a => (VName, [a]) -> Doc ann
pf [(VName, [MonoType])]
l
    where
      comment :: Doc ann -> Doc ann
comment = (Doc ann
"-- " Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<>)
      pf :: (VName, [a]) -> Doc ann
pf (VName
name, [a]
ts) =
        Doc ann -> Doc ann
comment (Name -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Name -> Doc ann
pretty (VName -> Name
baseName VName
name) Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> 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 (VName -> Int
baseTag VName
name))
          Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann
parens (Int -> Doc ann
forall ann. Int -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ts) Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Doc ann
"instantiations")
          Doc ann -> Doc ann -> Doc ann
forall a. Semigroup a => a -> a -> a
<> Doc ann
":"
            Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
</> [Doc ann] -> Doc ann
forall a. [Doc a] -> Doc a
stack ((a -> Doc ann) -> [a] -> [Doc ann]
forall a b. (a -> b) -> [a] -> [b]
map (Doc ann -> Doc ann
comment (Doc ann -> Doc ann) -> (a -> Doc ann) -> a -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
indent Int
2 (Doc ann -> Doc ann) -> (a -> Doc ann) -> a -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Doc ann
forall ann. a -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty) [a]
ts)

-- | Monomorphise a list of top-level value bindings.
transformProg :: (MonadFreshNames m) => [ValBind] -> m ([ValBind], MonoStats)
transformProg :: forall (m :: * -> *).
MonadFreshNames m =>
[ValBind] -> m ([ValBind], MonoStats)
transformProg [ValBind]
decs = do
  (a, b) <-
    (VNameSource -> (([ValBind], Lifts), VNameSource))
-> m ([ValBind], Lifts)
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> (([ValBind], Lifts), VNameSource))
 -> m ([ValBind], Lifts))
-> (VNameSource -> (([ValBind], Lifts), VNameSource))
-> m ([ValBind], Lifts)
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
      VNameSource -> MonoM () -> (([ValBind], Lifts), VNameSource)
runMonoM VNameSource
namesrc (MonoM () -> (([ValBind], Lifts), VNameSource))
-> MonoM () -> (([ValBind], Lifts), VNameSource)
forall a b. (a -> b) -> a -> b
$ [ValBind] -> MonoM ()
transformValBinds [ValBind]
decs
  pure
    ( toList a,
      MonoStats . sortBy (comparing (Down . length . snd)) $
        map (\NonEmpty (VName, MonoType)
l -> ((VName, MonoType) -> VName
forall a b. (a, b) -> a
fst (NonEmpty (VName, MonoType) -> (VName, MonoType)
forall a. NonEmpty a -> a
NE.head NonEmpty (VName, MonoType)
l), NonEmpty MonoType -> [MonoType]
forall a. NonEmpty a -> [a]
NE.toList (NonEmpty MonoType -> [MonoType])
-> NonEmpty MonoType -> [MonoType]
forall a b. (a -> b) -> a -> b
$ ((VName, MonoType) -> MonoType)
-> NonEmpty (VName, MonoType) -> NonEmpty MonoType
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (VName, MonoType) -> MonoType
forall a b. (a, b) -> b
snd NonEmpty (VName, MonoType)
l)) $
          NE.groupBy ((==) `on` fst) $
            sortBy (comparing fst) $
              map fst $
                M.toList b
    )

-- Note [Higher-Order Parameter Sizes]
--
-- Normally, when 'transformType' (via 'onDim') encounters a size that is some
-- non-trivial expression rather than a single variable or a constant (say, 'm *
-- n'), it hoists that expression into a fresh named size parameter, tracked via
-- 'ExpReplacements' ('replaceExp'). Any function that ends up with such a
-- parameter in its (monomorphised) signature has it turned into an explicit
-- leading argument, and every reference to that function must then be augmented
-- with a concrete argument for it. This augmentation happens in
-- 'inferSizeArgs', which tries to recover a concrete expression for the
-- parameter by structurally comparing the callee's own declared type against
-- the type at the point of reference ('dimMapping').
--
-- This works fine as long as the expression's free variables are all reliably
-- available, spelled out the same way, at every place the function is
-- referenced. But consider (#2230):
--
--   def f (n: i64) (m: i64) (g: f64 -> [m ** n]f64) = g 0
--   entry main n m = f n m (\x -> replicate (m ** n) x)
--
-- Here 'm ** n' only occurs as the return size of a *parameter* of 'f' (the
-- higher-order function 'g'), not as a size of 'f' itself. Because
-- 'transformFName' processes the bare reference to 'f' before its arguments
-- have been applied, 'dimMapping' has no way to know, at that point, that the
-- surrounding application happens to supply concrete values for 'n' and 'm'.
-- (In general it may not even be a full application - 'f' could be passed
-- onward as a value.) 'inferSizeArgs' then has no honest answer and used to
-- fall back to inserting a literal '0' as the argument, which will then be
-- treated as ground truth by the internaliser.
--
-- The same thing happens if the problematic size instead occurs in the *
-- argument* type of a returned or parameter-bound function, e.g.
--
--   def f (n: i64) (m: i64) =
--     \(g: [m ** n]f64 -> f64) -> g (replicate (m ** n) 0.0)
--   entry main n m = f n m (\x -> head x)
--
-- (see @test.fut@ in the repository root at the time of writing). 'd = m ** n'
-- is here unwitnessed by anything else in 'f's own signature, so it again
-- becomes an explicit parameter of 'f' that 'inferSizeArgs' cannot recover at
-- the unsaturated reference to 'f'.
--
-- The fix implemented here recognises that a size occurring anywhere in the
-- type of a higher-order (function-valued) parameter or return value - whether
-- describing what it consumes or what it produces - is different in kind from
-- an ordinary size: the function that declares it (here 'f') always has the
-- free variables it depends on in scope, since such a size can only be built
-- from parameters of the enclosing function itself. So nothing ever needs its
-- value threaded in externally as an explicit argument; the enclosing function
-- can always reconstruct it itself, wherever it actually needs it (to type the
-- result of calling 'g', to build an array to pass to 'g', or just to pass it
-- on to a further nested 'Arrow'). We erase the expression into a fresh
-- 'anySize' right where it occurs, instead of pretending we can reconstruct it
-- as an explicit argument (see Note [AnySize]). Since the size is never turned
-- into an explicit parameter of 'f' in the first place, there is also nothing
-- for 'inferSizeArgs' to fail to supply.
--
-- Deciding when we are in "the type of a higher-order parameter or return
-- value" is the job of the 'DimPos' that 'traverseDims' reports for each
-- size it visits: 'PosImmediate' means the size is not (yet) inside any
-- 'Arrow', while 'PosParam'/'PosReturn' mean it is on the argument or return
-- side of one (we treat the two identically). 'transformType' is thus
-- simply 'traverseDims' with a callback that hoists at 'PosImmediate' and
-- erases to an 'anySize' otherwise; unlike the hand-rolled traversal this
-- replaced, it gets the structural recursion through 'Record'/'Sum'/'Arrow'
-- for free. The one place that cannot use 'traverseDims' directly is
-- 'transformFNameType', which must stay precise through every nested
-- 'Arrow' regardless of position (see its docstring for why).
--
-- Repeated occurrences of the same original expression within one binding group
-- (e.g. the same higher-order parameter's size mentioned twice, or two
-- different higher-order parameters sharing a size) are mapped to the same
-- 'anySize' equivalence class via the 'sAnySizes' cache in 'MonoState', rather
-- than each minting a fresh one, mirroring 'sizesForPat' below. Two erasures
-- that end up in *different* equivalence classes (e.g. one introduced while
-- processing 'f's own parameter list, another introduced independently at some
-- call site's reconstructed argument type) are not a problem: 'anySize'
-- equivalence classes are only ever compared within a single binding group,
-- never across independent functions or call sites, so there is nothing to keep
-- synchronised between the two.