SegTree2d.hs

{-# LANGUAGE RecordWildCards #-}

module Tests.Extra.SegTree2d (tests) where

import AtCoder.Extra.SegTree2d qualified as Seg
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (RealWorld)
import Data.Foldable (for_)
import Data.Semigroup (Sum (..))
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import Test.QuickCheck.Monadic as QCM
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck as QC

data Init = Init
  { capacity :: {-# UNPACK #-} !Int,
    q :: {-# UNPACK #-} !Int,
    debugPoints :: VU.Vector (Int, Int),
    vecM :: !(IO (VUM.MVector RealWorld (Int, Int, Sum Int))),
    segM :: !(IO (Seg.SegTree2d RealWorld (Sum Int)))
  }

instance Show Init where
  show Init {..} = show ("Init", capacity, q, debugPoints)

instance QC.Arbitrary Init where
  arbitrary = do
    QC.Positive n <- QC.arbitrary
    q <- QC.chooseInt (1, 2 * n)
    yxs <- (VU.fromList <$>) $ QC.vectorOf n $ do
      x <- QC.chooseInt (-16, 16)
      y <- QC.chooseInt (-16, 16)
      pure (x, y)
    let vecM = VU.thaw $ VU.map (\(!x, !y) -> (x, y, mempty)) yxs
    pure $ Init n q yxs vecM (Seg.new yxs)

data Query
  = -- Read !Int |
    Write !Int !Int
  | ModifyAdd !Int !Int
  | Prod !(Int, Int, Int, Int)
  | AllProd
  | Count !(Int, Int, Int, Int)
  deriving
    ( Eq,
      Show
    )

genQuery :: Int -> QC.Gen Query
genQuery n = do
  QC.oneof
    [ -- Read <$> i,
      Write <$> i <*> val,
      ModifyAdd <$> val <*> i,
      Prod <$> rect,
      pure AllProd,
      Count <$> rect
    ]
  where
    i = QC.chooseInt (0, n - 1)
    rect = do
      xl <- QC.chooseInt (-18, 18)
      xr <- QC.chooseInt (xl, 18)
      yl <- QC.chooseInt (-18, 18)
      yr <- QC.chooseInt (yl, 18)
      pure (xl, xr, yl, yr)
    val = QC.arbitrary @Int

-- | Arbitrary return type for the `Query` result.
data Result
  = None
  | S !(Sum Int)
  | I !Int
  | MS !(Maybe (Sum Int))
  deriving (Show, Eq)

-- | containers. (referential implementation)
handleRef :: VUM.MVector RealWorld (Int, Int, Sum Int) -> Query -> IO Result
handleRef vec q = case q of
  -- Read i -> do
  --   S . (\(!_, !_, !w) -> w) <$> VGM.read vec i
  Write i v -> do
    (!x, !y, !_) <- VGM.read vec i
    VGM.write vec i (x, y, Sum v)
    pure None
  ModifyAdd w i -> do
    VGM.modify vec (\(!x, !y, !w0) -> (x, y, w0 + Sum w)) i
    pure None
  Prod (!x1, !x2, !y1, !y2) -> do
    vec' <- VU.unsafeFreeze vec
    pure
      . S
      . VU.sum
      . VU.map (\(!_, !_, !w) -> w)
      $ VU.filter (\(!x, !y, !_) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec'
  AllProd -> do
    vec' <- VU.unsafeFreeze vec
    pure . S . VU.sum $ VU.map (\(!_, !_, !w) -> w) vec'
  Count (!x1, !x2, !y1, !y2) -> do
    vec' <- VU.unsafeFreeze vec
    pure
      . I
      . VU.length
      $ VU.filter (\(!x, !y, !_) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec'

handleAcl :: (PrimMonad m) => Seg.SegTree2d (PrimState m) (Sum Int) -> Query -> m Result
handleAcl seg q = case q of
  -- Read i -> do
  --   S <$> Seg.read seg i
  Write i v -> do
    Seg.write seg i $ Sum v
    pure None
  ModifyAdd w i -> do
    Seg.modify seg (+ Sum w) i
    pure None
  Prod (!x1, !x2, !y1, !y2) -> do
    S <$> Seg.prod seg x1 x2 y1 y2
  AllProd -> do
    S <$> Seg.allProd seg
  Count (!x1, !x2, !y1, !y2) -> do
    I <$> Seg.count seg x1 x2 y1 y2

prop_randomTest :: Init -> QC.Property
prop_randomTest Init {..} = QCM.monadicIO $ do
  seg <- QCM.run segM
  vec <- QCM.run vecM
  qs <- QCM.pick $ QC.vectorOf q (genQuery capacity)
  for_ qs $ \query -> do
    expected <- QCM.run $ handleRef vec query
    actual <- QCM.run $ handleAcl seg query
    QCM.assertWith (expected == actual) $ show (query, expected, actual)

unit_zero :: TestTree
unit_zero = testCase "zero" $ do
  seg <- Seg.build @_ @(Sum Int) VU.empty VU.empty VU.empty
  (@?= mempty) =<< Seg.prod seg 0 10 0 10
  (@?= mempty) =<< Seg.allProd seg
  (@?= 0) =<< Seg.count seg 0 10 0 10

tests :: [TestTree]
tests =
  [ unit_zero,
    QC.testProperty "random test" prop_randomTest
  ]