Unification of Trees (doodle made while reading a tutorial).

Tree.hs

module Tree where

-- https://crypto.stanford.edu/~blynn/lambda/hm.html
-- http://hackage.haskell.org/package/base-4.12.0.0/docs/src/Control.Arrow.html#Arrow

import Control.Arrow ((***))
import Control.Monad (join)
import Control.Monad.State
import Text.Parsec hiding (State)
import Text.Parsec.String (Parser)

data Tree
  = Var Char
  | Leaf Integer
  | Branch Tree Tree
  deriving Show

type Constraint = (Char, Tree)

gather :: Tree -> Tree -> Maybe [(Tree, Tree)]
gather (Leaf x) (Leaf y)
  | x == y    = Just []
  | otherwise = Nothing
gather (Branch x y) (Branch x' y') = do
  left  <- gather x x'
  right <- gather y y'
  pure $ left ++ right
gather x@(Var _) y = Just [(x, y)]
gather x y@(Var _) = gather y x
gather _ _         = Nothing

unify :: [(Tree, Tree)] -> State [Constraint] (Maybe [Constraint])
unify [] = Just <$> get
unify ((Branch x y, Branch x' y'):rest) =
  unify $ (x, x'):(y, y'):rest
unify ((Leaf x, Leaf y):rest)
  | x == y = unify rest
unify ((Var x, Var y):rest)
  | x == y = unify rest
unify ((Var x, y):rest) =
  if occurs y
    then pure Nothing
    else modify ((x, y):) >> unify (join (***) (substitute x y) <$> rest)
  where occurs (Branch x' y')     = occurs x' || occurs y'
        occurs (Var x') | x == x' = True
        occurs _                  = False
unify ((x, y@(Var _)):rest) = unify ((y, x):rest)
unify _                     = pure Nothing

substitute :: Char -> Tree -> Tree -> Tree
substitute x y (Var x')
  | x == x' = y
substitute x y (Branch x' y') = Branch (substitute x y x') (substitute x y y')
substitute _ _ x              = x

solve :: Tree -> Tree -> Maybe [Constraint]
solve x y = gather x y >>= (flip evalState []) . unify

-- Parser

parseTree :: String -> Tree
parseTree s =
  case parse (whitespace *> tree <* eof) "" s of
    Left e  -> error $ show e
    Right x -> x

tree :: Parser Tree
tree = variable
   <|> leaf
   <|> branch

variable :: Parser Tree
variable = Var <$> letter <* whitespace

leaf :: Parser Tree
leaf = do
  sign   <- option ' ' (char '-')
  digits <- many1 digit
  whitespace
  pure $ Leaf $ read (sign:digits)

branch :: Parser Tree
branch = parens $ do
  x <- tree
  y <- tree
  pure $ Branch x y

parens :: Parser a -> Parser a
parens p = between open close p
  where open  = char '(' <* whitespace
        close = char ')' <* whitespace

whitespace :: Parser ()
whitespace = skipMany space

-- Main

main :: IO ()
main = do
  x <- getLine
  y <- getLine
  print $ solve (parseTree x) (parseTree y)