content: -- Copyright 2023 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  module Dex.Foreign.JAX where
  
  import Control.Monad.IO.Class
  import Data.Aeson (encode, eitherDecode')
  import qualified Data.ByteString.Lazy.Char8 as B
  import Foreign.C
  import Foreign.Ptr
  
  import Dex.Foreign.Context
  import Export
  import JAX.Concrete
  import JAX.Rename
  import JAX.ToSimp
  import Name
  
  -- TODO newCString just mallocs the string; we have to
  -- arrange for the caller to free it.
  dexRoundtripJaxprJson :: CString -> IO CString
  dexRoundtripJaxprJson jsonPtr = do
    json <- B.pack <$> peekCString jsonPtr
    let maybeJaxpr :: Either String (ClosedJaxpr VoidS) = eitherDecode' json
    case maybeJaxpr of
      Right jaxpr -> do
        let redumped = encode jaxpr
        newCString $ B.unpack redumped
      Left err -> newCString err
  
  dexCompileJaxpr :: Ptr Context -> CInt -> CString -> IO ExportNativeFunctionAddr
  dexCompileJaxpr ctxPtr ccInt jsonPtr = do
    json <- B.pack <$> peekCString jsonPtr
    let maybeJaxpr :: Either String (ClosedJaxpr VoidS) = eitherDecode' json
    case maybeJaxpr of
      Right jaxpr -> runTopperMFromContext ctxPtr do
        Distinct <- getDistinct
        jRename <- liftRenameM $ renameClosedJaxpr (unsafeCoerceE jaxpr)
        sLam <- liftJaxSimpM $ simplifyClosedJaxpr jRename
        func <- prepareSLamForExport (intAsCC ccInt) sLam
        liftIO $ emitExport ctxPtr func
      Left err -> error err
  