GHCi/TH.hs

    {-# LANGUAGE ScopedTypeVariables, StandaloneDeriving, DeriveGeneric,
        TupleSections, RecordWildCards, InstanceSigs #-}
    {-# OPTIONS_GHC -fno-warn-name-shadowing #-}

    -- |
    -- Running TH splices
    --
    module GHCi.TH (startTH, finishTH, runTH, GHCiQException(..)) where

    import GHCi.Message
    import GHCi.RemoteTypes
    import GHC.Serialized

    import Control.Exception
    import qualified Control.Monad.Fail as Fail
    import Data.Binary
    import Data.Binary.Put
    import Data.ByteString (ByteString)
    import qualified Data.ByteString as B
    import qualified Data.ByteString.Lazy as LB
    import Data.Data
    import Data.Dynamic
    import Data.IORef
    import Data.Map (Map)
    import qualified Data.Map as M
    import Data.Maybe
    import GHC.Desugar
    import qualified Language.Haskell.TH        as TH
    import qualified Language.Haskell.TH.Syntax as TH
    import Unsafe.Coerce

    initQState :: Pipe -> QState
    initQState p = QState M.empty [] Nothing p

    runModFinalizers :: GHCiQ ()
    runModFinalizers = go =<< getState
      where
        go s | (f:ff) <- qsFinalizers s = do
          putState (s { qsFinalizers = ff}) >> TH.runQ f >> getState >>= go
        go _ = return ()

    newtype GHCiQ a = GHCiQ { runGHCiQ :: QState -> IO (a, QState) }

    data GHCiQException = GHCiQException QState String
      deriving (Show, Typeable)

    instance Exception GHCiQException

    instance Functor GHCiQ where
      fmap f (GHCiQ s) = GHCiQ $ fmap (\(x,s') -> (f x,s')) . s

    instance Applicative GHCiQ where
      f <*> a = GHCiQ $ \s ->
        do (f',s')  <- runGHCiQ f s
           (a',s'') <- runGHCiQ a s'
           return (f' a', s'')
      pure x = GHCiQ (\s -> return (x,s))

    instance Monad GHCiQ where
      m >>= f = GHCiQ $ \s ->
        do (m', s')  <- runGHCiQ m s
           (a,  s'') <- runGHCiQ (f m') s'
           return (a, s'')

      fail = Fail.fail

    instance Fail.MonadFail GHCiQ where
      fail err  = GHCiQ $ \s -> throwIO (GHCiQException s err)

    getState :: GHCiQ QState
    getState = GHCiQ $ \s -> return (s,s)

    putState :: QState -> GHCiQ ()
    putState s = GHCiQ $ \_ -> return ((),s)

    noLoc :: TH.Loc
    noLoc = TH.Loc "<no file>" "<no package>" "<no module>" (0,0) (0,0)

    ghcCmd :: Binary a => Message (THResult a) -> GHCiQ a
    ghcCmd m = GHCiQ $ \s -> do
      r <- remoteCall (qsPipe s) m
      case r of
        THException str -> throwIO (GHCiQException s str)
        THComplete res -> return (res, s)

    instance TH.Quasi GHCiQ where
      qNewName str = ghcCmd (NewName str)
      qReport isError msg = ghcCmd (Report isError msg)

      -- See Note [TH recover with -fexternal-interpreter] in TcSplice
      qRecover (GHCiQ h) (GHCiQ a) = GHCiQ $ \s -> (do
        remoteCall (qsPipe s) StartRecover
        (r, s') <- a s
        remoteCall (qsPipe s) (EndRecover False)
        return (r,s'))
          `catch`
           \GHCiQException{} -> remoteCall (qsPipe s) (EndRecover True) >> h s
      qLookupName isType occ = ghcCmd (LookupName isType occ)
      qReify name = ghcCmd (Reify name)
      qReifyFixity name = ghcCmd (ReifyFixity name)
      qReifyInstances name tys = ghcCmd (ReifyInstances name tys)
      qReifyRoles name = ghcCmd (ReifyRoles name)

      -- To reify annotations, we send GHC the AnnLookup and also the
      -- TypeRep of the thing we're looking for, to avoid needing to
      -- serialize irrelevant annotations.
      qReifyAnnotations :: forall a . Data a => TH.AnnLookup -> GHCiQ [a]
      qReifyAnnotations lookup =
        map (deserializeWithData . B.unpack) <$>
          ghcCmd (ReifyAnnotations lookup typerep)
        where typerep = typeOf (undefined :: a)

      qReifyModule m = ghcCmd (ReifyModule m)
      qReifyConStrictness name = ghcCmd (ReifyConStrictness name)
      qLocation = fromMaybe noLoc . qsLocation <$> getState
      qRunIO m = GHCiQ $ \s -> fmap (,s) m
      qAddDependentFile file = ghcCmd (AddDependentFile file)
      qAddTopDecls decls = ghcCmd (AddTopDecls decls)
      qAddModFinalizer fin = GHCiQ $ \s ->
        return ((), s { qsFinalizers = fin : qsFinalizers s })
      qGetQ = GHCiQ $ \s ->
        let lookup :: forall a. Typeable a => Map TypeRep Dynamic -> Maybe a
            lookup m = fromDynamic =<< M.lookup (typeOf (undefined::a)) m
        in return (lookup (qsMap s), s)
      qPutQ k = GHCiQ $ \s ->
        return ((), s { qsMap = M.insert (typeOf k) (toDyn k) (qsMap s) })
      qIsExtEnabled x = ghcCmd (IsExtEnabled x)
      qExtsEnabled = ghcCmd ExtsEnabled

    startTH :: IO (RemoteRef (IORef QState))
    startTH = do
      r <- newIORef (initQState (error "startTH: no pipe"))
      mkRemoteRef r

    finishTH :: Pipe -> RemoteRef (IORef QState) -> IO ()
    finishTH pipe rstate = do
      qstateref <- localRef rstate
      qstate <- readIORef qstateref
      _ <- runGHCiQ runModFinalizers qstate { qsPipe = pipe }
      return ()

    runTH
      :: Pipe -> RemoteRef (IORef QState) -> HValueRef
      -> THResultType
      -> Maybe TH.Loc
      -> IO ByteString
    runTH pipe rstate rhv ty mb_loc = do
      hv <- localRef rhv
      case ty of
        THExp -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Exp)
        THPat -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Pat)
        THType -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q TH.Type)
        THDec -> runTHQ pipe rstate mb_loc (unsafeCoerce hv :: TH.Q [TH.Dec])
        THAnnWrapper -> do
          hv <- unsafeCoerce <$> localRef rhv
          case hv :: AnnotationWrapper of
            AnnotationWrapper thing -> return $!
              LB.toStrict (runPut (put (toSerialized serializeWithData thing)))

    runTHQ
      :: Binary a => Pipe -> RemoteRef (IORef QState) -> Maybe TH.Loc -> TH.Q a
      -> IO ByteString
    runTHQ pipe@Pipe{..} rstate mb_loc ghciq = do
      qstateref <- localRef rstate
      qstate <- readIORef qstateref
      let st = qstate { qsLocation = mb_loc, qsPipe = pipe }
      (r,new_state) <- runGHCiQ (TH.runQ ghciq) st
      writeIORef qstateref new_state
      return $! LB.toStrict (runPut (put r))