-- | Defunctionalization of typed, monomorphic Futhark programs without modules.
module Futhark.Internalise.Defunctionalise (transformProg) where

import Control.Monad
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State
import Data.Bifoldable (bifoldMap)
import Data.Bifunctor
import Data.Bitraversable
import Data.List (partition, sortOn)
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as M
import Data.Maybe
import Data.Set qualified as S
import Futhark.IR.Pretty ()
import Futhark.MonadFreshNames
import Futhark.Util (mapAccumLM, nubOrd)
import Language.Futhark
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types (Subst (..), applySubst)

-- | A static value stores additional information about the result of
-- defunctionalization of an expression, aside from the residual expression.
--
-- The Ord instance here is really important, as it is used for the memoisation
-- machinery that handles recursive functions.
data StaticVal
  = -- | A first-order value.  The freshness slot is meaningful only where the
    -- value is used as a parameter, and is converted at those boundaries; see
    -- 'paramTypeFromSV'.
    Dynamic ResType
  | -- | The Env is the lexical closure of the lambda.
    LambdaSV (Pat ParamType) ResRetType Exp Env
  | RecordSV [(Name, StaticVal)]
  | -- | The constructor that is actually present, plus
    -- the others that are not.
    SumSV Name [StaticVal] [(Name, [ResType])]
  | -- | The pair is the StaticVal and residual expression of this
    -- function as a whole, while the second StaticVal is its
    -- body. (Don't trust this too much, my understanding may have
    -- holes.)
    DynamicFun (Exp, StaticVal) StaticVal
  | IntrinsicSV
  | HoleSV StructType SrcLoc
  deriving (Int -> StaticVal -> ShowS
[StaticVal] -> ShowS
StaticVal -> [Char]
(Int -> StaticVal -> ShowS)
-> (StaticVal -> [Char])
-> ([StaticVal] -> ShowS)
-> Show StaticVal
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StaticVal -> ShowS
showsPrec :: Int -> StaticVal -> ShowS
$cshow :: StaticVal -> [Char]
show :: StaticVal -> [Char]
$cshowList :: [StaticVal] -> ShowS
showList :: [StaticVal] -> ShowS
Show, StaticVal -> StaticVal -> Bool
(StaticVal -> StaticVal -> Bool)
-> (StaticVal -> StaticVal -> Bool) -> Eq StaticVal
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StaticVal -> StaticVal -> Bool
== :: StaticVal -> StaticVal -> Bool
$c/= :: StaticVal -> StaticVal -> Bool
/= :: StaticVal -> StaticVal -> Bool
Eq, Eq StaticVal
Eq StaticVal =>
(StaticVal -> StaticVal -> Ordering)
-> (StaticVal -> StaticVal -> Bool)
-> (StaticVal -> StaticVal -> Bool)
-> (StaticVal -> StaticVal -> Bool)
-> (StaticVal -> StaticVal -> Bool)
-> (StaticVal -> StaticVal -> StaticVal)
-> (StaticVal -> StaticVal -> StaticVal)
-> Ord StaticVal
StaticVal -> StaticVal -> Bool
StaticVal -> StaticVal -> Ordering
StaticVal -> StaticVal -> StaticVal
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 :: StaticVal -> StaticVal -> Ordering
compare :: StaticVal -> StaticVal -> Ordering
$c< :: StaticVal -> StaticVal -> Bool
< :: StaticVal -> StaticVal -> Bool
$c<= :: StaticVal -> StaticVal -> Bool
<= :: StaticVal -> StaticVal -> Bool
$c> :: StaticVal -> StaticVal -> Bool
> :: StaticVal -> StaticVal -> Bool
$c>= :: StaticVal -> StaticVal -> Bool
>= :: StaticVal -> StaticVal -> Bool
$cmax :: StaticVal -> StaticVal -> StaticVal
max :: StaticVal -> StaticVal -> StaticVal
$cmin :: StaticVal -> StaticVal -> StaticVal
min :: StaticVal -> StaticVal -> StaticVal
Ord)

data Binding = Binding
  { -- | Just if this is a polymorphic binding that must be
    -- instantiated.
    Binding -> Maybe ([VName], StructType)
bindingType :: Maybe ([VName], StructType),
    Binding -> StaticVal
bindingSV :: StaticVal
  }
  deriving (Int -> Binding -> ShowS
[Binding] -> ShowS
Binding -> [Char]
(Int -> Binding -> ShowS)
-> (Binding -> [Char]) -> ([Binding] -> ShowS) -> Show Binding
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Binding -> ShowS
showsPrec :: Int -> Binding -> ShowS
$cshow :: Binding -> [Char]
show :: Binding -> [Char]
$cshowList :: [Binding] -> ShowS
showList :: [Binding] -> ShowS
Show, Binding -> Binding -> Bool
(Binding -> Binding -> Bool)
-> (Binding -> Binding -> Bool) -> Eq Binding
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Binding -> Binding -> Bool
== :: Binding -> Binding -> Bool
$c/= :: Binding -> Binding -> Bool
/= :: Binding -> Binding -> Bool
Eq, Eq Binding
Eq Binding =>
(Binding -> Binding -> Ordering)
-> (Binding -> Binding -> Bool)
-> (Binding -> Binding -> Bool)
-> (Binding -> Binding -> Bool)
-> (Binding -> Binding -> Bool)
-> (Binding -> Binding -> Binding)
-> (Binding -> Binding -> Binding)
-> Ord Binding
Binding -> Binding -> Bool
Binding -> Binding -> Ordering
Binding -> Binding -> Binding
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 :: Binding -> Binding -> Ordering
compare :: Binding -> Binding -> Ordering
$c< :: Binding -> Binding -> Bool
< :: Binding -> Binding -> Bool
$c<= :: Binding -> Binding -> Bool
<= :: Binding -> Binding -> Bool
$c> :: Binding -> Binding -> Bool
> :: Binding -> Binding -> Bool
$c>= :: Binding -> Binding -> Bool
>= :: Binding -> Binding -> Bool
$cmax :: Binding -> Binding -> Binding
max :: Binding -> Binding -> Binding
$cmin :: Binding -> Binding -> Binding
min :: Binding -> Binding -> Binding
Ord)

-- | Environment mapping variable names to their associated static
-- value.
type Env = M.Map VName Binding

localEnv :: (Env -> Env) -> DefM a -> DefM a
localEnv :: forall a. (Env -> Env) -> DefM a -> DefM a
localEnv Env -> Env
f = ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((Env, Env) -> (Env, Env)) -> DefM a -> DefM a)
-> ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall a b. (a -> b) -> a -> b
$ (Env -> Env) -> (Env, Env) -> (Env, Env)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Env -> Env
f

-- Even when using a "new" environment (for evaluating closures) we
-- still ram the global environment of DynamicFuns in there.
localNewEnv :: Env -> DefM a -> DefM a
localNewEnv :: forall a. Env -> DefM a -> DefM a
localNewEnv Env
env = ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((Env, Env) -> (Env, Env)) -> DefM a -> DefM a)
-> ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall a b. (a -> b) -> a -> b
$ \(Env
globals, Env
_old_env) ->
  (Env
globals, Env
globals Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env)

askEnv :: DefM Env
askEnv :: DefM Env
askEnv = ((Env, Env) -> Env) -> DefM Env
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Env, Env) -> Env
forall a b. (a, b) -> b
snd

replaceTypeSizes ::
  M.Map VName SizeSubst ->
  TypeBase Size o ->
  TypeBase Size o
replaceTypeSizes :: forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs = (Exp -> Exp) -> TypeBase Exp o -> TypeBase Exp o
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
onDim
  where
    onDim :: Exp -> Exp
onDim (Var QualName VName
v Info StructType
typ SrcLoc
loc) =
      case VName -> Map VName SizeSubst -> Maybe SizeSubst
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
        Just (SubstNamed QualName VName
v') -> QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info StructType
typ SrcLoc
loc
        Just (SubstConst Int64
d) -> Integer -> SrcLoc -> Exp
sizeFromInteger (Int64 -> Integer
forall a. Integral a => a -> Integer
toInteger Int64
d) SrcLoc
loc
        Maybe SizeSubst
Nothing -> QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v Info StructType
typ SrcLoc
loc
    onDim Exp
d = Exp
d

replaceStaticValSizes ::
  S.Set VName ->
  M.Map VName SizeSubst ->
  StaticVal ->
  StaticVal
replaceStaticValSizes :: Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv =
  case StaticVal
sv of
    StaticVal
_ | Map VName SizeSubst -> Bool
forall k a. Map k a -> Bool
M.null Map VName SizeSubst
orig_substs -> StaticVal
sv
    LambdaSV Pat ParamType
param (RetType [VName]
t_dims ResType
t) Exp
e Env
closure_env ->
      let substs :: Map VName SizeSubst
substs =
            (Map VName SizeSubst -> VName -> Map VName SizeSubst)
-> Map VName SizeSubst -> Set VName -> Map VName SizeSubst
forall b a. (b -> a -> b) -> b -> Set a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((VName -> Map VName SizeSubst -> Map VName SizeSubst)
-> Map VName SizeSubst -> VName -> Map VName SizeSubst
forall a b c. (a -> b -> c) -> b -> a -> c
flip VName -> Map VName SizeSubst -> Map VName SizeSubst
forall k a. Ord k => k -> Map k a -> Map k a
M.delete) Map VName SizeSubst
orig_substs (Set VName -> Map VName SizeSubst)
-> Set VName -> Map VName SizeSubst
forall a b. (a -> b) -> a -> b
$
              [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (Env -> [VName]
forall k a. Map k a -> [k]
M.keys Env
closure_env)
       in Pat ParamType -> ResRetType -> Exp -> Env -> StaticVal
LambdaSV
            ((ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
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 (Map VName SizeSubst -> ParamType -> ParamType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs) Pat ParamType
param)
            ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
t_dims (Map VName SizeSubst -> ResType -> ResType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs ResType
t))
            (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
            (Map VName SizeSubst -> Env -> Env
forall {k}.
Ord k =>
Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
orig_substs Env
closure_env) -- intentional
    Dynamic ResType
t ->
      ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> ResType -> ResType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
orig_substs ResType
t
    RecordSV [(Name, StaticVal)]
fs ->
      [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$ ((Name, StaticVal) -> (Name, StaticVal))
-> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall a b. (a -> b) -> [a] -> [b]
map ((StaticVal -> StaticVal) -> (Name, StaticVal) -> (Name, StaticVal)
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs)) [(Name, StaticVal)]
fs
    SumSV Name
c [StaticVal]
svs [(Name, [ResType])]
ts ->
      Name -> [StaticVal] -> [(Name, [ResType])] -> StaticVal
SumSV Name
c ((StaticVal -> StaticVal) -> [StaticVal] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs) [StaticVal]
svs) ([(Name, [ResType])] -> StaticVal)
-> [(Name, [ResType])] -> StaticVal
forall a b. (a -> b) -> a -> b
$
        ((Name, [ResType]) -> (Name, [ResType]))
-> [(Name, [ResType])] -> [(Name, [ResType])]
forall a b. (a -> b) -> [a] -> [b]
map (([ResType] -> [ResType]) -> (Name, [ResType]) -> (Name, [ResType])
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([ResType] -> [ResType])
 -> (Name, [ResType]) -> (Name, [ResType]))
-> ([ResType] -> [ResType])
-> (Name, [ResType])
-> (Name, [ResType])
forall a b. (a -> b) -> a -> b
$ (ResType -> ResType) -> [ResType] -> [ResType]
forall a b. (a -> b) -> [a] -> [b]
map ((ResType -> ResType) -> [ResType] -> [ResType])
-> (ResType -> ResType) -> [ResType] -> [ResType]
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> ResType -> ResType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
orig_substs) [(Name, [ResType])]
ts
    DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2 ->
      (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
orig_substs Exp
e, Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv1) (StaticVal -> StaticVal) -> StaticVal -> StaticVal
forall a b. (a -> b) -> a -> b
$
        Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
orig_substs StaticVal
sv2
    StaticVal
IntrinsicSV ->
      StaticVal
IntrinsicSV
    HoleSV StructType
t SrcLoc
loc ->
      StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc
  where
    tv :: Map VName SizeSubst -> ASTMapper Identity
tv Map VName SizeSubst
substs =
      ASTMapper
        { mapOnStructType :: StructType -> Identity StructType
mapOnStructType = StructType -> Identity StructType
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType -> Identity StructType)
-> (StructType -> StructType) -> StructType -> Identity StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> StructType -> StructType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs,
          mapOnParamType :: ParamType -> Identity ParamType
mapOnParamType = ParamType -> Identity ParamType
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ParamType -> Identity ParamType)
-> (ParamType -> ParamType) -> ParamType -> Identity ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> ParamType -> ParamType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs,
          mapOnResRetType :: ResRetType -> Identity ResRetType
mapOnResRetType = ResRetType -> Identity ResRetType
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure,
          mapOnExp :: Exp -> Identity Exp
mapOnExp = Exp -> Identity Exp
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Identity Exp) -> (Exp -> Exp) -> Exp -> Identity Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs,
          mapOnName :: QualName VName -> Identity (QualName VName)
mapOnName = QualName VName -> Identity (QualName VName)
forall a. a -> Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> Identity (QualName VName))
-> (QualName VName -> QualName VName)
-> QualName VName
-> Identity (QualName VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> VName) -> QualName VName -> QualName VName
forall a b. (a -> b) -> QualName a -> QualName b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs)
        }

    onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
      case VName -> Map VName SizeSubst -> Maybe SizeSubst
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
        Just (SubstNamed QualName VName
v') -> QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
        Maybe SizeSubst
_ -> VName
v

    onExp :: Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs (Var QualName VName
v Info StructType
t SrcLoc
loc) =
      case VName -> Map VName SizeSubst -> Maybe SizeSubst
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v) Map VName SizeSubst
substs of
        Just (SubstNamed QualName VName
v') ->
          QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' Info StructType
t SrcLoc
loc
        Just (SubstConst Int64
d) ->
          PrimValue -> SrcLoc -> Exp
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (Int64 -> IntValue
Int64Value (Int64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
d))) SrcLoc
loc
        Maybe SizeSubst
Nothing ->
          QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (Map VName SizeSubst -> StructType -> StructType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs (StructType -> StructType) -> Info StructType -> Info StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info StructType
t) SrcLoc
loc
    onExp Map VName SizeSubst
substs (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 (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e) TypeExp Exp VName
te (Map VName SizeSubst -> StructType -> StructType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs (StructType -> StructType) -> Info StructType -> Info StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Info StructType
t) SrcLoc
loc
    onExp Map VName SizeSubst
substs (Lambda [Pat ParamType]
params Exp
e Maybe (TypeExp Exp VName)
ret (Info (RetType [VName]
t_dims ResType
t)) SrcLoc
loc) =
      [Pat ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda
        ((Pat ParamType -> Pat ParamType)
-> [Pat ParamType] -> [Pat ParamType]
forall a b. (a -> b) -> [a] -> [b]
map ((ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
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 ((ParamType -> ParamType) -> Pat ParamType -> Pat ParamType)
-> (ParamType -> ParamType) -> Pat ParamType -> Pat ParamType
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> ParamType -> ParamType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs) [Pat ParamType]
params)
        (Map VName SizeSubst -> Exp -> Exp
onExp Map VName SizeSubst
substs Exp
e)
        Maybe (TypeExp Exp VName)
ret
        (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
t_dims (Map VName SizeSubst -> ResType -> ResType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs ResType
t)))
        SrcLoc
loc
    onExp Map VName SizeSubst
substs 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 (Map VName SizeSubst -> ASTMapper Identity
tv Map VName SizeSubst
substs) Exp
e

    onEnv :: Map VName SizeSubst -> Map k Binding -> Map k Binding
onEnv Map VName SizeSubst
substs =
      [(k, Binding)] -> Map k Binding
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
        ([(k, Binding)] -> Map k Binding)
-> (Map k Binding -> [(k, Binding)])
-> Map k Binding
-> Map k Binding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, Binding) -> (k, Binding)) -> [(k, Binding)] -> [(k, Binding)]
forall a b. (a -> b) -> [a] -> [b]
map ((Binding -> Binding) -> (k, Binding) -> (k, Binding)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs))
        ([(k, Binding)] -> [(k, Binding)])
-> (Map k Binding -> [(k, Binding)])
-> Map k Binding
-> [(k, Binding)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k Binding -> [(k, Binding)]
forall k a. Map k a -> [(k, a)]
M.toList

    onBinding :: Map VName SizeSubst -> Binding -> Binding
onBinding Map VName SizeSubst
substs (Binding Maybe ([VName], StructType)
t StaticVal
bsv) =
      Maybe ([VName], StructType) -> StaticVal -> Binding
Binding
        ((StructType -> StructType)
-> ([VName], StructType) -> ([VName], StructType)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Map VName SizeSubst -> StructType -> StructType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
substs) (([VName], StructType) -> ([VName], StructType))
-> Maybe ([VName], StructType) -> Maybe ([VName], StructType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe ([VName], StructType)
t)
        (Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
bsv)

-- | Returns the defunctionalization environment restricted
-- to the given set of variable names.
restrictEnvTo :: FV -> DefM Env
restrictEnvTo :: FV -> DefM Env
restrictEnvTo FV
fv = ((Env, Env) -> Env) -> DefM Env
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Env, Env) -> Env
forall {a}. (Map VName a, Env) -> Env
restrict
  where
    restrict :: (Map VName a, Env) -> Env
restrict (Map VName a
globals, Env
env) = (VName -> Binding -> Maybe Binding) -> Env -> Env
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
M.mapMaybeWithKey VName -> Binding -> Maybe Binding
keep Env
env
      where
        keep :: VName -> Binding -> Maybe Binding
keep VName
k (Binding Maybe ([VName], StructType)
t StaticVal
sv) = do
          Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (VName
k VName -> Map VName a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` Map VName a
globals) Bool -> Bool -> Bool
&& VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member VName
k (FV -> Set VName
fvVars FV
fv)
          Binding -> Maybe Binding
forall a. a -> Maybe a
Just (Binding -> Maybe Binding) -> Binding -> Maybe Binding
forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv
    restrict' :: StaticVal -> StaticVal
restrict' (Dynamic ResType
t) =
      ResType -> StaticVal
Dynamic ResType
t
    restrict' (LambdaSV Pat ParamType
pat ResRetType
t Exp
e Env
env) =
      Pat ParamType -> ResRetType -> Exp -> Env -> StaticVal
LambdaSV Pat ParamType
pat ResRetType
t Exp
e (Env -> StaticVal) -> Env -> StaticVal
forall a b. (a -> b) -> a -> b
$ (Binding -> Binding) -> Env -> Env
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Binding -> Binding
restrict'' Env
env
    restrict' (RecordSV [(Name, StaticVal)]
fields) =
      [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$ ((Name, StaticVal) -> (Name, StaticVal))
-> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall a b. (a -> b) -> [a] -> [b]
map ((StaticVal -> StaticVal) -> (Name, StaticVal) -> (Name, StaticVal)
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> StaticVal
restrict') [(Name, StaticVal)]
fields
    restrict' (SumSV Name
c [StaticVal]
svs [(Name, [ResType])]
fields) =
      Name -> [StaticVal] -> [(Name, [ResType])] -> StaticVal
SumSV Name
c ((StaticVal -> StaticVal) -> [StaticVal] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> StaticVal
restrict' [StaticVal]
svs) [(Name, [ResType])]
fields
    restrict' (DynamicFun (Exp
e, StaticVal
sv1) StaticVal
sv2) =
      (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp
e, StaticVal -> StaticVal
restrict' StaticVal
sv1) (StaticVal -> StaticVal) -> StaticVal -> StaticVal
forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv2
    restrict' StaticVal
IntrinsicSV = StaticVal
IntrinsicSV
    restrict' (HoleSV StructType
t SrcLoc
loc) = StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc
    restrict'' :: Binding -> Binding
restrict'' (Binding Maybe ([VName], StructType)
t StaticVal
sv) = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
t (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ StaticVal -> StaticVal
restrict' StaticVal
sv

-- | Maps a function specialisation - a function's 'LambdaSV' together with the
-- 'StaticVal' of the argument it is applied to - to the top-level function
-- lifted for it: its name, its return type, and the 'StaticVal' of the result.
-- Applying the same function to the same static argument again (in particular, a
-- recursive occurrence) reuses this lifting rather than lifting it anew, which
-- both deduplicates lifted functions and ties off recursion. See Note [Lifting
-- and recursion].
type LiftMemo = M.Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)

-- | Defunctionalization monad. The Reader environment tracks both the global
-- Env and the current Env. This is used to avoid unnecessarily large closure
-- environments (no need to capture the global one).
newtype DefM a
  = DefM (ReaderT (Env, Env) (State ([ValBind], VNameSource, LiftMemo)) a)
  deriving
    ( (forall a b. (a -> b) -> DefM a -> DefM b)
-> (forall a b. a -> DefM b -> DefM a) -> Functor DefM
forall a b. a -> DefM b -> DefM a
forall a b. (a -> b) -> DefM a -> DefM 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) -> DefM a -> DefM b
fmap :: forall a b. (a -> b) -> DefM a -> DefM b
$c<$ :: forall a b. a -> DefM b -> DefM a
<$ :: forall a b. a -> DefM b -> DefM a
Functor,
      Functor DefM
Functor DefM =>
(forall a. a -> DefM a)
-> (forall a b. DefM (a -> b) -> DefM a -> DefM b)
-> (forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c)
-> (forall a b. DefM a -> DefM b -> DefM b)
-> (forall a b. DefM a -> DefM b -> DefM a)
-> Applicative DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM 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 -> DefM a
pure :: forall a. a -> DefM a
$c<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
$cliftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
liftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
$c*> :: forall a b. DefM a -> DefM b -> DefM b
*> :: forall a b. DefM a -> DefM b -> DefM b
$c<* :: forall a b. DefM a -> DefM b -> DefM a
<* :: forall a b. DefM a -> DefM b -> DefM a
Applicative,
      Applicative DefM
Applicative DefM =>
(forall a b. DefM a -> (a -> DefM b) -> DefM b)
-> (forall a b. DefM a -> DefM b -> DefM b)
-> (forall a. a -> DefM a)
-> Monad DefM
forall a. a -> DefM a
forall a b. DefM a -> DefM b -> DefM b
forall a b. DefM a -> (a -> DefM b) -> DefM 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. DefM a -> (a -> DefM b) -> DefM b
>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
$c>> :: forall a b. DefM a -> DefM b -> DefM b
>> :: forall a b. DefM a -> DefM b -> DefM b
$creturn :: forall a. a -> DefM a
return :: forall a. a -> DefM a
Monad,
      MonadReader (Env, Env),
      MonadState ([ValBind], VNameSource, LiftMemo)
    )

instance MonadFreshNames DefM where
  putNameSource :: VNameSource -> DefM ()
putNameSource VNameSource
src = (([ValBind], VNameSource,
  Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
 -> ([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((([ValBind], VNameSource,
   Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
  -> ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
 -> DefM ())
-> (([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
    -> ([ValBind], VNameSource,
        Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall a b. (a -> b) -> a -> b
$ \([ValBind]
x, VNameSource
_, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m) -> ([ValBind]
x, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m)
  getNameSource :: DefM VNameSource
getNameSource = (([ValBind], VNameSource,
  Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
 -> VNameSource)
-> DefM VNameSource
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (\([ValBind]
_, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
_) -> VNameSource
src)

-- | Run a computation in the defunctionalization monad. Returns the result of
-- the computation, a new name source, and a list of lifted function declations.
runDefM :: VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM :: forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
src (DefM ReaderT
  (Env, Env)
  (State
     ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
  a
m) =
  let (a
x, ([ValBind]
vbs, VNameSource
src', Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
_)) = State
  ([ValBind], VNameSource,
   Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
  a
-> ([ValBind], VNameSource,
    Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
-> (a,
    ([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
forall s a. State s a -> s -> (a, s)
runState (ReaderT
  (Env, Env)
  (State
     ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
  a
-> (Env, Env)
-> State
     ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
     a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT
  (Env, Env)
  (State
     ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
  a
m (Env, Env)
forall a. Monoid a => a
mempty) ([ValBind]
forall a. Monoid a => a
mempty, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
forall a. Monoid a => a
mempty)
   in (a
x, VNameSource
src', [ValBind] -> [ValBind]
forall a. [a] -> [a]
reverse [ValBind]
vbs)

addValBind :: ValBind -> DefM ()
addValBind :: ValBind -> DefM ()
addValBind ValBind
vb = (([ValBind], VNameSource,
  Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
 -> ([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((([ValBind], VNameSource,
   Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
  -> ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
 -> DefM ())
-> (([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
    -> ([ValBind], VNameSource,
        Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall a b. (a -> b) -> a -> b
$ \([ValBind]
vbs, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m) -> (ValBind
vb ValBind -> [ValBind] -> [ValBind]
forall a. a -> [a] -> [a]
: [ValBind]
vbs, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m)

-- | The function lifted for the given specialisation, if any (see 'LiftMemo').
lookupLift :: (StaticVal, StaticVal) -> DefM (Maybe (VName, ResRetType, StaticVal))
lookupLift :: (StaticVal, StaticVal)
-> DefM (Maybe (VName, ResRetType, StaticVal))
lookupLift (StaticVal, StaticVal)
key = (([ValBind], VNameSource,
  Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
 -> Maybe (VName, ResRetType, StaticVal))
-> DefM (Maybe (VName, ResRetType, StaticVal))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((([ValBind], VNameSource,
   Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
  -> Maybe (VName, ResRetType, StaticVal))
 -> DefM (Maybe (VName, ResRetType, StaticVal)))
-> (([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
    -> Maybe (VName, ResRetType, StaticVal))
-> DefM (Maybe (VName, ResRetType, StaticVal))
forall a b. (a -> b) -> a -> b
$ \([ValBind]
_, VNameSource
_, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m) -> (StaticVal, StaticVal)
-> Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
-> Maybe (VName, ResRetType, StaticVal)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (StaticVal, StaticVal)
key Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m

-- | Record the function lifted for a specialisation (see 'LiftMemo').
insertLift :: (StaticVal, StaticVal) -> (VName, ResRetType, StaticVal) -> DefM ()
insertLift :: (StaticVal, StaticVal) -> (VName, ResRetType, StaticVal) -> DefM ()
insertLift (StaticVal, StaticVal)
key (VName, ResRetType, StaticVal)
v =
  (([ValBind], VNameSource,
  Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
 -> ([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((([ValBind], VNameSource,
   Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
  -> ([ValBind], VNameSource,
      Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
 -> DefM ())
-> (([ValBind], VNameSource,
     Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal))
    -> ([ValBind], VNameSource,
        Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)))
-> DefM ()
forall a b. (a -> b) -> a -> b
$ \([ValBind]
vbs, VNameSource
src, Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m) -> ([ValBind]
vbs, VNameSource
src, (StaticVal, StaticVal)
-> (VName, ResRetType, StaticVal)
-> Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
-> Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (StaticVal, StaticVal)
key (VName, ResRetType, StaticVal)
v Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)
m)

-- | Create a new top-level value declaration with the given function name,
-- return type, list of parameters, and body expression.
liftValDec :: VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec :: VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()
liftValDec VName
fname (RetType [VName]
ret_dims ResType
ret) [VName]
dims [Pat ParamType]
pats Exp
body = ValBind -> DefM ()
addValBind ValBind
dec
  where
    dims' :: [TypeParamBase VName]
dims' = (VName -> TypeParamBase VName) -> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> [a] -> [b]
map (VName -> SrcLoc -> TypeParamBase VName
forall vn. vn -> SrcLoc -> TypeParamBase vn
`TypeParamDim` SrcLoc
forall a. Monoid a => a
mempty) [VName]
dims
    -- FIXME: this pass is still not correctly size-preserving, so
    -- forget those return sizes that we forgot to propagate along
    -- the way.  Hopefully the internaliser is conservative and
    -- will insert reshapes...
    bound_here :: Set VName
bound_here = [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
$ [VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> (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]
pats
    mkExt :: VName -> Maybe VName
mkExt VName
v
      | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = VName -> Maybe VName
forall a. a -> Maybe a
Just VName
v
    mkExt VName
_ = Maybe VName
forall a. Maybe a
Nothing
    rettype_st :: ResRetType
rettype_st = [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType ((VName -> Maybe VName) -> [VName] -> [VName]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe VName -> Maybe VName
mkExt (Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [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
$ ResType -> FV
forall o. TypeBase Exp o -> FV
freeInType ResType
ret) [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
ret_dims) ResType
ret

    dec :: ValBind
dec =
      ValBind
        { valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing,
          valBindName :: VName
valBindName = VName
fname,
          valBindNameLoc :: SrcLoc
valBindNameLoc = SrcLoc
forall a. Monoid a => a
mempty,
          valBindRetDecl :: Maybe (TypeExp Exp VName)
valBindRetDecl = Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing,
          valBindRetType :: Info ResRetType
valBindRetType = ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
rettype_st,
          valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
dims',
          valBindParams :: [Pat ParamType]
valBindParams = [Pat ParamType]
pats,
          valBindBody :: Exp
valBindBody = Exp
body,
          valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing,
          valBindAttrs :: [AttrInfo VName]
valBindAttrs = [AttrInfo VName]
forall a. Monoid a => a
mempty,
          valBindLocation :: SrcLoc
valBindLocation = SrcLoc
forall a. Monoid a => a
mempty
        }

-- | Looks up the associated static value for a given name in the environment.
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar :: StructType -> VName -> DefM StaticVal
lookupVar StructType
t VName
x = do
  env <- DefM Env
askEnv
  case M.lookup x env of
    Just (Binding (Just ([VName]
dims, StructType
sv_t)) StaticVal
sv) -> do
      globals <- ((Env, Env) -> Env) -> DefM Env
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Env, Env) -> Env
forall a b. (a, b) -> a
fst
      instStaticVal (M.keysSet globals) dims t sv_t sv
    Just (Binding Maybe ([VName], StructType)
Nothing StaticVal
sv) ->
      StaticVal -> DefM StaticVal
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
sv
    Maybe Binding
Nothing -- If the variable is unknown, it may refer to the 'intrinsics'
    -- module, which we will have to treat specially.
      | VName -> Bool
isIntrinsic VName
x -> StaticVal -> DefM StaticVal
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure StaticVal
IntrinsicSV
      | Bool
otherwise ->
          -- Anything not in scope is going to be an existential size.
          StaticVal -> DefM StaticVal
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StaticVal -> DefM StaticVal) -> StaticVal -> DefM StaticVal
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Freshness
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Freshness)
-> PrimType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- Like freeInPat, but ignores sizes that are only found in
-- funtion types.
arraySizes :: StructType -> S.Set VName
arraySizes :: StructType -> Set VName
arraySizes (Scalar Arrow {}) = Set VName
forall a. Monoid a => a
mempty
arraySizes (Scalar (Record Map Name StructType
fields)) = (StructType -> Set VName) -> Map Name StructType -> Set VName
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes Map Name StructType
fields
arraySizes (Scalar (Sum Map Name [StructType]
cs)) = ([StructType] -> Set VName) -> Map Name [StructType] -> Set VName
forall m a. Monoid m => (a -> m) -> Map Name a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((StructType -> Set VName) -> [StructType] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StructType -> Set VName
arraySizes) Map Name [StructType]
cs
arraySizes (Scalar (TypeVar NoMode
_ QualName VName
_ [TypeArg Exp]
targs)) =
  [Set VName] -> Set VName
forall a. Monoid a => [a] -> a
mconcat ([Set VName] -> Set VName) -> [Set VName] -> Set VName
forall a b. (a -> b) -> a -> b
$ (TypeArg Exp -> Set VName) -> [TypeArg Exp] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg Exp -> Set VName
f [TypeArg Exp]
targs
  where
    f :: TypeArg Exp -> Set VName
f (TypeArgDim (Var QualName VName
d Info StructType
_ SrcLoc
_)) = VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    f TypeArgDim {} = Set VName
forall a. Monoid a => a
mempty
    f (TypeArgType StructType
t) = StructType -> Set VName
arraySizes StructType
t
arraySizes (Scalar Prim {}) = Set VName
forall a. Monoid a => a
mempty
arraySizes (Array NoMode
_ Shape Exp
shape ScalarTypeBase Exp NoMode
t) =
  StructType -> Set VName
arraySizes (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase Exp NoMode
t) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (Exp -> Set VName) -> [Exp] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Exp -> Set VName
dimName (Shape Exp -> [Exp]
forall dim. Shape dim -> [dim]
shapeDims Shape Exp
shape)
  where
    dimName :: Size -> S.Set VName
    dimName :: Exp -> Set VName
dimName (Var QualName VName
qn Info StructType
_ SrcLoc
_) = VName -> Set VName
forall a. a -> Set a
S.singleton (VName -> Set VName) -> VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
    dimName Exp
_ = Set VName
forall a. Monoid a => a
mempty

patternArraySizes :: Pat ParamType -> S.Set VName
patternArraySizes :: Pat ParamType -> Set VName
patternArraySizes = StructType -> Set VName
arraySizes (StructType -> Set VName)
-> (Pat ParamType -> StructType) -> Pat ParamType -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat ParamType -> StructType
forall o. Pat (TypeBase Exp o) -> StructType
patternStructType

data SizeSubst
  = SubstNamed (QualName VName)
  | SubstConst Int64
  deriving (SizeSubst -> SizeSubst -> Bool
(SizeSubst -> SizeSubst -> Bool)
-> (SizeSubst -> SizeSubst -> Bool) -> Eq SizeSubst
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SizeSubst -> SizeSubst -> Bool
== :: SizeSubst -> SizeSubst -> Bool
$c/= :: SizeSubst -> SizeSubst -> Bool
/= :: SizeSubst -> SizeSubst -> Bool
Eq, Eq SizeSubst
Eq SizeSubst =>
(SizeSubst -> SizeSubst -> Ordering)
-> (SizeSubst -> SizeSubst -> Bool)
-> (SizeSubst -> SizeSubst -> Bool)
-> (SizeSubst -> SizeSubst -> Bool)
-> (SizeSubst -> SizeSubst -> Bool)
-> (SizeSubst -> SizeSubst -> SizeSubst)
-> (SizeSubst -> SizeSubst -> SizeSubst)
-> Ord SizeSubst
SizeSubst -> SizeSubst -> Bool
SizeSubst -> SizeSubst -> Ordering
SizeSubst -> SizeSubst -> SizeSubst
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 :: SizeSubst -> SizeSubst -> Ordering
compare :: SizeSubst -> SizeSubst -> Ordering
$c< :: SizeSubst -> SizeSubst -> Bool
< :: SizeSubst -> SizeSubst -> Bool
$c<= :: SizeSubst -> SizeSubst -> Bool
<= :: SizeSubst -> SizeSubst -> Bool
$c> :: SizeSubst -> SizeSubst -> Bool
> :: SizeSubst -> SizeSubst -> Bool
$c>= :: SizeSubst -> SizeSubst -> Bool
>= :: SizeSubst -> SizeSubst -> Bool
$cmax :: SizeSubst -> SizeSubst -> SizeSubst
max :: SizeSubst -> SizeSubst -> SizeSubst
$cmin :: SizeSubst -> SizeSubst -> SizeSubst
min :: SizeSubst -> SizeSubst -> SizeSubst
Ord, Int -> SizeSubst -> ShowS
[SizeSubst] -> ShowS
SizeSubst -> [Char]
(Int -> SizeSubst -> ShowS)
-> (SizeSubst -> [Char])
-> ([SizeSubst] -> ShowS)
-> Show SizeSubst
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SizeSubst -> ShowS
showsPrec :: Int -> SizeSubst -> ShowS
$cshow :: SizeSubst -> [Char]
show :: SizeSubst -> [Char]
$cshowList :: [SizeSubst] -> ShowS
showList :: [SizeSubst] -> ShowS
Show)

dimMapping ::
  (Monoid o) =>
  TypeBase Size o ->
  TypeBase Size o ->
  M.Map VName SizeSubst
dimMapping :: forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> Map VName SizeSubst
dimMapping TypeBase Exp o
t1 TypeBase Exp o
t2 = State (Map VName SizeSubst) (TypeBase Exp o)
-> Map VName SizeSubst -> Map VName SizeSubst
forall s a. State s a -> s -> s
execState (([VName]
 -> Exp -> Exp -> StateT (Map VName SizeSubst) Identity Exp)
-> TypeBase Exp o
-> TypeBase Exp o
-> State (Map VName SizeSubst) (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 (Map VName SizeSubst) Identity Exp
forall {t :: * -> *} {f :: * -> *} {k} {f :: * -> *} {f :: * -> *}.
(Foldable t, MonadState (Map k SizeSubst) f, Ord k) =>
t VName -> ExpBase f k -> ExpBase f VName -> f (ExpBase f k)
f TypeBase Exp o
t1 TypeBase Exp o
t2) Map VName SizeSubst
forall a. Monoid a => a
mempty
  where
    f :: t VName -> ExpBase f k -> ExpBase f VName -> f (ExpBase f k)
f t VName
bound ExpBase f k
d1 (Var QualName VName
d2 f StructType
_ SrcLoc
_)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d2 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 = ExpBase f k -> f (ExpBase f k)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f k
d1
    f t VName
_ (Var QualName k
d1 f StructType
typ SrcLoc
loc) (Var QualName VName
d2 f StructType
_ SrcLoc
_) = do
      (Map k SizeSubst -> Map k SizeSubst) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k SizeSubst -> Map k SizeSubst) -> f ())
-> (Map k SizeSubst -> Map k SizeSubst) -> f ()
forall a b. (a -> b) -> a -> b
$ k -> SizeSubst -> Map k SizeSubst -> Map k SizeSubst
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d1) (SizeSubst -> Map k SizeSubst -> Map k SizeSubst)
-> SizeSubst -> Map k SizeSubst -> Map k SizeSubst
forall a b. (a -> b) -> a -> b
$ QualName VName -> SizeSubst
SubstNamed QualName VName
d2
      ExpBase f k -> f (ExpBase f k)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpBase f k -> f (ExpBase f k)) -> ExpBase f k -> f (ExpBase f k)
forall a b. (a -> b) -> a -> b
$ QualName k -> f StructType -> SrcLoc -> ExpBase f k
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName k
d1 f StructType
typ SrcLoc
loc
    f t VName
_ (Var QualName k
d1 f StructType
typ SrcLoc
loc) (IntLit Integer
d2 f StructType
_ SrcLoc
_) = do
      (Map k SizeSubst -> Map k SizeSubst) -> f ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map k SizeSubst -> Map k SizeSubst) -> f ())
-> (Map k SizeSubst -> Map k SizeSubst) -> f ()
forall a b. (a -> b) -> a -> b
$ k -> SizeSubst -> Map k SizeSubst -> Map k SizeSubst
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
d1) (SizeSubst -> Map k SizeSubst -> Map k SizeSubst)
-> SizeSubst -> Map k SizeSubst -> Map k SizeSubst
forall a b. (a -> b) -> a -> b
$ Int64 -> SizeSubst
SubstConst (Int64 -> SizeSubst) -> Int64 -> SizeSubst
forall a b. (a -> b) -> a -> b
$ Integer -> Int64
forall a. Num a => Integer -> a
fromInteger Integer
d2
      ExpBase f k -> f (ExpBase f k)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpBase f k -> f (ExpBase f k)) -> ExpBase f k -> f (ExpBase f k)
forall a b. (a -> b) -> a -> b
$ QualName k -> f StructType -> SrcLoc -> ExpBase f k
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName k
d1 f StructType
typ SrcLoc
loc
    f t VName
_ ExpBase f k
d ExpBase f VName
_ = ExpBase f k -> f (ExpBase f k)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f k
d

dimMapping' ::
  (Monoid o) =>
  TypeBase Size o ->
  TypeBase Size o ->
  M.Map VName VName
dimMapping' :: forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> Map VName VName
dimMapping' TypeBase Exp o
t1 TypeBase Exp o
t2 = (SizeSubst -> Maybe VName)
-> Map VName SizeSubst -> Map VName VName
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe SizeSubst -> Maybe VName
f (Map VName SizeSubst -> Map VName VName)
-> Map VName SizeSubst -> Map VName VName
forall a b. (a -> b) -> a -> b
$ TypeBase Exp o -> TypeBase Exp o -> Map VName SizeSubst
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> Map VName SizeSubst
dimMapping TypeBase Exp o
t1 TypeBase Exp o
t2
  where
    f :: SizeSubst -> Maybe VName
f (SubstNamed QualName VName
d) = VName -> Maybe VName
forall a. a -> Maybe a
Just (VName -> Maybe VName) -> VName -> Maybe VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
    f SizeSubst
_ = Maybe VName
forall a. Maybe a
Nothing

sizesToRename :: StaticVal -> S.Set VName
sizesToRename :: StaticVal -> Set VName
sizesToRename (DynamicFun (Exp
_, StaticVal
sv1) StaticVal
_sv2) =
  -- It is intentional that we do not look at sv2 here, as some names
  -- that are free in sv2 are actually bound by the parameters in sv1.
  -- See #2234.
  StaticVal -> Set VName
sizesToRename StaticVal
sv1
sizesToRename StaticVal
IntrinsicSV =
  Set VName
forall a. Monoid a => a
mempty
sizesToRename HoleSV {} =
  Set VName
forall a. Monoid a => a
mempty
sizesToRename Dynamic {} =
  Set VName
forall a. Monoid a => a
mempty
sizesToRename (RecordSV [(Name, StaticVal)]
fs) =
  ((Name, StaticVal) -> Set VName)
-> [(Name, StaticVal)] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StaticVal -> Set VName
sizesToRename (StaticVal -> Set VName)
-> ((Name, StaticVal) -> StaticVal)
-> (Name, StaticVal)
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StaticVal) -> StaticVal
forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fs
sizesToRename (SumSV Name
_ [StaticVal]
svs [(Name, [ResType])]
_) =
  (StaticVal -> Set VName) -> [StaticVal] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StaticVal -> Set VName
sizesToRename [StaticVal]
svs
sizesToRename (LambdaSV Pat ParamType
param ResRetType
_ Exp
_ Env
_) =
  -- We used to rename parameters here, but I don't understand why
  -- that was necessary and it caused some problems.
  FV -> Set VName
fvVars (Pat ParamType -> FV
forall o. Pat (TypeBase Exp o) -> FV
freeInPat Pat ParamType
param)

-- | Combine the shape information of types as much as possible. The first
-- argument is the orignal type and the second is the type of the transformed
-- expression. This is necessary since the original type may contain additional
-- information (e.g., shape restrictions) from the user given annotation.
combineTypeShapes ::
  (Monoid o) =>
  TypeBase Size o ->
  TypeBase Size o ->
  TypeBase Size o
combineTypeShapes :: forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes (Scalar (Record Map Name (TypeBase Exp o)
ts1)) (Scalar (Record Map Name (TypeBase Exp o)
ts2))
  | Map Name (TypeBase Exp o) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp o)
ts1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name (TypeBase Exp o) -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name (TypeBase Exp o)
ts2 =
      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
$
        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) -> TypeBase Exp o)
-> Map Name (TypeBase Exp o, 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 -> TypeBase Exp o)
-> (TypeBase Exp o, TypeBase Exp o) -> TypeBase Exp o
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes)
            ((TypeBase Exp o
 -> TypeBase Exp o -> (TypeBase Exp o, TypeBase Exp o))
-> Map Name (TypeBase Exp o)
-> Map Name (TypeBase Exp o)
-> Map Name (TypeBase Exp o, TypeBase Exp o)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name (TypeBase Exp o)
ts1 Map Name (TypeBase Exp o)
ts2)
combineTypeShapes (Scalar (Sum Map Name [TypeBase Exp o]
cs1)) (Scalar (Sum Map Name [TypeBase Exp o]
cs2))
  | Map Name [TypeBase Exp o] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Exp o]
cs1 [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name [TypeBase Exp o] -> [Name]
forall k a. Map k a -> [k]
M.keys Map Name [TypeBase Exp o]
cs2 =
      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
$
        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]) -> [TypeBase Exp o])
-> Map Name ([TypeBase Exp o], [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] -> [TypeBase Exp o])
-> ([TypeBase Exp o], [TypeBase Exp o]) -> [TypeBase Exp o]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (([TypeBase Exp o] -> [TypeBase Exp o] -> [TypeBase Exp o])
 -> ([TypeBase Exp o], [TypeBase Exp o]) -> [TypeBase Exp o])
-> ([TypeBase Exp o] -> [TypeBase Exp o] -> [TypeBase Exp o])
-> ([TypeBase Exp o], [TypeBase Exp o])
-> [TypeBase Exp o]
forall a b. (a -> b) -> a -> b
$ (TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o)
-> [TypeBase Exp o] -> [TypeBase Exp o] -> [TypeBase Exp o]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes)
            (([TypeBase Exp o]
 -> [TypeBase Exp o] -> ([TypeBase Exp o], [TypeBase Exp o]))
-> Map Name [TypeBase Exp o]
-> Map Name [TypeBase Exp o]
-> Map Name ([TypeBase Exp o], [TypeBase Exp o])
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith (,) Map Name [TypeBase Exp o]
cs1 Map Name [TypeBase Exp o]
cs2)
combineTypeShapes (Scalar (Arrow o
als1 PName
p1 Diet
d1 StructType
a1 (RetType [VName]
dims1 ResType
b1))) (Scalar (Arrow o
als2 PName
_p2 Diet
_d2 StructType
a2 (RetType [VName]
_ ResType
b2))) =
  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
$
    o
-> PName
-> Diet
-> StructType
-> ResRetType
-> ScalarTypeBase Exp o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow
      (o
als1 o -> o -> o
forall a. Semigroup a => a -> a -> a
<> o
als2)
      PName
p1
      Diet
d1
      (StructType -> StructType -> StructType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes StructType
a1 StructType
a2)
      ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
dims1 (ResType -> ResType -> ResType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes ResType
b1 ResType
b2))
combineTypeShapes (Scalar (TypeVar o
o QualName VName
v [TypeArg Exp]
targs1)) (Scalar (TypeVar o
_ QualName VName
_ [TypeArg Exp]
targs2)) =
  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
$ o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
v ([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] -> [TypeArg Exp] -> [TypeArg Exp]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeArg Exp -> TypeArg Exp -> TypeArg Exp
f [TypeArg Exp]
targs1 [TypeArg Exp]
targs2
  where
    f :: TypeArg Exp -> TypeArg Exp -> TypeArg Exp
f (TypeArgType StructType
t1) (TypeArgType StructType
t2) = StructType -> TypeArg Exp
forall dim. TypeBase dim NoMode -> TypeArg dim
TypeArgType (StructType -> StructType -> StructType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes StructType
t1 StructType
t2)
    f TypeArg Exp
targ TypeArg Exp
_ = TypeArg Exp
targ
combineTypeShapes (Array o
o Shape Exp
shape1 ScalarTypeBase Exp NoMode
et1) (Array o
_ Shape Exp
_shape2 ScalarTypeBase Exp NoMode
et2) =
  o -> Shape Exp -> TypeBase Exp o -> TypeBase Exp o
forall o dim o'.
o -> Shape dim -> TypeBase dim o' -> TypeBase dim o
arrayOfWithAliases
    o
o
    Shape Exp
shape1
    (TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes (StructType -> o -> TypeBase Exp o
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
setMode (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase Exp NoMode
et1) o
o) (StructType -> o -> TypeBase Exp o
forall dim o1 o2. TypeBase dim o1 -> o2 -> TypeBase dim o2
setMode (ScalarTypeBase Exp NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar ScalarTypeBase Exp NoMode
et2) o
o))
combineTypeShapes TypeBase Exp o
_ TypeBase Exp o
t = TypeBase Exp o
t

-- When we instantiate a polymorphic StaticVal, we rename all the
-- sizes to avoid name conflicts later on.  This is a bit of a hack...
instStaticVal ::
  (MonadFreshNames m) =>
  S.Set VName ->
  [VName] ->
  StructType ->
  StructType ->
  StaticVal ->
  m StaticVal
instStaticVal :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [VName] -> StructType -> StructType -> StaticVal -> m StaticVal
instStaticVal Set VName
globals [VName]
dims StructType
t StructType
sv_t StaticVal
sv = do
  fresh_substs <-
    [VName] -> m (Map VName SizeSubst)
forall {f :: * -> *}.
MonadFreshNames f =>
[VName] -> f (Map VName SizeSubst)
mkSubsts ([VName] -> m (Map VName SizeSubst))
-> (Set VName -> [VName]) -> Set VName -> m (Map VName SizeSubst)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (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
globals) ([VName] -> [VName])
-> (Set VName -> [VName]) -> Set VName -> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> m (Map VName SizeSubst))
-> Set VName -> m (Map VName SizeSubst)
forall a b. (a -> b) -> a -> b
$
      [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv
  let dims' = (VName -> VName) -> [VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
fresh_substs) [VName]
dims
      isDim VName
k p
_ = VName
k VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
dims'
      dim_substs =
        (VName -> SizeSubst -> Bool)
-> Map VName SizeSubst -> Map VName SizeSubst
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey VName -> SizeSubst -> Bool
forall {p}. VName -> p -> Bool
isDim (Map VName SizeSubst -> Map VName SizeSubst)
-> Map VName SizeSubst -> Map VName SizeSubst
forall a b. (a -> b) -> a -> b
$ StructType -> StructType -> Map VName SizeSubst
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> Map VName SizeSubst
dimMapping (Map VName SizeSubst -> StructType -> StructType
forall o. Map VName SizeSubst -> TypeBase Exp o -> TypeBase Exp o
replaceTypeSizes Map VName SizeSubst
fresh_substs StructType
sv_t) StructType
t
      replace (SubstNamed QualName VName
k) = SizeSubst -> Maybe SizeSubst -> SizeSubst
forall a. a -> Maybe a -> a
fromMaybe (QualName VName -> SizeSubst
SubstNamed QualName VName
k) (Maybe SizeSubst -> SizeSubst) -> Maybe SizeSubst -> SizeSubst
forall a b. (a -> b) -> a -> b
$ VName -> Map VName SizeSubst -> Maybe SizeSubst
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
k) Map VName SizeSubst
dim_substs
      replace SizeSubst
k = SizeSubst
k
      substs = (SizeSubst -> SizeSubst)
-> Map VName SizeSubst -> Map VName SizeSubst
forall a b k. (a -> b) -> Map k a -> Map k b
M.map SizeSubst -> SizeSubst
replace Map VName SizeSubst
fresh_substs Map VName SizeSubst -> Map VName SizeSubst -> Map VName SizeSubst
forall a. Semigroup a => a -> a -> a
<> Map VName SizeSubst
dim_substs

  pure $ replaceStaticValSizes globals substs sv
  where
    mkSubsts :: [VName] -> f (Map VName SizeSubst)
mkSubsts [VName]
names =
      [(VName, SizeSubst)] -> Map VName SizeSubst
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, SizeSubst)] -> Map VName SizeSubst)
-> ([VName] -> [(VName, SizeSubst)])
-> [VName]
-> Map VName SizeSubst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [SizeSubst] -> [(VName, SizeSubst)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
names ([SizeSubst] -> [(VName, SizeSubst)])
-> ([VName] -> [SizeSubst]) -> [VName] -> [(VName, SizeSubst)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> SizeSubst) -> [VName] -> [SizeSubst]
forall a b. (a -> b) -> [a] -> [b]
map (QualName VName -> SizeSubst
SubstNamed (QualName VName -> SizeSubst)
-> (VName -> QualName VName) -> VName -> SizeSubst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName)
        ([VName] -> Map VName SizeSubst)
-> f [VName] -> f (Map VName SizeSubst)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VName -> f VName) -> [VName] -> f [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 VName -> f VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName [VName]
names

    onName :: Map VName SizeSubst -> VName -> VName
onName Map VName SizeSubst
substs VName
v =
      case VName -> Map VName SizeSubst -> Maybe SizeSubst
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName SizeSubst
substs of
        Just (SubstNamed QualName VName
v') -> QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v'
        Maybe SizeSubst
_ -> VName
v

defuncFun ::
  [VName] ->
  [Pat ParamType] ->
  Exp ->
  ResRetType ->
  SrcLoc ->
  DefM (Exp, StaticVal)
defuncFun :: [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
tparams [Pat ParamType]
pats Exp
e0 ResRetType
ret SrcLoc
loc = do
  -- Extract the first parameter of the lambda and "push" the
  -- remaining ones (if there are any) into the body of the lambda.
  let (Pat ParamType
pat, ResRetType
ret', Exp
e0') = case [Pat ParamType]
pats of
        [] -> [Char] -> (Pat ParamType, ResRetType, Exp)
forall a. HasCallStack => [Char] -> a
error [Char]
"Received a lambda with no parameters."
        [Pat ParamType
pat'] -> (Pat ParamType
pat', ResRetType
ret, Exp
e0)
        (Pat ParamType
pat' : [Pat ParamType]
pats') ->
          ( Pat ParamType
pat',
            [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ (NoMode -> Freshness) -> StructType -> ResType
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) (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> ResRetType -> StructType
forall d.
[Pat (TypeBase d Diet)]
-> RetTypeBase d Freshness -> TypeBase d NoMode
funType [Pat ParamType]
pats' ResRetType
ret,
            [Pat ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats' Exp
e0 Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
ret) SrcLoc
loc
          )

  -- Construct a record literal that closes over the environment of
  -- the lambda.  Closed-over 'DynamicFun's are converted to their
  -- closure representation.
  let used :: FV
used =
        Exp -> FV
freeInExp ([Pat ParamType]
-> Exp
-> Maybe (TypeExp Exp VName)
-> Info ResRetType
-> SrcLoc
-> Exp
forall (f :: * -> *) vn.
[PatBase f vn ParamType]
-> ExpBase f vn
-> Maybe (TypeExp (ExpBase f vn) vn)
-> f ResRetType
-> SrcLoc
-> ExpBase f vn
Lambda [Pat ParamType]
pats Exp
e0 Maybe (TypeExp Exp VName)
forall a. Maybe a
Nothing (ResRetType -> Info ResRetType
forall a. a -> Info a
Info ResRetType
ret) SrcLoc
loc)
          FV -> Set VName -> FV
`freeWithout` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams
  used_env <- FV -> DefM Env
restrictEnvTo FV
used

  -- The closure parts that are sizes are proactively turned into size
  -- parameters.
  let sizes_of_arrays =
        (Binding -> Set VName) -> Env -> Set VName
forall m a. Monoid m => (a -> m) -> Map VName a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (StructType -> Set VName
arraySizes (StructType -> Set VName)
-> (Binding -> StructType) -> Binding -> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> StructType
structTypeFromSV (StaticVal -> StructType)
-> (Binding -> StaticVal) -> Binding -> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV) Env
used_env
          Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Pat ParamType -> Set VName
patternArraySizes Pat ParamType
pat
      notSize = Bool -> Bool
not (Bool -> Bool) -> (VName -> Bool) -> VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
sizes_of_arrays)
      (fields, env) =
        second M.fromList
          . unzip
          . map closureFromDynamicFun
          . filter (notSize . fst)
          $ M.toList used_env

  pure
    ( RecordLit fields loc,
      LambdaSV pat ret' e0' env
    )
  where
    closureFromDynamicFun :: (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ (DynamicFun (Exp
clsr_env, StaticVal
sv) StaticVal
_)) =
      let name :: Name
name = Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
vn
       in ( L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc Name
name) Exp
clsr_env SrcLoc
forall a. Monoid a => a
mempty,
            (VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing StaticVal
sv)
          )
    closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
      let name :: Name
name = Text -> Name
nameFromText (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
vn
          tp' :: StructType
tp' = StaticVal -> StructType
structTypeFromSV StaticVal
sv
       in ( L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
              (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc Name
name)
              (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
tp') SrcLoc
forall a. Monoid a => a
mempty)
              SrcLoc
forall a. Monoid a => a
mempty,
            (VName
vn, Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing StaticVal
sv)
          )

-- | Defunctionalization of an expression. Returns the residual expression and
-- the associated static value in the defunctionalization monad.
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp :: Exp -> DefM (Exp, StaticVal)
defuncExp e :: Exp
e@Literal {} =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@IntLit {} =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@FloatLit {} =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp e :: Exp
e@StringLit {} =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp (Parens Exp
e SrcLoc
loc) = do
  (e', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
  pure (Parens e' loc, sv)
defuncExp (QualParens (QualName VName, SrcLoc)
qn Exp
e SrcLoc
loc) = do
  (e', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
  pure (QualParens qn e' loc, sv)
defuncExp (TupLit [Exp]
es SrcLoc
loc) = do
  (es', svs) <- (Exp -> DefM (Exp, StaticVal))
-> [Exp] -> DefM ([Exp], [StaticVal])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
  pure (TupLit es' loc, RecordSV $ zip tupleFieldNames svs)
defuncExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
  (fs', names_svs) <- (FieldBase Info VName
 -> DefM (FieldBase Info VName, (Name, StaticVal)))
-> [FieldBase Info VName]
-> DefM ([FieldBase Info VName], [(Name, StaticVal)])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField [FieldBase Info VName]
fs
  pure (RecordLit fs' loc, RecordSV names_svs)
  where
    defuncField :: FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField (RecordFieldExplicit L Name
vn Exp
e SrcLoc
loc') = do
      (e', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
      pure (RecordFieldExplicit vn e' loc', (unLoc vn, sv))
    defuncField (RecordFieldImplicit (L Loc
_ VName
vn) (Info StructType
t) SrcLoc
loc') = do
      sv <- StructType -> VName -> DefM StaticVal
lookupVar (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct StructType
t) VName
vn
      case sv of
        -- If the implicit field refers to a dynamic function, we
        -- convert it to an explicit field with a record closing over
        -- the environment and bind the corresponding static value.
        DynamicFun (Exp
e, StaticVal
sv') StaticVal
_ ->
          let vn' :: Name
vn' = VName -> Name
baseName VName
vn
           in (FieldBase Info VName, (Name, StaticVal))
-> DefM (FieldBase Info VName, (Name, StaticVal))
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( L Name -> Exp -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit (Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc Name
vn') Exp
e SrcLoc
loc',
                  (Name
vn', StaticVal
sv')
                )
        -- The field may refer to a functional expression, so we get the
        -- type from the static value and not the one from the AST.
        StaticVal
_ ->
          let tp :: Info StructType
tp = StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ StaticVal -> StructType
structTypeFromSV StaticVal
sv
           in (FieldBase Info VName, (Name, StaticVal))
-> DefM (FieldBase Info VName, (Name, StaticVal))
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
                ( L VName -> Info StructType -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
L vn -> f StructType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit (Loc -> VName -> L VName
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc VName
vn) Info StructType
tp SrcLoc
loc',
                  (VName -> Name
baseName VName
vn, StaticVal
sv)
                )
defuncExp e :: Exp
e@(ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc) =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([PrimValue] -> PrimType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
[PrimValue] -> PrimType -> SrcLoc -> ExpBase f vn
ArrayVal [PrimValue]
vs PrimType
t SrcLoc
loc, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e)
defuncExp (ArrayLit [Exp]
es t :: Info StructType
t@(Info StructType
t') SrcLoc
loc) = do
  es' <- (Exp -> DefM Exp) -> [Exp] -> DefM [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 -> DefM Exp
defuncExp' [Exp]
es
  pure (ArrayLit es' t loc, Dynamic $ toRes Nonfresh t')
defuncExp (AppExp (Range Exp
e1 Maybe Exp
me Inclusiveness Exp
incl SrcLoc
loc) Info AppRes
res) = do
  e1' <- Exp -> DefM Exp
defuncExp' Exp
e1
  me' <- mapM defuncExp' me
  incl' <- mapM defuncExp' incl
  pure
    ( AppExp (Range e1' me' incl' loc) res,
      Dynamic $ toRes Nonfresh $ appResType $ unInfo res
    )
defuncExp e :: Exp
e@(Var QualName VName
qn (Info StructType
t) SrcLoc
loc) = do
  sv <- StructType -> VName -> DefM StaticVal
lookupVar (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct StructType
t) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
  case sv of
    -- If the variable refers to a dynamic function, we eta-expand it
    -- so that we do not have to duplicate its definition.
    DynamicFun {} -> do
      (params, body, ret) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t) Exp
e
      defuncFun [] params body ret mempty
    -- Intrinsic functions used as variables are eta-expanded, so we
    -- can get rid of them.
    StaticVal
IntrinsicSV -> do
      (pats, body, tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t) Exp
e
      defuncExp $ Lambda pats body Nothing (Info tp) mempty
    HoleSV StructType
_ SrcLoc
hole_loc ->
      (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
hole_loc, StaticVal
sv)
    StaticVal
_ ->
      (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (StructType -> Info StructType
forall a. a -> Info a
Info (StaticVal -> StructType
structTypeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncExp (Hole (Info StructType
t) SrcLoc
loc) =
  (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn. f StructType -> SrcLoc -> ExpBase f vn
Hole (StructType -> Info StructType
forall a. a -> Info a
Info StructType
t) SrcLoc
loc, StructType -> SrcLoc -> StaticVal
HoleSV StructType
t SrcLoc
loc)
defuncExp (Ascript Exp
e0 TypeExp Exp VName
tydecl SrcLoc
loc)
  | StructType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero (Exp -> StructType
typeOf Exp
e0) = do
      (e0', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
      pure (Ascript e0' tydecl loc, sv)
  | Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (Coerce Exp
e0 TypeExp Exp VName
tydecl Info StructType
t SrcLoc
loc)
  | StructType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero (Exp -> StructType
typeOf Exp
e0) = do
      (e0', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
      pure (Coerce e0' tydecl t loc, sv)
  | Bool
otherwise = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
defuncExp (AppExp (LetPat [SizeBinder VName]
sizes PatBase Info VName StructType
pat Exp
e1 Exp
e2 SrcLoc
loc) (Info (AppRes StructType
t [VName]
retext))) = do
  (e1', sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
  let extEnv Env
env = Env -> Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Env
env ((StructType -> ParamType)
-> PatBase Info VName StructType -> Pat ParamType
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 (Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
pat) StaticVal
sv1
      pat' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat ((StructType -> ParamType)
-> PatBase Info VName StructType -> Pat ParamType
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 (Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
pat) StaticVal
sv1
  (e2', sv2) <- localEnv extEnv $ defuncExp e2
  -- To maintain any sizes going out of scope, we need to compute the
  -- old size substitution induced by retext and also apply it to the
  -- newly computed body type.
  let mapping = StructType -> StructType -> Map VName VName
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> Map VName VName
dimMapping' (Exp -> StructType
typeOf Exp
e2) StructType
t
      subst VName
v = Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> (VName -> Exp) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
forall a. Monoid a => a
mempty (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName (VName -> Subst t) -> Maybe VName -> Maybe (Subst t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName VName -> Maybe VName
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v Map VName VName
mapping
      t' = 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
e2'
  pure (AppExp (LetPat sizes (fmap toStruct pat') e1' e2' loc) (Info (AppRes t' retext)), sv2)
defuncExp (AppExp (LetFun (VName, SrcLoc)
vn ([TypeParamBase VName], [Pat ParamType], Maybe (TypeExp Exp VName),
 Info ResRetType, Exp)
_ Exp
_ SrcLoc
_) Info AppRes
_) =
  [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM (Exp, StaticVal))
-> [Char] -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: Unexpected LetFun: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ (VName, SrcLoc) -> [Char]
forall a. Show a => a -> [Char]
show (VName, SrcLoc)
vn
defuncExp (AppExp (If Exp
e1 Exp
e2 Exp
e3 SrcLoc
loc) Info AppRes
res) = do
  (e1', _) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
  (e2', sv) <- defuncExp e2
  (e3', _) <- defuncExp e3
  pure (AppExp (If e1' e2' e3' loc) res, sv)
defuncExp (AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
args SrcLoc
loc) (Info AppRes
appres)) =
  Exp
-> NonEmpty (Maybe VName, Exp)
-> AppRes
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncApply Exp
f (((Info (Maybe VName), Exp) -> (Maybe VName, Exp))
-> NonEmpty (Info (Maybe VName), Exp)
-> NonEmpty (Maybe VName, Exp)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Info (Maybe VName) -> Maybe VName)
-> (Info (Maybe VName), Exp) -> (Maybe VName, Exp)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Info (Maybe VName) -> Maybe VName
forall a. Info a -> a
unInfo) NonEmpty (Info (Maybe VName), Exp)
args) AppRes
appres SrcLoc
loc
defuncExp (Negate Exp
e0 SrcLoc
loc) = do
  (e0', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
  pure (Negate e0' loc, sv)
defuncExp (Not Exp
e0 SrcLoc
loc) = do
  (e0', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
  pure (Not e0' loc, sv)
defuncExp (Lambda [Pat ParamType]
pats Exp
e0 Maybe (TypeExp Exp VName)
_ (Info ResRetType
ret) SrcLoc
loc) =
  [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [] [Pat ParamType]
pats Exp
e0 ResRetType
ret SrcLoc
loc
-- Operator sections are expected to be converted to lambda-expressions
-- by the monomorphizer, so they should no longer occur at this point.
defuncExp OpSection {} = [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionLeft {} = [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp OpSectionRight {} = [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected operator section."
defuncExp UpdateSection {} = [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected projection section."
defuncExp (AppExp (Loop [VName]
sparams Pat ParamType
pat LoopInitBase Info VName
loopinit LoopFormBase Info VName
form Exp
e3 SrcLoc
loc) Info AppRes
res) = do
  (e1', sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp (Exp -> DefM (Exp, StaticVal)) -> Exp -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ LoopInitBase Info VName -> Exp
loopInitExp LoopInitBase Info VName
loopinit
  env <- askEnv
  let env1 = Env -> Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Env
env Pat ParamType
pat StaticVal
sv1
  (form', env2) <- case form of
    For IdentBase Info VName StructType
v Exp
e2 -> do
      e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
      pure (For v e2', insertIdent v env1)
    ForIn PatBase Info VName StructType
pat2 Exp
e2 -> do
      e2' <- Exp -> DefM Exp
defuncExp' Exp
e2
      pure
        ( ForIn pat2 e2',
          envFromPat env1 (fmap (toParam Observe) pat2)
        )
    While Exp
e2 -> do
      e2' <- ((Env, Env) -> (Env, Env)) -> DefM Exp -> DefM Exp
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> (Env, Env) -> (Env, Env)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Env -> Env -> Env
forall a b. a -> b -> a
const Env
env1)) (DefM Exp -> DefM Exp) -> DefM Exp -> DefM Exp
forall a b. (a -> b) -> a -> b
$ Exp -> DefM Exp
defuncExp' Exp
e2
      pure (While e2', env1)
  (e3', sv) <- local (second (const env2)) $ defuncExp e3
  pure (AppExp (Loop sparams pat (LoopInitExplicit e1') form' e3' loc) res, sv)
  where
    insertIdent :: IdentBase Info k (TypeBase Exp o) -> Map k Binding -> Map k Binding
insertIdent (Ident k
vn (Info TypeBase Exp o
tp) SrcLoc
_) =
      k -> Binding -> Map k Binding -> Map k Binding
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
vn (Binding -> Map k Binding -> Map k Binding)
-> Binding -> Map k Binding -> Map k Binding
forall a b. (a -> b) -> a -> b
$ Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> TypeBase Exp o -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh TypeBase Exp o
tp
defuncExp e :: Exp
e@(AppExp BinOp {} Info AppRes
_) =
  [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM (Exp, StaticVal))
-> [Char] -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ [Char]
"defuncExp: unexpected binary operator: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
e
defuncExp (Project Name
vn Exp
e0 tp :: Info StructType
tp@(Info StructType
tp') SrcLoc
loc) = do
  (e0', sv0) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e0
  case sv0 of
    RecordSV [(Name, StaticVal)]
svs -> case Name -> [(Name, StaticVal)] -> Maybe StaticVal
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
vn [(Name, StaticVal)]
svs of
      Just StaticVal
sv -> (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ StaticVal -> StructType
structTypeFromSV StaticVal
sv) SrcLoc
loc, StaticVal
sv)
      Maybe StaticVal
Nothing -> [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
    Dynamic ResType
_ -> (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info StructType
tp SrcLoc
loc, ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
tp')
    HoleSV StructType
_ SrcLoc
hloc -> (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Exp -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f StructType -> SrcLoc -> ExpBase f vn
Project Name
vn Exp
e0' Info StructType
tp SrcLoc
loc, StructType -> SrcLoc -> StaticVal
HoleSV StructType
tp' SrcLoc
hloc)
    StaticVal
_ -> [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM (Exp, StaticVal))
-> [Char] -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ [Char]
"Projection of an expression with static value " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv0
defuncExp (AppExp LetWith {} Info AppRes
_) =
  [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp: unexpected LetWith"
defuncExp expr :: Exp
expr@(AppExp (Index Exp
e0 SliceBase Info VName
idxs SrcLoc
loc) Info AppRes
res) = do
  e0' <- Exp -> DefM Exp
defuncExp' Exp
e0
  idxs' <- mapM defuncDimIndex idxs
  pure
    ( AppExp (Index e0' idxs' loc) res,
      Dynamic $ toRes Nonfresh $ typeOf expr
    )

-- Note that we might change the type of the record field here.  This
-- is not permitted in the type checker due to problems with type
-- inference, but it actually works fine.
defuncExp (Update Exp
e1 [UpdateStep Info VName]
steps Exp
e2 Info StructType
t SrcLoc
loc) = do
  (e1', sv1) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
  (e2', sv2) <- defuncExp e2
  steps' <- mapM defuncStep steps
  let sv' = StaticVal -> [UpdateStep Info VName] -> StaticVal -> StaticVal
forall {f :: * -> *} {vn}.
StaticVal -> [UpdateStep f vn] -> StaticVal -> StaticVal
updateStatic StaticVal
sv1 [UpdateStep Info VName]
steps StaticVal
sv2
  pure (Update e1' steps' e2' t loc, sv')
  where
    defuncStep :: UpdateStep Info VName -> DefM (UpdateStep Info VName)
defuncStep (UpdateStepField Name
f) =
      UpdateStep Info VName -> DefM (UpdateStep Info VName)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UpdateStep Info VName -> DefM (UpdateStep Info VName))
-> UpdateStep Info VName -> DefM (UpdateStep Info VName)
forall a b. (a -> b) -> a -> b
$ Name -> UpdateStep Info VName
forall (f :: * -> *) vn. Name -> UpdateStep f vn
UpdateStepField Name
f
    defuncStep (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)
-> DefM (SliceBase Info VName) -> DefM (UpdateStep Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (DimIndexBase Info VName -> DefM (DimIndexBase Info VName))
-> SliceBase Info VName -> DefM (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 -> DefM (DimIndexBase Info VName)
defuncDimIndex SliceBase Info VName
idxs

    updateStatic :: StaticVal -> [UpdateStep f vn] -> StaticVal -> StaticVal
updateStatic StaticVal
_ [] StaticVal
newv = StaticVal
newv
    updateStatic (RecordSV [(Name, StaticVal)]
fs) (UpdateStepField Name
f : [UpdateStep f vn]
rest) StaticVal
newv =
      case Name -> [(Name, StaticVal)] -> Maybe StaticVal
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
f [(Name, StaticVal)]
fs of
        Just StaticVal
old ->
          [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$
            (Name
f, StaticVal -> [UpdateStep f vn] -> StaticVal -> StaticVal
updateStatic StaticVal
old [UpdateStep f vn]
rest StaticVal
newv) (Name, StaticVal) -> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall a. a -> [a] -> [a]
: ((Name, StaticVal) -> Bool)
-> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
f) (Name -> Bool)
-> ((Name, StaticVal) -> Name) -> (Name, StaticVal) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StaticVal) -> Name
forall a b. (a, b) -> a
fst) [(Name, StaticVal)]
fs
        Maybe StaticVal
Nothing ->
          [Char] -> StaticVal
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid record projection."
    updateStatic (Dynamic t' :: ResType
t'@(Scalar Record {})) [UpdateStep f vn]
steps0 StaticVal
newv =
      StaticVal -> [UpdateStep f vn] -> StaticVal -> StaticVal
updateStatic (ResType -> StaticVal
svFromType ResType
t') [UpdateStep f vn]
steps0 StaticVal
newv
    updateStatic StaticVal
cur (UpdateStepSlice SliceBase f vn
_ : [UpdateStep f vn]
_) StaticVal
_ =
      StaticVal
cur
    updateStatic StaticVal
_ [UpdateStep f vn]
_ StaticVal
_ =
      [Char] -> StaticVal
forall a. HasCallStack => [Char] -> a
error [Char]
"defuncExp Update: invalid update."
defuncExp (Assert Exp
e1 Exp
e2 Info Text
desc SrcLoc
loc) = do
  (e1', _) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
  (e2', sv) <- defuncExp e2
  pure (Assert e1' e2' desc loc, sv)
defuncExp (Constr Name
name [Exp]
es (Info sum_t :: StructType
sum_t@(Scalar (Sum Map Name [StructType]
all_fs))) SrcLoc
loc) = do
  (es', svs) <- (Exp -> DefM (Exp, StaticVal))
-> [Exp] -> DefM ([Exp], [StaticVal])
forall (m :: * -> *) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM Exp -> DefM (Exp, StaticVal)
defuncExp [Exp]
es
  let sv =
        Name -> [StaticVal] -> [(Name, [ResType])] -> StaticVal
SumSV Name
name [StaticVal]
svs ([(Name, [ResType])] -> StaticVal)
-> [(Name, [ResType])] -> StaticVal
forall a b. (a -> b) -> a -> b
$
          Map Name [ResType] -> [(Name, [ResType])]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name [ResType] -> [(Name, [ResType])])
-> Map Name [ResType] -> [(Name, [ResType])]
forall a b. (a -> b) -> a -> b
$
            Name
name Name -> Map Name [ResType] -> Map Name [ResType]
forall k a. Ord k => k -> Map k a -> Map k a
`M.delete` ([StructType] -> [ResType])
-> Map Name [StructType] -> Map Name [ResType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((StructType -> ResType) -> [StructType] -> [ResType]
forall a b. (a -> b) -> [a] -> [b]
map (Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType)
-> (StructType -> StructType) -> StructType -> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StructType -> StructType
forall o. Monoid o => TypeBase Exp o -> TypeBase Exp o
defuncType)) Map Name [StructType]
all_fs
      sum_t' = StructType -> StructType -> StructType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes StructType
sum_t (StaticVal -> StructType
structTypeFromSV StaticVal
sv)
  pure (Constr name es' (Info sum_t') loc, sv)
  where
    defuncType ::
      (Monoid o) =>
      TypeBase Size o ->
      TypeBase Size o
    defuncType :: forall o. Monoid o => TypeBase Exp o -> TypeBase Exp o
defuncType (Array o
o Shape Exp
shape ScalarTypeBase Exp NoMode
t) = 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 -> ScalarTypeBase Exp NoMode
forall o. Monoid o => ScalarTypeBase Exp o -> ScalarTypeBase Exp o
defuncScalar ScalarTypeBase Exp NoMode
t)
    defuncType (Scalar ScalarTypeBase Exp o
t) = 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
$ ScalarTypeBase Exp o -> ScalarTypeBase Exp o
forall o. Monoid o => ScalarTypeBase Exp o -> ScalarTypeBase Exp o
defuncScalar ScalarTypeBase Exp o
t

    defuncScalar ::
      (Monoid o) =>
      ScalarTypeBase Size o ->
      ScalarTypeBase Size o
    defuncScalar :: forall o. Monoid o => ScalarTypeBase Exp o -> ScalarTypeBase Exp o
defuncScalar (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 TypeBase Exp o -> TypeBase Exp o
forall o. Monoid o => TypeBase Exp o -> TypeBase Exp o
defuncType Map Name (TypeBase Exp o)
fs
    defuncScalar Arrow {} = 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)
forall a. Monoid a => a
mempty
    defuncScalar (Sum 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
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)
-> [TypeBase Exp o] -> [TypeBase Exp o]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase Exp o -> TypeBase Exp o
forall o. Monoid o => TypeBase Exp o -> TypeBase Exp o
defuncType) Map Name [TypeBase Exp o]
fs
    defuncScalar (Prim PrimType
t) = PrimType -> ScalarTypeBase Exp o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim PrimType
t
    defuncScalar (TypeVar o
o QualName VName
tn [TypeArg Exp]
targs) = o -> QualName VName -> [TypeArg Exp] -> ScalarTypeBase Exp o
forall dim o.
o -> QualName VName -> [TypeArg dim] -> ScalarTypeBase dim o
TypeVar o
o QualName VName
tn [TypeArg Exp]
targs
defuncExp (Constr Name
name [Exp]
_ (Info StructType
t) SrcLoc
loc) =
  [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM (Exp, StaticVal))
-> [Char] -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$
    [Char]
"Constructor "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
name
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" given type "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ StructType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString StructType
t
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" at "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SrcLoc -> [Char]
forall a. Located a => a -> [Char]
locStr SrcLoc
loc
defuncExp (AppExp (Match Exp
e NonEmpty (CaseBase Info VName)
cs SrcLoc
loc) Info AppRes
res) = do
  (e', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
  let bad = [Char] -> b
forall a. HasCallStack => [Char] -> a
error ([Char] -> b) -> [Char] -> b
forall a b. (a -> b) -> a -> b
$ [Char]
"No case matches StaticVal\n" [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv
  csPairs <-
    fromMaybe bad . NE.nonEmpty . catMaybes
      <$> mapM (defuncCase sv) (NE.toList cs)
  let cs' = ((CaseBase Info VName, StaticVal) -> CaseBase Info VName)
-> NonEmpty (CaseBase Info VName, StaticVal)
-> NonEmpty (CaseBase Info VName)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CaseBase Info VName, StaticVal) -> CaseBase Info VName
forall a b. (a, b) -> a
fst NonEmpty (CaseBase Info VName, StaticVal)
csPairs
      sv' = (CaseBase Info VName, StaticVal) -> StaticVal
forall a b. (a, b) -> b
snd ((CaseBase Info VName, StaticVal) -> StaticVal)
-> (CaseBase Info VName, StaticVal) -> StaticVal
forall a b. (a -> b) -> a -> b
$ NonEmpty (CaseBase Info VName, StaticVal)
-> (CaseBase Info VName, StaticVal)
forall a. NonEmpty a -> a
NE.head NonEmpty (CaseBase Info VName, StaticVal)
csPairs
  pure (AppExp (Match e' cs' loc) res, sv')
defuncExp (Attr AttrInfo VName
info Exp
e SrcLoc
loc) = do
  (e', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
  pure (Attr info e' loc, sv)

-- | Same as 'defuncExp', except it ignores the static value.
defuncExp' :: Exp -> DefM Exp
defuncExp' :: Exp -> DefM Exp
defuncExp' = ((Exp, StaticVal) -> Exp) -> DefM (Exp, StaticVal) -> DefM Exp
forall a b. (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Exp, StaticVal) -> Exp
forall a b. (a, b) -> a
fst (DefM (Exp, StaticVal) -> DefM Exp)
-> (Exp -> DefM (Exp, StaticVal)) -> Exp -> DefM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> DefM (Exp, StaticVal)
defuncExp

defuncCase :: StaticVal -> Case -> DefM (Maybe (Case, StaticVal))
defuncCase :: StaticVal
-> CaseBase Info VName
-> DefM (Maybe (CaseBase Info VName, StaticVal))
defuncCase StaticVal
sv (CasePat PatBase Info VName StructType
p Exp
e SrcLoc
loc) = do
  let p' :: Pat ParamType
p' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat ((StructType -> ParamType)
-> PatBase Info VName StructType -> Pat ParamType
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 (Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe) PatBase Info VName StructType
p) StaticVal
sv
  env <- DefM Env
askEnv
  case matchPatSV env (fmap (toParam Observe) p) sv of
    Just Env
env' -> do
      (e', sv') <- ((Env, Env) -> (Env, Env))
-> DefM (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> (Env, Env) -> (Env, Env)
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Env -> Env -> Env
forall a b. a -> b -> a
const Env
env')) (DefM (Exp, StaticVal) -> DefM (Exp, StaticVal))
-> DefM (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e
      pure $ Just (CasePat (fmap toStruct p') e' loc, sv')
    Maybe Env
Nothing ->
      Maybe (CaseBase Info VName, StaticVal)
-> DefM (Maybe (CaseBase Info VName, StaticVal))
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (CaseBase Info VName, StaticVal)
forall a. Maybe a
Nothing

-- | Defunctionalize the function argument to a SOAC by eta-expanding if
-- necessary and then defunctionalizing the body of the introduced lambda.
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp :: Exp -> DefM Exp
defuncSoacExp e :: Exp
e@OpSection {} = Exp -> DefM Exp
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionLeft {} = Exp -> DefM Exp
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@OpSectionRight {} = Exp -> DefM Exp
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp e :: Exp
e@UpdateSection {} = Exp -> DefM Exp
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e
defuncSoacExp (Parens Exp
e SrcLoc
loc) =
  Exp -> SrcLoc -> Exp
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (Exp -> SrcLoc -> Exp) -> DefM Exp -> DefM (SrcLoc -> Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM Exp
defuncSoacExp Exp
e DefM (SrcLoc -> Exp) -> DefM SrcLoc -> DefM Exp
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> DefM SrcLoc
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
defuncSoacExp (Lambda [Pat ParamType]
params Exp
e0 Maybe (TypeExp Exp VName)
decl Info ResRetType
tp SrcLoc
loc) = do
  env <- DefM Env
askEnv
  let env' = (Env -> Pat ParamType -> Env) -> Env -> [Pat ParamType] -> Env
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Env -> Pat ParamType -> Env
envFromPat Env
env [Pat ParamType]
params
  e0' <- local (second (const env')) $ defuncSoacExp e0
  pure $ Lambda params e0' decl tp loc
defuncSoacExp Exp
e
  | Scalar Arrow {} <- Exp -> StructType
typeOf Exp
e = do
      (pats, body, tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e) Exp
e
      env <- askEnv
      let env' = (Env -> Pat ParamType -> Env) -> Env -> [Pat ParamType] -> Env
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Env -> Pat ParamType -> Env
envFromPat Env
env [Pat ParamType]
pats
      body' <- local (second (const env')) $ defuncExp' body
      pure $ Lambda pats body' Nothing (Info tp) (srclocOf e)
  | Bool
otherwise = Exp -> DefM Exp
defuncExp' Exp
e

etaExpand :: ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand :: ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ResRetType
e_t Exp
e = do
  let ([(PName, (Diet, StructType))]
ps, ResRetType
ret) = ResRetType -> ([(PName, (Diet, StructType))], ResRetType)
forall {dim}.
RetTypeBase dim Freshness
-> ([(PName, (Diet, TypeBase dim NoMode))],
    RetTypeBase dim Freshness)
getType ResRetType
e_t
  -- Some careful hackery to avoid duplicate names.
  (_, (params, vars)) <- ([(Pat ParamType, Exp)] -> ([Pat ParamType], [Exp]))
-> ([VName], [(Pat ParamType, Exp)])
-> ([VName], ([Pat ParamType], [Exp]))
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second [(Pat ParamType, Exp)] -> ([Pat ParamType], [Exp])
forall a b. [(a, b)] -> ([a], [b])
unzip (([VName], [(Pat ParamType, Exp)])
 -> ([VName], ([Pat ParamType], [Exp])))
-> DefM ([VName], [(Pat ParamType, Exp)])
-> DefM ([VName], ([Pat ParamType], [Exp]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([VName]
 -> (PName, (Diet, StructType))
 -> DefM ([VName], (Pat ParamType, Exp)))
-> [VName]
-> [(PName, (Diet, StructType))]
-> DefM ([VName], [(Pat ParamType, Exp)])
forall (m :: * -> *) (t :: * -> *) acc x y.
(Monad m, Traversable t) =>
(acc -> x -> m (acc, y)) -> acc -> t x -> m (acc, t y)
mapAccumLM [VName]
-> (PName, (Diet, StructType))
-> DefM ([VName], (Pat ParamType, Exp))
forall {m :: * -> *} {o} {b}.
MonadFreshNames m =>
[VName]
-> (PName, (o, TypeBase Exp b))
-> m ([VName], (PatBase Info VName (TypeBase Exp o), Exp))
f [] [(PName, (Diet, StructType))]
ps
  -- Important that we synthesize new existential names and substitute
  -- them into the (body) return type.
  ext' <- mapM newName $ retDims ret
  let extsubst =
        [(VName, Subst t)] -> Map VName (Subst t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst t)] -> Map VName (Subst t))
-> ([Subst t] -> [(VName, Subst t)])
-> [Subst t]
-> Map VName (Subst t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [Subst t] -> [(VName, Subst t)]
forall a b. [a] -> [b] -> [(a, b)]
zip (ResRetType -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims ResRetType
ret) ([Subst t] -> Map VName (Subst t))
-> [Subst t] -> Map VName (Subst t)
forall a b. (a -> b) -> a -> b
$
          (VName -> Subst t) -> [VName] -> [Subst t]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> (VName -> Exp) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
forall a. Monoid a => a
mempty (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
ext'
      ret' = TypeSubs -> ResRetType -> ResRetType
forall a. Substitutable a => TypeSubs -> a -> a
applySubst (VName
-> Map VName (Subst StructRetType) -> Maybe (Subst StructRetType)
forall k a. Ord k => k -> Map k a -> Maybe a
`M.lookup` Map VName (Subst StructRetType)
forall {t}. Map VName (Subst t)
extsubst) ResRetType
ret
      e' = Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply Exp
e ((Exp -> (Maybe VName, Exp)) -> [Exp] -> [(Maybe VName, Exp)]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe VName
forall a. Maybe a
Nothing,) [Exp]
vars) (AppRes -> Exp) -> AppRes -> Exp
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (ResType -> StructType) -> ResType -> StructType
forall a b. (a -> b) -> a -> b
$ ResRetType -> ResType
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType ResRetType
ret') [VName]
ext'
  pure (params, e', ret)
  where
    getType :: RetTypeBase dim Freshness
-> ([(PName, (Diet, TypeBase dim NoMode))],
    RetTypeBase dim Freshness)
getType (RetType [VName]
_ (Scalar (Arrow Freshness
_ PName
p Diet
d TypeBase dim NoMode
t1 RetTypeBase dim Freshness
t2))) =
      let ([(PName, (Diet, TypeBase dim NoMode))]
ps, RetTypeBase dim Freshness
r) = RetTypeBase dim Freshness
-> ([(PName, (Diet, TypeBase dim NoMode))],
    RetTypeBase dim Freshness)
getType RetTypeBase dim Freshness
t2
       in ((PName
p, (Diet
d, TypeBase dim NoMode
t1)) (PName, (Diet, TypeBase dim NoMode))
-> [(PName, (Diet, TypeBase dim NoMode))]
-> [(PName, (Diet, TypeBase dim NoMode))]
forall a. a -> [a] -> [a]
: [(PName, (Diet, TypeBase dim NoMode))]
ps, RetTypeBase dim Freshness
r)
    getType RetTypeBase dim Freshness
t = ([], RetTypeBase dim Freshness
t)

    f :: [VName]
-> (PName, (o, TypeBase Exp b))
-> m ([VName], (PatBase Info VName (TypeBase Exp o), Exp))
f [VName]
prev (PName
p, (o
d, TypeBase Exp b
t)) = do
      let t' :: TypeBase Exp o
t' = (b -> o) -> TypeBase Exp b -> TypeBase Exp o
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 (o -> b -> o
forall a b. a -> b -> a
const o
d) TypeBase Exp b
t
      x <- case PName
p of
        Named VName
x | VName
x VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [VName]
prev -> VName -> m VName
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
x
        PName
_ -> Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"eta_p"
      pure
        ( x : prev,
          ( Id x (Info t') mempty,
            Var (qualName x) (Info $ toStruct t') mempty
          )
        )

-- | Defunctionalize an indexing of a single array dimension.
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex (DimFix Exp
e1) = Exp -> DimIndexBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (Exp -> DimIndexBase Info VName)
-> ((Exp, StaticVal) -> Exp)
-> (Exp, StaticVal)
-> DimIndexBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp, StaticVal) -> Exp
forall a b. (a, b) -> a
fst ((Exp, StaticVal) -> DimIndexBase Info VName)
-> DefM (Exp, StaticVal) -> DefM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e1
defuncDimIndex (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)
-> DefM (Maybe Exp)
-> DefM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me1 DefM (Maybe Exp -> Maybe Exp -> DimIndexBase Info VName)
-> DefM (Maybe Exp) -> DefM (Maybe Exp -> DimIndexBase Info VName)
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me2 DefM (Maybe Exp -> DimIndexBase Info VName)
-> DefM (Maybe Exp) -> DefM (DimIndexBase Info VName)
forall a b. DefM (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Exp -> DefM (Maybe Exp)
defunc' Maybe Exp
me3
  where
    defunc' :: Maybe Exp -> DefM (Maybe Exp)
defunc' = (Exp -> DefM Exp) -> Maybe Exp -> DefM (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 -> DefM Exp
defuncExp'

envFromDimNames :: [VName] -> Env
envFromDimNames :: [VName] -> Env
envFromDimNames = [(VName, Binding)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Binding)] -> Env)
-> ([VName] -> [(VName, Binding)]) -> [VName] -> Env
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([VName] -> [Binding] -> [(VName, Binding)])
-> [Binding] -> [VName] -> [(VName, Binding)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [VName] -> [Binding] -> [(VName, Binding)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Binding -> [Binding]
forall a. a -> [a]
repeat Binding
d)
  where
    d :: Binding
d = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Freshness
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Freshness)
-> PrimType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | Defunctionalize a let-bound function, while preserving parameters
-- that have order 0 types (i.e., non-functional).
defuncLet ::
  [VName] ->
  [Pat ParamType] ->
  Exp ->
  ResRetType ->
  DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
defuncLet :: [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
defuncLet [VName]
dims ps :: [Pat ParamType]
ps@(Pat ParamType
pat : [Pat ParamType]
pats) Exp
body (RetType [VName]
ret_dims ResType
rettype)
  | Pat ParamType -> Bool
forall d o. Pat (TypeBase d o) -> Bool
patternOrderZero Pat ParamType
pat = do
      env <- DefM Env
askEnv
      let bound_by_pat = (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` FV -> Set VName
fvVars (Pat ParamType -> FV
forall o. Pat (TypeBase Exp o) -> FV
freeInPat Pat ParamType
pat))
          -- Take care to not include more size parameters than necessary.
          (pat_dims, rest_dims) = partition bound_by_pat dims
          env' = Env -> Pat ParamType -> Env
envFromPat Env
env Pat ParamType
pat Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> [VName] -> Env
envFromDimNames [VName]
pat_dims
      (rest_dims', pats', body', sv, sv_t) <-
        local (second (const env')) $ defuncLet rest_dims pats body $ RetType ret_dims rettype
      closure <- defuncFun dims ps body (RetType ret_dims rettype) mempty
      pure
        ( pat_dims ++ rest_dims',
          pat : pats',
          body',
          DynamicFun closure sv,
          sv_t
        )
  | Bool
otherwise = do
      (e, sv) <- [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncFun [VName]
dims [Pat ParamType]
ps Exp
body ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ret_dims ResType
rettype) SrcLoc
forall a. Monoid a => a
mempty
      pure ([], [], e, sv, resTypeFromSV sv)
defuncLet [VName]
_ [] Exp
body (RetType [VName]
_ ResType
rettype) = do
  (body', sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
body
  pure
    ( [],
      [],
      body',
      imposeType sv rettype,
      resTypeFromSV sv
    )
  where
    imposeType :: StaticVal -> ResType -> StaticVal
imposeType Dynamic {} ResType
t =
      ResType -> StaticVal
Dynamic ResType
t
    imposeType (RecordSV [(Name, StaticVal)]
fs1) (Scalar (Record Map Name ResType
fs2)) =
      [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$ Map Name StaticVal -> [(Name, StaticVal)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name StaticVal -> [(Name, StaticVal)])
-> Map Name StaticVal -> [(Name, StaticVal)]
forall a b. (a -> b) -> a -> b
$ (StaticVal -> ResType -> StaticVal)
-> Map Name StaticVal -> Map Name ResType -> Map Name StaticVal
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith StaticVal -> ResType -> StaticVal
imposeType ([(Name, StaticVal)] -> Map Name StaticVal
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, StaticVal)]
fs1) Map Name ResType
fs2
    imposeType StaticVal
sv ResType
_ = StaticVal
sv

instAnySizes :: (MonadFreshNames m) => [Pat ParamType] -> m [Pat ParamType]
instAnySizes :: forall (m :: * -> *).
MonadFreshNames m =>
[Pat ParamType] -> m [Pat ParamType]
instAnySizes = (Pat ParamType -> m (Pat ParamType))
-> [Pat ParamType] -> 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) -> [a] -> f [b]
traverse ((Pat ParamType -> m (Pat ParamType))
 -> [Pat ParamType] -> m [Pat ParamType])
-> (Pat ParamType -> m (Pat ParamType))
-> [Pat ParamType]
-> m [Pat ParamType]
forall a b. (a -> b) -> a -> b
$ (ParamType -> m ParamType) -> Pat ParamType -> 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 ((ParamType -> m ParamType) -> Pat ParamType -> m (Pat ParamType))
-> (ParamType -> m ParamType) -> Pat ParamType -> m (Pat ParamType)
forall a b. (a -> b) -> a -> b
$ (Exp -> m Exp) -> (Diet -> m Diet) -> ParamType -> 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 -> m Exp
forall {m :: * -> *}. MonadFreshNames m => Exp -> m Exp
onDim Diet -> m Diet
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    onDim :: Exp -> m Exp
onDim Exp
d
      | Just Int
_ <- Exp -> Maybe Int
isAnySize Exp
d = do
          v <- Name -> m VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
"size"
          pure $ sizeFromName (qualName v) mempty
    onDim Exp
d = Exp -> m Exp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
d

unboundSizes :: S.Set VName -> [Pat ParamType] -> [VName]
unboundSizes :: Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
params = [VName] -> [VName]
forall a. Ord a => [a] -> [a]
nubOrd ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ State [VName] [Pat ParamType] -> [VName] -> [VName]
forall s a. State s a -> s -> s
execState ([Pat ParamType] -> State [VName] [Pat ParamType]
forall {f :: * -> *} {d}.
[PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
     [VName]
     Identity
     [PatBase Info VName (TypeBase (ExpBase f VName) d)]
f [Pat ParamType]
params) []
  where
    f :: [PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
     [VName]
     Identity
     [PatBase Info VName (TypeBase (ExpBase f VName) d)]
f = (PatBase Info VName (TypeBase (ExpBase f VName) d)
 -> StateT
      [VName]
      Identity
      (PatBase Info VName (TypeBase (ExpBase f VName) d)))
-> [PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
     [VName]
     Identity
     [PatBase Info VName (TypeBase (ExpBase f VName) d)]
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 ((PatBase Info VName (TypeBase (ExpBase f VName) d)
  -> StateT
       [VName]
       Identity
       (PatBase Info VName (TypeBase (ExpBase f VName) d)))
 -> [PatBase Info VName (TypeBase (ExpBase f VName) d)]
 -> StateT
      [VName]
      Identity
      [PatBase Info VName (TypeBase (ExpBase f VName) d)])
-> (PatBase Info VName (TypeBase (ExpBase f VName) d)
    -> StateT
         [VName]
         Identity
         (PatBase Info VName (TypeBase (ExpBase f VName) d)))
-> [PatBase Info VName (TypeBase (ExpBase f VName) d)]
-> StateT
     [VName]
     Identity
     [PatBase Info VName (TypeBase (ExpBase f VName) d)]
forall a b. (a -> b) -> a -> b
$ (TypeBase (ExpBase f VName) d
 -> StateT [VName] Identity (TypeBase (ExpBase f VName) d))
-> PatBase Info VName (TypeBase (ExpBase f VName) d)
-> StateT
     [VName]
     Identity
     (PatBase Info VName (TypeBase (ExpBase f VName) d))
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 (ExpBase f VName) d
  -> StateT [VName] Identity (TypeBase (ExpBase f VName) d))
 -> PatBase Info VName (TypeBase (ExpBase f VName) d)
 -> StateT
      [VName]
      Identity
      (PatBase Info VName (TypeBase (ExpBase f VName) d)))
-> (TypeBase (ExpBase f VName) d
    -> StateT [VName] Identity (TypeBase (ExpBase f VName) d))
-> PatBase Info VName (TypeBase (ExpBase f VName) d)
-> StateT
     [VName]
     Identity
     (PatBase Info VName (TypeBase (ExpBase f VName) d))
forall a b. (a -> b) -> a -> b
$ (ExpBase f VName -> StateT [VName] Identity (ExpBase f VName))
-> (d -> StateT [VName] Identity d)
-> TypeBase (ExpBase f VName) d
-> StateT [VName] Identity (TypeBase (ExpBase f VName) d)
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 ExpBase f VName -> StateT [VName] Identity (ExpBase f VName)
forall {m :: * -> *} {f :: * -> *}.
MonadState [VName] m =>
ExpBase f VName -> m (ExpBase f VName)
onDim d -> StateT [VName] Identity d
forall a. a -> StateT [VName] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    bound :: Set VName
bound = Set VName
bound_sizes Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)
    onDim :: ExpBase f VName -> m (ExpBase f VName)
onDim (Var QualName VName
d f StructType
typ SrcLoc
loc) = do
      Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ ([VName] -> [VName]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d VName -> [VName] -> [VName]
forall a. a -> [a] -> [a]
:)
      ExpBase f VName -> m (ExpBase f VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpBase f VName -> m (ExpBase f VName))
-> ExpBase f VName -> m (ExpBase f VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> f StructType -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
d f StructType
typ SrcLoc
loc
    onDim ExpBase f VName
d = ExpBase f VName -> m (ExpBase f VName)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ExpBase f VName
d

unRetType :: ResRetType -> DefM AppRes
unRetType :: ResRetType -> DefM AppRes
unRetType (RetType [] ResType
t) = AppRes -> DefM AppRes
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AppRes -> DefM AppRes) -> AppRes -> DefM AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct ResType
t) []
unRetType (RetType [VName]
ext ResType
t) = do
  ext' <- (VName -> DefM VName) -> [VName] -> DefM [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 VName -> DefM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName [VName]
ext
  let extsubst =
        [(VName, Subst t)] -> Map VName (Subst t)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Subst t)] -> Map VName (Subst t))
-> ([Subst t] -> [(VName, Subst t)])
-> [Subst t]
-> Map VName (Subst t)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [VName] -> [Subst t] -> [(VName, Subst t)]
forall a b. [a] -> [b] -> [(a, b)]
zip [VName]
ext ([Subst t] -> Map VName (Subst t))
-> [Subst t] -> Map VName (Subst t)
forall a b. (a -> b) -> a -> b
$
          (VName -> Subst t) -> [VName] -> [Subst t]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> Subst t
forall t. Exp -> Subst t
ExpSubst (Exp -> Subst t) -> (VName -> Exp) -> VName -> Subst t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (QualName VName -> SrcLoc -> Exp)
-> SrcLoc -> QualName VName -> Exp
forall a b c. (a -> b -> c) -> b -> a -> c
flip QualName VName -> SrcLoc -> Exp
sizeFromName SrcLoc
forall a. Monoid a => a
mempty (QualName VName -> Exp)
-> (VName -> QualName VName) -> VName -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> QualName VName
forall v. v -> QualName v
qualName) [VName]
ext'
  pure $ AppRes (applySubst (`M.lookup` extsubst) $ toStruct t) ext'

defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction e :: Exp
e@(Var QualName VName
qn (Info StructType
t) SrcLoc
loc) Int
num_args = do
  let ([ParamType]
argtypes, StructType
rettype) = ([(PName, ParamType)] -> [ParamType])
-> ([(PName, ParamType)], StructType) -> ([ParamType], StructType)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (((PName, ParamType) -> ParamType)
-> [(PName, ParamType)] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (PName, ParamType) -> ParamType
forall a b. (a, b) -> b
snd) (([(PName, ParamType)], StructType) -> ([ParamType], StructType))
-> ([(PName, ParamType)], StructType) -> ([ParamType], StructType)
forall a b. (a -> b) -> a -> b
$ StructType -> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType StructType
t
  sv <- StructType -> VName -> DefM StaticVal
lookupVar (StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct StructType
t) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)

  case sv of
    DynamicFun (Exp, StaticVal)
_ StaticVal
_
      | StaticVal -> Int -> Bool
fullyApplied StaticVal
sv Int
num_args -> do
          -- We still need to update the types in case the dynamic
          -- function returns a higher-order term.
          let ([ParamType]
argtypes', ResType
rettype') = StaticVal -> [ParamType] -> ([ParamType], ResType)
dynamicFunType StaticVal
sv [ParamType]
argtypes
          (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (StructType -> Info StructType
forall a. a -> Info a
Info ([ParamType] -> ResRetType -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [ParamType]
argtypes' (ResRetType -> StructType) -> ResRetType -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] ResType
rettype')) SrcLoc
loc, StaticVal
sv)
      | (ParamType -> Bool) -> [ParamType] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ParamType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero [ParamType]
argtypes,
        StructType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero StructType
rettype -> do
          (params, body, ret) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t) Exp
e
          defuncFun [] params body ret mempty
      | Bool
otherwise -> do
          fname <- Name -> DefM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName (Name -> DefM VName) -> Name -> DefM VName
forall a b. (a -> b) -> a -> b
$ Name
"dyn_" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
          let (pats, e0, sv') = liftDynFun (nameFromText (prettyText qn)) sv num_args
              (argtypes', rettype') = dynamicFunType sv' argtypes
              dims' = [VName]
forall a. Monoid a => a
mempty

          -- Ensure that no parameter sizes are AnySize.  The internaliser
          -- expects this.  This is easy, because they are all
          -- first-order.
          globals <- asks $ M.keysSet . fst
          let bound_sizes = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
dims' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
globals
          pats' <- instAnySizes pats
          let dims'' = [VName]
dims' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
pats'

          liftValDec fname (RetType [] rettype') dims'' pats' e0
          pure
            ( Var
                (qualName fname)
                (Info (foldFunType argtypes' $ RetType [] rettype'))
                loc,
              sv'
            )
    StaticVal
IntrinsicSV -> (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e, StaticVal
IntrinsicSV)
    StaticVal
_ -> (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (QualName VName -> Info StructType -> SrcLoc -> Exp
forall (f :: * -> *) vn.
QualName vn -> f StructType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (StructType -> Info StructType
forall a. a -> Info a
Info (StaticVal -> StructType
structTypeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncApplyFunction Exp
e Int
_ = Exp -> DefM (Exp, StaticVal)
defuncExp Exp
e

-- Embed some information about the original function
-- into the name of the lifted function, to make the
-- result slightly more human-readable.
liftedName :: Int -> Exp -> Name
liftedName :: Int -> Exp -> Name
liftedName Int
i (Var QualName VName
f Info StructType
_ SrcLoc
_) =
  Name
"defunc_" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> [Char] -> Name
nameFromString (Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i) Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> Name
"_" Name -> Name -> Name
forall a. Semigroup a => a -> a -> a
<> VName -> Name
baseName (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
f)
liftedName Int
i (AppExp (Apply Exp
f NonEmpty (Info (Maybe VName), Exp)
_ SrcLoc
_) Info AppRes
_) =
  Int -> Exp -> Name
liftedName (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Exp
f
liftedName Int
_ Exp
_ = Name
"defunc"

defuncApplyArg ::
  (Name, SrcLoc) ->
  (Exp, StaticVal) ->
  ((Maybe VName, Exp), [ParamType]) ->
  DefM (Exp, StaticVal)
defuncApplyArg :: (Name, SrcLoc)
-> (Exp, StaticVal)
-> ((Maybe VName, Exp), [ParamType])
-> DefM (Exp, StaticVal)
defuncApplyArg (Name
fname_s, SrcLoc
floc) (Exp
f', fsv :: StaticVal
fsv@(LambdaSV Pat ParamType
pat ResRetType
lam_e_t Exp
lam_e Env
closure_env)) ((Maybe VName
argext, Exp
arg), [ParamType]
_) = do
  (arg', arg_sv) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
arg

  -- Build a call to the lifted specialisation named 'lifted' with the
  -- given (first-order) return type.  Reuses the closure value 'f'' and
  -- the (already defunctionalised) argument.
  let mkCall VName
lifted ResRetType
lifted_ret = do
        let f_t :: StructType
f_t = StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
f'
            arg_t :: StructType
arg_t = StructType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
arg'
            fname_t :: StructType
fname_t = [ParamType] -> ResRetType -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe StructType
f_t, Diet -> StructType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam (ParamType -> Diet
forall shape. TypeBase shape Diet -> Diet
diet (Pat ParamType -> ParamType
forall d o. Pat (TypeBase d o) -> TypeBase d o
patternType Pat ParamType
pat)) StructType
arg_t] ResRetType
lifted_ret
            fname' :: Exp
fname' = 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
lifted) (StructType -> Info StructType
forall a. a -> Info a
Info StructType
fname_t) SrcLoc
floc
        callret <- ResRetType -> DefM AppRes
unRetType ResRetType
lifted_ret
        pure $ mkApply fname' [(Nothing, f'), (argext, arg')] callret

  -- Applying this function to this static argument may already have been
  -- lifted; if so, reuse that lifting. This is what ties off recursion,
  -- since a recursive occurrence re-applies the same function to the same
  -- static argument. See Note [Lifting and recursion].
  let key = (StaticVal
fsv, StaticVal
arg_sv)

  memhit <- lookupLift key
  case memhit of
    Just (VName
lifted, ResRetType
lifted_ret, StaticVal
sv) -> do
      e <- VName -> ResRetType -> DefM Exp
mkCall VName
lifted ResRetType
lifted_ret
      pure (e, sv)
    Maybe (VName, ResRetType, StaticVal)
Nothing -> do
      let env' :: Env
env' = Env -> Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Env
closure_env Pat ParamType
pat StaticVal
arg_sv
          dims :: [VName]
dims = [VName]
forall a. Monoid a => a
mempty

      -- This slot yields the first-order body of the function (rather
      -- than peeling off another function-typed parameter) exactly when
      -- its result is order-zero.  There we must record the lifted name
      -- before defunctionalising the body, so that a recursive
      -- occurrence inside it finds the entry rather than re-inlining
      -- forever.  The recorded return type is the declared one, since
      -- the body's inferred refinement is not available yet; we replace
      -- it with the refined one once the body has been lifted.
      let is_body :: Bool
is_body = ResType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero (ResType -> Bool) -> ResType -> Bool
forall a b. (a -> b) -> a -> b
$ ResRetType -> ResType
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType ResRetType
lam_e_t
      fname <- Name -> DefM VName
forall (m :: * -> *). MonadFreshNames m => Name -> m VName
newVName Name
fname_s
      let memo_ret = [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType (ResRetType -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims ResRetType
lam_e_t) (ResRetType -> ResType
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType ResRetType
lam_e_t)
      when is_body $
        insertLift key (fname, memo_ret, Dynamic $ retType memo_ret)
      (lam_e', sv) <-
        localNewEnv env' $
          defuncExp lam_e

      let closure_pat = [VName] -> Env -> Pat ParamType
buildEnvPat [VName]
dims Env
closure_env
          pat' = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
arg_sv

      globals <- asks $ M.keysSet . fst

      -- Lift lambda to top-level function definition.  We put in
      -- a lot of effort to try to infer the freshness attributes
      -- of the lifted function, but this is ultimately all a sham
      -- and a hack.  There is some piece we're missing.
      let params = [Pat ParamType
closure_pat, Pat ParamType
pat']
          lifted_rettype =
            [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType (ResRetType -> [VName]
forall dim o. RetTypeBase dim o -> [VName]
retDims ResRetType
lam_e_t) (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$
              ResType -> ResType -> ResType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes (ResRetType -> ResType
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType ResRetType
lam_e_t) (StaticVal -> ResType
resTypeFromSV StaticVal
sv)

          already_bound =
            Set VName
globals Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> (Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params)

          more_dims =
            Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [VName]
forall a b. (a -> b) -> a -> b
$
              (VName -> Bool) -> Set VName -> Set VName
forall a. (a -> Bool) -> Set a -> Set a
S.filter (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set VName
already_bound) (Set VName -> Set VName) -> Set VName -> Set VName
forall a b. (a -> b) -> a -> b
$
                (Pat ParamType -> Set VName) -> [Pat ParamType] -> Set VName
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> Set VName
patternArraySizes [Pat ParamType]
params

      -- Ensure that no parameter sizes are AnySize.  The internaliser
      -- expects this.  This is easy, because they are all
      -- first-order.
      let bound_sizes = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ([VName]
dims [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
more_dims) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
globals
      params' <- instAnySizes params

      liftValDec
        fname
        lifted_rettype
        (dims ++ more_dims ++ unboundSizes bound_sizes params')
        params'
        lam_e'

      insertLift key (fname, lifted_rettype, sv)

      e <- mkCall fname lifted_rettype
      pure (e, sv)
-- If 'f' is a dynamic function, we just leave the application in
-- place, but we update the types since it may be partially
-- applied or return a higher-order value.
defuncApplyArg (Name, SrcLoc)
_ (Exp
f', DynamicFun (Exp, StaticVal)
_ StaticVal
sv) ((Maybe VName
argext, Exp
arg), [ParamType]
argtypes) = do
  (arg', _) <- Exp -> DefM (Exp, StaticVal)
defuncExp Exp
arg
  let (argtypes', rettype) = dynamicFunType sv argtypes
      restype = [ParamType] -> ResRetType -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [ParamType]
argtypes' ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] ResType
rettype)
      callret = StructType -> [VName] -> AppRes
AppRes StructType
restype []
      apply_e = Exp -> [(Maybe VName, Exp)] -> AppRes -> Exp
forall vn.
ExpBase Info vn
-> [(Maybe VName, ExpBase Info vn)] -> AppRes -> ExpBase Info vn
mkApply Exp
f' [(Maybe VName
argext, Exp
arg')] AppRes
callret
  pure (apply_e, sv)
--
defuncApplyArg (Name
fname_s, SrcLoc
_) (Exp
_, StaticVal
sv) ((Maybe VName
_, Exp
arg), [ParamType]
_) =
  [Char] -> DefM (Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM (Exp, StaticVal))
-> [Char] -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$
    [Char]
"defuncApplyArg: cannot apply StaticVal\n"
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
"\nFunction name: "
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
fname_s
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Char]
"\nArgument: "
      [Char] -> ShowS
forall a. Semigroup a => a -> a -> a
<> Exp -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Exp
arg

updateReturn :: AppRes -> Exp -> Exp
updateReturn :: AppRes -> Exp -> Exp
updateReturn (AppRes StructType
ret1 [VName]
ext1) (AppExp AppExpBase Info VName
apply (Info (AppRes StructType
ret2 [VName]
ext2))) =
  AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp AppExpBase Info VName
apply (Info AppRes -> Exp) -> Info AppRes -> Exp
forall a b. (a -> b) -> a -> b
$ AppRes -> Info AppRes
forall a. a -> Info a
Info (AppRes -> Info AppRes) -> AppRes -> Info AppRes
forall a b. (a -> b) -> a -> b
$ StructType -> [VName] -> AppRes
AppRes (StructType -> StructType -> StructType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes StructType
ret1 StructType
ret2) ([VName]
ext1 [VName] -> [VName] -> [VName]
forall a. Semigroup a => a -> a -> a
<> [VName]
ext2)
updateReturn AppRes
_ Exp
e = Exp
e

defuncApply :: Exp -> NE.NonEmpty (Maybe VName, Exp) -> AppRes -> SrcLoc -> DefM (Exp, StaticVal)
defuncApply :: Exp
-> NonEmpty (Maybe VName, Exp)
-> AppRes
-> SrcLoc
-> DefM (Exp, StaticVal)
defuncApply Exp
f NonEmpty (Maybe VName, Exp)
args AppRes
appres SrcLoc
loc = do
  (f', f_sv) <- Exp -> Int -> DefM (Exp, StaticVal)
defuncApplyFunction Exp
f (NonEmpty (Maybe VName, Exp) -> Int
forall a. NonEmpty a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length NonEmpty (Maybe VName, Exp)
args)
  case f_sv of
    StaticVal
IntrinsicSV -> do
      args' <- ((Maybe VName, Exp) -> (Info (Maybe VName), Exp))
-> NonEmpty (Maybe VName, Exp)
-> NonEmpty (Info (Maybe VName), Exp)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Maybe VName -> Info (Maybe VName))
-> (Maybe VName, Exp) -> (Info (Maybe VName), Exp)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info) (NonEmpty (Maybe VName, Exp) -> NonEmpty (Info (Maybe VName), Exp))
-> DefM (NonEmpty (Maybe VName, Exp))
-> DefM (NonEmpty (Info (Maybe VName), Exp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Maybe VName, Exp) -> DefM (Maybe VName, Exp))
-> NonEmpty (Maybe VName, Exp)
-> DefM (NonEmpty (Maybe VName, 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) -> NonEmpty a -> f (NonEmpty b)
traverse ((Exp -> DefM Exp) -> (Maybe VName, Exp) -> DefM (Maybe VName, 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) -> (Maybe VName, a) -> f (Maybe VName, b)
traverse Exp -> DefM Exp
defuncSoacExp) NonEmpty (Maybe VName, Exp)
args
      let e' = AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (Info (Maybe VName), Exp)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Maybe VName), Exp)
args' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
appres)
      intrinsicOrHole e'
    HoleSV {} -> do
      args' <- ((Maybe VName, Exp) -> (Info (Maybe VName), Exp))
-> NonEmpty (Maybe VName, Exp)
-> NonEmpty (Info (Maybe VName), Exp)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Maybe VName -> Info (Maybe VName))
-> (Maybe VName, Exp) -> (Info (Maybe VName), Exp)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Maybe VName -> Info (Maybe VName)
forall a. a -> Info a
Info) (NonEmpty (Maybe VName, Exp) -> NonEmpty (Info (Maybe VName), Exp))
-> DefM (NonEmpty (Maybe VName, Exp))
-> DefM (NonEmpty (Info (Maybe VName), Exp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Maybe VName, Exp) -> DefM (Maybe VName, Exp))
-> NonEmpty (Maybe VName, Exp)
-> DefM (NonEmpty (Maybe VName, 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) -> NonEmpty a -> f (NonEmpty b)
traverse ((Exp -> DefM Exp) -> (Maybe VName, Exp) -> DefM (Maybe VName, 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) -> (Maybe VName, a) -> f (Maybe VName, b)
traverse ((Exp -> DefM Exp)
 -> (Maybe VName, Exp) -> DefM (Maybe VName, Exp))
-> (Exp -> DefM Exp)
-> (Maybe VName, Exp)
-> DefM (Maybe VName, Exp)
forall a b. (a -> b) -> a -> b
$ ((Exp, StaticVal) -> Exp) -> DefM (Exp, StaticVal) -> DefM Exp
forall a b. (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Exp, StaticVal) -> Exp
forall a b. (a, b) -> a
fst (DefM (Exp, StaticVal) -> DefM Exp)
-> (Exp -> DefM (Exp, StaticVal)) -> Exp -> DefM Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> DefM (Exp, StaticVal)
defuncExp) NonEmpty (Maybe VName, Exp)
args
      let e' = AppExpBase Info VName -> Info AppRes -> Exp
forall (f :: * -> *) vn.
AppExpBase f vn -> f AppRes -> ExpBase f vn
AppExp (Exp
-> NonEmpty (Info (Maybe VName), Exp)
-> SrcLoc
-> AppExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (f (Maybe VName), ExpBase f vn)
-> SrcLoc
-> AppExpBase f vn
Apply Exp
f' NonEmpty (Info (Maybe VName), Exp)
args' SrcLoc
loc) (AppRes -> Info AppRes
forall a. a -> Info a
Info AppRes
appres)
      intrinsicOrHole e'
    StaticVal
_ -> do
      let fname :: Name
fname = Int -> Exp -> Name
liftedName Int
0 Exp
f
          ([(PName, ParamType)]
argtypes, StructType
_) = StructType -> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType (StructType -> ([(PName, ParamType)], StructType))
-> StructType -> ([(PName, ParamType)], StructType)
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
f
      ((Exp, StaticVal) -> (Exp, StaticVal))
-> DefM (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> DefM a -> DefM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Exp -> Exp) -> (Exp, StaticVal) -> (Exp, StaticVal)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((Exp -> Exp) -> (Exp, StaticVal) -> (Exp, StaticVal))
-> (Exp -> Exp) -> (Exp, StaticVal) -> (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$ AppRes -> Exp -> Exp
updateReturn AppRes
appres) (DefM (Exp, StaticVal) -> DefM (Exp, StaticVal))
-> DefM (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$
        ((Exp, StaticVal)
 -> ((Maybe VName, Exp), [ParamType]) -> DefM (Exp, StaticVal))
-> (Exp, StaticVal)
-> NonEmpty ((Maybe VName, Exp), [ParamType])
-> DefM (Exp, StaticVal)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ((Name, SrcLoc)
-> (Exp, StaticVal)
-> ((Maybe VName, Exp), [ParamType])
-> DefM (Exp, StaticVal)
defuncApplyArg (Name
fname, SrcLoc
loc)) (Exp
f', StaticVal
f_sv) (NonEmpty ((Maybe VName, Exp), [ParamType])
 -> DefM (Exp, StaticVal))
-> NonEmpty ((Maybe VName, Exp), [ParamType])
-> DefM (Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$
          NonEmpty (Maybe VName, Exp)
-> NonEmpty [ParamType]
-> NonEmpty ((Maybe VName, Exp), [ParamType])
forall a b. NonEmpty a -> NonEmpty b -> NonEmpty (a, b)
NE.zip NonEmpty (Maybe VName, Exp)
args (NonEmpty [ParamType]
 -> NonEmpty ((Maybe VName, Exp), [ParamType]))
-> ([(PName, ParamType)] -> NonEmpty [ParamType])
-> [(PName, ParamType)]
-> NonEmpty ((Maybe VName, Exp), [ParamType])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ParamType] -> NonEmpty [ParamType]
forall (f :: * -> *) a. Foldable f => f a -> NonEmpty [a]
NE.tails ([ParamType] -> NonEmpty [ParamType])
-> ([(PName, ParamType)] -> [ParamType])
-> [(PName, ParamType)]
-> NonEmpty [ParamType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((PName, ParamType) -> ParamType)
-> [(PName, ParamType)] -> [ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (PName, ParamType) -> ParamType
forall a b. (a, b) -> b
snd ([(PName, ParamType)]
 -> NonEmpty ((Maybe VName, Exp), [ParamType]))
-> [(PName, ParamType)]
-> NonEmpty ((Maybe VName, Exp), [ParamType])
forall a b. (a -> b) -> a -> b
$
            [(PName, ParamType)]
argtypes
  where
    intrinsicOrHole :: Exp -> DefM (Exp, StaticVal)
intrinsicOrHole Exp
e' = do
      -- If the intrinsic is fully applied, then we are done.
      -- Otherwise we need to eta-expand it and recursively
      -- defunctionalise. XXX: might it be better to simply eta-expand
      -- immediately any time we encounter a non-fully-applied
      -- intrinsic?
      if [(PName, ParamType)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([(PName, ParamType)] -> Bool) -> [(PName, ParamType)] -> Bool
forall a b. (a -> b) -> a -> b
$ ([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a, b) -> a
fst (([(PName, ParamType)], StructType) -> [(PName, ParamType)])
-> ([(PName, ParamType)], StructType) -> [(PName, ParamType)]
forall a b. (a -> b) -> a -> b
$ StructType -> ([(PName, ParamType)], StructType)
forall dim o.
TypeBase dim o
-> ([(PName, TypeBase dim Diet)], TypeBase dim NoMode)
unfoldFunType (StructType -> ([(PName, ParamType)], StructType))
-> StructType -> ([(PName, ParamType)], StructType)
forall a b. (a -> b) -> a -> b
$ AppRes -> StructType
appResType AppRes
appres
        then (Exp, StaticVal) -> DefM (Exp, StaticVal)
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp
e', ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ AppRes -> StructType
appResType AppRes
appres)
        else do
          (pats, body, tp) <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ([VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (ResType -> ResRetType) -> ResType -> ResRetType
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh (StructType -> ResType) -> StructType -> ResType
forall a b. (a -> b) -> a -> b
$ Exp -> StructType
typeOf Exp
e') Exp
e'
          defuncExp $ Lambda pats body Nothing (Info tp) mempty

-- | Check if a 'StaticVal' and a given application depth corresponds
-- to a fully applied dynamic function.
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied :: StaticVal -> Int -> Bool
fullyApplied (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) Int
depth
  | Int
depth Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Bool
False
  | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = StaticVal -> Int -> Bool
fullyApplied StaticVal
sv (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
fullyApplied StaticVal
_ Int
_ = Bool
True

-- | Converts a dynamic function 'StaticVal' into a list of
-- dimensions, a list of parameters, a function body, and the
-- appropriate static value for applying the function at the given
-- depth of partial application.
liftDynFun :: Name -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun :: Name -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun Name
_ (DynamicFun (Exp
e, StaticVal
sv) StaticVal
_) Int
0 = ([], Exp
e, StaticVal
sv)
liftDynFun Name
s (DynamicFun clsr :: (Exp, StaticVal)
clsr@(Exp
_, LambdaSV Pat ParamType
pat ResRetType
_ Exp
_ Env
_) StaticVal
sv) Int
d
  | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 =
      let ([Pat ParamType]
pats, Exp
e', StaticVal
sv') = Name -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)
liftDynFun Name
s StaticVal
sv (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
       in (Pat ParamType
pat Pat ParamType -> [Pat ParamType] -> [Pat ParamType]
forall a. a -> [a] -> [a]
: [Pat ParamType]
pats, Exp
e', (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp, StaticVal)
clsr StaticVal
sv')
liftDynFun Name
s StaticVal
sv Int
d =
  [Char] -> ([Pat ParamType], Exp, StaticVal)
forall a. HasCallStack => [Char] -> a
error ([Char] -> ([Pat ParamType], Exp, StaticVal))
-> [Char] -> ([Pat ParamType], Exp, StaticVal)
forall a b. (a -> b) -> a -> b
$
    Name -> [Char]
nameToString Name
s
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" Tried to lift a StaticVal "
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
100 (StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv)
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
", but expected a dynamic function.\n"
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Int
d

-- | Converts a pattern to an environment that binds the individual names of the
-- pattern to their corresponding types wrapped in a 'Dynamic' static value.
envFromPat :: Env -> Pat ParamType -> Env
envFromPat :: Env -> Pat ParamType -> Env
envFromPat Env
env Pat ParamType
pat = case Pat ParamType
pat of
  TuplePat [Pat ParamType]
ps SrcLoc
_ -> (Env -> Pat ParamType -> Env) -> Env -> [Pat ParamType] -> Env
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Env -> Pat ParamType -> Env
envFromPat Env
env [Pat ParamType]
ps
  RecordPat [(L Name, Pat ParamType)]
fs SrcLoc
_ -> (Env -> Pat ParamType -> Env) -> Env -> [Pat ParamType] -> Env
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Env -> Pat ParamType -> Env
envFromPat Env
env ([Pat ParamType] -> Env) -> [Pat ParamType] -> Env
forall a b. (a -> b) -> a -> b
$ ((L Name, Pat ParamType) -> Pat ParamType)
-> [(L Name, Pat ParamType)] -> [Pat ParamType]
forall a b. (a -> b) -> [a] -> [b]
map (L Name, Pat ParamType) -> Pat ParamType
forall a b. (a, b) -> b
snd [(L Name, Pat ParamType)]
fs
  PatParens Pat ParamType
p SrcLoc
_ -> Env -> Pat ParamType -> Env
envFromPat Env
env Pat ParamType
p
  PatAttr AttrInfo VName
_ Pat ParamType
p SrcLoc
_ -> Env -> Pat ParamType -> Env
envFromPat Env
env Pat ParamType
p
  Id VName
vn (Info ParamType
t) SrcLoc
_ -> VName -> Binding -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ParamType -> ResType
paramToRes ParamType
t) Env
env
  Wildcard Info ParamType
_ SrcLoc
_ -> Env
env
  PatAscription Pat ParamType
p TypeExp Exp VName
_ SrcLoc
_ -> Env -> Pat ParamType -> Env
envFromPat Env
env Pat ParamType
p
  PatLit {} -> Env
env
  PatConstr Name
_ Info ParamType
_ [Pat ParamType]
ps SrcLoc
_ -> (Env -> Pat ParamType -> Env) -> Env -> [Pat ParamType] -> Env
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Env -> Pat ParamType -> Env
envFromPat Env
env [Pat ParamType]
ps

-- | Given a closure environment, construct a record pattern that
-- binds the closed over variables.  Insert wildcard for any patterns
-- that would otherwise clash with size parameters.
buildEnvPat :: [VName] -> Env -> Pat ParamType
buildEnvPat :: [VName] -> Env -> Pat ParamType
buildEnvPat [VName]
sizes Env
env = [(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat (((VName, Binding) -> (L Name, Pat ParamType))
-> [(VName, Binding)] -> [(L Name, Pat ParamType)]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (L Name, Pat ParamType)
buildField ([(VName, Binding)] -> [(L Name, Pat ParamType)])
-> [(VName, Binding)] -> [(L Name, Pat ParamType)]
forall a b. (a -> b) -> a -> b
$ Env -> [(VName, Binding)]
forall k a. Map k a -> [(k, a)]
M.toList Env
env) SrcLoc
forall a. Monoid a => a
mempty
  where
    buildField :: (VName, Binding) -> (L Name, Pat ParamType)
buildField (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
      ( Loc -> Name -> L Name
forall a. Loc -> a -> L a
L Loc
forall a. IsLocation a => a
noLoc (Name -> L Name) -> Name -> L Name
forall a b. (a -> b) -> a -> b
$ Text -> Name
nameFromText (VName -> Text
forall a. Pretty a => a -> Text
prettyText VName
vn),
        if VName
vn VName -> [VName] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
sizes
          then Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
forall a. Monoid a => a
mempty
          else VName -> Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
vn (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
forall a. Monoid a => a
mempty
      )

-- | Compute the corresponding type for the *representation* of a
-- given static value (not the original possibly higher-order value).
resTypeFromSV :: StaticVal -> ResType
resTypeFromSV :: StaticVal -> ResType
resTypeFromSV (Dynamic ResType
tp) =
  ResType
tp
resTypeFromSV (LambdaSV Pat ParamType
_ ResRetType
_ Exp
_ Env
env) =
  ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ([(Name, ResType)] -> ScalarTypeBase Exp Freshness)
-> [(Name, ResType)]
-> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name ResType -> ScalarTypeBase Exp Freshness
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name ResType -> ScalarTypeBase Exp Freshness)
-> ([(Name, ResType)] -> Map Name ResType)
-> [(Name, ResType)]
-> ScalarTypeBase Exp Freshness
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, ResType)] -> Map Name ResType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, ResType)] -> ResType) -> [(Name, ResType)] -> ResType
forall a b. (a -> b) -> a -> b
$
    ((VName, Binding) -> (Name, ResType))
-> [(VName, Binding)] -> [(Name, ResType)]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Name)
-> (Binding -> ResType) -> (VName, Binding) -> (Name, ResType)
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 (Text -> Name
nameFromText (Text -> Name) -> (VName -> Text) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Text
forall a. Pretty a => a -> Text
prettyText) (StaticVal -> ResType
resTypeFromSV (StaticVal -> ResType)
-> (Binding -> StaticVal) -> Binding -> ResType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV)) ([(VName, Binding)] -> [(Name, ResType)])
-> [(VName, Binding)] -> [(Name, ResType)]
forall a b. (a -> b) -> a -> b
$
      Env -> [(VName, Binding)]
forall k a. Map k a -> [(k, a)]
M.toList Env
env
resTypeFromSV (RecordSV [(Name, StaticVal)]
ls) =
  let ts :: [(Name, ResType)]
ts = ((Name, StaticVal) -> (Name, ResType))
-> [(Name, StaticVal)] -> [(Name, ResType)]
forall a b. (a -> b) -> [a] -> [b]
map ((StaticVal -> ResType) -> (Name, StaticVal) -> (Name, ResType)
forall a b. (a -> b) -> (Name, a) -> (Name, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> ResType
resTypeFromSV) [(Name, StaticVal)]
ls
   in ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Map Name ResType -> ScalarTypeBase Exp Freshness
forall dim o. Map Name (TypeBase dim o) -> ScalarTypeBase dim o
Record (Map Name ResType -> ScalarTypeBase Exp Freshness)
-> Map Name ResType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ [(Name, ResType)] -> Map Name ResType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, ResType)]
ts
resTypeFromSV (DynamicFun (Exp
_, StaticVal
sv) StaticVal
_) =
  StaticVal -> ResType
resTypeFromSV StaticVal
sv
resTypeFromSV (SumSV Name
name [StaticVal]
svs [(Name, [ResType])]
fields) =
  let svs' :: [ResType]
svs' = (StaticVal -> ResType) -> [StaticVal] -> [ResType]
forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> ResType
resTypeFromSV [StaticVal]
svs
   in ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ Map Name [ResType] -> ScalarTypeBase Exp Freshness
forall dim o. Map Name [TypeBase dim o] -> ScalarTypeBase dim o
Sum (Map Name [ResType] -> ScalarTypeBase Exp Freshness)
-> Map Name [ResType] -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ Name -> [ResType] -> Map Name [ResType] -> Map Name [ResType]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name [ResType]
svs' (Map Name [ResType] -> Map Name [ResType])
-> Map Name [ResType] -> Map Name [ResType]
forall a b. (a -> b) -> a -> b
$ [(Name, [ResType])] -> Map Name [ResType]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [ResType])]
fields
resTypeFromSV (HoleSV StructType
t SrcLoc
_) =
  Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t
resTypeFromSV StaticVal
IntrinsicSV =
  [Char] -> ResType
forall a. HasCallStack => [Char] -> a
error [Char]
"Tried to get the type from the static value of an intrinsic."

structTypeFromSV :: StaticVal -> StructType
structTypeFromSV :: StaticVal -> StructType
structTypeFromSV = ResType -> StructType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct (ResType -> StructType)
-> (StaticVal -> ResType) -> StaticVal -> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> ResType
resTypeFromSV

-- | The 'Diet' of the result means something else than the freshness it is
-- converted from: that the function may consume the parameter.  The two agree
-- on what matters, since a function owns what it consumes.
paramTypeFromSV :: StaticVal -> ParamType
paramTypeFromSV :: StaticVal -> ParamType
paramTypeFromSV = ResType -> ParamType
resToParam (ResType -> ParamType)
-> (StaticVal -> ResType) -> StaticVal -> ParamType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> ResType
resTypeFromSV

-- | Construct the type for a fully-applied dynamic function from its
-- static value and the original types of its arguments.
dynamicFunType :: StaticVal -> [ParamType] -> ([ParamType], ResType)
dynamicFunType :: StaticVal -> [ParamType] -> ([ParamType], ResType)
dynamicFunType (DynamicFun (Exp, StaticVal)
_ StaticVal
sv) (ParamType
p : [ParamType]
ps) =
  let ([ParamType]
ps', ResType
ret) = StaticVal -> [ParamType] -> ([ParamType], ResType)
dynamicFunType StaticVal
sv [ParamType]
ps
   in (ParamType
p ParamType -> [ParamType] -> [ParamType]
forall a. a -> [a] -> [a]
: [ParamType]
ps', ResType
ret)
dynamicFunType StaticVal
sv [ParamType]
_ = ([], StaticVal -> ResType
resTypeFromSV StaticVal
sv)

-- | Match a pattern with its static value. Returns an environment
-- with the identifier components of the pattern mapped to the
-- corresponding subcomponents of the static value.  If this function
-- returns 'Nothing', then it corresponds to an unmatchable case.
-- These should only occur for 'Match' expressions.
matchPatSV :: Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV :: Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env (TuplePat [Pat ParamType]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls) =
  (Env -> (Pat ParamType, (Name, StaticVal)) -> Maybe Env)
-> Env -> [(Pat ParamType, (Name, StaticVal))] -> Maybe Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat ParamType
p, (Name
_, StaticVal
sv)) -> Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env' Pat ParamType
p StaticVal
sv) Env
env ([(Pat ParamType, (Name, StaticVal))] -> Maybe Env)
-> [(Pat ParamType, (Name, StaticVal))] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [Pat ParamType]
-> [(Name, StaticVal)] -> [(Pat ParamType, (Name, StaticVal))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Pat ParamType]
ps [(Name, StaticVal)]
ls
matchPatSV Env
env (RecordPat [(L Name, Pat ParamType)]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls)
  | [(Name, Pat ParamType)]
ps' <- ((Name, Pat ParamType) -> Name)
-> [(Name, Pat ParamType)] -> [(Name, Pat ParamType)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, Pat ParamType) -> Name
forall a b. (a, b) -> a
fst ([(Name, Pat ParamType)] -> [(Name, Pat ParamType)])
-> [(Name, Pat ParamType)] -> [(Name, Pat ParamType)]
forall a b. (a -> b) -> a -> b
$ ((L Name, Pat ParamType) -> (Name, Pat ParamType))
-> [(L Name, Pat ParamType)] -> [(Name, Pat ParamType)]
forall a b. (a -> b) -> [a] -> [b]
map ((L Name -> Name)
-> (L Name, Pat ParamType) -> (Name, Pat ParamType)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first L Name -> Name
forall a. L a -> a
unLoc) [(L Name, Pat ParamType)]
ps,
    [(Name, StaticVal)]
ls' <- ((Name, StaticVal) -> Name)
-> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, StaticVal) -> Name
forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls,
    ((Name, Pat ParamType) -> Name)
-> [(Name, Pat ParamType)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Pat ParamType) -> Name
forall a b. (a, b) -> a
fst [(Name, Pat ParamType)]
ps' [Name] -> [Name] -> Bool
forall a. Eq a => a -> a -> Bool
== ((Name, StaticVal) -> Name) -> [(Name, StaticVal)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StaticVal) -> Name
forall a b. (a, b) -> a
fst [(Name, StaticVal)]
ls' =
      (Env -> ((Name, Pat ParamType), (Name, StaticVal)) -> Maybe Env)
-> Env -> [((Name, Pat ParamType), (Name, StaticVal))] -> Maybe Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' ((Name
_, Pat ParamType
p), (Name
_, StaticVal
sv)) -> Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env' Pat ParamType
p StaticVal
sv) Env
env ([((Name, Pat ParamType), (Name, StaticVal))] -> Maybe Env)
-> [((Name, Pat ParamType), (Name, StaticVal))] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [(Name, Pat ParamType)]
-> [(Name, StaticVal)]
-> [((Name, Pat ParamType), (Name, StaticVal))]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Name, Pat ParamType)]
ps' [(Name, StaticVal)]
ls'
matchPatSV Env
env (PatParens Pat ParamType
pat SrcLoc
_) StaticVal
sv = Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat StaticVal
sv
matchPatSV Env
env (PatAttr AttrInfo VName
_ Pat ParamType
pat SrcLoc
_) StaticVal
sv = Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat StaticVal
sv
matchPatSV Env
env (Id VName
vn (Info ParamType
t) SrcLoc
_) StaticVal
sv =
  -- When matching a zero-order pattern with a StaticVal, the type of
  -- the pattern wins out.  This is important for propagating sizes
  -- (but probably reveals a flaw in our bookkeeping).
  Env -> Maybe Env
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Env -> Maybe Env) -> Env -> Maybe Env
forall a b. (a -> b) -> a -> b
$
    if ParamType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero ParamType
t
      then Env
dim_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> VName -> Binding -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ParamType -> ResType
paramToRes ParamType
t) Env
env
      else Env
dim_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> VName -> Binding -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing StaticVal
sv) Env
env
  where
    -- Extract all sizes that are potentially bound here. This is
    -- different from all free variables (see #2040).
    dim_env :: Env
dim_env = (Exp -> Env) -> (Diet -> Env) -> ParamType -> Env
forall m a b. Monoid m => (a -> m) -> (b -> m) -> TypeBase a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap Exp -> Env
forall {k} {f :: * -> *}. Ord k => ExpBase f k -> Map k Binding
onDim (Env -> Diet -> Env
forall a b. a -> b -> a
const Env
forall a. Monoid a => a
mempty) ParamType
t
    onDim :: ExpBase f k -> Map k Binding
onDim (Var QualName k
v f StructType
_ SrcLoc
_) = k -> Binding -> Map k Binding
forall k a. k -> a -> Map k a
M.singleton (QualName k -> k
forall vn. QualName vn -> vn
qualLeaf QualName k
v) Binding
i64
    onDim ExpBase f k
_ = Map k Binding
forall a. Monoid a => a
mempty
    i64 :: Binding
i64 = Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing (StaticVal -> Binding) -> StaticVal -> Binding
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
Dynamic (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Exp Freshness -> ResType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Exp Freshness -> ResType)
-> ScalarTypeBase Exp Freshness -> ResType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Exp Freshness
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Exp Freshness)
-> PrimType -> ScalarTypeBase Exp Freshness
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
matchPatSV Env
env (Wildcard Info ParamType
_ SrcLoc
_) StaticVal
_ = Env -> Maybe Env
forall a. a -> Maybe a
Just Env
env
matchPatSV Env
env (PatAscription Pat ParamType
pat TypeExp Exp VName
_ SrcLoc
_) StaticVal
sv = Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat StaticVal
sv
matchPatSV Env
env PatLit {} StaticVal
_ = Env -> Maybe Env
forall a. a -> Maybe a
Just Env
env
matchPatSV Env
env (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
_) (SumSV Name
c2 [StaticVal]
ls [(Name, [ResType])]
fs)
  | Name
c1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
c2 =
      (Env -> (Pat ParamType, StaticVal) -> Maybe Env)
-> Env -> [(Pat ParamType, StaticVal)] -> Maybe Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat ParamType
p, StaticVal
l) -> Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env' Pat ParamType
p StaticVal
l) Env
env ([(Pat ParamType, StaticVal)] -> Maybe Env)
-> [(Pat ParamType, StaticVal)] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> [StaticVal] -> [(Pat ParamType, StaticVal)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Pat ParamType]
ps [StaticVal]
ls
  | Just [ResType]
_ <- Name -> [(Name, [ResType])] -> Maybe [ResType]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
c1 [(Name, [ResType])]
fs =
      Maybe Env
forall a. Maybe a
Nothing
  | Bool
otherwise =
      [Char] -> Maybe Env
forall a. HasCallStack => [Char] -> a
error ([Char] -> Maybe Env) -> [Char] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV Env
env (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
_) (Dynamic (Scalar (Sum Map Name [ResType]
fs)))
  | Just [ResType]
ts <- Name -> Map Name [ResType] -> Maybe [ResType]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c1 Map Name [ResType]
fs =
      -- A higher-order pattern can only match an appropriate SumSV.
      if (ResType -> Bool) -> [ResType] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ResType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero [ResType]
ts
        then (Env -> (Pat ParamType, StaticVal) -> Maybe Env)
-> Env -> [(Pat ParamType, StaticVal)] -> Maybe Env
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Env
env' (Pat ParamType
p, StaticVal
sv) -> Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env' Pat ParamType
p StaticVal
sv) Env
env ([(Pat ParamType, StaticVal)] -> Maybe Env)
-> [(Pat ParamType, StaticVal)] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> [StaticVal] -> [(Pat ParamType, StaticVal)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Pat ParamType]
ps ([StaticVal] -> [(Pat ParamType, StaticVal)])
-> [StaticVal] -> [(Pat ParamType, StaticVal)]
forall a b. (a -> b) -> a -> b
$ (ResType -> StaticVal) -> [ResType] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map ResType -> StaticVal
svFromType [ResType]
ts
        else Maybe Env
forall a. Maybe a
Nothing
  | Bool
otherwise =
      [Char] -> Maybe Env
forall a. HasCallStack => [Char] -> a
error ([Char] -> Maybe Env) -> [Char] -> Maybe Env
forall a b. (a -> b) -> a -> b
$ [Char]
"matchPatSV: missing constructor in type: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Name
c1
matchPatSV Env
env Pat ParamType
pat (Dynamic ResType
t) = Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat (StaticVal -> Maybe Env) -> StaticVal -> Maybe Env
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
svFromType ResType
t
matchPatSV Env
env Pat ParamType
pat (HoleSV StructType
t SrcLoc
_) = Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat (StaticVal -> Maybe Env) -> StaticVal -> Maybe Env
forall a b. (a -> b) -> a -> b
$ ResType -> StaticVal
svFromType (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t
matchPatSV Env
_ Pat ParamType
pat StaticVal
sv =
  [Char] -> Maybe Env
forall a. HasCallStack => [Char] -> a
error ([Char] -> Maybe Env) -> [Char] -> Maybe Env
forall a b. (a -> b) -> a -> b
$
    [Char]
"Tried to match pattern\n"
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Pat ParamType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"\n with static value\n"
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv

alwaysMatchPatSV :: Env -> Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV :: Env -> Pat ParamType -> StaticVal -> Env
alwaysMatchPatSV Env
env Pat ParamType
pat StaticVal
sv = Env -> Maybe Env -> Env
forall a. a -> Maybe a -> a
fromMaybe Env
forall {b}. b
bad (Maybe Env -> Env) -> Maybe Env -> Env
forall a b. (a -> b) -> a -> b
$ Env -> Pat ParamType -> StaticVal -> Maybe Env
matchPatSV Env
env Pat ParamType
pat StaticVal
sv
  where
    bad :: b
bad = [Char] -> b
forall a. HasCallStack => [Char] -> a
error ([Char] -> b) -> [Char] -> b
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Char]
unlines [Pat ParamType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat, [Char]
"cannot match StaticVal", StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv]

-- | Given a pattern and the static value for the defunctionalized argument,
-- update the pattern to reflect the changes in the types.
updatePat :: Pat ParamType -> StaticVal -> Pat ParamType
updatePat :: Pat ParamType -> StaticVal -> Pat ParamType
updatePat (TuplePat [Pat ParamType]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs) =
  [Pat ParamType] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[PatBase f vn t] -> SrcLoc -> PatBase f vn t
TuplePat ((Pat ParamType -> StaticVal -> Pat ParamType)
-> [Pat ParamType] -> [StaticVal] -> [Pat ParamType]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Pat ParamType -> StaticVal -> Pat ParamType
updatePat [Pat ParamType]
ps ([StaticVal] -> [Pat ParamType]) -> [StaticVal] -> [Pat ParamType]
forall a b. (a -> b) -> a -> b
$ ((Name, StaticVal) -> StaticVal)
-> [(Name, StaticVal)] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map (Name, StaticVal) -> StaticVal
forall a b. (a, b) -> b
snd [(Name, StaticVal)]
svs) SrcLoc
loc
updatePat (RecordPat [(L Name, Pat ParamType)]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs)
  | [(L Name, Pat ParamType)]
ps' <- ((L Name, Pat ParamType) -> L Name)
-> [(L Name, Pat ParamType)] -> [(L Name, Pat ParamType)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (L Name, Pat ParamType) -> L Name
forall a b. (a, b) -> a
fst [(L Name, Pat ParamType)]
ps,
    [(Name, StaticVal)]
svs' <- ((Name, StaticVal) -> Name)
-> [(Name, StaticVal)] -> [(Name, StaticVal)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, StaticVal) -> Name
forall a b. (a, b) -> a
fst [(Name, StaticVal)]
svs =
      [(L Name, Pat ParamType)] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
[(L Name, PatBase f vn t)] -> SrcLoc -> PatBase f vn t
RecordPat
        (((L Name, Pat ParamType)
 -> (Name, StaticVal) -> (L Name, Pat ParamType))
-> [(L Name, Pat ParamType)]
-> [(Name, StaticVal)]
-> [(L Name, Pat ParamType)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(L Name
n, Pat ParamType
p) (Name
_, StaticVal
sv) -> (L Name
n, Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
p StaticVal
sv)) [(L Name, Pat ParamType)]
ps' [(Name, StaticVal)]
svs')
        SrcLoc
loc
updatePat (PatParens Pat ParamType
pat SrcLoc
loc) StaticVal
sv =
  Pat ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
PatBase f vn t -> SrcLoc -> PatBase f vn t
PatParens (Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv) SrcLoc
loc
updatePat (PatAttr AttrInfo VName
attr Pat ParamType
pat SrcLoc
loc) StaticVal
sv =
  AttrInfo VName -> Pat ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
AttrInfo vn -> PatBase f vn t -> SrcLoc -> PatBase f vn t
PatAttr AttrInfo VName
attr (Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv) SrcLoc
loc
updatePat (Id VName
vn (Info ParamType
tp) SrcLoc
loc) StaticVal
sv =
  VName -> Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. vn -> f t -> SrcLoc -> PatBase f vn t
Id VName
vn (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ ParamType -> ParamType -> ParamType
forall {dim} {o}.
TypeBase dim o -> TypeBase dim o -> TypeBase dim o
comb ParamType
tp (ParamType -> ParamType) -> ParamType -> ParamType
forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
loc
  where
    -- Preserve any original zeroth-order types.
    comb :: TypeBase dim o -> TypeBase dim o -> TypeBase dim o
comb (Scalar Arrow {}) TypeBase dim o
t2 = TypeBase dim o
t2
    comb (Scalar (Record Map Name (TypeBase dim o)
m1)) (Scalar (Record Map Name (TypeBase dim o)
m2)) =
      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) -> ScalarTypeBase dim o)
-> Map Name (TypeBase dim o) -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ (TypeBase dim o -> TypeBase dim o -> TypeBase dim o)
-> Map Name (TypeBase dim o)
-> Map Name (TypeBase dim o)
-> Map Name (TypeBase dim o)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase dim o -> TypeBase dim o -> TypeBase dim o
comb Map Name (TypeBase dim o)
m1 Map Name (TypeBase dim o)
m2
    comb (Scalar (Sum Map Name [TypeBase dim o]
m1)) (Scalar (Sum Map Name [TypeBase dim o]
m2)) =
      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
Sum (Map Name [TypeBase dim o] -> ScalarTypeBase dim o)
-> Map Name [TypeBase dim o] -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim o] -> [TypeBase dim o] -> [TypeBase dim o])
-> Map Name [TypeBase dim o]
-> Map Name [TypeBase dim o]
-> Map Name [TypeBase dim o]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeBase dim o -> TypeBase dim o -> TypeBase dim o)
-> [TypeBase dim o] -> [TypeBase dim o] -> [TypeBase dim o]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim o -> TypeBase dim o -> TypeBase dim o
comb) Map Name [TypeBase dim o]
m1 Map Name [TypeBase dim o]
m2
    comb TypeBase dim o
t1 TypeBase dim o
_ = TypeBase dim o
t1 -- t1 must be array or prim.
updatePat pat :: Pat ParamType
pat@(Wildcard (Info ParamType
tp) SrcLoc
loc) StaticVal
sv
  | ParamType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero ParamType
tp = Pat ParamType
pat
  | Bool
otherwise = Info ParamType -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t. f t -> SrcLoc -> PatBase f vn t
Wildcard (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ StaticVal -> ParamType
paramTypeFromSV StaticVal
sv) SrcLoc
loc
updatePat (PatAscription Pat ParamType
pat TypeExp Exp VName
_ SrcLoc
_) StaticVal
sv =
  Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat StaticVal
sv
updatePat p :: Pat ParamType
p@PatLit {} StaticVal
_ = Pat ParamType
p
updatePat pat :: Pat ParamType
pat@(PatConstr Name
c1 (Info ParamType
t) [Pat ParamType]
ps SrcLoc
loc) sv :: StaticVal
sv@(SumSV Name
_ [StaticVal]
svs [(Name, [ResType])]
_)
  | ParamType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero ParamType
t = Pat ParamType
pat
  | Bool
otherwise = Name
-> Info ParamType -> [Pat ParamType] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
c1 (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ResType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe ResType
t') [Pat ParamType]
ps' SrcLoc
loc
  where
    t' :: ResType
t' = StaticVal -> ResType
resTypeFromSV StaticVal
sv
    ps' :: [Pat ParamType]
ps' = (Pat ParamType -> StaticVal -> Pat ParamType)
-> [Pat ParamType] -> [StaticVal] -> [Pat ParamType]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Pat ParamType -> StaticVal -> Pat ParamType
updatePat [Pat ParamType]
ps [StaticVal]
svs
updatePat (PatConstr Name
c1 Info ParamType
_ [Pat ParamType]
ps SrcLoc
loc) (Dynamic ResType
t) =
  Name
-> Info ParamType -> [Pat ParamType] -> SrcLoc -> Pat ParamType
forall (f :: * -> *) vn t.
Name -> f t -> [PatBase f vn t] -> SrcLoc -> PatBase f vn t
PatConstr Name
c1 (ParamType -> Info ParamType
forall a. a -> Info a
Info (ParamType -> Info ParamType) -> ParamType -> Info ParamType
forall a b. (a -> b) -> a -> b
$ Diet -> ResType -> ParamType
forall o. Diet -> TypeBase Exp o -> ParamType
toParam Diet
Observe ResType
t) [Pat ParamType]
ps SrcLoc
loc
updatePat Pat ParamType
pat (Dynamic ResType
t) = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat (ResType -> StaticVal
svFromType ResType
t)
updatePat Pat ParamType
pat (HoleSV StructType
t SrcLoc
_) = Pat ParamType -> StaticVal -> Pat ParamType
updatePat Pat ParamType
pat (ResType -> StaticVal
svFromType (ResType -> StaticVal) -> ResType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Freshness -> StructType -> ResType
forall o. Freshness -> TypeBase Exp o -> ResType
toRes Freshness
Nonfresh StructType
t)
updatePat Pat ParamType
pat StaticVal
sv =
  [Char] -> Pat ParamType
forall a. HasCallStack => [Char] -> a
error ([Char] -> Pat ParamType) -> [Char] -> Pat ParamType
forall a b. (a -> b) -> a -> b
$
    [Char]
"Tried to update pattern\n"
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Pat ParamType -> [Char]
forall a. Pretty a => a -> [Char]
prettyString Pat ParamType
pat
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"\nto reflect the static value\n"
      [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> [Char]
forall a. Show a => a -> [Char]
show StaticVal
sv

-- | Convert a record (or tuple) type to a record static value. This
-- is used for "unwrapping" tuples and records that are nested in
-- 'Dynamic' static values.
svFromType :: ResType -> StaticVal
svFromType :: ResType -> StaticVal
svFromType (Scalar (Record Map Name ResType
fs)) = [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> (Map Name StaticVal -> [(Name, StaticVal)])
-> Map Name StaticVal
-> StaticVal
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name StaticVal -> [(Name, StaticVal)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name StaticVal -> StaticVal)
-> Map Name StaticVal -> StaticVal
forall a b. (a -> b) -> a -> b
$ (ResType -> StaticVal) -> Map Name ResType -> Map Name StaticVal
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ResType -> StaticVal
svFromType Map Name ResType
fs
svFromType ResType
t = ResType -> StaticVal
Dynamic ResType
t

-- | The static value of a (top-level) binding within the scope of its own
-- definition, used to handle recursion. We construct a 'DynamicFun' spine
-- matching the value parameters, ending in a 'Dynamic' for the result. This
-- suffices for fully applied recursive calls: the residual expression simply
-- refers to the top-level binding by name.
--
-- This is possible only when the parameters and the result are order-zero, as
-- the static value would otherwise depend on the body. See Note [Lifting and
-- recursion] for why that suffices.
selfSV :: VName -> [Pat ParamType] -> ResType -> Maybe StaticVal
selfSV :: VName -> [Pat ParamType] -> ResType -> Maybe StaticVal
selfSV VName
name [Pat ParamType]
params ResType
rettype
  | (Pat ParamType -> Bool) -> [Pat ParamType] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Pat ParamType -> Bool
forall d o. Pat (TypeBase d o) -> Bool
patternOrderZero [Pat ParamType]
params,
    ResType -> Bool
forall dim o. TypeBase dim o -> Bool
orderZero ResType
rettype =
      StaticVal -> Maybe StaticVal
forall a. a -> Maybe a
Just (StaticVal -> Maybe StaticVal) -> StaticVal -> Maybe StaticVal
forall a b. (a -> b) -> a -> b
$ [Pat ParamType] -> StaticVal
forall {a}. [a] -> StaticVal
go [Pat ParamType]
params
  | Bool
otherwise = Maybe StaticVal
forall a. Maybe a
Nothing
  where
    ret_sv :: StaticVal
ret_sv = ResType -> StaticVal
Dynamic ResType
rettype
    go :: [a] -> StaticVal
go [] = StaticVal
ret_sv
    go (a
_ : [a]
ps) =
      let inner :: StaticVal
inner = [a] -> StaticVal
go [a]
ps
          self :: Exp
self = 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
name) (StructType -> Info StructType
forall a. a -> Info a
Info (StaticVal -> StructType
structTypeFromSV StaticVal
inner)) SrcLoc
forall a. Monoid a => a
mempty
       in (Exp, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Exp
self, StaticVal
inner) StaticVal
inner

-- | The self static value binding for a top-level binding, if it is first-order
-- (see 'selfSV'). We register these for all bindings before defunctionalising
-- any of them, so that mutually recursive references resolve. The real static
-- value of each binding, added as it is processed, takes precedence over the
-- bindings produced here. This is what lets a lambda-lifted function refer to
-- the binding it was lifted out of; see Note [Lifting and recursion].
selfBinding :: ValBind -> Env
selfBinding :: ValBind -> Env
selfBinding ValBind
valbind =
  case VName -> [Pat ParamType] -> ResType -> Maybe StaticVal
selfSV (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) (ValBind -> [Pat ParamType]
forall (f :: * -> *) vn.
ValBindBase f vn -> [PatBase f vn ParamType]
valBindParams ValBind
valbind) ResType
rettype of
    Just StaticVal
self_sv ->
      VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton (ValBind -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBind
valbind) (Binding -> Env) -> Binding -> Env
forall a b. (a -> b) -> a -> b
$
        Maybe ([VName], StructType) -> StaticVal -> Binding
Binding (([VName], StructType) -> Maybe ([VName], StructType)
forall a. a -> Maybe a
Just (([TypeParamBase VName] -> [VName])
-> ([TypeParamBase VName], StructType) -> ([VName], StructType)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName) (ValBind -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBind
valbind))) StaticVal
self_sv
    Maybe StaticVal
Nothing -> Env
forall a. Monoid a => a
mempty
  where
    Info (RetType [VName]
_ ResType
rettype) = ValBind -> Info ResRetType
forall (f :: * -> *) vn. ValBindBase f vn -> f ResRetType
valBindRetType ValBind
valbind

-- | Defunctionalize a top-level value binding. Returns the
-- transformed result as well as an environment that binds the name of
-- the value binding to the static value of the transformed body.
defuncValBind :: ValBind -> DefM (ValBind, Env)
-- Eta-expand a functional result into further parameters, as the core language
-- is first-order. A recursive binding never has one (see Note [Lifting and
-- recursion]), so this cannot affect the self static values.
defuncValBind :: ValBind -> DefM (ValBind, Env)
defuncValBind (ValBind Maybe (Info EntryPoint)
entry VName
name SrcLoc
name_loc Maybe (TypeExp Exp VName)
_ (Info ResRetType
rettype) [TypeParamBase VName]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
attrs SrcLoc
loc)
  | Scalar Arrow {} <- ResRetType -> ResType
forall dim o. RetTypeBase dim o -> TypeBase dim o
retType ResRetType
rettype = do
      (body_pats, body', rettype') <- ResRetType -> Exp -> DefM ([Pat ParamType], Exp, ResRetType)
etaExpand ((Freshness -> Freshness) -> ResRetType -> ResRetType
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 (Freshness -> Freshness -> Freshness
forall a b. a -> b -> a
const Freshness
forall a. Monoid a => a
mempty) ResRetType
rettype) Exp
body
      defuncValBind $
        ValBind
          entry
          name
          name_loc
          Nothing
          (Info rettype')
          tparams
          (params <> body_pats)
          body'
          Nothing
          attrs
          loc
defuncValBind valbind :: ValBind
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
name SrcLoc
_ Maybe (TypeExp Exp VName)
retdecl (Info (RetType [VName]
ret_dims ResType
rettype)) [TypeParamBase VName]
tparams [Pat ParamType]
params Exp
body Maybe DocComment
_ [AttrInfo VName]
_ SrcLoc
_) = do
  Bool -> DefM () -> DefM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ((TypeParamBase VName -> Bool) -> [TypeParamBase VName] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any TypeParamBase VName -> Bool
forall vn. TypeParamBase vn -> Bool
isTypeParam [TypeParamBase VName]
tparams) (DefM () -> DefM ()) -> DefM () -> DefM ()
forall a b. (a -> b) -> a -> b
$
    [Char] -> DefM ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> DefM ()) -> [Char] -> DefM ()
forall a b. (a -> b) -> a -> b
$
      VName -> [Char]
forall a. Show a => a -> [Char]
show VName
name
        [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" has type parameters, "
        [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"but the defunctionaliser expects a monomorphic input program."
  -- Bind the name in the scope of its own definition, so that first-order
  -- recursive references can be defunctionalised via a static value; see
  -- 'selfSV'. Higher-order recursive functions get no self static value: their
  -- body is not defunctionalised here (it is stored raw in a 'LambdaSV' and
  -- defunctionalised when applied), so recursive references resolve against the
  -- global scope, and the recursion is tied off by memoisation in
  -- 'defuncApplyArg' (see Note [Lifting and recursion]).
  let self :: Env -> Env
self = (ValBind -> Env
selfBinding ValBind
valbind Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>)
  -- The self static value goes into the *global* environment (as well as the
  -- local one) so that it is treated as a top-level function: recursive uses of
  -- it as a value are then eta-expanded and closure-converted like any other
  -- top-level function (see the 'DynamicFun' case of 'defuncExp'), rather than
  -- capturing the self static value into a closure, which 'restrictEnvTo' omits
  -- for globals.
  (tparams', params', body', sv, sv_t) <-
    ((Env, Env) -> (Env, Env))
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> (Env -> Env) -> (Env, Env) -> (Env, Env)
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 Env -> Env
self Env -> Env
self) (DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
 -> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType))
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
forall a b. (a -> b) -> a -> b
$
      [VName]
-> [Pat ParamType]
-> Exp
-> ResRetType
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
defuncLet ((TypeParamBase VName -> VName) -> [TypeParamBase VName] -> [VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeParamBase VName -> VName
forall vn. TypeParamBase vn -> vn
typeParamName [TypeParamBase VName]
tparams) [Pat ParamType]
params Exp
body (ResRetType
 -> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType))
-> ResRetType
-> DefM ([VName], [Pat ParamType], Exp, StaticVal, ResType)
forall a b. (a -> b) -> a -> b
$
        [VName] -> ResType -> ResRetType
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [VName]
ret_dims ResType
rettype
  globals <- asks $ M.keysSet . fst
  let bound_sizes = [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList ((Pat ParamType -> [VName]) -> [Pat ParamType] -> [VName]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Pat ParamType -> [VName]
forall t. Pat t -> [VName]
patNames [Pat ParamType]
params') Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams' Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> Set VName
globals
  params'' <- instAnySizes params'
  let rettype' = ResType -> ResType -> ResType
forall o.
Monoid o =>
TypeBase Exp o -> TypeBase Exp o -> TypeBase Exp o
combineTypeShapes ResType
rettype ResType
sv_t
      tparams'' = [VName]
tparams' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ Set VName -> [Pat ParamType] -> [VName]
unboundSizes Set VName
bound_sizes [Pat ParamType]
params''
      ret_dims' = (VName -> Bool) -> [VName] -> [VName]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Set VName -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Set VName
bound_sizes) ([VName] -> [VName]) -> [VName] -> [VName]
forall a b. (a -> b) -> a -> b
$ Set VName -> [VName]
forall a. Set a -> [a]
S.toList (Set VName -> [VName]) -> Set VName -> [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
$ ResType -> FV
forall o. TypeBase Exp o -> FV
freeInType ResType
rettype'

  pure
    ( valbind
        { valBindRetDecl = retdecl,
          valBindRetType =
            Info $
              if null params'
                then RetType ret_dims' $ rettype' `setMode` Nonfresh
                else RetType ret_dims' rettype',
          valBindTypeParams = map (`TypeParamDim` mempty) tparams'',
          valBindParams = params'',
          valBindBody = body'
        },
      M.singleton name $
        Binding
          (Just (first (map typeParamName) (valBindTypeScheme valbind)))
          sv
    )

-- | Defunctionalize a list of top-level declarations.
defuncVals :: [ValBind] -> DefM ()
defuncVals :: [ValBind] -> DefM ()
defuncVals [ValBind]
binds =
  ((Env, Env) -> (Env, Env)) -> DefM () -> DefM ()
forall a. ((Env, Env) -> (Env, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Env -> Env) -> (Env -> Env) -> (Env, Env) -> (Env, Env)
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 (Env
self_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>) (Env
self_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>)) (DefM () -> DefM ()) -> DefM () -> DefM ()
forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
go [ValBind]
binds
  where
    self_env :: Env
self_env = (ValBind -> Env) -> [ValBind] -> Env
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ValBind -> Env
selfBinding [ValBind]
binds
    go :: [ValBind] -> DefM ()
go [] = () -> DefM ()
forall a. a -> DefM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    go (ValBind
valbind : [ValBind]
ds) = do
      (valbind', env) <- ValBind -> DefM (ValBind, Env)
defuncValBind ValBind
valbind
      addValBind valbind'
      local (bimap (env <>) (env <>)) $ go ds

{-# NOINLINE transformProg #-}

-- | Transform a list of top-level value bindings. May produce new
-- lifted function definitions, which are placed in front of the
-- resulting list of declarations.
transformProg :: (MonadFreshNames m) => [ValBind] -> m [ValBind]
transformProg :: forall (m :: * -> *). MonadFreshNames m => [ValBind] -> m [ValBind]
transformProg [ValBind]
decs = (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([ValBind], VNameSource)) -> m [ValBind])
-> (VNameSource -> ([ValBind], VNameSource)) -> m [ValBind]
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
  let ((), VNameSource
namesrc', [ValBind]
decs') = VNameSource -> DefM () -> ((), VNameSource, [ValBind])
forall a. VNameSource -> DefM a -> (a, VNameSource, [ValBind])
runDefM VNameSource
namesrc (DefM () -> ((), VNameSource, [ValBind]))
-> DefM () -> ((), VNameSource, [ValBind])
forall a b. (a -> b) -> a -> b
$ [ValBind] -> DefM ()
defuncVals [ValBind]
decs
   in ([ValBind]
decs', VNameSource
namesrc')

-- Note [Lifting and recursion]
--
-- Defunctionalisation works by static interpretation: applying a higher-order
-- function ('defuncApplyArg' on a 'LambdaSV') is specialised by inlining the
-- function body with its static (function-valued) argument substituted in, and
-- lifting the result to a fresh top-level function. Recursion is not a special
-- case of this; it falls out of memoising the lifting.
--
-- The specialisation a given step produces is fully determined by the function
-- (the 'LambdaSV', which carries the body, the parameter, and the lexical
-- closure) together with the 'StaticVal' of the argument it is applied to.
-- Everything that varies at runtime (the captured values) is passed in the
-- closure, not baked into the lifted function. So we memoise ('LiftMemo') the
-- lifted function under exactly that pair, and reuse it whenever the same
-- function is applied to the same static argument again. This deduplicates
-- lifted functions, and it terminates recursion for free: a recursive
-- occurrence re-applies the same function to a statically identical argument
-- (order-zero arguments have the same 'StaticVal' regardless of their runtime
-- value, and an unchanged higher-order argument resolves to the same
-- 'StaticVal'), so it hits the memo and becomes a call to the function being
-- lifted rather than another round of inlining.
--
-- There is no distinction between recursive and non-recursive functions: an
-- ordinary call and a recursive call are the same memo lookup. Note that this
-- means a program whose recursion is not statically resolvable - an indirect
-- recursive call, or one that passes a *different* function at each step (e.g.
-- `go (\x -> g x) n = ... go (\x -> g (g x)) (n-1)`) - produces ever-changing
-- keys and does not terminate here. The type checker rejects such programs, by
-- requiring a recursive application to be saturated and to pass its
-- higher-order arguments unchanged.
--
-- One subtlety remains, and it is about knot-tying, not about recognising
-- recursion. A curried application peels off one parameter per step; the step
-- whose result is order-zero yields the actual first-order body, and its body
-- may contain the recursive occurrence. That entry must therefore be recorded
-- _before_ the body is defunctionalised, using the declared (order-zero) return
-- type and a 'Dynamic' static value, both known up front; it is replaced with
-- the refined return type and real static value once the body is lifted. The
-- earlier (function-typed) steps are recorded *after* lifting, since their
-- residual 'StaticVal' - which a memo hit must return so application can
-- continue - is not known until then. This is sound because those earlier steps
-- are all processed before the order-zero step (the body is innermost), so they
-- are in the memo by the time the body's recursive occurrence looks for them.
--
-- Another problem is scoping. Recursion is resolved by name, and 'defuncVals'
-- processes bindings in order, so a body being defunctionalised can normally
-- see the static value of everything it refers to. Lambda lifting is the only
-- thing that breaks this as it can introduce mutual recursion (this is not
-- allowed in the source language).
--
-- Whether this is a problem depends on the lifted function:
--
-- - A higher-order function is not defunctionalised where it is defined - its
--   body is stored in a 'LambdaSV' and defunctionalised when it is applied,
--   which is necessarily from a later binding. By then the binding it refers to
--   has been processed, and nothing more is needed.
--
-- - A first-order function is defunctionalised eagerly, and so does need the
--   static value of a binding that has not been processed yet. 'selfBinding'
--   provides it for all bindings up front, which 'selfSV' can derive from the
--   signature alone when the parameters and the result are order-zero.
--
-- The type rules ensure the result of a recursive binding is always order-zero,
-- which is crucial for making this work.
--
-- A binding with a higher-order parameter thus gets no self static value, and
-- needs none: the type checker also requires a recursive application to pass
-- such a parameter unchanged, so a lifted function containing the recursive
-- occurrence must capture the function-typed parameter, which makes it
-- higher-order itself - the first case above.