#lang racket

(require "unify.rkt")

(require rackunit)

(define (uc->v x)
  (define-values (x1 env) (upper-case-symbols->variables x))
  x1)

(define-syntax-rule (check-match expected pat)
  (check-pred (lambda (v) (match v [pat #t] [_ #f])) expected))

(check-equal? (list->set (set-map (variables-in (uc->v '(a b c))) variable-info)) (set))
(check-equal? (list->set (set-map (variables-in (uc->v '(a B B))) variable-info)) (set 'B))
(check-equal? (list->set (set-map (variables-in (uc->v '(a B C))) variable-info)) (set 'B 'C))

(check-equal? (unify/vars '(a B c) '(a b c)) '((B . b)))
(check-equal? (unify/vars '(a B B) '(a b c)) #f)
(check-equal? (unify/vars '(a B B) '(a b B)) '((B . b)))

(check-match (unify/vars '(a b C) '(a b C)) `((C . ,(variable 'C))))

(check-equal? (unify/vars '#(A 2 3) '#(1 2 B)) '((B . 3) (A . 1)))

(check-equal? (unify-match/vars '#(A 2 3) '#(1 2 B)) '((A . 1)))
(check-equal? (unify-match/vars '#(1 2 B) '#(A 2 3)) '((B . 3)))

(check-equal? (unify-match/vars '(a B B) '(a b C)) '((B . b)))
(check-equal? (unify-match/vars '(a b C) '(a B B)) '((C . b)))

(check-equal? (unify-match/vars '(a C C) '(a b c)) #f)
(check-equal? (unify-match/vars '(a b c) '(a C C)) #f)
(check-match (unify-match/vars '(a b C) '(a D B)) `((C . ,(variable 'B))))
(check-match (unify-match/vars '(a C C) '(a D B)) `((C . ,(variable 'D))))

(check-equal? (mgu-canonical (uc->v '(a b C)) (uc->v '(a D B))) `(a b ,(canonical-variable 0)))
(check-equal? (canonicalize (uc->v `(A (B A) D)))
	      (canonicalize (uc->v `(E (F E) H))))

(check-equal? (upper-case-symbols->canonical '(A (B A) D))
	      (list (canonical-variable 0)
		    (list (canonical-variable 1) (canonical-variable 0))
		    (canonical-variable 2)))

(let* ((left (uc->v `(A (b A) D)))
       (right (freshen left)))
  (check-true (not (eq? (car left) (car right))))
  (check-eq? (variable-info (car left)) (variable-info (car right)))
  (check-equal? (canonicalize left) (list (canonical-variable 0)
					  (list 'b (canonical-variable 0))
					  (canonical-variable 1)))
  (check-equal? (canonicalize left) (canonicalize right)))

(struct x (y) #:transparent)
(let ((A (variable 'A)))
  (check-equal? (unify (x A) (x 123)) `((,A . 123))))

(struct x1 (y z) #:prefab)
(check-equal? (unify/vars (x1 'A 'a) (x1 'b 'a)) '((A . b)))
(let ((A (variable 'A)))
  (check-equal? (apply-subst `((,A . b)) (x1 A 'a)) (x1 'b 'a)))