module FMap where

import Maybe

type FMap{|Unit|}		   vT  =  FMUnit (Maybe vT)
type FMap{|Char|}		   vT  =  FMChar (FMapChar vT)
type FMap{|Int|}	       vT  =  FMInt (PatriciaDict vT)
type FMap{| :+: |} fma fmb vT  =  FMEither (Pair (fma vT) (fmb vT))
type FMap{| :*: |} fma fmb vT  =  FMProd (fma (fmb vT)) 
type FMap{| Con |} fma     vT  =  FMCon (fma vT)

data Pair a b		  =  Null | Pair a b deriving (Show,Eq)
type FMapChar vT      =  [(Char,vT)]
type PatriciaDict vT  =  [(Int,vT)]

type Empty{[ * ]} t       =  forall vT . FMap{|t|} vT
type Empty{[kK -> lK]} t  =  forall a . Empty{[kK]} a  ->  Empty{[lK]} (t a) 

empty{|t :: kK|}	  :: Empty{[kK]} t 
empty{|Unit|}		  =  FMUnit Nothing
empty{|Char|}		  =  FMChar []
empty{|Int|}		  =  FMInt []
empty{| :+: |} ea eb  =  FMEither Null
empty{| :*: |} ea eb  =  FMProd ea
empty{| Con c |} ea   =  FMCon ea

type IsEmpty{[ * ]}    t    =  forall vT . FMap{|t|} vT -> Bool
type IsEmpty{[kK -> lK]} t  =  forall a . IsEmpty{[kK]} a  ->  IsEmpty{[lK]} (t a) 

isempty{|t :: kK|}	                      :: IsEmpty{[kK]} t 
isempty{|Unit|}		     (FMUnit v)       =  isNothing v
isempty{|Char|}		     (FMChar l)       =  null l
isempty{|Int|}		     (FMInt l)        =  null l
isempty{| :+: |} iea ieb (FMEither Null)  =  True
isempty{| :+: |} iea ieb (FMEither p)     =  False
isempty{| :*: |} iea ieb (FMProd p)       =  iea p
isempty{| Con c |} iea   (FMCon p)        =  iea p

lookupA                    :: Ord a => a -> [(a,vT)] -> Maybe vT
lookupA a []               =  Nothing
lookupA a ((a', v) : x)
  | a <  a'                =  Nothing
  | a == a'                =  Just v
  | a >  a'                =  lookupA a x

lookupInt  :: Int -> [(Int,vT)] -> Maybe vT
lookupInt  =  lookupA

lookupChar  :: Char -> [(Char,vT)] -> Maybe vT
lookupChar  =  lookupA

type Lookup{[ * ]}      t  =  forall v. t -> FMap{| t |} v -> Maybe v
type Lookup{[ k -> l ]} t  =  forall a. Lookup{[ k ]} a -> Lookup{[ l ]} (t a)

tlookup{| t :: k |}                                          :: Lookup{[ k ]} t
tlookup{| Unit |}       Unit      (FMUnit fm)                =  fm
tlookup{| Char |}       c         (FMChar fm)                =  lookupChar c fm
tlookup{| Int |}        i         (FMInt fm)                 =  lookupInt i fm
tlookup{| :+:  |} lA lB (Inl a)   (FMEither (Pair fma fmb))  =  lA a fma 
tlookup{| :+:  |} lA lB (Inr b)   (FMEither (Pair fma fmb))  =  lB b fmb
tlookup{| :*:  |} lA lB (a :*: b) (FMProd fma)               =  do fmb <- lA a fma
                                                                   lB b fmb
tlookup{| Con d |} l    (Con b)   (FMCon fm)                 =  l b fm

-- single can also be defined as a generic abstraction using add and empty.

dependency single <- single empty

type Single{[ * ]}      t  =  forall v . (t,v) -> FMap{| t |} v
type Single{[ k -> l ]} t  =  forall a . Single{[ k ]} a -> Empty{[ k ]} a -> Single{[ l ]} (t a)

single{| t :: k |}                       :: Single{[ k ]} t
single{| Unit |}            (Unit,v)     =  FMUnit (Just v)
single{| Char |}            (c,v)        =  FMChar [(c,v)]
single{| Int |}             (i,v)        =  FMInt [(i,v)]
single{| :+: |} sA eA sB eB (Inl a,v)    =  FMEither (Pair (sA (a,v)) eB)
single{| :+: |} sA eA sB eB (Inr b,v)    =  FMEither (Pair eA (sB (b,v)))
single{| :*: |} sA eA sB eB (a :*: b,v)  =  FMProd (sA (a,sB (b,v)))
single{| Con d |} sA eA     (Con b,v)    =  FMCon (sA (b,v))

type Merge{[ * ]}      t  =  forall v. 
  (v -> v -> v) -> FMap{| t |} v -> FMap{| t |} v -> FMap{| t |} v
type Merge{[ k -> l ]} t  =  forall a . Merge{[ k ]} a -> Merge{[ l ]} (t a)

merge{| t :: k |}                       :: Merge{[ k ]} t
merge{| Unit |} c      (FMUnit v)            (FMUnit v')            =  FMUnit (combine c v v')
merge{| Char |} c      (FMChar fm)           (FMChar fm')           =  FMChar (foldr insertChar fm' fm)
merge{| Int |} c       (FMInt fm)            (FMInt fm')            =  FMInt (foldr insertInt fm' fm)
merge{| :+: |} mA mB c d                     (FMEither Null)        =  d
merge{| :+: |} mA mB c (FMEither Null)       d                      =  d
merge{| :+: |} mA mB c (FMEither (Pair x y)) (FMEither (Pair v w))  =  
  FMEither (Pair (mA c x v) (mB c y w))
merge{| :*: |} mA mB c (FMProd d)            (FMProd d')            =  FMProd (mA (mB c) d d')
merge{| Con d |} mA c  (FMCon e)             (FMCon e')             =  FMCon (mA c e e')

combine                        :: (v -> v -> v) -> Maybe v -> Maybe v -> Maybe v
combine c Nothing Nothing      =  Nothing
combine c (Just v) Nothing     =  Just v
combine c Nothing (Just v)     =  Just v
combine c (Just v1) (Just v2)  =  Just (c v1 v2)

deleteA                 :: Ord a => a -> [(a,vT)] -> [(a,vT)]
deleteA a []            =  []
deleteA a r@((a', v') : x)
  | a <  a'             =  r
  | a == a'             =  x
  | a >  a'             =  (a',v'):deleteA a x

deleteInt  :: Int -> [(Int,vT)] -> [(Int,vT)]
deleteInt  =  deleteA

deleteChar  :: Char -> [(Char,vT)] -> [(Char,vT)]
deleteChar  =  deleteA

dependency delete <- delete tlookup insert isempty

type Delete{[ * ]} t = forall v . t -> FMap{| t |} v -> FMap{| t |} v
type Delete{[ k -> l ]} t  =  forall a. Delete{[ k ]} a -> Lookup{[ k ]} a -> Insert{[ k ]} a -> IsEmpty{[ k ]} a -> Delete{[ l ]} (t a)

delete{| t :: k |}                                                         :: Delete{[ k ]} t
delete{| Unit |}                          Unit      (FMUnit v)             =  FMUnit Nothing
delete{| Char |}                          c         (FMChar fm)            =  FMChar (deleteChar c fm)
delete{| Int |}                           i         (FMInt fm)             =  FMInt (deleteInt i fm)
delete{| :+: |} dA lA iA ieA dB lB iB ieB (Inl a)   (FMEither (Pair x y))  =  FMEither (Pair (dA a x) y)
delete{| :+: |} dA lA iA ieA dB lB iB ieB (Inr b)   (FMEither (Pair x y))  =  FMEither (Pair x (dB b y))
delete{| :*: |} dA lA iA ieA dB lB iB ieB (a :*: b) (FMProd d)             =  
  let Just d' = lA a d
      d'' = dB b d'        
  in if ieB d'' then FMProd (dA a d) else FMProd (iA (a,d'') d)
delete{| Con c |} dA lA iA ieA            (Con b)   (FMCon d)              =  FMCon (dA b d)


-- insert is defined as a generic abstraction using single and merge

insertga{| t :: * |}       :: (v -> v -> v) -> (t,v) -> FMap{| t |} v -> FMap{| t |} v 
insertga{| t |} c (x,v) d  =  merge{| t |} c (single{| t |} (x,v)) d

-- Can I define insert as a type-indexed function?

insertA                 :: Ord a => (a,vT) -> [(a,vT)] -> [(a,vT)]
insertA (a,v) []        =  [(a,v)]
insertA (a,v) r@((a', v') : x)
  | a <  a'             =  (a,v):r
  | a == a'             =  (a,v):x
  | a >  a'             =  (a',v'):insertA (a,v) x

insertInt  :: (Int,vT) -> [(Int,vT)] -> [(Int,vT)]
insertInt  =  insertA

insertChar  :: (Char,vT) -> [(Char,vT)] -> [(Char,vT)]
insertChar  =  insertA

dependency insert <- insert empty

type Insert{[ * ]} t = forall v . (t,v) -> FMap{| t |} v -> FMap{| t |} v
type Insert{[ k -> l ]} t  =  forall a. Insert{[ k ]} a -> Empty{[ k ]} a -> Insert{[ l ]} (t a)

insert{| t :: k |}                                             :: Insert{[ k ]} t
insert{| Unit |}            (Unit,v)    (FMUnit v')            =  FMUnit (Just v)
insert{| Char |}            (c,v)       (FMChar fm)            =  FMChar (insertChar (c,v) fm)
insert{| Int |}             (i,v)       (FMInt fm)             =  FMInt (insertInt (i,v) fm)
insert{| :+: |} iA eA iB eB (Inl a,v)   (FMEither Null)        =  FMEither (Pair (iA (a,v) eA) eB)
insert{| :+: |} iA eA iB eB (Inr b,v)   (FMEither Null)        =  FMEither (Pair eA (iB (b,v) eB))
insert{| :+: |} iA eA iB eB (Inl a,v)   (FMEither (Pair x y))  =  FMEither (Pair (iA (a,v) x) y)
insert{| :+: |} iA eA iB eB (Inr b,v)   (FMEither (Pair x y))  =  FMEither (Pair x (iB (b,v) y))
insert{| :*: |} iA eA iB eB (a :*: b,v) (FMProd d)             =  FMProd (iA (a,iB (b,v) eB) d)
insert{| Con c |} iA eA     (Con b,v)   (FMCon d)              =  FMCon (iA (b,v) d)

-- Note that the original Unit binding is lost. I tried to combine the
-- two values as in merge but didn't succeed in the product case.

data Nat = Zero | Succ Nat
data Tree a = Leaf a | Node (Tree a) a (Tree a)

-- Testing with trees leads to strange bugs
-- testfmt = empty{|Tree Int|}

testfm = let dict = (delete{|Nat|} (Succ Zero)
                    .insert{|Nat|} (Succ (Succ Zero),2)
                    .insert{|Nat|} (Succ Zero,1)
                    .insert{|Nat|} (Zero,0)
                    ) 
                    empty{|Nat|}
         in (tlookup{|Nat|} (Succ Zero) dict
            ,tlookup{|Nat|} (Succ (Succ Zero)) dict
            ,tlookup{|Nat|} Zero dict
            )