{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | The core Futhark AST does not contain type information when we
-- use a variable.  Therefore, most transformations expect to be able
-- to access some kind of symbol table that maps names to their types.
--
-- This module defines the concept of a type environment as a mapping
-- from variable names to 'NameInfo's.  Convenience facilities are
-- also provided to communicate that some monad or applicative functor
-- maintains type information.
module Futhark.IR.Prop.Scope
  ( HasScope (..),
    NameInfo (..),
    LocalScope (..),
    Scope,
    Scoped (..),
    inScopeOf,
    scopeOfLParams,
    scopeOfFParams,
    scopeOfPat,
    scopeOfPatElem,
    SameScope,
    castScope,

    -- * Extended type environment
    ExtendedScope,
    extendedScope,
  )
where

import Control.Monad.Except
import qualified Control.Monad.RWS.Lazy
import qualified Control.Monad.RWS.Strict
import Control.Monad.Reader
import qualified Data.Map.Strict as M
import Futhark.IR.Pretty ()
import Futhark.IR.Prop.Types
import Futhark.IR.Rep
import Futhark.IR.Syntax

-- | How some name in scope was bound.
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) = LetDec rep -> Type
forall t. Typed t => t -> Type
typeOf LetDec rep
dec
  typeOf (FParamName FParamInfo rep
dec) = FParamInfo rep -> Type
forall t. Typed t => t -> Type
typeOf FParamInfo rep
dec
  typeOf (LParamName LParamInfo rep
dec) = LParamInfo rep -> Type
forall t. Typed t => t -> Type
typeOf LParamInfo rep
dec
  typeOf (IndexName IntType
it) = PrimType -> Type
forall shape u. PrimType -> TypeBase shape u
Prim (PrimType -> Type) -> PrimType -> Type
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
IntType IntType
it

-- | A scope is a mapping from variable names to information about
-- that name.
type Scope rep = M.Map VName (NameInfo rep)

-- | The class of applicative functors (or more common in practice:
-- monads) that permit the lookup of variable types.  A default method
-- for 'lookupType' exists, which is sufficient (if not always
-- maximally efficient, and using 'error' to fail) when 'askScope'
-- is defined.
class (Applicative m, RepTypes rep) => HasScope rep m | m -> rep where
  -- | Return the type of the given variable, or fail if it is not in
  -- the type environment.
  lookupType :: VName -> m Type
  lookupType = (NameInfo rep -> Type) -> m (NameInfo rep) -> m Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameInfo rep -> Type
forall t. Typed t => t -> Type
typeOf (m (NameInfo rep) -> m Type)
-> (VName -> m (NameInfo rep)) -> VName -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> m (NameInfo rep)
forall rep (m :: * -> *).
HasScope rep m =>
VName -> m (NameInfo rep)
lookupInfo

  -- | Return the info of the given variable, or fail if it is not in
  -- the type environment.
  lookupInfo :: VName -> m (NameInfo rep)
  lookupInfo VName
name =
    (Scope rep -> NameInfo rep) -> m (NameInfo rep)
forall rep (m :: * -> *) a.
HasScope rep m =>
(Scope rep -> a) -> m a
asksScope (NameInfo rep -> VName -> Scope rep -> NameInfo rep
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault NameInfo rep
notFound VName
name)
    where
      notFound :: NameInfo rep
notFound =
        String -> NameInfo rep
forall a. HasCallStack => String -> a
error (String -> NameInfo rep) -> String -> NameInfo rep
forall a b. (a -> b) -> a -> b
$
          String
"Scope.lookupInfo: Name " String -> ShowS
forall a. [a] -> [a] -> [a]
++ VName -> String
forall a. Pretty a => a -> String
pretty VName
name
            String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" not found in type environment."

  -- | Return the type environment contained in the applicative
  -- functor.
  askScope :: m (Scope rep)

  -- | Return the result of applying some function to the type
  -- environment.
  asksScope :: (Scope rep -> a) -> m a
  asksScope Scope rep -> a
f = Scope rep -> a
f (Scope rep -> a) -> m (Scope rep) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope

instance
  (Applicative m, Monad m, RepTypes rep) =>
  HasScope rep (ReaderT (Scope rep) m)
  where
  askScope :: ReaderT (Scope rep) m (Scope rep)
askScope = ReaderT (Scope rep) m (Scope rep)
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 = m (Scope rep) -> ExceptT e m (Scope rep)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope

instance
  (Applicative m, 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 = RWST (Scope rep) w s m (Scope rep)
forall r (m :: * -> *). MonadReader r m => m r
ask

instance
  (Applicative m, 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 = RWST (Scope rep) w s m (Scope rep)
forall r (m :: * -> *). MonadReader r m => m r
ask

-- | The class of monads that not only provide a 'Scope', but also
-- the ability to locally extend it.  A 'Reader' containing a
-- 'Scope' is the prototypical example of such a monad.
class (HasScope rep m, Monad m) => LocalScope rep m where
  -- | Run a computation with an extended type environment.  Note that
  -- this is intended to *add* to the current type environment, it
  -- does not replace it.
  localScope :: Scope rep -> m a -> m a

instance (Monad m, LocalScope rep m) => LocalScope rep (ExceptT e m) where
  localScope :: Scope rep -> ExceptT e m a -> ExceptT e m a
localScope = (m (Either e a) -> m (Either e a))
-> ExceptT e m a -> ExceptT e m a
forall (m :: * -> *) e a (n :: * -> *) e' b.
(m (Either e a) -> n (Either e' b))
-> ExceptT e m a -> ExceptT e' n b
mapExceptT ((m (Either e a) -> m (Either e a))
 -> ExceptT e m a -> ExceptT e m a)
-> (Scope rep -> m (Either e a) -> m (Either e a))
-> Scope rep
-> ExceptT e m a
-> ExceptT e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope rep -> m (Either e a) -> m (Either e a)
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope

instance
  (Applicative m, Monad m, RepTypes rep) =>
  LocalScope rep (ReaderT (Scope rep) m)
  where
  localScope :: Scope rep -> ReaderT (Scope rep) m a -> ReaderT (Scope rep) m a
localScope = (Scope rep -> Scope rep)
-> ReaderT (Scope rep) m a -> ReaderT (Scope rep) m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Scope rep -> Scope rep)
 -> ReaderT (Scope rep) m a -> ReaderT (Scope rep) m a)
-> (Scope rep -> Scope rep -> Scope rep)
-> Scope rep
-> ReaderT (Scope rep) m a
-> ReaderT (Scope rep) m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope rep -> Scope rep -> Scope rep
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union

instance
  (Applicative m, Monad m, Monoid w, RepTypes rep) =>
  LocalScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m)
  where
  localScope :: Scope rep -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
localScope = (Scope rep -> Scope rep)
-> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Scope rep -> Scope rep)
 -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a)
-> (Scope rep -> Scope rep -> Scope rep)
-> Scope rep
-> RWST (Scope rep) w s m a
-> RWST (Scope rep) w s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope rep -> Scope rep -> Scope rep
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union

instance
  (Applicative m, Monad m, Monoid w, RepTypes rep) =>
  LocalScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m)
  where
  localScope :: Scope rep -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
localScope = (Scope rep -> Scope rep)
-> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Scope rep -> Scope rep)
 -> RWST (Scope rep) w s m a -> RWST (Scope rep) w s m a)
-> (Scope rep -> Scope rep -> Scope rep)
-> Scope rep
-> RWST (Scope rep) w s m a
-> RWST (Scope rep) w s m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope rep -> Scope rep -> Scope rep
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union

-- | The class of things that can provide a scope.  There is no
-- overarching rule for what this means.  For a 'Stm', it is the
-- corresponding pattern.  For a t'Lambda', is is the parameters.
class Scoped rep a | a -> rep where
  scopeOf :: a -> Scope rep

-- | Extend the monadic scope with the 'scopeOf' the given value.
inScopeOf :: (Scoped rep a, LocalScope rep m) => a -> m b -> m b
inScopeOf :: a -> m b -> m b
inScopeOf = Scope rep -> m b -> m b
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope (Scope rep -> m b -> m b) -> (a -> Scope rep) -> a -> m b -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf

instance Scoped rep a => Scoped rep [a] where
  scopeOf :: [a] -> Scope rep
scopeOf = [Scope rep] -> Scope rep
forall a. Monoid a => [a] -> a
mconcat ([Scope rep] -> Scope rep)
-> ([a] -> [Scope rep]) -> [a] -> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Scope rep) -> [a] -> [Scope rep]
forall a b. (a -> b) -> [a] -> [b]
map a -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf

instance Scoped rep (Stms rep) where
  scopeOf :: Stms rep -> Scope rep
scopeOf = (Stm rep -> Scope rep) -> Stms rep -> Scope rep
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Stm rep -> Scope rep
forall rep a. Scoped rep a => a -> Scope rep
scopeOf

instance Scoped rep (Stm rep) where
  scopeOf :: Stm rep -> Scope rep
scopeOf = PatT (LetDec rep) -> Scope rep
forall rep dec. (LetDec rep ~ dec) => PatT dec -> Scope rep
scopeOfPat (PatT (LetDec rep) -> Scope rep)
-> (Stm rep -> PatT (LetDec rep)) -> Stm rep -> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stm rep -> PatT (LetDec rep)
forall rep. Stm rep -> Pat rep
stmPat

instance Scoped rep (FunDef rep) where
  scopeOf :: FunDef rep -> Scope rep
scopeOf = [Param (FParamInfo rep)] -> Scope rep
forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams ([Param (FParamInfo rep)] -> Scope rep)
-> (FunDef rep -> [Param (FParamInfo rep)])
-> FunDef rep
-> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FunDef rep -> [Param (FParamInfo rep)]
forall rep. FunDef rep -> [FParam rep]
funDefParams

instance Scoped rep (VName, NameInfo rep) where
  scopeOf :: (VName, NameInfo rep) -> Scope rep
scopeOf = (VName -> NameInfo rep -> Scope rep)
-> (VName, NameInfo rep) -> Scope rep
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry VName -> NameInfo rep -> Scope rep
forall k a. k -> a -> Map k a
M.singleton

instance Scoped rep (LoopForm rep) where
  scopeOf :: LoopForm rep -> Scope rep
scopeOf (WhileLoop VName
_) = Scope rep
forall a. Monoid a => a
mempty
  scopeOf (ForLoop VName
i IntType
it SubExp
_ [(LParam rep, VName)]
xs) =
    VName -> NameInfo rep -> Scope rep -> Scope rep
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VName
i (IntType -> NameInfo rep
forall rep. IntType -> NameInfo rep
IndexName IntType
it) (Scope rep -> Scope rep) -> Scope rep -> Scope rep
forall a b. (a -> b) -> a -> b
$ [LParam rep] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams (((LParam rep, VName) -> LParam rep)
-> [(LParam rep, VName)] -> [LParam rep]
forall a b. (a -> b) -> [a] -> [b]
map (LParam rep, VName) -> LParam rep
forall a b. (a, b) -> a
fst [(LParam rep, VName)]
xs)

-- | The scope of a pattern.
scopeOfPat :: LetDec rep ~ dec => PatT dec -> Scope rep
scopeOfPat :: PatT dec -> Scope rep
scopeOfPat =
  [Scope rep] -> Scope rep
forall a. Monoid a => [a] -> a
mconcat ([Scope rep] -> Scope rep)
-> (PatT dec -> [Scope rep]) -> PatT dec -> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PatElemT dec -> Scope rep) -> [PatElemT dec] -> [Scope rep]
forall a b. (a -> b) -> [a] -> [b]
map PatElemT dec -> Scope rep
forall rep dec. (LetDec rep ~ dec) => PatElemT dec -> Scope rep
scopeOfPatElem ([PatElemT dec] -> [Scope rep])
-> (PatT dec -> [PatElemT dec]) -> PatT dec -> [Scope rep]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PatT dec -> [PatElemT dec]
forall dec. PatT dec -> [PatElemT dec]
patElems

-- | The scope of a pattern element.
scopeOfPatElem :: LetDec rep ~ dec => PatElemT dec -> Scope rep
scopeOfPatElem :: PatElemT dec -> Scope rep
scopeOfPatElem (PatElem VName
name dec
dec) = VName -> NameInfo rep -> Scope rep
forall k a. k -> a -> Map k a
M.singleton VName
name (NameInfo rep -> Scope rep) -> NameInfo rep -> Scope rep
forall a b. (a -> b) -> a -> b
$ LetDec rep -> NameInfo rep
forall rep. LetDec rep -> NameInfo rep
LetName dec
LetDec rep
dec

-- | The scope of some lambda parameters.
scopeOfLParams ::
  LParamInfo rep ~ dec =>
  [Param dec] ->
  Scope rep
scopeOfLParams :: [Param dec] -> Scope rep
scopeOfLParams = [(VName, NameInfo rep)] -> Scope rep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, NameInfo rep)] -> Scope rep)
-> ([Param dec] -> [(VName, NameInfo rep)])
-> [Param dec]
-> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param dec -> (VName, NameInfo rep))
-> [Param dec] -> [(VName, NameInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map Param dec -> (VName, NameInfo rep)
forall rep. Param (LParamInfo rep) -> (VName, NameInfo rep)
f
  where
    f :: Param (LParamInfo rep) -> (VName, NameInfo rep)
f Param (LParamInfo rep)
param = (Param (LParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (LParamInfo rep)
param, LParamInfo rep -> NameInfo rep
forall rep. LParamInfo rep -> NameInfo rep
LParamName (LParamInfo rep -> NameInfo rep) -> LParamInfo rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ Param (LParamInfo rep) -> LParamInfo rep
forall dec. Param dec -> dec
paramDec Param (LParamInfo rep)
param)

-- | The scope of some function or loop parameters.
scopeOfFParams ::
  FParamInfo rep ~ dec =>
  [Param dec] ->
  Scope rep
scopeOfFParams :: [Param dec] -> Scope rep
scopeOfFParams = [(VName, NameInfo rep)] -> Scope rep
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, NameInfo rep)] -> Scope rep)
-> ([Param dec] -> [(VName, NameInfo rep)])
-> [Param dec]
-> Scope rep
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param dec -> (VName, NameInfo rep))
-> [Param dec] -> [(VName, NameInfo rep)]
forall a b. (a -> b) -> [a] -> [b]
map Param dec -> (VName, NameInfo rep)
forall rep. Param (FParamInfo rep) -> (VName, NameInfo rep)
f
  where
    f :: Param (FParamInfo rep) -> (VName, NameInfo rep)
f Param (FParamInfo rep)
param = (Param (FParamInfo rep) -> VName
forall dec. Param dec -> VName
paramName Param (FParamInfo rep)
param, FParamInfo rep -> NameInfo rep
forall rep. FParamInfo rep -> NameInfo rep
FParamName (FParamInfo rep -> NameInfo rep) -> FParamInfo rep -> NameInfo rep
forall a b. (a -> b) -> a -> b
$ Param (FParamInfo rep) -> FParamInfo rep
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 = [Param (LParamInfo rep)] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams ([Param (LParamInfo rep)] -> Scope rep)
-> [Param (LParamInfo rep)] -> Scope rep
forall a b. (a -> b) -> a -> b
$ Lambda rep -> [Param (LParamInfo rep)]
forall rep. LambdaT rep -> [LParam rep]
lambdaParams Lambda rep
lam

-- | A constraint that indicates two representations have the same 'NameInfo'
-- representation.
type SameScope rep1 rep2 =
  ( LetDec rep1 ~ LetDec rep2,
    FParamInfo rep1 ~ FParamInfo rep2,
    LParamInfo rep1 ~ LParamInfo rep2
  )

-- | If two scopes are really the same, then you can convert one to
-- the other.
castScope ::
  SameScope fromrep torep =>
  Scope fromrep ->
  Scope torep
castScope :: Scope fromrep -> Scope torep
castScope = (NameInfo fromrep -> NameInfo torep)
-> Scope fromrep -> Scope torep
forall a b k. (a -> b) -> Map k a -> Map k b
M.map NameInfo fromrep -> NameInfo torep
forall fromrep torep.
SameScope fromrep torep =>
NameInfo fromrep -> NameInfo torep
castNameInfo

castNameInfo ::
  SameScope fromrep torep =>
  NameInfo fromrep ->
  NameInfo torep
castNameInfo :: NameInfo fromrep -> NameInfo torep
castNameInfo (LetName LetDec fromrep
dec) = LetDec torep -> NameInfo torep
forall rep. LetDec rep -> NameInfo rep
LetName LetDec fromrep
LetDec torep
dec
castNameInfo (FParamName FParamInfo fromrep
dec) = FParamInfo torep -> NameInfo torep
forall rep. FParamInfo rep -> NameInfo rep
FParamName FParamInfo fromrep
FParamInfo torep
dec
castNameInfo (LParamName LParamInfo fromrep
dec) = LParamInfo torep -> NameInfo torep
forall rep. LParamInfo rep -> NameInfo rep
LParamName LParamInfo fromrep
LParamInfo torep
dec
castNameInfo (IndexName IntType
it) = IntType -> NameInfo torep
forall rep. IntType -> NameInfo rep
IndexName IntType
it

-- | A monad transformer that carries around an extended 'Scope'.
-- Its 'lookupType' method will first look in the extended 'Scope',
-- and then use the 'lookupType' method of the underlying monad.
newtype ExtendedScope rep m a = ExtendedScope (ReaderT (Scope rep) m a)
  deriving
    ( a -> ExtendedScope rep m b -> ExtendedScope rep m a
(a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m b
(forall a b.
 (a -> b) -> ExtendedScope rep m a -> ExtendedScope rep m b)
-> (forall a b.
    a -> ExtendedScope rep m b -> ExtendedScope rep m a)
-> Functor (ExtendedScope rep m)
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
<$ :: 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 :: (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,
      Functor (ExtendedScope rep m)
a -> ExtendedScope rep m a
Functor (ExtendedScope rep m)
-> (forall a. a -> ExtendedScope rep m a)
-> (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 a b.
    ExtendedScope rep m a
    -> ExtendedScope rep m b -> ExtendedScope rep m b)
-> (forall a b.
    ExtendedScope rep m a
    -> ExtendedScope rep m b -> ExtendedScope rep m a)
-> Applicative (ExtendedScope rep m)
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m a
ExtendedScope rep m (a -> b)
-> ExtendedScope rep m a -> ExtendedScope rep m b
(a -> b -> c)
-> ExtendedScope rep m a
-> ExtendedScope rep m b
-> ExtendedScope rep m c
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
<* :: 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
*> :: 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 :: (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
<*> :: 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 :: a -> ExtendedScope rep m a
$cpure :: forall rep (m :: * -> *) a.
Applicative m =>
a -> ExtendedScope rep m a
$cp1Applicative :: forall rep (m :: * -> *).
Applicative m =>
Functor (ExtendedScope rep m)
Applicative,
      Applicative (ExtendedScope rep m)
a -> ExtendedScope rep m a
Applicative (ExtendedScope rep m)
-> (forall a b.
    ExtendedScope rep m a
    -> (a -> ExtendedScope rep m b) -> ExtendedScope rep m b)
-> (forall a b.
    ExtendedScope rep m a
    -> ExtendedScope rep m b -> ExtendedScope rep m b)
-> (forall a. a -> ExtendedScope rep m a)
-> Monad (ExtendedScope rep m)
ExtendedScope rep m a
-> (a -> ExtendedScope rep m b) -> ExtendedScope rep m b
ExtendedScope rep m a
-> ExtendedScope rep m b -> ExtendedScope rep m b
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 :: a -> ExtendedScope rep m a
$creturn :: forall rep (m :: * -> *) a. Monad m => a -> ExtendedScope rep m a
>> :: 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
>>= :: 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
$cp1Monad :: forall rep (m :: * -> *).
Monad m =>
Applicative (ExtendedScope rep m)
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 <- (Scope rep -> Maybe Type) -> ExtendedScope rep m (Maybe Type)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Scope rep -> Maybe Type) -> ExtendedScope rep m (Maybe Type))
-> (Scope rep -> Maybe Type) -> ExtendedScope rep m (Maybe Type)
forall a b. (a -> b) -> a -> b
$ (NameInfo rep -> Type) -> Maybe (NameInfo rep) -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NameInfo rep -> Type
forall t. Typed t => t -> Type
typeOf (Maybe (NameInfo rep) -> Maybe Type)
-> (Scope rep -> Maybe (NameInfo rep)) -> Scope rep -> Maybe Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VName -> Scope rep -> Maybe (NameInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name
    ExtendedScope rep m Type
-> (Type -> ExtendedScope rep m Type)
-> Maybe Type
-> ExtendedScope rep m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ReaderT (Scope rep) m Type -> ExtendedScope rep m Type
forall rep (m :: * -> *) a.
ReaderT (Scope rep) m a -> ExtendedScope rep m a
ExtendedScope (ReaderT (Scope rep) m Type -> ExtendedScope rep m Type)
-> ReaderT (Scope rep) m Type -> ExtendedScope rep m Type
forall a b. (a -> b) -> a -> b
$ m Type -> ReaderT (Scope rep) m Type
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Type -> ReaderT (Scope rep) m Type)
-> m Type -> ReaderT (Scope rep) m Type
forall a b. (a -> b) -> a -> b
$ VName -> m Type
forall rep (m :: * -> *). HasScope rep m => VName -> m Type
lookupType VName
name) Type -> ExtendedScope rep m Type
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Type
res
  askScope :: ExtendedScope rep m (Scope rep)
askScope = (Scope rep -> Scope rep -> Scope rep)
-> ExtendedScope rep m (Scope rep -> Scope rep)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Scope rep -> Scope rep -> Scope rep
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union ExtendedScope rep m (Scope rep -> Scope rep)
-> ExtendedScope rep m (Scope rep)
-> ExtendedScope rep m (Scope rep)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ReaderT (Scope rep) m (Scope rep)
-> ExtendedScope rep m (Scope rep)
forall rep (m :: * -> *) a.
ReaderT (Scope rep) m a -> ExtendedScope rep m a
ExtendedScope (m (Scope rep) -> ReaderT (Scope rep) m (Scope rep)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Scope rep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope)

-- | Run a computation in the extended type environment.
extendedScope ::
  ExtendedScope rep m a ->
  Scope rep ->
  m a
extendedScope :: ExtendedScope rep m a -> Scope rep -> m a
extendedScope (ExtendedScope ReaderT (Scope rep) m a
m) = ReaderT (Scope rep) m a -> Scope rep -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Scope rep) m a
m