-- |
-- Copyright: (C) 2013 Amgen, Inc.
--

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}

module Language.R.QQ
  ( r
  , rsafe
  ) where

import           Control.Memory.Region
import           Control.Monad.R.Class
import qualified Data.Vector.SEXP as Vector
import qualified Foreign.R as R
import qualified Foreign.R.Parse as R
import           Foreign.R (SEXP, SomeSEXP(..))
import           Foreign.R.Error
import           Internal.Error
import           Language.R (eval)
import           Language.R.Globals (nilValue, globalEnv)
import           Language.R.GC (automaticSome)
import           Language.R.HExp
import           Language.R.Instance
import           Language.R.Literal (mkSEXPIO)

import Language.Haskell.TH (Q, runIO)
import Language.Haskell.TH.Quote
import qualified Language.Haskell.TH.Syntax as TH
import qualified Language.Haskell.TH.Lib as TH

import Control.Concurrent (MVar, newMVar, takeMVar, putMVar)
import Control.Exception (throwIO)
import Control.Monad (unless)
import Data.List (intercalate, isSuffixOf)
import qualified Data.Set as Set
import Data.Set (Set)
import Foreign (alloca, peek)
import Foreign.C.String (withCString)
import System.IO.Unsafe (unsafePerformIO)

-------------------------------------------------------------------------------
-- Compile time Quasi-Quoter                                                 --
-------------------------------------------------------------------------------

-- | An R value, expressed as an R expression, in R's syntax.
r :: QuasiQuoter
r :: QuasiQuoter
r = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp = \txt :: String
txt -> [| eval =<< $(expQQ txt) |]
    , quotePat :: String -> Q Pat
quotePat  = String -> String -> Q Pat
forall a. String -> a
unimplemented "quotePat"
    , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall a. String -> a
unimplemented "quoteType"
    , quoteDec :: String -> Q [Dec]
quoteDec  = String -> String -> Q [Dec]
forall a. String -> a
unimplemented "quoteDec"
    }

-- | Quasiquoter for pure R code (no side effects) and that does not depend on
-- the global environment (referential transparency). This means that all
-- symbols must appear qualified with a package namespace (whose bindings are
-- locked by default), the code must not affect R shared state in any way,
-- including the global environment, and must not perform I/O.

-- TODO some of the above invariants can be checked statically. Do so.
rsafe :: QuasiQuoter
rsafe :: QuasiQuoter
rsafe = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
    { quoteExp :: String -> Q Exp
quoteExp  = \txt :: String
txt -> [| unsafePerformIO $ runRegion $ automaticSome =<< eval =<< $(expQQ txt) |]
    , quotePat :: String -> Q Pat
quotePat  = String -> String -> Q Pat
forall a. String -> a
unimplemented "quotePat"
    , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall a. String -> a
unimplemented "quoteType"
    , quoteDec :: String -> Q [Dec]
quoteDec  = String -> String -> Q [Dec]
forall a. String -> a
unimplemented "quoteDec"
    }

-- | Serialize quasiquotes using a global lock, because the compiler is allowed
-- in theory to run them in parallel, yet the R runtime is not reentrant.
qqLock :: MVar ()
qqLock :: MVar ()
qqLock = IO (MVar ()) -> MVar ()
forall a. IO a -> a
unsafePerformIO (IO (MVar ()) -> MVar ()) -> IO (MVar ()) -> MVar ()
forall a b. (a -> b) -> a -> b
$ () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
{-# NOINLINE qqLock #-}

parse :: String -> IO (R.SEXP V 'R.Expr)
parse :: String -> IO (SEXP V 'Expr)
parse txt :: String
txt = do
    Config -> IO ()
initialize Config
defaultConfig
    String -> (CString -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a. String -> (CString -> IO a) -> IO a
withCString String
txt ((CString -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr))
-> (CString -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a b. (a -> b) -> a -> b
$ \ctxt :: CString
ctxt ->
        IO (SEXP V 'String)
-> (SEXP V 'String -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall (a :: SEXPTYPE) s b.
IO (SEXP V a) -> (SEXP s a -> IO b) -> IO b
R.withProtected (CString -> IO (SEXP V 'String)
R.mkString CString
ctxt) ((SEXP V 'String -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr))
-> (SEXP V 'String -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a b. (a -> b) -> a -> b
$ \rtxt :: SEXP V 'String
rtxt ->
          (Ptr CInt -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr CInt -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr))
-> (Ptr CInt -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a b. (a -> b) -> a -> b
$ \status :: Ptr CInt
status -> do
            IO (SEXP V 'Expr)
-> (SEXP V 'Expr -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall (a :: SEXPTYPE) s b.
IO (SEXP V a) -> (SEXP s a -> IO b) -> IO b
R.withProtected (SEXP V 'String
-> Int -> Ptr CInt -> SEXP V 'Nil -> IO (SEXP V 'Expr)
forall (a :: SEXPTYPE) s.
In a '[ 'Nil, 'String] =>
SEXP s 'String -> Int -> Ptr CInt -> SEXP s a -> IO (SEXP s 'Expr)
R.parseVector SEXP V 'String
rtxt (-1) Ptr CInt
status (SEXP G 'Nil -> SEXP V 'Nil
forall t s (a :: SEXPTYPE). (t <= s) => SEXP s a -> SEXP t a
R.release SEXP G 'Nil
nilValue)) ((SEXP V 'Expr -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr))
-> (SEXP V 'Expr -> IO (SEXP V 'Expr)) -> IO (SEXP V 'Expr)
forall a b. (a -> b) -> a -> b
$ \exprs :: SEXP V 'Expr
exprs -> do
              Int
rc <- CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr CInt -> IO CInt
forall a. Storable a => Ptr a -> IO a
peek Ptr CInt
status
              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ParseStatus
R.PARSE_OK ParseStatus -> ParseStatus -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> ParseStatus
forall a. Enum a => Int -> a
toEnum Int
rc) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                RError -> IO ()
forall e a. Exception e => e -> IO a
throwIO (RError -> IO ()) -> (String -> RError) -> String -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> RError
RError (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ "Parse error in: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
txt
              SEXP V 'Expr -> IO (SEXP V 'Expr)
forall (m :: * -> *) a. Monad m => a -> m a
return SEXP V 'Expr
exprs

antiSuffix :: String
antiSuffix :: String
antiSuffix = "_hs"

isAnti :: SEXP s 'R.Char -> Bool
isAnti :: SEXP s 'Char -> Bool
isAnti (SEXP s 'Char -> HExp s 'Char
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> Char (Vector 'Char Word8 -> String
Vector.toString -> String
name)) = String
antiSuffix String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` String
name
isAnti _ = String -> Bool
forall a. HasCallStack => String -> a
error "Impossible"

-- | Chop antiquotation variable names to get the corresponding Haskell variable name.
chop :: String -> String
chop :: String -> String
chop name :: String
name = Int -> String -> String
forall a. Int -> [a] -> [a]
take (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
name Int -> Int -> Int
forall a. Num a => a -> a -> a
- String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
antiSuffix) String
name

-- | Traverse 'R.SEXP' structure and find all occurences of antiquotations.
collectAntis :: R.SEXP s a -> Set (SEXP s 'R.Char)
collectAntis :: SEXP s a -> Set (SEXP s 'Char)
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> Symbol (SEXP s a -> SEXP s 'Char
forall s (a :: SEXPTYPE) (b :: SEXPTYPE). SEXP s a -> SEXP s b
R.unsafeCoerce -> SEXP s 'Char
name) _ _)
  | SEXP s 'Char -> Bool
forall s. SEXP s 'Char -> Bool
isAnti SEXP s 'Char
name = SEXP s 'Char -> Set (SEXP s 'Char)
forall a. a -> Set a
Set.singleton SEXP s 'Char
name
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (List sxa :: SEXP s a
sxa sxb :: SEXP s b
sxb sxc :: SEXP s c
sxc)) = do
    [Set (SEXP s 'Char)] -> Set (SEXP s 'Char)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [SEXP s a -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s a
sxa, SEXP s b -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s b
sxb, SEXP s c -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s c
sxc]
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (Lang (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> Symbol (SEXP s a -> SEXP s 'Char
forall s (a :: SEXPTYPE) (b :: SEXPTYPE). SEXP s a -> SEXP s b
R.unsafeCoerce -> SEXP s 'Char
name) _ _) sxb :: SEXP s b
sxb))
  | SEXP s 'Char -> Bool
forall s. SEXP s 'Char -> Bool
isAnti SEXP s 'Char
name = SEXP s 'Char -> Set (SEXP s 'Char) -> Set (SEXP s 'Char)
forall a. Ord a => a -> Set a -> Set a
Set.insert SEXP s 'Char
name (SEXP s b -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s b
sxb)
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (Lang sxa :: SEXP s a
sxa sxb :: SEXP s b
sxb)) =
    Set (SEXP s 'Char) -> Set (SEXP s 'Char) -> Set (SEXP s 'Char)
forall a. Ord a => Set a -> Set a -> Set a
Set.union (SEXP s a -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s a
sxa) (SEXP s b -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s b
sxb)
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (Closure sxa :: SEXP s a
sxa sxb :: SEXP s b
sxb sxc :: SEXP s 'Env
sxc)) =
    [Set (SEXP s 'Char)] -> Set (SEXP s 'Char)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [SEXP s a -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s a
sxa, SEXP s b -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s b
sxb, SEXP s 'Env -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP s 'Env
sxc]
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (Vector _ sxv :: Vector 'Vector (SomeSEXP V)
sxv)) =
    [Set (SEXP s 'Char)] -> Set (SEXP s 'Char)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [SEXP s a -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis (SEXP V a -> SEXP s a
forall s (a :: SEXPTYPE) r. SEXP s a -> SEXP r a
R.unsafeRelease SEXP V a
sx) | SomeSEXP sx :: SEXP V a
sx <- Vector 'Vector (SomeSEXP V) -> [SomeSEXP V]
forall (ty :: SEXPTYPE) a. SVECTOR ty a => Vector ty a -> [a]
Vector.toList Vector 'Vector (SomeSEXP V)
sxv]
collectAntis (SEXP s a -> HExp s a
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> (Expr _ sxv :: Vector 'Expr (SomeSEXP V)
sxv)) =
    [Set (SEXP s 'Char)] -> Set (SEXP s 'Char)
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [SEXP s a -> Set (SEXP s 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis (SEXP V a -> SEXP s a
forall s (a :: SEXPTYPE) r. SEXP s a -> SEXP r a
R.unsafeRelease SEXP V a
sx) | SomeSEXP sx :: SEXP V a
sx <- Vector 'Expr (SomeSEXP V) -> [SomeSEXP V]
forall (ty :: SEXPTYPE) a. SVECTOR ty a => Vector ty a -> [a]
Vector.toList Vector 'Expr (SomeSEXP V)
sxv]
collectAntis _ = Set (SEXP s 'Char)
forall a. Set a
Set.empty

-- | An R quasiquote is syntactic sugar for a function that we
-- generate, which closes over all antiquotation variables, and applies the
-- function to the Haskell values to which those variables are bound. Example:
--
-- @
-- [r| x_hs + y_hs |] ==> apply (apply [r| function(x_hs, y_hs) x_hs + y_hs |] x) y
-- @
expQQ :: String -> Q TH.Exp
expQQ :: String -> Q Exp
expQQ input :: String
input = do
    ()
_ <- IO () -> Q ()
forall a. IO a -> Q a
runIO (IO () -> Q ()) -> IO () -> Q ()
forall a b. (a -> b) -> a -> b
$ MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
qqLock
    SEXP G 'Expr
expr <- IO (SEXP G 'Expr) -> Q (SEXP G 'Expr)
forall a. IO a -> Q a
runIO (IO (SEXP G 'Expr) -> Q (SEXP G 'Expr))
-> IO (SEXP G 'Expr) -> Q (SEXP G 'Expr)
forall a b. (a -> b) -> a -> b
$ SEXP V 'Expr -> IO (SEXP G 'Expr)
forall s (a :: SEXPTYPE). SEXP s a -> IO (SEXP G a)
R.protect (SEXP V 'Expr -> IO (SEXP G 'Expr))
-> IO (SEXP V 'Expr) -> IO (SEXP G 'Expr)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO (SEXP V 'Expr)
parse String
input
    let antis :: [String]
antis = [String
x | (SEXP G 'Char -> HExp G 'Char
forall s (a :: SEXPTYPE). SEXP s a -> HExp s a
hexp -> Char (Vector 'Char Word8 -> String
Vector.toString -> String
x))
                       <- Set (SEXP G 'Char) -> [SEXP G 'Char]
forall a. Set a -> [a]
Set.toList (SEXP G 'Expr -> Set (SEXP G 'Char)
forall s (a :: SEXPTYPE). SEXP s a -> Set (SEXP s 'Char)
collectAntis SEXP G 'Expr
expr)]
        args :: [Q Exp]
args = (String -> Q Exp) -> [String] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Q Exp
TH.dyn (String -> Q Exp) -> (String -> String) -> String -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
chop) [String]
antis
        closure :: String
closure = "function(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate "," [String]
antis String -> String -> String
forall a. [a] -> [a] -> [a]
++ "){" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
input String -> String -> String
forall a. [a] -> [a] -> [a]
++ "}"
        z :: Q Exp
z = [| return (R.release nilValue) |]
    [Name]
vars <- (String -> Q Name) -> [String] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\_ -> String -> Q Name
TH.newName "x") [String]
antis
    -- Abstract over antis using fresh vars, to avoid captures with names bound
    -- internally (such as 'f' below).
    Exp
x <- (\body :: Q Exp
body -> (Q Exp -> Q Exp -> Q Exp) -> Q Exp -> [Q Exp] -> Q Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Q Exp -> Q Exp -> Q Exp
TH.appE Q Exp
body [Q Exp]
args) (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Q Pat] -> Q Exp -> Q Exp
TH.lamE ((Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Pat
TH.varP [Name]
vars)
      [| -- Memoize the runtime parsing of the generated closure (provided the
         -- compiler notices that it can let-float to top-level).
         let sx = unsafePerformIO $ do
                    exprs <- parse closure
                    SomeSEXP e <- R.readVector exprs 0
                    clos <- R.eval e (R.release globalEnv)
                    R.unSomeSEXP clos R.preserveObject
                    return clos
         in io $ case sx of
           SomeSEXP f ->
             R.lcons f =<<
               $(foldr (\x xs -> [| R.withProtected $xs $ \cdr -> do
                                        car <- mkSEXPIO $(TH.varE x)
                                        R.lcons car cdr |]) z vars)
       |]
    IO () -> Q ()
forall a. IO a -> Q a
runIO (IO () -> Q ()) -> IO () -> Q ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
R.unprotect 1 -- Ptr expr
    IO () -> Q ()
forall a. IO a -> Q a
runIO (IO () -> Q ()) -> IO () -> Q ()
forall a b. (a -> b) -> a -> b
$ MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
qqLock ()
    Exp -> Q Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
x