-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at http://mozilla.org/MPL/2.0/.

{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}

module Test.Data.Trie.Pattern (tests) where

import Control.Monad
import Data.ByteString (ByteString)
import Data.Foldable
import Data.Functor.Identity
import Data.Functor.Compose
import Data.List (inits)
import Data.Sequence (Seq (..))
import Data.Trie.Pattern
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.QuickCheck

import qualified Data.ByteString.Char8 as C8
import qualified Data.Sequence         as Seq
import qualified Data.Trie.Pattern     as Trie

tests :: TestTree
tests = testGroup "Data.Trie.Pattern"
    [ testGroup "Semigroup"
        [ testProperty "Left preference" checkSemigroupLeftPref
        , testProperty "Associativity" checkSemigroupAssoc
        ]
    , testGroup "Monoid"
        [ testProperty "Identity" checkMonoidId
        ]
    , testGroup "Functor"
        [ testProperty "Identity" checkFunctorId
        , testProperty "Composition" checkFunctorComp
        ]
    , testGroup "Traversable"
        [ testProperty "Identity" checkTraversableId
        , testProperty "Composition" checkTraversableComp
        ]
    , testProperty "capture" checkCapture
    , testProperty "overlapping" checkOverlapping
    , testProperty "match" checkMatch
    , testProperty "match (overlapping)" checkMatchOverlapping
    , testProperty "match (partial overlap)" checkMatchPartialOverlap
    , testProperty "lookup" checkLookup
    , testProperty "insert" checkInsert
    , testProperty "from/to list" checkListConversion
    , testProperty "delete" checkDelete
    , testProperty "adjust" checkAdjust
    , testProperty "null" checkNull
    , testProperty "value" checkValue
    , testProperty "lookupPrefix" checkLookupPrefix
    , testProperty "matchPrefix" checkMatchPrefix
    ]

-------------------------------------------------------------------------------
-- Basic properties

checkNull :: Property
checkNull = once (null (mempty :: Trie ByteString Int))
    .&&. (forAll genTrie (not . null))

checkValue :: Property
checkValue = forAll genTrie $ \t ->
    Trie.value t == Trie.lookup Empty t

checkListConversion :: Property
checkListConversion = forAll genTrie $ \t ->
    Trie.fromAssocList (Trie.toAssocList t) == t

checkCapture :: Property
checkCapture = forAll genPatternMatch $ \(p, s) ->
    Trie.unapplyCapture p (Trie.applyCapture s p) == s

checkOverlapping :: Property
checkOverlapping = forAll genTwinPatterns $ \(p1, p2) ->
        -- Irreflexive
        not (Trie.overlapping p1 p1)
        .&&.
        -- Symmetric
        (Trie.overlapping p1 p2 ==> Trie.overlapping p2 p1)

-------------------------------------------------------------------------------
-- Properties of lookups and matching

checkMatch :: Property
checkMatch = forAll genPatterns check
  where
    check patterns =
        let t = Trie.fromAssocList patterns
        in conjoin . flip map patterns $ \(p, a) ->
            forAll (genStr p) $ \s ->
                let c = Trie.applyCapture s p
                in Trie.match s t == Just (a, c)

checkMatchOverlapping :: Property
checkMatchOverlapping = forAll genPatternMatch $ \(p, s) ->
    let -- Overlaps with 'p' w.r.t 's'
        p' = Seq.fromList (map Trie.EqStr s)
        t  = Trie.fromAssocList [(p, 1), (p', 2)] :: Trie ByteString Int
        c  = Trie.applyCapture s p
    in
        (Trie.match s t == Just (2, Seq.empty))
        .&&.
        (Trie.matchPrefix s t == Just (2, Seq.empty, []))
        .&&.
        (p /= p' ==> Trie.overlapping p p')
        .&&.
        (p /= p' ==> Trie.matchOrd p' > Trie.matchOrd p)
        .&&.
        (p /= p' ==> Trie.matchPrefix (s ++ ["x"]) t == Just (1, c, ["x"])
                    && Trie.matchPrefixOrd p > Trie.matchPrefixOrd p')

-- Partially overlapping patterns (i.e. those with an overlapping
-- proper prefix) are not ambiguous but require backtracking via
-- choice points, since the more specific path is explored first.
checkMatchPartialOverlap :: Property
checkMatchPartialOverlap = forAll genPatternMatch $ \(p, s) ->
    let -- A match for ../a/c requires backtracking to the choice point
        -- before ../a, since the /a/.. branches are explored first.
        pm  = p |> AnyStr    |> EqStr "c" -- matches
        p'  = p |> EqStr "a" |> EqStr "b" -- no match but explored first
        p'' = p |> EqStr "a" |> AnyStr |> AnyStr -- no match but explored first
        s'  = s ++ ["a","c"]
        t   = Trie.fromAssocList [(p',  1), (pm, 2), (p'', 3)] :: Trie ByteString Int
        c   = Trie.applyCapture s' pm
    in
        Trie.match s' t == Just (2, c)

checkLookup :: Property
checkLookup = forAll genPatterns check
  where
    check patterns =
        let t = Trie.fromAssocList patterns
        in conjoin . flip map patterns $ \(p, a) ->
            Trie.lookup p t == Just a

-- Generate a trie containing values for all prefixes of an
-- arbitrary pattern. Then iteratively remove the patterns
-- from the trie, starting with the longest prefix, each time
-- verifying that 'lookupPrefix' yields the right value and
-- remaining suffix.
checkLookupPrefix :: Property
checkLookupPrefix = forAll (genPattern Nothing) $ \p ->
    let
        patterns = toList (Seq.inits p) `zip` [(1::Int)..]
        trie     = Trie.fromAssocList patterns
        check (px, a) ~(t, props) =
            let
                p' = Seq.drop (Seq.length px) p
                ok = Trie.lookupPrefix p t == Just (a, p')
            in
                (Trie.delete px t, ok : props)
    in
        conjoin (snd (foldr check (trie, []) patterns))

-- Generate a trie containing values for all prefixes of an
-- arbitrary pattern. Then iteratively remove the patterns
-- from the trie, starting with the longest prefix, each time
-- verifying that 'matchPrefix' applied to an input string
-- matching the entire pattern yields the right value, captured
-- chunks and remaining suffix.
checkMatchPrefix :: Property
checkMatchPrefix = forAll genPatternMatch $ \(p, s) ->
    let
        patterns = toList (Seq.inits p) `zip` [(1::Int)..]
        trie     = Trie.fromAssocList patterns
        inputs   = patterns `zip` inits s
        check ((px, a), sx) ~(t, props) =
            let
                s' = drop (length sx) s
                cs = Trie.applyCapture sx px
                ok = Trie.matchPrefix s t == Just (a, cs, s')
            in
                (Trie.delete px t, ok : props)
    in
        conjoin (snd (foldr check (trie, []) inputs))

-------------------------------------------------------------------------------
-- Properties of modifications

checkInsert :: Property
checkInsert = forAll genTrie $ \t ->
    forAll (genPattern Nothing) $ \p ->
        let t' = Trie.insert p 42 t
        in Trie.lookup p t' == Just 42

-- For an arbitrary trie, iteratively delete every pattern, checking
-- the presence and absence of the pattern before and after each
-- deletion, respectively, as well as that the final trie is empty.
checkDelete :: Property
checkDelete = forAll genTrie $ \t ->
    let (t', props) = Trie.foldrWithKey reduce (t, []) t
    in null t' .&&. conjoin props
  where
    reduce p a (t, props) =
        let before = Trie.lookup p t  == Just a
            t'     = Trie.delete p t
            after  = Trie.lookup p t' == Nothing
        in (t', before : after : props)

checkAdjust :: Property
checkAdjust = forAll genTrie $ \t ->
    let (p, a) = head (Trie.toAssocList t)
        t' = Trie.adjust p (+ 1) t
    in Trie.lookup p t' == Just (a + 1)     &&
       Trie.delete p t  == Trie.delete p t'

-------------------------------------------------------------------------------
-- Semigroup and monoid properties of tries

checkSemigroupLeftPref :: Property
checkSemigroupLeftPref = forAll genTrie $ \t ->
    forAll (genPattern Nothing) $ \p ->
        let t'  = Trie.insert p 1 t
            t'' = Trie.insert p 2 t
        in (t' <> t'') == t'

checkSemigroupAssoc :: Property
checkSemigroupAssoc = forAll (replicateM 3 genTrie) $ \[t1,t2,t3] ->
    (t1 <> t2) <> t3 == t1 <> (t2 <> t3)

checkMonoidId :: Property
checkMonoidId = forAll genTrie $ \t ->
    t <> mempty == t && mempty <> t == t

-------------------------------------------------------------------------------
-- Functor laws for tries

checkFunctorId :: Property
checkFunctorId = forAll genTrie $ \t ->
    fmap id t == id t

checkFunctorComp :: Property
checkFunctorComp = forAll genTrie $ \t ->
    fmap (f . g) t == (fmap f . fmap g) t
  where
    f, g :: Int -> Int
    f x = x + 1
    g x = x * 2

-------------------------------------------------------------------------------
-- Traversable laws for tries

checkTraversableId :: Property
checkTraversableId = forAll genTrie $ \t ->
    traverse Identity t == Identity t

checkTraversableComp :: Property
checkTraversableComp = forAll genTrie $ \t ->
    traverse (Compose . fmap g . f) t
    ==
    (Compose . fmap (traverse g) . traverse f) t
  where
    f, g :: Int -> Maybe Int
    f x = Just (x + 1)
    g x = Just (x * 2)

-------------------------------------------------------------------------------
-- Generators

genByteString :: Gen ByteString
genByteString = C8.pack <$> listOf1 arbitraryASCIIChar

-- Generate an input string matching a given pattern.
genStr :: Pattern ByteString -> Gen (Str ByteString)
genStr p = mapM gen (toList p)
  where
    gen  Trie.AnyStr   = genByteString
    gen (Trie.EqStr s) = pure s

genPattern :: Maybe ByteString -> Gen (Pattern ByteString)
genPattern prefix = do
    n <- choose (1, 10)
    s <- vectorOf n genMatcher
    return $ Seq.fromList (maybe s ((:s) . Trie.EqStr) prefix)

-- | Generate two (possibly overlapping) patterns of the same length.
genTwinPatterns :: Gen (Pattern ByteString, Pattern ByteString)
genTwinPatterns = do
    p1 <- genPattern Nothing
    p2 <- Seq.fromList <$> vectorOf (Seq.length p1) genMatcher
    return (p1, p2)

-- Generate an arbitrary pattern together with a matching
-- input string.
genPatternMatch :: Gen (Pattern ByteString, Str ByteString)
genPatternMatch = do
    p <- genPattern Nothing
    s <- genStr p
    return (p, s)

genMatcher :: Gen (Matcher ByteString)
genMatcher = oneof [str, var]
  where
    str = Trie.EqStr <$> genByteString
    var = pure Trie.AnyStr

-- Generate 1-100 non-overlapping patterns
genPatterns :: Gen [(Pattern ByteString, Int)]
genPatterns = do
    n <- choose (1, 100)
    r <- mapM (genPattern . Just . C8.pack . show) [1..n]
    return $ r `zip` [1..n]

genTrie :: Gen (Trie ByteString Int)
genTrie = Trie.fromAssocList <$> genPatterns