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

import Data.Text (Text)
import GHC.Generics (Generic)
import Data.Aeson (ToJSON, FromJSON)

type Identifier = Text

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

infixr 5 :->

pretty :: Type -> Text
pretty = \case
  TyVar x -> x
  f :-> g ->
    let left = case f of
          TyVar _ -> pretty f
          _ -> "(" <> pretty f <> ")"
        right = pretty g
     in left <> " -> " <> right

-- 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
  = Circularity Identifier Type
  deriving (Show, Eq)

render :: UnifyError -> Text
render (Circularity x t) =
  "Cannot unify type variable "
    <> x
    <> " with type "
    <> pretty t
    <> " due to circularity: The variable already occurs in the type."

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