{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

{-# OPTIONS -Wall #-}

{- |
This module contains a (very simple) property based test for the RadixTree data
structure.

It uses QuickCheck to generate a list of operations (insert, delete, lookup),
executes them both against
 * "Data.Map" as a known-to-be-good specification and
 * RadixTree (with "Data.Map" as the backend)
records the resulting values in each case and compares these lists.

This is very rudimentary. In particular, 'generateOps' can be further tweaked
to produce “interesting” sequences -- there is little (but some) value in
deleting mostly entries that are not there.

You can sample what 'generateOps' produces by opening this file in ghci
(e.g. cabal new-repl property-tests) and then running @sample generateOps@.

-}
module Properties (tests) where

import qualified Data.Map as M
import Data.Maybe (isJust)
import Data.ByteString.Char8 (ByteString, pack)
import Data.ByteString.Short (fromShort)
import Test.QuickCheck
import Test.Tasty
import Test.Tasty.QuickCheck
import Data.List
import Data.Bifunctor
import Control.Monad.State.Strict

import DFINITY.RadixTree
import Types

{-
In this module we test the RadixTree module against a pure implementation of
a map (Data.Map). We generate random sequences of Insert, Delete and Lookup
calls, run them in both implementations, and compare the results.
 -}

generateOps :: Gen [Op]
generateOps = do
    -- Create a small, non-zero number of keys to consider
    len <- getSmall . getPositive <$> arbitrary
    keys <- vectorOf len arbBS
    let pickKey = elements keys
    listOf $ frequency
        [ (1, Insert <$> pickKey <*> arbBS)
        , (1, Delete <$> pickKey)
        , (2, Lookup <$> pickKey <*> pure Nothing)
        ]

arbBS :: Gen ByteString
arbBS = pack <$> arbitrary


runPure :: [Op] -> [Maybe ByteString]
runPure ops = snd $ mapAccumL go M.empty ops
  where
    go m (Insert k v) = (M.insert k v m, Nothing)
    go m (Delete k  ) = (M.delete k m,   Nothing)
    go m (Lookup k _) = (m,              M.lookup k m)


type M = State (M.Map ByteString ByteString)

runRadix :: [Op] -> [Maybe ByteString]
runRadix ops0 = evalState (initTree >>= go ops0) M.empty
  where
    go :: [Op] -> RadixTree () -> M [Maybe ByteString]
    go []                 _ = return []
    go (Insert k v : ops) t = do
        t' <- insertRadixTree k v t
        (Nothing :) <$> go ops t'
    go (Delete k   : ops) t = do
        t' <- deleteRadixTree k t
        (Nothing :) <$> go ops t'
    go (Lookup k _ : ops) t =
        lookupRadixTree k t >>= \case
            Nothing     -> (Nothing :) <$> go ops t
            Just (v,t') -> (Just v :)  <$> go ops t'

prop_lookup :: Property
prop_lookup  = forAll generateOps $ \ops ->
    runPure ops === runRadix ops

{-
We want to test if two different ways of generating the same map yield the same
state root. We first generate one final state (a Data.Map). Then we generate two
different set of operations that should yield that state. Then we run them and
compare the output.
-}

genMap :: Gen (M.Map ByteString ByteString)
genMap = M.fromList . map (bimap pack pack) . M.toList <$> arbitrary

genOpsForMap :: M.Map ByteString ByteString -> Gen [Op]
genOpsForMap m = do
    inserts <- forM (M.toList m) $ \(k,v) -> do
        ops <- genOpsForKey k
        return $ ops ++ [ Insert k v ]
    -- Extra keys (which we delete as the last action)
    extra_keys <- listOf $ arbBS `suchThat` (`M.notMember` m)
    deletes <- forM extra_keys $ \k -> do
        ops <- genOpsForKey k
        return $ ops ++ [ Delete k ]
    interleaves (inserts ++ deletes)
  where
    genOpsForKey k = listOf (oneof [ Insert k <$> arbBS , pure (Delete k) ])


-- TODO
interleaves :: [[a]] -> Gen [a]
interleaves xss = frequency'
    [ (length (x:xs), (x:) <$> interleaves (xs : rest))
    | ((x:xs):rest) <- rots xss ]
 where
   frequency' [] = return []
   frequency' xs = frequency xs
   rots xs = tail $ zipWith (++) (tails xs) (inits xs)


type TreeGenOps = [ (ByteString, [Maybe ByteString]) ]

prop_stateRoot :: Property
prop_stateRoot = forAll genMap $ \m ->
    forAll (genOpsForMap m) $ \ops1 ->
    forAll (genOpsForMap m) $ \ops2 ->
    run ops1 === run ops2
  where
    run :: [Op] -> ByteString
    run ops0 = evalState (initTree >>= go ops0) M.empty
      where
        go :: [Op] -> RadixTree () -> M ByteString
        go []                 t = fromShort . fst <$> merkleizeRadixTree t
        go (Insert k v : ops) t = insertRadixTree k v t >>= go ops
        go (Delete k   : ops) t = deleteRadixTree k t   >>= go ops
        go (Lookup _ _ : _  ) _ = error "no lookup in this test please"

-- | 'createRadixProof' should find the same value as 'lookupRadixTree'
prop_proofLookup :: Property
prop_proofLookup = forAll arbitrary $ \keyStr ->
  let key = pack keyStr
      proofRes = lookupProof key
      radixRes = lookupRadix key
   in isJust proofRes && isJust radixRes
      && proofRes == radixRes
  where
    run = flip evalState M.empty
    lookupProof k = run $ do
      tree <- initTree
      tree' <- insertRadixTree k (mappend k "suffix") tree
      (_, mTree) <- merkleizeRadixTree tree'
      fmap (getValue . fst) <$> createRadixProof k mTree 

    lookupRadix k = run $ do
      tree <- initTree
      tree' <- insertRadixTree k (mappend k "suffix") tree
      fmap fst <$> lookupRadixTree k tree'

prop_proofValid :: Property
prop_proofValid = forAll arbitrary $ \keyPrefixStr ->
  let keyPrefix = pack keyPrefixStr
      (key, proof, root) = lookupProof keyPrefix
   in case proof of
    Nothing -> False
    Just p -> verifyRadixProof key root p
  where
    run = flip evalState M.empty
    lookupProof kp = run $ do
      tree <- initTree
      tree' <- insertRadixTree (mappend kp "1") "body1" tree
      tree'' <- insertRadixTree (mappend kp "12") "body2" tree'
      let key = mappend kp "123"
      tree''' <- insertRadixTree key "body3" tree''
      (root, mTree) <- merkleizeRadixTree tree'''
      proof <- fmap fst <$> createRadixProof key mTree
      pure (key, proof, root)

tests :: TestTree
tests = testGroup "Property tests"
    [ testProperty "lookup" prop_lookup
    , testProperty "contents" prop_stateRoot
    , testProperty "createRadixProof ~ lookupRadixTree" prop_proofLookup
    , testProperty "createRadixProofs are always valid" prop_proofValid
    ]

initTree :: M (RadixTree ())
initTree = createRadixTree 262144 2028 Nothing ()