tweet_trie.hs

{-# LANGUAGE RankNTypes #-}

import Control.Comonad.Cofree
import Control.Lens hiding ((:<))
import qualified Data.Map as Map
import Data.Map (Map)
import Prelude hiding (lookup)
import Data.Maybe (isJust)
import Test.QuickCheck

type Trie a b = Cofree (Map a) (Maybe b)

string :: Ord a => [a] -> Lens' (Trie a b) (Maybe b)
string =
 foldr
   (\x r -> _unwrap . at x .
      anon
        (Nothing :< mempty)
        (\(v :< m) -> null v && null m) . r)
   _extract
    

insert :: Ord a => [a] -> b -> Trie a b -> Trie a b
insert xs x = string xs .~ Just x

lookup :: Ord a => [a] -> Trie a b -> Maybe b
lookup = view . string

delete :: Ord a => [a] -> Trie a b -> Trie a b
delete xs = string xs .~ Nothing

member :: Ord a => [a] -> Trie a b -> Bool
member xs = isJust . lookup xs

fromList :: Ord a => [([a], b)] -> Trie a b
fromList = foldr (uncurry insert) (Nothing :< mempty)

trie :: Gen (Trie Word Word)
trie = fmap fromList arbitrary

lookupInsert :: Property
lookupInsert =
  forAll arbitrary $ \s ->
  forAll arbitrary $ \x ->
  forAll trie      $ \t ->
  lookup s (insert s x t) === Just x
  
deleteInsert :: Property
deleteInsert = 
  forAll arbitrary $ \s ->
  forAll arbitrary $ \x ->
  forAll trie      $ \t ->
  not (member s t) ==> (delete s (insert s x t) === t)
  
memberDelete :: Property
memberDelete =
  forAll arbitrary $ \s ->
  forAll trie      $ \t ->
  not (member s (delete s t))