module TypeInference where

import Control.Monad.State.Strict
    ( gets, modify, evalState, State )
import qualified Data.Map.Strict as Map

import Terms ( Term(..) )
import Types ( TermInContext(..), Context, Type(TypeVar, Arrow) )
import Unification ( Unifier, unify, applyMgu )

-- state needed for algorithm W
newtype PTState = PTState {
        varCount :: Int
    }
type PTMonad = State PTState

-- returns a fresh type var in algorithm W
getNewTypeVar :: PTMonad String
getNewTypeVar = do
    count <- gets varCount
    modify (\s -> s {varCount = count + 1})
    return $ 'a' : show count

-- implementation of algorithm W, this builds a list of terms to be unified in the next step
pt :: Context -> Term -> Type -> PTMonad Unifier
pt gamma (Var x) alpha = return [(alpha, gamma Map.! x)]
pt gamma (Abs x t) alpha = do
    a <- TypeVar <$> getNewTypeVar
    b <- TypeVar <$> getNewTypeVar
    result <- pt (Map.insert x a gamma) t b
    return $ result ++ [(Arrow a b, alpha)]
pt gamma (App t s) alpha = do
    a <- TypeVar <$> getNewTypeVar
    result1 <- pt gamma t (Arrow a alpha)
    result2 <- pt gamma s a
    return $ result1 ++ result2

-- infer the type of a term in context (if any)
inferType :: Context -> Term -> Maybe TermInContext
inferType ctx term = do
    let unificationProblem = evalState (pt ctx term (TypeVar "a0")) (PTState 1)
    mgu <- unify unificationProblem
    let principalType = applyMgu (TypeVar "a0") mgu
    return $ TermInContext ctx term principalType