{-# LANGUAGE TypeFamilies, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses, FlexibleContexts #-}

module Data.TrieMap.Key where

import Control.Applicative
import Control.Arrow ((***))
import Data.TrieMap.Class
import Data.TrieMap.TrieKey
import Data.TrieMap.Rep

newtype Key k = Key {getKey :: k}
newtype KeyMap k a = KeyMap {getKeyMap :: TrieMap (Rep k) a}

instance (TKey k) => Eq (Key k) where
	Key k1 == Key k2 = toRep k1 == toRep k2

instance (TKey k) => Ord (Key k) where
	Key k1 `compare` Key k2 = compare (toRep k1) (toRep k2)

type instance TrieMap (Key k) = KeyMap k

instance TKey k => TrieKey (Key k) (KeyMap k) where
	emptyM = KeyMap emptyM
	nullM (KeyMap m) = nullM m
	lookupM (Key k) (KeyMap m) = lookupM (toRep k) m
	lookupIxM s (Key k) (KeyMap m) = onKey (Key . fromRep) (lookupIxM s (toRep k) m)
	assocAtM s i (KeyMap m) = onKey (Key . fromRep) (assocAtM s i m)
	alterM s f (Key k) (KeyMap m) = KeyMap (alterM s f (toRep k) m)
	alterLookupM s f (Key k) (KeyMap m) = KeyMap <$> alterLookupM s f (toRep k) m
	traverseWithKeyM s f (KeyMap m) = KeyMap <$> traverseWithKeyM s (f . Key . fromRep) m
	foldWithKeyM f (KeyMap m) = foldWithKeyM (f . Key . fromRep) m
	foldlWithKeyM f (KeyMap m) = foldlWithKeyM (f . Key . fromRep) m
	mapEitherM s1 s2 f (KeyMap m) = (KeyMap *** KeyMap) (mapEitherM s1 s2 (f . Key . fromRep) m)
	splitLookupM s f (Key k) (KeyMap m) = KeyMap `sides` splitLookupM s f (toRep k) m
	unionM s f (KeyMap m1) (KeyMap m2) = KeyMap (unionM s (f . Key . fromRep) m1 m2)
	isectM s f (KeyMap m1) (KeyMap m2) = KeyMap (isectM s (f . Key . fromRep) m1 m2)
	diffM s f (KeyMap m1) (KeyMap m2) = KeyMap (diffM s (f . Key . fromRep) m1 m2)
	extractM s f (KeyMap m) = fmap KeyMap <$> extractM s (f . Key . fromRep) m
	isSubmapM (<=) (KeyMap m1) (KeyMap m2) = isSubmapM (<=) m1 m2