{-# LANGUAGE TupleSections, TypeFamilies, UnboxedTuples, MagicHash #-}

module Data.TrieMap.TrieKey where

import Data.TrieMap.Sized

import Control.Applicative
import Control.Monad

import Data.Monoid
import Data.Foldable

import Prelude hiding (foldr, foldl)


import GHC.Exts

type EitherMap k a b c = k -> a -> (# Maybe b, Maybe c #)
type SplitMap a x = a -> (# Maybe a, Maybe x, Maybe a #)
type UnionFunc k a = k -> a -> a -> Maybe a
type IsectFunc k a b c = k -> a -> b -> Maybe c
type DiffFunc k a b = k -> a -> b -> Maybe a
type LEq a b = a -> b -> Bool

onUnboxed :: (c -> d) -> (a -> (# b, c #)) -> a -> (# b, d #)
onUnboxed g f a = case f a of
	(# b, c #) -> (# b, g c #)

instance TrieKey k => Foldable (TrieMap k) where
	foldr f z m = foldrWithKeyM (const f) m z
	foldl f z m = foldlWithKeyM (const f) m z

class Ord k => TrieKey k where
	data TrieMap k :: * -> *
	emptyM :: TrieMap k a
	singletonM :: Sized a => k -> a -> TrieMap k a
	nullM :: TrieMap k a -> Bool
	sizeM :: Sized a => TrieMap k a -> Int#
	lookupM :: k -> TrieMap k a -> Maybe a
	mapWithKeyM :: Sized b => (k -> a -> b) -> TrieMap k a -> TrieMap k b
	traverseWithKeyM :: (Applicative f, Sized b) =>
		(k -> a -> f b) -> TrieMap k a -> f (TrieMap k b)
	foldrWithKeyM :: (k -> a -> b -> b) -> TrieMap k a -> b -> b
	foldlWithKeyM :: (k -> b -> a -> b) -> TrieMap k a -> b -> b
	mapMaybeM :: Sized b => (k -> a -> Maybe b) -> TrieMap k a -> TrieMap k b
	mapEitherM :: (Sized b, Sized c) => EitherMap k a b c -> TrieMap k a -> (# TrieMap k b, TrieMap k c #)
	unionM :: Sized a => UnionFunc k a -> TrieMap k a -> TrieMap k a -> TrieMap k a
	isectM :: (Sized a, Sized b, Sized c) => IsectFunc k a b c -> TrieMap k a -> TrieMap k b -> TrieMap k c
	diffM :: Sized a => DiffFunc k a b -> TrieMap k a -> TrieMap k b -> TrieMap k a
	isSubmapM :: (Sized a, Sized b) => LEq a b -> LEq (TrieMap k a) (TrieMap k b)
	fromListM, fromAscListM :: Sized a => (k -> a -> a -> a) -> [(k, a)] -> TrieMap k a
	fromDistAscListM :: Sized a => [(k, a)] -> TrieMap k a
	
	data Hole k :: * -> *
	singleHoleM :: k -> Hole k a
	keyM :: Hole k a -> k
	beforeM :: Sized a => Maybe a -> Hole k a -> TrieMap k a
	afterM :: Sized a => Maybe a -> Hole k a -> TrieMap k a
	searchM :: k -> TrieMap k a -> (# Maybe a, Hole k a #)
	indexM :: Sized a => Int# -> TrieMap k a -> (# Int#, a, Hole k a #)
	{-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> First (a, Hole k a) #-}
	{-# SPECIALIZE extractHoleM :: Sized a => TrieMap k a -> Last (a, Hole k a) #-}
	extractHoleM :: MonadPlus m => Sized a => TrieMap k a -> m (a, Hole k a)
	assignM :: Sized a => a -> Hole k a -> TrieMap k a
	clearM :: Sized a => Hole k a -> TrieMap k a

	singletonM k a = assignM a (singleHoleM k)
	lookupM k m = case searchM k m of
		(# a, _ #)	-> a
	foldrWithKeyM f = appEndo . getConst . traverseWithKeyM (endofy f) where
		endofy :: (k -> a -> b -> b) -> k -> a -> Const (Endo b) (Elem ())
		endofy f k a = Const (Endo (f k a))
	foldlWithKeyM f m = foldrWithKeyM (\ k a g z -> g (f k z a)) m id
	fromListM f = foldr (uncurry (insertWithKeyM f)) emptyM
	fromAscListM = fromListM
	fromDistAscListM = fromAscListM (const const)

instance (TrieKey k, Sized a) => Sized (TrieMap k a) where
	getSize# = sizeM

{-# INLINE alterM #-}
alterM :: (TrieKey k, Sized a) => (Maybe a -> Maybe a) -> k -> TrieMap k a -> TrieMap k a
alterM f k m = case searchM k m of
	(# Nothing, hole #)	-> maybe m (\ a -> assignM a hole) (f Nothing)
	(# a, hole #)		-> fillHoleM (f a) hole

traverseM :: (Applicative f, TrieKey k, Sized b) => (a -> f b) -> TrieMap k a -> f (TrieMap k b)
traverseM f = traverseWithKeyM (const f)

guardNullM :: TrieKey k => TrieMap k a -> Maybe (TrieMap k a)
guardNullM m
	| nullM m	= Nothing
	| otherwise	= Just m

fillHoleM :: (TrieKey k, Sized a) => Maybe a -> Hole k a -> TrieMap k a
fillHoleM Nothing hole = clearM hole
fillHoleM (Just a) hole = assignM a hole

sides :: (b -> d) -> (a -> (# b, c, b #)) -> a -> (# d, c, d #)
sides g f a = case f a of
	(# x, y, z #) -> (# g x, y, g z #)

both :: (b -> b') -> (c -> c') -> (a -> (# b, c #)) -> a -> (# b', c' #)
both g1 g2 f a = case f a of
	(# x, y #) -> (# g1 x, g2 y #)

fmapM :: (TrieKey k, Sized b) => (a -> b) -> TrieMap k a -> TrieMap k b
fmapM = mapWithKeyM . const

assocsM :: TrieKey k => TrieMap k a -> [(k, a)]
assocsM m = build (\ f z -> foldrWithKeyM (\ k a xs -> (k, a) `f` xs) m z)

insertWithKeyM :: (TrieKey k, Sized a) => (k -> a -> a -> a) -> k -> a -> TrieMap k a -> TrieMap k a
insertWithKeyM f k a m = case searchM k m of
	(# Nothing, hole #)	-> assignM a hole
	(# Just a', hole #)	-> assignM (f k a a') hole

unionMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
unionMaybe _ Nothing y = y
unionMaybe _ x Nothing = x
unionMaybe f (Just x) (Just y) = f x y

isectMaybe :: (a -> b -> Maybe c) -> Maybe a -> Maybe b -> Maybe c
isectMaybe f (Just x) (Just y) = f x y
isectMaybe _ _ _ = Nothing

diffMaybe :: (a -> b -> Maybe a) -> Maybe a -> Maybe b -> Maybe a
diffMaybe _ Nothing _ = Nothing
diffMaybe _ (Just x) Nothing = Just x
diffMaybe f (Just x) (Just y) = f x y

subMaybe :: (a -> b -> Bool) -> Maybe a -> Maybe b -> Bool
subMaybe _ Nothing _ = True
subMaybe (<=) (Just a) (Just b) = a <= b
subMaybe _ _ _ = False