{-# LANGUAGE FlexibleInstances
  , BangPatterns
  , MagicHash 
  , ScopedTypeVariables
  , TypeFamilies 
  , UndecidableInstances
  , OverlappingInstances
  , MultiParamTypeClasses
  #-}

module GMap where

import Data.Word
import Data.Int
import Data.Bits
import Data.Char (ord,chr)
import qualified Data.Map as DM
import qualified Data.IntMap as DI
import qualified Data.List as L
import Prelude hiding (lookup)

import Debug.Trace

--------------------------------------------------------------------------------
-- Class of types that fit in a machine word.
--------------------------------------------------------------------------------

-- |All datatypes that can be packed into a single word, including
-- scalars and some tuple types.
class FitInWord v where 
  toWord   :: v -> Word
  fromWord :: Word -> v

intToWord :: Int -> Word
intToWord = fromIntegral

wordToInt :: Word -> Int
wordToInt = fromIntegral 

instance FitInWord Char where
  toWord   = intToWord . ord
  fromWord = chr . wordToInt

instance FitInWord Int where
  toWord   = fromIntegral
  fromWord = fromIntegral

instance FitInWord Int16 where
  toWord   = fromIntegral
  fromWord = fromIntegral

instance FitInWord Int8 where
  toWord   = fromIntegral
  fromWord = fromIntegral

instance FitInWord Word8 where
  toWord   = fromIntegral
  fromWord = fromIntegral

instance FitInWord Word16 where
  toWord   = fromIntegral
  fromWord = fromIntegral


#ifdef x86_64_HOST_ARCH
instance FitInWord Int64 where
  toWord   = fromIntegral
  fromWord = fromIntegral
instance FitInWord Word64 where
  toWord   = fromIntegral
  fromWord = fromIntegral
#endif

-- Some tuples can fit in words too!
-- FIXME TODO: Use some code generation method to generate instances for all
-- combinations of small words/ints that fit in a machine word (a lot).
instance FitInWord (Word16,Word16) where
  toWord (a,b) = shiftL (fromIntegral a) 16 + (fromIntegral b)
  fromWord n = (fromIntegral$ shiftR n 16, 
		fromIntegral$ n .&. 0xFFFF)

instance FitInWord (Int16,Int16) where
  toWord (a,b) = shiftL (fromIntegral a) 16 + (fromIntegral b)
  fromWord n = (fromIntegral$ shiftR n 16, 
		fromIntegral$ n .&. 0xFFFF)


--------------------------------------------------------------------------------
-- ADT definition for generic Maps:
--------------------------------------------------------------------------------

-- |Class for generic map key types.  By using indexed type families,
-- each key type may correspond to a different data structure that
-- implements it.
class (Ord k, Eq k, Show k) => GMapKey k where
  data GMap k :: * -> *
  empty       :: GMap k v
  lookup      :: k -> GMap k v -> Maybe v
  insert      :: k -> v -> GMap k v -> GMap k v
  alter       :: (Maybe a -> Maybe a) -> k -> GMap k a -> GMap k a
  toList      :: GMap k a -> [(k,a)]

--------------------------------------------------------------------------------

-- What problems was I running into here:
-- It's hard to avoid conflicting instances, for example with the tuple instance.
-- I think I may need a NotFitInWord class constraint..
#if 0
instance FitInWord t => GMapKey t where
  data GMap t v           = GMapInt (DI.IntMap v) deriving Show
  --empty                   = trace "\n <<<<< FitInWord Gmap... >>>>\n"$ GMapInt DI.empty
  empty                   = GMapInt DI.empty
  lookup k    (GMapInt m) = DI.lookup (wordToInt$ toWord k) m
  insert k v  (GMapInt m) = GMapInt (DI.insert (wordToInt$ toWord k) v m)
  alter  fn k (GMapInt m) = GMapInt (DI.alter fn (wordToInt$ toWord k) m)
  toList      (GMapInt m) = map (\ (i,v) -> (fromWord$ intToWord i, v)) $ 
			    DI.toList m

#else

-- Unit and Bool can have specialized implementations, but because
-- they also "FitInWord", these result in conflicts.
------------------------------------------------------------
instance GMapKey () where
  data GMap () v           = GMapUnit (Maybe v)
  empty                    = GMapUnit Nothing
  lookup ()   (GMapUnit v) = v
  insert () v (GMapUnit _) = GMapUnit $ Just v
  alter fn () (GMapUnit v) = GMapUnit $ fn v
  toList (GMapUnit Nothing) = []
  toList (GMapUnit (Just v)) = [((),v)]
instance GMapKey Bool where
  data GMap Bool v              = GMapBool (Maybe v) (Maybe v)
  empty                       = GMapBool Nothing Nothing
  lookup True  (GMapBool v _) = v
  lookup False (GMapBool _ v) = v
  insert True v  (GMapBool a b) = GMapBool (Just v) b
  insert False v (GMapBool a b) = GMapBool a (Just v)
  alter fn True  (GMapBool a b) = GMapBool (fn a) b
  alter fn False (GMapBool a b) = GMapBool a (fn b)
  toList (GMapBool Nothing Nothing)   = []
  toList (GMapBool (Just a) Nothing)  = [(True,a)]
  toList (GMapBool Nothing (Just b))  = [(False,b)]
  toList (GMapBool (Just a) (Just b)) = [(True,a),(False,b)]
------------------------------------------------------------

-- GMaps on Int keys become Data.IntMaps:
instance GMapKey Int where
  data GMap Int v         = GMapInt (DI.IntMap v) deriving Show
  empty                   = GMapInt DI.empty
  lookup k    (GMapInt m) = DI.lookup k m
  insert k v  (GMapInt m) = GMapInt (DI.insert k v m)
  alter  fn k (GMapInt m) = GMapInt (DI.alter fn k m)
  toList      (GMapInt m) = DI.toList m

-- Then other scalar keys can be converted to Ints:
-- CODE DUPLICATION
instance GMapKey Char where
  data GMap Char v         = GMapChar (GMap Int v) deriving Show
  empty                    = GMapChar empty
  lookup k (GMapChar m)    = lookup (ord k) m
  insert k v (GMapChar m)  = GMapChar (insert (ord k) v m)
  alter  fn k (GMapChar m) = GMapChar (alter fn (ord k) m)
  toList      (GMapChar m) = map (\ (i,v) -> (chr i,v)) (toList m)

instance GMapKey Word8 where
  data GMap Word8 v        = GMapWord8 (GMap Int v) deriving Show
  empty                    = GMapWord8 empty
  lookup k (GMapWord8 m)    = lookup (fromIntegral k) m
  insert k v (GMapWord8 m)  = GMapWord8 (insert (fromIntegral k) v m)
  alter  fn k (GMapWord8 m) = GMapWord8 (alter fn (fromIntegral k) m)
  toList      (GMapWord8 m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)

instance GMapKey Word16 where
  data GMap Word16 v         = GMapWord16 (GMap Int v) deriving Show
  empty                    = GMapWord16 empty
  lookup k (GMapWord16 m)    = lookup (fromIntegral k) m
  insert k v (GMapWord16 m)  = GMapWord16 (insert (fromIntegral k) v m)
  alter  fn k (GMapWord16 m) = GMapWord16 (alter fn (fromIntegral k) m)
  toList      (GMapWord16 m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)

instance GMapKey Word where
  data GMap Word v        = GMapWord (GMap Int v) deriving Show
  empty                    = GMapWord empty
  lookup k (GMapWord m)    = lookup (fromIntegral k) m
  insert k v (GMapWord m)  = GMapWord (insert (fromIntegral k) v m)
  alter  fn k (GMapWord m) = GMapWord (alter fn (fromIntegral k) m)
  toList      (GMapWord m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)

instance GMapKey Int8 where
  data GMap Int8 v        = GMapInt8 (GMap Int v) deriving Show
  empty                    = GMapInt8 empty
  lookup k (GMapInt8 m)    = lookup (fromIntegral k) m
  insert k v (GMapInt8 m)  = GMapInt8 (insert (fromIntegral k) v m)
  alter  fn k (GMapInt8 m) = GMapInt8 (alter fn (fromIntegral k) m)
  toList      (GMapInt8 m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)

instance GMapKey Int16 where
  data GMap Int16 v         = GMapInt16 (GMap Int v) deriving Show
  empty                    = GMapInt16 empty
  lookup k (GMapInt16 m)    = lookup (fromIntegral k) m
  insert k v (GMapInt16 m)  = GMapInt16 (insert (fromIntegral k) v m)
  alter  fn k (GMapInt16 m) = GMapInt16 (alter fn (fromIntegral k) m)
  toList      (GMapInt16 m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)

-- TODO: Int64 + Word64

-- Can't get past the "Conflicting family instance declarations"
-- Comment out this block and it works:
#if 1
instance GMapKey (Int16, Int16) where
  data GMap (Int16,Int16) v         = GMapInt16Pair (GMap Int v) deriving Show
  empty                     = trace "<<< Constructing Double-Int16 GMAP (Intmap) >>>> " $ 
			      GMapInt16Pair empty
  lookup k (GMapInt16Pair m)    = lookup (fromIntegral k) m
  insert k v (GMapInt16Pair m)  = GMapInt16Pair (insert (fromIntegral k) v m)
  alter  fn k (GMapInt16Pair m) = GMapInt16Pair (alter fn (fromIntegral k) m)
  toList      (GMapInt16Pair m) = map (\ (i,v) -> (fromIntegral i,v)) (toList m)
#endif

#endif


--------------------------------------------------------------------------------

-- |GMaps over pairs are implemented by nested GMaps.
instance (GMapKey a, GMapKey b) => GMapKey (a, b) where
  data GMap (a, b) v            = GMapPair (GMap a (GMap b v))
  empty		                = --trace "CONSTRUCTED GMAP USING NESTED MAPS!"$ 
                                  GMapPair empty
  lookup (a, b) (GMapPair gm)   = lookup a gm >>= lookup b 
  insert (a, b) v (GMapPair gm) = GMapPair $ case lookup a gm of
				    Nothing  -> insert a (insert b v empty) gm
				    Just gm2 -> insert a (insert b v gm2  ) gm
  alter fn (a, b) (GMapPair gm) = GMapPair $ alter newfun a gm
     where 
       newfun entry =
	   case entry of 
	    Nothing -> case fn Nothing of 
	                Nothing -> Nothing
	                Just v  -> Just $ insert b v empty
	    Just m -> Just$ alter fn b m

  toList (GMapPair gm) = L.foldl' (\ acc (a,m) -> map (\ (b,v) -> ((a,b),v)) (toList m) ++ acc) [] $ 
			 toList gm

-- |Sum types are represented by separate GMaps for the separate variants.
instance (GMapKey a, GMapKey b) => GMapKey (Either a b) where
  data GMap (Either a b) v                = GMapEither (GMap a v) (GMap b v)
  empty                                   = GMapEither empty empty
  lookup (Left  a) (GMapEither gm1  _gm2) = lookup a gm1
  lookup (Right b) (GMapEither _gm1 gm2 ) = lookup b gm2
  insert (Left  a) v (GMapEither gm1 gm2) = GMapEither (insert a v gm1) gm2
  insert (Right b) v (GMapEither gm1 gm2) = GMapEither gm1 (insert b v gm2)
  alter fn (Left  a) (GMapEither gm1 gm2) = GMapEither (alter fn a gm1) gm2
  alter fn (Right b) (GMapEither gm1 gm2) = GMapEither gm1 (alter fn b gm2)
  toList (GMapEither gm1 gm2) = 
      map (\ (a,v) -> (Left  a, v)) (toList gm1) ++ 
      map (\ (b,v) -> (Right b, v)) (toList gm2)

-- |GMaps with list indices could be treated like tuples (nested
-- maps).  Instead, we put them in a regular Data.Map.
instance (GMapKey a) => GMapKey [a] where
  data GMap [a] v         = GMapList (DM.Map [a] v) deriving Show
  empty                   = GMapList DM.empty
  lookup k    (GMapList m) = DM.lookup k m
  insert k v  (GMapList m) = GMapList (DM.insert k v m)
  alter  fn k (GMapList m) = GMapList (DM.alter fn k m)
  toList      (GMapList m) = DM.toList m