module Table where

import Test.QuickCheck
import Data.List

-------------------------------------------------------------------

data Table k v
  = Join (Table k v) k v (Table k v)
  | Empty
 deriving ( Eq, Show )

-------------------------------------------------------------------

lookupT :: Ord k => k -> Table k v -> Maybe v
lookupT key Empty = Nothing
lookupT key (Join left k v right)
  | key < k  = lookupT key left
  | key == k = Just v
  | key > k  = lookupT key right

-------------------------------------------------------------------

insertT :: Ord k => k -> v -> Table k v -> Table k v
insertT key val Empty = Join Empty key val Empty
insertT key val (Join left k v right)
  | key > k  = Join left k v (insertT key val right)
--   | key <= k = Join (insertT key val left) k v right
  | key < k  = Join (insertT key val left) k v right
  | key == k = Join left k val right

-------------------------------------------------------------------

-- properties

prop_lookupT :: Int -> Table Int Int -> Bool
prop_lookupT k t =
  lookupT k t == lookup k (contents t)

prop_insertT :: Int -> Int -> Table Int Int -> Bool
prop_insertT k v t =
--   insert (k,v) (contents t) ==
  insert (k,v) [(k',v') | (k',v') <- contents t, k' /= k] ==
     contents (insertT k v t)

prop_lookup_insert :: Int -> Int -> Int -> Table Int Int -> Bool
prop_lookup_insert k' k v t =
  lookupT k' (insertT k v t) ==
    if k == k' then Just v else lookupT k' t

-------------------------------------------------------------------

contents :: Table k v -> [(k,v)]
contents Empty          = []
contents (Join l k v r) = contents l ++ [(k,v)] ++ contents r

-------------------------------------------------------------------

arbTable :: (Arbitrary k, Arbitrary v) => Int -> Gen (Table k v)
arbTable s =
  frequency
  [ (1, return Empty)
  , (s, do x <- arbitrary
           y <- arbitrary
           l <- arbTable (s `div` 2)
           r <- arbTable (s `div` 2)
           return (Join l x y r))
  ]

-------------------------------------------------------------------

ordered :: Ord a => [a] -> Bool
ordered []       = True
ordered [x]      = True
ordered (x:y:xs) = x <= y && ordered (y:xs)

prop_invTable :: Table Integer Integer -> Bool
prop_invTable tab =
--     ordered ks
    ordered ks && ks == nub ks
 where
  ks = [ k | (k,v) <- contents tab ]

-------------------------------------------------------------------

instance (Arbitrary k, Ord k, Arbitrary v) => Arbitrary (Table k v) where
--   arbitrary = sized arbTable
  arbitrary =
    do kvs <- arbitrary
       return (table kvs)

-- table kvs converts a list of key-value pairs into a Table
-- satisfying the ordering invariant
table :: Ord k => [(k,v)] -> Table k v
table []          = Empty
table ((k,v):kvs) = Join (table smaller) k v (table larger)
 where
--   smaller = [(k',v') | (k',v') <- kvs, k' <= k]
  smaller = [(k',v') | (k',v') <- kvs, k' < k]
  larger  = [(k',v') | (k',v') <- kvs, k' > k]

-------------------------------------------------------------------

main = do
    quickCheck prop_lookupT
    quickCheck prop_insertT
    quickCheck prop_lookup_insert
    quickCheck prop_invTable