Binder.hs

{-# LANGUAGE FlexibleContexts, GeneralizedNewtypeDeriving, TypeFamilies, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Trustworthy #-}
-- | This module defines a convenience monad/typeclass for creating
-- normalised programs.  The fundamental building block is 'BinderT'
-- and its execution functions, but it is usually easier to use
-- 'Binder'.
--
-- See "Futhark.Construct" for a high-level description.
module Futhark.Binder
  ( -- * A concrete @MonadBinder@ monad.
    BinderT
  , runBinderT, runBinderT_
  , runBinderT', runBinderT'_
  , BinderOps (..)
  , Binder
  , runBinder
  , runBinder_
  , runBodyBinder

  -- * The 'MonadBinder' typeclass
  , module Futhark.Binder.Class
  )
where

import Control.Arrow (second)
import Control.Monad.Writer
import Control.Monad.State.Strict
import Control.Monad.Reader
import Control.Monad.Error.Class
import qualified Data.Map.Strict as M

import Futhark.Binder.Class
import Futhark.IR

-- | A 'BinderT' (and by extension, a 'Binder') is only an instance of
-- 'MonadBinder' for lores that implement this type class, which
-- contains methods for constructing statements.
class ASTLore lore => BinderOps lore where
  mkExpDecB :: (MonadBinder m, Lore m ~ lore) =>
                Pattern lore -> Exp lore -> m (ExpDec lore)
  mkBodyB :: (MonadBinder m, Lore m ~ lore) =>
             Stms lore -> Result -> m (Body lore)
  mkLetNamesB :: (MonadBinder m, Lore m ~ lore) =>
                 [VName] -> Exp lore -> m (Stm lore)

  default mkExpDecB :: (MonadBinder m, Bindable lore) =>
                       Pattern lore -> Exp lore -> m (ExpDec lore)
  mkExpDecB pat e = return $ mkExpDec pat e

  default mkBodyB :: (MonadBinder m, Bindable lore) =>
                     Stms lore -> Result -> m (Body lore)
  mkBodyB stms res = return $ mkBody stms res

  default mkLetNamesB :: (MonadBinder m, Lore m ~ lore, Bindable lore) =>
                         [VName] -> Exp lore -> m (Stm lore)
  mkLetNamesB = mkLetNames

-- | A monad transformer that tracks statements and provides a
-- 'MonadBinder' instance, assuming that the underlying monad provides
-- a name source.  In almost all cases, this is what you will use for
-- constructing statements (possibly as part of a larger monad stack).
-- If you find yourself needing to implement 'MonadBinder' from
-- scratch, then it is likely that you are making a mistake.
newtype BinderT lore m a = BinderT (StateT (Stms lore, Scope lore) m a)
  deriving (Functor, Monad, Applicative)

instance MonadTrans (BinderT lore) where
  lift = BinderT . lift

-- | The most commonly used binder monad.
type Binder lore = BinderT lore (State VNameSource)

instance MonadFreshNames m => MonadFreshNames (BinderT lore m) where
  getNameSource = lift getNameSource
  putNameSource = lift . putNameSource

instance (ASTLore lore, Monad m) =>
         HasScope lore (BinderT lore m) where
  lookupType name = do
    t <- BinderT $ gets $ M.lookup name . snd
    case t of
      Nothing -> error $ "BinderT.lookupType: unknown variable " ++ pretty name
      Just t' -> return $ typeOf t'
  askScope = BinderT $ gets snd

instance (ASTLore lore, Monad m) =>
         LocalScope lore (BinderT lore m) where
  localScope types (BinderT m) = BinderT $ do
    modify $ second (M.union types)
    x <- m
    modify $ second (`M.difference` types)
    return x

instance (ASTLore lore, MonadFreshNames m, BinderOps lore) =>
         MonadBinder (BinderT lore m) where
  type Lore (BinderT lore m) = lore
  mkExpDecM = mkExpDecB
  mkBodyM = mkBodyB
  mkLetNamesM = mkLetNamesB

  addStms stms =
    BinderT $ modify $ \(cur_stms,scope) ->
    (cur_stms<>stms, scope `M.union` scopeOf stms)

  collectStms m = do
    (old_stms, old_scope) <- BinderT get
    BinderT $ put (mempty, old_scope)
    x <- m
    (new_stms, _) <- BinderT get
    BinderT $ put (old_stms, old_scope)
    return (x, new_stms)

-- | Run a binder action given an initial scope, returning a value and
-- the statements added ('addStm') during the action.
runBinderT :: MonadFreshNames m =>
              BinderT lore m a
           -> Scope lore
           -> m (a, Stms lore)
runBinderT (BinderT m) scope = do
  (x, (stms, _)) <- runStateT m (mempty, scope)
  return (x, stms)

-- | Like 'runBinderT', but return only the statements.
runBinderT_ :: MonadFreshNames m =>
               BinderT lore m () -> Scope lore -> m (Stms lore)
runBinderT_ m = fmap snd . runBinderT m

-- | Like 'runBinderT', but get the initial scope from the current
-- monad.
runBinderT' :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) =>
               BinderT lore m a
            -> m (a, Stms lore)
runBinderT' m = do
  scope <- askScope
  runBinderT m $ castScope scope

-- | Like 'runBinderT_', but get the initial scope from the current
-- monad.
runBinderT'_ :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) =>
                BinderT lore m a -> m (Stms lore)
runBinderT'_ = fmap snd . runBinderT'

-- | Run a binder action, returning a value and the statements added
-- ('addStm') during the action.  Assumes that the current monad
-- provides initial scope and name source.
runBinder :: (MonadFreshNames m,
              HasScope somelore m, SameScope somelore lore) =>
              Binder lore a
           -> m (a, Stms lore)
runBinder m = do
  types <- askScope
  modifyNameSource $ runState $ runBinderT m $ castScope types

-- | Like 'runBinder', but throw away the result and just return the
-- added statements.
runBinder_ :: (MonadFreshNames m,
               HasScope somelore m, SameScope somelore lore) =>
              Binder lore a
           -> m (Stms lore)
runBinder_ = fmap snd . runBinder

-- | Run a binder that produces a t'Body', and prefix that t'Body' by
-- the statements produced during execution of the action.
runBodyBinder :: (Bindable lore, MonadFreshNames m,
                  HasScope somelore m, SameScope somelore lore) =>
                 Binder lore (Body lore) -> m (Body lore)
runBodyBinder = fmap (uncurry $ flip insertStms) . runBinder

-- Utility instance defintions for MTL classes.  These require
-- UndecidableInstances, but save on typing elsewhere.

mapInner :: Monad m =>
            (m (a, (Stms lore, Scope lore))
             -> m (b, (Stms lore, Scope lore)))
         -> BinderT lore m a -> BinderT lore m b
mapInner f (BinderT m) = BinderT $ do
  s <- get
  (x, s') <- lift $ f $ runStateT m s
  put s'
  return x

instance MonadReader r m => MonadReader r (BinderT lore m) where
  ask = BinderT $ lift ask
  local f = mapInner $ local f

instance MonadState s m => MonadState s (BinderT lore m) where
  get = BinderT $ lift get
  put = BinderT . lift . put

instance MonadWriter w m => MonadWriter w (BinderT lore m) where
  tell = BinderT . lift . tell
  pass = mapInner $ \m -> pass $ do
    ((x, f), s) <- m
    return ((x, s), f)
  listen = mapInner $ \m -> do
    ((x, s), y) <- listen m
    return ((x, y), s)

instance MonadError e m => MonadError e (BinderT lore m) where
  throwError = lift . throwError
  catchError (BinderT m) f =
    BinderT $ catchError m $ unBinder . f
    where unBinder (BinderT m') = m'