{-# OPTIONS -fglasgow-exts -fallow-undecidable-instances #-}
-- By Jesse Tov (ccs.neu.edu!tov), November 2007
-- This is a DRAFT.  Later versions at http://www.ccs.neu.edu/home/tov/.
module Env (
  Key, Env, Empty(Empty),
  Extend(extend), Lookup(lookup), Replace(replace), Remove(remove),
  Extend1(extend1), Extend2(extend2),
  Use(use), Modify(modify, modify_, modifyM, modifyM_),
  Partition, peelEnv, peelEnvM,
  outEnv, outEnvM, inEnv, inEnvM,
  eitherEnv, leftEnv, rightEnv,
  List, ListMark,
  introList, ConsList(consList), CdrList(cdrList), ElimList(elimList),
  foldList, foldListM,
) where

import Prelude hiding (lookup)
import Types

-- Errors:

data KeyNotFound n
data KeyCollision n

-- Representation:

data Key n

class EnvList e
newtype Env t e = Env e
newtype Link m n v = Link v

infixr 5 :::
data a ::: b = a ::: b
data Empty   = Empty

instance EnvList Empty
instance (Nat n, EnvList e) => EnvList (Link m n v:::e)

-- Printing:

class EnvList e => ShowEnv e where
  showsEnv :: e -> (String -> String)

instance ShowEnv Empty where
  showsEnv _ = id

instance (Nat n, Show n, Show v) => ShowEnv (Link m n v:::Empty) where
  showsEnv (Link v:::Empty) =
     shows (bot::n) . (" => " ++) . shows v

instance (Nat n, Show n, Show v, ShowEnv (q:::e)) =>
         ShowEnv (Link m n v:::q:::e) where
  showsEnv (p:::e) =
     showsEnv (p:::Empty) . (", " ++) . showsEnv e

instance Show (Env t Empty) where
  showsPrec _ (Env e) = ("Env{" ++) . showsEnv e . ('}' :)

instance ShowEnv (p:::e) => Show (Env t (p:::e)) where
  showsPrec _ (Env e) = ("Env{" ++) . showsEnv e . ('}' :)

-- Operations:

-- extend:
class (EnvList e, Nat n) => Extend m v e e' n | m v e -> e' n where
  extend' :: m -> v -> e -> (e', n)
  extend  :: m -> v -> Env t e -> (Env t e', Key n)
  extend m v (Env e) = let (e', n) = extend' m v e in (Env e', bot)

instance Extend m v Empty (Link m Z v:::Empty) Z where
  extend' m v e = (Link v:::e, bot)
instance (Nat n0, EnvList e) =>
         Extend m v (Link m0 n0 v0:::e)
                    (Link m (S n0) v:::Link m0 n0 v0:::e) (S n0) where
  extend' m v e = (Link v:::e, bot)

class (EnvList e, Nat n) => Extend1 m c e e' n
                            | m c e -> e' n where
  extend1' :: m -> (forall a. c a) -> e -> (e', n)
  extend1  :: m -> (forall a. c a) -> Env t e -> (Env t e', Key n)
  extend1 m v (Env e) = let (e', n) = extend1' m v e in (Env e', bot)

instance Extend1 m c Empty (Link m Z (c a):::Empty) Z where
  extend1' m v e = (Link v:::e, bot)
instance (Nat n0, EnvList e) =>
         Extend1 m c (Link m0 n0 v0:::e)
                         (Link m (S n0) (c a):::Link m0 n0 v0:::e)
                         (S n0) where
  extend1' m v e = (Link v:::e, bot)

class (EnvList e, Nat n) => Extend2 m c e e' n
                            | m c e -> e' n where
  extend2' :: m -> (forall a b. c a b) -> e -> (e', n)
  extend2  :: m -> (forall a b. c a b) -> Env t e -> (Env t e', Key n)
  extend2 m v (Env e) = let (e', n) = extend2' m v e in (Env e', bot)

instance Extend2 m c Empty (Link m Z (c a b):::Empty) Z where
  extend2' m v e = (Link v:::e, bot)
instance (Nat n0, EnvList e) =>
         Extend2 m c (Link m0 n0 v0:::e)
                     (Link m (S n0) (c a b):::Link m0 n0 v0:::e)
                     (S n0) where
  extend2' m v e = (Link v:::e, bot)

-- lookup:
class (EnvList e, Nat n) => Lookup m n e v | n e -> m v where
  lookup' :: m -> n -> e -> v
  lookup  :: m -> Key n -> Env t e -> v
  lookup m _ (Env e) = lookup' m (bot::n) e

instance (Nat n, Fail (KeyNotFound n)) => Lookup m n Empty () where
  lookup' _ _ _ = bot

instance (Nat n0,
          Equal n n0 b,
          LookupEnvCase b m0 v0 m n e v) =>
         Lookup m n (Link m0 n0 v0:::e) v where
  lookup' m n (Link v0:::e) =
    lookupEnvCase (equal n (bot::n0)) (bot::m0) v0 m n e

class (Prop b, EnvList e, Nat n) =>
      LookupEnvCase b m0 v0 m n e v | b v0 n e -> m v where
  lookupEnvCase :: b -> m0 -> v0 -> m -> n -> e -> v

instance (EnvList e, Nat n) => LookupEnvCase T m v0 m n e v0 where
  lookupEnvCase _ _ v _ _ _ = v

instance Lookup m n e v => LookupEnvCase F m0 v0 m n e v where
  lookupEnvCase _ _ _ m = lookup' m

-- remove:
class (EnvList e, Nat n, EnvList e') =>
      Remove m n e e' | n e -> m e' where
  remove' :: m -> n -> e -> e'
  remove  :: m -> Key n -> Env t e -> Env t e'
  remove m _ (Env e) = Env (remove' m (bot::n) e)

instance (Nat n, Fail (KeyNotFound n)) => Remove m n Empty Empty where
  remove' _ _ _ = bot

instance (Nat n0,
          Equal n n0 b,
          RemoveEnvCase b (Link m0 n0 v0) m n e e') =>
         Remove m n (Link m0 n0 v0:::e) e' where
  remove' m n (link@(Link v0):::e)
    = removeEnvCase (equal n (bot::n0)) link m n e

class (Prop b, EnvList e, Nat n, EnvList e') =>
      RemoveEnvCase b p m n e e' | b p n e -> m e' where
  removeEnvCase :: b -> p -> m -> n -> e -> e'

instance (EnvList e, Nat n) => RemoveEnvCase T (Link m n0 v0) m n e e where
  removeEnvCase _ _ _ _ e = e

instance (EnvList (p:::e'), Remove m n e e') =>
         RemoveEnvCase F p m n e (p:::e') where
  removeEnvCase _ p m n e = p ::: remove' m n e

-- replace
class (EnvList e, Nat n, EnvList e') =>
      Replace m n v e e' | n v e -> m e' where
  replace' :: m -> n -> v -> e -> e'
  replace  :: m -> Key n -> v -> Env t e -> Env t e'
  replace m _ v (Env e) = Env (replace' m (bot::n) v e)

instance (Nat n, Fail (KeyNotFound n)) => Replace m n v Empty Empty where
  replace' _ _ _ _ = bot

instance (Nat n0,
          Equal n n0 b,
          ReplaceEnvCase b (Link m0 n0 v0) m n v e e') =>
         Replace m n v (Link m0 n0 v0:::e) e' where
  replace' m n v (link@(Link v0):::e)
    = replaceEnvCase (equal n (bot::n0)) link m n v e

class (Prop b, EnvList e, Nat n, EnvList e') =>
      ReplaceEnvCase b p m n v e e' | b p n v e -> m e' where
  replaceEnvCase :: b -> p -> m -> n -> v -> e -> e'

instance (EnvList e, Nat n) =>
         ReplaceEnvCase T (Link m n0 v0) m n v e (Link m n v:::e) where
  replaceEnvCase _ _ m n v e = Link v:::e

instance (EnvList (p:::e'), Replace m n v e e') =>
         ReplaceEnvCase F p m n v e (p:::e') where
  replaceEnvCase _ p m n v e = p ::: replace' m n v e

-- partition
class (EnvList e, EnvList e1, EnvList e2) =>
      Partition sel e e1 e2 | sel e -> e1 e2 where
  partition' :: sel -> e -> (e1, e2)

instance Partition sel Empty Empty Empty where
  partition' _ Empty = (Empty, Empty)

instance (EnvList e, EnvList e1, EnvList e2, Nat n,
          NatMemberHelper n sel b,
          PartitionCase b sel (Link m n v) e e1 e2) =>
         Partition sel (Link m n v:::e) e1 e2 where
  partition' sel (p@(Link v):::e) =
    partitionCase (natMemberHelper (bot::n) sel) sel p e

class PartitionCase b sel p e e1 e2 | b sel p e -> e1 e2 where
  partitionCase :: b -> sel -> p -> e -> (e1, e2)

instance Partition sel e e1 e2 =>
         PartitionCase T sel p e (p:::e1) e2 where
  partitionCase _ sel p e = (p:::e1, e2) where
      (e1, e2) = partition' sel e

instance Partition sel e e1 e2 =>
         PartitionCase F sel p e e1 (p:::e2) where
  partitionCase _ sel p e = (e1, p:::e2) where
      (e1, e2) = partition' sel e

class NatMemberHelper n ks b | n ks -> b where
  natMemberHelper :: n -> ks -> b

instance NatMemberHelper n () F where
  natMemberHelper _ _ = bot

instance (Equal n n0 b1, NatMemberHelper n ks b2, OrProp b1 b2 b) =>
         NatMemberHelper n (Key n0, ks) b where
  natMemberHelper _ _ = bot

-- Environment wrappers
outEnv :: (forall t. Env t e -> (Env t e', a)) ->
          e -> (e', a)
outEnv f e = (e', a) where (Env e', a) = f (Env e)

outEnvM :: Monad m => (forall t. Env t e -> m (Env t e', a)) ->
           e -> m (e', a)
outEnvM f e = f (Env e) >>= \(Env e', a) -> return (e', a)

inEnv :: (e -> (e', a)) -> Env t e -> (Env t e', a)
inEnv f (Env e) = (Env e', a) where (e', a) = f e

inEnvM :: Monad m => (e -> m (e', a)) -> Env t e -> m (Env t e', a)
inEnvM f (Env e) = f e >>= \(e', a) -> return (Env e', a)

peelEnv :: Partition ks e e1 e2 =>
        ks -> (forall t. Env t e1 -> (Env t Empty, a)) ->
        e -> (e2, a)
peelEnv ks f e =
  let (e1, e2)       = partition' ks e
      (Env Empty, a) = f (Env e1)
   in (e2, a)

peelEnvM :: (Partition ks e e1 e2, Monad m0) =>
         ks -> (forall t. Env t e1 -> m0 (Env t Empty, a)) ->
         e -> (e2, m0 a)
peelEnvM ks f e =
  let (e1, e2) = partition' ks e
   in (e2, f (Env e1) >>= \(Env Empty, a) -> return a)

-- derived classes
class (Remove m n e e', Lookup m n e v) =>
      Use m n e e' v | m n e -> e' v where
  use :: m -> Key n -> Env t e -> (Env t e', v)
  use m k e = (remove m k e, lookup m k e)
instance (Remove m n e e', Lookup m n e v) => Use m n e e' v where

class (Lookup m n e v, Replace m n v' e e') =>
      Modify m n v e v' e' | m n e -> v, m n v' e -> e' where
  modify   :: m -> Key n -> (v -> (v', a)) -> Env t e -> (Env t e', a)
  modify_  :: m -> Key n -> (v -> v') -> Env t e -> Env t e'
  modifyM  :: Monad m0 => m -> Key n -> (v -> m0 (v', a)) ->
              Env t e -> m0 (Env t e', a)
  modifyM_ :: Monad m0 => m -> Key n -> (v -> m0 v') ->
              Env t e -> m0 (Env t e')
  modify m k f e   = let (v, a) = f (lookup m k e) in
                       (replace m k v e, a)
  modify_ m k f e  = replace m k (f (lookup m k e)) e
  modifyM m k f e  = f (lookup m k e) >>= \(v', a) ->
                     return (replace m k v' e, a)
  modifyM_ m k f e = f (lookup m k e) >>= \v' ->
                     return (replace m k v' e)
instance (Lookup m n e v, Replace m n v' e e') =>
         Modify m n v e v' e' where

-- Either:

eitherEnv :: Env t (Either e1 e2) -> Either (Env t e1) (Env t e2)
leftEnv   :: Env t e1 -> Env t (Either e1 e2)
rightEnv  :: Env t e2 -> Env t (Either e1 e2)
eitherEnv (Env e) = either (Left . Env) (Right . Env) e
leftEnv (Env e)   = Env (Left e)
rightEnv (Env e)  = Env (Right e)

-- Lists:

data ListMark    = ListMark
newtype List m a = List [a]

introList :: Extend2 ListMark List e e' n =>
             Env t e -> (Env t e', Key n)
introList e = extend2 ListMark (List []) e

class (Lookup ListMark n e (List m a),
       Replace ListMark n (List m a) e e,
       Extend m a e e' n') =>
      CdrList n e m a n' e' | n e -> m a n' e' where
  cdrList :: Key n -> Env t e -> Either (Env t e', Key n') (Env t e)
  cdrList k e =
    case lookup ListMark k e of
      List (x:xs) -> Left $
                       extend (bot::m) x $
                         replace ListMark k (List xs::List m a) e
      List []     -> Right e
instance (Lookup ListMark n e (List m a),
          Replace ListMark n (List m a) e e,
          Extend m a e e' n') =>
         CdrList n e m a n' e'

class (Lookup ListMark n e (List m a),
       Replace ListMark n (List m a) e e,
       Extend m a e e' n',
       Remove ListMark n e e'') =>
      ElimList n e m a n' e' e'' | n e -> m a n' e' e'' where
  elimList :: Key n -> Env t e -> Either (Env t e', Key n') (Env t e'')
  elimList k e =
    either Left (Right . remove ListMark k) (cdrList k e)
instance (Lookup ListMark n e (List m a),
          Replace ListMark n (List m a) e e,
          Extend m a e e' n',
          Remove ListMark n e e'') =>
         ElimList n e m a n' e' e''

class (Lookup m na e a,
       Lookup ListMark nl e (List m a),
       Remove m na e e',
       Replace ListMark nl (List m a) e' e') =>
      ConsList na nl e m a e' | na nl e -> m a e' where
  consList :: Key na -> Key nl -> Env t e -> Env t e'
  consList ka kl e =
    let m      = bot::m
        a      = lookup m ka e
        List l = lookup ListMark kl e
        e'     = remove m ka e
     in replace ListMark kl (List (a:l) :: List m a) e'
instance (Lookup m na e a,
          Lookup ListMark nl e (List m a),
          Remove m na e e',
          Replace ListMark nl (List m a) e' e') =>
         ConsList na nl e m a e'

foldList :: ElimList nl e m a na e' e'' =>
            (z -> Key na -> Env t e' -> (Env t e, z)) ->
            z -> Key nl -> Env t e -> (Env t e'', z)
foldList f z kl e = case elimList kl e of
  Left (e', ka) -> let (e'', z') = f z ka e'
                     in foldList f z' kl e''
  Right e'      -> (e', z)

foldListM :: (Monad m0, ElimList nl e m a na e' e'') =>
            (z -> Key na -> Env t e' -> m0 (Env t e, z)) ->
            z -> Key nl -> Env t e -> m0 (Env t e'', z)
foldListM f z kl e = case elimList kl e of
  Left (e', ka) -> do (e'', z') <- f z ka e'
                      foldListM f z' kl e''
  Right e'      -> return (e', z)