content: -- Copyright 2022 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
  
  {-# OPTIONS_GHC -Wno-orphans #-}
  
  module ConstantCastingSpec (spec) where
  
  import Control.Monad
  import Test.Hspec
  import Test.Hspec.QuickCheck
  import Test.QuickCheck
  
  import Builder
  import Core
  import Name
  import Optimize
  import Runtime
  import IRVariants
  import TopLevel
  import Types.Core
  import Types.Imp
  import Types.Primitives
  import Types.Source
  import QueryType
  
  castOp :: ScalarBaseType -> (SAtom n) -> PrimOp SimpIR n
  castOp ty x = MiscOp $ CastOp (BaseTy (Scalar ty)) x
  
  castLam :: EnvExtender m => ScalarBaseType -> ScalarBaseType -> m n (SLam n)
  castLam fromTy toTy = do
    withFreshBinder noHint (BaseTy (Scalar fromTy)) \x -> do
      body <- exprToBlock $ PrimOp $ castOp toTy $ Var $ binderVar x
      return $ LamExpr (Nest x Empty) body
  
  exprToBlock :: EnvReader m => SExpr n -> m n (SBlock n)
  exprToBlock expr = do
    liftBuilder $ buildBlock $ do
      v <- emit $ sink expr
      return $ Var v
  
  compile :: (Topper m, Mut n)
    => ScalarBaseType -> ScalarBaseType -> m n LLVMCallable
  compile fromTy toTy = do
    sLam <- liftEnvReaderM (castLam fromTy toTy) >>= asTopLam
    compileTopLevelFun (EntryFunCC CUDANotRequired) sLam >>= packageLLVMCallable
  
  arbLitVal :: ScalarBaseType -> Gen LitVal
  arbLitVal ty = case ty of
    Int64Type   -> Int64Lit   <$> arbitrary
    Int32Type   -> Int32Lit   <$> arbitrary
    Word64Type  -> Word64Lit  <$> arbitrary
    Word32Type  -> Word32Lit  <$> arbitrary
    Word8Type   -> Word8Lit   <$> arbitrary
    Float64Type -> Float64Lit <$> arbitrary
    Float32Type -> Float32Lit <$> arbitrary
  
  constant_folding_and_runtime_casts_agree :: LLVMCallable
    -> ScalarBaseType -> LitVal -> IO ()
  constant_folding_and_runtime_casts_agree llvmFunc toTy lit =
    case foldCast toTy lit of
      Nothing -> return ()
      Just folded -> do
        ~[evaled] <- callEntryFun llvmFunc [lit]
        -- The failure message will list `evaled` as "expected" and `folded` as "got"
        folded `shouldBe` evaled
  
  spec :: Spec
  spec = do
    let cfg = EvalConfig LLVM [LibBuiltinPath] Nothing Nothing Nothing Optimize PrintCodegen
    -- or just initTopState, to always compile the prelude during unit tests?
    init_env <- runIO loadCache
    describe "constant-folding casts" do
      forM_ [ Int64Type, Int32Type, Word8Type, Word32Type
            , Word64Type, Float64Type, Float32Type] \fromTy ->
        forM_ [ Int64Type, Int32Type, Word8Type, Word32Type
              , Word64Type, Float64Type, Float32Type] \toTy -> do
          llvmFunc <- runIO $ fst <$> runTopperM cfg init_env (compile fromTy toTy)
          -- TODO: We'd really like to run 10,000 or 1,000,000 examples here, but
          -- taking 500 as a compromise between test suite speed and test
          -- coverage.
          -- I ran this offline with 10,000 before checking in, and it passed.
          modifyMaxSuccess (const 500)
            $ prop ("agrees with runtime when casting from "
                    ++ show fromTy ++ " to " ++ show toTy)
            $ forAll (arbLitVal fromTy)
            $ \lit -> constant_folding_and_runtime_casts_agree llvmFunc toTy lit
      llvm_w32_f32 <- runIO $ fst <$> runTopperM cfg init_env
        (compile Word32Type Float32Type)
      it "agrees with runtime on rounding mode" $
        -- These values are chosen to tickle the difference between different
        -- rounding modes when rounding to float32, and specifically between
        -- breaking ties to even, to zero, or to +infinity when rounding to
        -- nearest.
        --
        -- Specifically, these are 32-bit integers that are large enough not to be
        -- exactly representable in 32-bit floating-point, whose low-order bits go
        -- through every configuration that rounding behavior is sensitive to
        -- (i.e., nearest float is larger, nearest float is smaller, exactly
        -- between two floats with the previous bit being even, or odd).
        forM_ [ Word32Lit 0x3000000
              , Word32Lit 0x3000001
              , Word32Lit 0x3000002
              , Word32Lit 0x3000003
              , Word32Lit 0x3000004
              , Word32Lit 0x3000005
              , Word32Lit 0x3000006
              , Word32Lit 0x3000007
              , Word32Lit 0x3000008
              , Word32Lit 0x3000009
              , Word32Lit 0x300000A
              , Word32Lit 0x300000B
              , Word32Lit 0x300000C
              , Word32Lit 0x300000D
              , Word32Lit 0x300000E
              ] \val ->
                  constant_folding_and_runtime_casts_agree llvm_w32_f32 Float32Type val
      llvm_w64_f64 <- runIO $ fst <$> runTopperM cfg init_env
        (compile Word64Type Float64Type)
      it "rounds to nearest even" $
        -- This particular literal was found by QuickCheck as a counter-example
        -- where a previous version of `foldCast.fixUlp` did not round to nearest
        -- even correctly.
        constant_folding_and_runtime_casts_agree llvm_w64_f64 Float64Type $
          Word64Lit 10633092062366456832
  