-- | The foreign interface to LLVM, through the C shim in @cbits/jit_llvm.c@.
-- LLVM isn't linked in: the shim loads it with @dlopen@ when the JIT is
-- turned on ('initLLVM'), so that ucm builds and runs without LLVM, and the
-- JIT is unavailable, with a message saying where it looked, on a machine
-- that has no LLVM 20 or newer.
--
-- The shim is ours, over LLVM's C API, because the @llvm-hs@ bindings lag
-- LLVM releases and we need six functions. Modules are handed over as IR
-- /text/, which LLVM parses in memory: text can be built by string
-- building, read by a person and pasted into LLVM's command-line tools,
-- where bitcode is a bit-level stream only LLVM's own libraries write, and
-- building IR through the C API would mean many more foreign calls for IR
-- that can't be inspected without asking LLVM to print it. If parsing ever
-- shows in the compile times, the generator can be retargeted to the C API
-- without touching the code generator's logic.
module Unison.Runtime.JIT.LLVM
  ( initLLVM,
    LoadFailure (..),
    addModule,
    lookupSymbol,
    defineSymbol,
    targetTriple,
  )
where

import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Data.Text qualified as T
import Data.Text.Encoding (encodeUtf8)
import Data.Word (Word64)
import Foreign.C.String (CString, peekCString, withCString)
import Foreign.C.Types (CInt (..), CSize (..))
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (FunPtr, Ptr, WordPtr (..), castPtrToFunPtr, nullPtr, wordPtrToPtr)
import Foreign.Storable (peek, poke)

foreign import ccall safe "unison_jit_init" c_init :: IO CInt

foreign import ccall unsafe "unison_jit_load_attempts" c_loadAttempts :: IO CString

foreign import ccall unsafe "unison_jit_llvm_path" c_llvmPath :: IO CString

foreign import ccall unsafe "unison_jit_llvm_version" c_llvmVersion :: IO CString

foreign import ccall safe "unison_jit_add_module" c_addModule :: CString -> CSize -> CString -> Ptr CString -> IO CInt

foreign import ccall unsafe "unison_jit_free_string" c_freeString :: CString -> IO ()

foreign import ccall safe "unison_jit_lookup" c_lookup :: CString -> IO Word64

foreign import ccall safe "unison_jit_define_symbol" c_defineSymbol :: CString -> Word64 -> IO CInt

foreign import ccall unsafe "unison_jit_last_error" c_lastError :: IO CString

foreign import ccall unsafe "unison_jit_triple" c_triple :: IO CString

lastError :: String -> IO String
lastError :: [Char] -> IO [Char]
lastError [Char]
what = do
  [Char]
e <- Ptr CChar -> IO [Char]
peekCString (Ptr CChar -> IO [Char]) -> IO (Ptr CChar) -> IO [Char]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Ptr CChar)
c_lastError
  pure ([Char]
"JIT: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
what [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
": " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
e)

-- | Why LLVM couldn't be started.
data LoadFailure
  = -- | no usable library: each location tried, with what was found there
    NotFound [String]
  | -- | a library was loaded but the JIT couldn't be set up
    InitError String

-- | Loads LLVM and starts the JIT. Safe to call more than once. On success,
-- the version loaded and the path it came from.
initLLVM :: IO (Either LoadFailure (String, FilePath))
initLLVM :: IO (Either LoadFailure ([Char], [Char]))
initLLVM = do
  CInt
r <- IO CInt
c_init
  [Char]
path <- Ptr CChar -> IO [Char]
peekCString (Ptr CChar -> IO [Char]) -> IO (Ptr CChar) -> IO [Char]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Ptr CChar)
c_llvmPath
  if CInt
r CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
== CInt
0
    then do
      [Char]
v <- Ptr CChar -> IO [Char]
peekCString (Ptr CChar -> IO [Char]) -> IO (Ptr CChar) -> IO [Char]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Ptr CChar)
c_llvmVersion
      pure (([Char], [Char]) -> Either LoadFailure ([Char], [Char])
forall a b. b -> Either a b
Right ([Char]
v, [Char]
path))
    else
      if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
path
        then LoadFailure -> Either LoadFailure ([Char], [Char])
forall a b. a -> Either a b
Left (LoadFailure -> Either LoadFailure ([Char], [Char]))
-> ([Char] -> LoadFailure)
-> [Char]
-> Either LoadFailure ([Char], [Char])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Char]] -> LoadFailure
NotFound ([[Char]] -> LoadFailure)
-> ([Char] -> [[Char]]) -> [Char] -> LoadFailure
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> [[Char]]
lines ([Char] -> Either LoadFailure ([Char], [Char]))
-> IO [Char] -> IO (Either LoadFailure ([Char], [Char]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ptr CChar -> IO [Char]
peekCString (Ptr CChar -> IO [Char]) -> IO (Ptr CChar) -> IO [Char]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Ptr CChar)
c_loadAttempts)
        else LoadFailure -> Either LoadFailure ([Char], [Char])
forall a b. a -> Either a b
Left (LoadFailure -> Either LoadFailure ([Char], [Char]))
-> ([Char] -> LoadFailure)
-> [Char]
-> Either LoadFailure ([Char], [Char])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> LoadFailure
InitError ([Char] -> Either LoadFailure ([Char], [Char]))
-> IO [Char] -> IO (Either LoadFailure ([Char], [Char]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> IO [Char]
lastError [Char]
"initializing LLVM"

-- | Parses a module from IR text, runs the given pass pipeline
-- (such as @default<O2>@, or @""@ for none) and hands it to the JIT.
-- When asked, also returns the module's text after the passes.
addModule :: Bool -> String -> T.Text -> IO (Either String (Maybe String))
addModule :: Bool -> [Char] -> Text -> IO (Either [Char] (Maybe [Char]))
addModule Bool
wantOptimized [Char]
passes Text
ir =
  ByteString
-> (CStringLen -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen (Text -> ByteString
encodeUtf8 Text
ir) ((CStringLen -> IO (Either [Char] (Maybe [Char])))
 -> IO (Either [Char] (Maybe [Char])))
-> (CStringLen -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
p, Int
n) -> [Char]
-> (Ptr CChar -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a. [Char] -> (Ptr CChar -> IO a) -> IO a
withCString [Char]
passes ((Ptr CChar -> IO (Either [Char] (Maybe [Char])))
 -> IO (Either [Char] (Maybe [Char])))
-> (Ptr CChar -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a b. (a -> b) -> a -> b
$ \Ptr CChar
ps -> (Ptr (Ptr CChar) -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr (Ptr CChar) -> IO (Either [Char] (Maybe [Char])))
 -> IO (Either [Char] (Maybe [Char])))
-> (Ptr (Ptr CChar) -> IO (Either [Char] (Maybe [Char])))
-> IO (Either [Char] (Maybe [Char]))
forall a b. (a -> b) -> a -> b
$ \Ptr (Ptr CChar)
out -> do
    Ptr (Ptr CChar) -> Ptr CChar -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr (Ptr CChar)
out Ptr CChar
forall a. Ptr a
nullPtr
    CInt
r <- Ptr CChar -> CSize -> Ptr CChar -> Ptr (Ptr CChar) -> IO CInt
c_addModule Ptr CChar
p (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Ptr CChar
ps (if Bool
wantOptimized then Ptr (Ptr CChar)
out else Ptr (Ptr CChar)
forall a. Ptr a
nullPtr)
    if CInt
r CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
/= CInt
0
      then [Char] -> Either [Char] (Maybe [Char])
forall a b. a -> Either a b
Left ([Char] -> Either [Char] (Maybe [Char]))
-> IO [Char] -> IO (Either [Char] (Maybe [Char]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> IO [Char]
lastError [Char]
"adding a module"
      else do
        Ptr CChar
cs <- Ptr (Ptr CChar) -> IO (Ptr CChar)
forall a. Storable a => Ptr a -> IO a
peek Ptr (Ptr CChar)
out
        if Ptr CChar
cs Ptr CChar -> Ptr CChar -> Bool
forall a. Eq a => a -> a -> Bool
== Ptr CChar
forall a. Ptr a
nullPtr
          then Either [Char] (Maybe [Char]) -> IO (Either [Char] (Maybe [Char]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe [Char] -> Either [Char] (Maybe [Char])
forall a b. b -> Either a b
Right Maybe [Char]
forall a. Maybe a
Nothing)
          else do
            [Char]
txt <- Ptr CChar -> IO [Char]
peekCString Ptr CChar
cs
            Ptr CChar -> IO ()
c_freeString Ptr CChar
cs
            pure (Maybe [Char] -> Either [Char] (Maybe [Char])
forall a b. b -> Either a b
Right ([Char] -> Maybe [Char]
forall a. a -> Maybe a
Just [Char]
txt))

-- | The address of a compiled function. Code is generated on first lookup.
lookupSymbol :: String -> IO (Either String (FunPtr a))
lookupSymbol :: forall a. [Char] -> IO (Either [Char] (FunPtr a))
lookupSymbol [Char]
name = do
  Word64
a <- [Char] -> (Ptr CChar -> IO Word64) -> IO Word64
forall a. [Char] -> (Ptr CChar -> IO a) -> IO a
withCString [Char]
name Ptr CChar -> IO Word64
c_lookup
  if Word64
a Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0
    then [Char] -> Either [Char] (FunPtr a)
forall a b. a -> Either a b
Left ([Char] -> Either [Char] (FunPtr a))
-> IO [Char] -> IO (Either [Char] (FunPtr a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> IO [Char]
lastError ([Char]
"looking up " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
name)
    else Either [Char] (FunPtr a) -> IO (Either [Char] (FunPtr a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FunPtr a -> Either [Char] (FunPtr a)
forall a b. b -> Either a b
Right (Ptr (ZonkAny 0) -> FunPtr a
forall a b. Ptr a -> FunPtr b
castPtrToFunPtr (WordPtr -> Ptr (ZonkAny 0)
forall a. WordPtr -> Ptr a
wordPtrToPtr (Word -> WordPtr
WordPtr (Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
a)))))

-- | Makes @name@ in generated code refer to @addr@.
defineSymbol :: String -> Word64 -> IO (Either String ())
defineSymbol :: [Char] -> Word64 -> IO (Either [Char] ())
defineSymbol [Char]
name Word64
addr = do
  CInt
r <- [Char] -> (Ptr CChar -> IO CInt) -> IO CInt
forall a. [Char] -> (Ptr CChar -> IO a) -> IO a
withCString [Char]
name (\Ptr CChar
n -> Ptr CChar -> Word64 -> IO CInt
c_defineSymbol Ptr CChar
n Word64
addr)
  if CInt
r CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
== CInt
0 then Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either [Char] ()
forall a b. b -> Either a b
Right ()) else [Char] -> Either [Char] ()
forall a b. a -> Either a b
Left ([Char] -> Either [Char] ()) -> IO [Char] -> IO (Either [Char] ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> IO [Char]
lastError ([Char]
"defining " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
name)

targetTriple :: IO String
targetTriple :: IO [Char]
targetTriple = Ptr CChar -> IO [Char]
peekCString (Ptr CChar -> IO [Char]) -> IO (Ptr CChar) -> IO [Char]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Ptr CChar)
c_triple