module CheckModule where

import Control.Monad (liftM, liftM2, liftM3)
import Control.Monad.State (State, evalState, get, put)
import Data.List (intersect)

import AbsDep
import ToExp
import PrintDep

import TypeCheck
import Util

type TypeDecl = (Id, Exp)
type Funs = [(Id,FunDef)]
data FunDef = FunDef Env Int [([Pattern],Exp)]
type PattDecl = (Id, [Pattern], Exp)


checkModules :: (Env,Env,Funs) -> [(String,Module)] -> (Env,Env,Funs)
checkModules c [] = c
checkModules c@(r,g,fs) ((n,m):ms) = checkModules (r',g',fs') ms
    where (tds,pds) = decls m
	  (r',g') = checkTypeDecls (r,g) tds
	  fs' = checkPattDecls (r',g',fs) pds

-- | Get all type declarations from a module. Numbers all
--   metavariables in the declarations.
decls :: Module -> ([TypeDecl],[PattDecl])
decls (Module _ ds) = catEithers $ evalState (mapM decl_ ds) 0
    where
    decl_ (TypeDecl (Ident i) e) = liftM (Left . (,) i) $ nm (toExp e)
    decl_ (PattDecl (Ident i) ps e) = liftM (Right . (,,) i ps) $ nm (toExp e)

-- | Number all metavariables in an expression.
nm :: Exp -> State Int Exp
nm x = case x of
	      Let id e1 e2 e3 -> liftM3 (Let id) (nm e1) (nm e2) (nm e3)
	      Abs id e1 -> liftM (Abs id) (nm e1)
	      Pi id e1 e2 -> liftM2 (Pi id) (nm e1) (nm e2)
	      App e1 e2 -> liftM2 App (nm e1) (nm e2)
	      Var _ -> return x
	      Type -> return x
	      Meta _ -> do
			n <- get
			put (n+1)
			return (Meta n)

-- | Check that the types in a list of type declarations are
--   all of type Type. Extends the given environments with
--   new declarations where the values of the declared variables
--   are constants, and the types of the variables are calculated
--   in the extened value environment.
checkTypeDecls :: (Env,Env) -> [TypeDecl] -> (Env,Env)
checkTypeDecls (rho,gamma) ds 
    | null conflicts && checkTypes = (rho',gamma')
    | otherwise = error $ "checkTypeDecls: Duplicate type declarations: "-- ++ show conflicts
    where
    conflicts = map fst rho `intersect` map fst ds
    rho' = rho ++ [ (n,VConst n) | (n,_) <- ds ]
    gamma' = gamma ++ [ (n,eval rho' e) | (n,e) <- ds ]
    checkTypes = all (checkType' (0,rho',gamma') . snd) ds


-- simple version that only checks declarations with no pattern
checkPattDecls :: (Env,Env,Funs) -> [PattDecl] -> Funs
checkPattDecls (_,_,fs) [] = fs
checkPattDecls (r,g,fs) ((i,[],e):ds) 
	       | checkExp' (0,r,g) e t = 
		   -- FIXME: resolve metavariables
		   checkPattDecls (r,g, addFunCase r i [] e fs) ds
	       | otherwise = error $ i ++ " does not have type " -- ++ show t
    where t = lookUp i g
-- FIXME: check these too
checkPattDecls (r,g,fs) ((i,ps,e):ds) = 
    checkPattDecls (r,g, addFunCase r i ps e fs) ds


{-
-- FIXME: need to carry constraints between patterns for same fun
-- FIXME: need to use constraints to refine fun type after checking all patterns
checkPattDecls :: Int -> (Env,Env,Funs) -> [PattDecl] -> Funs
checkPattDecls m (_,_,fs) [] = fs
checkPattDecls m (r,g,fs) ((i,ps,e):ds) = 
    case checkPattDecl m (0,r,g) ps e t of
         Failed err -> error $ "checking " ++ i ++ " " ++ printTree ps ++ ": " ++ err
	 Ok m' cs -> let (ss,cs') = solve cs
			 e' = refine ss e
			 in checkPattDecls m' (r,g,addFunCase i (FunCase ps e') fs) ds
    where t = lookUp i g

checkPattDecl :: Int -> (Int,Env,Env) -> [Pattern] -> Exp -> Val 
	      -> CheckResult Int -- ^ returns meta base
checkPattDecl m (k,r,g) [] e t = checkExp (k,r,g) e t `join` Ok m []
checkPattDecl m (k,r,g) (p:ps) e (VPi t f) =
    checkPattType m (k,r,g) p t `thenR` (\ (m,k,r,g) -> checkPattDecl m (k+1,r,g) ps e (fapp f (VGen k)))
checkPattDecl _ _ ps _ t = Failed $ "checkPattDecl: Couldn't match patterns " ++ show ps ++ " with type " ++ show t

consType :: Int -> (Env,Env) -> Val -> ([Val],Val)
consType m (r,g) (VPi t f) = let (as,rt) = consType (m+1) (r,g) (fapp f (VMeta m)) in (t:as,rt)
consType m (r,g) t = ([],t)

checkPattType :: Int           -- ^ meta base
	      -> (Int,Env,Env) -- ^ (generic value base, rho, gamma)
	      -> Pattern 
	      -> Val -> CheckResult (Int,Int,Env,Env)
checkPattType m (k,r,g) p t = 
    case p of 
	   PCons (Ident c) ps -> let (as,rt) = consType m (r,g) (lookUp c g)
			     in checkPattTypes m (k,r,g) ps as `thenR` \x -> Ok x [EqC rt t]
	   PVar (Ident x) -> case lookup x r of
				   Nothing -> 
				       let v = VMeta m
					   in Ok (m+1, k, (x,v):r, (x,t):g) []
				   Just v -> 
				       let vt = lookUp x g 
					   in Ok (m,k,r,g) [EqC vt t]
	   PType -> Ok (m,k,r,g) [EqC t VType]
	   PWild -> Ok (m,k,r,g) []

checkPattTypes :: Int -> (Int,Env,Env) -> [Pattern] -> [Val] -> CheckResult (Int,Int,Env,Env)
checkPattTypes m (k,r,g) [] [] = Ok (m,k,r,g) []
checkPattTypes m (k,r,g) (p:ps) (t:ts) = 
    checkPattType m (k,r,g) p t `thenR` (\ (m,k,r,g) -> checkPattTypes m (k,r,g) ps ts)
-}



addFunCase :: Env -> Id -> [Pattern] -> Exp -> Funs -> Funs
addFunCase env i ps e [] = [(i,FunDef env (length ps) [(ps,e)])]
addFunCase env i ps e ((i',d@(FunDef _ a cs)):fs) 
    | i == i' = if length ps /= a then
		  error $ "Arity mis-match for " ++ i
		 else
                  (i',FunDef env a ((ps,e):cs)) : fs
    | True = (i',d) : addFunCase env i ps e fs