module Language.Haskell.Liquid.Bag where

import qualified Data.Map      as M

{-@ embed   Data.Map.Map as Map_t                                         @-}

{-@ measure Map_default :: Int -> Bag a                                   @-}
{-@ measure Map_union   :: Bag a -> Bag a -> Bag a                        @-}
{-@ measure Map_select  :: Data.Map.Map k v -> k -> v                     @-}
{-@ measure Map_store   :: Data.Map.Map k v -> k -> v -> Data.Map.Map k v @-}
{-@ measure bagSize     :: Bag k -> Int                                   @-}

-- if I just write measure fromList the measure definition is not imported
{-@ measure fromList :: [k] -> Bag k
      fromList ([])   = (Map_default 0)
      fromList (x:xs) = Map_store (fromList xs) x (1 + (Map_select (fromList xs) x))
@-}


type Bag a = M.Map a Int

{-@ assume empty :: {v:Bag k | v = Map_default 0} @-}
empty :: Bag k
empty :: Bag k
empty = Bag k
forall k a. Map k a
M.empty

{-@ assume bagSize :: b:Bag k -> {i:Nat | i == bagSize b} @-}
bagSize :: Bag k -> Int 
bagSize :: Bag k -> Int
bagSize Bag k
b = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Bag k -> [Int]
forall k a. Map k a -> [a]
M.elems Bag k
b) 

{-@ fromList :: (Ord k) => xs:[k] -> {v:Bag k | v == fromList xs } @-} 
fromList :: (Ord k) => [k] -> Bag k 
fromList :: [k] -> Bag k
fromList []     = Bag k
forall k. Bag k
empty
fromList (k
x:[k]
xs) = k -> Bag k -> Bag k
forall k. Ord k => k -> Bag k -> Bag k
put k
x ([k] -> Bag k
forall k. Ord k => [k] -> Bag k
fromList [k]
xs)

{-@ assume get :: (Ord k) => k:k -> b:Bag k -> {v:Nat | v = Map_select b k}  @-}
get :: (Ord k) => k -> Bag k -> Int
get :: k -> Bag k -> Int
get k
k Bag k
m = Int -> k -> Bag k -> Int
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault Int
0 k
k Bag k
m

{-@ assume put :: (Ord k) => k:k -> b:Bag k -> {v:Bag k | v = Map_store b k (1 + (Map_select b k))} @-}
put :: (Ord k) => k -> Bag k -> Bag k
put :: k -> Bag k -> Bag k
put k
k Bag k
m = k -> Int -> Bag k -> Bag k
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ k -> Bag k -> Int
forall k. Ord k => k -> Bag k -> Int
get k
k Bag k
m) Bag k
m

{-@ assume union :: (Ord k) => m1:Bag k -> m2:Bag k -> {v:Bag k | v = Map_union m1 m2} @-}
union :: (Ord k) => Bag k -> Bag k -> Bag k
union :: Bag k -> Bag k -> Bag k
union Bag k
m1 Bag k
m2 = Bag k -> Bag k -> Bag k
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Bag k
m1 Bag k
m2

{-@ thm_emp :: x:k -> xs:Bag k ->  { Language.Haskell.Liquid.Bag.empty /= put x xs }  @-}
thm_emp :: (Ord k) => k -> Bag k -> ()  
thm_emp :: k -> Bag k -> ()
thm_emp k
x Bag k
xs = () -> Int -> ()
forall a b. a -> b -> a
const () (k -> Bag k -> Int
forall k. Ord k => k -> Bag k -> Int
get k
x Bag k
xs)

{-@ assume thm_size :: xs:[k] ->  { bagSize (fromList xs) == len xs }  @-}
thm_size :: (Ord k) => [k] -> ()  
thm_size :: [k] -> ()
thm_size [k]
_ = ()