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

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

import qualified Control.Arrow as Arrow
import Control.Monad.Identity
import Control.Monad.RWS hiding (Sum)
import Control.Monad.State
import Data.Bifunctor
import Data.Bitraversable
import Data.Foldable
import Data.List (partition, sortOn, tails)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Sequence as Seq
import qualified Data.Set as S
import Futhark.IR.Pretty ()
import qualified Futhark.Internalise.FreeVars as FV
import Futhark.MonadFreshNames
import Language.Futhark
import Language.Futhark.Traversals

-- | An expression or an extended 'Lambda' (with size parameters,
-- which AST lambdas do not support).
data ExtExp
  = ExtLambda [Pattern] Exp StructType SrcLoc
  | ExtExp Exp
  deriving (Int -> ExtExp -> ShowS
[ExtExp] -> ShowS
ExtExp -> String
(Int -> ExtExp -> ShowS)
-> (ExtExp -> String) -> ([ExtExp] -> ShowS) -> Show ExtExp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ExtExp] -> ShowS
$cshowList :: [ExtExp] -> ShowS
show :: ExtExp -> String
$cshow :: ExtExp -> String
showsPrec :: Int -> ExtExp -> ShowS
$cshowsPrec :: Int -> ExtExp -> ShowS
Show)

-- | A static value stores additional information about the result of
-- defunctionalization of an expression, aside from the residual expression.
data StaticVal
  = Dynamic PatternType
  | LambdaSV Pattern StructType ExtExp Env
  | RecordSV [(Name, StaticVal)]
  | -- | The constructor that is actually present, plus
    -- the others that are not.
    SumSV Name [StaticVal] [(Name, [PatternType])]
  | -- | 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
  deriving (Int -> StaticVal -> ShowS
[StaticVal] -> ShowS
StaticVal -> String
(Int -> StaticVal -> ShowS)
-> (StaticVal -> String)
-> ([StaticVal] -> ShowS)
-> Show StaticVal
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [StaticVal] -> ShowS
$cshowList :: [StaticVal] -> ShowS
show :: StaticVal -> String
$cshow :: StaticVal -> String
showsPrec :: Int -> StaticVal -> ShowS
$cshowsPrec :: Int -> StaticVal -> ShowS
Show)

-- | The type is Just if this is a polymorphic binding that must be
-- instantiated.
data Binding = Binding (Maybe ([VName], StructType)) StaticVal
  deriving (Int -> Binding -> ShowS
[Binding] -> ShowS
Binding -> String
(Int -> Binding -> ShowS)
-> (Binding -> String) -> ([Binding] -> ShowS) -> Show Binding
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Binding] -> ShowS
$cshowList :: [Binding] -> ShowS
show :: Binding -> String
$cshow :: Binding -> String
showsPrec :: Int -> Binding -> ShowS
$cshowsPrec :: Int -> Binding -> ShowS
Show)

bindingSV :: Binding -> StaticVal
bindingSV :: Binding -> StaticVal
bindingSV (Binding Maybe ([VName], StructType)
_ StaticVal
sv) = StaticVal
sv

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

localEnv :: Env -> DefM a -> DefM a
localEnv :: forall a. Env -> DefM a -> DefM a
localEnv Env
env = ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a)
-> ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall a b. (a -> b) -> a -> b
$ (Env -> Env) -> (Set VName, Env) -> (Set VName, Env)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
Arrow.second (Env
env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<>)

-- 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 = ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a)
-> ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall a b. (a -> b) -> a -> b
$ \(Set VName
globals, Env
old_env) ->
  (Set VName
globals, (VName -> Binding -> Bool) -> Env -> Env
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\VName
k Binding
_ -> VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals) Env
old_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env)

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

isGlobal :: VName -> DefM a -> DefM a
isGlobal :: forall a. VName -> DefM a -> DefM a
isGlobal VName
v = ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a)
-> ((Set VName, Env) -> (Set VName, Env)) -> DefM a -> DefM a
forall a b. (a -> b) -> a -> b
$ (Set VName -> Set VName) -> (Set VName, Env) -> (Set VName, Env)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
Arrow.first (VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
v)

replaceTypeSizes ::
  M.Map VName SizeSubst ->
  TypeBase (DimDecl VName) als ->
  TypeBase (DimDecl VName) als
replaceTypeSizes :: forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs = (DimDecl VName -> DimDecl VName)
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
onDim
  where
    onDim :: DimDecl VName -> DimDecl VName
onDim (NamedDim QualName VName
v) =
      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 -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v'
        Just (SubstConst Int
d) -> Int -> DimDecl VName
forall vn. Int -> DimDecl vn
ConstDim Int
d
        Maybe SizeSubst
Nothing -> QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v
    onDim DimDecl VName
d = DimDecl VName
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 PatternBase Info VName
param StructType
t ExtExp
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 (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 PatternBase Info VName -> StructType -> ExtExp -> Env -> StaticVal
LambdaSV
            (Map VName SizeSubst
-> PatternBase Info VName -> PatternBase Info VName
forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs PatternBase Info VName
param)
            (Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs StructType
t)
            (Map VName SizeSubst -> ExtExp -> ExtExp
onExtExp Map VName SizeSubst
substs ExtExp
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 PatternType
t ->
      PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> PatternType -> PatternType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
orig_substs PatternType
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 (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, [PatternType])]
ts ->
      Name -> [StaticVal] -> [(Name, [PatternType])] -> 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, [PatternType])] -> StaticVal)
-> [(Name, [PatternType])] -> StaticVal
forall a b. (a -> b) -> a -> b
$
        ((Name, [PatternType]) -> (Name, [PatternType]))
-> [(Name, [PatternType])] -> [(Name, [PatternType])]
forall a b. (a -> b) -> [a] -> [b]
map (([PatternType] -> [PatternType])
-> (Name, [PatternType]) -> (Name, [PatternType])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([PatternType] -> [PatternType])
 -> (Name, [PatternType]) -> (Name, [PatternType]))
-> ([PatternType] -> [PatternType])
-> (Name, [PatternType])
-> (Name, [PatternType])
forall a b. (a -> b) -> a -> b
$ (PatternType -> PatternType) -> [PatternType] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map ((PatternType -> PatternType) -> [PatternType] -> [PatternType])
-> (PatternType -> PatternType) -> [PatternType] -> [PatternType]
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> PatternType -> PatternType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
orig_substs) [(Name, [PatternType])]
ts
    DynamicFun (ExpBase Info VName
e, StaticVal
sv1) StaticVal
sv2 ->
      (ExpBase Info VName, StaticVal) -> StaticVal -> StaticVal
DynamicFun (Map VName SizeSubst -> ExpBase Info VName -> ExpBase Info VName
forall {f :: * -> *}.
(ASTMappable (ExpBase f VName), Functor f) =>
Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
orig_substs ExpBase Info VName
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
  where
    tv :: Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs =
      ASTMapper m
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper
        { mapOnPatternType :: PatternType -> m PatternType
mapOnPatternType = PatternType -> m PatternType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PatternType -> m PatternType)
-> (PatternType -> PatternType) -> PatternType -> m PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> PatternType -> PatternType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs,
          mapOnStructType :: StructType -> m StructType
mapOnStructType = StructType -> m StructType
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StructType -> m StructType)
-> (StructType -> StructType) -> StructType -> m StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs,
          mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName)
mapOnExp = ExpBase Info VName -> m (ExpBase Info VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpBase Info VName -> m (ExpBase Info VName))
-> (ExpBase Info VName -> ExpBase Info VName)
-> ExpBase Info VName
-> m (ExpBase Info VName)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map VName SizeSubst -> ExpBase Info VName -> ExpBase Info VName
forall {f :: * -> *}.
(ASTMappable (ExpBase f VName), Functor f) =>
Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
substs,
          mapOnName :: VName -> m VName
mapOnName = VName -> m VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure (VName -> m VName) -> (VName -> VName) -> VName -> m VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
substs (Var QualName VName
v f PatternType
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 -> f PatternType -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
v' f PatternType
t SrcLoc
loc
        Just (SubstConst Int
d) ->
          PrimValue -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn. PrimValue -> SrcLoc -> ExpBase f vn
Literal (IntValue -> PrimValue
SignedValue (Int64 -> IntValue
Int64Value (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d))) SrcLoc
loc
        Maybe SizeSubst
Nothing ->
          QualName VName -> f PatternType -> SrcLoc -> ExpBase f VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
v (Map VName SizeSubst -> PatternType -> PatternType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs (PatternType -> PatternType) -> f PatternType -> f PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f PatternType
t) SrcLoc
loc
    onExp Map VName SizeSubst
substs (Coerce ExpBase f VName
e TypeDeclBase f VName
tdecl (f PatternType, f [VName])
t SrcLoc
loc) =
      ExpBase f VName
-> TypeDeclBase f VName
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f VName
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeDeclBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Coerce (Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
substs ExpBase f VName
e) TypeDeclBase f VName
tdecl' ((f PatternType -> f PatternType)
-> (f PatternType, f [VName]) -> (f PatternType, f [VName])
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((PatternType -> PatternType) -> f PatternType -> f PatternType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map VName SizeSubst -> PatternType -> PatternType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs)) (f PatternType, f [VName])
t) SrcLoc
loc
      where
        tdecl' :: TypeDeclBase f VName
tdecl' =
          TypeDecl :: forall (f :: * -> *) vn.
TypeExp vn -> f StructType -> TypeDeclBase f vn
TypeDecl
            { declaredType :: TypeExp VName
declaredType = Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs (TypeExp VName -> TypeExp VName) -> TypeExp VName -> TypeExp VName
forall a b. (a -> b) -> a -> b
$ TypeDeclBase f VName -> TypeExp VName
forall (f :: * -> *) vn. TypeDeclBase f vn -> TypeExp vn
declaredType TypeDeclBase f VName
tdecl,
              expandedType :: f StructType
expandedType = Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs (StructType -> StructType) -> f StructType -> f StructType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeDeclBase f VName -> f StructType
forall (f :: * -> *) vn. TypeDeclBase f vn -> f StructType
expandedType TypeDeclBase f VName
tdecl
            }
    onExp Map VName SizeSubst
substs ExpBase f VName
e = Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs ExpBase f VName
e

    onTypeExpDim :: Map VName SizeSubst -> DimExp VName -> DimExp VName
onTypeExpDim Map VName SizeSubst
substs d :: DimExp VName
d@(DimExpNamed QualName VName
v 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 -> SrcLoc -> DimExp VName
forall vn. QualName vn -> SrcLoc -> DimExp vn
DimExpNamed QualName VName
v' SrcLoc
loc
        Just (SubstConst Int
x) ->
          Int -> SrcLoc -> DimExp VName
forall vn. Int -> SrcLoc -> DimExp vn
DimExpConst Int
x SrcLoc
loc
        Maybe SizeSubst
Nothing ->
          DimExp VName
d
    onTypeExpDim Map VName SizeSubst
_ DimExp VName
d = DimExp VName
d

    onTypeArgExp :: Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs (TypeArgExpDim DimExp VName
d SrcLoc
loc) =
      DimExp VName -> SrcLoc -> TypeArgExp VName
forall vn. DimExp vn -> SrcLoc -> TypeArgExp vn
TypeArgExpDim (Map VName SizeSubst -> DimExp VName -> DimExp VName
onTypeExpDim Map VName SizeSubst
substs DimExp VName
d) SrcLoc
loc
    onTypeArgExp Map VName SizeSubst
substs (TypeArgExpType TypeExp VName
te) =
      TypeExp VName -> TypeArgExp VName
forall vn. TypeExp vn -> TypeArgExp vn
TypeArgExpType (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te)

    onTypeExp :: Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs (TEArray TypeExp VName
te DimExp VName
d SrcLoc
loc) =
      TypeExp VName -> DimExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> DimExp vn -> SrcLoc -> TypeExp vn
TEArray (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
te) (Map VName SizeSubst -> DimExp VName -> DimExp VName
onTypeExpDim Map VName SizeSubst
substs DimExp VName
d) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TEUnique TypeExp VName
t SrcLoc
loc) =
      TypeExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> SrcLoc -> TypeExp vn
TEUnique (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TEApply TypeExp VName
t1 TypeArgExp VName
t2 SrcLoc
loc) =
      TypeExp VName -> TypeArgExp VName -> SrcLoc -> TypeExp VName
forall vn. TypeExp vn -> TypeArgExp vn -> SrcLoc -> TypeExp vn
TEApply (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeArgExp VName -> TypeArgExp VName
onTypeArgExp Map VName SizeSubst
substs TypeArgExp VName
t2) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TEArrow Maybe VName
p TypeExp VName
t1 TypeExp VName
t2 SrcLoc
loc) =
      Maybe VName
-> TypeExp VName -> TypeExp VName -> SrcLoc -> TypeExp VName
forall vn.
Maybe vn -> TypeExp vn -> TypeExp vn -> SrcLoc -> TypeExp vn
TEArrow Maybe VName
p (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t1) (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs TypeExp VName
t2) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TETuple [TypeExp VName]
ts SrcLoc
loc) =
      [TypeExp VName] -> SrcLoc -> TypeExp VName
forall vn. [TypeExp vn] -> SrcLoc -> TypeExp vn
TETuple ((TypeExp VName -> TypeExp VName)
-> [TypeExp VName] -> [TypeExp VName]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [TypeExp VName]
ts) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TERecord [(Name, TypeExp VName)]
ts SrcLoc
loc) =
      [(Name, TypeExp VName)] -> SrcLoc -> TypeExp VName
forall vn. [(Name, TypeExp vn)] -> SrcLoc -> TypeExp vn
TERecord (((Name, TypeExp VName) -> (Name, TypeExp VName))
-> [(Name, TypeExp VName)] -> [(Name, TypeExp VName)]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeExp VName -> TypeExp VName)
-> (Name, TypeExp VName) -> (Name, TypeExp VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TypeExp VName -> TypeExp VName)
 -> (Name, TypeExp VName) -> (Name, TypeExp VName))
-> (TypeExp VName -> TypeExp VName)
-> (Name, TypeExp VName)
-> (Name, TypeExp VName)
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, TypeExp VName)]
ts) SrcLoc
loc
    onTypeExp Map VName SizeSubst
substs (TESum [(Name, [TypeExp VName])]
ts SrcLoc
loc) =
      [(Name, [TypeExp VName])] -> SrcLoc -> TypeExp VName
forall vn. [(Name, [TypeExp vn])] -> SrcLoc -> TypeExp vn
TESum (((Name, [TypeExp VName]) -> (Name, [TypeExp VName]))
-> [(Name, [TypeExp VName])] -> [(Name, [TypeExp VName])]
forall a b. (a -> b) -> [a] -> [b]
map (([TypeExp VName] -> [TypeExp VName])
-> (Name, [TypeExp VName]) -> (Name, [TypeExp VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([TypeExp VName] -> [TypeExp VName])
 -> (Name, [TypeExp VName]) -> (Name, [TypeExp VName]))
-> ([TypeExp VName] -> [TypeExp VName])
-> (Name, [TypeExp VName])
-> (Name, [TypeExp VName])
forall a b. (a -> b) -> a -> b
$ (TypeExp VName -> TypeExp VName)
-> [TypeExp VName] -> [TypeExp VName]
forall a b. (a -> b) -> [a] -> [b]
map ((TypeExp VName -> TypeExp VName)
 -> [TypeExp VName] -> [TypeExp VName])
-> (TypeExp VName -> TypeExp VName)
-> [TypeExp VName]
-> [TypeExp VName]
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> TypeExp VName -> TypeExp VName
onTypeExp Map VName SizeSubst
substs) [(Name, [TypeExp VName])]
ts) SrcLoc
loc
    onTypeExp Map VName SizeSubst
_ (TEVar QualName VName
v SrcLoc
loc) =
      QualName VName -> SrcLoc -> TypeExp VName
forall vn. QualName vn -> SrcLoc -> TypeExp vn
TEVar QualName VName
v SrcLoc
loc

    onExtExp :: Map VName SizeSubst -> ExtExp -> ExtExp
onExtExp Map VName SizeSubst
substs (ExtExp ExpBase Info VName
e) =
      ExpBase Info VName -> ExtExp
ExtExp (ExpBase Info VName -> ExtExp) -> ExpBase Info VName -> ExtExp
forall a b. (a -> b) -> a -> b
$ Map VName SizeSubst -> ExpBase Info VName -> ExpBase Info VName
forall {f :: * -> *}.
(ASTMappable (ExpBase f VName), Functor f) =>
Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
substs ExpBase Info VName
e
    onExtExp Map VName SizeSubst
substs (ExtLambda [PatternBase Info VName]
params ExpBase Info VName
e StructType
t SrcLoc
loc) =
      [PatternBase Info VName]
-> ExpBase Info VName -> StructType -> SrcLoc -> ExtExp
ExtLambda ((PatternBase Info VName -> PatternBase Info VName)
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a b. (a -> b) -> [a] -> [b]
map (Map VName SizeSubst
-> PatternBase Info VName -> PatternBase Info VName
forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs) [PatternBase Info VName]
params) (Map VName SizeSubst -> ExpBase Info VName -> ExpBase Info VName
forall {f :: * -> *}.
(ASTMappable (ExpBase f VName), Functor f) =>
Map VName SizeSubst -> ExpBase f VName -> ExpBase f VName
onExp Map VName SizeSubst
substs ExpBase Info VName
e) (Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
substs StructType
t) SrcLoc
loc

    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 (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 (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
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)

    onAST :: ASTMappable x => M.Map VName SizeSubst -> x -> x
    onAST :: forall x. ASTMappable x => Map VName SizeSubst -> x -> x
onAST Map VName SizeSubst
substs = Identity x -> x
forall a. Identity a -> a
runIdentity (Identity x -> x) -> (x -> Identity x) -> x -> x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASTMapper Identity -> x -> Identity x
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap (Map VName SizeSubst -> ASTMapper Identity
forall {m :: * -> *}. Monad m => Map VName SizeSubst -> ASTMapper m
tv Map VName SizeSubst
substs)

-- | Returns the defunctionalization environment restricted
-- to the given set of variable names and types.
restrictEnvTo :: FV.NameSet -> DefM Env
restrictEnvTo :: NameSet -> DefM Env
restrictEnvTo (FV.NameSet Map VName StructType
m) = (Set VName, Env) -> Env
restrict ((Set VName, Env) -> Env) -> DefM (Set VName, Env) -> DefM Env
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DefM (Set VName, Env)
forall r (m :: * -> *). MonadReader r m => m r
ask
  where
    restrict :: (Set VName, Env) -> Env
restrict (Set VName
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 (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ VName
k VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
globals
          Uniqueness
u <- StructType -> Uniqueness
forall shape as. TypeBase shape as -> Uniqueness
uniqueness (StructType -> Uniqueness) -> Maybe StructType -> Maybe Uniqueness
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName StructType -> Maybe StructType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
k Map VName StructType
m
          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
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv
    restrict' :: Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
Nonunique (Dynamic PatternType
t) =
      PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ PatternType
t PatternType -> Uniqueness -> PatternType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
    restrict' Uniqueness
_ (Dynamic PatternType
t) =
      PatternType -> StaticVal
Dynamic PatternType
t
    restrict' Uniqueness
u (LambdaSV PatternBase Info VName
pat StructType
t ExtExp
e Env
env) =
      PatternBase Info VName -> StructType -> ExtExp -> Env -> StaticVal
LambdaSV PatternBase Info VName
pat StructType
t ExtExp
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 (Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u) Env
env
    restrict' Uniqueness
u (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 (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((StaticVal -> StaticVal)
 -> (Name, StaticVal) -> (Name, StaticVal))
-> (StaticVal -> StaticVal)
-> (Name, StaticVal)
-> (Name, StaticVal)
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [(Name, StaticVal)]
fields
    restrict' Uniqueness
u (SumSV Name
c [StaticVal]
svs [(Name, [PatternType])]
fields) =
      Name -> [StaticVal] -> [(Name, [PatternType])] -> StaticVal
SumSV Name
c ((StaticVal -> StaticVal) -> [StaticVal] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map (Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u) [StaticVal]
svs) [(Name, [PatternType])]
fields
    restrict' Uniqueness
u (DynamicFun (ExpBase Info VName
e, StaticVal
sv1) StaticVal
sv2) =
      (ExpBase Info VName, StaticVal) -> StaticVal -> StaticVal
DynamicFun (ExpBase Info VName
e, Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv1) (StaticVal -> StaticVal) -> StaticVal -> StaticVal
forall a b. (a -> b) -> a -> b
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv2
    restrict' Uniqueness
_ StaticVal
IntrinsicSV = StaticVal
IntrinsicSV
    restrict'' :: Uniqueness -> Binding -> Binding
restrict'' Uniqueness
u (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
$ Uniqueness -> StaticVal -> StaticVal
restrict' Uniqueness
u StaticVal
sv

-- | Defunctionalization monad.  The Reader environment tracks both
-- the current Env as well as the set of globally defined dynamic
-- functions.  This is used to avoid unnecessarily large closure
-- environments.
newtype DefM a = DefM (RWS (S.Set VName, Env) (Seq.Seq ValBind) VNameSource 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
<$ :: forall a b. a -> DefM b -> DefM a
$c<$ :: forall a b. a -> DefM b -> DefM a
fmap :: forall a b. (a -> b) -> DefM a -> DefM b
$cfmap :: forall a b. (a -> b) -> DefM a -> DefM b
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
<* :: forall a b. DefM a -> DefM b -> DefM a
$c<* :: forall a b. DefM a -> DefM b -> DefM a
*> :: forall a b. DefM a -> DefM b -> DefM b
$c*> :: forall a b. DefM a -> DefM b -> DefM b
liftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
$cliftA2 :: forall a b c. (a -> b -> c) -> DefM a -> DefM b -> DefM c
<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
$c<*> :: forall a b. DefM (a -> b) -> DefM a -> DefM b
pure :: forall a. a -> DefM a
$cpure :: forall a. a -> 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
return :: forall a. a -> DefM a
$creturn :: forall a. a -> DefM a
>> :: forall a b. DefM a -> DefM b -> DefM b
$c>> :: forall a b. DefM a -> DefM b -> DefM b
>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
$c>>= :: forall a b. DefM a -> (a -> DefM b) -> DefM b
Monad,
      MonadReader (S.Set VName, Env),
      MonadWriter (Seq.Seq ValBind),
      Monad DefM
Applicative DefM
DefM VNameSource
Applicative DefM
-> Monad DefM
-> DefM VNameSource
-> (VNameSource -> DefM ())
-> MonadFreshNames DefM
VNameSource -> DefM ()
forall (m :: * -> *).
Applicative m
-> Monad m
-> m VNameSource
-> (VNameSource -> m ())
-> MonadFreshNames m
putNameSource :: VNameSource -> DefM ()
$cputNameSource :: VNameSource -> DefM ()
getNameSource :: DefM VNameSource
$cgetNameSource :: DefM VNameSource
MonadFreshNames
    )

-- | 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, Seq.Seq ValBind)
runDefM :: forall a.
VNameSource
-> DefM a -> (a, VNameSource, Seq (ValBindBase Info VName))
runDefM VNameSource
src (DefM RWS (Set VName, Env) (Seq (ValBindBase Info VName)) VNameSource a
m) = RWS (Set VName, Env) (Seq (ValBindBase Info VName)) VNameSource a
-> (Set VName, Env)
-> VNameSource
-> (a, VNameSource, Seq (ValBindBase Info VName))
forall r w s a. RWS r w s a -> r -> s -> (a, s, w)
runRWS RWS (Set VName, Env) (Seq (ValBindBase Info VName)) VNameSource a
m (Set VName, Env)
forall a. Monoid a => a
mempty VNameSource
src

collectFuns :: DefM a -> DefM (a, Seq.Seq ValBind)
collectFuns :: forall a. DefM a -> DefM (a, Seq (ValBindBase Info VName))
collectFuns DefM a
m = DefM
  ((a, Seq (ValBindBase Info VName)),
   Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName))
-> DefM (a, Seq (ValBindBase Info VName))
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (DefM
   ((a, Seq (ValBindBase Info VName)),
    Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName))
 -> DefM (a, Seq (ValBindBase Info VName)))
-> DefM
     ((a, Seq (ValBindBase Info VName)),
      Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName))
-> DefM (a, Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$ do
  (a
x, Seq (ValBindBase Info VName)
decs) <- DefM a -> DefM (a, Seq (ValBindBase Info VName))
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen DefM a
m
  ((a, Seq (ValBindBase Info VName)),
 Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName))
-> DefM
     ((a, Seq (ValBindBase Info VName)),
      Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName))
forall (m :: * -> *) a. Monad m => a -> m a
return ((a
x, Seq (ValBindBase Info VName)
decs), Seq (ValBindBase Info VName)
-> Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName)
forall a b. a -> b -> a
const Seq (ValBindBase Info VName)
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
env <- DefM Env
askEnv
  case VName -> Env -> Maybe Binding
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
x Env
env of
    Just (Binding (Just ([VName]
dims, StructType
sv_t)) StaticVal
sv) -> do
      Set VName
globals <- ((Set VName, Env) -> Set VName) -> DefM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Set VName
forall a b. (a, b) -> a
fst
      Set VName
-> [VName]
-> StructType
-> StructType
-> StaticVal
-> DefM StaticVal
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
    Just (Binding Maybe ([VName], StructType)
Nothing StaticVal
sv) ->
      StaticVal -> DefM StaticVal
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 -> Int
baseTag VName
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag -> StaticVal -> DefM StaticVal
forall (m :: * -> *) a. Monad m => a -> m a
return StaticVal
IntrinsicSV
      | Bool
otherwise ->
        -- Anything not in scope is going to be an existential size.
        StaticVal -> DefM StaticVal
forall (m :: * -> *) a. Monad m => a -> m a
return (StaticVal -> DefM StaticVal) -> StaticVal -> DefM StaticVal
forall a b. (a -> b) -> a -> b
$ PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- Like patternDimNames, 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 (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 (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((StructType -> Set VName) -> [StructType] -> Set VName
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 ()
_ Uniqueness
_ TypeName
_ [TypeArg (DimDecl VName)]
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 (DimDecl VName) -> Set VName)
-> [TypeArg (DimDecl VName)] -> [Set VName]
forall a b. (a -> b) -> [a] -> [b]
map TypeArg (DimDecl VName) -> Set VName
f [TypeArg (DimDecl VName)]
targs
  where
    f :: TypeArg (DimDecl VName) -> Set VName
f (TypeArgDim (NamedDim QualName VName
d) 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 SrcLoc
_) = StructType -> Set VName
arraySizes StructType
t
arraySizes (Scalar Prim {}) = Set VName
forall a. Monoid a => a
mempty
arraySizes (Array ()
_ Uniqueness
_ ScalarTypeBase (DimDecl VName) ()
t ShapeDecl (DimDecl VName)
shape) =
  StructType -> Set VName
arraySizes (ScalarTypeBase (DimDecl VName) () -> StructType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar ScalarTypeBase (DimDecl VName) ()
t) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (DimDecl VName -> Set VName) -> [DimDecl VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap DimDecl VName -> Set VName
dimName (ShapeDecl (DimDecl VName) -> [DimDecl VName]
forall dim. ShapeDecl dim -> [dim]
shapeDims ShapeDecl (DimDecl VName)
shape)
  where
    dimName :: DimDecl VName -> S.Set VName
    dimName :: DimDecl VName -> Set VName
dimName (NamedDim QualName VName
qn) = 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 DimDecl VName
_ = Set VName
forall a. Monoid a => a
mempty

patternArraySizes :: Pattern -> S.Set VName
patternArraySizes :: PatternBase Info VName -> Set VName
patternArraySizes = StructType -> Set VName
arraySizes (StructType -> Set VName)
-> (PatternBase Info VName -> StructType)
-> PatternBase Info VName
-> Set VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> StructType
patternStructType

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

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

dimMapping' ::
  Monoid a =>
  TypeBase (DimDecl VName) a ->
  TypeBase (DimDecl VName) a ->
  M.Map VName VName
dimMapping' :: forall a.
Monoid a =>
TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a -> Map VName VName
dimMapping' TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
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 (DimDecl VName) a
-> TypeBase (DimDecl VName) a -> Map VName SizeSubst
forall a.
Monoid a =>
TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a -> Map VName SizeSubst
dimMapping TypeBase (DimDecl VName) a
t1 TypeBase (DimDecl VName) a
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 (ExpBase Info VName
_, StaticVal
sv1) StaticVal
sv2) =
  StaticVal -> Set VName
sizesToRename StaticVal
sv1 Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> StaticVal -> Set VName
sizesToRename StaticVal
sv2
sizesToRename StaticVal
IntrinsicSV =
  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 (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, [PatternType])]
_) =
  (StaticVal -> Set VName) -> [StaticVal] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap StaticVal -> Set VName
sizesToRename [StaticVal]
svs
sizesToRename (LambdaSV PatternBase Info VName
param StructType
_ ExtExp
_ Env
_) =
  PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
param
    Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ((IdentBase Info VName -> Bool)
-> Set (IdentBase Info VName) -> Set (IdentBase Info VName)
forall a. (a -> Bool) -> Set a -> Set a
S.filter IdentBase Info VName -> Bool
forall {vn}. IdentBase Info vn -> Bool
couldBeSize (Set (IdentBase Info VName) -> Set (IdentBase Info VName))
-> Set (IdentBase Info VName) -> Set (IdentBase Info VName)
forall a b. (a -> b) -> a -> b
$ PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents PatternBase Info VName
param)
  where
    couldBeSize :: IdentBase Info vn -> Bool
couldBeSize IdentBase Info vn
ident =
      Info PatternType -> PatternType
forall a. Info a -> a
unInfo (IdentBase Info vn -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
identType IdentBase Info vn
ident) PatternType -> PatternType -> Bool
forall a. Eq a => a -> a -> Bool
== ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (IntType -> PrimType
Signed IntType
Int64))

-- 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
  Map VName SizeSubst
fresh_substs <- [VName] -> m (Map VName SizeSubst)
forall {f :: * -> *}.
MonadFreshNames f =>
[VName] -> f (Map VName SizeSubst)
mkSubsts ([VName] -> m (Map VName SizeSubst))
-> [VName] -> m (Map VName SizeSubst)
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
$ [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]
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 -> p -> Bool
isDim VName
k p
_ = VName
k VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
dims'
      dim_substs :: Map VName SizeSubst
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 a.
Monoid a =>
TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a -> Map VName SizeSubst
dimMapping (Map VName SizeSubst -> StructType -> StructType
forall als.
Map VName SizeSubst
-> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
replaceTypeSizes Map VName SizeSubst
fresh_substs StructType
sv_t) StructType
t
      replace :: SizeSubst -> SizeSubst
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 :: Map VName SizeSubst
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

  StaticVal -> m StaticVal
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StaticVal -> m StaticVal) -> StaticVal -> m StaticVal
forall a b. (a -> b) -> a -> b
$ Set VName -> Map VName SizeSubst -> StaticVal -> StaticVal
replaceStaticValSizes Set VName
globals Map VName SizeSubst
substs StaticVal
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)
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] ->
  [Pattern] ->
  Exp ->
  StructType ->
  SrcLoc ->
  DefM (Exp, StaticVal)
defuncFun :: [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> SrcLoc
-> DefM (ExpBase Info VName, StaticVal)
defuncFun [VName]
tparams [PatternBase Info VName]
pats ExpBase Info VName
e0 StructType
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 (PatternBase Info VName
pat, StructType
ret', ExtExp
e0') = case [PatternBase Info VName]
pats of
        [] -> String -> (PatternBase Info VName, StructType, ExtExp)
forall a. HasCallStack => String -> a
error String
"Received a lambda with no parameters."
        [PatternBase Info VName
pat'] -> (PatternBase Info VName
pat', StructType
ret, ExpBase Info VName -> ExtExp
ExtExp ExpBase Info VName
e0)
        (PatternBase Info VName
pat' : [PatternBase Info VName]
pats') ->
          ( PatternBase Info VName
pat',
            [StructType] -> StructType -> StructType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType ((PatternBase Info VName -> StructType)
-> [PatternBase Info VName] -> [StructType]
forall a b. (a -> b) -> [a] -> [b]
map (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> StructType)
-> (PatternBase Info VName -> PatternType)
-> PatternBase Info VName
-> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatternBase Info VName -> PatternType
patternType) [PatternBase Info VName]
pats') StructType
ret,
            [PatternBase Info VName]
-> ExpBase Info VName -> StructType -> SrcLoc -> ExtExp
ExtLambda [PatternBase Info VName]
pats' ExpBase Info VName
e0 StructType
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 :: NameSet
used =
        ExpBase Info VName -> NameSet
FV.freeVars ([PatternBase Info VName]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Set Alias, StructType)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
pats ExpBase Info VName
e0 Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Set Alias, StructType) -> Info (Set Alias, StructType)
forall a. a -> Info a
Info (Set Alias
forall a. Monoid a => a
mempty, StructType
ret)) SrcLoc
loc)
          NameSet -> Set VName -> NameSet
`FV.without` [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList [VName]
tparams
  Env
used_env <- NameSet -> DefM Env
restrictEnvTo NameSet
used

  -- The closure parts that are sizes are proactively turned into size
  -- parameters.
  let sizes_of_arrays :: Set VName
sizes_of_arrays =
        (Binding -> Set VName) -> Env -> Set VName
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
. PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> StructType)
-> (Binding -> PatternType) -> Binding -> StructType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticVal -> PatternType
typeFromSV (StaticVal -> PatternType)
-> (Binding -> StaticVal) -> Binding -> PatternType
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
<> PatternBase Info VName -> Set VName
patternArraySizes PatternBase Info VName
pat
      notSize :: VName -> Bool
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)
      ([FieldBase Info VName]
fields, Env
env) =
        ([(VName, Binding)] -> Env)
-> ([FieldBase Info VName], [(VName, Binding)])
-> ([FieldBase Info VName], Env)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second [(VName, Binding)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (([FieldBase Info VName], [(VName, Binding)])
 -> ([FieldBase Info VName], Env))
-> ([(VName, Binding)]
    -> ([FieldBase Info VName], [(VName, Binding)]))
-> [(VName, Binding)]
-> ([FieldBase Info VName], Env)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(FieldBase Info VName, (VName, Binding))]
-> ([FieldBase Info VName], [(VName, Binding)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(FieldBase Info VName, (VName, Binding))]
 -> ([FieldBase Info VName], [(VName, Binding)]))
-> ([(VName, Binding)]
    -> [(FieldBase Info VName, (VName, Binding))])
-> [(VName, Binding)]
-> ([FieldBase Info VName], [(VName, Binding)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Binding) -> (FieldBase Info VName, (VName, Binding)))
-> [(VName, Binding)] -> [(FieldBase Info VName, (VName, Binding))]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun
          ([(VName, Binding)] -> [(FieldBase Info VName, (VName, Binding))])
-> ([(VName, Binding)] -> [(VName, Binding)])
-> [(VName, Binding)]
-> [(FieldBase Info VName, (VName, Binding))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((VName, Binding) -> Bool)
-> [(VName, Binding)] -> [(VName, Binding)]
forall a. (a -> Bool) -> [a] -> [a]
filter (VName -> Bool
notSize (VName -> Bool)
-> ((VName, Binding) -> VName) -> (VName, Binding) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VName, Binding) -> VName
forall a b. (a, b) -> a
fst)
          ([(VName, Binding)] -> ([FieldBase Info VName], Env))
-> [(VName, Binding)] -> ([FieldBase Info VName], Env)
forall a b. (a -> b) -> a -> b
$ Env -> [(VName, Binding)]
forall k a. Map k a -> [(k, a)]
M.toList Env
used_env

  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( [FieldBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fields SrcLoc
loc,
      PatternBase Info VName -> StructType -> ExtExp -> Env -> StaticVal
LambdaSV PatternBase Info VName
pat StructType
ret' ExtExp
e0' Env
env
    )
  where
    closureFromDynamicFun :: (VName, Binding) -> (FieldBase Info VName, (VName, Binding))
closureFromDynamicFun (VName
vn, Binding Maybe ([VName], StructType)
_ (DynamicFun (ExpBase Info VName
clsr_env, StaticVal
sv) StaticVal
_)) =
      let name :: Name
name = String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
pretty VName
vn
       in ( Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
name ExpBase Info VName
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 = String -> Name
nameFromString (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ VName -> String
forall a. Pretty a => a -> String
pretty VName
vn
          tp' :: PatternType
tp' = StaticVal -> PatternType
typeFromSV StaticVal
sv
       in ( Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit
              Name
name
              (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
vn) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
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 :: ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp e :: ExpBase Info VName
e@Literal {} =
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e, PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e)
defuncExp e :: ExpBase Info VName
e@IntLit {} =
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e, PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e)
defuncExp e :: ExpBase Info VName
e@FloatLit {} =
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e, PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e)
defuncExp e :: ExpBase Info VName
e@StringLit {} =
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e, PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e)
defuncExp (Parens ExpBase Info VName
e SrcLoc
loc) = do
  (ExpBase Info VName
e', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens ExpBase Info VName
e' SrcLoc
loc, StaticVal
sv)
defuncExp (QualParens (QualName VName, SrcLoc)
qn ExpBase Info VName
e SrcLoc
loc) = do
  (ExpBase Info VName
e', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ((QualName VName, SrcLoc)
-> ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
(QualName vn, SrcLoc) -> ExpBase f vn -> SrcLoc -> ExpBase f vn
QualParens (QualName VName, SrcLoc)
qn ExpBase Info VName
e' SrcLoc
loc, StaticVal
sv)
defuncExp (TupLit [ExpBase Info VName]
es SrcLoc
loc) = do
  ([ExpBase Info VName]
es', [StaticVal]
svs) <- [(ExpBase Info VName, StaticVal)]
-> ([ExpBase Info VName], [StaticVal])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(ExpBase Info VName, StaticVal)]
 -> ([ExpBase Info VName], [StaticVal]))
-> DefM [(ExpBase Info VName, StaticVal)]
-> DefM ([ExpBase Info VName], [StaticVal])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal))
-> [ExpBase Info VName] -> DefM [(ExpBase Info VName, StaticVal)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp [ExpBase Info VName]
es
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [ExpBase Info VName]
es' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$ [Name] -> [StaticVal] -> [(Name, StaticVal)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tupleFieldNames [StaticVal]
svs)
defuncExp (RecordLit [FieldBase Info VName]
fs SrcLoc
loc) = do
  ([FieldBase Info VName]
fs', [(Name, StaticVal)]
names_svs) <- [(FieldBase Info VName, (Name, StaticVal))]
-> ([FieldBase Info VName], [(Name, StaticVal)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(FieldBase Info VName, (Name, StaticVal))]
 -> ([FieldBase Info VName], [(Name, StaticVal)]))
-> DefM [(FieldBase Info VName, (Name, StaticVal))]
-> DefM ([FieldBase Info VName], [(Name, StaticVal)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FieldBase Info VName
 -> DefM (FieldBase Info VName, (Name, StaticVal)))
-> [FieldBase Info VName]
-> DefM [(FieldBase Info VName, (Name, StaticVal))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField [FieldBase Info VName]
fs
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([FieldBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [FieldBase f vn] -> SrcLoc -> ExpBase f vn
RecordLit [FieldBase Info VName]
fs' SrcLoc
loc, [(Name, StaticVal)] -> StaticVal
RecordSV [(Name, StaticVal)]
names_svs)
  where
    defuncField :: FieldBase Info VName
-> DefM (FieldBase Info VName, (Name, StaticVal))
defuncField (RecordFieldExplicit Name
vn ExpBase Info VName
e SrcLoc
loc') = do
      (ExpBase Info VName
e', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
      (FieldBase Info VName, (Name, StaticVal))
-> DefM (FieldBase Info VName, (Name, StaticVal))
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn ExpBase Info VName
e' SrcLoc
loc', (Name
vn, StaticVal
sv))
    defuncField (RecordFieldImplicit VName
vn (Info PatternType
t) SrcLoc
loc') = do
      StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t) VName
vn
      case StaticVal
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 (ExpBase Info VName
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 (m :: * -> *) a. Monad m => a -> m a
return
                ( Name -> ExpBase Info VName -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> SrcLoc -> FieldBase f vn
RecordFieldExplicit Name
vn' ExpBase Info VName
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 PatternType
tp = PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv
           in (FieldBase Info VName, (Name, StaticVal))
-> DefM (FieldBase Info VName, (Name, StaticVal))
forall (m :: * -> *) a. Monad m => a -> m a
return (VName -> Info PatternType -> SrcLoc -> FieldBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> FieldBase f vn
RecordFieldImplicit VName
vn Info PatternType
tp SrcLoc
loc', (VName -> Name
baseName VName
vn, StaticVal
sv))
defuncExp (ArrayLit [ExpBase Info VName]
es t :: Info PatternType
t@(Info PatternType
t') SrcLoc
loc) = do
  [ExpBase Info VName]
es' <- (ExpBase Info VName -> DefM (ExpBase Info VName))
-> [ExpBase Info VName] -> DefM [ExpBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' [ExpBase Info VName]
es
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([ExpBase Info VName]
-> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
[ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
ArrayLit [ExpBase Info VName]
es' Info PatternType
t SrcLoc
loc, PatternType -> StaticVal
Dynamic PatternType
t')
defuncExp (Range ExpBase Info VName
e1 Maybe (ExpBase Info VName)
me Inclusiveness (ExpBase Info VName)
incl t :: (Info PatternType, Info [VName])
t@(Info PatternType
t', Info [VName]
_) SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e1
  Maybe (ExpBase Info VName)
me' <- (ExpBase Info VName -> DefM (ExpBase Info VName))
-> Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' Maybe (ExpBase Info VName)
me
  Inclusiveness (ExpBase Info VName)
incl' <- (ExpBase Info VName -> DefM (ExpBase Info VName))
-> Inclusiveness (ExpBase Info VName)
-> DefM (Inclusiveness (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' Inclusiveness (ExpBase Info VName)
incl
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> Maybe (ExpBase Info VName)
-> Inclusiveness (ExpBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> Maybe (ExpBase f vn)
-> Inclusiveness (ExpBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Range ExpBase Info VName
e1' Maybe (ExpBase Info VName)
me' Inclusiveness (ExpBase Info VName)
incl' (Info PatternType, Info [VName])
t SrcLoc
loc, PatternType -> StaticVal
Dynamic PatternType
t')
defuncExp e :: ExpBase Info VName
e@(Var QualName VName
qn (Info PatternType
t) SrcLoc
loc) = do
  StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)
  case StaticVal
sv of
    -- If the variable refers to a dynamic function, we return its closure
    -- representation (i.e., a record expression capturing the free variables
    -- and a 'LambdaSV' static value) instead of the variable itself.
    DynamicFun (ExpBase Info VName, StaticVal)
closure StaticVal
_ -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName, StaticVal)
closure
    -- Intrinsic functions used as variables are eta-expanded, so we
    -- can get rid of them.
    StaticVal
IntrinsicSV -> do
      ([PatternBase Info VName]
pats, ExpBase Info VName
body, StructType
tp) <- PatternType
-> ExpBase Info VName
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
etaExpand (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e) ExpBase Info VName
e
      ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp (ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal))
-> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ [PatternBase Info VName]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Set Alias, StructType)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
pats ExpBase Info VName
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Set Alias, StructType) -> Info (Set Alias, StructType)
forall a. a -> Info a
Info (Set Alias
forall a. Monoid a => a
mempty, StructType
tp)) SrcLoc
forall a. Monoid a => a
mempty
    StaticVal
_ ->
      let tp :: PatternType
tp = StaticVal -> PatternType
typeFromSV StaticVal
sv
       in (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
tp) SrcLoc
loc, StaticVal
sv)
defuncExp (Ascript ExpBase Info VName
e0 TypeDeclBase Info VName
tydecl SrcLoc
loc)
  | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e0) = do
    (ExpBase Info VName
e0', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
    (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> TypeDeclBase f vn -> SrcLoc -> ExpBase f vn
Ascript ExpBase Info VName
e0' TypeDeclBase Info VName
tydecl SrcLoc
loc, StaticVal
sv)
  | Bool
otherwise = ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
defuncExp (Coerce ExpBase Info VName
e0 TypeDeclBase Info VName
tydecl (Info PatternType, Info [VName])
t SrcLoc
loc)
  | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e0) = do
    (ExpBase Info VName
e0', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
    (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> TypeDeclBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> TypeDeclBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Coerce ExpBase Info VName
e0' TypeDeclBase Info VName
tydecl (Info PatternType, Info [VName])
t SrcLoc
loc, StaticVal
sv)
  | Bool
otherwise = ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
defuncExp (LetPat PatternBase Info VName
pat ExpBase Info VName
e1 ExpBase Info VName
e2 (Info PatternType
t, Info [VName]
retext) SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
sv1) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  let env :: Env
env = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat StaticVal
sv1
      pat' :: PatternBase Info VName
pat' = PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat StaticVal
sv1
  (ExpBase Info VName
e2', StaticVal
sv2) <- Env
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM (ExpBase Info VName, StaticVal)
 -> DefM (ExpBase Info VName, StaticVal))
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
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 :: Map VName VName
mapping = PatternType -> PatternType -> Map VName VName
forall a.
Monoid a =>
TypeBase (DimDecl VName) a
-> TypeBase (DimDecl VName) a -> Map VName VName
dimMapping' (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e2) PatternType
t
      subst :: VName -> VName
subst VName
v = VName -> Maybe VName -> VName
forall a. a -> Maybe a -> a
fromMaybe VName
v (Maybe VName -> VName) -> Maybe VName -> VName
forall a b. (a -> b) -> a -> 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' :: PatternType
t' = (DimDecl VName -> DimDecl VName) -> PatternType -> PatternType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((VName -> VName) -> DimDecl VName -> DimDecl VName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap VName -> VName
subst) (PatternType -> PatternType) -> PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e2'
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
LetPat PatternBase Info VName
pat' ExpBase Info VName
e1' ExpBase Info VName
e2' (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t', Info [VName]
retext) SrcLoc
loc, StaticVal
sv2)
defuncExp (LetFun VName
vn ([TypeParamBase VName], [PatternBase Info VName],
 Maybe (TypeExp VName), Info StructType, ExpBase Info VName)
_ ExpBase Info VName
_ Info PatternType
_ SrcLoc
_) =
  String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String -> DefM (ExpBase Info VName, StaticVal))
-> String -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ String
"defuncExp: Unexpected LetFun: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall v. IsName v => v -> String
prettyName VName
vn
defuncExp (If ExpBase Info VName
e1 ExpBase Info VName
e2 ExpBase Info VName
e3 (Info PatternType, Info [VName])
tp SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
_) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  (ExpBase Info VName
e2', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e2
  (ExpBase Info VName
e3', StaticVal
_) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e3
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> ExpBase Info VName
-> ExpBase Info VName
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> ExpBase f vn
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
If ExpBase Info VName
e1' ExpBase Info VName
e2' ExpBase Info VName
e3' (Info PatternType, Info [VName])
tp SrcLoc
loc, StaticVal
sv)
defuncExp e :: ExpBase Info VName
e@(Apply f :: ExpBase Info VName
f@(Var QualName VName
f' Info PatternType
_ SrcLoc
_) ExpBase Info VName
arg Info (Diet, Maybe VName)
d (Info PatternType
t, Info [VName]
ext) SrcLoc
loc)
  | VName -> Int
baseTag (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
f') Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxIntrinsicTag,
    TupLit [ExpBase Info VName]
es SrcLoc
tuploc <- ExpBase Info VName
arg = do
    -- defuncSoacExp also works fine for non-SOACs.
    [ExpBase Info VName]
es' <- (ExpBase Info VName -> DefM (ExpBase Info VName))
-> [ExpBase Info VName] -> DefM [ExpBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName)
defuncSoacExp [ExpBase Info VName]
es
    (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply ExpBase Info VName
f ([ExpBase Info VName] -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. [ExpBase f vn] -> SrcLoc -> ExpBase f vn
TupLit [ExpBase Info VName]
es' SrcLoc
tuploc) Info (Diet, Maybe VName)
d (Info PatternType
t, Info [VName]
ext) SrcLoc
loc,
        PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e
      )
defuncExp e :: ExpBase Info VName
e@Apply {} = Int -> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncApply Int
0 ExpBase Info VName
e
defuncExp (Negate ExpBase Info VName
e0 SrcLoc
loc) = do
  (ExpBase Info VName
e0', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Negate ExpBase Info VName
e0' SrcLoc
loc, StaticVal
sv)
defuncExp (Lambda [PatternBase Info VName]
pats ExpBase Info VName
e0 Maybe (TypeExp VName)
_ (Info (Set Alias
_, StructType
ret)) SrcLoc
loc) =
  [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> SrcLoc
-> DefM (ExpBase Info VName, StaticVal)
defuncFun [] [PatternBase Info VName]
pats ExpBase Info VName
e0 StructType
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 {} = String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"defuncExp: unexpected operator section."
defuncExp OpSectionLeft {} = String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"defuncExp: unexpected operator section."
defuncExp OpSectionRight {} = String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"defuncExp: unexpected operator section."
defuncExp ProjectSection {} = String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"defuncExp: unexpected projection section."
defuncExp IndexSection {} = String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"defuncExp: unexpected projection section."
defuncExp (DoLoop [VName]
sparams PatternBase Info VName
pat ExpBase Info VName
e1 LoopFormBase Info VName
form ExpBase Info VName
e3 Info (PatternType, [VName])
ret SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
sv1) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  let env1 :: Env
env1 = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat StaticVal
sv1
  (LoopFormBase Info VName
form', Env
env2) <- case LoopFormBase Info VName
form of
    For IdentBase Info VName
v ExpBase Info VName
e2 -> do
      ExpBase Info VName
e2' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e2
      (LoopFormBase Info VName, Env)
-> DefM (LoopFormBase Info VName, Env)
forall (m :: * -> *) a. Monad m => a -> m a
return (IdentBase Info VName
-> ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn -> ExpBase f vn -> LoopFormBase f vn
For IdentBase Info VName
v ExpBase Info VName
e2', IdentBase Info VName -> Env
forall {k}. IdentBase Info k -> Map k Binding
envFromIdent IdentBase Info VName
v)
    ForIn PatternBase Info VName
pat2 ExpBase Info VName
e2 -> do
      ExpBase Info VName
e2' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e2
      (LoopFormBase Info VName, Env)
-> DefM (LoopFormBase Info VName, Env)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternBase Info VName
-> ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> LoopFormBase f vn
ForIn PatternBase Info VName
pat2 ExpBase Info VName
e2', PatternBase Info VName -> Env
envFromPattern PatternBase Info VName
pat2)
    While ExpBase Info VName
e2 -> do
      ExpBase Info VName
e2' <- Env -> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a. Env -> DefM a -> DefM a
localEnv Env
env1 (DefM (ExpBase Info VName) -> DefM (ExpBase Info VName))
-> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e2
      (LoopFormBase Info VName, Env)
-> DefM (LoopFormBase Info VName, Env)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> LoopFormBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> LoopFormBase f vn
While ExpBase Info VName
e2', Env
forall a. Monoid a => a
mempty)
  (ExpBase Info VName
e3', StaticVal
sv) <- Env
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localEnv (Env
env1 Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
env2) (DefM (ExpBase Info VName, StaticVal)
 -> DefM (ExpBase Info VName, StaticVal))
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e3
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([VName]
-> PatternBase Info VName
-> ExpBase Info VName
-> LoopFormBase Info VName
-> ExpBase Info VName
-> Info (PatternType, [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[VName]
-> PatternBase f vn
-> ExpBase f vn
-> LoopFormBase f vn
-> ExpBase f vn
-> f (PatternType, [VName])
-> SrcLoc
-> ExpBase f vn
DoLoop [VName]
sparams PatternBase Info VName
pat ExpBase Info VName
e1' LoopFormBase Info VName
form' ExpBase Info VName
e3' Info (PatternType, [VName])
ret SrcLoc
loc, StaticVal
sv)
  where
    envFromIdent :: IdentBase Info k -> Map k Binding
envFromIdent (Ident k
vn (Info PatternType
tp) SrcLoc
_) =
      k -> Binding -> Map k Binding
forall k a. k -> a -> Map k a
M.singleton k
vn (Binding -> Map k Binding) -> 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
$ PatternType -> StaticVal
Dynamic PatternType
tp
defuncExp e :: ExpBase Info VName
e@BinOp {} =
  String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String -> DefM (ExpBase Info VName, StaticVal))
-> String -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ String
"defuncExp: unexpected binary operator: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e
defuncExp (Project Name
vn ExpBase Info VName
e0 tp :: Info PatternType
tp@(Info PatternType
tp') SrcLoc
loc) = do
  (ExpBase Info VName
e0', StaticVal
sv0) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e0
  case StaticVal
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 -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
vn ExpBase Info VName
e0' (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv) SrcLoc
loc, StaticVal
sv)
      Maybe StaticVal
Nothing -> String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error String
"Invalid record projection."
    Dynamic PatternType
_ -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> ExpBase f vn -> f PatternType -> SrcLoc -> ExpBase f vn
Project Name
vn ExpBase Info VName
e0' Info PatternType
tp SrcLoc
loc, PatternType -> StaticVal
Dynamic PatternType
tp')
    StaticVal
_ -> String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String -> DefM (ExpBase Info VName, StaticVal))
-> String -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ String
"Projection of an expression with static value " String -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> String
forall a. Show a => a -> String
show StaticVal
sv0
defuncExp (LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs ExpBase Info VName
e1 ExpBase Info VName
body Info PatternType
t SrcLoc
loc) = do
  ExpBase Info VName
e1' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e1
  [DimIndexBase Info VName]
idxs' <- (DimIndexBase Info VName -> DefM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> DefM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex [DimIndexBase Info VName]
idxs
  let id1_binding :: Binding
id1_binding = 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
$ PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ Info PatternType -> PatternType
forall a. Info a -> a
unInfo (Info PatternType -> PatternType)
-> Info PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ IdentBase Info VName -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
identType IdentBase Info VName
id1
  (ExpBase Info VName
body', StaticVal
sv) <-
    Env
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localEnv (VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton (IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName IdentBase Info VName
id1) Binding
id1_binding) (DefM (ExpBase Info VName, StaticVal)
 -> DefM (ExpBase Info VName, StaticVal))
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$
      ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
body
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (IdentBase Info VName
-> IdentBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
IdentBase f vn
-> IdentBase f vn
-> [DimIndexBase f vn]
-> ExpBase f vn
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
LetWith IdentBase Info VName
id1 IdentBase Info VName
id2 [DimIndexBase Info VName]
idxs' ExpBase Info VName
e1' ExpBase Info VName
body' Info PatternType
t SrcLoc
loc, StaticVal
sv)
defuncExp expr :: ExpBase Info VName
expr@(Index ExpBase Info VName
e0 [DimIndexBase Info VName]
idxs (Info PatternType, Info [VName])
info SrcLoc
loc) = do
  ExpBase Info VName
e0' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e0
  [DimIndexBase Info VName]
idxs' <- (DimIndexBase Info VName -> DefM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> DefM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex [DimIndexBase Info VName]
idxs
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> [DimIndexBase Info VName]
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn]
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Index ExpBase Info VName
e0' [DimIndexBase Info VName]
idxs' (Info PatternType, Info [VName])
info SrcLoc
loc, PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
expr)
defuncExp (Update ExpBase Info VName
e1 [DimIndexBase Info VName]
idxs ExpBase Info VName
e2 SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  [DimIndexBase Info VName]
idxs' <- (DimIndexBase Info VName -> DefM (DimIndexBase Info VName))
-> [DimIndexBase Info VName] -> DefM [DimIndexBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DimIndexBase Info VName -> DefM (DimIndexBase Info VName)
defuncDimIndex [DimIndexBase Info VName]
idxs
  ExpBase Info VName
e2' <- ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e2
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> [DimIndexBase Info VName]
-> ExpBase Info VName
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [DimIndexBase f vn] -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Update ExpBase Info VName
e1' [DimIndexBase Info VName]
idxs' ExpBase Info VName
e2' SrcLoc
loc, StaticVal
sv)

-- 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 (RecordUpdate ExpBase Info VName
e1 [Name]
fs ExpBase Info VName
e2 Info PatternType
_ SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
sv1) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  (ExpBase Info VName
e2', StaticVal
sv2) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e2
  let sv :: StaticVal
sv = StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv1 StaticVal
sv2 [Name]
fs
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( ExpBase Info VName
-> [Name]
-> ExpBase Info VName
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> [Name]
-> ExpBase f vn
-> f PatternType
-> SrcLoc
-> ExpBase f vn
RecordUpdate ExpBase Info VName
e1' [Name]
fs ExpBase Info VName
e2' (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv1) SrcLoc
loc,
      StaticVal
sv
    )
  where
    staticField :: StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (RecordSV [(Name, StaticVal)]
svs) StaticVal
sv2 (Name
f : [Name]
fs') =
      case Name -> [(Name, StaticVal)] -> Maybe StaticVal
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
f [(Name, StaticVal)]
svs of
        Just StaticVal
sv ->
          [(Name, StaticVal)] -> StaticVal
RecordSV ([(Name, StaticVal)] -> StaticVal)
-> [(Name, StaticVal)] -> StaticVal
forall a b. (a -> b) -> a -> b
$
            (Name
f, StaticVal -> StaticVal -> [Name] -> StaticVal
staticField StaticVal
sv StaticVal
sv2 [Name]
fs') (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)]
svs
        Maybe StaticVal
Nothing -> String -> StaticVal
forall a. HasCallStack => String -> a
error String
"Invalid record projection."
    staticField (Dynamic t :: PatternType
t@(Scalar Record {})) StaticVal
sv2 fs' :: [Name]
fs'@(Name
_ : [Name]
_) =
      StaticVal -> StaticVal -> [Name] -> StaticVal
staticField (PatternType -> StaticVal
svFromType PatternType
t) StaticVal
sv2 [Name]
fs'
    staticField StaticVal
_ StaticVal
sv2 [Name]
_ = StaticVal
sv2
defuncExp (Assert ExpBase Info VName
e1 ExpBase Info VName
e2 Info String
desc SrcLoc
loc) = do
  (ExpBase Info VName
e1', StaticVal
_) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
  (ExpBase Info VName
e2', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e2
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> ExpBase Info VName
-> Info String
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn -> ExpBase f vn -> f String -> SrcLoc -> ExpBase f vn
Assert ExpBase Info VName
e1' ExpBase Info VName
e2' Info String
desc SrcLoc
loc, StaticVal
sv)
defuncExp (Constr Name
name [ExpBase Info VName]
es (Info (Scalar (Sum Map Name [PatternType]
all_fs))) SrcLoc
loc) = do
  ([ExpBase Info VName]
es', [StaticVal]
svs) <- [(ExpBase Info VName, StaticVal)]
-> ([ExpBase Info VName], [StaticVal])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(ExpBase Info VName, StaticVal)]
 -> ([ExpBase Info VName], [StaticVal]))
-> DefM [(ExpBase Info VName, StaticVal)]
-> DefM ([ExpBase Info VName], [StaticVal])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal))
-> [ExpBase Info VName] -> DefM [(ExpBase Info VName, StaticVal)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp [ExpBase Info VName]
es
  let sv :: StaticVal
sv =
        Name -> [StaticVal] -> [(Name, [PatternType])] -> StaticVal
SumSV Name
name [StaticVal]
svs ([(Name, [PatternType])] -> StaticVal)
-> [(Name, [PatternType])] -> StaticVal
forall a b. (a -> b) -> a -> b
$
          Map Name [PatternType] -> [(Name, [PatternType])]
forall k a. Map k a -> [(k, a)]
M.toList (Map Name [PatternType] -> [(Name, [PatternType])])
-> Map Name [PatternType] -> [(Name, [PatternType])]
forall a b. (a -> b) -> a -> b
$
            Name
name Name -> Map Name [PatternType] -> Map Name [PatternType]
forall k a. Ord k => k -> Map k a -> Map k a
`M.delete` ([PatternType] -> [PatternType])
-> Map Name [PatternType] -> Map Name [PatternType]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((PatternType -> PatternType) -> [PatternType] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> PatternType
forall als.
Monoid als =>
TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
defuncType) Map Name [PatternType]
all_fs
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
-> [ExpBase Info VName]
-> Info PatternType
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
Name -> [ExpBase f vn] -> f PatternType -> SrcLoc -> ExpBase f vn
Constr Name
name [ExpBase Info VName]
es' (PatternType -> Info PatternType
forall a. a -> Info a
Info (StaticVal -> PatternType
typeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
  where
    defuncType ::
      Monoid als =>
      TypeBase (DimDecl VName) als ->
      TypeBase (DimDecl VName) als
    defuncType :: forall als.
Monoid als =>
TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
defuncType (Array als
as Uniqueness
u ScalarTypeBase (DimDecl VName) ()
t ShapeDecl (DimDecl VName)
shape) = als
-> Uniqueness
-> ScalarTypeBase (DimDecl VName) ()
-> ShapeDecl (DimDecl VName)
-> TypeBase (DimDecl VName) als
forall dim as.
as
-> Uniqueness
-> ScalarTypeBase dim ()
-> ShapeDecl dim
-> TypeBase dim as
Array als
as Uniqueness
u (ScalarTypeBase (DimDecl VName) ()
-> ScalarTypeBase (DimDecl VName) ()
forall als.
Monoid als =>
ScalarTypeBase (DimDecl VName) als
-> ScalarTypeBase (DimDecl VName) als
defuncScalar ScalarTypeBase (DimDecl VName) ()
t) ShapeDecl (DimDecl VName)
shape
    defuncType (Scalar ScalarTypeBase (DimDecl VName) als
t) = ScalarTypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) als
 -> TypeBase (DimDecl VName) als)
-> ScalarTypeBase (DimDecl VName) als
-> TypeBase (DimDecl VName) als
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) als
-> ScalarTypeBase (DimDecl VName) als
forall als.
Monoid als =>
ScalarTypeBase (DimDecl VName) als
-> ScalarTypeBase (DimDecl VName) als
defuncScalar ScalarTypeBase (DimDecl VName) als
t

    defuncScalar ::
      Monoid als =>
      ScalarTypeBase (DimDecl VName) als ->
      ScalarTypeBase (DimDecl VName) als
    defuncScalar :: forall als.
Monoid als =>
ScalarTypeBase (DimDecl VName) als
-> ScalarTypeBase (DimDecl VName) als
defuncScalar (Record Map Name (TypeBase (DimDecl VName) als)
fs) = Map Name (TypeBase (DimDecl VName) als)
-> ScalarTypeBase (DimDecl VName) als
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase (DimDecl VName) als)
 -> ScalarTypeBase (DimDecl VName) als)
-> Map Name (TypeBase (DimDecl VName) als)
-> ScalarTypeBase (DimDecl VName) als
forall a b. (a -> b) -> a -> b
$ (TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als)
-> Map Name (TypeBase (DimDecl VName) als)
-> Map Name (TypeBase (DimDecl VName) als)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
forall als.
Monoid als =>
TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
defuncType Map Name (TypeBase (DimDecl VName) als)
fs
    defuncScalar Arrow {} = Map Name (TypeBase (DimDecl VName) als)
-> ScalarTypeBase (DimDecl VName) als
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record Map Name (TypeBase (DimDecl VName) als)
forall a. Monoid a => a
mempty
    defuncScalar (Sum Map Name [TypeBase (DimDecl VName) als]
fs) = Map Name [TypeBase (DimDecl VName) als]
-> ScalarTypeBase (DimDecl VName) als
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase (DimDecl VName) als]
 -> ScalarTypeBase (DimDecl VName) als)
-> Map Name [TypeBase (DimDecl VName) als]
-> ScalarTypeBase (DimDecl VName) als
forall a b. (a -> b) -> a -> b
$ ([TypeBase (DimDecl VName) als] -> [TypeBase (DimDecl VName) als])
-> Map Name [TypeBase (DimDecl VName) als]
-> Map Name [TypeBase (DimDecl VName) als]
forall a b k. (a -> b) -> Map k a -> Map k b
M.map ((TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als)
-> [TypeBase (DimDecl VName) als] -> [TypeBase (DimDecl VName) als]
forall a b. (a -> b) -> [a] -> [b]
map TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
forall als.
Monoid als =>
TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als
defuncType) Map Name [TypeBase (DimDecl VName) als]
fs
    defuncScalar (Prim PrimType
t) = PrimType -> ScalarTypeBase (DimDecl VName) als
forall dim as. PrimType -> ScalarTypeBase dim as
Prim PrimType
t
    defuncScalar (TypeVar als
as Uniqueness
u TypeName
tn [TypeArg (DimDecl VName)]
targs) = als
-> Uniqueness
-> TypeName
-> [TypeArg (DimDecl VName)]
-> ScalarTypeBase (DimDecl VName) als
forall dim as.
as
-> Uniqueness -> TypeName -> [TypeArg dim] -> ScalarTypeBase dim as
TypeVar als
as Uniqueness
u TypeName
tn [TypeArg (DimDecl VName)]
targs
defuncExp (Constr Name
name [ExpBase Info VName]
_ (Info PatternType
t) SrcLoc
loc) =
  String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String -> DefM (ExpBase Info VName, StaticVal))
-> String -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$
    String
"Constructor " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" given type "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatternType -> String
forall a. Pretty a => a -> String
pretty PatternType
t
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" at "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ SrcLoc -> String
forall a. Located a => a -> String
locStr SrcLoc
loc
defuncExp (Match ExpBase Info VName
e NonEmpty (CaseBase Info VName)
cs (Info PatternType, Info [VName])
t SrcLoc
loc) = do
  (ExpBase Info VName
e', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
  NonEmpty (CaseBase Info VName, StaticVal)
csPairs <- (CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal))
-> NonEmpty (CaseBase Info VName)
-> DefM (NonEmpty (CaseBase Info VName, StaticVal))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv) NonEmpty (CaseBase Info VName)
cs
  let cs' :: NonEmpty (CaseBase Info VName)
cs' = ((CaseBase Info VName, StaticVal) -> CaseBase Info VName)
-> NonEmpty (CaseBase Info VName, StaticVal)
-> NonEmpty (CaseBase Info VName)
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' :: StaticVal
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
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
-> NonEmpty (CaseBase Info VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> NonEmpty (CaseBase f vn)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Match ExpBase Info VName
e' NonEmpty (CaseBase Info VName)
cs' (Info PatternType, Info [VName])
t SrcLoc
loc, StaticVal
sv')
defuncExp (Attr AttrInfo
info ExpBase Info VName
e SrcLoc
loc) = do
  (ExpBase Info VName
e', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
  (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (AttrInfo -> ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
AttrInfo -> ExpBase f vn -> SrcLoc -> ExpBase f vn
Attr AttrInfo
info ExpBase Info VName
e' SrcLoc
loc, StaticVal
sv)

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

defuncExtExp :: ExtExp -> DefM (Exp, StaticVal)
defuncExtExp :: ExtExp -> DefM (ExpBase Info VName, StaticVal)
defuncExtExp (ExtExp ExpBase Info VName
e) = ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
defuncExtExp (ExtLambda [PatternBase Info VName]
pats ExpBase Info VName
e0 StructType
ret SrcLoc
loc) =
  [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> SrcLoc
-> DefM (ExpBase Info VName, StaticVal)
defuncFun [] [PatternBase Info VName]
pats ExpBase Info VName
e0 StructType
ret SrcLoc
loc

defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal)
defuncCase :: StaticVal
-> CaseBase Info VName -> DefM (CaseBase Info VName, StaticVal)
defuncCase StaticVal
sv (CasePat PatternBase Info VName
p ExpBase Info VName
e SrcLoc
loc) = do
  let p' :: PatternBase Info VName
p' = PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
p StaticVal
sv
      env :: Env
env = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
p StaticVal
sv
  (ExpBase Info VName
e', StaticVal
sv') <- Env
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM (ExpBase Info VName, StaticVal)
 -> DefM (ExpBase Info VName, StaticVal))
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e
  (CaseBase Info VName, StaticVal)
-> DefM (CaseBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (PatternBase Info VName
-> ExpBase Info VName -> SrcLoc -> CaseBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> ExpBase f vn -> SrcLoc -> CaseBase f vn
CasePat PatternBase Info VName
p' ExpBase Info VName
e' SrcLoc
loc, StaticVal
sv')

-- | 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 :: ExpBase Info VName -> DefM (ExpBase Info VName)
defuncSoacExp e :: ExpBase Info VName
e@OpSection {} = ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
defuncSoacExp e :: ExpBase Info VName
e@OpSectionLeft {} = ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
defuncSoacExp e :: ExpBase Info VName
e@OpSectionRight {} = ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
defuncSoacExp e :: ExpBase Info VName
e@ProjectSection {} = ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return ExpBase Info VName
e
defuncSoacExp (Parens ExpBase Info VName
e SrcLoc
loc) =
  ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens (ExpBase Info VName -> SrcLoc -> ExpBase Info VName)
-> DefM (ExpBase Info VName) -> DefM (SrcLoc -> ExpBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> DefM (ExpBase Info VName)
defuncSoacExp ExpBase Info VName
e DefM (SrcLoc -> ExpBase Info VName)
-> DefM SrcLoc -> DefM (ExpBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SrcLoc -> DefM SrcLoc
forall (f :: * -> *) a. Applicative f => a -> f a
pure SrcLoc
loc
defuncSoacExp (Lambda [PatternBase Info VName]
params ExpBase Info VName
e0 Maybe (TypeExp VName)
decl Info (Set Alias, StructType)
tp SrcLoc
loc) = do
  let env :: Env
env = (PatternBase Info VName -> Env) -> [PatternBase Info VName] -> Env
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Env
envFromPattern [PatternBase Info VName]
params
  ExpBase Info VName
e0' <- Env -> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM (ExpBase Info VName) -> DefM (ExpBase Info VName))
-> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName)
defuncSoacExp ExpBase Info VName
e0
  ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> DefM (ExpBase Info VName))
-> ExpBase Info VName -> DefM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [PatternBase Info VName]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Set Alias, StructType)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
params ExpBase Info VName
e0' Maybe (TypeExp VName)
decl Info (Set Alias, StructType)
tp SrcLoc
loc
defuncSoacExp ExpBase Info VName
e
  | Scalar Arrow {} <- ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e = do
    ([PatternBase Info VName]
pats, ExpBase Info VName
body, StructType
tp) <- PatternType
-> ExpBase Info VName
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
etaExpand (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e) ExpBase Info VName
e
    let env :: Env
env = (PatternBase Info VName -> Env) -> [PatternBase Info VName] -> Env
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Env
envFromPattern [PatternBase Info VName]
pats
    ExpBase Info VName
body' <- Env -> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM (ExpBase Info VName) -> DefM (ExpBase Info VName))
-> DefM (ExpBase Info VName) -> DefM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
body
    ExpBase Info VName -> DefM (ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName -> DefM (ExpBase Info VName))
-> ExpBase Info VName -> DefM (ExpBase Info VName)
forall a b. (a -> b) -> a -> b
$ [PatternBase Info VName]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Set Alias, StructType)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
pats ExpBase Info VName
body' Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Set Alias, StructType) -> Info (Set Alias, StructType)
forall a. a -> Info a
Info (Set Alias
forall a. Monoid a => a
mempty, StructType
tp)) SrcLoc
forall a. Monoid a => a
mempty
  | Bool
otherwise = ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp' ExpBase Info VName
e

etaExpand :: PatternType -> Exp -> DefM ([Pattern], Exp, StructType)
etaExpand :: PatternType
-> ExpBase Info VName
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
etaExpand PatternType
e_t ExpBase Info VName
e = do
  let ([(PName, PatternType)]
ps, PatternType
ret) = PatternType -> ([(PName, PatternType)], PatternType)
forall {dim} {as}.
TypeBase dim as -> ([(PName, TypeBase dim as)], TypeBase dim as)
getType PatternType
e_t
  ([PatternBase Info VName]
pats, [ExpBase Info VName]
vars) <- ([(PatternBase Info VName, ExpBase Info VName)]
 -> ([PatternBase Info VName], [ExpBase Info VName]))
-> DefM [(PatternBase Info VName, ExpBase Info VName)]
-> DefM ([PatternBase Info VName], [ExpBase Info VName])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(PatternBase Info VName, ExpBase Info VName)]
-> ([PatternBase Info VName], [ExpBase Info VName])
forall a b. [(a, b)] -> ([a], [b])
unzip (DefM [(PatternBase Info VName, ExpBase Info VName)]
 -> DefM ([PatternBase Info VName], [ExpBase Info VName]))
-> (((PName, PatternType)
     -> DefM (PatternBase Info VName, ExpBase Info VName))
    -> DefM [(PatternBase Info VName, ExpBase Info VName)])
-> ((PName, PatternType)
    -> DefM (PatternBase Info VName, ExpBase Info VName))
-> DefM ([PatternBase Info VName], [ExpBase Info VName])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(PName, PatternType)]
-> ((PName, PatternType)
    -> DefM (PatternBase Info VName, ExpBase Info VName))
-> DefM [(PatternBase Info VName, ExpBase Info VName)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(PName, PatternType)]
ps (((PName, PatternType)
  -> DefM (PatternBase Info VName, ExpBase Info VName))
 -> DefM ([PatternBase Info VName], [ExpBase Info VName]))
-> ((PName, PatternType)
    -> DefM (PatternBase Info VName, ExpBase Info VName))
-> DefM ([PatternBase Info VName], [ExpBase Info VName])
forall a b. (a -> b) -> a -> b
$ \(PName
p, PatternType
t) -> do
    VName
x <- case PName
p of
      Named VName
x -> VName -> DefM VName
forall (f :: * -> *) a. Applicative f => a -> f a
pure VName
x
      PName
Unnamed -> String -> DefM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString String
"x"
    (PatternBase Info VName, ExpBase Info VName)
-> DefM (PatternBase Info VName, ExpBase Info VName)
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
x (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
forall a. Monoid a => a
mempty,
        QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var (VName -> QualName VName
forall v. v -> QualName v
qualName VName
x) (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) SrcLoc
forall a. Monoid a => a
mempty
      )
  let e' :: ExpBase Info VName
e' =
        (ExpBase Info VName
 -> (ExpBase Info VName, PatternType, [PatternType])
 -> ExpBase Info VName)
-> ExpBase Info VName
-> [(ExpBase Info VName, PatternType, [PatternType])]
-> ExpBase Info VName
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          ( \ExpBase Info VName
e1 (ExpBase Info VName
e2, PatternType
t2, [PatternType]
argtypes) ->
              ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
                ExpBase Info VName
e1
                ExpBase Info VName
e2
                ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (PatternType -> Diet
forall shape as. TypeBase shape as -> Diet
diet PatternType
t2, Maybe VName
forall a. Maybe a
Nothing))
                (PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [PatternType]
argtypes PatternType
ret), [VName] -> Info [VName]
forall a. a -> Info a
Info [])
                SrcLoc
forall a. Monoid a => a
mempty
          )
          ExpBase Info VName
e
          ([(ExpBase Info VName, PatternType, [PatternType])]
 -> ExpBase Info VName)
-> [(ExpBase Info VName, PatternType, [PatternType])]
-> ExpBase Info VName
forall a b. (a -> b) -> a -> b
$ [ExpBase Info VName]
-> [PatternType]
-> [[PatternType]]
-> [(ExpBase Info VName, PatternType, [PatternType])]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [ExpBase Info VName]
vars (((PName, PatternType) -> PatternType)
-> [(PName, PatternType)] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map (PName, PatternType) -> PatternType
forall a b. (a, b) -> b
snd [(PName, PatternType)]
ps) (Int -> [[PatternType]] -> [[PatternType]]
forall a. Int -> [a] -> [a]
drop Int
1 ([[PatternType]] -> [[PatternType]])
-> [[PatternType]] -> [[PatternType]]
forall a b. (a -> b) -> a -> b
$ [PatternType] -> [[PatternType]]
forall a. [a] -> [[a]]
tails ([PatternType] -> [[PatternType]])
-> [PatternType] -> [[PatternType]]
forall a b. (a -> b) -> a -> b
$ ((PName, PatternType) -> PatternType)
-> [(PName, PatternType)] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map (PName, PatternType) -> PatternType
forall a b. (a, b) -> b
snd [(PName, PatternType)]
ps)
  ([PatternBase Info VName], ExpBase Info VName, StructType)
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
forall (m :: * -> *) a. Monad m => a -> m a
return ([PatternBase Info VName]
pats, ExpBase Info VName
e', PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
ret)
  where
    getType :: TypeBase dim as -> ([(PName, TypeBase dim as)], TypeBase dim as)
getType (Scalar (Arrow as
_ PName
p TypeBase dim as
t1 TypeBase dim as
t2)) =
      let ([(PName, TypeBase dim as)]
ps, TypeBase dim as
r) = TypeBase dim as -> ([(PName, TypeBase dim as)], TypeBase dim as)
getType TypeBase dim as
t2 in ((PName
p, TypeBase dim as
t1) (PName, TypeBase dim as)
-> [(PName, TypeBase dim as)] -> [(PName, TypeBase dim as)]
forall a. a -> [a] -> [a]
: [(PName, TypeBase dim as)]
ps, TypeBase dim as
r)
    getType TypeBase dim as
t = ([], TypeBase dim as
t)

-- | 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 ExpBase Info VName
e1) = ExpBase Info VName -> DimIndexBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> DimIndexBase f vn
DimFix (ExpBase Info VName -> DimIndexBase Info VName)
-> ((ExpBase Info VName, StaticVal) -> ExpBase Info VName)
-> (ExpBase Info VName, StaticVal)
-> DimIndexBase Info VName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ExpBase Info VName, StaticVal) -> ExpBase Info VName
forall a b. (a, b) -> a
fst ((ExpBase Info VName, StaticVal) -> DimIndexBase Info VName)
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e1
defuncDimIndex (DimSlice Maybe (ExpBase Info VName)
me1 Maybe (ExpBase Info VName)
me2 Maybe (ExpBase Info VName)
me3) =
  Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> Maybe (ExpBase Info VName)
-> DimIndexBase Info VName
forall (f :: * -> *) vn.
Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> Maybe (ExpBase f vn)
-> DimIndexBase f vn
DimSlice (Maybe (ExpBase Info VName)
 -> Maybe (ExpBase Info VName)
 -> Maybe (ExpBase Info VName)
 -> DimIndexBase Info VName)
-> DefM (Maybe (ExpBase Info VName))
-> DefM
     (Maybe (ExpBase Info VName)
      -> Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
defunc' Maybe (ExpBase Info VName)
me1 DefM
  (Maybe (ExpBase Info VName)
   -> Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
-> DefM (Maybe (ExpBase Info VName))
-> DefM (Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
defunc' Maybe (ExpBase Info VName)
me2 DefM (Maybe (ExpBase Info VName) -> DimIndexBase Info VName)
-> DefM (Maybe (ExpBase Info VName))
-> DefM (DimIndexBase Info VName)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
defunc' Maybe (ExpBase Info VName)
me3
  where
    defunc' :: Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
defunc' = (ExpBase Info VName -> DefM (ExpBase Info VName))
-> Maybe (ExpBase Info VName) -> DefM (Maybe (ExpBase Info VName))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ExpBase Info VName -> DefM (ExpBase Info VName)
defuncExp'

-- | Defunctionalize a let-bound function, while preserving parameters
-- that have order 0 types (i.e., non-functional).
defuncLet ::
  [VName] ->
  [Pattern] ->
  Exp ->
  StructType ->
  DefM ([VName], [Pattern], Exp, StaticVal)
defuncLet :: [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
defuncLet [VName]
dims ps :: [PatternBase Info VName]
ps@(PatternBase Info VName
pat : [PatternBase Info VName]
pats) ExpBase Info VName
body StructType
rettype
  | PatternBase Info VName -> Bool
forall vn. PatternBase Info vn -> Bool
patternOrderZero PatternBase Info VName
pat = do
    let bound_by_pat :: VName -> Bool
bound_by_pat = (VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` PatternBase Info VName -> Set VName
patternDimNames PatternBase Info VName
pat)
        -- Take care to not include more size parameters than necessary.
        ([VName]
pat_dims, [VName]
rest_dims) = (VName -> Bool) -> [VName] -> ([VName], [VName])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition VName -> Bool
bound_by_pat [VName]
dims
        env :: Env
env = PatternBase Info VName -> Env
envFromPattern PatternBase Info VName
pat Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> [VName] -> Env
envFromDimNames [VName]
pat_dims
    ([VName]
rest_dims', [PatternBase Info VName]
pats', ExpBase Info VName
body', StaticVal
sv) <- Env
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM
   ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
 -> DefM
      ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal))
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
defuncLet [VName]
rest_dims [PatternBase Info VName]
pats ExpBase Info VName
body StructType
rettype
    (ExpBase Info VName, StaticVal)
closure <- [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> SrcLoc
-> DefM (ExpBase Info VName, StaticVal)
defuncFun [VName]
dims [PatternBase Info VName]
ps ExpBase Info VName
body StructType
rettype SrcLoc
forall a. Monoid a => a
mempty
    ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
      ( [VName]
pat_dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
rest_dims',
        PatternBase Info VName
pat PatternBase Info VName
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. a -> [a] -> [a]
: [PatternBase Info VName]
pats',
        ExpBase Info VName
body',
        (ExpBase Info VName, StaticVal) -> StaticVal -> StaticVal
DynamicFun (ExpBase Info VName, StaticVal)
closure StaticVal
sv
      )
  | Bool
otherwise = do
    (ExpBase Info VName
e, StaticVal
sv) <- [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> SrcLoc
-> DefM (ExpBase Info VName, StaticVal)
defuncFun [VName]
dims [PatternBase Info VName]
ps ExpBase Info VName
body StructType
rettype SrcLoc
forall a. Monoid a => a
mempty
    ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [], ExpBase Info VName
e, StaticVal
sv)
defuncLet [VName]
_ [] ExpBase Info VName
body StructType
rettype = do
  (ExpBase Info VName
body', StaticVal
sv) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
body
  ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return ([], [], ExpBase Info VName
body', StaticVal -> StructType -> StaticVal
forall {as}. StaticVal -> TypeBase (DimDecl VName) as -> StaticVal
imposeType StaticVal
sv StructType
rettype)
  where
    imposeType :: StaticVal -> TypeBase (DimDecl VName) as -> StaticVal
imposeType Dynamic {} TypeBase (DimDecl VName) as
t =
      PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ TypeBase (DimDecl VName) as -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase (DimDecl VName) as
t
    imposeType (RecordSV [(Name, StaticVal)]
fs1) (Scalar (Record Map Name (TypeBase (DimDecl VName) as)
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 -> TypeBase (DimDecl VName) as -> StaticVal)
-> Map Name StaticVal
-> Map Name (TypeBase (DimDecl VName) as)
-> 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 -> TypeBase (DimDecl VName) as -> StaticVal
imposeType ([(Name, StaticVal)] -> Map Name StaticVal
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, StaticVal)]
fs1) Map Name (TypeBase (DimDecl VName) as)
fs2
    imposeType StaticVal
sv TypeBase (DimDecl VName) as
_ = StaticVal
sv

sizesForAll :: MonadFreshNames m => S.Set VName -> [Pattern] -> m ([VName], [Pattern])
sizesForAll :: forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatternBase Info VName]
-> m ([VName], [PatternBase Info VName])
sizesForAll Set VName
bound_sizes [PatternBase Info VName]
params = do
  ([PatternBase Info VName]
params', Set VName
sizes) <- StateT (Set VName) m [PatternBase Info VName]
-> Set VName -> m ([PatternBase Info VName], Set VName)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((PatternBase Info VName
 -> StateT (Set VName) m (PatternBase Info VName))
-> [PatternBase Info VName]
-> StateT (Set VName) m [PatternBase Info VName]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ASTMapper (StateT (Set VName) m)
-> PatternBase Info VName
-> StateT (Set VName) m (PatternBase Info VName)
forall x (m :: * -> *).
(ASTMappable x, Monad m) =>
ASTMapper m -> x -> m x
astMap ASTMapper (StateT (Set VName) m)
tv) [PatternBase Info VName]
params) Set VName
forall a. Monoid a => a
mempty
  ([VName], [PatternBase Info VName])
-> m ([VName], [PatternBase Info VName])
forall (m :: * -> *) a. Monad m => a -> m a
return (Set VName -> [VName]
forall a. Set a -> [a]
S.toList Set VName
sizes, [PatternBase Info VName]
params')
  where
    bound :: Set VName
bound = Set VName
bound_sizes Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set vn
patternNames [PatternBase Info VName]
params
    tv :: ASTMapper (StateT (Set VName) m)
tv = ASTMapper (StateT (Set VName) m)
forall (m :: * -> *). Monad m => ASTMapper m
identityMapper {mapOnPatternType :: PatternType -> StateT (Set VName) m PatternType
mapOnPatternType = (DimDecl VName -> StateT (Set VName) m (DimDecl VName))
-> (Set Alias -> StateT (Set VName) m (Set Alias))
-> PatternType
-> StateT (Set VName) m PatternType
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse DimDecl VName -> StateT (Set VName) m (DimDecl VName)
forall {t :: (* -> *) -> * -> *} {m :: * -> *}.
(MonadTrans t, MonadFreshNames m, MonadState (Set VName) (t m)) =>
DimDecl VName -> t m (DimDecl VName)
onDim Set Alias -> StateT (Set VName) m (Set Alias)
forall (f :: * -> *) a. Applicative f => a -> f a
pure}
    onDim :: DimDecl VName -> t m (DimDecl VName)
onDim DimDecl VName
AnyDim = do
      VName
v <- m VName -> t m VName
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m VName -> t m VName) -> m VName -> t m VName
forall a b. (a -> b) -> a -> b
$ String -> m VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newVName String
"size"
      (Set VName -> Set VName) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName -> Set VName) -> t m ())
-> (Set VName -> Set VName) -> t m ()
forall a b. (a -> b) -> a -> b
$ VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert VName
v
      DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim (QualName VName -> DimDecl VName)
-> QualName VName -> DimDecl VName
forall a b. (a -> b) -> a -> b
$ VName -> QualName VName
forall v. v -> QualName v
qualName VName
v
    onDim (NamedDim QualName VName
d) = do
      Bool -> t m () -> t 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) (t m () -> t m ()) -> t m () -> t m ()
forall a b. (a -> b) -> a -> b
$
        (Set VName -> Set VName) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Set VName -> Set VName) -> t m ())
-> (Set VName -> Set VName) -> t m ()
forall a b. (a -> b) -> a -> b
$ VName -> Set VName -> Set VName
forall a. Ord a => a -> Set a -> Set a
S.insert (VName -> Set VName -> Set VName)
-> VName -> Set VName -> Set VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
d
      DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DimDecl VName -> t m (DimDecl VName))
-> DimDecl VName -> t m (DimDecl VName)
forall a b. (a -> b) -> a -> b
$ QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
d
    onDim DimDecl VName
d = DimDecl VName -> t m (DimDecl VName)
forall (f :: * -> *) a. Applicative f => a -> f a
pure DimDecl VName
d

-- | Defunctionalize an application expression at a given depth of application.
-- Calls to dynamic (first-order) functions are preserved at much as possible,
-- but a new lifted function is created if a dynamic function is only partially
-- applied.
defuncApply :: Int -> Exp -> DefM (Exp, StaticVal)
defuncApply :: Int -> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncApply Int
depth e :: ExpBase Info VName
e@(Apply ExpBase Info VName
e1 ExpBase Info VName
e2 Info (Diet, Maybe VName)
d t :: (Info PatternType, Info [VName])
t@(Info PatternType
ret, Info [VName]
ext) SrcLoc
loc) = do
  let ([PatternType]
argtypes, PatternType
_) = PatternType -> ([PatternType], PatternType)
forall dim as.
TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType PatternType
ret
  (ExpBase Info VName
e1', StaticVal
sv1) <- Int -> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncApply (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ExpBase Info VName
e1
  (ExpBase Info VName
e2', StaticVal
sv2) <- ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
e2
  let e' :: ExpBase Info VName
e' = ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply ExpBase Info VName
e1' ExpBase Info VName
e2' Info (Diet, Maybe VName)
d (Info PatternType, Info [VName])
t SrcLoc
loc
  case StaticVal
sv1 of
    LambdaSV PatternBase Info VName
pat StructType
e0_t ExtExp
e0 Env
closure_env -> do
      let env' :: Env
env' = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat StaticVal
sv2
          dims :: [VName]
dims = [VName]
forall a. Monoid a => a
mempty
      (ExpBase Info VName
e0', StaticVal
sv) <-
        Env
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a. Env -> DefM a -> DefM a
localNewEnv (Env
env' Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> Env
closure_env) (DefM (ExpBase Info VName, StaticVal)
 -> DefM (ExpBase Info VName, StaticVal))
-> DefM (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$
          ExtExp -> DefM (ExpBase Info VName, StaticVal)
defuncExtExp ExtExp
e0

      let closure_pat :: PatternBase Info VName
closure_pat = [VName] -> Env -> PatternBase Info VName
buildEnvPattern [VName]
dims Env
closure_env
          pat' :: PatternBase Info VName
pat' = PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat StaticVal
sv2

      Set VName
globals <- ((Set VName, Env) -> Set VName) -> DefM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Set VName
forall a b. (a, b) -> a
fst

      -- Lift lambda to top-level function definition.  We put in
      -- a lot of effort to try to infer the uniqueness attributes
      -- of the lifted function, but this is ultimately all a sham
      -- and a hack.  There is some piece we're missing.
      let params :: [PatternBase Info VName]
params = [PatternBase Info VName
closure_pat, PatternBase Info VName
pat']
          params_for_rettype :: [PatternBase Info VName]
params_for_rettype = [PatternBase Info VName]
params [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatternBase Info VName]
svParams StaticVal
sv1 [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. [a] -> [a] -> [a]
++ StaticVal -> [PatternBase Info VName]
svParams StaticVal
sv2
          svParams :: StaticVal -> [PatternBase Info VName]
svParams (LambdaSV PatternBase Info VName
sv_pat StructType
_ ExtExp
_ Env
_) = [PatternBase Info VName
sv_pat]
          svParams StaticVal
_ = []
          rettype :: PatternType
rettype = Env
-> [PatternBase Info VName]
-> StructType
-> PatternType
-> PatternType
buildRetType Env
closure_env [PatternBase Info VName]
params_for_rettype StructType
e0_t (PatternType -> PatternType) -> PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e0'

          already_bound :: Set VName
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
              Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ((PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> Set (IdentBase Info VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
params)

          more_dims :: [VName]
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
$
                (PatternBase Info VName -> Set VName)
-> [PatternBase Info VName] -> Set VName
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set VName
patternArraySizes [PatternBase Info VName]
params

          -- Embed some information about the original function
          -- into the name of the lifted function, to make the
          -- result slightly more human-readable.
          liftedName :: t -> ExpBase f VName -> String
liftedName t
i (Var QualName VName
f f PatternType
_ SrcLoc
_) =
            String
"defunc_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ t -> String
forall a. Show a => a -> String
show t
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"_" String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
baseString (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
f)
          liftedName t
i (Apply ExpBase f VName
f ExpBase f VName
_ f (Diet, Maybe VName)
_ (f PatternType, f [VName])
_ SrcLoc
_) =
            t -> ExpBase f VName -> String
liftedName (t
i t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) ExpBase f VName
f
          liftedName t
_ ExpBase f VName
_ = String
"defunc"

      -- Ensure that no parameter sizes are AnyDim.  The internaliser
      -- expects this.  This is easy, because they are all
      -- first-order.
      let bound_sizes :: Set VName
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
      ([VName]
missing_dims, [PatternBase Info VName]
params') <- Set VName
-> [PatternBase Info VName]
-> DefM ([VName], [PatternBase Info VName])
forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatternBase Info VName]
-> m ([VName], [PatternBase Info VName])
sizesForAll Set VName
bound_sizes [PatternBase Info VName]
params

      VName
fname <- String -> DefM VName
forall (m :: * -> *). MonadFreshNames m => String -> m VName
newNameFromString (String -> DefM VName) -> String -> DefM VName
forall a b. (a -> b) -> a -> b
$ Int -> ExpBase Info VName -> String
forall {t} {f :: * -> *}.
(Show t, Num t) =>
t -> ExpBase f VName -> String
liftedName (Int
0 :: Int) ExpBase Info VName
e1
      VName
-> PatternType
-> [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> DefM ()
liftValDec
        VName
fname
        PatternType
rettype
        ([VName]
dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
more_dims [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims)
        [PatternBase Info VName]
params'
        ExpBase Info VName
e0'

      let t1 :: StructType
t1 = PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> StructType) -> PatternType -> StructType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e1'
          t2 :: StructType
t2 = PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> StructType) -> PatternType -> StructType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e2'
          fname' :: QualName VName
fname' = VName -> QualName VName
forall v. v -> QualName v
qualName VName
fname
          fname'' :: ExpBase Info VName
fname'' =
            QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var
              QualName VName
fname'
              ( PatternType -> Info PatternType
forall a. a -> Info a
Info
                  ( ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$
                      Set Alias
-> PName
-> PatternType
-> PatternType
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow Set Alias
forall a. Monoid a => a
mempty PName
Unnamed (StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
t1) (PatternType -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> PatternType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$
                        ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ Set Alias
-> PName
-> PatternType
-> PatternType
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow Set Alias
forall a. Monoid a => a
mempty PName
Unnamed (StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
t2) PatternType
rettype
                  )
              )
              SrcLoc
loc

          -- FIXME: what if this application returns both a function
          -- and a value?
          callret :: (Info PatternType, Info [VName])
callret
            | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
ret = (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
ret, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
ext)
            | Bool
otherwise = (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
rettype, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
ext)

      (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
        ( ExpBase Info VName -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn. ExpBase f vn -> SrcLoc -> ExpBase f vn
Parens
            ( ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
                ( ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply
                    ExpBase Info VName
fname''
                    ExpBase Info VName
e1'
                    ((Diet, Maybe VName) -> Info (Diet, Maybe VName)
forall a. a -> Info a
Info (Diet
Observe, Maybe VName
forall a. Maybe a
Nothing))
                    ( PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ Set Alias
-> PName
-> PatternType
-> PatternType
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as.
as
-> PName
-> TypeBase dim as
-> TypeBase dim as
-> ScalarTypeBase dim as
Arrow Set Alias
forall a. Monoid a => a
mempty PName
Unnamed (StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
t2) PatternType
rettype,
                      [VName] -> Info [VName]
forall a. a -> Info a
Info []
                    )
                    SrcLoc
loc
                )
                ExpBase Info VName
e2'
                Info (Diet, Maybe VName)
d
                (Info PatternType, Info [VName])
callret
                SrcLoc
loc
            )
            SrcLoc
forall a. Monoid a => a
mempty,
          StaticVal
sv
        )

    -- If e1 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 term.
    DynamicFun (ExpBase Info VName, StaticVal)
_ StaticVal
sv -> do
      let ([PatternType]
argtypes', PatternType
rettype) = StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType StaticVal
sv [PatternType]
argtypes
          restype :: PatternType
restype = [PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [PatternType]
argtypes' PatternType
rettype PatternType -> Set Alias -> PatternType
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` PatternType -> Set Alias
forall as shape. Monoid as => TypeBase shape as -> as
aliases PatternType
ret
          -- FIXME: what if this application returns both a function
          -- and a value?
          callret :: (Info PatternType, Info [VName])
callret
            | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
ret = (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
ret, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
ext)
            | Bool
otherwise = (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
restype, [VName] -> Info [VName]
forall a. a -> Info a
Info [VName]
ext)
          apply_e :: ExpBase Info VName
apply_e = ExpBase Info VName
-> ExpBase Info VName
-> Info (Diet, Maybe VName)
-> (Info PatternType, Info [VName])
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
ExpBase f vn
-> ExpBase f vn
-> f (Diet, Maybe VName)
-> (f PatternType, f [VName])
-> SrcLoc
-> ExpBase f vn
Apply ExpBase Info VName
e1' ExpBase Info VName
e2' Info (Diet, Maybe VName)
d (Info PatternType, Info [VName])
callret SrcLoc
loc
      (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
apply_e, StaticVal
sv)
    -- Propagate the 'IntrinsicsSV' until we reach the outermost application,
    -- where we construct a dynamic static value with the appropriate type.
    StaticVal
IntrinsicSV
      | Int
depth Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 ->
        -- 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 [PatternType] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatternType]
argtypes
          then (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e', PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e)
          else do
            ([PatternBase Info VName]
pats, ExpBase Info VName
body, StructType
tp) <- PatternType
-> ExpBase Info VName
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
etaExpand (ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
e') ExpBase Info VName
e'
            ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp (ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal))
-> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$ [PatternBase Info VName]
-> ExpBase Info VName
-> Maybe (TypeExp VName)
-> Info (Set Alias, StructType)
-> SrcLoc
-> ExpBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn]
-> ExpBase f vn
-> Maybe (TypeExp vn)
-> f (Set Alias, StructType)
-> SrcLoc
-> ExpBase f vn
Lambda [PatternBase Info VName]
pats ExpBase Info VName
body Maybe (TypeExp VName)
forall a. Maybe a
Nothing ((Set Alias, StructType) -> Info (Set Alias, StructType)
forall a. a -> Info a
Info (Set Alias
forall a. Monoid a => a
mempty, StructType
tp)) SrcLoc
forall a. Monoid a => a
mempty
      | Bool
otherwise -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e', StaticVal
IntrinsicSV)
    StaticVal
_ ->
      String -> DefM (ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String -> DefM (ExpBase Info VName, StaticVal))
-> String -> DefM (ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$
        String
"Application of an expression\n"
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ ExpBase Info VName -> String
forall a. Pretty a => a -> String
pretty ExpBase Info VName
e1
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\nthat is neither a static lambda "
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"nor a dynamic function, but has static value:\n"
          String -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> String
forall a. Show a => a -> String
show StaticVal
sv1
defuncApply Int
depth e :: ExpBase Info VName
e@(Var QualName VName
qn (Info PatternType
t) SrcLoc
loc) = do
  let ([PatternType]
argtypes, PatternType
_) = PatternType -> ([PatternType], PatternType)
forall dim as.
TypeBase dim as -> ([TypeBase dim as], TypeBase dim as)
unfoldFunType PatternType
t
  StaticVal
sv <- StructType -> VName -> DefM StaticVal
lookupVar (PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
t) (QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn)

  case StaticVal
sv of
    DynamicFun (ExpBase Info VName, StaticVal)
_ StaticVal
_
      | StaticVal -> Int -> Bool
fullyApplied StaticVal
sv Int
depth -> do
        -- We still need to update the types in case the dynamic
        -- function returns a higher-order term.
        let ([PatternType]
argtypes', PatternType
rettype) = StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType StaticVal
sv [PatternType]
argtypes
        (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [PatternType]
argtypes' PatternType
rettype)) SrcLoc
loc, StaticVal
sv)
      | Bool
otherwise -> do
        VName
fname <- VName -> DefM VName
forall (m :: * -> *). MonadFreshNames m => VName -> m VName
newName (VName -> DefM VName) -> VName -> DefM VName
forall a b. (a -> b) -> a -> b
$ QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
qn
        let ([PatternBase Info VName]
pats, ExpBase Info VName
e0, StaticVal
sv') = String
-> StaticVal
-> Int
-> ([PatternBase Info VName], ExpBase Info VName, StaticVal)
liftDynFun (QualName VName -> String
forall a. Pretty a => a -> String
pretty QualName VName
qn) StaticVal
sv Int
depth
            ([PatternType]
argtypes', PatternType
rettype) = StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType StaticVal
sv' [PatternType]
argtypes
            dims' :: [VName]
dims' = [VName]
forall a. Monoid a => a
mempty

        -- Ensure that no parameter sizes are AnyDim.  The internaliser
        -- expects this.  This is easy, because they are all
        -- first-order.
        Set VName
globals <- ((Set VName, Env) -> Set VName) -> DefM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Set VName
forall a b. (a, b) -> a
fst
        let bound_sizes :: Set VName
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
        ([VName]
missing_dims, [PatternBase Info VName]
pats') <- Set VName
-> [PatternBase Info VName]
-> DefM ([VName], [PatternBase Info VName])
forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatternBase Info VName]
-> m ([VName], [PatternBase Info VName])
sizesForAll Set VName
bound_sizes [PatternBase Info VName]
pats

        VName
-> PatternType
-> [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> DefM ()
liftValDec VName
fname (PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct PatternType
rettype) ([VName]
dims' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims) [PatternBase Info VName]
pats' ExpBase Info VName
e0
        (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return
          ( QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var
              (VName -> QualName VName
forall v. v -> QualName v
qualName VName
fname)
              (PatternType -> Info PatternType
forall a. a -> Info a
Info ([PatternType] -> PatternType -> PatternType
forall as dim.
Monoid as =>
[TypeBase dim as] -> TypeBase dim as -> TypeBase dim as
foldFunType [PatternType]
argtypes' (PatternType -> PatternType) -> PatternType -> PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct PatternType
rettype))
              SrcLoc
loc,
            StaticVal
sv'
          )
    StaticVal
IntrinsicSV -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (ExpBase Info VName
e, StaticVal
IntrinsicSV)
    StaticVal
_ -> (ExpBase Info VName, StaticVal)
-> DefM (ExpBase Info VName, StaticVal)
forall (m :: * -> *) a. Monad m => a -> m a
return (QualName VName -> Info PatternType -> SrcLoc -> ExpBase Info VName
forall (f :: * -> *) vn.
QualName vn -> f PatternType -> SrcLoc -> ExpBase f vn
Var QualName VName
qn (PatternType -> Info PatternType
forall a. a -> Info a
Info (StaticVal -> PatternType
typeFromSV StaticVal
sv)) SrcLoc
loc, StaticVal
sv)
defuncApply Int
depth (Parens ExpBase Info VName
e SrcLoc
_) = Int -> ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncApply Int
depth ExpBase Info VName
e
defuncApply Int
_ ExpBase Info VName
expr = ExpBase Info VName -> DefM (ExpBase Info VName, StaticVal)
defuncExp ExpBase Info VName
expr

-- | 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 (ExpBase Info VName, 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 :: String -> StaticVal -> Int -> ([Pattern], Exp, StaticVal)
liftDynFun :: String
-> StaticVal
-> Int
-> ([PatternBase Info VName], ExpBase Info VName, StaticVal)
liftDynFun String
_ (DynamicFun (ExpBase Info VName
e, StaticVal
sv) StaticVal
_) Int
0 = ([], ExpBase Info VName
e, StaticVal
sv)
liftDynFun String
s (DynamicFun clsr :: (ExpBase Info VName, StaticVal)
clsr@(ExpBase Info VName
_, LambdaSV PatternBase Info VName
pat StructType
_ ExtExp
_ Env
_) StaticVal
sv) Int
d
  | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 =
    let ([PatternBase Info VName]
pats, ExpBase Info VName
e', StaticVal
sv') = String
-> StaticVal
-> Int
-> ([PatternBase Info VName], ExpBase Info VName, StaticVal)
liftDynFun String
s StaticVal
sv (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
     in (PatternBase Info VName
pat PatternBase Info VName
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. a -> [a] -> [a]
: [PatternBase Info VName]
pats, ExpBase Info VName
e', (ExpBase Info VName, StaticVal) -> StaticVal -> StaticVal
DynamicFun (ExpBase Info VName, StaticVal)
clsr StaticVal
sv')
liftDynFun String
s StaticVal
sv Int
d =
  String -> ([PatternBase Info VName], ExpBase Info VName, StaticVal)
forall a. HasCallStack => String -> a
error (String
 -> ([PatternBase Info VName], ExpBase Info VName, StaticVal))
-> String
-> ([PatternBase Info VName], ExpBase Info VName, StaticVal)
forall a b. (a -> b) -> a -> b
$
    String
s
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" Tried to lift a StaticVal "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> ShowS
forall a. Int -> [a] -> [a]
take Int
100 (StaticVal -> String
forall a. Show a => a -> String
show StaticVal
sv)
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", but expected a dynamic function.\n"
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Pretty a => a -> String
pretty 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.
envFromPattern :: Pattern -> Env
envFromPattern :: PatternBase Info VName -> Env
envFromPattern PatternBase Info VName
pat = case PatternBase Info VName
pat of
  TuplePattern [PatternBase Info VName]
ps SrcLoc
_ -> (PatternBase Info VName -> Env) -> [PatternBase Info VName] -> Env
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Env
envFromPattern [PatternBase Info VName]
ps
  RecordPattern [(Name, PatternBase Info VName)]
fs SrcLoc
_ -> ((Name, PatternBase Info VName) -> Env)
-> [(Name, PatternBase Info VName)] -> Env
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (PatternBase Info VName -> Env
envFromPattern (PatternBase Info VName -> Env)
-> ((Name, PatternBase Info VName) -> PatternBase Info VName)
-> (Name, PatternBase Info VName)
-> Env
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, PatternBase Info VName) -> PatternBase Info VName
forall a b. (a, b) -> b
snd) [(Name, PatternBase Info VName)]
fs
  PatternParens PatternBase Info VName
p SrcLoc
_ -> PatternBase Info VName -> Env
envFromPattern PatternBase Info VName
p
  Id VName
vn (Info PatternType
t) SrcLoc
_ -> VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton VName
vn (Binding -> Env) -> Binding -> Env
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
$ PatternType -> StaticVal
Dynamic PatternType
t
  Wildcard Info PatternType
_ SrcLoc
_ -> Env
forall a. Monoid a => a
mempty
  PatternAscription PatternBase Info VName
p TypeDeclBase Info VName
_ SrcLoc
_ -> PatternBase Info VName -> Env
envFromPattern PatternBase Info VName
p
  PatternLit {} -> Env
forall a. Monoid a => a
mempty
  PatternConstr Name
_ Info PatternType
_ [PatternBase Info VName]
ps SrcLoc
_ -> (PatternBase Info VName -> Env) -> [PatternBase Info VName] -> Env
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Env
envFromPattern [PatternBase Info VName]
ps

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
$ PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64

-- | Create a new top-level value declaration with the given function name,
-- return type, list of parameters, and body expression.
liftValDec :: VName -> PatternType -> [VName] -> [Pattern] -> Exp -> DefM ()
liftValDec :: VName
-> PatternType
-> [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> DefM ()
liftValDec VName
fname PatternType
rettype [VName]
dims [PatternBase Info VName]
pats ExpBase Info VName
body = Seq (ValBindBase Info VName) -> DefM ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (Seq (ValBindBase Info VName) -> DefM ())
-> Seq (ValBindBase Info VName) -> DefM ()
forall a b. (a -> b) -> a -> b
$ ValBindBase Info VName -> Seq (ValBindBase Info VName)
forall a. a -> Seq a
Seq.singleton ValBindBase Info VName
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]
dims Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ((PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> Set (IdentBase Info VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
pats)
    anyDimIfNotBound :: DimDecl VName -> DimDecl VName
anyDimIfNotBound (NamedDim QualName VName
v)
      | QualName VName -> VName
forall vn. QualName vn -> vn
qualLeaf QualName VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound_here = QualName VName -> DimDecl VName
forall vn. QualName vn -> DimDecl vn
NamedDim QualName VName
v
      | Bool
otherwise = DimDecl VName
forall vn. DimDecl vn
AnyDim
    anyDimIfNotBound DimDecl VName
d = DimDecl VName
d
    rettype_st :: StructType
rettype_st = (DimDecl VName -> DimDecl VName) -> StructType -> StructType
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first DimDecl VName -> DimDecl VName
anyDimIfNotBound (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct PatternType
rettype

    dec :: ValBindBase Info VName
dec =
      ValBind :: forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructType, [VName])
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
ValBind
        { valBindEntryPoint :: Maybe (Info EntryPoint)
valBindEntryPoint = Maybe (Info EntryPoint)
forall a. Maybe a
Nothing,
          valBindName :: VName
valBindName = VName
fname,
          valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
forall a. Maybe a
Nothing,
          valBindRetType :: Info (StructType, [VName])
valBindRetType = (StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info (StructType
rettype_st, []),
          valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams = [TypeParamBase VName]
dims',
          valBindParams :: [PatternBase Info VName]
valBindParams = [PatternBase Info VName]
pats,
          valBindBody :: ExpBase Info VName
valBindBody = ExpBase Info VName
body,
          valBindDoc :: Maybe DocComment
valBindDoc = Maybe DocComment
forall a. Maybe a
Nothing,
          valBindAttrs :: [AttrInfo]
valBindAttrs = [AttrInfo]
forall a. Monoid a => a
mempty,
          valBindLocation :: SrcLoc
valBindLocation = SrcLoc
forall a. Monoid a => a
mempty
        }

-- | 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.
buildEnvPattern :: [VName] -> Env -> Pattern
buildEnvPattern :: [VName] -> Env -> PatternBase Info VName
buildEnvPattern [VName]
sizes Env
env = [(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern (((VName, Binding) -> (Name, PatternBase Info VName))
-> [(VName, Binding)] -> [(Name, PatternBase Info VName)]
forall a b. (a -> b) -> [a] -> [b]
map (VName, Binding) -> (Name, PatternBase Info VName)
buildField ([(VName, Binding)] -> [(Name, PatternBase Info VName)])
-> [(VName, Binding)] -> [(Name, PatternBase Info VName)]
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) -> (Name, PatternBase Info VName)
buildField (VName
vn, Binding Maybe ([VName], StructType)
_ StaticVal
sv) =
      ( String -> Name
nameFromString (VName -> String
forall a. Pretty a => a -> String
pretty VName
vn),
        if VName
vn VName -> [VName] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [VName]
sizes
          then Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv) SrcLoc
forall a. Monoid a => a
mempty
          else VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
vn (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv) SrcLoc
forall a. Monoid a => a
mempty
      )

-- | Given a closure environment pattern and the type of a term,
-- construct the type of that term, where uniqueness is set to
-- `Nonunique` for those arrays that are bound in the environment or
-- pattern (except if they are unique there).  This ensures that a
-- lifted function can create unique arrays as long as they do not
-- alias any of its parameters.  XXX: it is not clear that this is a
-- sufficient property, unfortunately.
buildRetType :: Env -> [Pattern] -> StructType -> PatternType -> PatternType
buildRetType :: Env
-> [PatternBase Info VName]
-> StructType
-> PatternType
-> PatternType
buildRetType Env
env [PatternBase Info VName]
pats = StructType -> PatternType -> PatternType
forall {t :: * -> *} {shape} {as}.
(Foldable t, Monoid (t Alias)) =>
TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb
  where
    bound :: Set VName
bound =
      [VName] -> Set VName
forall a. Ord a => [a] -> Set a
S.fromList (Env -> [VName]
forall k a. Map k a -> [k]
M.keys Env
env) Set VName -> Set VName -> Set VName
forall a. Semigroup a => a -> a -> a
<> (IdentBase Info VName -> VName)
-> Set (IdentBase Info VName) -> Set VName
forall b a. Ord b => (a -> b) -> Set a -> Set b
S.map IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName ((PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> Set (IdentBase Info VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
pats)
    boundAsUnique :: VName -> Bool
boundAsUnique VName
v =
      Bool
-> (IdentBase Info VName -> Bool)
-> Maybe (IdentBase Info VName)
-> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
unique (PatternType -> Bool)
-> (IdentBase Info VName -> PatternType)
-> IdentBase Info VName
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info PatternType -> PatternType
forall a. Info a -> a
unInfo (Info PatternType -> PatternType)
-> (IdentBase Info VName -> Info PatternType)
-> IdentBase Info VName
-> PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase Info VName -> Info PatternType
forall (f :: * -> *) vn. IdentBase f vn -> f PatternType
identType) (Maybe (IdentBase Info VName) -> Bool)
-> Maybe (IdentBase Info VName) -> Bool
forall a b. (a -> b) -> a -> b
$
        (IdentBase Info VName -> Bool)
-> [IdentBase Info VName] -> Maybe (IdentBase Info VName)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((VName -> VName -> Bool
forall a. Eq a => a -> a -> Bool
== VName
v) (VName -> Bool)
-> (IdentBase Info VName -> VName) -> IdentBase Info VName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentBase Info VName -> VName
forall (f :: * -> *) vn. IdentBase f vn -> vn
identName) ([IdentBase Info VName] -> Maybe (IdentBase Info VName))
-> [IdentBase Info VName] -> Maybe (IdentBase Info VName)
forall a b. (a -> b) -> a -> b
$ Set (IdentBase Info VName) -> [IdentBase Info VName]
forall a. Set a -> [a]
S.toList (Set (IdentBase Info VName) -> [IdentBase Info VName])
-> Set (IdentBase Info VName) -> [IdentBase Info VName]
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> Set (IdentBase Info VName))
-> [PatternBase Info VName] -> Set (IdentBase Info VName)
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PatternBase Info VName -> Set (IdentBase Info VName)
forall (f :: * -> *) vn.
(Functor f, Ord vn) =>
PatternBase f vn -> Set (IdentBase f vn)
patternIdents [PatternBase Info VName]
pats
    problematic :: VName -> Bool
problematic VName
v = (VName
v VName -> Set VName -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set VName
bound) Bool -> Bool -> Bool
&& Bool -> Bool
not (VName -> Bool
boundAsUnique VName
v)
    comb :: TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb (Scalar (Record Map Name (TypeBase shape as)
fs_annot)) (Scalar (Record Map Name (TypeBase shape (t Alias))
fs_got)) =
      ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias))
-> ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase shape (t Alias))
-> ScalarTypeBase shape (t Alias)
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase shape (t Alias))
 -> ScalarTypeBase shape (t Alias))
-> Map Name (TypeBase shape (t Alias))
-> ScalarTypeBase shape (t Alias)
forall a b. (a -> b) -> a -> b
$ (TypeBase shape as
 -> TypeBase shape (t Alias) -> TypeBase shape (t Alias))
-> Map Name (TypeBase shape as)
-> Map Name (TypeBase shape (t Alias))
-> Map Name (TypeBase shape (t Alias))
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb Map Name (TypeBase shape as)
fs_annot Map Name (TypeBase shape (t Alias))
fs_got
    comb (Scalar (Sum Map Name [TypeBase shape as]
cs_annot)) (Scalar (Sum Map Name [TypeBase shape (t Alias)]
cs_got)) =
      ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias))
-> ScalarTypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase shape (t Alias)]
-> ScalarTypeBase shape (t Alias)
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase shape (t Alias)]
 -> ScalarTypeBase shape (t Alias))
-> Map Name [TypeBase shape (t Alias)]
-> ScalarTypeBase shape (t Alias)
forall a b. (a -> b) -> a -> b
$ ([TypeBase shape as]
 -> [TypeBase shape (t Alias)] -> [TypeBase shape (t Alias)])
-> Map Name [TypeBase shape as]
-> Map Name [TypeBase shape (t Alias)]
-> Map Name [TypeBase shape (t Alias)]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeBase shape as
 -> TypeBase shape (t Alias) -> TypeBase shape (t Alias))
-> [TypeBase shape as]
-> [TypeBase shape (t Alias)]
-> [TypeBase shape (t Alias)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase shape as
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
comb) Map Name [TypeBase shape as]
cs_annot Map Name [TypeBase shape (t Alias)]
cs_got
    comb (Scalar Arrow {}) TypeBase shape (t Alias)
t =
      TypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend TypeBase shape (t Alias)
t
    comb TypeBase shape as
got TypeBase shape (t Alias)
et =
      TypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall {t :: * -> *} {dim}.
(Foldable t, Monoid (t Alias)) =>
TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend (TypeBase shape (t Alias) -> TypeBase shape (t Alias))
-> TypeBase shape (t Alias) -> TypeBase shape (t Alias)
forall a b. (a -> b) -> a -> b
$ TypeBase shape as -> TypeBase shape (Set Alias)
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct TypeBase shape as
got TypeBase shape (Set Alias) -> t Alias -> TypeBase shape (t Alias)
forall dim asf ast. TypeBase dim asf -> ast -> TypeBase dim ast
`setAliases` TypeBase shape (t Alias) -> t Alias
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase shape (t Alias)
et

    descend :: TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend t :: TypeBase dim (t Alias)
t@Array {}
      | (Alias -> Bool) -> t Alias -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (VName -> Bool
problematic (VName -> Bool) -> (Alias -> VName) -> Alias -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alias -> VName
aliasVar) (TypeBase dim (t Alias) -> t Alias
forall as shape. Monoid as => TypeBase shape as -> as
aliases TypeBase dim (t Alias)
t) = TypeBase dim (t Alias)
t TypeBase dim (t Alias) -> Uniqueness -> TypeBase dim (t Alias)
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
    descend (Scalar (Record Map Name (TypeBase dim (t Alias))
t)) = ScalarTypeBase dim (t Alias) -> TypeBase dim (t Alias)
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim (t Alias) -> TypeBase dim (t Alias))
-> ScalarTypeBase dim (t Alias) -> TypeBase dim (t Alias)
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim (t Alias)) -> ScalarTypeBase dim (t Alias)
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim (t Alias)) -> ScalarTypeBase dim (t Alias))
-> Map Name (TypeBase dim (t Alias))
-> ScalarTypeBase dim (t Alias)
forall a b. (a -> b) -> a -> b
$ (TypeBase dim (t Alias) -> TypeBase dim (t Alias))
-> Map Name (TypeBase dim (t Alias))
-> Map Name (TypeBase dim (t Alias))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TypeBase dim (t Alias) -> TypeBase dim (t Alias)
descend Map Name (TypeBase dim (t Alias))
t
    descend TypeBase dim (t Alias)
t = TypeBase dim (t Alias)
t

-- | Compute the corresponding type for the *representation* of a
-- given static value (not the original possibly higher-order value).
typeFromSV :: StaticVal -> PatternType
typeFromSV :: StaticVal -> PatternType
typeFromSV (Dynamic PatternType
tp) =
  PatternType
tp
typeFromSV (LambdaSV PatternBase Info VName
_ StructType
_ ExtExp
_ Env
env) =
  ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$
    Map Name PatternType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatternType
 -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> Map Name PatternType
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$
      [(Name, PatternType)] -> Map Name PatternType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, PatternType)] -> Map Name PatternType)
-> [(Name, PatternType)] -> Map Name PatternType
forall a b. (a -> b) -> a -> b
$
        ((VName, Binding) -> (Name, PatternType))
-> [(VName, Binding)] -> [(Name, PatternType)]
forall a b. (a -> b) -> [a] -> [b]
map ((VName -> Name)
-> (Binding -> PatternType)
-> (VName, Binding)
-> (Name, PatternType)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (String -> Name
nameFromString (String -> Name) -> (VName -> String) -> VName -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> String
forall a. Pretty a => a -> String
pretty) (StaticVal -> PatternType
typeFromSV (StaticVal -> PatternType)
-> (Binding -> StaticVal) -> Binding -> PatternType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Binding -> StaticVal
bindingSV)) ([(VName, Binding)] -> [(Name, PatternType)])
-> [(VName, Binding)] -> [(Name, PatternType)]
forall a b. (a -> b) -> a -> b
$
          Env -> [(VName, Binding)]
forall k a. Map k a -> [(k, a)]
M.toList Env
env
typeFromSV (RecordSV [(Name, StaticVal)]
ls) =
  let ts :: [(Name, PatternType)]
ts = ((Name, StaticVal) -> (Name, PatternType))
-> [(Name, StaticVal)] -> [(Name, PatternType)]
forall a b. (a -> b) -> [a] -> [b]
map ((StaticVal -> PatternType)
-> (Name, StaticVal) -> (Name, PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap StaticVal -> PatternType
typeFromSV) [(Name, StaticVal)]
ls
   in ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ Map Name PatternType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name PatternType
 -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> Map Name PatternType
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$ [(Name, PatternType)] -> Map Name PatternType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, PatternType)]
ts
typeFromSV (DynamicFun (ExpBase Info VName
_, StaticVal
sv) StaticVal
_) =
  StaticVal -> PatternType
typeFromSV StaticVal
sv
typeFromSV (SumSV Name
name [StaticVal]
svs [(Name, [PatternType])]
fields) =
  let svs' :: [PatternType]
svs' = (StaticVal -> PatternType) -> [StaticVal] -> [PatternType]
forall a b. (a -> b) -> [a] -> [b]
map StaticVal -> PatternType
typeFromSV [StaticVal]
svs
   in ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ Map Name [PatternType]
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [PatternType]
 -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> Map Name [PatternType]
-> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$ Name
-> [PatternType]
-> Map Name [PatternType]
-> Map Name [PatternType]
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
name [PatternType]
svs' (Map Name [PatternType] -> Map Name [PatternType])
-> Map Name [PatternType] -> Map Name [PatternType]
forall a b. (a -> b) -> a -> b
$ [(Name, [PatternType])] -> Map Name [PatternType]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name, [PatternType])]
fields
typeFromSV StaticVal
IntrinsicSV =
  String -> PatternType
forall a. HasCallStack => String -> a
error String
"Tried to get the type from the static value of an intrinsic."

-- | Construct the type for a fully-applied dynamic function from its
-- static value and the original types of its arguments.
dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType (DynamicFun (ExpBase Info VName, StaticVal)
_ StaticVal
sv) (PatternType
p : [PatternType]
ps) =
  let ([PatternType]
ps', PatternType
ret) = StaticVal -> [PatternType] -> ([PatternType], PatternType)
dynamicFunType StaticVal
sv [PatternType]
ps in (PatternType
p PatternType -> [PatternType] -> [PatternType]
forall a. a -> [a] -> [a]
: [PatternType]
ps', PatternType
ret)
dynamicFunType StaticVal
sv [PatternType]
_ = ([], StaticVal -> PatternType
typeFromSV 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.
matchPatternSV :: PatternBase Info VName -> StaticVal -> Env
matchPatternSV :: PatternBase Info VName -> StaticVal -> Env
matchPatternSV (TuplePattern [PatternBase Info VName]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls) =
  [Env] -> Env
forall a. Monoid a => [a] -> a
mconcat ([Env] -> Env) -> [Env] -> Env
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> (Name, StaticVal) -> Env)
-> [PatternBase Info VName] -> [(Name, StaticVal)] -> [Env]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\PatternBase Info VName
p (Name
_, StaticVal
sv) -> PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
p StaticVal
sv) [PatternBase Info VName]
ps [(Name, StaticVal)]
ls
matchPatternSV (RecordPattern [(Name, PatternBase Info VName)]
ps SrcLoc
_) (RecordSV [(Name, StaticVal)]
ls)
  | [(Name, PatternBase Info VName)]
ps' <- ((Name, PatternBase Info VName) -> Name)
-> [(Name, PatternBase Info VName)]
-> [(Name, PatternBase Info VName)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, PatternBase Info VName) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternBase Info VName)]
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, PatternBase Info VName) -> Name)
-> [(Name, PatternBase Info VName)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, PatternBase Info VName) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternBase Info VName)]
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] -> Env
forall a. Monoid a => [a] -> a
mconcat ([Env] -> Env) -> [Env] -> Env
forall a b. (a -> b) -> a -> b
$ ((Name, PatternBase Info VName) -> (Name, StaticVal) -> Env)
-> [(Name, PatternBase Info VName)] -> [(Name, StaticVal)] -> [Env]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(Name
_, PatternBase Info VName
p) (Name
_, StaticVal
sv) -> PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
p StaticVal
sv) [(Name, PatternBase Info VName)]
ps' [(Name, StaticVal)]
ls'
matchPatternSV (PatternParens PatternBase Info VName
pat SrcLoc
_) StaticVal
sv = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat StaticVal
sv
matchPatternSV (Id VName
vn (Info PatternType
t) SrcLoc
_) StaticVal
sv =
  -- When matching a pattern with a zero-order STaticVal, the type of
  -- the pattern wins out.  This is important when matching a
  -- nonunique pattern with a unique value.
  if StaticVal -> Bool
orderZeroSV StaticVal
sv
    then Env
dim_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton 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
$ PatternType -> StaticVal
Dynamic PatternType
t)
    else Env
dim_env Env -> Env -> Env
forall a. Semigroup a => a -> a -> a
<> VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton VName
vn (Maybe ([VName], StructType) -> StaticVal -> Binding
Binding Maybe ([VName], StructType)
forall a. Maybe a
Nothing StaticVal
sv)
  where
    dim_env :: Env
dim_env =
      [(VName, Binding)] -> Env
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, Binding)] -> Env) -> [(VName, Binding)] -> Env
forall a b. (a -> b) -> a -> b
$ (VName -> (VName, Binding)) -> [VName] -> [(VName, Binding)]
forall a b. (a -> b) -> [a] -> [b]
map (,Binding
i64) ([VName] -> [(VName, Binding)]) -> [VName] -> [(VName, Binding)]
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
$ PatternType -> Set VName
forall als. TypeBase (DimDecl VName) als -> Set VName
typeDimNames PatternType
t
    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
$ PatternType -> StaticVal
Dynamic (PatternType -> StaticVal) -> PatternType -> StaticVal
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType)
-> ScalarTypeBase (DimDecl VName) (Set Alias) -> PatternType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall dim as. PrimType -> ScalarTypeBase dim as
Prim (PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias))
-> PrimType -> ScalarTypeBase (DimDecl VName) (Set Alias)
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64
matchPatternSV (Wildcard Info PatternType
_ SrcLoc
_) StaticVal
_ = Env
forall a. Monoid a => a
mempty
matchPatternSV (PatternAscription PatternBase Info VName
pat TypeDeclBase Info VName
_ SrcLoc
_) StaticVal
sv = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat StaticVal
sv
matchPatternSV PatternLit {} StaticVal
_ = Env
forall a. Monoid a => a
mempty
matchPatternSV (PatternConstr Name
c1 Info PatternType
_ [PatternBase Info VName]
ps SrcLoc
_) (SumSV Name
c2 [StaticVal]
ls [(Name, [PatternType])]
fs)
  | Name
c1 Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
c2 =
    [Env] -> Env
forall a. Monoid a => [a] -> a
mconcat ([Env] -> Env) -> [Env] -> Env
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> StaticVal -> Env)
-> [PatternBase Info VName] -> [StaticVal] -> [Env]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatternBase Info VName -> StaticVal -> Env
matchPatternSV [PatternBase Info VName]
ps [StaticVal]
ls
  | Just [PatternType]
ts <- Name -> [(Name, [PatternType])] -> Maybe [PatternType]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
c1 [(Name, [PatternType])]
fs =
    [Env] -> Env
forall a. Monoid a => [a] -> a
mconcat ([Env] -> Env) -> [Env] -> Env
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> StaticVal -> Env)
-> [PatternBase Info VName] -> [StaticVal] -> [Env]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatternBase Info VName -> StaticVal -> Env
matchPatternSV [PatternBase Info VName]
ps ([StaticVal] -> [Env]) -> [StaticVal] -> [Env]
forall a b. (a -> b) -> a -> b
$ (PatternType -> StaticVal) -> [PatternType] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StaticVal
svFromType [PatternType]
ts
  | Bool
otherwise =
    String -> Env
forall a. HasCallStack => String -> a
error (String -> Env) -> String -> Env
forall a b. (a -> b) -> a -> b
$ String
"matchPatternSV: missing constructor in type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
c1
matchPatternSV (PatternConstr Name
c1 Info PatternType
_ [PatternBase Info VName]
ps SrcLoc
_) (Dynamic (Scalar (Sum Map Name [PatternType]
fs)))
  | Just [PatternType]
ts <- Name -> Map Name [PatternType] -> Maybe [PatternType]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
c1 Map Name [PatternType]
fs =
    [Env] -> Env
forall a. Monoid a => [a] -> a
mconcat ([Env] -> Env) -> [Env] -> Env
forall a b. (a -> b) -> a -> b
$ (PatternBase Info VName -> StaticVal -> Env)
-> [PatternBase Info VName] -> [StaticVal] -> [Env]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatternBase Info VName -> StaticVal -> Env
matchPatternSV [PatternBase Info VName]
ps ([StaticVal] -> [Env]) -> [StaticVal] -> [Env]
forall a b. (a -> b) -> a -> b
$ (PatternType -> StaticVal) -> [PatternType] -> [StaticVal]
forall a b. (a -> b) -> [a] -> [b]
map PatternType -> StaticVal
svFromType [PatternType]
ts
  | Bool
otherwise =
    String -> Env
forall a. HasCallStack => String -> a
error (String -> Env) -> String -> Env
forall a b. (a -> b) -> a -> b
$ String
"matchPatternSV: missing constructor in type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Pretty a => a -> String
pretty Name
c1
matchPatternSV PatternBase Info VName
pat (Dynamic PatternType
t) = PatternBase Info VName -> StaticVal -> Env
matchPatternSV PatternBase Info VName
pat (StaticVal -> Env) -> StaticVal -> Env
forall a b. (a -> b) -> a -> b
$ PatternType -> StaticVal
svFromType PatternType
t
matchPatternSV PatternBase Info VName
pat StaticVal
sv =
  String -> Env
forall a. HasCallStack => String -> a
error (String -> Env) -> String -> Env
forall a b. (a -> b) -> a -> b
$
    String
"Tried to match pattern " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatternBase Info VName -> String
forall a. Pretty a => a -> String
pretty PatternBase Info VName
pat
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" with static value "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> String
forall a. Show a => a -> String
show StaticVal
sv
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"."

orderZeroSV :: StaticVal -> Bool
orderZeroSV :: StaticVal -> Bool
orderZeroSV Dynamic {} = Bool
True
orderZeroSV (RecordSV [(Name, StaticVal)]
fields) = ((Name, StaticVal) -> Bool) -> [(Name, StaticVal)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (StaticVal -> Bool
orderZeroSV (StaticVal -> Bool)
-> ((Name, StaticVal) -> StaticVal) -> (Name, StaticVal) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, StaticVal) -> StaticVal
forall a b. (a, b) -> b
snd) [(Name, StaticVal)]
fields
orderZeroSV StaticVal
_ = Bool
False

-- | Given a pattern and the static value for the defunctionalized argument,
-- update the pattern to reflect the changes in the types.
updatePattern :: Pattern -> StaticVal -> Pattern
updatePattern :: PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern (TuplePattern [PatternBase Info VName]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs) =
  [PatternBase Info VName] -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[PatternBase f vn] -> SrcLoc -> PatternBase f vn
TuplePattern ((PatternBase Info VName -> StaticVal -> PatternBase Info VName)
-> [PatternBase Info VName]
-> [StaticVal]
-> [PatternBase Info VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern [PatternBase Info VName]
ps ([StaticVal] -> [PatternBase Info VName])
-> [StaticVal] -> [PatternBase Info VName]
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
updatePattern (RecordPattern [(Name, PatternBase Info VName)]
ps SrcLoc
loc) (RecordSV [(Name, StaticVal)]
svs)
  | [(Name, PatternBase Info VName)]
ps' <- ((Name, PatternBase Info VName) -> Name)
-> [(Name, PatternBase Info VName)]
-> [(Name, PatternBase Info VName)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Name, PatternBase Info VName) -> Name
forall a b. (a, b) -> a
fst [(Name, PatternBase Info VName)]
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 =
    [(Name, PatternBase Info VName)]
-> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
[(Name, PatternBase f vn)] -> SrcLoc -> PatternBase f vn
RecordPattern
      ( ((Name, PatternBase Info VName)
 -> (Name, StaticVal) -> (Name, PatternBase Info VName))
-> [(Name, PatternBase Info VName)]
-> [(Name, StaticVal)]
-> [(Name, PatternBase Info VName)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          ( \(Name
n, PatternBase Info VName
p) (Name
_, StaticVal
sv) ->
              (Name
n, PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
p StaticVal
sv)
          )
          [(Name, PatternBase Info VName)]
ps'
          [(Name, StaticVal)]
svs'
      )
      SrcLoc
loc
updatePattern (PatternParens PatternBase Info VName
pat SrcLoc
loc) StaticVal
sv =
  PatternBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> SrcLoc -> PatternBase f vn
PatternParens (PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat StaticVal
sv) SrcLoc
loc
updatePattern (Id VName
vn (Info PatternType
tp) SrcLoc
loc) StaticVal
sv =
  VName -> Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
vn -> f PatternType -> SrcLoc -> PatternBase f vn
Id VName
vn (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ PatternType -> PatternType -> PatternType
forall {dim} {as}.
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb PatternType
tp (StaticVal -> PatternType
typeFromSV StaticVal
sv PatternType -> Uniqueness -> PatternType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique)) SrcLoc
loc
  where
    -- Preserve any original zeroth-order types.
    comb :: TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb (Scalar Arrow {}) TypeBase dim as
t2 = TypeBase dim as
t2
    comb (Scalar (Record Map Name (TypeBase dim as)
m1)) (Scalar (Record Map Name (TypeBase dim as)
m2)) =
      ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall dim as. Map Name (TypeBase dim as) -> ScalarTypeBase dim as
Record (Map Name (TypeBase dim as) -> ScalarTypeBase dim as)
-> Map Name (TypeBase dim as) -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ (TypeBase dim as -> TypeBase dim as -> TypeBase dim as)
-> Map Name (TypeBase dim as)
-> Map Name (TypeBase dim as)
-> Map Name (TypeBase dim as)
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb Map Name (TypeBase dim as)
m1 Map Name (TypeBase dim as)
m2
    comb (Scalar (Sum Map Name [TypeBase dim as]
m1)) (Scalar (Sum Map Name [TypeBase dim as]
m2)) =
      ScalarTypeBase dim as -> TypeBase dim as
forall dim as. ScalarTypeBase dim as -> TypeBase dim as
Scalar (ScalarTypeBase dim as -> TypeBase dim as)
-> ScalarTypeBase dim as -> TypeBase dim as
forall a b. (a -> b) -> a -> b
$ Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall dim as. Map Name [TypeBase dim as] -> ScalarTypeBase dim as
Sum (Map Name [TypeBase dim as] -> ScalarTypeBase dim as)
-> Map Name [TypeBase dim as] -> ScalarTypeBase dim as
forall a b. (a -> b) -> a -> b
$ ([TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as])
-> Map Name [TypeBase dim as]
-> Map Name [TypeBase dim as]
-> Map Name [TypeBase dim as]
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith ((TypeBase dim as -> TypeBase dim as -> TypeBase dim as)
-> [TypeBase dim as] -> [TypeBase dim as] -> [TypeBase dim as]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith TypeBase dim as -> TypeBase dim as -> TypeBase dim as
comb) Map Name [TypeBase dim as]
m1 Map Name [TypeBase dim as]
m2
    comb TypeBase dim as
t1 TypeBase dim as
_ = TypeBase dim as
t1 -- t1 must be array or prim.
updatePattern pat :: PatternBase Info VName
pat@(Wildcard (Info PatternType
tp) SrcLoc
loc) StaticVal
sv
  | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
tp = PatternBase Info VName
pat
  | Bool
otherwise = Info PatternType -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
f PatternType -> SrcLoc -> PatternBase f vn
Wildcard (PatternType -> Info PatternType
forall a. a -> Info a
Info (PatternType -> Info PatternType)
-> PatternType -> Info PatternType
forall a b. (a -> b) -> a -> b
$ StaticVal -> PatternType
typeFromSV StaticVal
sv) SrcLoc
loc
updatePattern (PatternAscription PatternBase Info VName
pat TypeDeclBase Info VName
tydecl SrcLoc
loc) StaticVal
sv
  | StructType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero (StructType -> Bool)
-> (Info StructType -> StructType) -> Info StructType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Info StructType -> StructType
forall a. Info a -> a
unInfo (Info StructType -> Bool) -> Info StructType -> Bool
forall a b. (a -> b) -> a -> b
$ TypeDeclBase Info VName -> Info StructType
forall (f :: * -> *) vn. TypeDeclBase f vn -> f StructType
expandedType TypeDeclBase Info VName
tydecl =
    PatternBase Info VName
-> TypeDeclBase Info VName -> SrcLoc -> PatternBase Info VName
forall (f :: * -> *) vn.
PatternBase f vn -> TypeDeclBase f vn -> SrcLoc -> PatternBase f vn
PatternAscription (PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat StaticVal
sv) TypeDeclBase Info VName
tydecl SrcLoc
loc
  | Bool
otherwise = PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat StaticVal
sv
updatePattern p :: PatternBase Info VName
p@PatternLit {} StaticVal
_ = PatternBase Info VName
p
updatePattern pat :: PatternBase Info VName
pat@(PatternConstr Name
c1 (Info PatternType
t) [PatternBase Info VName]
ps SrcLoc
loc) sv :: StaticVal
sv@(SumSV Name
_ [StaticVal]
svs [(Name, [PatternType])]
_)
  | PatternType -> Bool
forall dim as. TypeBase dim as -> Bool
orderZero PatternType
t = PatternBase Info VName
pat
  | Bool
otherwise = Name
-> Info PatternType
-> [PatternBase Info VName]
-> SrcLoc
-> PatternBase Info VName
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
c1 (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t') [PatternBase Info VName]
ps' SrcLoc
loc
  where
    t' :: PatternType
t' = StaticVal -> PatternType
typeFromSV StaticVal
sv PatternType -> Uniqueness -> PatternType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
    ps' :: [PatternBase Info VName]
ps' = (PatternBase Info VName -> StaticVal -> PatternBase Info VName)
-> [PatternBase Info VName]
-> [StaticVal]
-> [PatternBase Info VName]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern [PatternBase Info VName]
ps [StaticVal]
svs
updatePattern (PatternConstr Name
c1 Info PatternType
_ [PatternBase Info VName]
ps SrcLoc
loc) (Dynamic PatternType
t) =
  Name
-> Info PatternType
-> [PatternBase Info VName]
-> SrcLoc
-> PatternBase Info VName
forall (f :: * -> *) vn.
Name
-> f PatternType
-> [PatternBase f vn]
-> SrcLoc
-> PatternBase f vn
PatternConstr Name
c1 (PatternType -> Info PatternType
forall a. a -> Info a
Info PatternType
t) [PatternBase Info VName]
ps SrcLoc
loc
updatePattern PatternBase Info VName
pat (Dynamic PatternType
t) = PatternBase Info VName -> StaticVal -> PatternBase Info VName
updatePattern PatternBase Info VName
pat (PatternType -> StaticVal
svFromType PatternType
t)
updatePattern PatternBase Info VName
pat StaticVal
sv =
  String -> PatternBase Info VName
forall a. HasCallStack => String -> a
error (String -> PatternBase Info VName)
-> String -> PatternBase Info VName
forall a b. (a -> b) -> a -> b
$
    String
"Tried to update pattern " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PatternBase Info VName -> String
forall a. Pretty a => a -> String
pretty PatternBase Info VName
pat
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"to reflect the static value "
      String -> ShowS
forall a. [a] -> [a] -> [a]
++ StaticVal -> String
forall a. Show a => a -> String
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 :: PatternType -> StaticVal
svFromType :: PatternType -> StaticVal
svFromType (Scalar (Record Map Name PatternType
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
$ (PatternType -> StaticVal)
-> Map Name PatternType -> Map Name StaticVal
forall a b k. (a -> b) -> Map k a -> Map k b
M.map PatternType -> StaticVal
svFromType Map Name PatternType
fs
svFromType PatternType
t = PatternType -> StaticVal
Dynamic PatternType
t

-- | 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.  The
-- boolean is true if the function is a 'DynamicFun'.
defuncValBind :: ValBind -> DefM (ValBind, Env, Bool)
-- Eta-expand entry points with a functional return type.
defuncValBind :: ValBindBase Info VName -> DefM (ValBindBase Info VName, Env, Bool)
defuncValBind (ValBind Maybe (Info EntryPoint)
entry VName
name Maybe (TypeExp VName)
_ (Info (StructType
rettype, [VName]
retext)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
attrs SrcLoc
loc)
  | Scalar Arrow {} <- StructType
rettype = do
    ([PatternBase Info VName]
body_pats, ExpBase Info VName
body', StructType
rettype') <- PatternType
-> ExpBase Info VName
-> DefM ([PatternBase Info VName], ExpBase Info VName, StructType)
etaExpand (StructType -> PatternType
forall dim as. TypeBase dim as -> TypeBase dim (Set Alias)
fromStruct StructType
rettype) ExpBase Info VName
body
    ValBindBase Info VName -> DefM (ValBindBase Info VName, Env, Bool)
defuncValBind (ValBindBase Info VName
 -> DefM (ValBindBase Info VName, Env, Bool))
-> ValBindBase Info VName
-> DefM (ValBindBase Info VName, Env, Bool)
forall a b. (a -> b) -> a -> b
$
      Maybe (Info EntryPoint)
-> VName
-> Maybe (TypeExp VName)
-> Info (StructType, [VName])
-> [TypeParamBase VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase Info VName
forall (f :: * -> *) vn.
Maybe (f EntryPoint)
-> vn
-> Maybe (TypeExp vn)
-> f (StructType, [VName])
-> [TypeParamBase vn]
-> [PatternBase f vn]
-> ExpBase f vn
-> Maybe DocComment
-> [AttrInfo]
-> SrcLoc
-> ValBindBase f vn
ValBind
        Maybe (Info EntryPoint)
entry
        VName
name
        Maybe (TypeExp VName)
forall a. Maybe a
Nothing
        ((StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info (StructType
rettype', [VName]
retext))
        [TypeParamBase VName]
tparams
        ([PatternBase Info VName]
params [PatternBase Info VName]
-> [PatternBase Info VName] -> [PatternBase Info VName]
forall a. Semigroup a => a -> a -> a
<> [PatternBase Info VName]
body_pats)
        ExpBase Info VName
body'
        Maybe DocComment
forall a. Maybe a
Nothing
        [AttrInfo]
attrs
        SrcLoc
loc
defuncValBind valbind :: ValBindBase Info VName
valbind@(ValBind Maybe (Info EntryPoint)
_ VName
name Maybe (TypeExp VName)
retdecl (Info (StructType
rettype, [VName]
retext)) [TypeParamBase VName]
tparams [PatternBase Info VName]
params ExpBase Info VName
body Maybe DocComment
_ [AttrInfo]
_ 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
$
    String -> DefM ()
forall a. HasCallStack => String -> a
error (String -> DefM ()) -> String -> DefM ()
forall a b. (a -> b) -> a -> b
$
      VName -> String
forall v. IsName v => v -> String
prettyName VName
name String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" has type parameters, "
        String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"but the defunctionaliser expects a monomorphic input program."
  ([VName]
tparams', [PatternBase Info VName]
params', ExpBase Info VName
body', StaticVal
sv) <-
    [VName]
-> [PatternBase Info VName]
-> ExpBase Info VName
-> StructType
-> DefM
     ([VName], [PatternBase Info VName], ExpBase Info VName, StaticVal)
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) [PatternBase Info VName]
params ExpBase Info VName
body StructType
rettype
  let rettype' :: StructType
rettype' = StructType -> StructType -> StructType
forall as dim.
(Monoid as, ArrayDim dim) =>
TypeBase dim as -> TypeBase dim as -> TypeBase dim as
combineTypeShapes StructType
rettype (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ StructType -> StructType
forall vn as. TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as
anySizes (StructType -> StructType) -> StructType -> StructType
forall a b. (a -> b) -> a -> b
$ PatternType -> StructType
forall dim as. TypeBase dim as -> TypeBase dim ()
toStruct (PatternType -> StructType) -> PatternType -> StructType
forall a b. (a -> b) -> a -> b
$ ExpBase Info VName -> PatternType
typeOf ExpBase Info VName
body'
  Set VName
globals <- ((Set VName, Env) -> Set VName) -> DefM (Set VName)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Set VName, Env) -> Set VName
forall a b. (a, b) -> a
fst
  let bound_sizes :: Set VName
bound_sizes = [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
  ([VName]
missing_dims, [PatternBase Info VName]
params'') <- Set VName
-> [PatternBase Info VName]
-> DefM ([VName], [PatternBase Info VName])
forall (m :: * -> *).
MonadFreshNames m =>
Set VName
-> [PatternBase Info VName]
-> m ([VName], [PatternBase Info VName])
sizesForAll Set VName
bound_sizes [PatternBase Info VName]
params'
  (ValBindBase Info VName, Env, Bool)
-> DefM (ValBindBase Info VName, Env, Bool)
forall (m :: * -> *) a. Monad m => a -> m a
return
    ( ValBindBase Info VName
valbind
        { valBindRetDecl :: Maybe (TypeExp VName)
valBindRetDecl = Maybe (TypeExp VName)
retdecl,
          valBindRetType :: Info (StructType, [VName])
valBindRetType =
            (StructType, [VName]) -> Info (StructType, [VName])
forall a. a -> Info a
Info
              ( if [PatternBase Info VName] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PatternBase Info VName]
params'
                  then StructType
rettype' StructType -> Uniqueness -> StructType
forall dim as. TypeBase dim as -> Uniqueness -> TypeBase dim as
`setUniqueness` Uniqueness
Nonunique
                  else StructType
rettype',
                [VName]
retext
              ),
          valBindTypeParams :: [TypeParamBase VName]
valBindTypeParams =
            (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] -> [TypeParamBase VName])
-> [VName] -> [TypeParamBase VName]
forall a b. (a -> b) -> a -> b
$ [VName]
tparams' [VName] -> [VName] -> [VName]
forall a. [a] -> [a] -> [a]
++ [VName]
missing_dims,
          valBindParams :: [PatternBase Info VName]
valBindParams = [PatternBase Info VName]
params'',
          valBindBody :: ExpBase Info VName
valBindBody = ExpBase Info VName
body'
        },
      VName -> Binding -> Env
forall k a. k -> a -> Map k a
M.singleton VName
name (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 (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)
                  (ValBindBase Info VName -> ([TypeParamBase VName], StructType)
valBindTypeScheme ValBindBase Info VName
valbind)
              )
          )
          StaticVal
sv,
      case StaticVal
sv of
        DynamicFun {} -> Bool
True
        Dynamic {} -> Bool
True
        StaticVal
_ -> Bool
False
    )

-- | Defunctionalize a list of top-level declarations.
defuncVals :: [ValBind] -> DefM (Seq.Seq ValBind)
defuncVals :: [ValBindBase Info VName] -> DefM (Seq (ValBindBase Info VName))
defuncVals [] = Seq (ValBindBase Info VName) -> DefM (Seq (ValBindBase Info VName))
forall (m :: * -> *) a. Monad m => a -> m a
return Seq (ValBindBase Info VName)
forall a. Monoid a => a
mempty
defuncVals (ValBindBase Info VName
valbind : [ValBindBase Info VName]
ds) = do
  ((ValBindBase Info VName
valbind', Env
env, Bool
dyn), Seq (ValBindBase Info VName)
defs) <- DefM (ValBindBase Info VName, Env, Bool)
-> DefM
     ((ValBindBase Info VName, Env, Bool), Seq (ValBindBase Info VName))
forall a. DefM a -> DefM (a, Seq (ValBindBase Info VName))
collectFuns (DefM (ValBindBase Info VName, Env, Bool)
 -> DefM
      ((ValBindBase Info VName, Env, Bool),
       Seq (ValBindBase Info VName)))
-> DefM (ValBindBase Info VName, Env, Bool)
-> DefM
     ((ValBindBase Info VName, Env, Bool), Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$ ValBindBase Info VName -> DefM (ValBindBase Info VName, Env, Bool)
defuncValBind ValBindBase Info VName
valbind
  Seq (ValBindBase Info VName)
ds' <-
    Env
-> DefM (Seq (ValBindBase Info VName))
-> DefM (Seq (ValBindBase Info VName))
forall a. Env -> DefM a -> DefM a
localEnv Env
env (DefM (Seq (ValBindBase Info VName))
 -> DefM (Seq (ValBindBase Info VName)))
-> DefM (Seq (ValBindBase Info VName))
-> DefM (Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$
      if Bool
dyn
        then VName
-> DefM (Seq (ValBindBase Info VName))
-> DefM (Seq (ValBindBase Info VName))
forall a. VName -> DefM a -> DefM a
isGlobal (ValBindBase Info VName -> VName
forall (f :: * -> *) vn. ValBindBase f vn -> vn
valBindName ValBindBase Info VName
valbind') (DefM (Seq (ValBindBase Info VName))
 -> DefM (Seq (ValBindBase Info VName)))
-> DefM (Seq (ValBindBase Info VName))
-> DefM (Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$ [ValBindBase Info VName] -> DefM (Seq (ValBindBase Info VName))
defuncVals [ValBindBase Info VName]
ds
        else [ValBindBase Info VName] -> DefM (Seq (ValBindBase Info VName))
defuncVals [ValBindBase Info VName]
ds
  Seq (ValBindBase Info VName) -> DefM (Seq (ValBindBase Info VName))
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq (ValBindBase Info VName)
 -> DefM (Seq (ValBindBase Info VName)))
-> Seq (ValBindBase Info VName)
-> DefM (Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$ Seq (ValBindBase Info VName)
defs Seq (ValBindBase Info VName)
-> Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName)
forall a. Semigroup a => a -> a -> a
<> ValBindBase Info VName -> Seq (ValBindBase Info VName)
forall a. a -> Seq a
Seq.singleton ValBindBase Info VName
valbind' Seq (ValBindBase Info VName)
-> Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName)
forall a. Semigroup a => a -> a -> a
<> Seq (ValBindBase Info VName)
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 =>
[ValBindBase Info VName] -> m [ValBindBase Info VName]
transformProg [ValBindBase Info VName]
decs = (VNameSource -> ([ValBindBase Info VName], VNameSource))
-> m [ValBindBase Info VName]
forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource ((VNameSource -> ([ValBindBase Info VName], VNameSource))
 -> m [ValBindBase Info VName])
-> (VNameSource -> ([ValBindBase Info VName], VNameSource))
-> m [ValBindBase Info VName]
forall a b. (a -> b) -> a -> b
$ \VNameSource
namesrc ->
  let (Seq (ValBindBase Info VName)
decs', VNameSource
namesrc', Seq (ValBindBase Info VName)
liftedDecs) = VNameSource
-> DefM (Seq (ValBindBase Info VName))
-> (Seq (ValBindBase Info VName), VNameSource,
    Seq (ValBindBase Info VName))
forall a.
VNameSource
-> DefM a -> (a, VNameSource, Seq (ValBindBase Info VName))
runDefM VNameSource
namesrc (DefM (Seq (ValBindBase Info VName))
 -> (Seq (ValBindBase Info VName), VNameSource,
     Seq (ValBindBase Info VName)))
-> DefM (Seq (ValBindBase Info VName))
-> (Seq (ValBindBase Info VName), VNameSource,
    Seq (ValBindBase Info VName))
forall a b. (a -> b) -> a -> b
$ [ValBindBase Info VName] -> DefM (Seq (ValBindBase Info VName))
defuncVals [ValBindBase Info VName]
decs
   in (Seq (ValBindBase Info VName) -> [ValBindBase Info VName]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq (ValBindBase Info VName) -> [ValBindBase Info VName])
-> Seq (ValBindBase Info VName) -> [ValBindBase Info VName]
forall a b. (a -> b) -> a -> b
$ Seq (ValBindBase Info VName)
liftedDecs Seq (ValBindBase Info VName)
-> Seq (ValBindBase Info VName) -> Seq (ValBindBase Info VName)
forall a. Semigroup a => a -> a -> a
<> Seq (ValBindBase Info VName)
decs', VNameSource
namesrc')