module Unify where

import Data.Either
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer

data Type = Simple String |
            Variable String |
            Function Type Type
            deriving Show

instance Eq Type where
    (Simple s1) == (Simple s2)           = s1 == s2
    (Simple s)  == (Variable v)          = s  == v
    (Variable v) == (Simple s)           = v  == s
    (Variable v1) == (Variable v2)       = v1 == v2
    (Function t1 t2) == (Function s1 s2) = (t1 == t2) && (s1 == s2)
    _ == _                               = False

data Constraint = Equiv Type Type deriving Show
data Substitute = S Type Type     deriving Show

type Substitutes = [Substitute]

substitute :: Substitute -> [Constraint] -> [Constraint]
substitute _ []                 = []
substitute s ((Equiv t1 t2):xs) = (Equiv (substitute1 s t1) (substitute1 s t2)) : (substitute s xs)
  where
    substitute1 s@(S from to) (Function t1 t2) = Function (substitute1 s t1) (substitute1 s t2)
    substitute1 (S from to) t                = if from == t then to else t

unify :: [Constraint] -> Writer Substitutes [Constraint]
unify []                                             = return []
unify ((Equiv v@(Variable _) t):xs)                  = let s = (S v t)   in tell [s] >> (return $ substitute s xs)
unify ((Equiv t v@(Variable _)):xs)                  = let s = (S v t)   in tell [s] >> (return $ substitute s xs)
unify ((Equiv (Function t1 t2) (Function s1 s2)):xs) = return $ (Equiv t1 s1) : (Equiv t2 s2) : xs
unify c                                              = error ("Type check mismatch" ++ show c)


-- unifiable

sampleConstraints1 = [ Equiv (Variable "X") (Simple "Nat")
                     , Equiv (Variable "Y") (Function (Variable "X") (Variable "X"))
                     ]

sampleConstraints2 = [ Equiv (Function (Simple "Nat") (Simple "Nat")) (Function (Variable "X") (Variable "Y"))
                     ]

sampleConstraints3 = [ Equiv (Function (Variable "X") (Variable "Y")) (Function (Variable "Y") (Variable "Z"))
                     , Equiv (Variable "Z") (Function (Variable "U") (Variable "W"))
                     ]

sampleConstraints4 = []


-- not unifiable

sampleConstraints5 = [ Equiv (Function (Simple "Nat") (Variable "Y")) (Simple "Nat")]

sampleConstraints6 = [ Equiv (Variable "Y") (Function (Simple "Nat") (Variable "Y"))] -- TODO: must check free type variable