1 files changed, 99 insertions, 0 deletions
diff --git a/src/CHAD/Compile/Exec.hs b/src/CHAD/Compile/Exec.hs
new file mode 100644
index 0000000..5b4afc8
--- /dev/null
+++ b/src/CHAD/Compile/Exec.hs
@@ -0,0 +1,99 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TupleSections #-}
+module CHAD.Compile.Exec (
+  KernelLib,
+  buildKernel,
+  callKernelFun,
+
+  -- * misc
+  lineNumbers,
+) where
+
+import Control.Monad (when)
+import Data.IORef
+import Foreign (Ptr)
+import Foreign.Ptr (FunPtr)
+import System.Directory (removeDirectoryRecursive)
+import System.Environment (lookupEnv)
+import System.Exit (ExitCode(..))
+import System.IO (hPutStrLn, stderr)
+import System.IO.Error (mkIOError, userErrorType)
+import System.IO.Unsafe (unsafePerformIO)
+import System.Posix.DynamicLinker
+import System.Posix.Temp (mkdtemp)
+import System.Process (readProcessWithExitCode)
+
+
+debug :: Bool
+debug = False
+
+-- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs)
+data KernelLib = KernelLib !(IORef (FunPtr (Ptr () -> IO ())))
+
+buildKernel :: String -> String -> IO (KernelLib, String)
+buildKernel csource funname = do
+  template <- (++ "/tmp.chad.") <$> getTempDir
+  path <- mkdtemp template
+
+  let outso = path ++ "/out.so"
+  let args = ["-O3", "-march=native"
+             ,"-shared", "-fPIC"
+             ,"-std=c99", "-x", "c"
+             ,"-o", outso, "-"
+             ,"-Wall", "-Wextra"
+             ,"-Wno-unused-variable", "-Wno-unused-but-set-variable"
+             ,"-Wno-unused-parameter", "-Wno-unused-function"
+             ,"-Wno-alloc-size-larger-than"  -- ideally we'd keep this, but gcc reports false positives
+             ,"-Wno-maybe-uninitialized"]  -- maximum1i goes out of range if its input is empty, yes, don't complain
+  (ec, gccStdout, gccStderr) <- readProcessWithExitCode "gcc" args csource
+
+  -- Print the source before the GCC output.
+  case ec of
+    ExitSuccess -> return ()
+    ExitFailure{} -> hPutStrLn stderr $ "[chad] Kernel compilation failed! Source: <<<\n" ++ lineNumbers csource ++ ">>>"
+
+  case ec of
+    ExitSuccess -> return ()
+    ExitFailure{} -> do
+      removeDirectoryRecursive path
+      ioError (mkIOError userErrorType "chad kernel compilation failed" Nothing Nothing)
+
+  numLoaded <- atomicModifyIORef' numLoadedCounter (\n -> (n+1, n+1))
+  when debug $ hPutStrLn stderr $ "[chad] loading kernel " ++ path ++ " (" ++ show numLoaded ++ " total)"
+  dl <- dlopen outso [RTLD_LAZY, RTLD_LOCAL]
+
+  removeDirectoryRecursive path  -- we keep a reference anyway because we have the file open now
+
+  ref <- newIORef =<< dlsym dl funname
+  _ <- mkWeakIORef ref (do numLeft <- atomicModifyIORef' numLoadedCounter (\n -> (n-1, n-1))
+                           when debug $ hPutStrLn stderr $ "[chad] unloading kernel " ++ path ++ " (" ++ show numLeft ++ " left)"
+                           dlclose dl)
+  return (KernelLib ref, gccStdout ++ (if null gccStdout then "" else "\n") ++ gccStderr)
+
+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 :: KernelLib -> Ptr () -> IO ()
+callKernelFun (KernelLib ref) arg = do
+  ptr <- readIORef ref
+  wrapKernelFun ptr arg
+
+getTempDir :: IO FilePath
+getTempDir =
+  lookupEnv "TMPDIR" >>= \case
+    Just s | not (null s) -> return s
+    _ -> return "/tmp"
+
+{-# NOINLINE numLoadedCounter #-}
+numLoadedCounter :: IORef Int
+numLoadedCounter = unsafePerformIO $ newIORef 0
+
+lineNumbers :: String -> String
+lineNumbers str =
+  let lns = lines str
+      numlines = length lns
+      width = length (show numlines)
+      pad s = replicate (width - length s) ' ' ++ s
+  in unlines (zipWith (\i ln -> pad (show i) ++ " | " ++ ln) [1::Int ..] lns)