#lang rosette

(require
  (prefix-in $ (only-in rosette bveq bvslt bvsgt bvsle bvsge bvult bvugt bvule bvuge))
  rosette/lib/angelic rosette/lib/match)

; ----- BV semantics -----;

; BV comparison operators return 1/0 instead of #t/#f.
; The language is similar to the one defined in this paper:
; https://dl.acm.org/citation.cfm?id=1993506

(define register? integer?)
(define int32? (bitvector 32))
(define (int32 c) (bv c int32?))

(define-syntax-rule (define-comparators [id op] ...)
  (begin (define (id x y) (if (op x y) (int32 1) (int32 0))) ...))

(define-comparators
  [bveq $bveq] [bvslt $bvslt] [bvult $bvult] [bvsle $bvsle] [bvule $bvule]
  [bvsgt $bvsgt] [bvugt $bvugt] [bvsge $bvsge] [bvuge $bvuge])

(define ops (list bveq bvslt bvult bvsle bvule bvsgt bvugt bvsge
                  bvnot bvand bvor bvxor bvshl bvlshr bvashr
                  bvneg bvadd bvsub bvmul bvsdiv bvudiv bvsrem bvurem bvsmod))

(define optable (for/list ([op ops]) (cons (object-name op) op)))

; Returns the procedure corresponding to the given opcode.
(define (lookup op)
  (cdr (assoc op optable)))

; Global registers.
(define memory          
  (let ([m (vector)])
    (case-lambda [() m]
                 [(size) (set! m (make-vector size #f))])))

; Returns the contents of the register idx if idx is a register,
; otherwise returns idx itself.
(define (load idx)
  (if (register? idx)
      (vector-ref (memory) idx)
      idx))

; Stores val in the register idx.
(define (store idx val) 
  (vector-set! (memory) idx val))

; Returns the contents of the last register.   
(define (last)          
  (sub1 (vector-length (memory))))

; Creates the registers for the given program and input.
(define (make-registers prog inputs) 
  (memory (+ (length prog) (length inputs)))
  (for ([(in idx) (in-indexed inputs)])
    (store idx in)))

; The BV interpreter.
(define (interpret prog inputs)
  (make-registers prog inputs) 
  (for ([stmt prog])
    (match stmt
      [(list out opcode in ...)
       (define op (lookup opcode))
       (define args (map load in))
       (store out (apply op args))]))  
  (load (last)))

; The BV verifier.
(define (ver impl spec)
  (define-symbolic* in int32? #:length (procedure-arity spec))
  (define cex
    (verify
     (assert (equal? (interpret (prog-body impl) in) (apply spec in)))))
  (if (sat? cex)
      (evaluate in cex)
      cex))

; The BV synthesizer. 
(define (syn impl spec)
  (define-symbolic* in int32? #:length (procedure-arity spec))
  (define sol
    (synthesize
     #:forall in
     #:guarantee (assert (equal? (interpret (prog-body impl) in) (apply spec in)))))
  (if (sat? sol)
      (evaluate impl sol)
      sol))

; ----- BV syntax -----;

(define (make-sketch args len ops)
  (define-values (unop binop)
    (partition (or/c 'bvneg 'bvnot) ops))
  (for/list ([r len])
    (define-symbolic* c0 c1 int32?)
    (define out (+ r args))
    (define ins (build-list out identity))
    (define i0 (apply choose* c0 ins))
    (define i1 (apply choose* c1 ins))
    (define inst1 (and (not (null? unop)) `(,(apply choose* unop) ,i0)))
    (define inst2 (and (not (null? binop)) `(,(apply choose* binop) ,i0 ,i1)))
    (cons out (apply choose* (filter identity (list inst1 inst2))))))

(struct prog (name args body) #:transparent) 

; The def macro turns a list of BV instructions, or a BV #:sketch specification,
; into a prog that invokes the V interpreter on the provided inputs.  
(define-syntax def  
  (syntax-rules () 
    [(_ id (idx ...) (out op in ...) ...)
     (define id
       (prog 'id '(idx ...) '((out op in ...) ...)))]
    [(_ id (idx ...) #:sketch len (op ...))
    (define id
      (prog 'id '(idx ...) (make-sketch (length '(idx ...)) len '(op ...))))]))
            

; ----- BV demo -----;
(define (bvmax0 x y)
  (if (equal? (bvsge x y) (int32 1)) x y))

(def bvmax1 (0 1) 
  (2 bvsge 0 1)
  (3 bvneg 2)
  (4 bvxor 0 2)
  (5 bvand 3 4)
  (6 bvxor 1 5))

(def bvmax2 (0 1)
  #:sketch 5 (bvsge bvneg bvxor bvand))

; Interaction script: comment out the following two
; lines if you don't have boolector 3.2.1 installed.
; Z3 will work but it's slower.
(require rosette/solver/smt/boolector)
(current-solver (boolector))

(current-bitwidth 4)

(time (ver bvmax1 bvmax0))
(time (syn bvmax2 bvmax0))

; z3:
; cpu time: 61 real time: 480 gc time: 0
; (list (bv #x00010000 32) (bv #x00000000 32))
; cpu time: 13002 real time: 58496 gc time: 995
; (prog 'bvmax2 '(0 1) '((2 bvsge 0 1) (3 bvneg 2) (4 bvxor 0 1) (5 bvand 3 4) (6 bvxor 1 5)))

; boolector 3.2.1 with CaDiCal:
; cpu time: 46 real time: 325 gc time: 0
; (list (bv #xffffffff 32) (bv #xffffffff 32))
; cpu time: 1106 real time: 5611 gc time: 48
; (prog 'bvmax2 '(0 1) '((2 bvxor 1 0) (3 bvsge 0 1) (4 bvneg 3) (5 bvand 2 4) (6 bvxor 1 5)))