{-# 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
    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
    expr <- g (loop ref)
    bindings <- runIO $ readIORef ref
    mkLet bindings 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
        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 Map.lookup y memo of
            Maybe (Name, Exp)
Nothing -> do
                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)
                _ <- mfix_ $ \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
                varE 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
        m <- IO (MVar a) -> Q (MVar a)
forall a. IO a -> Q a
runIO IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
        ans <- runIO (unsafeInterleaveIO (takeMVar m))
        result <- k ans
        runIO (putMVar m result)
        pure result
    {-# INLINE mfix_ #-}