GMapAssoc.hs 2.41 KB
Newer Older
Ian Lynagh's avatar
Ian Lynagh committed
1
{-# LANGUAGE TypeFamilies #-}
2 3 4 5

module Main where

import Prelude hiding (lookup)
6
import Data.Char (ord)
7
import qualified Data.Map as Map
8
import Data.Kind (Type)
9 10 11 12 13

-- Generic maps as ATs
-- -------------------

class GMapKey k where
14
  data GMap k :: Type -> Type
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38
  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 () where
  data GMap () v           = GMapUnit (Maybe v)
  empty                    = GMapUnit Nothing
  lookup () (GMapUnit v)   = v
  insert () 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))
thomie's avatar
thomie committed
39 40
  empty                         = GMapPair empty
  lookup (a, b) (GMapPair gm)   = lookup a gm >>= lookup b
41
  insert (a, b) v (GMapPair gm) = GMapPair $ case lookup a gm of
thomie's avatar
thomie committed
42 43
                                    Nothing  -> insert a (insert b v empty) gm
                                    Just gm2 -> insert a (insert b v gm2  ) gm
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60

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 a) v (GMapEither gm1 gm2) = GMapEither gm1 (insert a v gm2)

-- Test code
-- ---------

nonsence :: GMap Bool String
nonsence = undefined

myGMap :: GMap (Int, Either Char ()) String
myGMap = insert (5, Left 'c') "(5, Left 'c')"    $
thomie's avatar
thomie committed
61 62 63 64 65 66
         insert (4, Right ()) "(4, Right ())"    $
         insert (5, Right ()) "This is the one!" $
         insert (5, Right ()) "This is the two!" $
         insert (6, Right ()) "(6, Right ())"    $
         insert (5, Left 'a') "(5, Left 'a')"    $
         empty
67
main = putStrLn $ maybe "Couldn't find key!" id $ lookup (5, Right ()) myGMap