module Haskell.TypeCheck (typeCheck) where

import qualified Data.Map.Strict as Map
import Haskell.AST


typeCheck :: (Eq a, TypeCheck a) => a -> Either String a
typeCheck initVal = go initVal Map.empty
  where
    go value mapping = do
        value2 <- annotate mapping value
        if value2 == value then return value else go value2 mapping

class TypeCheck a where
    annotate :: Map.Map String Type -> a -> Either String a

instance TypeCheck AST where
    annotate mapping (AST tops) = AST <$> mapM (annotate mapping) tops

instance TypeCheck Toplevel where
    annotate mapping (TopDef x)   = TopDef <$> annotate mapping x
    annotate _       (TopDecl x)  = return $ TopDecl x
    annotate _       (TopData x)  = return $ TopData x
    annotate _       (TopClass x) = return $ TopClass x
    annotate mapping (TopInst x)  = TopInst <$> annotate mapping x

instance TypeCheck Def where
    annotate mapping (Def name expr) = Def name <$> annotate mapping expr

instance TypeCheck Inst where
    annotate mapping (Inst name ty defs) = Inst name ty <$> mapM (annotate mapping) defs

-- TODO: Just passing mappings around is not sufficient; the typechecking
-- needs to be done on the whole environment at once, and dynamically
-- updating the mapping as it goes. In fact, I'd probably do better just
-- looking up how the simplest possible implementation of Hindley-Miller
-- goes, stripping away most classes, and then implementing that.
instance TypeCheck Expr where
    annotate mapping (App func' [] _) = annotate mapping func'
    annotate mapping (App func' [arg'] _) = do
        func <- annotate mapping func'
        arg <- annotate mapping arg'
        case typeApply (typeOf func) (typeOf arg) of
            Just ty -> return $ App func [arg] (Just ty)
            Nothing -> Left $ "Cannot apply argument of type " ++ show (typeOf arg) ++
                              " to function of type " ++ show (typeOf func)
    annotate mapping (App func' (arg':args') _) = 
        combineApps <$> annotate mapping (App (App func' [arg'] Nothing) args' Nothing)

    -- annotate mapping (Ref )


combineApps :: Expr -> Expr
combineApps (App (App f as1 _) as2 ty) = combineApps (App f (as1 ++ as2) ty)
combineApps e = recurse id combineApps e