From 9b09365fd5db7925e98856af3993bcbde47c5d90 Mon Sep 17 00:00:00 2001
From: Sorawee Porncharoenwase <sorawee.pwase@gmail.com>
Date: Thu, 10 Nov 2022 11:42:41 -0800
Subject: [PATCH] optimize symbolic heap with hasheq

---
 rosette/base/core/term-cache.rkt | 144 +++++++++++++++++++++++
 rosette/base/core/term.rkt       |  59 ++++------
 test/base/smt.rkt                |  84 +++++++++++++
 test/base/term-cache.rkt         | 194 +++++++++++++++++++++++++++++++
 4 files changed, 447 insertions(+), 34 deletions(-)
 create mode 100644 rosette/base/core/term-cache.rkt
 create mode 100644 test/base/smt.rkt
 create mode 100644 test/base/term-cache.rkt

diff --git a/rosette/base/core/term-cache.rkt b/rosette/base/core/term-cache.rkt
new file mode 100644
index 00000000..7eae3e3c
--- /dev/null
+++ b/rosette/base/core/term-cache.rkt
@@ -0,0 +1,144 @@
+#lang racket/base
+
+(provide prop:term-cachable
+         term-cache-weak?
+         make-term-cache
+         make-weak-term-cache
+         term-cache-ref
+         term-cache-set!
+         term-cache-count
+         term-cache-copy
+         term-cache-copy-clear
+         term-cache->hash)
+
+(require racket/match)
+
+(struct term-cache (nullary unary binary ternary nary))
+
+(define-values (prop:term-cachable term-cachable? term-cachable-ref)
+  (make-struct-type-property 'term-cachable))
+
+(define (term-cache-weak? x)
+  (and (hash? x) (hash-ephemeron? x)))
+
+(define (make-term-cache [assocs '()])
+  (define out
+    (term-cache (make-hasheq)
+                (make-hasheq)
+                (make-hasheq)
+                (make-hasheq)
+                (make-hash)))
+  (for ([pair (in-list assocs)])
+    (term-cache-set! out (car pair) (cdr pair)))
+  out)
+
+(define make-weak-term-cache make-ephemeron-hash)
+
+(define (term-like? x)
+  (or (term-cachable? x) (syntax? x) (fixnum? x) (boolean? x) (procedure? x)))
+
+(define (term-cache-ref h k default)
+  (define (proc) (if (procedure? default) (default) default))
+  (match h
+    [(term-cache nullary unary binary ternary nary)
+     (match k
+       [_
+        #:when (term-like? k)
+        (hash-ref nullary k proc)]
+       [(list op a)
+        #:when (and (term-like? op) (term-like? a))
+        (match (hash-ref unary op #f)
+          [#f (proc)]
+          [h (hash-ref h a proc)])]
+       [(list op a b)
+        #:when (and (term-like? op) (term-like? a) (term-like? b))
+        (match (hash-ref binary op #f)
+          [#f (proc)]
+          [h
+           (match (hash-ref h a #f)
+             [#f (proc)]
+             [h (hash-ref h b proc)])])]
+       [(list op a b c)
+        #:when (and (term-like? op) (term-like? a) (term-like? b) (term-like? c))
+        (match (hash-ref ternary op #f)
+          [#f (proc)]
+          [h
+           (match (hash-ref h a #f)
+             [#f (proc)]
+             [h
+              (match (hash-ref h b #f)
+                [#f (proc)]
+                [h (hash-ref h c proc)])])])]
+       [_ (hash-ref nary k proc)])]
+    [_ (hash-ref h k proc)]))
+
+(define (term-cache-set! h k v)
+  (match h
+    [(term-cache nullary unary binary ternary nary)
+     (match k
+       [_
+        #:when (term-like? k)
+        (hash-set! nullary k v)]
+       [(list op a)
+        #:when (and (term-like? op) (term-like? a))
+        (hash-set! (hash-ref! unary op make-hasheq) a v)]
+       [(list op a b)
+        #:when (and (term-like? op) (term-like? a) (term-like? b))
+        (hash-set! (hash-ref! (hash-ref! binary op make-hasheq) a make-hasheq) b v)]
+       [(list op a b c)
+        #:when (and (term-like? op) (term-like? a) (term-like? b) (term-like? c))
+        (hash-set! (hash-ref! (hash-ref! (hash-ref! ternary op make-hasheq) a make-hasheq) b make-hasheq) c v)]
+       [_ (hash-set! nary k v)])]
+    [_ (hash-set! h k v)]))
+
+(define (term-cache-count h)
+  (match h
+    [(term-cache nullary unary binary ternary nary)
+     (+ (hash-count nullary)
+        (for/sum ([(_ h) (in-hash unary)])
+          (hash-count h))
+        (for/sum ([(_ h) (in-hash binary)])
+          (for/sum ([(_ h) (in-hash h)])
+            (hash-count h)))
+        (for/sum ([(_ h) (in-hash ternary)])
+          (for/sum ([(_ h) (in-hash h)])
+            (for/sum ([(_ h) (in-hash h)])
+              (hash-count h))))
+        (hash-count nary))]
+    [_ (hash-count h)]))
+
+(define (term-cache-copy h)
+  (match h
+    [(term-cache nullary unary binary ternary nary)
+     (term-cache (hash-copy nullary)
+                 (hash-copy unary)
+                 (hash-copy binary)
+                 (hash-copy ternary)
+                 (hash-copy nary))]
+    [_ (hash-copy h)]))
+
+(define (term-cache-copy-clear h)
+  (cond
+    [(term-cache-weak? h) (make-weak-term-cache)]
+    [else (make-term-cache)]))
+
+(define (term-cache->hash h term-val)
+  (match h
+    [(term-cache nullary unary binary ternary nary)
+     (define h* (hash-copy nary))
+     (for ([v (in-hash-values nullary)])
+       (hash-set! h* (term-val v) v))
+     (for ([h (in-hash-values unary)])
+       (for ([v (in-hash-values h)])
+         (hash-set! h* (term-val v) v)))
+     (for ([h (in-hash-values binary)])
+       (for ([h (in-hash-values h)])
+         (for ([v (in-hash-values h)])
+           (hash-set! h* (term-val v) v))))
+     (for ([h (in-hash-values ternary)])
+       (for ([h (in-hash-values h)])
+         (for ([h (in-hash-values h)])
+           (for ([v (in-hash-values h)])
+             (hash-set! h* (term-val v) v)))))
+     h*]
+    [_ h]))
diff --git a/rosette/base/core/term.rkt b/rosette/base/core/term.rkt
index 19c62033..a729f6da 100644
--- a/rosette/base/core/term.rkt
+++ b/rosette/base/core/term.rkt
@@ -2,7 +2,8 @@
 
 (require racket/syntax (for-syntax racket racket/syntax syntax/parse)
          racket/generic syntax/parse 
-         "type.rkt" "reporter.rkt")
+         "type.rkt" "reporter.rkt"
+         "term-cache.rkt")
 
 (provide
  terms terms-count terms-ref with-terms clear-terms! gc-terms!
@@ -10,6 +11,7 @@
  (rename-out [a-term term] [an-expression expression] [a-constant constant] [term-ord term-id]) 
  term-type term<? sublist? @app
  define-operator operator? operator-unsafe
+ prop:term-cachable
  (all-from-out "type.rkt"))
 
 #|-----------------------------------------------------------------------------------|#
@@ -23,7 +25,7 @@
 ;; Initialize with #f so that the hash table cooperates with garbage collector.
 ;; See #247
 (define current-terms (make-parameter #f))
-(current-terms (make-hash))
+(current-terms (make-term-cache))
 
 (define current-index (make-parameter 0))
 
@@ -31,52 +33,42 @@
 ; it clears all terms reachable from the given set of leaf terms.
 (define (clear-terms! [terms #f])
   (if (false? terms)
-      (hash-clear! (current-terms))
-      (let ([cache (current-terms)]
+      (current-terms (term-cache-copy-clear (current-terms)))
+      (let ([cache (term-cache->hash (current-terms) term-val)]
             [evicted (list->mutable-set terms)])
-        (for ([t terms])
+        (for ([t (in-list terms)])
           (hash-remove! cache (term-val t)))
         (let loop ()
-          (define delta  
-            (for/list ([(k t) cache] #:when (and (list? k) (for/or ([c k]) (set-member? evicted c))))
+          (define delta
+            (for/list ([(k t) (in-hash cache)] #:when (and (list? k) (for/or ([c (in-list k)]) (set-member? evicted c))))
               t))
           (unless (null? delta)
-            (for ([t delta])
+            (for ([t (in-list delta)])
               (hash-remove! cache (term-val t))
               (set-add! evicted t))
-            (loop))))))
+            (loop)))
+        (unless (term-cache-weak? (current-terms))
+          (current-terms (make-term-cache (hash->list cache)))))))
 
 ; Sets the current term cache to a garbage-collected (weak) hash.
 ; The setting preserves all reachable terms from (current-terms).
 (define (gc-terms!)
-  (unless (hash-weak? (current-terms)) ; Already a weak hash.
-    (define cache
-      (impersonate-hash
-       (make-weak-hash)
-       (lambda (h k)
-         (values k (lambda (h k e) (ephemeron-value e #f))))
-       (lambda (h k v)
-         (values k (make-ephemeron k v)))
-       (lambda (h k) k)
-       (lambda (h k) k)
-       hash-clear!))
-    (for ([(k v) (current-terms)])
-      (hash-set! cache k v))
-    (current-terms cache)))
+  (unless (term-cache-weak? (current-terms)) ; Already a weak hash.
+    (current-terms (make-weak-term-cache (hash->list (term-cache->hash (current-terms) term-val))))))
 
 ; Returns the term from current-terms that has the given contents. If
 ; no such term exists, failure-result is returned, unless it is a procedure.
 ; If failure-result is a procedure, it is called and its result is returned instead.
 (define (terms-ref contents [failure-result (lambda () (error 'terms-ref "no term for ~a" contents))])
-  (hash-ref (current-terms) contents failure-result))
+  (term-cache-ref (current-terms) contents failure-result))
 
 ; Returns a list of all terms in the current-term scache, in an unspecified order.
 (define (terms)
-  (hash-values (current-terms)))
+  (hash-values (term-cache->hash (current-terms) term-val)))
 
 ; Returns the size of the current-terms cache.
 (define (terms-count)
-  (hash-count (current-terms)))
+  (term-cache-count (current-terms)))
 
 ; Evaluates expr with (terms) set to terms-expr, returns the result, and
 ; restores (terms) to its old value. If terms-expr is not given, it defaults to
@@ -88,20 +80,17 @@
     [(_ expr)
      #'(let ([orig-terms (current-terms)])
          (parameterize ([current-terms #f])
-           (current-terms (hash-copy orig-terms))
+           (current-terms (term-cache-copy orig-terms))
            expr))]
     [(_ terms-expr expr)
      #'(let ([orig-terms (current-terms)])
          (parameterize ([current-terms #f])
-           (current-terms (hash-copy-clear orig-terms))
+           (current-terms (term-cache-copy-clear orig-terms))
            (let ([ts terms-expr]
                  [cache (current-terms)])
              (for ([t ts])
-               (hash-set! cache (term-val t) t))
+               (term-cache-set! cache (term-val t) t))
              expr)))]))
-           
-         
-    
 
 #|-----------------------------------------------------------------------------------|#
 ; The term structure defines a symbolic value, which can be a variable or an expression.
@@ -115,6 +104,7 @@
   (val                 ; (or/c any/c (cons/c function? (non-empty-listof any/c)))
    type                ; type?  
    ord)                ; integer?  
+  #:property prop:term-cachable #t
   #:methods gen:typed 
   [(define (get-type v) (term-type v))]
   #:property prop:custom-print-quotable 'never
@@ -133,8 +123,9 @@
 
 (define-syntax-rule (make-term term-constructor args type rest ...) 
   (let ([val args]
+        [the-terms (current-terms)]
         [ty type])
-    (define cached (hash-ref (current-terms) val #f))
+    (define cached (term-cache-ref the-terms val #f))
     (cond
       [cached
        (unless (equal? (term-type cached) ty)
@@ -145,7 +136,7 @@
        (define out (term-constructor val ty ord rest ...))
        (current-index (add1 ord))
        ((current-reporter) 'new-term out)
-       (hash-set! (current-terms) val out)
+       (term-cache-set! the-terms val out)
        out])))
            
 (define (make-const id t) 
diff --git a/test/base/smt.rkt b/test/base/smt.rkt
new file mode 100644
index 00000000..3a52da72
--- /dev/null
+++ b/test/base/smt.rkt
@@ -0,0 +1,84 @@
+#lang rosette
+
+(require rackunit rackunit/text-ui rosette/lib/roseunit)
+
+(define (extract fml)
+  (define-values (in out) (make-pipe))
+  (parameterize ([output-smt out])
+    (solve (assert fml)))
+  (close-output-port out)
+
+  ;; drop 5 for
+  ;; (reset)
+  ;; (set-option :auto-config true)
+  ;; (set-option :produce-unsat-cores false)
+  ;; (set-option :smt.mbqi.max_iterations 10000000)
+  ;; (set-option :smt.relevancy 2)
+  ;;
+  ;; drop-right 7 for
+  ;; (check-sat)
+  ;; (get-model)
+  ;; and the other 5 mentioned above
+
+  (drop-right (drop (port->list read in) 5) 7))
+
+(define smt-tests
+  (test-suite+
+   "SMT tests"
+
+   ;; a dummy call so that next tests start with (reset)
+   (solve #t)
+
+   (define-symbolic a b c d integer?)
+
+   (check-equal?
+    (extract (<= a b))
+    '((declare-fun c0 () Int)
+      (declare-fun c1 () Int)
+      (define-fun e2 () Bool (<= c0 c1))
+      (assert e2)))
+
+   (check-equal?
+    (extract (<= (+ a b) (- c)))
+    '((declare-fun c0 () Int)
+      (declare-fun c1 () Int)
+      (define-fun e2 () Int (+ c0 c1))
+      (declare-fun c3 () Int)
+      (define-fun e4 () Int (- c3))
+      (define-fun e5 () Bool (<= e2 e4))
+      (assert e5)))
+
+   (check-equal?
+    (extract (<= (+ a b) (- (+ a b))))
+    '((declare-fun c0 () Int)
+      (declare-fun c1 () Int)
+      (define-fun e2 () Int (+ c0 c1))
+      (define-fun e3 () Int (- e2))
+      (define-fun e4 () Bool (<= e2 e3))
+      (assert e4)))
+
+   (check-equal?
+    (extract (<= (+ a b c d) (- (+ a b c d))))
+    '((declare-fun c0 () Int)
+      (declare-fun c1 () Int)
+      (declare-fun c2 () Int)
+      (declare-fun c3 () Int)
+      (define-fun e4 () Int (+ c0 c1 c2 c3))
+      (define-fun e5 () Int (- e4))
+      (define-fun e6 () Bool (<= e4 e5))
+      (assert e6)))
+
+   (check-equal?
+    (extract (<= (if (= a b) c d) (if (= a b) d c)))
+    '((declare-fun c0 () Int)
+      (declare-fun c1 () Int)
+      (define-fun e2 () Bool (= c0 c1))
+      (declare-fun c3 () Int)
+      (declare-fun c4 () Int)
+      (define-fun e5 () Int (ite e2 c3 c4))
+      (define-fun e6 () Int (ite e2 c4 c3))
+      (define-fun e7 () Bool (<= e5 e6))
+      (assert e7)))))
+
+(module+ test
+  (time (run-tests smt-tests)))
diff --git a/test/base/term-cache.rkt b/test/base/term-cache.rkt
new file mode 100644
index 00000000..4eeb69f9
--- /dev/null
+++ b/test/base/term-cache.rkt
@@ -0,0 +1,194 @@
+#lang racket
+
+(require rackunit rackunit/text-ui rosette/base/core/term-cache)
+
+(define test-ref
+  (test-suite "term-cache-ref"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+      (check-equal? (term-cache-ref h 1 0) 1)
+      (check-equal? (term-cache-ref h (list xor #t #f) 0) #t)
+      (check-equal? (term-cache-ref h (list - 5) 0) -5)
+      (check-equal? (term-cache-ref h x 0) x)
+      (check-equal? (term-cache-ref h (list + 1 2 3 4) 0) 10)
+      (check-equal? (term-cache-ref h (list + 1 2 3) 0) 6)
+      (check-equal? (term-cache-ref h (dummy 1) 0) (dummy 1))
+      (check-equal? (term-cache-ref h (dummy 2) 0) 0) ; missing
+      )))
+
+(define test-set
+  (test-suite "term-cache-set"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      (define y #'y)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+      (term-cache-set! h 1 2)
+      (check-equal? (term-cache-ref h 1 0) 2)
+
+      (term-cache-set! h (list xor #t #f) #f)
+      (check-equal? (term-cache-ref h (list xor #t #f) 0) #f)
+
+      (term-cache-set! h (list xor #t #f) #f)
+      (check-equal? (term-cache-ref h (list - 5) 0) -5)
+
+      (term-cache-set! h x y)
+      (check-equal? (term-cache-ref h x 0) y)
+
+      (term-cache-set! h (list + 1 2 3 4) 11)
+      (check-equal? (term-cache-ref h (list + 1 2 3 4) 0) 11)
+
+      (term-cache-set! h (list + 1 2 3) 7)
+      (check-equal? (term-cache-ref h (list + 1 2 3) 0) 7)
+
+      (term-cache-set! h (dummy 1) (dummy 2))
+      (check-equal? (term-cache-ref h (dummy 1) 0) (dummy 2))
+
+      (term-cache-set! h (dummy 3) (dummy 3)) ; missing
+      (check-equal? (term-cache-ref h (dummy 3) 0) (dummy 3))
+      )))
+
+(define test-count
+  (test-suite "term-cache-count"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+
+      (check-equal? (term-cache-count h) 7))))
+
+(define test-copy
+  (test-suite "term-cache-copy"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+
+      (check-equal? (term-cache-count (term-cache-copy h)) 7))))
+
+(define test-copy-clear
+  (test-suite "term-cache-copy-clear"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+
+      (check-equal? (term-cache-count (term-cache-copy-clear h)) 0)
+      (check-false (term-cache-weak? (term-cache-copy-clear h)))
+
+      (define h* (make-weak-term-cache
+                  (list (cons 1 1) ; a number
+                        (cons (list xor #t #f) #t) ; binary
+                        (cons (list - 5) -5) ; unary
+                        (cons x x) ; identifier
+                        (cons (list + 1 2 3 4) 10) ; too many operands
+                        (cons (list + 1 2 3) 6) ; ternary
+                        (cons (dummy 1) (dummy 1)) ; unknown value
+                        )))
+      (check-equal? (term-cache-count (term-cache-copy-clear h*)) 0)
+      (check-true (term-cache-weak? (term-cache-copy-clear h*))))))
+
+(define test-strong-weak
+  (test-suite "make-term-cache + make-weak-term-cache"
+    (let ()
+      (struct dummy (v) #:transparent)
+      (define x #'x)
+      ;; h should functionally be a hash-equal?
+      (define h (make-term-cache
+                 (list (cons 1 1) ; a number
+                       (cons (list xor #t #f) #t) ; binary
+                       (cons (list - 5) -5) ; unary
+                       (cons x x) ; identifier
+                       (cons (list + 1 2 3 4) 10) ; too many operands
+                       (cons (list + 1 2 3) 6) ; ternary
+                       (cons (dummy 1) (dummy 1)) ; unknown value
+                       )))
+
+      (check-equal? (term-cache-count h) 7)
+      (check-false (term-cache-weak? h))
+
+      (define dummy-val (dummy 1))
+
+      (set! h (make-weak-term-cache
+               (list (cons 1 1) ; a number
+                     (cons (list xor #t #f) #t) ; binary
+                     (cons (list - 5) -5) ; unary
+                     (cons x x) ; identifier
+                     (cons (list + 1 2 3 4) 10) ; too many operands
+                     (cons (list + 1 2 3) 6) ; ternary
+                     (cons dummy-val dummy-val) ; unknown value
+                     )))
+
+      (collect-garbage)
+      (collect-garbage)
+      (collect-garbage)
+
+      ;; only 1, x and, dummy-val should be left
+      (check-equal? (term-cache-count h) 3)
+
+      (set! dummy-val #f)
+      (collect-garbage)
+      (collect-garbage)
+      (collect-garbage)
+
+      ;; only 1 and x should be left
+      (check-equal? (term-cache-count h) 2)
+
+      (check-true (term-cache-weak? h)))))
+
+
+(module+ test
+  (run-tests test-ref)
+  (run-tests test-set)
+  (run-tests test-count)
+  (run-tests test-copy)
+  (run-tests test-copy-clear)
+  (run-tests test-strong-weak)
+  )