1 files changed, 59 insertions, 0 deletions
diff --git a/src/Compile/Exec.hs b/src/Compile/Exec.hs
new file mode 100644
index 0000000..163be2b
--- /dev/null
+++ b/src/Compile/Exec.hs
@@ -0,0 +1,59 @@
+{-# 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"