;; miniKanren type inferencer

(load "mkdefs.scm")

;; Use unification with occur check:
;; (== `(,x) x) fails.
(define == ==-check)

(define-syntax test
  (syntax-rules ()
    ((_ title tested-expression expected-result)
     (let* ((expected expected-result)
            (produced tested-expression))
       (if (equal? expected produced)
           (printf "~s works!\n" title)
           (error
             'test
             "Failed ~s: ~a~%Expected: ~a~%Computed: ~a"
             title 'tested-expression expected produced))))))

(define parse
  (lambda (e)
    (cond
      [(symbol? e) `(var ,e)]
      [(number? e) `(intc ,e)]
      [(boolean? e) `(boolc ,e)]
      [else
	(case (car e)
	  [(sub1) `(sub1 ,(parse (cadr e)))]
	  [(+) `(+ ,(parse (cadr e)) ,(parse (caddr e)))]
	  [(zero?) `(zero? ,(parse (cadr e)))]
	  [(if) `(if ,(parse (cadr e)) ,(parse (caddr e)) ,(parse (cadddr e)))]
	  [(lambda) `(lambda ,(cadr e) ,(parse (caddr e)))]
	  [else `(app ,(parse (car e)) ,(parse (cadr e)))])])))

(define unparse
  (lambda (e)
    (if (var? e) e
      (case (car e)
	[(sub1) `(sub1 ,(unparse (cadr e)))]
	[(+) `(+ ,(unparse (cadr e)) ,(unparse (caddr e)))]
	[(var) e]
	[(intc) e]
	[(boolc) e]
	[(zero?) `(zero? ,(unparse (cadr e)))]
	[(if) `(if ,(unparse (cadr e))
		 ,(unparse (caddr e))
		 ,(unparse (cadddr e)))]
	[(lambda) `(lambda ,(cadr e) ,(unparse (caddr e)))]
	[(app) `(,(unparse (cadr e)) ,(unparse (caddr e)))]))))

(define !-
  (lambda (gamma exp type)
    (conde
      [(exist (n)
         (== `(intc ,n) exp)
         (== 'int type))]
      [(exist (b)
         (== `(boolc ,b) exp)
         (== 'bool type))]
      [(exist (e)
         (== `(zero? ,e) exp)
         (== 'bool type)
         (!- gamma e 'int))]
      [(exist (e)
         (== `(sub1 ,e) exp)
         (== 'int type)
         (!- gamma e 'int))]
      [(exist (n m)
         (== `(+ ,n ,m) exp)
         (== 'int type)
         (!- gamma n 'int)
         (!- gamma m 'int))]
      [(exist (t c a)
         (== `(if ,t ,c ,a) exp)
         (!- gamma t 'bool)
         (!- gamma c type)
         (!- gamma a type))])))

(define lookupo
  (lambda (x gamma type)
    (exist (y t rest)
      (== `((,y . ,t) . ,rest) gamma)
      (conde
        ((== x y) (== type t))
        ;; (=/= x y) is a disequality constraint, asserting that x and y can never be unified
        ((=/= x y) (lookupo x rest type))))))

(test "1"
  (run* (q) (!- '() '(intc 17) q))
  '(int))

(test "2"
  (run* (q) (!- '() '(zero? (intc 24)) q))
  '(bool))

(test "3"
  (run* (q) (!- '() '(zero? (sub1 (intc 24))) q))
  '(bool))

(test "4"
  (run* (q)
    (!- '() '(zero? (sub1 (sub1 (intc 18)))) q))
  '(bool))

(test "5"
  (run* (q)
    (!- '() (parse '(lambda (n) (if (zero? n) n n))) q))
  '((-> int int)))

(test "6"
  (run* (q)
    (!- '() (parse '((lambda (n) (zero? n)) 5)) q))  
  '(bool))

(test "7"
  (run* (q)
    (!- '() '(if (zero? (intc 24)) (intc 3) (intc 4)) q))
  '(int))

(test "8"
  (run* (q)
    (!- '() '(if (zero? (intc 24)) (zero? (intc 3)) (zero? (intc 4))) q))
  '(bool))

(test "9"
  (run* (q)
    (!- '() '(lambda (x) (sub1 (var x))) q))
  '((-> int int)))

(test "10"
  (run* (q)
    (!- '() '(lambda (a) (lambda (x) (+ (var a) (var x)))) q))
  '((-> int (-> int int))))

(test "11"
  (run* (q)
    (!- '() (parse '(lambda (f)
                      (lambda (x)
                        ((f x) x))))
      q))
  '((->
     (-> _.0 (-> _.0 _.1))
     (-> _.0 _.1))))

(test "12"
  (run* (q)
    (!- '() (parse '(sub1 (sub1 (sub1 6)))) q))
  '(int))

(test "13"
  (run 1 (q)
    (exist (t)
      (!- '() (parse '(lambda (f) (f f))) t)))
  '())

(test "24"
  (let ([v (run 20 (q)
             (exist (lam a b)
               (!- '() `(app (,lam (,a) ,b) (intc 5)) 'int)
               (== `(,lam (,a) ,b) q)))])
    (pretty-print v)
    (length v))
  20)

(test "25"
  (let ([v (run 30 (q)
             (!- '() q 'int))])
    (pretty-print v)
    (length v))
  30)

(test "26"
  (let ([v (run 30 (q)
             (!- '() q '(-> int int)))])
    (pretty-print v)
    (length v))
  30)

(test "27"
  (let ([v (run 30 (q)
             (!- '() q '(-> bool int)))])
    (pretty-print v)
    (length v))
  30)

(test "28"
  (let ([v (run 30 (q)
             (!- '() q '(-> int (-> int int))))])
    (pretty-print v)
    (length v))
  30)

(test "29"
  (let ([v (run 100 (q)
             (exist (e t)
               (!- '() e t)
               (== `(,e ,t) q)))])
    (pretty-print v)
    (length v))
  100)

(test "30"
  (let ([v (run 100 (q)
             (exist (g e t)
               (!- g e t)
               (== `(,g ,e ,t) q)))])
    (pretty-print v)
    (length v))
  100)

(test "31"
  (length
    (run 100 (q)
      (exist (g v)
        (!- g `(var ,v) 'int)
        (== `(,g ,v) q))))
  100)