;;
;; Unification
;; - Main program is originally written by Yuichi Nishiwaki
;; - Utity functions are originally written by Momoko Hattori
;;

(define $term
  (matcher
    {[<var $> integer
      {[<Var $i> {i}]
       [_ {}]}]
     [<compound $ $> [string (list term)]
      {[<Compound $s $l> {[s l]}]
       [_ {}]}]
     [<unify ,$t $> something
      {[$s (match (unify t s) (maybe something)
             {[(just $σ) {σ}]
              [(nothing) {}]})]}]
     [<subterm $ $> [term something]
      {[$s (subterm s)]}]
     [$ something
      {[$tgt {tgt}]}]}))

(define $var (lambda [$n] <Var n>))

(define $app
  (cambda $xs
    (match xs (list something)
      {[<cons $x $xs> <Compound x xs>]})))

(define $occur
  (pattern-function [$v]
    (| <var v>
       <compound _ <join _ <cons (occur v) _>>>)))

(define $fv
  (match-all-lambda [term]
    {[(occur $v) v]}))

(define $tsubst
  (match-lambda [something term]
    {[[$σ <var $n>]
       (match σ (multiset [integer term])
         {[<cons [,n $t] _> t]
          [_ <Var n>]})]
     [[$σ <compound $f $xs>]
       <Compound f (map (tsubst σ $) xs)>]}))

(define $unify
  (match-lambda (unordered-pair term)
    {[[<var $x> <var ,x>]
      (Just {})]
     [[<var $x> (& $t !(occur ,x))]
      (Just {[x t]})]
     [[<compound $f <nil>> <compound ,f <nil>>]
      (Just {})]
     [[<compound $f $xs> <compound ,f $ys>]
      (unify-list xs ys)]
     [_ Nothing]}))

(define $unify-list
  (match-lambda [(list term) (list term)]
    {[[<nil> <nil>] (Just {})]
     [[<cons $x $xs> <cons $y $ys>]
      (match (unify x y) (maybe something)
       {[(nothing) Nothing]
        [(just $σ1)
         (match (unify-list (map (tsubst σ1 $) xs) (map (tsubst σ1 $) ys)) (maybe something)
           {[(nothing) Nothing]
            [(just $σ2) (Just {@σ1 @σ2})]})]})]
     [_ Nothing]}))

;;
;; Utility for tests
;;

; variables
(define $x (var 0))
(define $y (var 1))
(define $z (var 2))
(define $w (var 3))

; constants
(define $a (app "a"))
(define $b (app "b"))
(define $c (app "c"))
(define $d (app "d"))

; function
(define $f "f")
(define $g "g")
(define $h "h")

(define $show-σ
  (lambda [$σ]
    (S.concat {"{" (show-σ' σ) "}"})))

(define $show-σ'
  (match-lambda (list [something something])
    {[<nil> ""]
     [<cons [$v $t] <nil>>
      (S.concat {"[" (show-var v) ", " (show-term t) "]"})]
     [<cons [$v $t] $σ>
      (S.concat {"[" (show-var v) ", " (show-term t) "], " (show-σ' σ)})]}))

(define $show-var
  (match-lambda integer
    {[,0 "x"]
     [,1 "y"]
     [,2 "z"]
     [,3 "w"]
     }))

(define $show-term
  (match-lambda term
    {[<var ,0> "x"]
     [<var ,1> "y"]
     [<var ,2> "z"]
     [<var ,3> "w"]
     [<var $x> (S.concat {"x" (show x)})]
     [<compound $f ,{}> f]
     [<compound ,"+" <cons (& <compound ,"+" _> $x) <cons $y <nil>>>>
       (S.concat {"(" (show-term x) ") + " (show-term y)})]
     [<compound ,"+" <cons $x <cons $y <nil>>>>
       (S.concat {(show-term x) " + " (show-term y)})]
     [<compound ,"*" <cons (& <compound ,"*" _> $x) <cons $y <nil>>>>
       (S.concat {"(" (show-term x) ") * " (show-term y)})]
     [<compound ,"*" <cons $x <cons $y <nil>>>>
       (S.concat {(show-term x) " * " (show-term y)})]
     [<compound $f $xs>
       (S.concat {f "(" (S.intercalate ", " (map show-term xs)) ")"})]
    }))

;;
;; Test
;;

(show-σ (car (unify (app "+" a b) x)))
; "{[x, a + b]}"

(show-σ (car (unify x (app "+" y z))))
; "{[x, y + z]}"

(show-σ (car (unify (app f x (app g y z) (app h x)) (app f a w y))))
; "{[x, a], [w, g(y, z)], [y, h(a)]}"