{-# OPTIONS_GHC -fno-warn-incomplete-patterns -fno-warn-name-shadowing -fno-warn-orphans #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Collections
-- Copyright   :  (c) Jean-Philippe Bernardy, 2006
-- License     :  BSD3
-- Maintainer  :  jeanphilippe.bernardy; google mail.
-- Stability   :  experimental
--
-- This module defines a class framework for collection types. It provides:
--
-- * Classes for the most common type of collections
--
-- * /View/ types to change the type of a collection, so it implements other classes.
-- This allows to use types for purposes that they are not originally designed for. (eg. 'ElemsView')
--
-- * A few generic functions for handling collections.
--
-- * Infix (operator) version of common functions.
-- 
-- Should you need a more precise documentation, "Data.Collections.Properties" lists laws that
-- implementations are entitled to assume.
--
-- The classes defined in this module are intended to give hints about performance.
-- eg. if a function has a @'Map' c k v@ context, this indicates that the function
-- will perform better if @c@ has an efficitent lookup function.
--
-- This class framework is based on ideas found in Simon Peyton Jones, \"/Bulk types with class/\".
-- <http://research.microsoft.com/Users/simonpj/Papers/collections.ps.gz>
-- 
-- Another inspiration source are the examples of MPTC and fuctional dependencies in Oleg Kiselyov's
-- many articles posted to the haskell mailing list.
-- 
--
-- This module name-clashes with a lot of Prelude functions, subsuming those.
-- The user is encouraged to import Prelude hiding the clashing functions.
-- Alternatively, it can be imported @qualified@.
--


{-


Selling points:
  * Unification of Map and Set (required by the below)
  * inclusion of Arrays
  * Good integration with existing base libraries
  * Relative simplicity: few classes, not too many methods, very little redundancy.
  * Reuses the same identifiers as other standard hierarchy modules.
    Conversion from the module-based API to this class-based one should be easy.
  * Comprehensive set of properties that define the behaviour of the classes.
  * Compatibility with GHC and Hugs.

Bad points
  * Extra complexity due to heavy usage of MTPC (although imho it's a matter of getting used to it)

TODO:
 * test with nhc98/hugs
 * add missing functions (partition, ..., ?)
 * optimizations (rules pragmas)
 * see how multimap/multiset fits this scheme.
 * Think about class Map' :: (* -> *) -> * -> $
 * Fix infelicity about null map test; (== mempty).
-}

module Data.Collections 
    (
-- * Classes
-- ** Unfoldable
     Unfoldable(..),
-- ** Collection
     Collection(..),
     SortingCollection(..),
-- ** Map
     Map(..),
     lookupWithDefault,
     unionsWith,
     intersectionWith',
     differenceWith',
     mapWithKey',
     (!),
-- ** Set
     Set(..),
     unions,
     notMember,
     (\\),
-- ** Sequence
     Sequence(..),
     head, 
     tail,
     append,
     concat,
     concatMap,
--     length,
     (<|),
     (|>),
     (><),
-- ** Others
     Array(..),
     Indexed(..),


-- * Conversions
     fromFoldable,
     fromAscFoldable,
     fromList,
     fromListWith,
     fromAscList,

-- * Views
     KeysView(..), ElemsView(..),
     withKeys, withElems,
-- * Foldable
     module Data.Collections.Foldable,


    ) where 

-- import Prelude (Bool(..), Int, Maybe(..),
--                 (==), (.), (+), ($), (-), (&&), (||),
--                 Eq, Ord, 
--                 error, const, not, fst, snd, maybe, head, otherwise, curry, uncurry, flip,
--                 min, max, Show)

import Prelude hiding (sum,concat,lookup,map,filter,foldr,foldr1,foldl,null,reverse,(++),minimum,maximum,all,elem,concatMap,drop,head,tail,init)

import Data.Monoid
import Data.Collections.Foldable

import qualified Data.Array as Array
import qualified Data.List as List
import qualified Data.Maybe as Maybe

infixl 9 !
infixl 9 \\ --

infixr 5 ><
infixr 5 <|
infixl 5 |>

------------------------------------------------------------------------
-- * Type classes

-- | Class of collection types.

class (Foldable c a, Unfoldable c a) => Collection c a | c -> a where
    -- | @filter f c@ returns the collection of those elements that satisfy the predicate @f@.
    filter :: (a -> Bool) -> c -> c       

-- | Class of collection with unobservable elements. It is the dual of the 'Foldable' class.

class Unfoldable c i | c -> i where

    -- | \'natural\' insertion of an element into a collection.
    insert :: i -> c -> c
    --insert i c = cofold (\Right c -> Right c; Left (i,c) -> Left (i,Right c)) (Left (i,c)) 
    -- | The empty collection.
    empty :: c 
    empty = unfold (const Nothing) undefined

    -- | Creates a collection with a single element.
    singleton :: i -> c 
    singleton i = insert i empty
                           
    -- | Insert all the elements of a foldable.
    insertMany :: Foldable c' i => c' -> c -> c
    insertMany c' c = foldr insert c c'
    -- At first sight, it looks like the above could just use List instead of any Foldable.
    -- However, it would then be more difficult to ensure that the conversion could be made
    -- very efficient between certain types.

    -- | Same as insertMany, but with the unchecked precondition that the input 'Foldable' is sorted.
    insertManySorted :: Foldable c' i => c' -> c -> c
    insertManySorted = insertMany

unfold :: Unfoldable c a => (b -> Maybe (a, b)) -> b -> c
unfold f b = insertMany (List.unfoldr f b) empty

class Collection c o => SortingCollection c o where
    minView :: c -> Maybe (o,c)

-- isSorted :: (Ord a, Foldable c a) => c -> Bool
-- isSorted = fst . foldr cmp (True, Nothing)
--    where curr `cmp` (acc, prev) = (acc && maybe True (curr <=) prev, Just curr)

-- | Conversion from a Foldable to a Collection.
fromFoldable :: (Foldable f a, Collection c' a) => f -> c'
fromFoldable = flip insertMany empty

-- TODO: Should be specialized (RULE pragmas) so it's efficient when converting from/to set/maps


-- | Conversion from a Foldable to a Collection, with the /unchecked/ precondition that the input is sorted 
fromAscFoldable :: (Foldable f a, Collection c' a) => f -> c'
fromAscFoldable = flip insertManySorted empty

-- | Converts a list into a collection.
fromList :: Collection c a => [a] -> c
fromList = fromFoldable

-- | Converts a list into a collection, with the precondition that the input is sorted.
fromAscList :: Collection c a => [a] -> c
fromAscList = fromAscFoldable


-- | Class of sequential-access types. 
-- In addition of the 'Collection' services, it provides deconstruction and concatenation.
class (Monoid c, Collection c a) => Sequence c a where
    -- | The first @i@ elements of a sequence.
    take :: Int -> c -> c
    -- | Elements of a sequence after the first @i@.
    drop :: Int -> c -> c
    -- | Split a sequence at a given index.
    splitAt :: Int -> c -> (c,c)
    -- | Reverse a sequence.
    reverse :: c -> c
    -- | Analyse the left end of a sequence.
    front :: c -> Maybe (a,c)
    -- | Analyse the right end of a sequence.
    back :: c -> Maybe (c,a)
    -- | Add an element to the left end of a sequence.    
    cons :: a -> c -> c
    -- | Add an element to the right end of a sequence.
    snoc :: c -> a -> c
    -- | The 'isPrefix' function takes two seqences and returns True iff 
    -- the first is a prefix of the second.
    isPrefix :: Eq a => c -> c -> Bool
            
    cons = insert
    isPrefix s1 s2 
        = case front s1 of
            Nothing -> True
            Just (x,xs) -> 
               case front s2 of
                 Nothing -> False
                 Just (y,ys) -> x == y && isPrefix xs ys


-- -- | Length of a sequence
-- length :: Sequence c i o => c -> Int
-- length = size

-- | Concatenate two sequences.
append :: Sequence c a => c -> c -> c
append = mappend

-- TODO: span ?

-- | Infix version of 'cons': add an element to the left end of a sequence.
-- Mnemonic: a triangle with the single element at the pointy end.
(<|) :: Sequence c i => i -> c -> c
(<|) = cons

-- | Infix version of 'snoc': add an element to the right end of a sequence.
-- Mnemonic: a triangle with the single element at the pointy end. 
(|>) :: Sequence c i => c -> i -> c
(|>) = snoc

-- | Infix verion of 'append'. Concatenate two sequences.
(><) :: Sequence c a => c -> c -> c
(><) = append


-- | The concatenation of all the elements of a container of sequences.
concat :: (Sequence s a, Foldable t s) => t -> s
concat = fold

-- | Map a function over all the elements of a container and concatenate
-- the resulting sequences.
concatMap :: (Sequence s b, Foldable t a) => (a -> s) -> t -> s
concatMap = foldMap

head :: Sequence s a => s -> a
head = fst . Maybe.fromJust . front
 
tail :: Sequence s a => s -> s
tail = drop 1

-- | Class of indexed types. 
-- The collection is 'dense': there is no way to /remove/ an element nor for lookup 
-- to return "not found".
--
-- In practice however, most shallow collection types will instanciate this
-- class in addition of 'Map', and leave the responsibility of failure to the caller.
class Indexed c k v | c -> k v where
    -- | @index c k@ returns element associated to @k@
    index :: k -> c -> v                 
    -- | @adjust f k c@ applies @f@ to element associated to @k@ and returns the resulting collection.
    adjust :: (v -> v) -> k -> c -> c 
    -- | if @inDomain k c@, then @index c k@ is guaranteed not to fail.
    inDomain :: k -> c -> Bool
    -- | Constructs a collection identical to the first argument except that it has
    -- been updated by the associations in the right argument.
    -- For example, if @m@ is a 1-origin, @n@ by @n@ matrix, then
    --
    -- > m//[((i,i), 0) | i <- [1..n]]
    --
    -- is the same matrix, except with the diagonal zeroed.    
    (//) :: Foldable l (k,v) => c -> l -> c
    (//) = foldr replace
        where replace (k,v) = adjust (const v) k
    -- | @'accum' f@ takes an array and an association list and accumulates
    -- pairs from the list into the array with the accumulating function @f@.
    -- Thus 'accumArray' can be defined using 'accum':
    accum :: Foldable l (k,v') => (v -> v' -> v) -> c -> l -> c
    accum f = foldr adjust'
        where adjust' (k,v') = adjust (\v->f v v') k

-- | Infix version of 'index', with arguments swapped.
(!) :: Indexed c k v => c -> k -> v
(!) = flip index

class (Array.Ix k, Foldable c (k,v), Indexed c k v) => Array c k v | c -> k v where
    -- | if @(l,r) = bounds c@, then @inDomain k c <==> l <= k <= r@
    bounds :: c -> (k,k)
    -- | Construct an array with the specified bounds and containing values
    -- for given indices within these bounds.
    --
    -- The array is undefined (i.e. bottom) if any index in the list is
    -- out of bounds.  The Haskell 98 Report further specifies that if any
    -- two associations in the list have the same index, the value at that
    -- index is undefined (i.e. bottom).  However in GHC's implementation,
    -- the value at such an index is the value part of the last association
    -- with that index in the list.
    --
    -- Because the indices must be checked for these errors, 'array' is
    -- strict in the bounds argument and in the indices of the association
    -- list, but nonstrict in the values.  Thus, recurrences such as the
    -- following are possible:
    --
    -- > a = array (1,100) ((1,1) : [(i, i * a!(i-1)) | i <- [2..100]])
    --
    -- Not every index within the bounds of the array need appear in the
    -- association list, but the values associated with indices that do not
    -- appear will be undefined (i.e. bottom).
    --
    -- If, in any dimension, the lower bound is greater than the upper bound,
    -- then the array is legal, but empty.  Indexing an empty array always
    -- gives an array-bounds error, but 'bounds' still yields the bounds
    -- with which the array was constructed.
    array :: Foldable l (k,v) => (k,k) -> l -> c


-- | Class of map-like types. (aka. for sparse associative types).
--
-- In opposition of Indexed, Map supports unexisting value for some indices.
class Monoid c => Map c k a | c -> k a where
    -- | Remove a key from the keySet (and therefore the associated value in the Map).
    delete :: k -> c -> c
    delete = alter (const Nothing)

    -- | Tells whether an key is member of the keySet.
    member :: k -> c -> Bool
    member k = Maybe.isJust . lookup k

    -- | Union of two keySets.
    -- When duplicates are encountered, the keys may come from any of the two input sets.
    -- Values come from the map given as first arguement.
    union :: c -> c -> c
    union = unionWith const

    -- | Intersection of two keySets.
    --
    -- When duplicates are encountered, the keys may come from any of the two input sets.
    -- Values come from the map given as first arguement.
    intersection :: c -> c -> c
    intersection = intersectionWith const

    -- | Difference of two keySets.
    -- Difference is to be read infix: @a `difference` b@ returns a set containing the 
    -- elements of @a@ that are also absent from @b@.
    difference :: c -> c -> c
    difference = differenceWith (\_ _-> Nothing)


    -- | @s1 `isSubset` s2@ returns True iff. the keys in s1 form a subset of the keys in s2.
    isSubset :: c -> c -> Bool
    isSubset = isSubmapBy (\_ _->True)

    -- | @s1 `isProperSubset` s2@ returns True iff. @s1 `isProperSubset` s2@ and @s1 /= s2@
    isProperSubset :: c -> c -> Bool
    isProperSubset = isProperSubmapBy (\_ _->True)

-- Follows functions for fully-fledged maps.
    -- | Lookup the value at a given key.
    lookup :: k -> c -> Maybe a

    -- | Change the value associated to a given key. 'Nothing' represents no associated value.
    alter :: (Maybe a -> Maybe a) -> k -> c -> c
    alter f k m = case lookup k m of
                  j@(Just _) -> case f j of
                                Just a'  -> insertWith (\a _ -> a) k a' m
                                Nothing -> delete k m
                  Nothing    -> case f Nothing of
                                Just a'  -> insertWith (\a _ -> a) k a' m
                                Nothing -> m

    -- | Insert with a combining function.
    --
    -- @insertWith f key value m@ 
    -- will insert the pair @(key, value)@ into @m@ if @key@ does
    -- not exist in the map. If the key does exist, the function will
    -- insert the pair @(key, f new_value old_value)@.
    insertWith :: (a -> a -> a) -> k -> a -> c -> c
    insertWith f k a c = alter (\x -> Just $ case x of {Nothing->a;Just a' -> f a a'}) k c

    -- | Convert a 'Foldable' to a 'Map', with a combining function. 
    -- Note the applications of the combining function: 
    -- @fromFoldableWith (+) [(k,x1), (k,x2), ..., (k,xn)] = fromFoldable [(k, xn + (... + (x2 + x1)))]@
    -- or more generally @fromFoldableWith f [(k,x) | x <- l] = fromFoldable [(k,foldl1 (flip f) l)]@
    -- 'foldGroups' is probably less surprising, so use it.
    fromFoldableWith :: Foldable l (k,a) => (a -> a -> a) -> l -> c 
    fromFoldableWith f = foldr (uncurry (insertWith f)) mempty 

    -- | Convert a 'Foldable' to a 'Map', with a combining function.
    -- @foldGroups f a l = let mkGroup g = (fst $ head g, foldr f a (map snd g)) in fromList . map mkGroup . groupBy ((==) `on` fst)) . toList@
    foldGroups :: Foldable l (k,b) => (b -> a -> a) -> a -> l -> c
    foldGroups f a = foldr' (\(k,b) c -> (alter (\x -> Just $ case x of {Nothing->f b a;Just a' -> f b a'}) k c)) mempty

    -- | Apply a function over all values in the map.
    mapWithKey :: (k -> a -> a) -> c -> c

    -- | Union with a combining function. 
    unionWith :: (a -> a -> a) -> c -> c -> c

    -- | Intersection with a combining function.
    intersectionWith :: (a -> a -> a) -> c -> c -> c

    -- | Difference with a combining function.
    differenceWith :: (a -> a -> Maybe a) -> c -> c -> c

    -- | isSubmapBy
    isSubmapBy :: (a -> a -> Bool) -> c -> c -> Bool

    -- | isProperSubmapBy
    isProperSubmapBy :: (a -> a -> Bool) -> c -> c -> Bool     
    -- isProperSubmapBy f m1 m2 = isSubmapBy f m1 m2 && not (isEmpty (differenceWith (\_ _->Nothing) m1 m2))

-- | Tells whether a key is not a member of the keySet.
notMember :: (Map c k a) => k -> c -> Bool
notMember k s = not $ member k s

-- | The expression @('lookupWithDefault' def k map)@ returns
-- the value at key @k@ or returns @def@ when the key is not in the map.
lookupWithDefault :: (Map c k a) => a -> k -> c -> a
lookupWithDefault a k c = Maybe.fromMaybe a (lookup k c)

-- | Specialized version of fromFoldableWith for lists.
fromListWith :: (Map c k a) => (a -> a -> a) -> [(k,a)] -> c
fromListWith = fromFoldableWith

data O a b c = L !a | R !b | O !c

-- | Same as 'intersectionWith', but with a more general type.
intersectionWith' :: (Functor m, Map (m (O a b c)) k (O a b c)) => 
                     (a->b->c) -> m a -> m b -> m c
intersectionWith' f m1 m2 = fmap extract (intersectionWith combine (fmap L m1) (fmap R m2))
    where combine (L l) (R r) = O (f l r)
          extract (O a) = a

-- | Same as 'differenceWith', but with a more general type.
differenceWith' :: (Functor m, Map (m (O a b c)) k (O a b c)) => 
                   (a->b->Maybe c) -> m a -> m b -> m c
differenceWith' f m1 m2 = fmap extract (differenceWith combine (fmap L m1) (fmap R m2))
    where combine (L l) (R r) = fmap O (f l r)
          extract (O a) = a

mapWithKey' :: (Functor m, Map (m (Either a b)) k (Either a b)) => 
              (k -> a -> b) -> m a -> m b
mapWithKey' f = fmap (either (error "mapWithKey': bug.") id) . mapWithKey f' . fmap Left
    where f' k (Left x) = Right (f k x)

-- | Class for set-like collection types. A set is really a map 
-- with no value associated to the keys,
-- so Set just states so.

-- Note that this should be a class alias, if it existed.
-- See: http://repetae.net/john/recent/out/classalias.html
class Map c k () => Set c k | c -> k where
    -- | Dummy method for haddock to accept the class.
    haddock_candy :: c -> k 

-- | Infix version of 'difference'. Difference of two (key) sets.
(\\) :: Map c k a => c -> c -> c
(\\) = difference

-- NOTE: the following two are only tentative, and thus not exported.

-- | Infix version of 'union'. Union of two (key) sets.
(\/) :: Map c k a => c -> c -> c
(\/) = union

-- | Infix version of 'intersection'. Intersection of two (key) sets.
(/\) :: Map c k a => c -> c -> c
(/\) = intersection


{-

Maybe it would be a good idea to bite the bullet and use a Lattice class for intersection and union.
Maybe leave it unrelated to the Map class. In a separate module/package? Something like:


class Lattice a where
    (/\) :: a -> a -> a
    (\/) :: a -> a -> a

instance Lattice () where
    _ /\ _ = ()
    _ \/ _ = ()

instance Lattice Bool where
    (/\) = (&&)
    (\/) = (||)

instance (Lattice a, Map c k a) => Lattice c where
    (/\) = intersectionWith (/\) 
    (\/) = unionWith (\/)   

-}

    

-- | Union of many (key) sets.
unions :: (Unfoldable s i, Map s k a, Foldable cs s) => cs -> s
unions sets = foldl' union empty sets

-- | Union of many (key) sets, with combining function
unionsWith :: (Unfoldable s i, Map s k a, Foldable cs s) => (a->a->a) -> cs -> s
unionsWith f sets = foldl' (unionWith f) empty sets

------------------------------------------------------------------------
-- Trickier stuff for alternate dictionnary usages

-- | "View" to the keys of a dictionnary
newtype KeysView m k v = KeysView {fromKeysView :: m}

-- | "View" to the elements of a dictionnary
newtype ElemsView m k v = ElemsView {fromElemsView :: m}

-- The following requires undecidable instances. An alternative
-- implementation is to define these instances directly on the
-- concrete map types and drop the requirement for the aforementioned
-- extension.

type T a = a->a

withKeys :: Collection m (k,v) => T (KeysView m k v) -> T m
withKeys f c = fromKeysView $ f (KeysView c)

withElems :: Collection m (k,v) => T (ElemsView m k v) -> T m
withElems f c = fromElemsView $ f (ElemsView c)

instance Foldable m (k,v) => Foldable (KeysView m k v) k where
    foldr f i (KeysView c) = foldr (f . fst) i c
    null (KeysView c) = null c

instance Unfoldable m (k,v) => Unfoldable (KeysView m k v) (k,v) where
    empty = KeysView empty
    insert x (KeysView m) = KeysView $ insert x m
    singleton x = KeysView (singleton x)

instance Foldable m (k,v) => Foldable (ElemsView m k v) v where
    foldr f i (ElemsView c) = foldr (f . snd) i c
    null (ElemsView c) = null c

instance Unfoldable m (k,v) => Unfoldable (ElemsView m k v) (k,v) where
    empty = ElemsView empty
    insert x (ElemsView m) = ElemsView $ insert x m
    singleton x = ElemsView (singleton x)

instance (Monoid m, Map m k v) => Monoid (KeysView m k v) where
    mempty = KeysView mempty
    mappend = union
 
instance Map m k v => Map (KeysView m k v) k v where
    isSubmapBy f (KeysView m) (KeysView m') = isSubmapBy f m m'
    isProperSubmapBy f (KeysView m) (KeysView m') = isProperSubmapBy f m m'
    member k (KeysView m) = Maybe.isJust $ lookup k m
    union (KeysView m) (KeysView m') = KeysView $ union m m'
    difference (KeysView m) (KeysView m') = KeysView $ difference m m'
    intersection (KeysView m) (KeysView m') = KeysView $ intersection m m'
    delete k (KeysView m) = KeysView $ delete k m
    insertWith f k a (KeysView m) = KeysView $ insertWith f k a m
    lookup k (KeysView m) = lookup k m
    alter f k (KeysView m) = KeysView $ alter f k m
    unionWith f (KeysView m) (KeysView m') = KeysView $ unionWith f m m'
    differenceWith f (KeysView m) (KeysView m') = KeysView $ differenceWith f m m'
    intersectionWith f (KeysView m) (KeysView m') = KeysView $ intersectionWith f m m'
    mapWithKey f (KeysView m) = KeysView $ mapWithKey f m