Scope.hs

{-# 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,
    scopeOfPattern,
    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.Decorations
import Futhark.IR.Pretty ()
import Futhark.IR.Prop.Patterns
import Futhark.IR.Prop.Types
import Futhark.IR.Syntax

-- | How some name in scope was bound.
data NameInfo lore
  = LetName (LetDec lore)
  | FParamName (FParamInfo lore)
  | LParamName (LParamInfo lore)
  | IndexName IntType

deriving instance Decorations lore => Show (NameInfo lore)

instance Decorations lore => Typed (NameInfo lore) where
  typeOf (LetName dec) = typeOf dec
  typeOf (FParamName dec) = typeOf dec
  typeOf (LParamName dec) = typeOf dec
  typeOf (IndexName it) = Prim $ IntType it

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

-- | 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, Decorations lore) => HasScope lore m | m -> lore where
  -- | Return the type of the given variable, or fail if it is not in
  -- the type environment.
  lookupType :: VName -> m Type
  lookupType = fmap typeOf . lookupInfo

  -- | Return the info of the given variable, or fail if it is not in
  -- the type environment.
  lookupInfo :: VName -> m (NameInfo lore)
  lookupInfo name =
    asksScope (M.findWithDefault notFound name)
    where
      notFound =
        error $
          "Scope.lookupInfo: Name " ++ pretty name
            ++ " not found in type environment."

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

  -- | Return the result of applying some function to the type
  -- environment.
  asksScope :: (Scope lore -> a) -> m a
  asksScope f = f <$> askScope

instance
  (Applicative m, Monad m, Decorations lore) =>
  HasScope lore (ReaderT (Scope lore) m)
  where
  askScope = ask

instance (Monad m, HasScope lore m) => HasScope lore (ExceptT e m) where
  askScope = lift askScope

instance
  (Applicative m, Monad m, Monoid w, Decorations lore) =>
  HasScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m)
  where
  askScope = ask

instance
  (Applicative m, Monad m, Monoid w, Decorations lore) =>
  HasScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m)
  where
  askScope = 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 lore m, Monad m) => LocalScope lore 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 lore -> m a -> m a

instance (Monad m, LocalScope lore m) => LocalScope lore (ExceptT e m) where
  localScope = mapExceptT . localScope

instance
  (Applicative m, Monad m, Decorations lore) =>
  LocalScope lore (ReaderT (Scope lore) m)
  where
  localScope = local . M.union

instance
  (Applicative m, Monad m, Monoid w, Decorations lore) =>
  LocalScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m)
  where
  localScope = local . M.union

instance
  (Applicative m, Monad m, Monoid w, Decorations lore) =>
  LocalScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m)
  where
  localScope = local . 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 lore a | a -> lore where
  scopeOf :: a -> Scope lore

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

instance Scoped lore a => Scoped lore [a] where
  scopeOf = mconcat . map scopeOf

instance Scoped lore (Stms lore) where
  scopeOf = foldMap scopeOf

instance Scoped lore (Stm lore) where
  scopeOf = scopeOfPattern . stmPattern

instance Scoped lore (FunDef lore) where
  scopeOf = scopeOfFParams . funDefParams

instance Scoped lore (VName, NameInfo lore) where
  scopeOf = uncurry M.singleton

instance Scoped lore (LoopForm lore) where
  scopeOf (WhileLoop _) = mempty
  scopeOf (ForLoop i it _ xs) =
    M.insert i (IndexName it) $ scopeOfLParams (map fst xs)

-- | The scope of a pattern.
scopeOfPattern :: LetDec lore ~ dec => PatternT dec -> Scope lore
scopeOfPattern =
  mconcat . map scopeOfPatElem . patternElements

-- | The scope of a pattern element.
scopeOfPatElem :: LetDec lore ~ dec => PatElemT dec -> Scope lore
scopeOfPatElem (PatElem name dec) = M.singleton name $ LetName dec

-- | The scope of some lambda parameters.
scopeOfLParams ::
  LParamInfo lore ~ dec =>
  [Param dec] ->
  Scope lore
scopeOfLParams = M.fromList . map f
  where
    f param = (paramName param, LParamName $ paramDec param)

-- | The scope of some function or loop parameters.
scopeOfFParams ::
  FParamInfo lore ~ dec =>
  [Param dec] ->
  Scope lore
scopeOfFParams = M.fromList . map f
  where
    f param = (paramName param, FParamName $ paramDec param)

instance Scoped lore (Lambda lore) where
  scopeOf lam = scopeOfLParams $ lambdaParams lam

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

-- | If two scopes are really the same, then you can convert one to
-- the other.
castScope ::
  SameScope fromlore tolore =>
  Scope fromlore ->
  Scope tolore
castScope = M.map castNameInfo

castNameInfo ::
  SameScope fromlore tolore =>
  NameInfo fromlore ->
  NameInfo tolore
castNameInfo (LetName dec) = LetName dec
castNameInfo (FParamName dec) = FParamName dec
castNameInfo (LParamName dec) = LParamName dec
castNameInfo (IndexName it) = IndexName 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 lore m a = ExtendedScope (ReaderT (Scope lore) m a)
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadReader (Scope lore)
    )

instance
  (HasScope lore m, Monad m) =>
  HasScope lore (ExtendedScope lore m)
  where
  lookupType name = do
    res <- asks $ fmap typeOf . M.lookup name
    maybe (ExtendedScope $ lift $ lookupType name) return res
  askScope = asks M.union <*> ExtendedScope (lift askScope)

-- | Run a computation in the extended type environment.
extendedScope ::
  ExtendedScope lore m a ->
  Scope lore ->
  m a
extendedScope (ExtendedScope m) = runReaderT m