{-# LANGUAGE TypeFamilies, UnboxedTuples #-}

module Data.TrieMap.Key (Key(..)) where

import Control.Applicative
import Data.TrieMap.Class
import Data.TrieMap.TrieKey
import Data.TrieMap.Rep
import Data.TrieMap.Modifiers

instance TKey k => TrieKey (Key k) where
	newtype TrieMap (Key k) a = KeyMap (TrieMap (Rep k) a)
	newtype Hole (Key k) a = KeyHole (Hole (Rep k) a)
	
	emptyM = KeyMap emptyM
	singletonM (Key k) a = KeyMap (singletonM (toRep k) a)
	nullM (KeyMap m) = nullM m
	sizeM (KeyMap m) = sizeM m
	lookupM (Key k) (KeyMap m) = lookupM (toRep k) m
	traverseWithKeyM f (KeyMap m) = KeyMap <$> traverseWithKeyM (f . Key . fromRep) m
	foldrWithKeyM f (KeyMap m) = foldrWithKeyM (f . Key . fromRep) m
	foldlWithKeyM f (KeyMap m) = foldlWithKeyM (f . Key . fromRep) m
	mapWithKeyM f (KeyMap m) = KeyMap (mapWithKeyM (f . Key . fromRep) m)
	mapMaybeM f (KeyMap m) = KeyMap (mapMaybeM (f . Key . fromRep) m)
	mapEitherM f (KeyMap m) = both KeyMap KeyMap (mapEitherM (f . Key . fromRep)) m
	unionM f (KeyMap m1) (KeyMap m2) = KeyMap (unionM (f . Key . fromRep) m1 m2)
	isectM f (KeyMap m1) (KeyMap m2) = KeyMap (isectM (f . Key . fromRep) m1 m2)
	diffM f (KeyMap m1) (KeyMap m2) = KeyMap (diffM (f . Key . fromRep) m1 m2)
	isSubmapM (<=) (KeyMap m1) (KeyMap m2) = isSubmapM (<=) m1 m2

	singleHoleM (Key k) = KeyHole (singleHoleM (toRep k))
	keyM (KeyHole hole) = Key (fromRep (keyM hole))
	beforeM a (KeyHole hole) = KeyMap (beforeM a hole)
	afterM a (KeyHole hole) = KeyMap (afterM a hole)
	searchM (Key k) (KeyMap m) = onUnboxed KeyHole (searchM (toRep k)) m
	indexM i (KeyMap m) = case indexM i m of
		(# i', v, hole #) -> (# i', v, KeyHole hole #)
	extractHoleM (KeyMap m) = do
		(v, hole) <- extractHoleM m
		return (v, KeyHole hole)
	assignM v (KeyHole hole) = KeyMap (assignM v hole)
	clearM (KeyHole hole) = KeyMap (clearM hole)