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 JaxADTSpec (spec) where
  
  import Data.Aeson (encode, decode)
  import Test.Hspec
  
  import Name
  import JAX.Concrete
  import JAX.Rename
  import JAX.ToSimp
  import Runtime
  import TopLevel
  import Types.Imp
  import Types.Primitives hiding (Sin)
  import Types.Source hiding (SourceName)
  import QueryType
  
  x_nm, y_nm :: JSourceName
  x_nm = JSourceName 0 0 "x"
  y_nm = JSourceName 1 0 "y"
  
  float :: JVarType
  float = (JArrayName [] F32)
  
  ten_vec :: JVarType
  ten_vec = (JArrayName [DimSize 10] F32)
  
  a_jaxpr :: JVarType -> Jaxpr VoidS
  a_jaxpr ty = Jaxpr
    (Nest (JBindSource x_nm ty) Empty)
    Empty
    (Nest (JEqn
      (Nest (JBindSource y_nm ty) Empty)
      Sin
      [JVariable $ JVar (SourceName x_nm) ty]) Empty)
    [JVariable $ JVar (SourceName y_nm) ty]
  
  compile :: Jaxpr VoidS -> IO LLVMCallable
  compile jaxpr = do
    let cfg = EvalConfig LLVM [LibBuiltinPath] Nothing Nothing Nothing NoOptimize PrintCodegen
    env <- initTopState
    fst <$> runTopperM cfg env do
      -- TODO Implement GenericE for jaxprs, derive SinkableE, and properly sink
      -- the jaxpr instead of just coercing it.
      Distinct <- getDistinct
      jRename <- liftRenameM $ renameJaxpr (unsafeCoerceE jaxpr)
      jSimp <- liftJaxSimpM (simplifyJaxpr jRename) >>= asTopLam
      compileTopLevelFun (EntryFunCC CUDANotRequired) jSimp >>= packageLLVMCallable
  
  spec :: Spec
  spec = do
    describe "JaxADT" do
      it "round-trips to json" do
        let first = encode $ a_jaxpr ten_vec
        let (Just decoded) = (decode first :: Maybe (Jaxpr VoidS))
        let second = encode decoded
        second `shouldBe` first
      it "executes" do
        jLLVM <- compile $ a_jaxpr float
        result <- callEntryFun jLLVM [Float32Lit 3.0]
        result `shouldBe` [Float32Lit $ sin 3.0]
  