Gensym.hs

{- | Source of fresh names.
-}

module Gensym where

data T s v = Gensym
  { gState :: s
  , gStep  :: s -> (v, s)
  }

fromList :: [v] -> T [v] v
fromList vs = make vs \(v : vs) -> (v, vs)

make :: s -> (s -> (v, s)) -> T s v
make = Gensym

run :: T s v -> (v, T s v)
run (Gensym s step) =
  let (v, s') = step s
  in  (v, Gensym s' step)

instance Show v => Show (T s v) where
  show gen = do
    let (v, _) = run gen
    show v

Subst.hs

{- | Type variable substitution.
-}
module Subst where

import Data.Map qualified as Map
import Data.Map (Map)
import Data.Set qualified as Set
import Data.Set (Set)
import Data.List (intercalate)

import Term

{- | Substitution is implemented as a map. It is not left-biased, and there's
     no composition - but you must check for cyclic terms before adding,
     because application of a substitution is "apply until no variable is known".
-}
data Subst t v = Subst
  { getSubst :: Map v (Term t v)
  }

instance (Show v, Show (t (Term t v))) => Show (Subst t v) where
  show (Subst sub) =
    let list = Map.toList sub
        pairs = map showPair list
    in "[" <> intercalate ", " pairs <> "]"
    where
      showPair (v, t) = show v <> ": " <> show t

{- | Empty substitution list.
-}
empty :: Ord v => Subst t v
empty = Subst mempty

{- | Add a substitution, unguarded.
-}
add :: Ord v => v -> Term t v -> Subst t v -> Subst t v
add v t (Subst s) = Subst (Map.insert v t s)

{- | Add substitutions, unguarded.
-}
addAll :: Ord v => [(v, Term t v)] -> Subst t v -> Subst t v
addAll d (Subst s) = Subst (Map.fromList d <> s)

{- | Recall a possible substitution for a type variable.
-}
get :: Ord v => v -> Subst t v -> Maybe (Term t v)
get v (Subst s) = Map.lookup v s

{- | Check if term is cyclic.
-}
occurs :: HasVars t v f => v -> f -> Bool
occurs v f = Set.member v (freeVars f)

{- | A property of containing type variables and ability to substitute them.
-}
class Ord v => HasVars t v f | f -> t v where
  freeVars :: f -> Set v
  apply    :: Subst t v -> f -> f

{- | A type is substitutable.

     Note again that if substitutions produces a variable anywhere in output,
     this variable will be searched and substituted as well.
-}
instance (Ord v, Foldable t, Functor t) => HasVars t v (Term t v) where
  freeVars = \case
    Var  v -> Set.singleton v
    Term t -> foldMap freeVars t

  apply sub = go
    where
      go = \case
        Var  v -> maybe (Var v) go (get v sub)
        Term t -> Term (fmap go t)

Term.hs

{- | Type skeleton.
-}

module Term where

{- | The type.
-}
data Term t v
  = Var v                -- ^ Type variable.
  | Term (t (Term t v))  -- ^ Any other structure.

instance (Show v, Show (t (Term t v))) => Show (Term t v) where
  show = \case
    Var  v -> show v
    Term t -> "(" <> show t <> ")"

Unifiable.hs

module Unifiable where

import Data.Traversable (for)
import GHC.Generics

{- | Match two types and produce a unification plan for its parts.

     Preferred way to get an instance is to derive with @anyclass@.
-}
class Traversable t => Unifiable t where
  match :: t a -> t a -> Maybe (t (Either a (a, a)))

  default
    match :: (Generic1 t, Unifiable (Rep1 t)) => t a -> t a -> Maybe (t (Either a (a, a)))
  match a b = to1 <$> match (from1 a) (from1 b)

instance Unifiable V1 where
  match a _ = return $ Left <$> a

instance Unifiable U1 where
  match a _ = return $ Left <$> a

instance Unifiable Par1 where
  match (Par1 a) (Par1 b) = return $ Par1 $ Right (a, b)

instance (Unifiable f) => Unifiable (Rec1 f) where
  match (Rec1 a) (Rec1 b) = Rec1 <$> match a b

instance (Eq c) => Unifiable (K1 i c) where
  match (K1 a) (K1 b)
    | a == b    = return (K1 a)
    | otherwise = Nothing

instance (Unifiable f) => Unifiable (M1 i c f) where
  match (M1 a) (M1 b) = M1 <$> match a b

instance (Unifiable f, Unifiable g) => Unifiable (f :+: g) where
  match a b = case (a, b) of
    (L1 q, L1 w) -> L1 <$> match q w
    (R1 q, R1 w) -> R1 <$> match q w
    _            -> Nothing

instance (Unifiable f, Unifiable g) => Unifiable (f :*: g) where
  match (a :*: c) (b :*: d) = pure (:*:) <*> match a b <*> match c d

instance (Unifiable f, Unifiable g) => Unifiable (f :.: g) where
  match (Comp1 a) (Comp1 b) = do
    res' <- match a b
    res'' <- for res' \case
      Left ga -> return $ Left <$> ga
      Right (ga, gb) -> match ga gb
    return $ Comp1 res''

Unify.hs

module Unify where

import Control.Monad.State (StateT(..))
import Control.Monad (when)
import Data.Function ((&))
import GHC.Generics (Generic1)

import Term
import Subst
import Gensym qualified
import Unifiable

data UnificationError t v
  = Cyclic v (Term t v)
  | Mismatch (Term t v) (Term t v)
  | Unifying (Term t v) (Term t v) (UnificationError t v)

toString :: (Show v, Show (t (Term t v))) => UnificationError t v -> String
toString = \case
  Cyclic v t -> unlines
    [ "found cyclic type"
    , "    " <> show v
    , "  ~"
    , "    " <> show t
    ]

  Mismatch l r -> unlines
    [ "found mismatch"
    , "    " <> show l
    , "  and"
    , "    " <> show r
    ]

  Unifying l r err -> unlines
    [ "while unifying"
    , "    " <> show l
    , "  and"
    , "    " <> show r
    , ""
    ] <> toString err

{- | Add a substitution, guarded. If neccessary, unify with existing knowledge.
-}
addSubst
  :: (Ord v, Unifiable t)
  => v
  -> Term t v
  -> (Subst t v, Gensym.T s v)
  -> Either (UnificationError t v) (Subst t v, Gensym.T s v)
addSubst v t (sub, gen) = do
  let fullT = apply sub t

  when (occurs v fullT) do
    Left (Cyclic v fullT)

  (t'', (sub', gen')) <- case Subst.get v sub of  -- do we know smth?
    Just t' -> unify (t, t') (sub, gen)             -- unify with it
    Nothing -> return (t, (sub, gen))               -- add as-is

  return (Subst.add v t'' sub', gen')

unify
  :: (Ord v, Unifiable t)
  => (Term t v, Term t v)
  -> (Subst t v, Gensym.T s v)
  -> Either (UnificationError t v) (Term t v, (Subst t v, Gensym.T s v))
unify (l0, r0) (sub, gen) = do
  let l = apply sub l0  -- "prune"
  let r = apply sub r0  -- "prune"
  case (l, r) of

    -- In @a ~ b@ case, transform to @a ~ t0, b ~ t0@, so neither of @a@ or @b@
    -- appear on the right side.
    --
    (Var lv, Var rv) -> do
      let (v, gen') = Gensym.run gen
      let sub'      = Subst.addAll [(lv, Var v), (rv, Var v)] sub
      return (Var v, (sub', gen'))

    (Var v, t) -> (t,) <$> addSubst v t (sub, gen)
    (t, Var v) -> (t,) <$> addSubst v t (sub, gen)

    (Term t, Term u) -> do
      case match t u of
        Just matcher -> do
          (t', (sub', gen')) <-
            unifying (l0, r0) do
              let unify' = either return (StateT . unify)
              runStateT (traverse unify' matcher) (sub, gen)

          return (Term t', (sub', gen'))

        Nothing -> do
          Left (Mismatch l r)
  where
    -- Catch inner unification error and wrap it with current pair of terms,
    -- so it is possible to see why @t72 ~ List t72@ is cyclic.
    --
    unifying :: (Term t v, Term t v) -> Either (UnificationError t v) a -> Either (UnificationError t v) a
    unifying (l, r) (Left  err) = Left (Unifying l r err)
    unifying  _     (Right res) = Right res

---- TESTING AREA --------------------------------------------------------------

data Type t
  = Int
  | String
  | Fun t t
  deriving stock (Show, Functor, Foldable, Traversable, Generic1)
  deriving anyclass (Unifiable)

sub0 :: Subst Type String
sub0 = Subst.empty & Subst.add "a" (Term (Fun (Var "b") (Var "b")))

gen0 :: Gensym.T [String] String
gen0 = Gensym.fromList ["t" <> show n | n <- [0 :: Int ..]]

pp :: Either
        (UnificationError Type String)
        (Term Type String, (Subst Type String, Gensym.T [String] String))
   -> IO ()
pp = putStrLn . either toString show

test :: IO ()
test = pp $ unify (Var "a", Var "b") (sub0, gen0)

test2 :: IO ()
test2 = pp $ unify (Var "a", Term Int) (sub0, gen0)

test3 :: IO ()
test3 = pp $ unify (Var "a", Term (Fun (Term Int) (Var "c"))) (sub0, gen0)

test4 :: IO ()
test4 = pp $ unify (Var "a", Term (Fun (Var "b") (Var "a"))) (sub0, gen0)