{-# LANGUAGE QuasiQuotes     #-}
{-# LANGUAGE TemplateHaskell #-}
-- |
module Python.Internal.EvalQQ
  ( -- * Evaluators and QQ
    evaluatorPyf
    -- * Code generation
  , expQQ
  , Mode(..)
  ) where

import Control.Monad.IO.Class
import Control.Monad.Catch
import Control.Monad.Trans.Cont (ContT(..))
import Data.Bits
import Data.Char
import Data.List                 (intercalate)
import Data.ByteString           qualified as BS
import Data.ByteString.Unsafe    qualified as BS
import Data.Text                 qualified as T
import Data.Text.Encoding        qualified as T
import Foreign.C.Types
import Foreign.Ptr
import System.Exit
import System.Process (readProcessWithExitCode)

import Language.C.Inline          qualified as C
import Language.C.Inline.Unsafe   qualified as CU
import Language.Haskell.TH.Lib    qualified as TH
import Language.Haskell.TH.Syntax qualified as TH

import Python.Internal.Types
import Python.Internal.Program
import Python.Internal.Eval
import Python.Internal.CAPI
import Python.Inline.Literal


----------------------------------------------------------------
C.context (C.baseCtx <> pyCtx)
C.include "<inline-python.h>"
----------------------------------------------------------------

-- | Python's variable name encoded using UTF-8. It exists in order to
--   avoid working with @String@ at runtime.
newtype PyVarName = PyVarName BS.ByteString
  deriving stock (Int -> PyVarName -> ShowS
[PyVarName] -> ShowS
PyVarName -> String
(Int -> PyVarName -> ShowS)
-> (PyVarName -> String)
-> ([PyVarName] -> ShowS)
-> Show PyVarName
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PyVarName -> ShowS
showsPrec :: Int -> PyVarName -> ShowS
$cshow :: PyVarName -> String
show :: PyVarName -> String
$cshowList :: [PyVarName] -> ShowS
showList :: [PyVarName] -> ShowS
Show, (forall (m :: * -> *). Quote m => PyVarName -> m Exp)
-> (forall (m :: * -> *). Quote m => PyVarName -> Code m PyVarName)
-> Lift PyVarName
forall t.
(forall (m :: * -> *). Quote m => t -> m Exp)
-> (forall (m :: * -> *). Quote m => t -> Code m t) -> Lift t
forall (m :: * -> *). Quote m => PyVarName -> m Exp
forall (m :: * -> *). Quote m => PyVarName -> Code m PyVarName
$clift :: forall (m :: * -> *). Quote m => PyVarName -> m Exp
lift :: forall (m :: * -> *). Quote m => PyVarName -> m Exp
$cliftTyped :: forall (m :: * -> *). Quote m => PyVarName -> Code m PyVarName
liftTyped :: forall (m :: * -> *). Quote m => PyVarName -> Code m PyVarName
TH.Lift)

varName :: String -> PyVarName
varName :: String -> PyVarName
varName = ByteString -> PyVarName
PyVarName (ByteString -> PyVarName)
-> (String -> ByteString) -> String -> PyVarName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> (String -> Text) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack

unsafeWithPyVarName :: PyVarName -> Program r (Ptr CChar)
unsafeWithPyVarName :: forall r. PyVarName -> Program r (Ptr CChar)
unsafeWithPyVarName (PyVarName ByteString
bs)
  = ((Ptr CChar -> IO r) -> IO r) -> Program r (Ptr CChar)
forall a r. ((a -> IO r) -> IO r) -> Program r a
progIOBracket (ByteString -> (Ptr CChar -> IO r) -> IO r
forall a. ByteString -> (Ptr CChar -> IO a) -> IO a
BS.unsafeUseAsCString ByteString
bs)


bindVar :: ToPy a => PyVarName -> a -> DictBinder
bindVar :: forall a. ToPy a => PyVarName -> a -> DictBinder
bindVar PyVarName
var a
a = (Ptr PyObject -> Py ()) -> DictBinder
DictBinder ((Ptr PyObject -> Py ()) -> DictBinder)
-> (Ptr PyObject -> Py ()) -> DictBinder
forall a b. (a -> b) -> a -> b
$ \Ptr PyObject
p_dict -> Program () () -> Py ()
forall a. Program a a -> Py a
runProgram (Program () () -> Py ()) -> Program () () -> Py ()
forall a b. (a -> b) -> a -> b
$ do
  p_key <- PyVarName -> Program () (Ptr CChar)
forall r. PyVarName -> Program r (Ptr CChar)
unsafeWithPyVarName PyVarName
var
  p_obj <- takeOwnership =<< progPy (throwOnNULL =<< basicToPy a)
  progPy $ do
    r <- Py [CU.block| int {
      PyObject* p_obj = $(PyObject* p_obj);
      return PyDict_SetItemString($(PyObject* p_dict), $(char* p_key), p_obj);
      } |]
    case r of
      CInt
0 -> () -> Py ()
forall a. a -> Py a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      CInt
_ -> Py ()
forall a. Py a
mustThrowPyError



----------------------------------------------------------------
-- Evaluators
----------------------------------------------------------------

evaluatorPyf :: PyQuote -> Py PyObject
evaluatorPyf :: PyQuote -> Py PyObject
evaluatorPyf (PyQuote Code
code DictBinder
binder) = Program PyObject PyObject -> Py PyObject
forall a. Program a a -> Py a
runProgram (Program PyObject PyObject -> Py PyObject)
-> Program PyObject PyObject -> Py PyObject
forall a b. (a -> b) -> a -> b
$ do
  p_locals <- Ptr PyObject -> Program PyObject (Ptr PyObject)
forall r. Ptr PyObject -> Program r (Ptr PyObject)
takeOwnership (Ptr PyObject -> Program PyObject (Ptr PyObject))
-> Program PyObject (Ptr PyObject)
-> Program PyObject (Ptr PyObject)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Py (Ptr PyObject) -> Program PyObject (Ptr PyObject)
forall a r. Py a -> Program r a
progPy Py (Ptr PyObject)
basicNewDict
  p_kwargs <- takeOwnership =<< progPy basicNewDict
  progPy $ do
    -- Create function in p_locals
    exec Main (DictPtr p_locals) (PyQuote code mempty)
    -- Look up function
    p_fun <- getFunctionObject p_locals >>= \case
      Ptr PyObject
NULL -> PyInternalError -> Py (Ptr PyObject)
forall e a. (HasCallStack, Exception e) => e -> Py a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (PyInternalError -> Py (Ptr PyObject))
-> PyInternalError -> Py (Ptr PyObject)
forall a b. (a -> b) -> a -> b
$ String -> PyInternalError
PyInternalError String
"_inline_python_ must be present"
      Ptr PyObject
p    -> Ptr PyObject -> Py (Ptr PyObject)
forall a. a -> Py a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ptr PyObject
p
    -- Call python function we just constructed
    binder.bind p_kwargs
    newPyObject =<< throwOnNULL =<< basicCallKwdOnly p_fun p_kwargs

getFunctionObject :: Ptr PyObject -> Py (Ptr PyObject)
getFunctionObject :: Ptr PyObject -> Py (Ptr PyObject)
getFunctionObject Ptr PyObject
p_dict = do
  IO (Ptr PyObject) -> Py (Ptr PyObject)
forall a. IO a -> Py a
Py [CU.exp| PyObject* { PyDict_GetItemString($(PyObject *p_dict), "_inline_python_") } |]



----------------------------------------------------------------
-- TH generator
----------------------------------------------------------------

script :: String
script :: String
script = $( do let path = "py/bound-vars.py"
               TH.addDependentFile path
               TH.lift =<< TH.runIO (readFile path)
          )

data Mode
  = Eval
  | Exec
  | Fun

-- | Generate TH splice which updates python environment dictionary
--   and returns python source code.
expQQ :: Mode   -- ^ Python evaluation mode: @exec@/@eval@
      -> String -- ^ Python source code
      -> TH.Q TH.Exp
expQQ :: Mode -> String -> Q Exp
expQQ Mode
mode String
qq_src = do
  -- We need to preprocess before passing it to python.
  let src :: String
src     = Mode -> ShowS
prepareSource       Mode
mode String
qq_src
      src_var :: String
src_var = Mode -> ShowS
prepareForVarLookup Mode
mode String
src
  antis  <- IO [String] -> Q [String]
forall a. IO a -> Q a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [String] -> Q [String]) -> IO [String] -> Q [String]
forall a b. (a -> b) -> a -> b
$ do
    -- We've embedded script into library and we need to pass source
    -- code of QQ to a script. It can contain whatever symbols so to
    -- be safe it's base16 encode. This encoding is very simple and we
    -- don't care much about efficiency here
    (code, stdout, stderr) <- String -> [String] -> String -> IO (ExitCode, String, String)
readProcessWithExitCode String
"python"
        [ String
"-"
        , case Mode
mode of Mode
Eval -> String
"eval"
                       Mode
Exec -> String
"exec"
                       Mode
Fun  -> String
"exec"
        ]
      (String -> IO (ExitCode, String, String))
-> String -> IO (ExitCode, String, String)
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines [ String
script
                , String
"decode_and_print('" String -> ShowS
forall a. Semigroup a => a -> a -> a
<>
                  [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ [ Int -> Char
intToDigit (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
w Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftR` Int
4)
                           , Int -> Char
intToDigit (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
w Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
15) ]
                         | Word8
w <- ByteString -> [Word8]
BS.unpack (ByteString -> [Word8]) -> ByteString -> [Word8]
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
src_var
                         ]
                  String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"')"
                ]
    case code of
      ExitCode
ExitSuccess   -> [String] -> IO [String]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([String] -> IO [String]) -> [String] -> IO [String]
forall a b. (a -> b) -> a -> b
$ String -> [String]
words String
stdout
      ExitFailure{} -> String -> IO [String]
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
stderr
  let args = [ [| bindVar $(PyVarName -> Q Exp
forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
forall (m :: * -> *). Quote m => PyVarName -> m Exp
TH.lift (String -> PyVarName
varName String
nm)) $(String -> Q Exp
forall (m :: * -> *). Quote m => String -> m Exp
TH.dyn (ShowS
chop String
nm)) |]
             | String
nm <- [String]
antis
             ]
      src_eval = Mode -> [String] -> ShowS
prepareForEval Mode
mode [String]
antis String
src
  --
  [| PyQuote ($(TH.lift $ codeFromString src_eval))
             (mconcat $(TH.listE args))
   |]


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

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


----------------------------------------------------------------
-- Python source code transform
----------------------------------------------------------------

prepareSource :: Mode -> String -> String
prepareSource :: Mode -> ShowS
prepareSource = \case
  Mode
Eval -> (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace
  Mode
Exec -> ShowS
unindent
  Mode
Fun  -> ShowS
unindent

prepareForVarLookup :: Mode -> String -> String
prepareForVarLookup :: Mode -> ShowS
prepareForVarLookup = \case
  Mode
Eval -> ShowS
forall a. a -> a
id
  Mode
Exec -> ShowS
forall a. a -> a
id
  Mode
Fun  -> (String
"def __dummy__():\n"String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
indent

prepareForEval :: Mode -> [String] -> String -> String
prepareForEval :: Mode -> [String] -> ShowS
prepareForEval Mode
mode [String]
vars String
src = case Mode
mode of
  Mode
Eval -> String
src
  Mode
Exec -> String
src
  Mode
Fun  -> String
"def _inline_python_("String -> ShowS
forall a. Semigroup a => a -> a -> a
<>String
argsString -> ShowS
forall a. Semigroup a => a -> a -> a
<>String
"):\n"
    String -> ShowS
forall a. Semigroup a => a -> a -> a
<> ShowS
indent String
src
  where
    args :: String
args = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," [String]
vars

-- Python is indentation based and quasiquotes do not strip leading
-- space. We have to do that ourself
unindent :: String -> String
unindent :: ShowS
unindent String
py_src = case String -> [String]
lines String
py_src of
  []  -> String
""
  -- Strip all leading space for 1-line scripts
  [String
l] -> (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace String
l
  -- For multiline script we require that first line should be empty
  String
l:[String]
ls
    | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Bool
isSpace) String
l -> ShowS
forall a. HasCallStack => String -> a
error String
"First line of multiline quasiquote must be empty"
    -- FIXME: We break multiline strings here. Badly. We need proper python lexer
    -- FIXME: We probably should just forbid tabs
    | Bool
otherwise ->
      let non_empty :: [String]
non_empty = (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Bool
isSpace)) [String]
ls
          n :: Int
n         = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [ String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==Char
' ') String
s) | String
s <- [String]
non_empty ]
      in [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ Int -> ShowS
forall a. Int -> [a] -> [a]
drop Int
n ShowS -> [String] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
ls

indent :: String -> String
indent :: ShowS
indent = [String] -> String
unlines
       ([String] -> String) -> (String -> [String]) -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String
"    "String -> ShowS
forall a. [a] -> [a] -> [a]
++)
       ([String] -> [String])
-> (String -> [String]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String]
lines