{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Main (main) where

import Test.ChasingBottoms.IsBottom
import Test.Tasty (TestTree, TestName, defaultMain, testGroup)
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck (testProperty, Arbitrary(arbitrary), Fun)
#if __GLASGOW_HASKELL__ >= 806
import Test.Tasty.QuickCheck (Property)
#endif
import Test.QuickCheck.Function (apply)

import Data.Strict.Map.Autogen.Strict (Map)
import qualified Data.Strict.Map.Autogen.Strict as M
import qualified Data.Map.Lazy as L

import Utils.IsUnit
#if __GLASGOW_HASKELL__ >= 806
import Utils.NoThunks
#endif

instance (Arbitrary k, Arbitrary v, Ord k) =>
         Arbitrary (Map k v) where
    arbitrary = M.fromList `fmap` arbitrary

apply2 :: Fun (a, b) c -> a -> b -> c
apply2 f a b = apply f (a, b)

apply3 :: Fun (a, b, c) d -> a -> b -> c -> d
apply3 f a b c = apply f (a, b, c)

------------------------------------------------------------------------
-- * Properties

------------------------------------------------------------------------
-- ** Strict module

pSingletonKeyStrict :: Int -> Bool
pSingletonKeyStrict v = isBottom $ M.singleton (bottom :: Int) v

pSingletonValueStrict :: Int -> Bool
pSingletonValueStrict k = isBottom $ (M.singleton k (bottom :: Int))

pFindWithDefaultKeyStrict :: Int -> Map Int Int -> Bool
pFindWithDefaultKeyStrict def m = isBottom $ M.findWithDefault def bottom m

pFindWithDefaultValueStrict :: Int -> Map Int Int -> Bool
pFindWithDefaultValueStrict k m =
    M.member k m || (isBottom $ M.findWithDefault bottom k m)

pAdjustKeyStrict :: Fun Int Int -> Map Int Int -> Bool
pAdjustKeyStrict f m = isBottom $ M.adjust (apply f) bottom m

pAdjustValueStrict :: Int -> Map Int Int -> Bool
pAdjustValueStrict k m
    | k `M.member` m = isBottom $ M.adjust (const bottom) k m
    | otherwise       = case M.keys m of
        []     -> True
        (k':_) -> isBottom $ M.adjust (const bottom) k' m

pInsertKeyStrict :: Int -> Map Int Int -> Bool
pInsertKeyStrict v m = isBottom $ M.insert bottom v m

pInsertValueStrict :: Int -> Map Int Int -> Bool
pInsertValueStrict k m = isBottom $ M.insert k bottom m

pInsertWithKeyStrict :: Fun (Int, Int) Int -> Int -> Map Int Int -> Bool
pInsertWithKeyStrict f v m = isBottom $ M.insertWith (apply2 f) bottom v m

pInsertWithValueStrict :: Fun (Int, Int) Int -> Int -> Int -> Map Int Int
                       -> Bool
pInsertWithValueStrict f k v m
    | M.member k m = (isBottom $ M.insertWith (const2 bottom) k v m) &&
                     not (isBottom $ M.insertWith (const2 1) k bottom m)
    | otherwise    = isBottom $ M.insertWith (apply2 f) k bottom m

pInsertLookupWithKeyKeyStrict :: Fun (Int, Int, Int) Int -> Int
                              -> Map Int Int -> Bool
pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey (apply3 f) bottom v m

pInsertLookupWithKeyValueStrict :: Fun (Int, Int, Int) Int -> Int -> Int
                                -> Map Int Int -> Bool
pInsertLookupWithKeyValueStrict f k v m
    | M.member k m = (isBottom $ M.insertLookupWithKey (const3 bottom) k v m) &&
                     not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m)
    | otherwise    = isBottom $ M.insertLookupWithKey (apply3 f) k bottom m

#if __GLASGOW_HASKELL__ >= 806
pStrictFoldr' :: Map Int Int -> Property
pStrictFoldr' m = whnfHasNoThunks (M.foldr' (:) [] m)
#endif

#if __GLASGOW_HASKELL__ >= 806
pStrictFoldl' :: Map Int Int -> Property
pStrictFoldl' m = whnfHasNoThunks (M.foldl' (flip (:)) [] m)
#endif

#if __GLASGOW_HASKELL__ >= 806
pStrictFoldrWithKey' :: Map Int Int -> Property
pStrictFoldrWithKey' m = whnfHasNoThunks (M.foldrWithKey' (\_ a as -> a : as) [] m)
#endif

#if __GLASGOW_HASKELL__ >= 806
pStrictFoldlWithKey' :: Map Int Int -> Property
pStrictFoldlWithKey' m = whnfHasNoThunks (M.foldlWithKey' (\as _ a -> a : as) [] m)
#endif

------------------------------------------------------------------------
-- check for extra thunks
--
-- These tests distinguish between `()`, a fully evaluated value, and
-- things like `id ()` which are extra thunks that should be avoided
-- in most cases. An exception is `L.fromListWith const`, which cannot
-- evaluate the `const` calls.

tExtraThunksM :: TestTree
tExtraThunksM = testGroup "Map.Strict - extra thunks" $
    if not isUnitSupported then [] else
    -- for strict maps, all the values should be evaluated to ()
    [ check "singleton"           $ m0
    , check "insert"              $ M.insert 42 () m0
    , check "insertWith"          $ M.insertWith const 42 () m0
    , check "fromList"            $ M.fromList [(42,()),(42,())]
    , check "fromListWith"        $ M.fromListWith const [(42,()),(42,())]
    , check "fromAscList"         $ M.fromAscList [(42,()),(42,())]
    , check "fromAscListWith"     $ M.fromAscListWith const [(42,()),(42,())]
    , check "fromDistinctAscList" $ M.fromAscList [(42,())]
    ]
  where
    m0 = M.singleton 42 ()
    check :: TestName -> M.Map Int () -> TestTree
    check n m = testCase n $ case M.lookup 42 m of
        Just v -> assertBool msg (isUnit v)
        _      -> assertBool "key not found" False
      where
        msg = "too lazy -- expected fully evaluated ()"

tExtraThunksL :: TestTree
tExtraThunksL = testGroup "Map.Lazy - extra thunks" $
    if not isUnitSupported then [] else
    -- for lazy maps, the *With functions should leave `const () ()` thunks,
    -- but the other functions should produce fully evaluated ().
    [ check "singleton"       True  $ m0
    , check "insert"          True  $ L.insert 42 () m0
    , check "insertWith"      False $ L.insertWith const 42 () m0
    , check "fromList"        True  $ L.fromList [(42,()),(42,())]
    , check "fromListWith"    False $ L.fromListWith const [(42,()),(42,())]
    , check "fromAscList"     True  $ L.fromAscList [(42,()),(42,())]
    , check "fromAscListWith" False $ L.fromAscListWith const [(42,()),(42,())]
    , check "fromDistinctAscList" True $ L.fromAscList [(42,())]
    ]
  where
    m0 = L.singleton 42 ()
    check :: TestName -> Bool -> L.Map Int () -> TestTree
    check n e m = testCase n $ case L.lookup 42 m of
        Just v -> assertBool msg (e == isUnit v)
        _      -> assertBool "key not found" False
      where
        msg | e         = "too lazy -- expected fully evaluated ()"
            | otherwise = "too strict -- expected a thunk"

------------------------------------------------------------------------
-- * Test list

tests :: [TestTree]
tests =
    [
    -- Basic interface
      testGroup "Map.Strict"
      [ testProperty "singleton is key-strict" pSingletonKeyStrict
      , testProperty "singleton is value-strict" pSingletonValueStrict
      , testProperty "member is key-strict" $ keyStrict M.member
      , testProperty "lookup is key-strict" $ keyStrict M.lookup
      , testProperty "findWithDefault is key-strict" pFindWithDefaultKeyStrict
      , testProperty "findWithDefault is value-strict" pFindWithDefaultValueStrict
      , testProperty "! is key-strict" $ keyStrict (flip (M.!))
      , testProperty "delete is key-strict" $ keyStrict M.delete
      , testProperty "adjust is key-strict" pAdjustKeyStrict
      , testProperty "adjust is value-strict" pAdjustValueStrict
      , testProperty "insert is key-strict" pInsertKeyStrict
      , testProperty "insert is value-strict" pInsertValueStrict
      , testProperty "insertWith is key-strict" pInsertWithKeyStrict
      , testProperty "insertWith is value-strict" pInsertWithValueStrict
      , testProperty "insertLookupWithKey is key-strict"
        pInsertLookupWithKeyKeyStrict
      , testProperty "insertLookupWithKey is value-strict"
        pInsertLookupWithKeyValueStrict
#if __GLASGOW_HASKELL__ >= 806
      , testProperty "strict foldr'" pStrictFoldr'
      , testProperty "strict foldl'" pStrictFoldl'
      , testProperty "strict foldrWithKey'" pStrictFoldrWithKey'
      , testProperty "strict foldlWithKey'" pStrictFoldlWithKey'
#endif
      ]
      , tExtraThunksM
      , tExtraThunksL
    ]

------------------------------------------------------------------------
-- * Test harness

main :: IO ()
main = defaultMain $ testGroup "map-strictness" tests

------------------------------------------------------------------------
-- * Utilities

keyStrict :: (Int -> Map Int Int -> a) -> Map Int Int -> Bool
keyStrict f m = isBottom $ f bottom m

const2 :: a -> b -> c -> a
const2 x _ _ = x

const3 :: a -> b -> c -> d -> a
const3 x _ _ _ = x