{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, RankNTypes #-}

----------------------------------------------------------------------
-- |
-- Module      : PGF.TypeCheck
-- Maintainer  : Krasimir Angelov
-- Stability   : (stable)
-- Portability : (portable)
--
-- Type checking in abstract syntax with dependent types.
-- The type checker also performs renaming and checking for unknown
-- functions. The variable references are replaced by de Bruijn indices.
--
-----------------------------------------------------------------------------

module PGF.TypeCheck ( checkType, checkExpr, inferExpr

                     , ppTcError, TcError(..)

                     -- internals needed for the typechecking of forests
                     , MetaStore, emptyMetaStore, newMeta, newGuardedMeta
                     , getMeta, setMeta, lookupMeta, MetaValue(..)
                     , Scope, emptyScope, scopeSize, scopeEnv, addScopedVar
                     , TcM(..), runTcM, TType(..), Selector(..)
                     , tcExpr, infExpr, eqType, eqValue
                     , lookupFunType, typeGenerators, eval
                     , generateForMetas, generateForForest, checkResolvedMetaStore
                     ) where

import PGF.Data
import PGF.Expr hiding (eval, apply, applyValue, value2expr)
import qualified PGF.Expr as Expr
import PGF.Macros (typeOfHypo, cidInt, cidFloat, cidString)
import PGF.CId

import Data.Map as Map
import Data.IntMap as IntMap
import Data.Maybe as Maybe
import Data.List as List
import Control.Monad
import Control.Monad.Identity
import Control.Monad.State
import Control.Monad.Error
import Text.PrettyPrint

-----------------------------------------------------
-- The Scope
-----------------------------------------------------

data    TType = TTyp Env Type
newtype Scope = Scope [(CId,TType)]

emptyScope = Scope []

addScopedVar :: CId -> TType -> Scope -> Scope
addScopedVar x tty (Scope gamma) = Scope ((x,tty):gamma)

-- | returns the type and the De Bruijn index of a local variable
lookupVar :: CId -> Scope -> Maybe (Int,TType)
lookupVar x (Scope gamma) = listToMaybe [(i,tty) | ((y,tty),i) <- zip gamma [0..], x == y]

-- | returns the type and the name of a local variable
getVar :: Int -> Scope -> (CId,TType)
getVar i (Scope gamma) = gamma !! i

scopeEnv :: Scope -> Env
scopeEnv (Scope gamma) = let n = length gamma
                         in [VGen (n-i-1) [] | i <- [0..n-1]]

scopeVars :: Scope -> [CId]
scopeVars (Scope gamma) = List.map fst gamma

scopeSize :: Scope -> Int
scopeSize (Scope gamma) = length gamma

-----------------------------------------------------
-- The Monad
-----------------------------------------------------

type MetaStore s = IntMap (MetaValue s)
data MetaValue s
  = MUnbound s Scope TType [Expr -> TcM s ()]
  | MBound   Expr
  | MGuarded Expr  [Expr -> TcM s ()] {-# UNPACK #-} !Int   -- the Int is the number of constraints that have to be solved 
                                                            -- to unlock this meta variable

newtype TcM s a = TcM {unTcM :: forall b . Abstr -> (a -> MetaStore s -> s -> b -> b) 
                                                 -> (TcError          -> s -> b -> b) 
                                                 -> (MetaStore s      -> s -> b -> b)}

class Selector s where
  splitSelector :: s -> (s,s)
  select        :: CId -> Scope -> Maybe Int -> TcM s (Expr,TType)

instance Monad (TcM s) where
  return x = TcM (\abstr k h -> k x)
  f >>= g  = TcM (\abstr k h -> unTcM f abstr (\x -> unTcM (g x) abstr k h) h)

instance Selector s => MonadPlus (TcM s) where
  mzero = TcM (\abstr k h ms s -> id)
  mplus f g = TcM (\abstr k h ms s -> let (s1,s2) = splitSelector s
                                      in unTcM f abstr k h ms s1 . unTcM g abstr k h ms s2)

instance MonadState s (TcM s) where
  get = TcM (\abstr k h ms s -> k s ms s)
  put s = TcM (\abstr k h ms _ -> k () ms s)

instance MonadError TcError (TcM s) where
  throwError e    = TcM (\abstr k h ms -> h e)
  catchError f fh = TcM (\abstr k h ms -> unTcM f abstr k (\e s -> unTcM (fh e) abstr k h ms s) ms)

instance Functor (TcM s) where
  fmap f m = TcM (\abstr k h -> unTcM m abstr (k . f) h)

runTcM :: Abstr -> TcM s a -> MetaStore s -> s -> ([(s,TcError)],[(MetaStore s,s,a)])
runTcM abstr f ms s = unTcM f abstr (\x ms s cp b -> let (es,xs) = cp b
                                                     in (es,(ms,s,x) : xs))
                                    (\e    s cp b -> let (es,xs) = cp b
                                                     in ((s,e) : es,xs))
                                    ms s id ([],[])

lookupCatHyps :: CId -> TcM s [Hypo]
lookupCatHyps cat = TcM (\abstr k h ms -> case Map.lookup cat (cats abstr) of
                                            Just (hyps,_) -> k hyps ms
                                            Nothing       -> h (UnknownCat cat))

lookupFunType :: CId -> TcM s Type
lookupFunType fun = TcM (\abstr k h ms -> case Map.lookup fun (funs abstr) of
                                            Just (ty,_,_,_) -> k ty ms
                                            Nothing         -> h (UnknownFun fun))

typeGenerators :: Scope -> CId -> TcM s [(Double,Expr,TType)]
typeGenerators scope cat = fmap normalize (liftM2 (++) x y)
  where
    x = return
           [(0.25,EVar i,tty) | (i,(_,tty@(TTyp _ (DTyp _ cat' _)))) <- zip [0..] gamma
                              , cat == cat']
         where
           Scope gamma = scope

    y | cat == cidInt    = return [(1.0,ELit (LInt 999),  TTyp [] (DTyp [] cat []))]
      | cat == cidFloat  = return [(1.0,ELit (LFlt 3.14), TTyp [] (DTyp [] cat []))]
      | cat == cidString = return [(1.0,ELit (LStr "Foo"),TTyp [] (DTyp [] cat []))]
      | otherwise        = TcM (\abstr k h ms ->
                                    case Map.lookup cat (cats abstr) of
                                      Just (_,fns) -> unTcM (mapM helper fns) abstr k h ms
                                      Nothing      -> h (UnknownCat cat))

    helper (p,fn) = do
      ty <- lookupFunType fn
      return (p,EFun fn,TTyp [] ty)
      
    normalize gens = [(p/s,e,tty) | (p,e,tty) <- gens]
      where
        s = sum [p | (p,_,_) <- gens]

emptyMetaStore :: MetaStore s
emptyMetaStore = IntMap.empty

newMeta :: Scope -> TType -> TcM s MetaId
newMeta scope tty = TcM (\abstr k h ms s -> let metaid = IntMap.size ms + 1
                                            in k metaid (IntMap.insert metaid (MUnbound s scope tty []) ms) s)

newGuardedMeta :: Expr -> TcM s MetaId
newGuardedMeta e = TcM (\abstr k h ms s -> let metaid = IntMap.size ms + 1
                                           in k metaid (IntMap.insert metaid (MGuarded e [] 0) ms) s)

getMeta :: MetaId -> TcM s (MetaValue s)
getMeta i = TcM (\abstr k h ms -> case IntMap.lookup i ms of
                                    Just mv -> k mv ms)

setMeta :: MetaId -> MetaValue s -> TcM s ()
setMeta i mv = TcM (\abstr k h ms -> k () (IntMap.insert i mv ms))

lookupMeta ms i =
  case IntMap.lookup i ms of
    Just (MBound   t)                 -> Just t
    Just (MGuarded t _ x) | x == 0    -> Just t
                          | otherwise -> Nothing
    Just (MUnbound _ _ _ _)           -> Nothing
    Nothing                           -> Nothing

addConstraint :: MetaId -> MetaId -> (Expr -> TcM s ()) -> TcM s ()
addConstraint i j c = do
  mv   <- getMeta j
  case mv of
    MUnbound s scope tty cs     -> addRef >> setMeta j (MUnbound s scope tty ((\e -> release >> c e) : cs))
    MBound   e                  -> c e
    MGuarded e cs x | x == 0    -> c e
                    | otherwise -> addRef >> setMeta j (MGuarded e ((\e -> release >> c e) : cs) x)
  where
    addRef  = TcM (\abstr k h ms -> case IntMap.lookup i ms of
                                      Just (MGuarded e cs x) -> k () $! IntMap.insert i (MGuarded e cs (x+1)) ms)

    release = TcM (\abstr k h ms -> case IntMap.lookup i ms of
                                      Just (MGuarded e cs x) -> if x == 1
                                                                  then unTcM (sequence_ [c e | c <- cs]) abstr k h $! IntMap.insert i (MGuarded e [] 0) ms
                                                                  else k () $! IntMap.insert i (MGuarded e cs (x-1)) ms)

-----------------------------------------------------
-- Type errors
-----------------------------------------------------

-- | If an error occurs in the typechecking phase
-- the type checker returns not a plain text error message
-- but a 'TcError' structure which describes the error.
data TcError
  = UnknownCat      CId                            -- ^ Unknown category name was found.
  | UnknownFun      CId                            -- ^ Unknown function name was found.
  | WrongCatArgs    [CId] Type CId  Int Int        -- ^ A category was applied to wrong number of arguments.
                                                   -- The first integer is the number of expected arguments and
                                                   -- the second the number of given arguments.
                                                   -- The @[CId]@ argument is the list of free variables
                                                   -- in the type. It should be used for the 'showType' function.
  | TypeMismatch    [CId] Expr Type Type           -- ^ The expression is not of the expected type.
                                                   -- The first type is the expected type, while
                                                   -- the second is the inferred. The @[CId]@ argument is the list
                                                   -- of free variables in both the expression and the type. 
                                                   -- It should be used for the 'showType' and 'showExpr' functions.
  | NotFunType      [CId] Expr Type                -- ^ Something that is not of function type was applied to an argument.
  | CannotInferType [CId] Expr                     -- ^ It is not possible to infer the type of an expression.
  | UnresolvedMetaVars [CId] Expr [MetaId]         -- ^ Some metavariables have to be instantiated in order to complete the typechecking.
  | UnexpectedImplArg [CId] Expr                   -- ^ Implicit argument was passed where the type doesn't allow it
  | UnsolvableGoal [CId] MetaId Type               -- ^ There is a goal that cannot be solved
  deriving Eq

-- | Renders the type checking error to a document. See 'Text.PrettyPrint'.
ppTcError :: TcError -> Doc
ppTcError (UnknownCat cat)             = text "Category" <+> ppCId cat <+> text "is not in scope"
ppTcError (UnknownFun fun)             = text "Function" <+> ppCId fun <+> text "is not in scope"
ppTcError (WrongCatArgs xs ty cat m n) = text "Category" <+> ppCId cat <+> text "should have" <+> int m <+> text "argument(s), but has been given" <+> int n $$
                                         text "In the type:" <+> ppType 0 xs ty
ppTcError (TypeMismatch xs e ty1 ty2)  = text "Couldn't match expected type" <+> ppType 0 xs ty1 $$
                                         text "       against inferred type" <+> ppType 0 xs ty2 $$
                                         text "In the expression:" <+> ppExpr 0 xs e
ppTcError (NotFunType xs e ty)         = text "A function type is expected for the expression" <+> ppExpr 0 xs e <+> text "instead of type" <+> ppType 0 xs ty
ppTcError (CannotInferType xs e)       = text "Cannot infer the type of expression" <+> ppExpr 0 xs e
ppTcError (UnresolvedMetaVars xs e ms) = text "Meta variable(s)" <+> fsep (List.map ppMeta ms) <+> text "should be resolved" $$
                                         text "in the expression:" <+> ppExpr 0 xs e
ppTcError (UnexpectedImplArg xs e)     = braces (ppExpr 0 xs e) <+> text "is implicit argument but not implicit argument is expected here"
ppTcError (UnsolvableGoal xs metaid ty)= text "The goal:" <+> ppMeta metaid <+> colon <+> ppType 0 xs ty $$
                                         text "cannot be solved"

-----------------------------------------------------
-- checkType
-----------------------------------------------------

-- | Check whether a given type is consistent with the abstract
-- syntax of the grammar.
checkType :: PGF -> Type -> Either TcError Type
checkType pgf ty = 
  unTcM (do ty <- tcType emptyScope ty
            refineType ty)
        (abstract pgf)
        (\ty ms s _ -> Right ty)
        (\err  s _ -> Left err)
        emptyMetaStore () (error "checkType")

tcType :: Scope -> Type -> TcM s Type
tcType scope ty@(DTyp hyps cat es) = do
  (scope,hyps) <- tcHypos scope hyps
  c_hyps <- lookupCatHyps cat
  let m = length es
      n = length [ty | (Explicit,x,ty) <- c_hyps]
  (delta,es) <- tcCatArgs scope es [] c_hyps ty n m
  return (DTyp hyps cat es)

tcHypos :: Scope -> [Hypo] -> TcM s (Scope,[Hypo])
tcHypos scope []     = return (scope,[])
tcHypos scope (h:hs) = do
  (scope,h ) <- tcHypo  scope h
  (scope,hs) <- tcHypos scope hs
  return (scope,h:hs)

tcHypo :: Scope -> Hypo -> TcM s (Scope,Hypo)
tcHypo scope (b,x,ty) = do
  ty <- tcType scope ty
  if x == wildCId
    then return (scope,(b,x,ty))
    else return (addScopedVar x (TTyp (scopeEnv scope) ty) scope,(b,x,ty))

tcCatArgs scope []              delta []                   ty0 n m = return (delta,[])
tcCatArgs scope (EImplArg e:es) delta ((Explicit,x,ty):hs) ty0 n m = throwError (UnexpectedImplArg (scopeVars scope) e)
tcCatArgs scope (EImplArg e:es) delta ((Implicit,x,ty):hs) ty0 n m = do
  e <- tcExpr scope e (TTyp delta ty)
  (delta,es) <- if x == wildCId
                  then tcCatArgs scope es delta  hs ty0 n m
                  else do v <- eval (scopeEnv scope) e
                          tcCatArgs scope es (v:delta) hs ty0 n m
  return (delta,EImplArg e:es)
tcCatArgs scope es delta ((Implicit,x,ty):hs) ty0 n m = do
  i <- newMeta scope (TTyp delta ty)
  (delta,es) <- if x == wildCId
                  then tcCatArgs scope es                                delta  hs ty0 n m
                  else tcCatArgs scope es (VMeta i (scopeEnv scope) [] : delta) hs ty0 n m
  return (delta,EImplArg (EMeta i) : es)
tcCatArgs scope (e:es) delta ((Explicit,x,ty):hs) ty0 n m = do
  e <- tcExpr scope e (TTyp delta ty)
  (delta,es) <- if x == wildCId
                  then tcCatArgs scope es                               delta  hs ty0 n m
                  else do v <- eval (scopeEnv scope) e
                          tcCatArgs scope es (v:delta) hs ty0 n m
  return (delta,e:es)
tcCatArgs scope _ delta _ ty0@(DTyp _ cat _) n m = do
  throwError (WrongCatArgs (scopeVars scope) ty0 cat n m)

-----------------------------------------------------
-- checkExpr
-----------------------------------------------------

-- | Checks an expression against a specified type.
checkExpr :: PGF -> Expr -> Type -> Either TcError Expr
checkExpr pgf e ty =
  unTcM (do e <- tcExpr emptyScope e (TTyp [] ty)
            checkResolvedMetaStore emptyScope e)
        (abstract pgf)
        (\e ms s _ -> Right e)
        (\err  s _ -> Left err)
        emptyMetaStore () (error "checkExpr")

tcExpr :: Scope -> Expr -> TType -> TcM s Expr
tcExpr scope e0@(EAbs Implicit x e) tty =
  case tty of
    TTyp delta (DTyp ((Implicit,y,ty):hs) c es) -> do e <- if y == wildCId
                                                             then tcExpr (addScopedVar x (TTyp delta ty) scope)
                                                                         e (TTyp delta (DTyp hs c es))
                                                             else tcExpr (addScopedVar x (TTyp delta ty) scope)
                                                                         e (TTyp ((VGen (scopeSize scope) []):delta) (DTyp hs c es))
                                                      return (EAbs Implicit x e)
    _                                           -> do ty <- evalType (scopeSize scope) tty
                                                      throwError (NotFunType (scopeVars scope) e0 ty)
tcExpr scope e0 (TTyp delta (DTyp ((Implicit,y,ty):hs) c es)) = do
  e0 <- if y == wildCId
          then tcExpr (addScopedVar wildCId (TTyp delta ty) scope)
                      e0 (TTyp delta (DTyp hs c es))
          else tcExpr (addScopedVar wildCId (TTyp delta ty) scope)
                      e0 (TTyp ((VGen (scopeSize scope) []):delta) (DTyp hs c es))
  return (EAbs Implicit wildCId e0)
tcExpr scope e0@(EAbs Explicit x e) tty =
  case tty of
    TTyp delta (DTyp ((Explicit,y,ty):hs) c es) -> do e <- if y == wildCId
                                                             then tcExpr (addScopedVar x (TTyp delta ty) scope)
                                                                         e (TTyp delta (DTyp hs c es))
                                                             else tcExpr (addScopedVar x (TTyp delta ty) scope)
                                                                         e (TTyp ((VGen (scopeSize scope) []):delta) (DTyp hs c es))
                                                      return (EAbs Explicit x e)
    _                                           -> do ty <- evalType (scopeSize scope) tty
                                                      throwError (NotFunType (scopeVars scope) e0 ty)
tcExpr scope (EMeta _) tty = do
  i <- newMeta scope tty
  return (EMeta i)
tcExpr scope e0        tty = do
  (e0,tty0) <- infExpr scope e0
  (e0,tty0) <- appImplArg scope e0 tty0
  i <- newGuardedMeta e0
  eqType scope (scopeSize scope) i tty tty0
  return (EMeta i)


-----------------------------------------------------
-- inferExpr
-----------------------------------------------------

-- | Tries to infer the type of a given expression. Note that
-- even if the expression is type correct it is not always
-- possible to infer its type in the GF type system.
-- In this case the function returns the 'CannotInferType' error.
inferExpr :: PGF -> Expr -> Either TcError (Expr,Type)
inferExpr pgf e =
  unTcM (do (e,tty) <- infExpr emptyScope e
            e <- checkResolvedMetaStore emptyScope e
            ty <- evalType 0 tty
            return (e,ty))
        (abstract pgf)
        (\e_ty ms s _ -> Right e_ty)
        (\err     s _ -> Left err)
        emptyMetaStore () (error "inferExpr")

infExpr :: Scope -> Expr -> TcM s (Expr,TType)
infExpr scope e0@(EApp e1 e2) = do
  (e1,TTyp delta ty) <- infExpr scope e1
  (e0,delta,ty) <- tcArg scope e1 e2 delta ty
  return (e0,TTyp delta ty)
infExpr scope e0@(EFun x) = do
  case lookupVar x scope of
    Just (i,tty) -> return (EVar i,tty)
    Nothing      -> do ty <- lookupFunType x
                       return (e0,TTyp [] ty)
infExpr scope e0@(EVar i) = do
  return (e0,snd (getVar i scope))
infExpr scope e0@(ELit l) = do
  let cat = case l of
              LStr _ -> mkCId "String"
              LInt _ -> mkCId "Int"
              LFlt _ -> mkCId "Float"
  return (e0,TTyp [] (DTyp [] cat []))
infExpr scope (ETyped e ty) = do
  ty <- tcType scope ty
  e  <- tcExpr scope e (TTyp (scopeEnv scope) ty)
  return (ETyped e ty,TTyp (scopeEnv scope) ty)
infExpr scope (EImplArg e) = do
  (e,tty)  <- infExpr scope e
  return (EImplArg e,tty)
infExpr scope e = throwError (CannotInferType (scopeVars scope) e)

tcArg scope e1 e2 delta ty0@(DTyp [] c es) = do
  ty1 <- evalType (scopeSize scope) (TTyp delta ty0)
  throwError (NotFunType (scopeVars scope) e1 ty1)
tcArg scope e1 (EImplArg e2) delta ty0@(DTyp ((Explicit,x,ty):hs) c es) = throwError (UnexpectedImplArg (scopeVars scope) e2)
tcArg scope e1 (EImplArg e2) delta ty0@(DTyp ((Implicit,x,ty):hs) c es) = do
  e2 <- tcExpr scope e2 (TTyp delta ty)
  if x == wildCId
    then return (EApp e1 (EImplArg e2),                              delta,DTyp hs c es)
    else do v2 <- eval (scopeEnv scope) e2
            return (EApp e1 (EImplArg e2),v2:delta,DTyp hs c es)
tcArg scope e1 e2 delta ty0@(DTyp ((Explicit,x,ty):hs) c es) = do
  e2 <- tcExpr scope e2 (TTyp delta ty)
  if x == wildCId
    then return (EApp e1 e2,delta,DTyp hs c es)
    else do v2 <- eval (scopeEnv scope) e2
            return (EApp e1 e2,v2:delta,DTyp hs c es)
tcArg scope e1 e2 delta ty0@(DTyp ((Implicit,x,ty):hs) c es) = do
  i <- newMeta scope (TTyp delta ty)
  if x == wildCId
    then tcArg scope (EApp e1 (EImplArg (EMeta i))) e2                                delta  (DTyp hs c es)
    else tcArg scope (EApp e1 (EImplArg (EMeta i))) e2 (VMeta i (scopeEnv scope) [] : delta) (DTyp hs c es)

appImplArg scope e (TTyp delta (DTyp ((Implicit,x,ty1):hypos) cat es)) = do
  i <- newMeta scope (TTyp delta ty1)
  let delta' = if x == wildCId
                 then                               delta  
                 else VMeta i (scopeEnv scope) [] : delta
  appImplArg scope (EApp e (EImplArg (EMeta i))) (TTyp delta' (DTyp hypos cat es))
appImplArg scope e tty                                                 = return (e,tty)

-----------------------------------------------------
-- eqType
-----------------------------------------------------

eqType :: Scope -> Int -> MetaId -> TType -> TType -> TcM s ()
eqType scope k i0 tty1@(TTyp delta1 ty1@(DTyp hyps1 cat1 es1)) tty2@(TTyp delta2 ty2@(DTyp hyps2 cat2 es2))
  | cat1 == cat2 = do (k,delta1,delta2) <- eqHyps k delta1 hyps1 delta2 hyps2
                      sequence_ [eqExpr raiseTypeMatchError (addConstraint i0) k delta1 e1 delta2 e2 | (e1,e2) <- zip es1 es2]
  | otherwise    = raiseTypeMatchError
  where
    raiseTypeMatchError = do ty1 <- evalType k tty1
                             ty2 <- evalType k tty2
                             e   <- refineExpr (EMeta i0)
                             throwError (TypeMismatch (scopeVars scope) e ty1 ty2)

    eqHyps :: Int -> Env -> [Hypo] -> Env -> [Hypo] -> TcM s (Int,Env,Env)
    eqHyps k delta1 []                 delta2 []                 =
      return (k,delta1,delta2)
    eqHyps k delta1 ((_,x,ty1) : h1s) delta2 ((_,y,ty2) : h2s) = do
      eqType scope k i0 (TTyp delta1 ty1) (TTyp delta2 ty2)
      if x == wildCId && y == wildCId
        then eqHyps k delta1 h1s delta2 h2s
        else if x /= wildCId && y /= wildCId
               then eqHyps (k+1) ((VGen k []):delta1) h1s ((VGen k []):delta2) h2s
               else raiseTypeMatchError
    eqHyps k delta1               h1s  delta2               h2s  = raiseTypeMatchError

eqExpr :: (forall a . TcM s a) -> (MetaId -> (Expr -> TcM s ()) -> TcM s ()) -> Int -> Env -> Expr -> Env -> Expr -> TcM s ()
eqExpr fail suspend k env1 e1 env2 e2 = do
  v1 <- eval env1 e1
  v2 <- eval env2 e2
  eqValue fail suspend k v1 v2

eqValue :: (forall a . TcM s a) -> (MetaId -> (Expr -> TcM s ()) -> TcM s ()) -> Int -> Value -> Value -> TcM s ()
eqValue fail suspend k v1 v2 = do
  v1 <- deRef v1
  v2 <- deRef v2
  eqValue' k v1 v2
  where
    deRef v@(VMeta i env vs) = do
      mv   <- getMeta i
      case mv of
        MBound   e                 -> apply env e vs
        MGuarded e _ x | x == 0    -> apply env e vs
                       | otherwise -> return v
        MUnbound _ _ _ _           -> return v
    deRef v = return v

    eqValue' k (VSusp i env vs1 c)            v2                             = suspend i (\e -> apply env e vs1 >>= \v1 -> eqValue fail suspend k (c v1) v2)
    eqValue' k v1                             (VSusp i env vs2 c)            = suspend i (\e -> apply env e vs2 >>= \v2 -> eqValue fail suspend k v1 (c v2))
    eqValue' k (VMeta i env1 vs1)             (VMeta j env2 vs2) | i  == j   = zipWithM_ (eqValue fail suspend k) vs1 vs2
    eqValue' k (VMeta i env1 vs1)             v2                             = do mv <- getMeta i
                                                                                  case mv of
                                                                                    MUnbound _ scopei _ cs -> bind i scopei cs env1 vs1 v2
                                                                                    MGuarded e cs x        -> setMeta i (MGuarded e ((\e -> apply env1 e vs1 >>= \v1 -> eqValue' k v1 v2) : cs) x)
    eqValue' k v1                             (VMeta i env2 vs2)             = do mv <- getMeta i
                                                                                  case mv of
                                                                                    MUnbound _ scopei _ cs -> bind i scopei cs env2 vs2 v1
                                                                                    MGuarded e cs x        -> setMeta i (MGuarded e ((\e -> apply env2 e vs2 >>= \v2 -> eqValue' k v1 v2) : cs) x)
    eqValue' k (VApp f1 vs1)                  (VApp f2 vs2)   | f1 == f2     = zipWithM_ (eqValue fail suspend k) vs1 vs2
    eqValue' k (VConst f1 vs1)                (VConst f2 vs2) | f1 == f2     = zipWithM_ (eqValue fail suspend k) vs1 vs2
    eqValue' k (VLit l1)                      (VLit l2    )   | l1 == l2     = return ()
    eqValue' k (VGen  i vs1)                  (VGen  j vs2)   | i  == j      = zipWithM_ (eqValue fail suspend k) vs1 vs2
    eqValue' k (VClosure env1 (EAbs _ x1 e1)) (VClosure env2 (EAbs _ x2 e2)) = let v = VGen k []
                                                                               in eqExpr fail suspend (k+1) (v:env1) e1 (v:env2) e2
    eqValue' k (VClosure env1 (EAbs _ x1 e1)) v2                             = let v = VGen k []
                                                                               in do v1 <- eval (v:env1) e1
                                                                                     v2 <- applyValue v2 [v]
                                                                                     eqValue fail suspend (k+1) v1 v2
    eqValue' k v1                             (VClosure env2 (EAbs _ x2 e2)) = let v = VGen k []
                                                                               in do v1 <- applyValue v1 [v]
                                                                                     v2 <- eval (v:env2) e2
                                                                                     eqValue fail suspend (k+1) v1 v2
    eqValue' k v1                             v2                             = fail

    bind i scope cs env vs0 v = do
      let k  = scopeSize scope
          vs  = reverse (take k env) ++ vs0
          xs  = nub [i | VGen i [] <- vs]
      if length vs /= length xs
        then suspend i (\e -> apply env e vs0 >>= \iv -> eqValue fail suspend k iv v)
        else do v <- occurCheck i k xs v
                e0 <- value2expr (length xs) v
                let e = addLam vs0 e0
                setMeta i (MBound e)
                sequence_ [c e | c <- cs]
      where
        addLam []     e = e
        addLam (v:vs) e = EAbs Explicit var (addLam vs e)

        var = mkCId "v"

    occurCheck i0 k xs (VApp f vs)      = do vs <- mapM (occurCheck i0 k xs) vs
                                             return (VApp f vs)
    occurCheck i0 k xs (VLit l)         = return (VLit l)
    occurCheck i0 k xs (VMeta i env vs) = do if i == i0
                                               then fail
                                               else return ()
                                             mv <- getMeta i
                                             case mv of
                                               MBound   e                                   -> apply env e vs >>= occurCheck i0 k xs
                                               MGuarded e _ _                               -> apply env e vs >>= occurCheck i0 k xs
                                               MUnbound _ scopei _ _ | scopeSize scopei > k -> fail
                                                                     | otherwise            -> do vs <- mapM (occurCheck i0 k xs) vs
                                                                                                  return (VMeta i env vs)
    occurCheck i0 k xs (VSusp i env vs cnt) = do suspend i (\e -> apply env e vs >>= \v -> occurCheck i0 k xs (cnt v) >> return ())
                                                 return (VSusp i env vs cnt)
    occurCheck i0 k xs (VGen  i vs)     = case List.findIndex (==i) xs of
                                            Just i  -> do vs <- mapM (occurCheck i0 k xs) vs
                                                          return (VGen i vs)
                                            Nothing -> fail
    occurCheck i0 k xs (VConst f vs)    = do vs <- mapM (occurCheck i0 k xs) vs
                                             return (VConst f vs)
    occurCheck i0 k xs (VClosure env e) = do env <- mapM (occurCheck i0 k xs) env
                                             return (VClosure env e)
    occurCheck i0 k xs (VImplArg e)     = do e <- occurCheck i0 k xs e
                                             return (VImplArg e)


-----------------------------------------------------------
-- three ways of dealing with meta variables that
-- still have to be resolved
-----------------------------------------------------------

checkResolvedMetaStore :: Scope -> Expr -> TcM s Expr
checkResolvedMetaStore scope e = do
  e <- refineExpr e
  TcM (\abstr k h ms -> 
           let xs = [i | (i,mv) <- IntMap.toList ms, not (isResolved mv)]
           in if List.null xs
                then k () ms
                else h (UnresolvedMetaVars (scopeVars scope) e xs))
  return e
  where
    isResolved (MUnbound _ _ _ []) = True
    isResolved (MGuarded  _ _ _)   = True
    isResolved (MBound    _)       = True
    isResolved _                   = False

generateForMetas :: Selector s => (Scope -> TType -> TcM s Expr) -> Expr -> TcM s Expr
generateForMetas prove e = do
  (e,_) <- infExpr emptyScope e
  fillinVariables
  refineExpr e
  where
    fillinVariables = do
      fvs <- TcM (\abstr k h ms -> k [(i,s,scope,tty,cs) | (i,MUnbound s scope tty cs) <- IntMap.toList ms] ms)
      case fvs of
        []                   -> return ()
        (i,_,scope,tty,cs):_ -> do e <- prove scope tty
                                   setMeta i (MBound e)
                                   sequence_ [c e | c <- cs]
                                   fillinVariables

generateForForest :: (Scope -> TType -> TcM FId Expr) -> Expr -> TcM FId Expr
generateForForest prove e = do
--  fillinVariables
  refineExpr e
  where
    fillinVariables = do
      fvs <- TcM (\abstr k h ms -> k [(i,s,scope,tty,cs) | (i,MUnbound s scope tty cs) <- IntMap.toList ms] ms)
      case fvs of
        []                   -> return ()
        (i,s,scope,tty,cs):_ -> TcM (\abstr k h ms s0 ->
                                        case snd $ runTcM abstr (prove scope tty) ms s of
                                          []           -> unTcM (do ty <- evalType (scopeSize scope) tty
                                                                    throwError (UnsolvableGoal (scopeVars scope) s ty)
                                                                ) abstr k h ms s
                                          ((ms,_,e):_) -> unTcM (do setMeta i (MBound e)
                                                                    sequence_ [c e | c <- cs]
                                                                    fillinVariables
                                                                ) abstr k h ms s)

-----------------------------------------------------
-- evalType
-----------------------------------------------------

evalType :: Int -> TType -> TcM s Type
evalType k (TTyp delta ty) = evalTy funs k delta ty
  where
    evalTy sig k delta (DTyp hyps cat es) = do
      (k,delta,hyps) <- evalHypos sig k delta hyps
      es <- mapM (\e -> eval delta e >>= value2expr k) es
      return (DTyp hyps cat es)

    evalHypos sig k delta []              = return (k,delta,[])
    evalHypos sig k delta ((b,x,ty):hyps) = do
      ty <- evalTy sig k delta ty
      (k,delta,hyps) <- if x == wildCId
                          then evalHypos sig k                  delta  hyps
                          else evalHypos sig (k+1) ((VGen k []):delta) hyps
      return (k,delta,(b,x,ty) : hyps)


-----------------------------------------------------
-- refinement
-----------------------------------------------------

refineExpr :: Expr -> TcM s Expr
refineExpr e = TcM (\abstr k h ms -> k (refineExpr_ ms e) ms)

refineExpr_ ms e = refine e
  where
    refine (EAbs b x e)  = EAbs b x (refine e)
    refine (EApp e1 e2)  = EApp (refine e1) (refine e2)
    refine (ELit l)      = ELit l
    refine (EMeta i)     = case IntMap.lookup i ms of
                             Just (MBound   e    ) -> refine e
                             Just (MGuarded e _ _) -> refine e
                             _                     -> EMeta i
    refine (EFun f)      = EFun f
    refine (EVar i)      = EVar i
    refine (ETyped e ty) = ETyped (refine e) (refineType_ ms ty)
    refine (EImplArg e)  = EImplArg (refine e)

refineType :: Type -> TcM s Type
refineType ty = TcM (\abstr k h ms -> k (refineType_ ms ty) ms)

refineType_ ms (DTyp hyps cat es) = DTyp [(b,x,refineType_ ms ty) | (b,x,ty) <- hyps] cat (List.map (refineExpr_ ms) es)

eval :: Env -> Expr -> TcM s Value
eval env e = TcM (\abstr k h ms -> k (Expr.eval (funs abstr,lookupMeta ms) env e) ms)

apply :: Env -> Expr -> [Value] -> TcM s Value
apply env e vs = TcM (\abstr k h ms -> k (Expr.apply (funs abstr,lookupMeta ms) env e vs) ms)

applyValue :: Value -> [Value] -> TcM s Value
applyValue v vs = TcM (\abstr k h ms -> k (Expr.applyValue (funs abstr,lookupMeta ms) v vs) ms)

value2expr :: Int -> Value -> TcM s Expr
value2expr i v = TcM (\abstr k h ms -> k (Expr.value2expr (funs abstr,lookupMeta ms) i v) ms)