{-#OPTIONS -fglasgow-exts #-}
module GenMap where

{-
Self-optimizing finite-maps

derives heavily from Ralf Hinze's "Generics for the masses"
<http://www.informatik.uni-bonn.de/~ralf/publications.html#P21>
-}


import Prelude hiding (lookup)
import qualified Data.Map as Prim
import qualified Data.IntMap as Patricia
import Data.Char (chr, ord)

__ = undefined

data EP  a b = EP { to :: a -> b, from :: b -> a }
type FEP f g = forall a. EP (f a) (g a)

idEP = EP id id

newtype Comp f g a = C { unC :: f (g a) }
data    Pair f g a = f a :*: g a

fold (*) f g (x :*: y) = f x * g y
proj1 (x :*: y) = x            -- proj1  = fold const id id
proj2 (x :*: y) = y            -- proj2  = fold (const id) id id
map1 f (x :*: y) = f x :*: y   -- map1 f = fold (:*:) f id
map2 f (x :*: y) = x :*: f y   -- map2 f = fold (:*:) id f


class Map k m | k -> m, m -> k where
	func :: (MapFunc g) => g k m

class MapFunc g where
	unit     :: g () Maybe
	pair     :: (Map k1 m1, Map k2 m2) => g (k1,k2) (Comp m1 m2)
	plus     :: (Map k1 m1, Map k2 m2) => g (Either k1 k2) (Pair m1 m2)
	datatype :: (Map k' m') => EP k k' -> FEP m m' -> g k m
	prim     :: (Ord k) => g k (Prim.Map k)
	int      :: g Int Patricia.IntMap

instance Map () Maybe where func = unit
instance (Map k1 m1, Map k2 m2) => Map (k1,k2) (Comp m1 m2) where func = pair
instance (Map k1 m1, Map k2 m2) => Map (Either k1 k2) (Pair m1 m2) where func = plus

instance Map Int Patricia.IntMap where func = int
--instance Map Char (Prim.Map Char) where func = prim

newtype CharMap a = CharMap { unCharMap :: Patricia.IntMap a }

instance Map Char CharMap where
	func = datatype (EP ord chr) (EP unCharMap CharMap)

----

newtype MaybeMap m a = MaybeMap { unMaybeMap :: Pair Maybe m a }

instance (Map k m) => Map (Maybe k) (MaybeMap m) where
	func = datatype (EP to from) (EP unMaybeMap MaybeMap)
		where
			to   = maybe (Left ()) Right
			from = either (const Nothing) Just

--

newtype Trie m a = Trie { unTrie :: Pair Maybe (Comp m (Trie m)) a }

instance (Map k m) => Map [k] (Trie m) where
	func = datatype (EP to from) (EP unTrie Trie)
		where
			to []     = Left ()
			to (x:xs) = Right (x,xs)
			
			from (Left ())      = []
			from (Right (x,xs)) = x:xs

type Dict = Trie CharMap

--

data Name = Uri String | Lit String String | Dat String String

newtype NameMap a = NameMap { unNameMap :: Pair Dict (Pair (Comp Dict Dict) (Comp Dict Dict)) a }
-- more readable: Dict :*: Dict :@: Dict :*: Dict :@: Dict

-- can you do this?

instance Map Name NameMap where
	func = datatype (EP to from) (EP unNameMap NameMap)
		where
			to (Uri u)   = Left u
			to (Lit s l) = Right (Left (s,l))
			to (Dat u s) = Right (Right (u,s))
			
			from (Left u)              = Uri u
			from (Right (Left (s,l)))  = (Lit s l)
			from (Right (Right (u,s))) = Dat u s




data Perfect a = Zero a | Succ (Perfect (a,a))

-- PTrie m = m :*: PTrie (m :@: m)
newtype PTrie m a = PTrie { unPTrie :: Pair m (PTrie (Comp m m)) a }
instance Map k m => Map (Perfect k) (PTrie m) where
	func = datatype (EP to from) (EP unPTrie PTrie)
		where
			to (Zero k) = Left k
			to (Succ p) = Right p
			
			from = either Zero Succ



----

empty :: Map k m => m a
empty = runEmpty func

newtype Empty k m = Empty { runEmpty :: forall a. m a }

instance MapFunc Empty where
	unit = Empty (Nothing)
	plus = Empty (empty :*: empty)
	pair = Empty (C empty)
	
	datatype kep mep = Empty (from mep empty)
	
	prim = Empty (Prim.empty)
	int  = Empty (Patricia.empty)
	
--

singleton :: Map k m => k -> a -> m a
singleton = runSingle func

newtype Single k m = Single { runSingle :: forall a. k -> a -> m a }

instance MapFunc Single where
	unit = Single (const Just)

	plus = Single (\k x -> case k of
		Left k1  -> singleton k1 x :*: empty
		Right k2 -> empty :*: singleton k2 x)

	pair = Single (\(k1,k2) -> C . singleton k1 . singleton k2)
	
	datatype kep mep = Single (\k -> from mep . singleton (to kep k))
	
	prim = Single Prim.singleton
	int  = Single Patricia.singleton

--

{-
fromList :: Map k m => (a -> a -> a) -> [(k,a)] -> m a
fromList = runFromList func

newtype FromList k m = FromList { runFromList :: forall a. (a -> a -> a) -> [(k,a)] -> m a }

instance MapFunc FromList where
	unit = FromList (\c xs -> case xs of
		[]       -> Nothing
		(_,x):xs -> Just (foldr (c . snd) x xs))
	
	plus 
-}

fromList c xs = foldr (\(k,x) -> insertWith c k x) empty xs

--

insertWith :: Map k m => (a -> a -> a) -> k -> a -> m a -> m a
insertWith = runInsertWith func

newtype InsertWith k m = InsertWith
	{ runInsertWith :: forall a. (a -> a -> a) -> k -> a -> m a -> m a }

instance MapFunc InsertWith where
	unit = InsertWith (\c () x -> Just . maybe x (c x))
	plus = InsertWith (\c k x -> case k of
		Left k1  -> map1 (insertWith c k1 x)
		Right k2 -> map2 (insertWith c k2 x))
	
	pair = InsertWith (\c (k1,k2) x (C m1) ->
		C (insertWith (\_ -> insertWith c k2 x) k1 (singleton k2 x) m1))
	
	datatype kep mep = InsertWith (\c k x -> from mep . insertWith c (to kep k) x . to mep)
	
	prim = InsertWith Prim.insertWith
	int  = InsertWith Patricia.insertWith

--

lookup :: Map k m => k -> m a -> Maybe a
lookup = runLookup func

newtype Lookup k m = Lookup { runLookup :: forall a. k -> m a -> Maybe a }

instance MapFunc Lookup where
	unit = Lookup (\() -> id)
	plus = Lookup (either (\k1 -> lookup k1 . proj1) (\k2 -> lookup k2 . proj2))
	pair = Lookup (\(k1,k2) (C m) -> lookup k1 m >>= lookup k2)
	
	datatype kep mep = Lookup (\k -> lookup (to kep k) . to mep)
	
	prim = Lookup Prim.lookup
	int  = Lookup Patricia.lookup




--

isEmpty :: Map k m => m a -> Bool
isEmpty = runIsEmpty func

newtype IsEmpty k m = IsEmpty { runIsEmpty :: forall a. m a -> Bool }

instance MapFunc IsEmpty where
	unit = IsEmpty (maybe False (const True))
	plus = IsEmpty (fold (&&) isEmpty isEmpty)
	pair = IsEmpty (isEmpty . unC)
	
	datatype (kep::EP k k') mep = IsEmpty (isEmpty . to mep)
	
	prim = IsEmpty (Prim.null)
	int  = IsEmpty (Patricia.null)

--

update :: Map k m => (a -> Maybe a) -> k -> m a -> m a
update = runUpdate func

newtype Update k m = Update { runUpdate :: forall a. (a -> Maybe a) -> k ->
m a -> m a }

instance MapFunc Update where
	unit = Update (\f () m -> m >>= f)
	plus = Update (\f -> either (map1 . update f) (map2 . update f))
	pair = Update (\f (k1,k2) -> C . update (justNot isEmpty . update f k2) k1 . unC)

	datatype kep mep = Update (\f k -> from mep . update f (to kep k) . to mep)
	
	prim = Update Prim.update
	int  = Update Patricia.update

justNot :: (a -> Bool) -> a -> Maybe a
justNot p x = if p x then Nothing else Just x

--

delete :: Map k m => k -> m a -> m a
delete = runDelete func

newtype Delete k m = Delete { runDelete :: forall a. k -> m a -> m a }

instance MapFunc Delete where
	unit = Delete (\_ _ -> Nothing)
	plus = Delete (either (map1 . delete) (map2 . delete))
	pair = Delete (\(k1,k2) -> C . update (justNot isEmpty . delete k2) k1 . unC)
	
	datatype kep mep = Delete (\k -> from mep . delete (to kep k) . to mep)
	
	prim = Delete Prim.delete
	int  = Delete Patricia.delete

--

toList :: Map k m => m a -> [(k,a)]
toList = runToList func

newtype ToList k m = ToList { runToList :: forall a. m a -> [(k,a)] }

instance MapFunc ToList where
	unit = ToList (maybe [] (\x -> [((),x)]))
	plus = ToList (fold (++) (mapFst Left . toList) (mapFst Right . toList))
		where
			
	pair = ToList (\(C m) -> do
		(k,m') <- toList m
		mapFst ((,) k) (toList m'))
	
	datatype kep mep = ToList (mapFst (from kep) . toList . to mep)
	
	prim = ToList Prim.toList
	int  = ToList Patricia.toList

mapFst f = map (\(k,x) -> (f k, x))

--

elems :: Map k m => m a -> [a]
elems = runElems func

newtype Elems k m = Elems { runElems :: forall a. m a -> [a] }

instance MapFunc Elems where
	unit = Elems (maybe [] (:[]))
	plus = Elems (fold (++) elems elems)
	pair = Elems (\(C m) -> elems m >>= elems)
	
	datatype kep mep = Elems (elems . to mep)
	
	prim = Elems (Prim.elems)
	int  = Elems (Patricia.elems)

--

keys :: Map k m => m a -> [k]
keys = runKeys func

newtype Keys k m = Keys { runKeys :: forall a. m a -> [k] }

instance MapFunc Keys where
	unit = Keys (maybe [] (const [()]))
	plus = Keys (fold (++) (map Left . keys) (map Right . keys))
	pair = Keys (\(C m1) -> do
		(k1,m2) <- toList m1
		k2 <- keys m2
		return (k1,k2))
	
	datatype kep mep = Keys (map (from kep) . keys . to mep)
	prim = Keys Prim.keys
	int  = Keys Patricia.keys
	
-- 

size :: Map k m => m a -> Int
size = runSize func

newtype Size k m = Size { runSize :: forall a. m a -> Int }

instance MapFunc Size where
	unit = Size (maybe 0 (const 1))
	plus = Size (fold (+) size size)
	pair = Size (sum . map size . elems . unC)
	
	datatype kep mep = Size (size . to mep)
	prim = Size (Prim.size)
	int  = Size (Patricia.size)