{-# LANGUAGE RecordWildCards #-} module Tests.Extra.KdTree where import AtCoder.Extra.KdTree qualified as Kt import Data.Foldable (for_) import Data.Vector.Algorithms.Intro qualified as VAI import Data.Vector.Unboxed qualified as VU import Test.QuickCheck.Monadic as QCM import Test.Tasty import Test.Tasty.QuickCheck as QC import Tests.Util import Test.Tasty.HUnit import Debug.Trace data Init = Init { n :: {-# UNPACK #-} !Int, q :: {-# UNPACK #-} !Int, refVec :: !(VU.Vector (Int, Int)), kt :: !Kt.KdTree } rng :: (Int, Int) rng = (-rngI, rngI) rngI :: Int rngI = 10 instance Show Init where show Init {..} = show n instance QC.Arbitrary Init where arbitrary = do -- n <- QC.chooseInt (1, 256) n <- QC.chooseInt (1, 8) -- q <- QC.chooseInt (1, 256) q <- QC.chooseInt (1, 4) refVec <- (VU.fromList <$>) $ QC.vectorOf n $ do x <- QC.chooseInt rng y <- QC.chooseInt rng pure (x, y) let !_ = traceShow refVec () let kt = Kt.build2 refVec pure Init {..} data Query = FindPointsIn ((Int, Int), (Int, Int)) | FindNearestPoint (Int, Int) deriving (Show) -- | Arbitrary return type for the `Query` result. data Result = MI !(Maybe Int) | V !(VU.Vector Int) deriving (Show, Eq) queryGen :: Int -> QC.Gen Query queryGen _n = do QC.oneof [ FindPointsIn <$> r, FindNearestPoint <$> xy ] where r = (,) <$> intervalGen' (-rngI) rngI <*> intervalGen' (-rngI) rngI xy = intervalGen' (-rngI) rngI -- | containers. (referencial implementation) handleRef :: VU.Vector (Int, Int) -> Query -> Result handleRef vec q = case q of FindPointsIn ((!x1, !x2), (!y1, !y2)) -> do V $ VU.findIndices (\(!x, !y) -> x1 <= x && x < x2 && y1 <= y && y < y2) vec FindNearestPoint (!x, !y) | VU.null vec -> MI Nothing | otherwise -> MI . Just . VU.minIndex $ VU.map (\(!x', !y') -> (x - x') * (x - x') + (y - y') * (y - y')) vec -- | ACL handleAcl :: Kt.KdTree -> Query -> Result handleAcl kt q = case q of FindPointsIn ((!x1, !x2), (!y1, !y2)) -> V $ Kt.findPointsIn kt x1 x2 y1 y2 (Kt.nKt kt) FindNearestPoint (!x, !y) -> MI $ Kt.findNearestPoint kt x y prop_randomTest :: Init -> QC.Property prop_randomTest Init {..} = QCM.monadicIO $ do qs <- QCM.pick $ QC.vectorOf q (queryGen n) for_ qs $ \query -> do let expected = handleRef refVec query let res = handleAcl kt query QCM.assertWith (p query expected res) $ show (query, expected, res) where p (FindNearestPoint (!x, !y)) (MI (Just a)) (MI (Just b)) = let (!x1, !y1) = refVec VU.! a (!x2, !y2) = refVec VU.! b d1 = (x1 - x) * (x1 - x) + (y1 - y) * (y1 - y) d2 = (x2 - x) * (x2 - x) + (y2 - y) * (y2 - y) in a == b || d1 == d2 p (FindPointsIn _) (V a) (V b) = VU.modify VAI.sort a == VU.modify VAI.sort b p _ a b = a == b unit_zero :: TestTree unit_zero = testCase "zero" $ do let kt = Kt.build VU.empty VU.empty (@?= VU.empty) $ Kt.findPointsIn kt 0 10 0 10 5 (@?= Nothing) $ Kt.findNearestPoint kt 10 10 pure () tests :: [TestTree] tests = [ QC.testProperty "randomTest" prop_randomTest, unit_zero ]