{-# LANGUAGE RankNTypes, FlexibleContexts, UndecidableInstances, TypeFamilies, TypeOperators, TemplateHaskell, NPlusKPatterns #-}
{-# OPTIONS -funbox-strict-fields #-}

module Data.TrieMap.Rep.Instances()  where

import Data.TrieMap.Rep
import Data.TrieMap.Rep.TH
import Data.TrieMap.Modifiers

import Data.Char
import Data.Int
import Data.Word
import Data.Foldable (toList)
import Data.Bits
import qualified Data.IntSet as ISet
import qualified Data.IntMap as IMap
import Data.ByteString hiding (map)
import qualified Data.ByteString as BS

import Data.Sequence ((|>))
import qualified Data.Sequence as Seq
import qualified Data.Foldable as Fold

import qualified Data.Map as Map
import qualified Data.Set as Set

import Prelude hiding (concat, take, length)

type Pair a = (,) a
type Sum a = Either a

instance ReprT Rev where
  type RepT Rev = Rev
  toRepTMap = fmap
  fromRepTMap = fmap

genRepr [t| Rev |]

instance ReprT [] where
	type RepT [] = []
	toRepTMap = map
	fromRepTMap = map

genRepr [t| [] |]

genTupleRepr 2
genTupleRepr 3
genTupleRepr 4
genTupleRepr 5
genTupleRepr 6
genTupleRepr 7
genTupleRepr 8

instance (Repr a, Repr b) => Repr (Either a b) where
  type Rep (Either a b) = Either (Rep a) (Rep b)
  toRep (Left a) = Left (toRep a)
  toRep (Right b) = Right (toRep b)
  fromRep (Left a) = Left (fromRep a)
  fromRep (Right b) = Right (fromRep b)

instance Repr Char where
	type Rep Char = Word32
	toRep = fromIntegral . ord
	fromRep = chr . fromIntegral

instance Repr () where
	type Rep () = ()
	toRep _ = ()
	fromRep _ = ()

instance Repr Int where
	type Rep Int = Rep Int32
	toRep = toSigned
	fromRep = fromSigned

instance Repr Word8 where
	type Rep Word8 = Word32
	toRep = fromIntegral
	fromRep = fromIntegral

instance Repr Word16 where
	type Rep Word16 = Word32
	toRep = fromIntegral
	fromRep = fromIntegral

instance Repr Word where
	type Rep Word = Word32
	toRep = fromIntegral
	fromRep = fromIntegral

instance Repr Int8 where
	type Rep Int8 = Rep Int32
	toRep = toSigned
	fromRep = fromSigned

instance Repr Int16 where
	type Rep Int16 = Rep Int32
	toRep = toSigned
	fromRep = fromSigned

instance Repr Int32 where
	type Rep Int32 = Sum (Rev Word32) Word32
	toRep = toSigned
	fromRep = fromSigned

instance Repr Word64 where
	type Rep Word64 = Pair Word32 Word32
	toRep x = (fromIntegral (x `shiftR` 32), fromIntegral x)
	fromRep (x, y) = fromIntegral x `shiftL` 32 .|. fromIntegral y

instance Repr Int64 where
	type Rep Int64 = Sum (Rev (Rep Word64)) (Rep Word64)
	toRep x | x < 0	= Left (Rev (toRep' (fromIntegral (-x))))
		| otherwise = Right (toRep' (fromIntegral x))
		where toRep' = toRep :: Word64 -> Rep Word64
	fromRep (Left (Rev x)) = - fromIntegral ((fromRep :: Rep Word64 -> Word64) x)
	fromRep (Right x) = fromIntegral ((fromRep :: Rep Word64 -> Word64) x)

{-# INLINE toSigned #-}
toSigned :: Integral a => a -> Sum (Rev Word32) Word32
toSigned x
	| x < 0		= Left (Rev (fromIntegral (-x)))
	| otherwise	= Right (fromIntegral x)

{-# INLINE fromSigned #-}
fromSigned :: Integral a => Sum (Rev Word32) Word32 -> a
fromSigned = either (\ (Rev x) -> - fromIntegral x) fromIntegral

instance Repr Word32 where
	type Rep Word32 = Word32
	toRep = id
	fromRep = id

instance Repr ByteString where
	type Rep ByteString = ([Word32], Word32)
	toRep xs = (toList64 xs, fromIntegral (length xs))
	fromRep (xs, n) = case xs of
		[]	-> BS.empty
		(x:xs) -> fst (unfoldrN (fromIntegral n) toBlock (W (Words 3 x) xs))

data Words = Words !Int !Word32
data Words' = W !Words [Word32]

toList64 :: ByteString -> [Word32]
toList64 xs = case BS.foldl' c (Words 4 0, Seq.empty) xs of
	(Words _ w32, ys) -> toList ys ++ [w32]
	where	(Words 0 w, xs) `c` w8
			= (Words 3 (w .|. sL w8 24), xs |> w)
		(Words i' w, xs) `c` w8
			= let !i = i' - 1 in (Words i (w .|. sL w8 (8 * i)), xs)
		sL :: Word8 -> Int -> Word32
		w `sL` x = fromIntegral w `shiftL` x

toBlock :: Words' -> Maybe (Word8, Words')
toBlock (W (Words i0@(i+1) w) xs) = Just (extract w (8 * i0), (W (Words i w) xs))
	where	extract :: Word32 -> Int -> Word8
		extract w x = fromIntegral (w `shiftR` x)
toBlock (W (Words 0 w) (x:xs)) = Just (fromIntegral w, (W (Words 3 x) xs))
toBlock _ = Nothing

instance ReprT Set.Set where
  type RepT Set.Set = []
  toRepTMap f s = Fold.foldr ((:) . f) [] s
  fromRepTMap f xs = Set.fromDistinctAscList [f x | x <- xs] 

genRepr [t| Set.Set |]

instance (Repr k, Repr a) => Repr (Map.Map k a) where
  type Rep (Map.Map k a) = [(Rep k, Rep a)]
  toRep m = [(toRep k, toRep a) | (k, a) <- Map.assocs m]
  fromRep xs = Map.fromDistinctAscList [(fromRep k, fromRep a) | (k, a) <- xs]

instance Repr ISet.IntSet where
  type Rep ISet.IntSet = Rep [Int]
  toRep = toRep . ISet.toList
  fromRep = ISet.fromDistinctAscList . fromRep

instance Repr a => Repr (IMap.IntMap a) where
  type Rep (IMap.IntMap a) = [(Rep Int, Rep a)]
  toRep m = [(toRep i, toRep a) | (i, a) <- IMap.assocs m]
  fromRep xs = IMap.fromDistinctAscList [(fromRep i, fromRep a) | (i, a) <- xs]

instance ReprT Seq.Seq where
  type RepT Seq.Seq = []
  toRepTMap f = Fold.foldr (\ a xs -> f a:xs) []
  fromRepTMap f = Fold.foldl (\ xs a -> xs |> f a) Seq.empty

genRepr [t| Seq.Seq |]