From e9c4cad143d483e29213e9c121574d1d46c2d56a Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Tue, 6 May 2025 10:30:38 +0200
Subject: Compile: special-case to a single kernel function

---
 src/Compile.hs      |  4 ++--
 src/Compile/Exec.hs | 19 ++++++++-----------
 2 files changed, 10 insertions(+), 13 deletions(-)

(limited to 'src')

diff --git a/src/Compile.hs b/src/Compile.hs
index cd10831..9ed5a27 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -77,7 +77,7 @@ compile = \env expr -> do
   let (source, offsets) = compileToString codeID env expr
   when debugPrintAST $ hPutStrLn stderr $ "Compiled AST: <<<\n" ++ ppExpr env expr ++ "\n>>>"
   when debugCSource $ hPutStrLn stderr $ "Generated C source: <<<\n\x1B[2m" ++ lineNumbers source ++ "\x1B[0m>>>"
-  lib <- buildKernel source ["kernel"]
+  lib <- buildKernel source "kernel"
 
   let result_type = typeOf expr
       result_size = sizeofSTy result_type
@@ -86,7 +86,7 @@ compile = \env expr -> do
     allocaBytes (koResultOffset offsets + result_size) $ \ptr -> do
       let args = zip (reverse (unSList Some (slistZip env val))) (koArgOffsets offsets)
       serialiseArguments args ptr $ do
-        callKernelFun "kernel" lib ptr
+        callKernelFun lib ptr
         ok <- peekByteOff @Word8 ptr (koOkResOffset offsets)
         when (ok /= 1) $
           ioError (mkIOError userErrorType "fatal error detected during chad kernel execution (memory has been leaked)" Nothing Nothing)
diff --git a/src/Compile/Exec.hs b/src/Compile/Exec.hs
index d708fc0..9b9fb15 100644
--- a/src/Compile/Exec.hs
+++ b/src/Compile/Exec.hs
@@ -11,8 +11,6 @@ module Compile.Exec (
 
 import Control.Monad (when)
 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)
@@ -30,10 +28,10 @@ debug :: Bool
 debug = False
 
 -- The IORef wrapper is required for the finalizer to attach properly (see the 'Weak' docs)
-data KernelLib = KernelLib !(IORef (Map String (FunPtr (Ptr () -> IO ()))))
+data KernelLib = KernelLib !(IORef (FunPtr (Ptr () -> IO ())))
 
-buildKernel :: String -> [String] -> IO KernelLib
-buildKernel csource funnames = do
+buildKernel :: String -> String -> IO KernelLib
+buildKernel csource funname = do
   template <- (++ "/tmp.chad.") <$> getTempDir
   path <- mkdtemp template
 
@@ -69,8 +67,7 @@ buildKernel csource funnames = do
 
   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
+  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)
@@ -81,10 +78,10 @@ foreign import ccall "dynamic"
 
 -- 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
+callKernelFun :: KernelLib -> Ptr () -> IO ()
+callKernelFun (KernelLib ref) arg = do
+  ptr <- readIORef ref
+  wrapKernelFun ptr arg
 
 getTempDir :: IO FilePath
 getTempDir =
-- 
cgit v1.2.3-70-g09d2