{-# LANGUAGE LambdaCase #-} {-# LANGUAGE TupleSections #-} module Compile.Exec ( KernelLib, buildKernel, callKernelFun, ) where import Data.IORef import qualified Data.Map.Strict as Map import Data.Map.Strict (Map) import Foreign (Ptr) import Foreign.Ptr (FunPtr) import System.Directory (removeDirectoryRecursive) import System.Environment (lookupEnv) import System.IO (hPutStrLn, stderr) import System.Posix.DynamicLinker import System.Posix.Temp (mkdtemp) import System.Process (readProcess) -- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs) data KernelLib = KernelLib !(IORef (Map String (FunPtr (Ptr () -> IO ())))) buildKernel :: String -> [String] -> IO KernelLib buildKernel csource funnames = do template <- (++ "/tmp.chad.") <$> getTempDir path <- mkdtemp template let outso = path ++ "/out.so" let args = ["-O3", "-march=native", "-shared", "-fPIC", "-x", "c", "-o", outso, "-"] _ <- readProcess "gcc" args csource hPutStrLn stderr $ "[chad] loading kernel " ++ path dl <- dlopen outso [RTLD_LAZY, RTLD_LOCAL] removeDirectoryRecursive path -- we keep a reference anyway because we have the file open now ptrs <- Map.fromList <$> sequence [(name,) <$> dlsym dl name | name <- funnames] ref <- newIORef ptrs _ <- mkWeakIORef ref (do hPutStrLn stderr $ "[chad] unloading kernel " ++ path dlclose dl) return (KernelLib ref) foreign import ccall "dynamic" wrapKernelFun :: FunPtr (Ptr () -> IO ()) -> Ptr () -> IO () -- Ensure that keeping a reference to the returned function also keeps the 'KernelLib' alive {-# NOINLINE callKernelFun #-} callKernelFun :: String -> KernelLib -> Ptr () -> IO () callKernelFun key (KernelLib ref) arg = do mp <- readIORef ref wrapKernelFun (mp Map.! key) arg getTempDir :: IO FilePath getTempDir = lookupEnv "TMPDIR" >>= \case Just s | not (null s) -> return s _ -> return "/tmp"