{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Futhark.IR.Prop.Scope
( HasScope (..),
NameInfo (..),
LocalScope (..),
Scope,
Scoped (..),
inScopeOf,
scopeOfLParams,
scopeOfFParams,
scopeOfPat,
scopeOfPatElem,
SameScope,
castScope,
ExtendedScope,
extendedScope,
)
where
import Control.Monad.Except
import Control.Monad.RWS.Lazy qualified
import Control.Monad.RWS.Strict qualified
import Control.Monad.Reader
import Data.Map.Strict qualified as M
import Futhark.IR.Pretty ()
import Futhark.IR.Prop.Types
import Futhark.IR.Rep
import Futhark.IR.Syntax
data NameInfo rep
= LetName (LetDec rep)
| FParamName (FParamInfo rep)
| LParamName (LParamInfo rep)
| IndexName IntType
deriving instance RepTypes rep => Show (NameInfo rep)
instance RepTypes rep => Typed (NameInfo rep) where
typeOf :: NameInfo rep -> Type
typeOf (LetName LetDec rep
dec) = forall t. Typed t => t -> Type
typeOf LetDec rep
dec
typeOf (FParamName FParamInfo rep
dec) = forall t. Typed t => t -> Type
typeOf FParamInfo rep
dec
typeOf (LParamName LParamInfo rep
dec) = forall t. Typed t => t -> Type
typeOf LParamInfo rep
dec
typeOf (IndexName IntType
it) = forall shape u. PrimType -> TypeBase shape u
Prim forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it
type Scope rep = M.Map VName (NameInfo rep)
class (Applicative m, RepTypes rep) => HasScope rep m | m -> rep where
lookupType :: VName -> m Type
lookupType = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall t. Typed t => t -> Type
typeOf forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (NameInfo rep)
lookupInfo
lookupInfo :: VName -> m (NameInfo rep)
lookupInfo VName
name =
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope (forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault NameInfo rep
notFound VName
name)
where
notFound :: NameInfo rep
notFound =
forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$
String
"Scope.lookupInfo: Name "
forall a. [a] -> [a] -> [a]
++ forall a. Pretty a => a -> String
prettyString VName
name
forall a. [a] -> [a] -> [a]
++ String
" not found in type environment."
askScope :: m (Scope rep)
asksScope :: (Scope rep -> a) -> m a
asksScope Scope rep -> a
f = Scope rep -> a
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
instance
(Monad m, RepTypes rep) =>
HasScope rep (ReaderT (Scope rep) m)
where
askScope :: ReaderT (Scope rep) m (Scope rep)
askScope = forall r (m :: * -> *). MonadReader r m => m r
ask
instance (Monad m, HasScope rep m) => HasScope rep (ExceptT e m) where
askScope :: ExceptT e m (Scope rep)
askScope = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
instance
(Monad m, Monoid w, RepTypes rep) =>
HasScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m)
where
askScope :: RWST (Scope rep) w s m (Scope rep)
askScope = forall r (m :: * -> *). MonadReader r m => m r
ask
instance
(Monad m, Monoid w, RepTypes rep) =>
HasScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m)
where
askScope :: RWST (Scope rep) w s m (Scope rep)
askScope = forall r (m :: * -> *). MonadReader r m => m r
ask
class (HasScope rep m, Monad m) => LocalScope rep m where
localScope :: Scope rep -> m a -> m a
instance (LocalScope rep m) => LocalScope rep (ExceptT e m) where
localScope :: forall a. Scope rep -> ExceptT e m a -> ExceptT e m a
localScope = forall (m :: * -> *) e a (n :: * -> *) e' b.
(m (Either e a) -> n (Either e' b))
-> ExceptT e m a -> ExceptT e' n b
mapExceptT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope
instance
(Monad m, RepTypes rep) =>
LocalScope rep (ReaderT (Scope rep) m)
where
localScope :: forall a.
Scope rep -> ReaderT (Scope rep) m a -> ReaderT (Scope rep) m a
localScope = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union
instance
(Monad m, Monoid w, RepTypes rep) =>
LocalScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m)
where
localScope :: forall a.
Scope rep -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
localScope = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union
instance
(Monad m, Monoid w, RepTypes rep) =>
LocalScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m)
where
localScope :: forall a.
Scope rep -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
localScope = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union
class Scoped rep a | a -> rep where
scopeOf :: a -> Scope rep
inScopeOf :: (Scoped rep a, LocalScope rep m) => a -> m b -> m b
inScopeOf :: forall rep a (m :: * -> *) b.
(Scoped rep a, LocalScope rep m) =>
a -> m b -> m b
inScopeOf = forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep a. Scoped rep a => a -> Scope rep
scopeOf
instance Scoped rep a => Scoped rep [a] where
scopeOf :: [a] -> Scope rep
scopeOf = forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall rep a. Scoped rep a => a -> Scope rep
scopeOf
instance Scoped rep (Stms rep) where
scopeOf :: Stms rep -> Scope rep
scopeOf = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall rep a. Scoped rep a => a -> Scope rep
scopeOf
instance Scoped rep (Stm rep) where
scopeOf :: Stm rep -> Scope rep
scopeOf = forall rep dec. (LetDec rep ~ dec) => Pat dec -> Scope rep
scopeOfPat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. Stm rep -> Pat (LetDec rep)
stmPat
instance Scoped rep (FunDef rep) where
scopeOf :: FunDef rep -> Scope rep
scopeOf = forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall rep. FunDef rep -> [FParam rep]
funDefParams
instance Scoped rep (VName, NameInfo rep) where
scopeOf :: (VName, NameInfo rep) -> Scope rep
scopeOf = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall k a. k -> a -> Map k a
M.singleton
instance Scoped rep (LoopForm rep) where
scopeOf :: LoopForm rep -> Scope rep
scopeOf (WhileLoop VName
_) = forall a. Monoid a => a
mempty
scopeOf (ForLoop VName
i IntType
it SubExp
_ [(LParam rep, VName)]
xs) =
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
i (forall rep. IntType -> NameInfo rep
IndexName IntType
it) forall a b. (a -> b) -> a -> b
$ forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(LParam rep, VName)]
xs)
scopeOfPat :: LetDec rep ~ dec => Pat dec -> Scope rep
scopeOfPat :: forall rep dec. (LetDec rep ~ dec) => Pat dec -> Scope rep
scopeOfPat =
forall a. Monoid a => [a] -> a
mconcat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall rep dec. (LetDec rep ~ dec) => PatElem dec -> Scope rep
scopeOfPatElem forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall dec. Pat dec -> [PatElem dec]
patElems
scopeOfPatElem :: LetDec rep ~ dec => PatElem dec -> Scope rep
scopeOfPatElem :: forall rep dec. (LetDec rep ~ dec) => PatElem dec -> Scope rep
scopeOfPatElem (PatElem VName
name dec
dec) = forall k a. k -> a -> Map k a
M.singleton VName
name forall a b. (a -> b) -> a -> b
$ forall rep. LetDec rep -> NameInfo rep
LetName dec
dec
scopeOfLParams ::
LParamInfo rep ~ dec =>
[Param dec] ->
Scope rep
scopeOfLParams :: forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {rep}. Param (LParamInfo rep) -> (VName, NameInfo rep)
f
where
f :: Param (LParamInfo rep) -> (VName, NameInfo rep)
f Param (LParamInfo rep)
param = (forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
param, forall rep. LParamInfo rep -> NameInfo rep
LParamName forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> dec
paramDec Param (LParamInfo rep)
param)
scopeOfFParams ::
FParamInfo rep ~ dec =>
[Param dec] ->
Scope rep
scopeOfFParams :: forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall {rep}. Param (FParamInfo rep) -> (VName, NameInfo rep)
f
where
f :: Param (FParamInfo rep) -> (VName, NameInfo rep)
f Param (FParamInfo rep)
param = (forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
param, forall rep. FParamInfo rep -> NameInfo rep
FParamName forall a b. (a -> b) -> a -> b
$ forall dec. Param dec -> dec
paramDec Param (FParamInfo rep)
param)
instance Scoped rep (Lambda rep) where
scopeOf :: Lambda rep -> Scope rep
scopeOf Lambda rep
lam = forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams forall a b. (a -> b) -> a -> b
$ forall rep. Lambda rep -> [LParam rep]
lambdaParams Lambda rep
lam
type SameScope rep1 rep2 =
( LetDec rep1 ~ LetDec rep2,
FParamInfo rep1 ~ FParamInfo rep2,
LParamInfo rep1 ~ LParamInfo rep2
)
castScope ::
SameScope fromrep torep =>
Scope fromrep ->
Scope torep
castScope :: forall fromrep torep.
SameScope fromrep torep =>
Scope fromrep -> Scope torep
castScope = forall a b k. (a -> b) -> Map k a -> Map k b
M.map forall fromrep torep.
SameScope fromrep torep =>
NameInfo fromrep -> NameInfo torep
castNameInfo
castNameInfo ::
SameScope fromrep torep =>
NameInfo fromrep ->
NameInfo torep
castNameInfo :: forall fromrep torep.
SameScope fromrep torep =>
NameInfo fromrep -> NameInfo torep
castNameInfo (LetName LetDec fromrep
dec) = forall rep. LetDec rep -> NameInfo rep
LetName LetDec fromrep
dec
castNameInfo (FParamName FParamInfo fromrep
dec) = forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo fromrep
dec
castNameInfo (LParamName LParamInfo fromrep
dec) = forall rep. LParamInfo rep -> NameInfo rep
LParamName LParamInfo fromrep
dec
castNameInfo (IndexName IntType
it) = forall rep. IntType -> NameInfo rep
IndexName IntType
it
newtype ExtendedScope rep m a = ExtendedScope (ReaderT (Scope rep) m a)
deriving
( forall a b. a -> ExtendedScope rep m b -> ExtendedScope rep m a
forall a b.
(a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m b
forall rep (m :: * -> *) a b.
Functor m =>
a -> ExtendedScope rep m b -> ExtendedScope rep m a
forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m 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 -> ExtendedScope rep m b -> ExtendedScope rep m a
$c<$ :: forall rep (m :: * -> *) a b.
Functor m =>
a -> ExtendedScope rep m b -> ExtendedScope rep m a
fmap :: forall a b.
(a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m b
$cfmap :: forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m b
Functor,
forall a. a -> ExtendedScope rep m a
forall a b.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m a
forall a b.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
forall a b.
ExtendedScope rep m (a -> b)
-> ExtendedScope rep m a -> ExtendedScope rep m b
forall a b c.
(a -> b -> c)
-> ExtendedScope rep m a
-> ExtendedScope rep m b
-> ExtendedScope rep m c
forall {rep} {m :: * -> *}.
Applicative m =>
Functor (ExtendedScope rep m)
forall rep (m :: * -> *) a.
Applicative m =>
a -> ExtendedScope rep m a
forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m a
forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m (a -> b)
-> ExtendedScope rep m a -> ExtendedScope rep m b
forall rep (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c)
-> ExtendedScope rep m a
-> ExtendedScope rep m b
-> ExtendedScope rep m 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.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m a
$c<* :: forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m a
*> :: forall a b.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
$c*> :: forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
liftA2 :: forall a b c.
(a -> b -> c)
-> ExtendedScope rep m a
-> ExtendedScope rep m b
-> ExtendedScope rep m c
$cliftA2 :: forall rep (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c)
-> ExtendedScope rep m a
-> ExtendedScope rep m b
-> ExtendedScope rep m c
<*> :: forall a b.
ExtendedScope rep m (a -> b)
-> ExtendedScope rep m a -> ExtendedScope rep m b
$c<*> :: forall rep (m :: * -> *) a b.
Applicative m =>
ExtendedScope rep m (a -> b)
-> ExtendedScope rep m a -> ExtendedScope rep m b
pure :: forall a. a -> ExtendedScope rep m a
$cpure :: forall rep (m :: * -> *) a.
Applicative m =>
a -> ExtendedScope rep m a
Applicative,
forall a. a -> ExtendedScope rep m a
forall a b.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
forall a b.
ExtendedScope rep m a
-> (a -> ExtendedScope rep m b) -> ExtendedScope rep m b
forall {rep} {m :: * -> *}.
Monad m =>
Applicative (ExtendedScope rep m)
forall rep (m :: * -> *) a. Monad m => a -> ExtendedScope rep m a
forall rep (m :: * -> *) a b.
Monad m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
forall rep (m :: * -> *) a b.
Monad m =>
ExtendedScope rep m a
-> (a -> ExtendedScope rep m b) -> ExtendedScope rep m 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 -> ExtendedScope rep m a
$creturn :: forall rep (m :: * -> *) a. Monad m => a -> ExtendedScope rep m a
>> :: forall a b.
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
$c>> :: forall rep (m :: * -> *) a b.
Monad m =>
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
>>= :: forall a b.
ExtendedScope rep m a
-> (a -> ExtendedScope rep m b) -> ExtendedScope rep m b
$c>>= :: forall rep (m :: * -> *) a b.
Monad m =>
ExtendedScope rep m a
-> (a -> ExtendedScope rep m b) -> ExtendedScope rep m b
Monad,
MonadReader (Scope rep)
)
instance
(HasScope rep m, Monad m) =>
HasScope rep (ExtendedScope rep m)
where
lookupType :: VName -> ExtendedScope rep m Type
lookupType VName
name = do
Maybe Type
res <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall t. Typed t => t -> Type
typeOf forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall rep (m :: * -> *) a.
ReaderT (Scope rep) m a -> ExtendedScope rep m a
ExtendedScope forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
name) forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Type
res
askScope :: ExtendedScope rep m (Scope rep)
askScope = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall rep (m :: * -> *) a.
ReaderT (Scope rep) m a -> ExtendedScope rep m a
ExtendedScope (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope)
extendedScope ::
ExtendedScope rep m a ->
Scope rep ->
m a
extendedScope :: forall rep (m :: * -> *) a.
ExtendedScope rep m a -> Scope rep -> m a
extendedScope (ExtendedScope ReaderT (Scope rep) m a
m) = forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope rep) m a
m