{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
module Data.Aeson.Internal.TH (
    letrecE,
    autoletE,
) where

import Data.IORef              (IORef, atomicModifyIORef, newIORef, readIORef)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Language.Haskell.TH     (varE, ExpQ, Name, Q, newName, runIO)
import System.IO.Unsafe        (unsafeInterleaveIO)

import qualified Data.Map as Map
import qualified Language.Haskell.TH.Syntax as TH

letrecE
    :: forall a. Ord a
    => ((a -> ExpQ) -> (a -> ExpQ))
    -> ((a -> ExpQ) -> ExpQ)
    -> ExpQ
letrecE :: forall a.
Ord a =>
((a -> ExpQ) -> a -> ExpQ) -> ((a -> ExpQ) -> ExpQ) -> ExpQ
letrecE (a -> ExpQ) -> a -> ExpQ
f (a -> ExpQ) -> ExpQ
g = do
    IORef (Map a (Name, Exp))
ref <- forall a. IO a -> Q a
runIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall k a. Map k a
Map.empty
    Exp
expr <- (a -> ExpQ) -> ExpQ
g (IORef (Map a (Name, Exp)) -> a -> ExpQ
loop IORef (Map a (Name, Exp))
ref)
    Map a (Name, Exp)
bindings <- forall a. IO a -> Q a
runIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef (Map a (Name, Exp))
ref
    Map a (Name, Exp) -> Exp -> ExpQ
mkLet Map a (Name, Exp)
bindings Exp
expr
  where
    mkLet :: Map.Map a (Name, TH.Exp) -> TH.Exp -> ExpQ
    mkLet :: Map a (Name, Exp) -> Exp -> ExpQ
mkLet Map a (Name, Exp)
bindings Exp
expr = do
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Dec] -> Exp -> Exp
TH.LetE
            [ Pat -> Body -> [Dec] -> Dec
TH.ValD (Name -> Pat
TH.VarP Name
name) (Exp -> Body
TH.NormalB Exp
code) []
            | (a
_, (Name
name, Exp
code)) <- forall k a. Map k a -> [(k, a)]
Map.toList Map a (Name, Exp)
bindings
            ]
            Exp
expr

    loop :: IORef (Map.Map a (Name, TH.Exp)) -> a -> ExpQ
    loop :: IORef (Map a (Name, Exp)) -> a -> ExpQ
loop IORef (Map a (Name, Exp))
ref a
y = do
        Map a (Name, Exp)
memo <- forall a. IO a -> Q a
runIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef (Map a (Name, Exp))
ref
        case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
y Map a (Name, Exp)
memo of
            Maybe (Name, Exp)
Nothing -> do
                Name
name <- forall (m :: * -> *). Quote m => String -> m Name
newName forall a b. (a -> b) -> a -> b
$ String
"_let" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall k a. Map k a -> Int
Map.size Map a (Name, Exp)
memo)
                Exp
_ <- forall (m :: * -> *) a. MonadFix_ m => (a -> m a) -> m a
mfix_ forall a b. (a -> b) -> a -> b
$ \Exp
yCode -> do
                    forall a. IO a -> Q a
runIO forall a b. (a -> b) -> a -> b
$ forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map a (Name, Exp))
ref forall a b. (a -> b) -> a -> b
$ \Map a (Name, Exp)
m -> (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
y (Name
name, Exp
yCode) Map a (Name, Exp)
m, ())
                    (a -> ExpQ) -> a -> ExpQ
f (IORef (Map a (Name, Exp)) -> a -> ExpQ
loop IORef (Map a (Name, Exp))
ref) a
y
                forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
name

            Just (Name
name, Exp
_) ->
                forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
name

-- | Better 'letE'.
autoletE
    :: Ord a
    => (a -> ExpQ)            -- ^ what bindings are
    -> ((a -> ExpQ) -> ExpQ)  -- ^ expression with a function to generate bindings
    -> ExpQ
autoletE :: forall a. Ord a => (a -> ExpQ) -> ((a -> ExpQ) -> ExpQ) -> ExpQ
autoletE a -> ExpQ
f = forall a.
Ord a =>
((a -> ExpQ) -> a -> ExpQ) -> ((a -> ExpQ) -> ExpQ) -> ExpQ
letrecE (forall a b. a -> b -> a
const a -> ExpQ
f)

-------------------------------------------------------------------------------
-- MonadFix Q is not always there
-------------------------------------------------------------------------------

class MonadFix_ m where
    mfix_ :: (a -> m a) -> m a

instance MonadFix_ Q where
    mfix_ :: forall a. (a -> Q a) -> Q a
mfix_ a -> Q a
k = do
        MVar a
m <- forall a. IO a -> Q a
runIO forall a. IO (MVar a)
newEmptyMVar
        a
ans <- forall a. IO a -> Q a
runIO (forall a. IO a -> IO a
unsafeInterleaveIO (forall a. MVar a -> IO a
takeMVar MVar a
m))
        a
result <- a -> Q a
k a
ans
        forall a. IO a -> Q a
runIO (forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
result)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure a
result
    {-# INLINE mfix_ #-}