--  ported from https://www.cs.cornell.edu/courses/cs3110/2011sp/Lectures/lec26-type-inference/type-inference.htm
module Minilang.Lambda.Unify where

import Data.Text (Text)

type Identifier = Text

data Type
  = TyVar Text
  | Type :-> Type
  deriving (Show, Eq)

infixr 5 :->

-- invariant: no identifier on the left hand-side can appear in
-- an earlier term in the list, ie. the list is in "dependency order"
-- and forms a DAG.
type Substitution = [(Identifier, Type)]

occurs :: Identifier -> Type -> Bool
occurs x (TyVar y) = x == y
occurs x (f :-> g) = occurs x f || occurs x g

subst :: Type -> Identifier -> Type -> Type
subst s x t@(TyVar y)
  | x == y = s
  | otherwise = t
subst s x (f :-> g) = subst s x f :-> subst s x g

apply :: Substitution -> Type -> Type
apply subs t =
  foldr (\(x, s) -> subst s x) t subs

data UnifyError
  = MismatchHead Text Text [Type] [Type]
  | Circularity Identifier Type
  deriving (Show, Eq)

unifyOne :: Type -> Type -> Either UnifyError Substitution
unifyOne (TyVar x) (TyVar y) = if x == y then Right [] else Right [(x, TyVar y)]
unifyOne (f :-> g) (f' :-> g') =
  unify [(f, f'), (g, g')]
unifyOne (TyVar x) t =
  if occurs x t
    then Left $ Circularity x t
    else Right [(x, t)]
unifyOne t (TyVar x) =
  if occurs x t
    then Left $ Circularity x t
    else Right [(x, t)]

unify :: [(Type, Type)] -> Either UnifyError Substitution
unify [] = Right []
unify ((s, t) : rest) = do
  sub2 <- unify rest
  sub1 <- unifyOne (apply sub2 s) (apply sub2 t)
  return $ sub1 ++ sub2