{-# 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 <- IO (IORef (Map a (Name, Exp))) -> Q (IORef (Map a (Name, Exp)))
forall a. IO a -> Q a
runIO (IO (IORef (Map a (Name, Exp))) -> Q (IORef (Map a (Name, Exp))))
-> IO (IORef (Map a (Name, Exp))) -> Q (IORef (Map a (Name, Exp)))
forall a b. (a -> b) -> a -> b
$ Map a (Name, Exp) -> IO (IORef (Map a (Name, Exp)))
forall a. a -> IO (IORef a)
newIORef Map a (Name, Exp)
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 <- IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp))
forall a. IO a -> Q a
runIO (IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp)))
-> IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp))
forall a b. (a -> b) -> a -> b
$ IORef (Map a (Name, Exp)) -> IO (Map a (Name, Exp))
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
        Exp -> ExpQ
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> ExpQ) -> Exp -> ExpQ
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)) <- Map a (Name, Exp) -> [(a, (Name, Exp))]
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 <- IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp))
forall a. IO a -> Q a
runIO (IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp)))
-> IO (Map a (Name, Exp)) -> Q (Map a (Name, Exp))
forall a b. (a -> b) -> a -> b
$ IORef (Map a (Name, Exp)) -> IO (Map a (Name, Exp))
forall a. IORef a -> IO a
readIORef IORef (Map a (Name, Exp))
ref
        case a -> Map a (Name, Exp) -> Maybe (Name, Exp)
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 <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"_let" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Map a (Name, Exp) -> Int
forall k a. Map k a -> Int
Map.size Map a (Name, Exp)
memo)
                Exp
_ <- (Exp -> ExpQ) -> ExpQ
forall a. (a -> Q a) -> Q a
forall (m :: * -> *) a. MonadFix_ m => (a -> m a) -> m a
mfix_ ((Exp -> ExpQ) -> ExpQ) -> (Exp -> ExpQ) -> ExpQ
forall a b. (a -> b) -> a -> b
$ \Exp
yCode -> do
                    IO () -> Q ()
forall a. IO a -> Q a
runIO (IO () -> Q ()) -> IO () -> Q ()
forall a b. (a -> b) -> a -> b
$ IORef (Map a (Name, Exp))
-> (Map a (Name, Exp) -> (Map a (Name, Exp), ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Map a (Name, Exp))
ref ((Map a (Name, Exp) -> (Map a (Name, Exp), ())) -> IO ())
-> (Map a (Name, Exp) -> (Map a (Name, Exp), ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Map a (Name, Exp)
m -> (a -> (Name, Exp) -> Map a (Name, Exp) -> Map a (Name, Exp)
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
                Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
name

            Just (Name
name, Exp
_) ->
                Name -> ExpQ
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 = ((a -> ExpQ) -> a -> ExpQ) -> ((a -> ExpQ) -> ExpQ) -> ExpQ
forall a.
Ord a =>
((a -> ExpQ) -> a -> ExpQ) -> ((a -> ExpQ) -> ExpQ) -> ExpQ
letrecE ((a -> ExpQ) -> (a -> ExpQ) -> a -> ExpQ
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 <- IO (MVar a) -> Q (MVar a)
forall a. IO a -> Q a
runIO IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
        a
ans <- IO a -> Q a
forall a. IO a -> Q a
runIO (IO a -> IO a
forall a. IO a -> IO a
unsafeInterleaveIO (MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
m))
        a
result <- a -> Q a
k a
ans
        IO () -> Q ()
forall a. IO a -> Q a
runIO (MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
result)
        a -> Q a
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
result
    {-# INLINE mfix_ #-}