{-# OPTIONS_GHC -fglasgow-exts -fallow-undecidable-instances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Collections
-- Copyright   :  (c) Jean-Philippe Bernardy 2006
-- License     :  BSD-style
-- Maintainer  :  jeanphilippe.bernardy; google mail.
-- Stability   :  experimental
-- Portability :  MPTC, FD, undecidable instances
--
-- 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. 'AssocView')
--
-- * A few generic functions for handling collections.
--
-- 
--
-- The classes defined in this module are intended to give hints about performance.
-- eg. if a function has a @MapLike c k v@ context, this indicates that the function
-- will perform better if @c@ has an efficitent lookup function.
--
-- 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@.

{-


TODO:
 * write instances for the new Seq type, following List.[]
-- fix union comment. Better semantics generally.

-- foldr/l in sequence

 * See how the new Foldable class superseeds any of this. Remove stuff as needed.

-}

module Data.Collections 
    (

-- * Classes
     Collection(..),
     Indexed(..),
     MapLike(..),
     SetLike,
     Sequence(..),

-- * Extra generic functions
     findWithDefault,
     unions,

-- ** Aliases
     (\\),

-- ** Unfolding
     unfold,

-- * Conversions
     convert,
     toList,
     fromList,

-- * Views
     AssocView(..),
     KeysView(..), ElemsView(..),
             
     Void
    ) where 

import Prelude hiding (sum,concat,lookup,map,filter,foldr,foldl,null,reverse,(++))

import Control.Monad
import qualified Data.Map as Map
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Array as Array
import qualified Data.Maybe as Maybe

infixl 9 !,\\ --

-- | Type with no value; ideally it should be a strict type.
data Void
noVoidValue = error "Don't use bottom to populate the Nothing type."

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

-- | Class of collection types.
--
-- * 'i' values are inserted into the collection.
--
-- * 'o' values are extracted out of the collection.
-- 
-- Having two extra parameters allows for:
--
--  * unobservable collections when 'o' = @()@
-- 
--  * \"readonly\" collections when 'i' = 'Void'
--
--  * Views over only some projection of the element (see 'KeysView' and 'ElemsView')
--
-- Also, please note that:
--
--  * There is no notion of order in this class. ('fold', 'toList', etc. provide specific order no guarantee)
--
--  * neither 'map' nor 'fmap' is in here, use Functor for that purpose.
--
--  * @extract :: c -> Maybe (o,c)@ to take a random element is not there either. 
--- Use 'front', possibly converting to 'Data.List' if needed. (you don't know if the collection implements a fast linear access)

class Collection c i o | c -> i o where
    -- | The empty collection.
    empty :: c                            
    -- | Tells whether the collection contains a single element.
    isSingleton :: c -> Bool              
    -- | 'filter', applied to a predicate and a list, returns the collection of those elements that satisfy the predicate.
    filter :: (o -> Bool) -> c -> c       
    -- | \'natural\' traversal of all elements of a collection. No particular order is guaranteed.
    fold :: (o -> b -> b) -> b -> c -> b  
    -- | \'natural\' insertion into the collection.
    fold' :: (o -> b -> b) -> b -> c -> b  
    -- | \'natural\' insertion into the collection, in a strict fashion
    insert :: i -> c -> c                 
    -- | Tells whether the collection is empty
    null :: c -> Bool                    
    -- | Creates a collection with a single element.
    singleton :: i -> c 
    -- | Returns the size of the collection
    size :: c -> Int   
    
    isSingleton = (1 ==) . size           
    singleton i = insert i empty
    size = fold (const (+1)) 0

unfold :: (Collection c a a) => (b -> Maybe (a, b)) -> b -> c
unfold f s = convert $ List.unfoldr f s
-- in the above List.unfoldr should be deforested away.

-- | Conversion between two collection types.
convert :: (Collection c i o, Collection c' o o) => c -> c'
convert = fold insert empty

-- | Converts a collection into a list.
toList :: Collection c i o => c -> [o]
toList = convert

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


-- | Class of sequential-access types.
class Collection c i o => Sequence c i o where
    foldl :: (b -> o -> b) -> b -> c -> b  
    take :: Int -> c -> c         
    drop :: Int -> c -> c
    splitAt :: Int -> c -> (c,c)
    reverse :: c -> c
    front :: Monad m => c -> m (o,c)
    back :: Monad m => c -> m (c,o)
    (<|) :: i -> c -> c
    (|>) :: c -> i -> c
    (><) :: c -> c -> c

foldr :: Sequence c i o => (o -> b -> b) -> b -> c -> b  
foldr = fold


-- | 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 sparse poplutated indexed collection will instanciate this
-- class, and leave the responsibility of failure to the caller.
class Indexed c k v | c -> k v where
    -- | @c!k@ returns element associated to 'k'
    (!) :: c -> k -> v                 
    -- | @adjust f k c@ applies 'f' to element associated to 'k'
    adjust :: (v -> v) -> k -> c -> c 


-- TODO: bounds as in the class array would be a nice addition. However, this does not fit well with Map being an instance of Indexed.
-- Have a separate class for that ?


-- | Class of map-like types. (aka. for sparse associative types).
--
-- In opposition of Indexed, MapLike supports unexisting value for some indices.

class MapLike c k a | c -> k a where
    -- | Remove an element from the keySet.
    delete :: k -> c -> c
    delete = update (const Nothing)

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

    -- | Union of two keySets.
    -- When duplicates are encountered, the elements 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

    -- | 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 (\x y -> Nothing)

    -- | Intersection of two keySets.
    --
    -- When duplicates are encountered, the elements may come from any of the two input sets.
    -- Intersection is commutative: @intersection a b == intersection b a@
    intersection :: c -> c -> c
    intersection = intersectionWith const

-- Follows functions for fully-fledged maps.

                                                               
    -- | 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 = update (\x -> Just $ case x of {Nothing->a;Just a' -> f a a'}) k 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

    -- NOTE: there's an infelicity here because Map difference has type: 
    -- Map k a -> Map k b -> Map k a -- (same infelicity for intersection)

    -- | Lookup the value at a given key.
    lookup :: Monad m => k -> c -> m a

    -- | Change the value at a given key. Nothing represents no associated value.
    update :: (Maybe a -> Maybe a) -> k -> c -> c

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



-- | Class for set-like collection types. A set is really a map with no value associated to the keys,
-- so SetLike just states so.
-- 
-- Note that this should be a context alias or something.
class MapLike c k () => SetLike c k where
    -- | Dummy method for haddock to accept the class.
    haddock_candy :: c -> k 

-- | Difference of two (key) sets.
(\\) :: MapLike c k a => c -> c -> c
(\\) = difference


-- | Union of many (key) sets.
unions :: (Collection s i o, MapLike s k a, Collection cs i' s) => cs -> s
unions sets = fold union empty sets

-- NOTE: Should be specialized (RULE pragma) so it's not horribly inefficient in the common cases


-----------------------------------------------------------------------------
-- Instances
-----------------------------------------------------------------------------


-- We follow with (sample) instances of the classes.

-----------------------------------------------------------------------------
-- Data.List

instance Collection [a] a a where
    null = List.null
    fold = List.foldr
    fold' f = List.foldl' (flip f)
    empty = []
    singleton = return
    insert = (:)
    filter = List.filter

instance Sequence [a] a a where
    foldl = List.foldl
    take = List.take
    drop = List.drop
    splitAt = List.splitAt
    reverse = List.reverse
    front (x:xs) = return (x,xs)
    front [] = fail "front: empty sequence"
    back s = return swap `ap` front (reverse s)
        where swap (a,b) = (b,a)
    (<|) = (:)
    xs |> x = xs List.++ [x]
    (><) = (List.++)

(++) s1 s2 = (><) s1 s2
-- Deprecate ?
    
-- For convenience, List is made and instance of Indexed.
instance Indexed [a] Int a where
    (!) = (List.!!)
    adjust f k l = l >< (f x:r)
        where (l,x:r) = List.splitAt (k-1) l


-- For "compatibility" with the Prelude, List is made and instance of SetLike.
-- This however conflicts with the below above declaration: Indexed [a] Int a.
-- Note: I wonder how ghc can accept this.

instance Eq a => SetLike [a] a where
    haddock_candy = haddock_candy

instance Eq a => MapLike [a] a () where
    difference = (List.\\)
    delete = List.delete
    member = List.elem
    union = List.union
    intersection = List.intersect
    insertWith f k () = insert k
    unionWith f = union
    intersectionWith f = intersection
    differenceWith f = difference
    lookup k l = if member k l then return () else fail "element not found"
    update f k l = let lk = lookup k l in
        case lk of
           Nothing -> case lk of
                         Nothing -> l
                         Just _ -> insert k l
           Just _ -> case lk of
                         Nothing -> delete k l
                         Just _ -> l
    
-- | View a list of @(key,value)@ pairs as a 'MapLike' collection.
--
-- This allows to feed sequences into algorithms that require a map without building a full-fledged map.
-- Most of the time this will be used only when the parameter list is known to be very small, such that
-- conversion to a Map would be to costly.

newtype AssocView s k v = AssocView {fromAssocView :: s} -- k and v parameters will become useful if we generalize to sequences.

association :: [(k,v)] -> AssocView [(k,v)] k v
association = AssocView

instance Collection (AssocView [(k,v)] k v) (k,v) (k,v) where
    empty = AssocView []
    fold f i (AssocView l) = fold f i l
    fold' f i (AssocView l) = fold' f i l
    null (AssocView l) = null l
    filter f (AssocView l) = AssocView $ filter f l
    insert x (AssocView l) = AssocView $ insert x l
    
instance Eq k => Indexed (AssocView [(k,v)] k v) k v where
    (AssocView c) ! k = Maybe.fromJust (List.lookup k c)
    adjust f k (AssocView c) = AssocView $ List.map (\a@(k',v) -> if k == k' then (k, f v) else a) c

instance Eq k => MapLike (AssocView [(k,v)] k v) k v where
    delete k c = update (const Nothing) k c
    member k c = Maybe.isJust (lookup k c)
    union = unionWith const
    intersection = intersectionWith const
    difference = differenceWith (\x y->Nothing)

    lookup k (AssocView l) = if List.null result then fail "Key not found" else return . snd . head $ result
        where result = [x | x <- l, fst x == k]
    insertWith f k a c = 
        case lookup k c of
           Nothing -> insert (k,a) c
           Just b -> insert (k, f a b) (delete k c)
    intersectionWith f (AssocView m1) (AssocView m2) = AssocView [(k,f x y) 
                                                                        | (k,x) <- m1, 
                                                                           y <- Maybe.maybeToList $ List.lookup k m2]
    unionWith f (AssocView m1) (AssocView m2) = AssocView $ List.map unionOne $ List.groupBy (testing fst) $ m1 >< m2
        where unionOne list = (fst (head list), foldr1 f (List.map snd list))
    differenceWith f (AssocView m1) (AssocView m2) = AssocView $ Maybe.catMaybes 
                                                         [newEl k x (List.lookup k m2) | (k,x) <- m1]
        where newEl k x Nothing = Just (k,x)
              newEl k x (Just y) = fmap (\x->(k,x)) (f x y)
    update f k (AssocView m) = AssocView $ case f $ fmap snd $ Maybe.listToMaybe eq of
                                              Nothing -> neq
                                              Just x -> (k,x):neq
        where (eq,neq) = List.partition (\x->fst x == k) m

testing :: Eq b => (a -> b) -> a -> a -> Bool
testing f x y = (==) (f x) (f y)


--------------------------------------
-- Data.Array

instance Array.Ix i => Collection (Array.Array i e) Void (i,e) where    
    fold f i c = List.foldr f i (Array.assocs c)
    fold' f i c = List.foldl' (flip f) i (Array.assocs c)
    insert = noVoidValue
    filter = noVoidValue
    empty = noVoidValue
    null c = null $ Array.range $ Array.bounds c

instance Array.Ix i => Indexed (Array.Array i e) i e where
    (!) = (Array.!)
    adjust f k a = a Array.// [(k,f (a!k))]

-----------------------------------------------------------------------------
-- Data.Map
instance Ord k => Collection (Map.Map k a) (k,a) (k,a) where
    filter f = Map.filterWithKey (curry f)
    insert = uncurry Map.insert
    null = Map.null
    singleton (k,a) = Map.singleton k a
    fold f i m = Map.foldWithKey (curry f) i m
    empty = Map.empty

instance Ord k => Indexed (Map.Map k a) k a where
    (!) = (Map.!)
    adjust = Map.adjust

instance Ord k => MapLike (Map.Map k a) k a where
    member = Map.member
    union = Map.union
    difference = Map.difference
    delete = Map.delete
    intersection = Map.intersection
    lookup = Map.lookup
    update f k m = case f (lookup k m) of
                      Just a -> Map.insert k a m
                      Nothing -> Map.delete k m
             -- TODO: add support for this in Data.Map
    insertWith = Map.insertWith
    unionWith = Map.unionWith
    intersectionWith = Map.intersectionWith
    differenceWith = Map.differenceWith

-----------------------------------------------------------------------------
-- Data.Set

instance Ord a => Collection (Set.Set a) a a where
    filter = Set.filter
    insert = Set.insert
    null = Set.null
    singleton = Set.singleton
    fold f i s = Set.fold f i s
    empty = Set.empty

instance Ord a => SetLike (Set.Set a) a where
    haddock_candy = haddock_candy

instance Ord a => MapLike (Set.Set a) a () where
    member = Set.member
    union = Set.union
    difference = Set.difference
    intersection = Set.intersection
    delete = Set.delete
    insertWith f k () = insert k
    unionWith f = union
    intersectionWith f = intersection
    differenceWith f = difference
    lookup k l = if member k l then return () else fail "element not found"    
    update f k m = case f (lookup k m) of
                      Just a -> insert k m
                      Nothing -> delete k m

------------------------------------------------------------------------
-- 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.

instance Collection m (k,v) (k,v) => Collection (KeysView m k v) (k,v) k where
    empty = KeysView empty
    filter f (KeysView m) = KeysView $ filter (f . fst) m
    fold f i (KeysView c) = fold (f . fst) i c
    fold' f i (KeysView c) = fold' (f . fst) i c
    insert x (KeysView m) = KeysView $ insert x m
    null (KeysView c) = null c
    singleton x = KeysView (singleton x)
    
instance Collection m (k,v) (k,v) => Collection (ElemsView m k v) (k,v) v where
    empty = ElemsView empty
    filter f (ElemsView m) = ElemsView $ filter (f . snd) m
    fold f i (ElemsView c) = fold (f . snd) i c
    fold' f i (ElemsView c) = fold' (f . snd) i c
    insert x (ElemsView m) = ElemsView $ insert x m
    null (ElemsView c) = null c
    singleton x = ElemsView (singleton x)
 
instance MapLike m k v => MapLike (KeysView m k v) k v where
    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
    update f k (KeysView m) = KeysView $ update 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'




-----------------------------
-- examples of use/test code

sum c = fold (+) 0 c

concat c = fold (><) [] c

origList = [("one", 1), ("two", 2)]

someMap :: Map.Map String Int
someMap = convert origList
 
test1 = sum $ ElemsView someMap
test1a = sum $ ElemsView origList

test2 = concat $ KeysView someMap
test2a = concat $ KeysView someMap

test3 = someMap ! "one"
test3a :: Int
test3a = association origList ! "one"