LetRec.hs

{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecursiveDo         #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Language.Haskell.TH.LetRec (
    letrecE,
    typedLetrecE,
) where

import Control.Monad.Fix              (MonadFix)
import Control.Monad.Trans.Class      (lift)
import Control.Monad.Trans.State.Lazy (StateT, get, modify, runStateT)
import Language.Haskell.TH.Lib        (letE, normalB, sigD, valD, varE, varP)
import Language.Haskell.TH.Syntax     (Exp, Name, Quote (newName), Type)

import qualified Data.Map.Lazy as Map

-- $setup
-- >>> :set -XTemplateHaskell -XRankNTypes -XScopedTypeVariables
-- >>> import Language.Haskell.TH.Syntax as TH
-- >>> import Language.Haskell.TH.Lib    as TH
-- >>> import Language.Haskell.TH.Ppr    as TH
-- >>> let mangleName :: TH.Name -> TH.Name; mangleName n | n == ''()  = TH.Name (TH.OccName "Unit") TH.NameS; mangleName (TH.Name occ TH.NameG {}) = TH.Name occ TH.NameS; mangleName n = n
-- >>> import Data.Data (Data, Typeable, cast, gmapT)
-- >>> mkT :: (Typeable a, Typeable b) => (b -> b) -> a -> a; mkT f x = maybe x (maybe x id . cast . f) (cast x)
-- >>> let everywhere :: (forall a. Data a => a -> a) -> (forall b. Data b => b -> b); everywhere f = go where go :: forall c. Data c => c -> c; go = f . gmapT go

-- | Generate potentially recursive let expression.
--
-- The 'Monad' constraint in generators forces to sequence
-- binding generation calls, thus allowing to do lazy binding generation.
--
-- Example of generating a list of alternating 'True' and 'False' values:
--
-- >>> let trueFalse = letrecE (\tag -> "go" ++ show tag) (\rec tag -> rec (not tag) >>= \next -> return [| $(TH.lift tag) : $next |]) ($ True)
--
-- The generated let-bindings look like:
--
-- >>> TH.ppr . everywhere (mkT mangleName) <$> trueFalse
-- let {goFalse_0 = False : goTrue_1; goTrue_1 = True : goFalse_0}
--  in goTrue_1
--
-- And when spliced it produces a list of alternative 'True' and 'False' values:
--
-- >>> take 10 $trueFalse
-- [True,False,True,False,True,False,True,False,True,False]
--
-- Another example where dynamic nature is visible is generating
-- fibonacci numbers:
--
-- >>> let fibRec rec tag = case tag of { 0 -> return [| 1 |]; 1 -> return [| 1 |]; _ -> do { minus1 <- rec (tag - 1); minus2 <- rec (tag - 2); return [| $minus1 + $minus2 |] }}
-- >>> let fib n = letrecE (\tag -> "fib" ++ show tag) fibRec ($ n)
--
-- The generated let-bindings look like:
-- >>> TH.ppr . everywhere (mkT mangleName) <$> fib 7
-- let {fib0_0 = 1;
--      fib1_1 = 1;
--      fib2_2 = fib1_1 + fib0_0;
--      fib3_3 = fib2_2 + fib1_1;
--      fib4_4 = fib3_3 + fib2_2;
--      fib5_5 = fib4_4 + fib3_3;
--      fib6_6 = fib5_5 + fib4_4;
--      fib7_7 = fib6_6 + fib5_5}
--  in fib7_7
--
-- And the result is expected:
--
-- >>> $(fib 7)
-- 21
--
letrecE
    :: forall q tag. (Ord tag, Quote q, MonadFix q)
    => (tag -> String)                                                   -- ^ tag naming function
    -> (forall m. Monad m => (tag -> m (q Exp)) -> (tag -> m (q Exp)))   -- ^ bindings generator (with recursive function)
    -> (forall m. Monad m => (tag -> m (q Exp)) -> m (q Exp))            -- ^ final expression generator
    -> q Exp                                                             -- ^ generated let expression.
letrecE nameOf = typedLetrecE nameOf (const Nothing)

-- | Generate potentially recursive let expression, with optional type annotations.
--
-- A fibonacci example:
-- >>> let fibRec rec tag = case tag of { 0 -> return [| 1 |]; 1 -> return [| 1 |]; _ -> do { minus1 <- rec (tag - 1); minus2 <- rec (tag - 2); return [| $minus1 + $minus2 |] }}
-- >>> let fib n = typedLetrecE (\tag -> "fib" ++ show tag) (\_ -> Just [t| Int |]) fibRec ($ n)
--
-- The generated let-bindings look like:
-- >>> TH.ppr . everywhere (mkT mangleName) <$> fib 7
-- let {fib0_0 :: Int;
--      fib0_0 = 1;
--      fib1_1 :: Int;
--      fib1_1 = 1;
--      fib2_2 :: Int;
--      fib2_2 = fib1_1 + fib0_0;
--      fib3_3 :: Int;
--      fib3_3 = fib2_2 + fib1_1;
--      fib4_4 :: Int;
--      fib4_4 = fib3_3 + fib2_2;
--      fib5_5 :: Int;
--      fib5_5 = fib4_4 + fib3_3;
--      fib6_6 :: Int;
--      fib6_6 = fib5_5 + fib4_4;
--      fib7_7 :: Int;
--      fib7_7 = fib6_6 + fib5_5}
--  in fib7_7
--
-- >>> $(fib 7)
-- 21
--
-- @since 0.1.1
--
typedLetrecE
    :: forall q tag. (Ord tag, Quote q, MonadFix q)
    => (tag -> String)                                                   -- ^ tag naming function
    -> (tag -> Maybe (q Type))                                           -- ^ binding type
    -> (forall m. Monad m => (tag -> m (q Exp)) -> (tag -> m (q Exp)))   -- ^ bindings generator (with recursive function)
    -> (forall m. Monad m => (tag -> m (q Exp)) -> m (q Exp))            -- ^ final expression generator
    -> q Exp                                                             -- ^ generated let expression.
typedLetrecE nameOf typeOf recf exprf = do
    (expr0, bindings) <- runStateT (exprf loop) Map.empty
    letE (concat
        [ [ sigD name ty' | Just ty' <- [ty] ] ++
          [ valD (varP name) (normalB expr) [] ]
        | (_tag, (name, ty, expr)) <- Map.toList bindings
        ])
        expr0
  where
    loop :: tag -> StateT (Map.Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
    loop tag = do
        m <- get
        case Map.lookup tag m of
            -- if name is already generated, return it.
            Just (name, _ty, _exp) -> return (varE name)

            -- otherwise generate new name, and insert it into the loop.
            Nothing -> mdo
                name <- lift (newName (nameOf tag))
                modify (Map.insert tag (name, typeOf tag, expr))
                expr <- recf loop tag
                return (varE name)