{-# LANGUAGE TypeFamilies, TypeOperators #-}

module GMap where

import Prelude hiding (lookup)
import Char (ord)
import qualified Data.Map as Map


-- Product-sum type representations
-- --------------------------------

data Unit = Unit
  deriving Show
data (:+:) a b = Inl a | Inr b
  deriving Show
data (:*:) a b = a :*: b
  deriving Show

class Representable a where
  type Repr a
  toRepr   :: a -> Repr a
  fromRepr :: Repr a -> a

instance Representable Unit where
  type Repr Unit  = Unit
  toRepr   = id
  fromRepr = id

instance Representable () where
  type Repr ()  = Unit
  toRepr ()     = Unit
  fromRepr Unit = () 

instance Representable Int where
  type Repr Int = Int
  toRepr   = id
  fromRepr = id

instance Representable Char where
  type Repr Char = Char
  toRepr   = id
  fromRepr = id

instance (Representable a, Representable b) => Representable (a :*: b) where
  type Repr (a :*: b) = a :*: b
  toRepr   = id
  fromRepr = id

instance (Representable a, Representable b) => Representable (a, b) where
  type Repr (a, b) = a :*: b
  toRepr (x, y)      = x :*: y
  fromRepr (x :*: y) = (x, y) 

instance (Representable a, Representable b) => Representable (a :+: b) where
  type Repr (a :+: b) = a :+: b
  toRepr   = id
  fromRepr = id

instance (Representable a, Representable b) => Representable (Either a b) where
  type Repr (Either a b) = a :+: b
  toRepr (Left x)  = Inl x
  toRepr (Right y) = Inr y
  fromRepr (Inl x) = Left x
  fromRepr (Inr x) = Right x


-- Generalised maps
-- ----------------

class Representable 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

instance GMapKey Int where
  data GMap Int v        = GMapInt (Map.Map Int v)
  empty                  = GMapInt Map.empty
  lookup k (GMapInt m)   = Map.lookup k m
  insert k v (GMapInt m) = GMapInt (Map.insert k v m)

instance GMapKey Char where
  data GMap Char v        = GMapChar (GMap Int v)
  empty                   = GMapChar empty
  lookup k (GMapChar m)   = lookup (ord k) m
  insert k v (GMapChar m) = GMapChar (insert (ord k) v m)

instance GMapKey Unit where
  data GMap Unit v           = GMapUnit (Maybe v)
  empty                      = GMapUnit Nothing
  lookup Unit (GMapUnit v)   = v
  insert Unit v (GMapUnit _) = GMapUnit $ Just v

instance (GMapKey a, GMapKey b) => GMapKey (a :*: b) where
  data GMap (a :*: b) v            = GMapPair (GMap a (GMap b v))
  empty		                   = 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

instance (GMapKey a, GMapKey b) => GMapKey (a :+: b) where
  data GMap (a :+: b) v                  = GMapEither (GMap a v) (GMap b v)
  empty                                  = GMapEither empty empty
  lookup (Inl  a) (GMapEither gm1  _gm2) = lookup a gm1
  lookup (Inr b) (GMapEither _gm1 gm2 )  = lookup b gm2
  insert (Inl  a) v (GMapEither gm1 gm2) = GMapEither (insert a v gm1) gm2
  insert (Inr a) v (GMapEither gm1 gm2)  = GMapEither gm1 (insert a v gm2)


-- Derived instances
-- -----------------

-- Generic list representation
--
-- * Could be added to Typeable
-- * Could use n-ary sums and products
--
instance Representable [a] where
  type Repr [a] = Unit :+: (a :*: [a]) 
  toRepr []     = Inl Unit
  toRepr (x:xs) = Inr (x :*: xs)
  fromRepr (Inl Unit)       = []
  fromRepr (Inr (x :*: xs)) = x:xs

-- List-indexed maps
--
-- * Could also be generated
--
instance GMapKey a => GMapKey [a] where
  newtype GMap [a] v = GMapList (GMap (Repr [a]) v)
  empty                    = GMapList empty
  lookup k (GMapList gm)   = lookup (toRepr k) gm
  insert k v (GMapList gm) = GMapList $ insert (toRepr k) v gm


-- Local Variables: 
-- haskell-ghci-program-name: "/Users/chak/Code/ghc-test/compiler/stage2/ghc-inplace"
-- haskell-ghci-program-args: ("--interactive")
-- End: