#lang racket

; This file contains a simple implementation of the congruence
; closure algorithm for deciding the satisfiability of a
; conjunction of literals in the theory of equality with uninterpreted functions.

; A node representing a subterm of a formula in the theory of
; uninterpreted functions with equality.
(struct node (id fn children [find #:mutable] [ccp #:mutable])
  #:methods gen:custom-write
  [(define (write-proc self port mode)
    (match self
      [(node id label children find ccp)
       (fprintf port "{~a label=~a children=~a find=~a ccp=~a}"
                id label (map node-id children) (node-id find) (set-map ccp node-id))]))])

; Returns a dictionary mapping each subterm in F to its node representation.
; The procedure assumes F is a list of literals of the form (= term1 term2) or (!= term1 term2).
(define (formula->graph F)
  (define g (make-hash))
  (define (term->node t)
    (hash-ref!
     g t
     (lambda ()
       (let ([n (match t
                  [(list fn child ...) (node t fn (map term->node child) #f (set))]
                  [name (node t name '() #f (set))])])
         (set-node-find! n n)
         (for ([c (node-children n)])
           (set-node-ccp! c (set-add (node-ccp c) n)))
         n))))
  (for ([lit F])
    (term->node (second lit))
    (term->node (third lit)))
  g)

; Returns the representative for node n.
(define (find n) 
  (let ([f (node-find n)])
    (if (equal? n f) n (find f))))

; Performs the union operation over the representatives of n1 and n2.
(define (union n1 n2)
  (let ([r1 (find n1)]
        [r2 (find n2)])
    (set-node-find! r1 r2)
    (set-node-ccp! r2 (set-union (node-ccp r2) (node-ccp r1)))
    (set-node-ccp! r1 (set))))

; Returns true if n1 and n2 have the label and their arguments
; are in the same congruence classes.
(define (congruent? n1 n2)
  (and (equal? (node-fn n1) (node-fn n2))
       (= (length (node-children n1)) (length (node-children n2)))
       (for/and ([c1 (node-children n1)] [c2 (node-children n2)])
         (equal? (find c1) (find c2)))))

; Merges the equivalence classes of n1 and n2.
(define (merge n1 n2)
  (let ([r1 (find n1)]
        [r2 (find n2)])
    (unless (equal? r1 r2)
      (let ([p1 (node-ccp r1)]
            [p2 (node-ccp r2)])
        (union r1 r2)
        (for* ([t1 p1][t2 p2]
               #:when (and (not (equal? (find t1) (find t2)))
                           (congruent? t1 t2)))
          (merge t1 t2))))))

; Returns a model if F is satisfiable. Otherwise returns #f.
; The procedure assumes F is a list of literals of the form
; (= term1 term2) or (!= term1 term).
(define (decide F)
  (define g (formula->graph F))
  (define (get t) (hash-ref g t))
  (for ([lit F]) ; Compute the congruence closure over F's positive literals.
    (match lit
      [(list '= s t) (merge (get s) (get t))]
      [_ (void)]))
  (and
   (for/and ([lit F]) 
     (match lit
       [(list '!= s t) (not (equal? (find (get s)) (find (get t))))]
       [_ #t]))
   (ccc->model g))) ; model if all negative literals are in different congruence classes.

; Takes dictionary representing congruence closure classes and constructs a model over a
; universe that is a finite set of naturals {0, ..., N}
(define (ccc->model g)
  (define H
    (for/hash ([(k v) g]) ; terms -> congruence representatives 
      (values k (node-id (find v))))) 
  (define U               ; Herbrand universe -> {0, ..., N}
    (for/hash ([(e i) (in-indexed (apply set (hash-values H)))]) 
      (values e i)))
  (for/hash ([(t v) H])   ; Herbrand model -> model over U
    (values
     (match t
       [(list fn args ...) (cons fn (for/list ([a args]) (hash-ref U (hash-ref H a))))]
       [_ t])
     (hash-ref U v))))

; Interprets the T= formula F with respect to the interpretation
; I given as a dictionary. The universe of interpretation is implicit
; in the range of I. The interpretation of functions is enumerated explicitly.
; For example, if the universe is {e1, e2}, then the meaning of a function f over
; this universe is given as the mappings (f e1) -> ... and (f e2) -> ...
(define (interpret F I)
  (define cache (make-hash))
  (define (interpret-term t)
    (hash-ref!
     cache t 
     (lambda ()
       (match t
         [(list fn args ...) (hash-ref I (cons fn (map interpret-term args)))]
         [_ (hash-ref I t)]))))
  (for/and ([lit F])
    (match lit
      [(list '= left right)  (equal? (interpret-term left) (interpret-term right))]
      [(list '!= left right) (not (equal? (interpret-term left) (interpret-term right)))])))

(define F0 '((= (f a b) a) (!= (f (f a b) b) a)))
(define F1 '((= (f (f (f a))) a) (= (f (f (f (f (f a))))) a) (!= (f a) a)))
(define F2 '((= (f x) (f y)) (!= x y)))
(define F3 '((= (f (g x)) (g (f x))) (= (f (g (f y))) x) (= (f y) x) (!= (g x) x)))

(printf "F0: ~a is unsat, so (decide F0) = ~a\n" F0 (decide F0))
(printf "F1: ~a is unsat, so (decide F1) = ~a\n" F1 (decide F1))

(printf "F2: ~a is sat, so (decide F2) =\n ~a\n" F2 (pretty-format (decide F2)))
(printf " Interpret F2 w.r.t. the model: ~a\n" (interpret F2 (decide F2))) 

(printf "F3: ~a is sat, so (decide F3) =\n ~a\n" F3 (pretty-format (decide F3)))
(printf " Interpret F3 w.r.t. the model: ~a\n" (interpret F3 (decide F3)))