Prim.hs

{-# LANGUAGE CPP #-}

module Test.RecoverRTTI.Prim (
    -- * Equality
    canComparePrim
    -- * Arbitrary
  , Wrap(..)
  , primSatisfiesArbitrary
  , arbitraryPrimClassifier
  ) where

import Control.Monad (replicateM)
import Data.Aeson qualified as Aeson
import Data.ByteString qualified as BS.Strict
import Data.ByteString.Lazy qualified as BS.Lazy
import Data.Int
import Data.IntSet (IntSet)
import Data.Primitive.ByteArray qualified as Prim (ByteArray)
import Data.SOP (Compose)
import Data.SOP.Dict
import Data.String (fromString)
import Data.Text qualified as Text.Strict
import Data.Text.Lazy qualified as Text.Lazy
import Data.Vector qualified as Vector.Boxed
import Data.Vector.Primitive qualified as Vector.Primitive
import Data.Vector.Storable qualified as Vector.Storable
import Data.Word
import Unsafe.Coerce (unsafeCoerce)

#if MIN_VERSION_base(4,17,0)
import GHC.IsList qualified as IsList
#else
import GHC.Exts qualified as IsList (fromList)
#endif

#if !MIN_VERSION_bytestring(0,12,0)
import Data.ByteString.Short qualified as BS.Short
#endif

import Debug.RecoverRTTI

import Test.QuickCheck (Arbitrary(..), Gen)
import Test.QuickCheck qualified as QC

import Test.RecoverRTTI.Globals
import Test.RecoverRTTI.Orphans ()

{-------------------------------------------------------------------------------
  Equality
-------------------------------------------------------------------------------}

canComparePrim :: PrimClassifier a -> Dict Eq a
canComparePrim = primSatisfies

{-------------------------------------------------------------------------------
  Arbitrary support for the primitive types
-------------------------------------------------------------------------------}

primSatisfiesArbitrary :: PrimClassifier a -> Dict (Compose Arbitrary Wrap) a
primSatisfiesArbitrary = primSatisfies

arbitraryPrimClassifier :: Gen (Some PrimClassifier)
arbitraryPrimClassifier = QC.elements [
    -- Primitive types

      Some C_Bool
    , Some C_Char
    , Some C_Double
    , Some C_Float
    , Some C_Int
    , Some C_Int16
    , Some C_Int8
    , Some C_Int32
    , Some C_Int64
    , Some C_Integer
    , Some C_Ordering
    , Some C_Unit
    , Some C_Word
    , Some C_Word8
    , Some C_Word16
    , Some C_Word32
    , Some C_Word64

    -- String types

    , Some C_BS_Strict
    , Some C_BS_Lazy
    , Some C_Text_Strict
    , Some C_Text_Lazy

#if !MIN_VERSION_bytestring(0,12,0)
    , Some C_BS_Short
#endif

    -- Aeson

    , Some C_Value

    -- Reference cells

    , Some C_STRef
    , Some C_TVar
    , Some C_MVar

    -- Functions

    , Some C_Fun

    -- Containers with no type arguments

    , Some C_IntSet
    , Some C_Prim_ArrayM
    , Some C_Vector_Storable
    , Some C_Vector_StorableM
    , Some C_Vector_Primitive
    , Some C_Vector_PrimitiveM
    , Some C_ByteArray
    , Some C_MutableByteArray
    ]
  where
    _checkAllCases :: PrimClassifier a -> ()
    _checkAllCases = \case
        -- Primitive types

        C_Bool     -> ()
        C_Char     -> ()
        C_Double   -> ()
        C_Float    -> ()
        C_Int      -> ()
        C_Int16    -> ()
        C_Int8     -> ()
        C_Int32    -> ()
        C_Int64    -> ()
        C_Integer  -> ()
        C_Ordering -> ()
        C_Unit     -> ()
        C_Word     -> ()
        C_Word8    -> ()
        C_Word16   -> ()
        C_Word32   -> ()
        C_Word64   -> ()

        -- String types

        C_BS_Strict   -> ()
        C_BS_Lazy     -> ()
        C_Text_Strict -> ()
        C_Text_Lazy   -> ()

#if !MIN_VERSION_bytestring(0,12,0)
        C_BS_Short    -> ()
#endif

        -- Aeson

        C_Value -> ()

        -- Reference cells

        C_STRef -> ()
        C_TVar  -> ()
        C_MVar  -> ()

        -- Functions

        C_Fun -> ()

        -- Containers with no type arguments

        C_IntSet            -> ()
        C_Prim_ArrayM       -> ()
        C_Vector_Storable   -> ()
        C_Vector_StorableM  -> ()
        C_Vector_Primitive  -> ()
        C_Vector_PrimitiveM -> ()
        C_ByteArray         -> ()
        C_MutableByteArray  -> ()

{-------------------------------------------------------------------------------
  Arbitrary instances for specific types
  -------------------------------------------------------------------------------}

-- | 'Wrap' makes it possible to override an 'Arbitrary' instance if needed.
newtype Wrap a = Wrap { unwrap :: a }

deriving newtype instance Arbitrary (Wrap ())
deriving newtype instance Arbitrary (Wrap Bool)
deriving newtype instance Arbitrary (Wrap Char)
deriving newtype instance Arbitrary (Wrap Double)
deriving newtype instance Arbitrary (Wrap Float)
deriving newtype instance Arbitrary (Wrap Int)
deriving newtype instance Arbitrary (Wrap Int16)
deriving newtype instance Arbitrary (Wrap Int32)
deriving newtype instance Arbitrary (Wrap Int64)
deriving newtype instance Arbitrary (Wrap Int8)
deriving newtype instance Arbitrary (Wrap Integer)
deriving newtype instance Arbitrary (Wrap IntSet)
deriving newtype instance Arbitrary (Wrap Ordering)
deriving newtype instance Arbitrary (Wrap String)
deriving newtype instance Arbitrary (Wrap Word)
deriving newtype instance Arbitrary (Wrap Word16)
deriving newtype instance Arbitrary (Wrap Word32)
deriving newtype instance Arbitrary (Wrap Word64)
deriving newtype instance Arbitrary (Wrap Word8)

instance Arbitrary (Wrap BS.Strict.ByteString) where
  arbitrary = Wrap . BS.Strict.pack <$> arbitrary

instance Arbitrary (Wrap BS.Lazy.ByteString) where
  arbitrary = Wrap . BS.Lazy.pack <$> arbitrary

#if !MIN_VERSION_bytestring(0,12,0)
instance Arbitrary (Wrap BS.Short.ShortByteString) where
  arbitrary = Wrap . BS.Short.pack <$> arbitrary
#endif

instance Arbitrary (Wrap Text.Strict.Text) where
  arbitrary = Wrap . Text.Strict.pack <$> arbitrary

instance Arbitrary (Wrap Text.Lazy.Text) where
  arbitrary = Wrap . Text.Lazy.pack <$> arbitrary

-- | aeson >= 2.0.3.0 does define an 'Arbitrary' instance for 'Aeson.Value',
-- but it generates values that are too big, which cause the size sanity check
-- on the generator 'prop_showGenerated' to start to fail.
instance Arbitrary (Wrap Aeson.Value) where
  arbitrary = QC.choose (0, 10) >>= fmap Wrap . go
    where
      go :: Int -> Gen Aeson.Value
      go 0  = QC.oneof nonRecursive
      go sz = QC.oneof (nonRecursive ++ recursive sz)

      nonRecursive :: [Gen Aeson.Value]
      nonRecursive = [
            Aeson.String . Text.Strict.pack <$> arbitrary
          , Aeson.Number . fromInteger <$> arbitrary
          , Aeson.Bool <$> arbitrary
          , return Aeson.Null
          ]

      recursive :: Int -> [Gen Aeson.Value]
      recursive sz = [
            do n <- QC.choose (0, 5)
               Aeson.Array . Vector.Boxed.fromList <$> replicateM n (go (sz `div` n))
          , do n <- QC.choose (0, 5)
               Aeson.object <$> replicateM n (
                       (Aeson..=)
                   <$> (fromString <$> fieldName)
                   <*> go (sz `div` n)
                 )
          ]

      -- We're not interested in testing crazy values
      fieldName :: Gen String
      fieldName = QC.elements ["a", "b", "c"]

-- | Rather than trying to be clever here, we just generate a handful of
-- examples in different categories.
instance Arbitrary (Wrap SomeFun) where
  arbitrary = fmap Wrap $ QC.elements [
        -- Parametrically polymorphic function
        fun (id    :: Int -> Int)
      , fun (const :: Int -> Bool -> Int)
        -- Ad-hoc polymorphic function
      , fun (negate :: Int -> Int)
      , fun ((+)    :: Int -> Int -> Int)
        -- Partial application
      , fun (const 1 :: Bool -> Int)
      , fun ((+)   1 :: Int -> Int)
      ]
    where
      fun :: (a -> b) -> SomeFun
      fun = unsafeCoerce

instance Arbitrary (Wrap SomeStorableVector) where
  arbitrary = fmap Wrap $ QC.elements [
        some $ Vector.Storable.fromList ([1, 2, 3] :: [Int])
      , some $ Vector.Storable.fromList ("abc"     :: String)
      ]
    where
      some :: Vector.Storable.Vector a -> SomeStorableVector
      some = unsafeCoerce

instance Arbitrary (Wrap SomePrimitiveVector) where
  arbitrary = fmap Wrap $ QC.elements [
        some $ Vector.Primitive.fromList ([1, 2, 3] :: [Int])
      , some $ Vector.Primitive.fromList ("abc"     :: String)
      ]
    where
      some :: Vector.Primitive.Vector a -> SomePrimitiveVector
      some = unsafeCoerce

instance Arbitrary (Wrap Prim.ByteArray) where
  arbitrary = Wrap . IsList.fromList <$> arbitrary

{-------------------------------------------------------------------------------
  For the mutable variables, we just use the one global example
-------------------------------------------------------------------------------}

instance Arbitrary (Wrap SomeSTRef) where
  arbitrary = return $ Wrap exampleSTRef

instance Arbitrary (Wrap SomeTVar) where
  arbitrary = return $ Wrap exampleTVar

instance Arbitrary (Wrap SomeMVar) where
  arbitrary = return $ Wrap exampleMVar

instance Arbitrary (Wrap SomePrimArrayM) where
  arbitrary = return $ Wrap examplePrimArrayM

instance Arbitrary (Wrap SomeStorableVectorM) where
  arbitrary = return $ Wrap exampleStorableVectorM

instance Arbitrary (Wrap SomePrimitiveVectorM) where
  arbitrary = return $ Wrap examplePrimitiveVectorM

instance Arbitrary (Wrap SomeMutableByteArray) where
  arbitrary = return $ Wrap exampleMutableByteArray