content: {-# LANGUAGE RecordWildCards #-}
  
  module Tests.Extra.IntMap where
  
  import AtCoder.Extra.IntMap qualified as IM
  import Control.Monad (foldM_)
  import Control.Monad.Primitive (PrimMonad, PrimState)
  import Control.Monad.ST (RealWorld)
  import Data.Map.Strict qualified as IMR -- R: referential implementation
  import Data.Vector.Unboxed qualified as VU
  import GHC.Stack (HasCallStack)
  import Test.QuickCheck.Monadic as QCM
  import Test.Tasty
  import Test.Tasty.QuickCheck as QC
  
  data Init = Init
    { n :: {-# UNPACK #-} !Int,
      ref0 :: !(IMR.Map Int Int),
      imM :: !(IO (IM.IntMap RealWorld Int))
    }
  
  instance Show Init where
    show Init {..} = show (n, ref0)
  
  instance QC.Arbitrary Init where
    arbitrary = do
      n <- QC.chooseInt (1, 10)
      pure $ Init n IMR.empty (IM.new n)
  
  data Query
    = Size
    | Member Int
    | NotMember Int
    | Null
    | Lookup Int
    | Insert Int Int
    | InsertWithAdd Int Int
    | Delete Int
    | Delete_ Int
    | LookupGE Int
    | LookupGT Int
    | LookupLE Int
    | LookupLT Int
    | LookupMin
    | LookupMax
    | DeleteMin
    | DeleteMax
    deriving (Show)
  
  -- | Arbitrary return type for the `Query` result.
  data Result
    = None
    | B Bool
    | I Int
    | M (Maybe Int)
    | KV (Maybe (Int, Int))
    | XS (VU.Vector Int)
    | KVS (VU.Vector (Int, Int))
    deriving (Show, Eq)
  
  queryGen :: Int -> QC.Gen Query
  queryGen n = do
    QC.oneof
      [ pure Size,
        Member <$> lookupKeyGen,
        NotMember <$> lookupKeyGen,
        pure Null,
        Lookup <$> lookupKeyGen,
        -- insert is partial function
        Insert <$> insertKeyGen <*> valGen,
        InsertWithAdd <$> insertKeyGen <*> valGen,
        Delete <$> lookupKeyGen,
        Delete_ <$> lookupKeyGen,
        LookupGE <$> lookupKeyGen,
        LookupGT <$> lookupKeyGen,
        LookupLE <$> lookupKeyGen,
        LookupLT <$> lookupKeyGen,
        pure LookupMin,
        pure LookupMax
      ]
    where
      -- for total functions
      lookupKeyGen = QC.chooseInt (-1, n)
      -- for partial functions
      insertKeyGen = QC.chooseInt (0, n - 1)
      valGen = QC.chooseInt (-10, 10)
  
  -- | containers. (referential implementation)
  handleRef :: IMR.Map Int Int -> Query -> (IMR.Map Int Int, Result)
  handleRef im q = case q of
    Size -> (im, I $ IMR.size im)
    Member k -> (im, B $ IMR.member k im)
    NotMember k -> (im, B . not $ IMR.member k im)
    Null -> (im, B $ IMR.null im)
    Lookup k -> (im, M $ IMR.lookup k im)
    Insert k v -> (IMR.insert k v im, None)
    InsertWithAdd k v -> (IMR.insertWith (+) k v im, None)
    Delete k -> (IMR.delete k im, B $ IMR.member k im)
    Delete_ k -> (IMR.delete k im, None)
    LookupGE k -> (im, KV (IMR.lookupGE k im))
    LookupGT k -> (im, KV (IMR.lookupGT k im))
    LookupLE k -> (im, KV (IMR.lookupLE k im))
    LookupLT k -> (im, KV (IMR.lookupLT k im))
    LookupMin -> (im, KV (IMR.lookupMin im))
    LookupMax -> (im, KV (IMR.lookupMax im))
    DeleteMin -> wrapKV IMR.deleteFindMin
    DeleteMax -> wrapKV IMR.deleteFindMax
    where
      wrapKV f
        | IMR.null im = (im, KV Nothing)
        | otherwise = let (!kv, !im') = f im in (im', KV (Just kv))
  
  -- | ac-library-hs.
  handleAcl :: (HasCallStack, PrimMonad m) => IM.IntMap (PrimState m) Int -> Query -> m Result
  handleAcl im q = case q of
    Size -> I <$> IM.size im
    Member k -> B <$> IM.member im k
    NotMember k -> B <$> IM.notMember im k
    Null -> B <$> IM.null im
    Lookup k -> M <$> IM.lookup im k
    Insert k v -> do
      IM.insert im k v
      pure None
    InsertWithAdd k v -> do
      IM.insertWith im (+) k v
      pure None
    Delete k -> B <$> IM.delete im k
    Delete_ k -> do
      IM.delete_ im k
      pure None
    LookupGE k -> KV <$> IM.lookupGE im k
    LookupGT k -> KV <$> IM.lookupGT im k
    LookupLE k -> KV <$> IM.lookupLE im k
    LookupLT k -> KV <$> IM.lookupLT im k
    LookupMin -> KV <$> IM.lookupMin im
    LookupMax -> KV <$> IM.lookupMax im
    DeleteMin -> KV <$> IM.deleteMin im
    DeleteMax -> KV <$> IM.deleteMax im
  
  prop_randomTest :: Init -> QC.Property
  prop_randomTest Init {..} = QCM.monadicIO $ do
    im <- QCM.run imM
    q <- QCM.pick $ QC.chooseInt (1, 5 * n)
    qs <- QCM.pick $ QC.vectorOf q (queryGen n)
    foldM_
      ( \ref query -> do
          let (!ref', !expected) = handleRef ref query
          res <- QCM.run $ handleAcl im query
          QCM.assertWith (expected == res) $ show (query, expected, res)
  
          -- check the map contents:
          let assocsE = VU.fromList $ IMR.assocs ref'
          assocs <- QCM.run $ IM.assocs im
          QCM.assertWith (assocsE == assocs) $ show ("- assocs", assocsE, assocs)
  
          let keysE = VU.fromList $ IMR.keys ref'
          keys <- QCM.run $ IM.keys im
          QCM.assertWith (keysE == keys) $ show ("- keys", keysE, keys)
  
          let elemsE = VU.fromList $ IMR.elems ref'
          elems <- QCM.run $ IM.elems im
          QCM.assertWith (elemsE == elems) $ show ("- elems", elemsE, elems)
  
          let sizeE = IMR.size ref'
          size <- QCM.run $ IM.size im
          QCM.assertWith (sizeE == size) $ show ("- size", sizeE, size)
  
          pure ref'
      )
      ref0
      qs
  
  tests :: [TestTree]
  tests =
    [ QC.testProperty "randomTest" prop_randomTest
    ]
  