{-# 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
-- >>> import Language.Haskell.TH.Syntax as TH
-- >>> import Language.Haskell.TH.Lib    as TH
-- >>> import Language.Haskell.TH.Ppr    as TH

-- | 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 <$> trueFalse
-- let {goFalse_0 = GHC.Types.False GHC.Types.: goTrue_1;
--      goTrue_1 = GHC.Types.True GHC.Types.: 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 <$> fib 7
-- let {fib0_0 = 1;
--      fib1_1 = 1;
--      fib2_2 = fib1_1 GHC.Num.+ fib0_0;
--      fib3_3 = fib2_2 GHC.Num.+ fib1_1;
--      fib4_4 = fib3_3 GHC.Num.+ fib2_2;
--      fib5_5 = fib4_4 GHC.Num.+ fib3_3;
--      fib6_6 = fib5_5 GHC.Num.+ fib4_4;
--      fib7_7 = fib6_6 GHC.Num.+ 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 :: forall (q :: * -> *) tag.
(Ord tag, Quote q, MonadFix q) =>
(tag -> String)
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> tag -> m (q Exp))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> m (q Exp))
-> q Exp
letrecE tag -> String
nameOf = (tag -> String)
-> (tag -> Maybe (q Type))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> tag -> m (q Exp))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> m (q Exp))
-> q Exp
forall (q :: * -> *) tag.
(Ord tag, Quote q, MonadFix q) =>
(tag -> String)
-> (tag -> Maybe (q Type))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> tag -> m (q Exp))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> m (q Exp))
-> q Exp
typedLetrecE tag -> String
nameOf (Maybe (q Type) -> tag -> Maybe (q Type)
forall a b. a -> b -> a
const Maybe (q Type)
forall a. Maybe a
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 <$> fib 7
-- let {fib0_0 :: GHC.Types.Int;
--      fib0_0 = 1;
--      fib1_1 :: GHC.Types.Int;
--      fib1_1 = 1;
--      fib2_2 :: GHC.Types.Int;
--      fib2_2 = fib1_1 GHC.Num.+ fib0_0;
--      fib3_3 :: GHC.Types.Int;
--      fib3_3 = fib2_2 GHC.Num.+ fib1_1;
--      fib4_4 :: GHC.Types.Int;
--      fib4_4 = fib3_3 GHC.Num.+ fib2_2;
--      fib5_5 :: GHC.Types.Int;
--      fib5_5 = fib4_4 GHC.Num.+ fib3_3;
--      fib6_6 :: GHC.Types.Int;
--      fib6_6 = fib5_5 GHC.Num.+ fib4_4;
--      fib7_7 :: GHC.Types.Int;
--      fib7_7 = fib6_6 GHC.Num.+ 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 :: forall (q :: * -> *) tag.
(Ord tag, Quote q, MonadFix q) =>
(tag -> String)
-> (tag -> Maybe (q Type))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> tag -> m (q Exp))
-> (forall (m :: * -> *).
    Monad m =>
    (tag -> m (q Exp)) -> m (q Exp))
-> q Exp
typedLetrecE tag -> String
nameOf tag -> Maybe (q Type)
typeOf forall (m :: * -> *).
Monad m =>
(tag -> m (q Exp)) -> tag -> m (q Exp)
recf forall (m :: * -> *). Monad m => (tag -> m (q Exp)) -> m (q Exp)
exprf = do
    (q Exp
expr0, Map tag (Name, Maybe (q Type), q Exp)
bindings) <- StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
-> Map tag (Name, Maybe (q Type), q Exp)
-> q (q Exp, Map tag (Name, Maybe (q Type), q Exp))
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp))
-> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
forall (m :: * -> *). Monad m => (tag -> m (q Exp)) -> m (q Exp)
exprf tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
loop) Map tag (Name, Maybe (q Type), q Exp)
forall k a. Map k a
Map.empty
    [q Dec] -> q Exp -> q Exp
forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE ([[q Dec]] -> [q Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [ [ Name -> q Type -> q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
name q Type
ty' | Just q Type
ty' <- [Maybe (q Type)
ty] ] [q Dec] -> [q Dec] -> [q Dec]
forall a. [a] -> [a] -> [a]
++
          [ q Pat -> q Body -> [q Dec] -> q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
name) (q Exp -> q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB q Exp
expr) [] ]
        | (tag
_tag, (Name
name, Maybe (q Type)
ty, q Exp
expr)) <- Map tag (Name, Maybe (q Type), q Exp)
-> [(tag, (Name, Maybe (q Type), q Exp))]
forall k a. Map k a -> [(k, a)]
Map.toList Map tag (Name, Maybe (q Type), q Exp)
bindings
        ])
        q Exp
expr0
  where
    loop :: tag -> StateT (Map.Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
    loop :: tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
loop tag
tag = do
        Map tag (Name, Maybe (q Type), q Exp)
m <- StateT
  (Map tag (Name, Maybe (q Type), q Exp))
  q
  (Map tag (Name, Maybe (q Type), q Exp))
forall (m :: * -> *) s. Monad m => StateT s m s
get
        case tag
-> Map tag (Name, Maybe (q Type), q Exp)
-> Maybe (Name, Maybe (q Type), q Exp)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup tag
tag Map tag (Name, Maybe (q Type), q Exp)
m of
            -- if name is already generated, return it.
            Just (Name
name, Maybe (q Type)
_ty, q Exp
_exp) -> q Exp -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
forall a. a -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
name)

            -- otherwise generate new name, and insert it into the loop.
            Maybe (Name, Maybe (q Type), q Exp)
Nothing -> mdo
                Name
name <- q Name -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q Name
forall (m :: * -> *) a.
Monad m =>
m a -> StateT (Map tag (Name, Maybe (q Type), q Exp)) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (String -> q Name
forall (m :: * -> *). Quote m => String -> m Name
newName (tag -> String
nameOf tag
tag))
                (Map tag (Name, Maybe (q Type), q Exp)
 -> Map tag (Name, Maybe (q Type), q Exp))
-> StateT (Map tag (Name, Maybe (q Type), q Exp)) q ()
forall (m :: * -> *) s. Monad m => (s -> s) -> StateT s m ()
modify (tag
-> (Name, Maybe (q Type), q Exp)
-> Map tag (Name, Maybe (q Type), q Exp)
-> Map tag (Name, Maybe (q Type), q Exp)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert tag
tag (Name
name, tag -> Maybe (q Type)
typeOf tag
tag, q Exp
expr))
                q Exp
expr <- (tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp))
-> tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
forall (m :: * -> *).
Monad m =>
(tag -> m (q Exp)) -> tag -> m (q Exp)
recf tag -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
loop tag
tag
                q Exp -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q (q Exp)
forall a. a -> StateT (Map tag (Name, Maybe (q Type), q Exp)) q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
name)