{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | This module defines a convenience monad/typeclass for building
-- ASTs.  The fundamental building block is 'BuilderT' and its
-- execution functions, but it is usually easier to use 'Builder'.
--
-- See "Futhark.Construct" for a high-level description.
module Futhark.Builder
  ( -- * A concrete @MonadBuilder@ monad.
    BuilderT,
    runBuilderT,
    runBuilderT_,
    runBuilderT',
    runBuilderT'_,
    BuilderOps (..),
    Builder,
    runBuilder,
    runBuilder_,
    runBodyBuilder,
    runLambdaBuilder,

    -- * The 'MonadBuilder' typeclass
    module Futhark.Builder.Class,
  )
where

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

-- | A 'BuilderT' (and by extension, a 'Builder') is only an instance of
-- 'MonadBuilder' for representations that implement this type class,
-- which contains methods for constructing statements.
class (ASTRep rep) => BuilderOps rep where
  mkExpDecB ::
    (MonadBuilder m, Rep m ~ rep) =>
    Pat (LetDec rep) ->
    Exp rep ->
    m (ExpDec rep)
  mkBodyB ::
    (MonadBuilder m, Rep m ~ rep) =>
    Stms rep ->
    Result ->
    m (Body rep)
  mkLetNamesB ::
    (MonadBuilder m, Rep m ~ rep) =>
    [VName] ->
    Exp rep ->
    m (Stm rep)

  default mkExpDecB ::
    (MonadBuilder m, Buildable rep) =>
    Pat (LetDec rep) ->
    Exp rep ->
    m (ExpDec rep)
  mkExpDecB Pat (LetDec rep)
pat Exp rep
e = ExpDec rep -> m (ExpDec rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ExpDec rep -> m (ExpDec rep)) -> ExpDec rep -> m (ExpDec rep)
forall a b. (a -> b) -> a -> b
$ Pat (LetDec rep) -> Exp rep -> ExpDec rep
forall rep.
Buildable rep =>
Pat (LetDec rep) -> Exp rep -> ExpDec rep
mkExpDec Pat (LetDec rep)
pat Exp rep
e

  default mkBodyB ::
    (MonadBuilder m, Buildable rep) =>
    Stms rep ->
    Result ->
    m (Body rep)
  mkBodyB Stms rep
stms Result
res = Body rep -> m (Body rep)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Body rep -> m (Body rep)) -> Body rep -> m (Body rep)
forall a b. (a -> b) -> a -> b
$ Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms rep
stms Result
res

  default mkLetNamesB ::
    (MonadBuilder m, Rep m ~ rep, Buildable rep) =>
    [VName] ->
    Exp rep ->
    m (Stm rep)
  mkLetNamesB = [VName] -> Exp rep -> m (Stm rep)
forall rep (m :: * -> *).
(Buildable rep, MonadFreshNames m, HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
forall (m :: * -> *).
(MonadFreshNames m, HasScope rep m) =>
[VName] -> Exp rep -> m (Stm rep)
mkLetNames

-- | A monad transformer that tracks statements and provides a
-- 'MonadBuilder' 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 'MonadBuilder' from
-- scratch, then it is likely that you are making a mistake.
newtype BuilderT rep m a = BuilderT (StateT (Stms rep, Scope rep) m a)
  deriving ((forall a b. (a -> b) -> BuilderT rep m a -> BuilderT rep m b)
-> (forall a b. a -> BuilderT rep m b -> BuilderT rep m a)
-> Functor (BuilderT rep m)
forall a b. a -> BuilderT rep m b -> BuilderT rep m a
forall a b. (a -> b) -> BuilderT rep m a -> BuilderT rep m b
forall rep (m :: * -> *) a b.
Functor m =>
a -> BuilderT rep m b -> BuilderT rep m a
forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> BuilderT rep m a -> BuilderT rep m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall rep (m :: * -> *) a b.
Functor m =>
(a -> b) -> BuilderT rep m a -> BuilderT rep m b
fmap :: forall a b. (a -> b) -> BuilderT rep m a -> BuilderT rep m b
$c<$ :: forall rep (m :: * -> *) a b.
Functor m =>
a -> BuilderT rep m b -> BuilderT rep m a
<$ :: forall a b. a -> BuilderT rep m b -> BuilderT rep m a
Functor, Applicative (BuilderT rep m)
Applicative (BuilderT rep m) =>
(forall a b.
 BuilderT rep m a -> (a -> BuilderT rep m b) -> BuilderT rep m b)
-> (forall a b.
    BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b)
-> (forall a. a -> BuilderT rep m a)
-> Monad (BuilderT rep m)
forall a. a -> BuilderT rep m a
forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
forall a b.
BuilderT rep m a -> (a -> BuilderT rep m b) -> BuilderT rep m b
forall rep (m :: * -> *). Monad m => Applicative (BuilderT rep m)
forall rep (m :: * -> *) a. Monad m => a -> BuilderT rep m a
forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> (a -> BuilderT rep m b) -> BuilderT 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
$c>>= :: forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> (a -> BuilderT rep m b) -> BuilderT rep m b
>>= :: forall a b.
BuilderT rep m a -> (a -> BuilderT rep m b) -> BuilderT rep m b
$c>> :: forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
>> :: forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
$creturn :: forall rep (m :: * -> *) a. Monad m => a -> BuilderT rep m a
return :: forall a. a -> BuilderT rep m a
Monad, Functor (BuilderT rep m)
Functor (BuilderT rep m) =>
(forall a. a -> BuilderT rep m a)
-> (forall a b.
    BuilderT rep m (a -> b) -> BuilderT rep m a -> BuilderT rep m b)
-> (forall a b c.
    (a -> b -> c)
    -> BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m c)
-> (forall a b.
    BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b)
-> (forall a b.
    BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m a)
-> Applicative (BuilderT rep m)
forall a. a -> BuilderT rep m a
forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m a
forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
forall a b.
BuilderT rep m (a -> b) -> BuilderT rep m a -> BuilderT rep m b
forall a b c.
(a -> b -> c)
-> BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m c
forall rep (m :: * -> *). Monad m => Functor (BuilderT rep m)
forall rep (m :: * -> *) a. Monad m => a -> BuilderT rep m a
forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m a
forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m (a -> b) -> BuilderT rep m a -> BuilderT rep m b
forall rep (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> BuilderT rep m a -> BuilderT rep m b -> BuilderT 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
$cpure :: forall rep (m :: * -> *) a. Monad m => a -> BuilderT rep m a
pure :: forall a. a -> BuilderT rep m a
$c<*> :: forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m (a -> b) -> BuilderT rep m a -> BuilderT rep m b
<*> :: forall a b.
BuilderT rep m (a -> b) -> BuilderT rep m a -> BuilderT rep m b
$cliftA2 :: forall rep (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m c
liftA2 :: forall a b c.
(a -> b -> c)
-> BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m c
$c*> :: forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
*> :: forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m b
$c<* :: forall rep (m :: * -> *) a b.
Monad m =>
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m a
<* :: forall a b.
BuilderT rep m a -> BuilderT rep m b -> BuilderT rep m a
Applicative)

instance MonadTrans (BuilderT rep) where
  lift :: forall (m :: * -> *) a. Monad m => m a -> BuilderT rep m a
lift = StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m a -> BuilderT rep m a)
-> (m a -> StateT (Stms rep, Scope rep) m a)
-> m a
-> BuilderT rep m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> StateT (Stms rep, Scope rep) m a
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Stms rep, Scope rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

-- | The most commonly used builder monad.
type Builder rep = BuilderT rep (State VNameSource)

instance (MonadFreshNames m) => MonadFreshNames (BuilderT rep m) where
  getNameSource :: BuilderT rep m VNameSource
getNameSource = m VNameSource -> BuilderT rep m VNameSource
forall (m :: * -> *) a. Monad m => m a -> BuilderT rep m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m VNameSource
forall (m :: * -> *). MonadFreshNames m => m VNameSource
getNameSource
  putNameSource :: VNameSource -> BuilderT rep m ()
putNameSource = m () -> BuilderT rep m ()
forall (m :: * -> *) a. Monad m => m a -> BuilderT rep m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> BuilderT rep m ())
-> (VNameSource -> m ()) -> VNameSource -> BuilderT rep m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VNameSource -> m ()
forall (m :: * -> *). MonadFreshNames m => VNameSource -> m ()
putNameSource

instance (ASTRep rep, Monad m) => HasScope rep (BuilderT rep m) where
  lookupType :: VName -> BuilderT rep m Type
lookupType VName
name = do
    t <- StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep))
-> BuilderT rep m (Maybe (NameInfo rep))
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep))
 -> BuilderT rep m (Maybe (NameInfo rep)))
-> StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep))
-> BuilderT rep m (Maybe (NameInfo rep))
forall a b. (a -> b) -> a -> b
$ ((Stms rep, Scope rep) -> Maybe (NameInfo rep))
-> StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((Stms rep, Scope rep) -> Maybe (NameInfo rep))
 -> StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep)))
-> ((Stms rep, Scope rep) -> Maybe (NameInfo rep))
-> StateT (Stms rep, Scope rep) m (Maybe (NameInfo rep))
forall a b. (a -> b) -> a -> b
$ VName -> Scope rep -> Maybe (NameInfo rep)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
name (Scope rep -> Maybe (NameInfo rep))
-> ((Stms rep, Scope rep) -> Scope rep)
-> (Stms rep, Scope rep)
-> Maybe (NameInfo rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms rep, Scope rep) -> Scope rep
forall a b. (a, b) -> b
snd
    case t of
      Maybe (NameInfo rep)
Nothing -> do
        known <- StateT (Stms rep, Scope rep) m [VName] -> BuilderT rep m [VName]
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m [VName] -> BuilderT rep m [VName])
-> StateT (Stms rep, Scope rep) m [VName] -> BuilderT rep m [VName]
forall a b. (a -> b) -> a -> b
$ ((Stms rep, Scope rep) -> [VName])
-> StateT (Stms rep, Scope rep) m [VName]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((Stms rep, Scope rep) -> [VName])
 -> StateT (Stms rep, Scope rep) m [VName])
-> ((Stms rep, Scope rep) -> [VName])
-> StateT (Stms rep, Scope rep) m [VName]
forall a b. (a -> b) -> a -> b
$ Scope rep -> [VName]
forall k a. Map k a -> [k]
M.keys (Scope rep -> [VName])
-> ((Stms rep, Scope rep) -> Scope rep)
-> (Stms rep, Scope rep)
-> [VName]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Stms rep, Scope rep) -> Scope rep
forall a b. (a, b) -> b
snd
        error . unlines $
          [ "BuilderT.lookupType: unknown variable " ++ prettyString name,
            "Known variables: ",
            unwords $ map prettyString known
          ]
      Just NameInfo rep
t' -> Type -> BuilderT rep m Type
forall a. a -> BuilderT rep m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> BuilderT rep m Type) -> Type -> BuilderT rep m Type
forall a b. (a -> b) -> a -> b
$ NameInfo rep -> Type
forall t. Typed t => t -> Type
typeOf NameInfo rep
t'
  askScope :: BuilderT rep m (Scope rep)
askScope = StateT (Stms rep, Scope rep) m (Scope rep)
-> BuilderT rep m (Scope rep)
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m (Scope rep)
 -> BuilderT rep m (Scope rep))
-> StateT (Stms rep, Scope rep) m (Scope rep)
-> BuilderT rep m (Scope rep)
forall a b. (a -> b) -> a -> b
$ ((Stms rep, Scope rep) -> Scope rep)
-> StateT (Stms rep, Scope rep) m (Scope rep)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (Stms rep, Scope rep) -> Scope rep
forall a b. (a, b) -> b
snd

instance (ASTRep rep, Monad m) => LocalScope rep (BuilderT rep m) where
  localScope :: forall a. Scope rep -> BuilderT rep m a -> BuilderT rep m a
localScope Scope rep
types (BuilderT StateT (Stms rep, Scope rep) m a
m) = StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m a -> BuilderT rep m a)
-> StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
forall a b. (a -> b) -> a -> b
$ do
    ((Stms rep, Scope rep) -> (Stms rep, Scope rep))
-> StateT (Stms rep, Scope rep) m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((Stms rep, Scope rep) -> (Stms rep, Scope rep))
 -> StateT (Stms rep, Scope rep) m ())
-> ((Stms rep, Scope rep) -> (Stms rep, Scope rep))
-> StateT (Stms rep, Scope rep) m ()
forall a b. (a -> b) -> a -> b
$ (Scope rep -> Scope rep)
-> (Stms rep, Scope rep) -> (Stms rep, Scope rep)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Scope rep -> Scope rep -> Scope rep
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Scope rep
types)
    x <- StateT (Stms rep, Scope rep) m a
m
    modify $ second (`M.difference` types)
    pure x

instance
  (MonadFreshNames m, BuilderOps rep) =>
  MonadBuilder (BuilderT rep m)
  where
  type Rep (BuilderT rep m) = rep
  mkExpDecM :: Pat (LetDec (Rep (BuilderT rep m)))
-> Exp (Rep (BuilderT rep m))
-> BuilderT rep m (ExpDec (Rep (BuilderT rep m)))
mkExpDecM = Pat (LetDec rep) -> Exp rep -> BuilderT rep m (ExpDec rep)
Pat (LetDec (Rep (BuilderT rep m)))
-> Exp (Rep (BuilderT rep m))
-> BuilderT rep m (ExpDec (Rep (BuilderT rep m)))
forall rep (m :: * -> *).
(BuilderOps rep, MonadBuilder m, Rep m ~ rep) =>
Pat (LetDec rep) -> Exp rep -> m (ExpDec rep)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ rep) =>
Pat (LetDec rep) -> Exp rep -> m (ExpDec rep)
mkExpDecB
  mkBodyM :: Stms (Rep (BuilderT rep m))
-> Result -> BuilderT rep m (Body (Rep (BuilderT rep m)))
mkBodyM = Stms rep -> Result -> BuilderT rep m (Body rep)
Stms (Rep (BuilderT rep m))
-> Result -> BuilderT rep m (Body (Rep (BuilderT rep m)))
forall rep (m :: * -> *).
(BuilderOps rep, MonadBuilder m, Rep m ~ rep) =>
Stms rep -> Result -> m (Body rep)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ rep) =>
Stms rep -> Result -> m (Body rep)
mkBodyB
  mkLetNamesM :: [VName]
-> Exp (Rep (BuilderT rep m))
-> BuilderT rep m (Stm (Rep (BuilderT rep m)))
mkLetNamesM = [VName] -> Exp rep -> BuilderT rep m (Stm rep)
[VName]
-> Exp (Rep (BuilderT rep m))
-> BuilderT rep m (Stm (Rep (BuilderT rep m)))
forall rep (m :: * -> *).
(BuilderOps rep, MonadBuilder m, Rep m ~ rep) =>
[VName] -> Exp rep -> m (Stm rep)
forall (m :: * -> *).
(MonadBuilder m, Rep m ~ rep) =>
[VName] -> Exp rep -> m (Stm rep)
mkLetNamesB

  addStms :: Stms (Rep (BuilderT rep m)) -> BuilderT rep m ()
addStms Stms (Rep (BuilderT rep m))
stms =
    StateT (Stms rep, Map VName (NameInfo rep)) m ()
-> BuilderT rep m ()
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Map VName (NameInfo rep)) m ()
 -> BuilderT rep m ())
-> StateT (Stms rep, Map VName (NameInfo rep)) m ()
-> BuilderT rep m ()
forall a b. (a -> b) -> a -> b
$
      ((Stms rep, Map VName (NameInfo rep))
 -> (Stms rep, Map VName (NameInfo rep)))
-> StateT (Stms rep, Map VName (NameInfo rep)) m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (((Stms rep, Map VName (NameInfo rep))
  -> (Stms rep, Map VName (NameInfo rep)))
 -> StateT (Stms rep, Map VName (NameInfo rep)) m ())
-> ((Stms rep, Map VName (NameInfo rep))
    -> (Stms rep, Map VName (NameInfo rep)))
-> StateT (Stms rep, Map VName (NameInfo rep)) m ()
forall a b. (a -> b) -> a -> b
$ \(Stms rep
cur_stms, Map VName (NameInfo rep)
scope) ->
        (Stms rep
cur_stms Stms rep -> Stms rep -> Stms rep
forall a. Semigroup a => a -> a -> a
<> Stms rep
Stms (Rep (BuilderT rep m))
stms, Map VName (NameInfo rep)
scope Map VName (NameInfo rep)
-> Map VName (NameInfo rep) -> Map VName (NameInfo rep)
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Stms rep -> Map VName (NameInfo rep)
forall rep a. Scoped rep a => a -> Scope rep
scopeOf Stms rep
Stms (Rep (BuilderT rep m))
stms)

  collectStms :: forall a.
BuilderT rep m a -> BuilderT rep m (a, Stms (Rep (BuilderT rep m)))
collectStms BuilderT rep m a
m = do
    (old_stms, old_scope) <- StateT
  (Stms rep, Map VName (NameInfo rep))
  m
  (Stms rep, Map VName (NameInfo rep))
-> BuilderT rep m (Stms rep, Map VName (NameInfo rep))
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT StateT
  (Stms rep, Map VName (NameInfo rep))
  m
  (Stms rep, Map VName (NameInfo rep))
forall s (m :: * -> *). MonadState s m => m s
get
    BuilderT $ put (mempty, old_scope)
    x <- m
    (new_stms, _) <- BuilderT get
    BuilderT $ put (old_stms, old_scope)
    pure (x, new_stms)

-- | Run a builder action given an initial scope, returning a value and
-- the statements added ('addStm') during the action.
runBuilderT ::
  (MonadFreshNames m) =>
  BuilderT rep m a ->
  Scope rep ->
  m (a, Stms rep)
runBuilderT :: forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT (BuilderT StateT (Stms rep, Scope rep) m a
m) Scope rep
scope = do
  (x, (stms, _)) <- StateT (Stms rep, Scope rep) m a
-> (Stms rep, Scope rep) -> m (a, (Stms rep, Scope rep))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (Stms rep, Scope rep) m a
m (Stms rep
forall a. Monoid a => a
mempty, Scope rep
scope)
  pure (x, stms)

-- | Like 'runBuilderT', but return only the statements.
runBuilderT_ ::
  (MonadFreshNames m) =>
  BuilderT rep m () ->
  Scope rep ->
  m (Stms rep)
runBuilderT_ :: forall (m :: * -> *) rep.
MonadFreshNames m =>
BuilderT rep m () -> Scope rep -> m (Stms rep)
runBuilderT_ BuilderT rep m ()
m = (((), Stms rep) -> Stms rep) -> m ((), Stms rep) -> m (Stms rep)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((), Stms rep) -> Stms rep
forall a b. (a, b) -> b
snd (m ((), Stms rep) -> m (Stms rep))
-> (Scope rep -> m ((), Stms rep)) -> Scope rep -> m (Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuilderT rep m () -> Scope rep -> m ((), Stms rep)
forall (m :: * -> *) rep a.
MonadFreshNames m =>
BuilderT rep m a -> Scope rep -> m (a, Stms rep)
runBuilderT BuilderT rep m ()
m

-- | Like 'runBuilderT', but get the initial scope from the current
-- monad.
runBuilderT' ::
  (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
  BuilderT rep m a ->
  m (a, Stms rep)
runBuilderT' :: forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BuilderT rep m a -> m (a, Stms rep)
runBuilderT' BuilderT rep m a
m = do
  scope <- m (Scope somerep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  runBuilderT m $ castScope scope

-- | Like 'runBuilderT_', but get the initial scope from the current
-- monad.
runBuilderT'_ ::
  (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
  BuilderT rep m a ->
  m (Stms rep)
runBuilderT'_ :: forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BuilderT rep m a -> m (Stms rep)
runBuilderT'_ = ((a, Stms rep) -> Stms rep) -> m (a, Stms rep) -> m (Stms rep)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Stms rep) -> Stms rep
forall a b. (a, b) -> b
snd (m (a, Stms rep) -> m (Stms rep))
-> (BuilderT rep m a -> m (a, Stms rep))
-> BuilderT rep m a
-> m (Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BuilderT rep m a -> m (a, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
BuilderT rep m a -> m (a, Stms rep)
runBuilderT'

-- | Run a builder action, returning a value and the statements added
-- ('addStm') during the action.  Assumes that the current monad
-- provides initial scope and name source.
runBuilder ::
  (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
  Builder rep a ->
  m (a, Stms rep)
runBuilder :: forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder Builder rep a
m = do
  types <- m (Scope somerep)
forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
askScope
  modifyNameSource $ runState $ runBuilderT m $ castScope types

-- | Like 'runBuilder', but throw away the result and just return the
-- added statements.
runBuilder_ ::
  (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
  Builder rep a ->
  m (Stms rep)
runBuilder_ :: forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (Stms rep)
runBuilder_ = ((a, Stms rep) -> Stms rep) -> m (a, Stms rep) -> m (Stms rep)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Stms rep) -> Stms rep
forall a b. (a, b) -> b
snd (m (a, Stms rep) -> m (Stms rep))
-> (Builder rep a -> m (a, Stms rep))
-> Builder rep a
-> m (Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder rep a -> m (a, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder

-- | Run a builder that produces a 'Result' and construct a body that
-- contains that result alongside the statements produced during the
-- builder.
runBodyBuilder ::
  ( Buildable rep,
    MonadFreshNames m,
    HasScope somerep m,
    SameScope somerep rep
  ) =>
  Builder rep Result ->
  m (Body rep)
runBodyBuilder :: forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
Builder rep Result -> m (Body rep)
runBodyBuilder =
  ((Body rep, Stms rep) -> Body rep)
-> m (Body rep, Stms rep) -> m (Body rep)
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Body rep -> Stms rep -> Body rep)
-> (Body rep, Stms rep) -> Body rep
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((Body rep -> Stms rep -> Body rep)
 -> (Body rep, Stms rep) -> Body rep)
-> (Body rep -> Stms rep -> Body rep)
-> (Body rep, Stms rep)
-> Body rep
forall a b. (a -> b) -> a -> b
$ (Stms rep -> Body rep -> Body rep)
-> Body rep -> Stms rep -> Body rep
forall a b c. (a -> b -> c) -> b -> a -> c
flip Stms rep -> Body rep -> Body rep
forall rep. Buildable rep => Stms rep -> Body rep -> Body rep
insertStms) (m (Body rep, Stms rep) -> m (Body rep))
-> (Builder rep Result -> m (Body rep, Stms rep))
-> Builder rep Result
-> m (Body rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder rep (Body rep) -> m (Body rep, Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder rep (Body rep) -> m (Body rep, Stms rep))
-> (Builder rep Result -> Builder rep (Body rep))
-> Builder rep Result
-> m (Body rep, Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Result -> Body rep)
-> Builder rep Result -> Builder rep (Body rep)
forall a b.
(a -> b)
-> BuilderT rep (StateT VNameSource Identity) a
-> BuilderT rep (StateT VNameSource Identity) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Stms rep -> Result -> Body rep
forall rep. Buildable rep => Stms rep -> Result -> Body rep
mkBody Stms rep
forall a. Monoid a => a
mempty)

-- | Given lambda parameters, Run a builder action that produces the
-- statements and returns the 'Result' of the lambda body.
runLambdaBuilder ::
  ( Buildable rep,
    MonadFreshNames m,
    HasScope somerep m,
    SameScope somerep rep
  ) =>
  [LParam rep] ->
  Builder rep Result ->
  m (Lambda rep)
runLambdaBuilder :: forall rep (m :: * -> *) somerep.
(Buildable rep, MonadFreshNames m, HasScope somerep m,
 SameScope somerep rep) =>
[LParam rep] -> Builder rep Result -> m (Lambda rep)
runLambdaBuilder [LParam rep]
params Builder rep Result
m = do
  ((res, ret), stms) <- Builder rep (Result, [Type]) -> m ((Result, [Type]), Stms rep)
forall (m :: * -> *) somerep rep a.
(MonadFreshNames m, HasScope somerep m, SameScope somerep rep) =>
Builder rep a -> m (a, Stms rep)
runBuilder (Builder rep (Result, [Type]) -> m ((Result, [Type]), Stms rep))
-> (Builder rep (Result, [Type]) -> Builder rep (Result, [Type]))
-> Builder rep (Result, [Type])
-> m ((Result, [Type]), Stms rep)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Scope rep
-> Builder rep (Result, [Type]) -> Builder rep (Result, [Type])
forall a.
Scope rep
-> BuilderT rep (StateT VNameSource Identity) a
-> BuilderT rep (StateT VNameSource Identity) a
forall rep (m :: * -> *) a.
LocalScope rep m =>
Scope rep -> m a -> m a
localScope ([LParam rep] -> Scope rep
forall rep dec. (LParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfLParams [LParam rep]
params) (Builder rep (Result, [Type]) -> m ((Result, [Type]), Stms rep))
-> Builder rep (Result, [Type]) -> m ((Result, [Type]), Stms rep)
forall a b. (a -> b) -> a -> b
$ do
    res <- Builder rep Result
m
    ret <- mapM subExpResType res
    pure (res, ret)
  pure $ Lambda params ret $ mkBody stms res

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

mapInner ::
  (Monad m) =>
  ( m (a, (Stms rep, Scope rep)) ->
    m (b, (Stms rep, Scope rep))
  ) ->
  BuilderT rep m a ->
  BuilderT rep m b
mapInner :: forall (m :: * -> *) a rep b.
Monad m =>
(m (a, (Stms rep, Scope rep)) -> m (b, (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m b
mapInner m (a, (Stms rep, Scope rep)) -> m (b, (Stms rep, Scope rep))
f (BuilderT StateT (Stms rep, Scope rep) m a
m) = StateT (Stms rep, Scope rep) m b -> BuilderT rep m b
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m b -> BuilderT rep m b)
-> StateT (Stms rep, Scope rep) m b -> BuilderT rep m b
forall a b. (a -> b) -> a -> b
$ do
  s <- StateT (Stms rep, Scope rep) m (Stms rep, Scope rep)
forall s (m :: * -> *). MonadState s m => m s
get
  (x, s') <- lift $ f $ runStateT m s
  put s'
  pure x

instance (MonadReader r m) => MonadReader r (BuilderT rep m) where
  ask :: BuilderT rep m r
ask = StateT (Stms rep, Scope rep) m r -> BuilderT rep m r
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m r -> BuilderT rep m r)
-> StateT (Stms rep, Scope rep) m r -> BuilderT rep m r
forall a b. (a -> b) -> a -> b
$ m r -> StateT (Stms rep, Scope rep) m r
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Stms rep, Scope rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m r
forall r (m :: * -> *). MonadReader r m => m r
ask
  local :: forall a. (r -> r) -> BuilderT rep m a -> BuilderT rep m a
local r -> r
f = (m (a, (Stms rep, Scope rep)) -> m (a, (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m a
forall (m :: * -> *) a rep b.
Monad m =>
(m (a, (Stms rep, Scope rep)) -> m (b, (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m b
mapInner ((m (a, (Stms rep, Scope rep)) -> m (a, (Stms rep, Scope rep)))
 -> BuilderT rep m a -> BuilderT rep m a)
-> (m (a, (Stms rep, Scope rep)) -> m (a, (Stms rep, Scope rep)))
-> BuilderT rep m a
-> BuilderT rep m a
forall a b. (a -> b) -> a -> b
$ (r -> r)
-> m (a, (Stms rep, Scope rep)) -> m (a, (Stms rep, Scope rep))
forall a. (r -> r) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f

instance (MonadState s m) => MonadState s (BuilderT rep m) where
  get :: BuilderT rep m s
get = StateT (Stms rep, Scope rep) m s -> BuilderT rep m s
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m s -> BuilderT rep m s)
-> StateT (Stms rep, Scope rep) m s -> BuilderT rep m s
forall a b. (a -> b) -> a -> b
$ m s -> StateT (Stms rep, Scope rep) m s
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Stms rep, Scope rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m s
forall s (m :: * -> *). MonadState s m => m s
get
  put :: s -> BuilderT rep m ()
put = StateT (Stms rep, Scope rep) m () -> BuilderT rep m ()
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m () -> BuilderT rep m ())
-> (s -> StateT (Stms rep, Scope rep) m ())
-> s
-> BuilderT rep m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m () -> StateT (Stms rep, Scope rep) m ()
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Stms rep, Scope rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateT (Stms rep, Scope rep) m ())
-> (s -> m ()) -> s -> StateT (Stms rep, Scope rep) m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put

instance (MonadWriter w m) => MonadWriter w (BuilderT rep m) where
  tell :: w -> BuilderT rep m ()
tell = StateT (Stms rep, Scope rep) m () -> BuilderT rep m ()
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m () -> BuilderT rep m ())
-> (w -> StateT (Stms rep, Scope rep) m ())
-> w
-> BuilderT rep m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m () -> StateT (Stms rep, Scope rep) m ()
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Stms rep, Scope rep) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateT (Stms rep, Scope rep) m ())
-> (w -> m ()) -> w -> StateT (Stms rep, Scope rep) m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. w -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell
  pass :: forall a. BuilderT rep m (a, w -> w) -> BuilderT rep m a
pass = (m ((a, w -> w), (Stms rep, Scope rep))
 -> m (a, (Stms rep, Scope rep)))
-> BuilderT rep m (a, w -> w) -> BuilderT rep m a
forall (m :: * -> *) a rep b.
Monad m =>
(m (a, (Stms rep, Scope rep)) -> m (b, (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m b
mapInner ((m ((a, w -> w), (Stms rep, Scope rep))
  -> m (a, (Stms rep, Scope rep)))
 -> BuilderT rep m (a, w -> w) -> BuilderT rep m a)
-> (m ((a, w -> w), (Stms rep, Scope rep))
    -> m (a, (Stms rep, Scope rep)))
-> BuilderT rep m (a, w -> w)
-> BuilderT rep m a
forall a b. (a -> b) -> a -> b
$ \m ((a, w -> w), (Stms rep, Scope rep))
m -> m ((a, (Stms rep, Scope rep)), w -> w)
-> m (a, (Stms rep, Scope rep))
forall a. m (a, w -> w) -> m a
forall w (m :: * -> *) a. MonadWriter w m => m (a, w -> w) -> m a
pass (m ((a, (Stms rep, Scope rep)), w -> w)
 -> m (a, (Stms rep, Scope rep)))
-> m ((a, (Stms rep, Scope rep)), w -> w)
-> m (a, (Stms rep, Scope rep))
forall a b. (a -> b) -> a -> b
$ do
    ((x, f), s) <- m ((a, w -> w), (Stms rep, Scope rep))
m
    pure ((x, s), f)
  listen :: forall a. BuilderT rep m a -> BuilderT rep m (a, w)
listen = (m (a, (Stms rep, Scope rep)) -> m ((a, w), (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m (a, w)
forall (m :: * -> *) a rep b.
Monad m =>
(m (a, (Stms rep, Scope rep)) -> m (b, (Stms rep, Scope rep)))
-> BuilderT rep m a -> BuilderT rep m b
mapInner ((m (a, (Stms rep, Scope rep))
  -> m ((a, w), (Stms rep, Scope rep)))
 -> BuilderT rep m a -> BuilderT rep m (a, w))
-> (m (a, (Stms rep, Scope rep))
    -> m ((a, w), (Stms rep, Scope rep)))
-> BuilderT rep m a
-> BuilderT rep m (a, w)
forall a b. (a -> b) -> a -> b
$ \m (a, (Stms rep, Scope rep))
m -> do
    ((x, s), y) <- m (a, (Stms rep, Scope rep)) -> m ((a, (Stms rep, Scope rep)), w)
forall a. m a -> m (a, w)
forall w (m :: * -> *) a. MonadWriter w m => m a -> m (a, w)
listen m (a, (Stms rep, Scope rep))
m
    pure ((x, y), s)

instance (MonadError e m) => MonadError e (BuilderT rep m) where
  throwError :: forall a. e -> BuilderT rep m a
throwError = m a -> BuilderT rep m a
forall (m :: * -> *) a. Monad m => m a -> BuilderT rep m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> BuilderT rep m a) -> (e -> m a) -> e -> BuilderT rep m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> m a
forall a. e -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError
  catchError :: forall a.
BuilderT rep m a -> (e -> BuilderT rep m a) -> BuilderT rep m a
catchError (BuilderT StateT (Stms rep, Scope rep) m a
m) e -> BuilderT rep m a
f =
    StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
forall rep (m :: * -> *) a.
StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
BuilderT (StateT (Stms rep, Scope rep) m a -> BuilderT rep m a)
-> StateT (Stms rep, Scope rep) m a -> BuilderT rep m a
forall a b. (a -> b) -> a -> b
$ StateT (Stms rep, Scope rep) m a
-> (e -> StateT (Stms rep, Scope rep) m a)
-> StateT (Stms rep, Scope rep) m a
forall a.
StateT (Stms rep, Scope rep) m a
-> (e -> StateT (Stms rep, Scope rep) m a)
-> StateT (Stms rep, Scope rep) m a
forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
catchError StateT (Stms rep, Scope rep) m a
m ((e -> StateT (Stms rep, Scope rep) m a)
 -> StateT (Stms rep, Scope rep) m a)
-> (e -> StateT (Stms rep, Scope rep) m a)
-> StateT (Stms rep, Scope rep) m a
forall a b. (a -> b) -> a -> b
$ BuilderT rep m a -> StateT (Stms rep, Scope rep) m a
forall {rep} {m :: * -> *} {a}.
BuilderT rep m a -> StateT (Stms rep, Scope rep) m a
unBuilder (BuilderT rep m a -> StateT (Stms rep, Scope rep) m a)
-> (e -> BuilderT rep m a) -> e -> StateT (Stms rep, Scope rep) m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> BuilderT rep m a
f
    where
      unBuilder :: BuilderT rep m a -> StateT (Stms rep, Scope rep) m a
unBuilder (BuilderT StateT (Stms rep, Scope rep) m a
m') = StateT (Stms rep, Scope rep) m a
m'