#lang rosette/safe

(require rosette/lib/angelic
         rosette/lib/match)

; Most of the setup taken from Whitequark's synthesis demo
; https://gist.github.com/whitequark/1f7cc43c1cc2b179a1d4a668f0089804

(provide ??prog verify-prog synthesize-prog optimize-prog)

(define (??prog ??insn fuel)
  (if (= fuel 0) null
      (cons (??insn) (??prog ??insn (- fuel 1)))))

; program verifier
(define (verify-prog make-state run-insn prog asserts)
  (define S  (make-state))
  (define S* (make-state))
  (define solution
    (verify
     #:guarantee
     (begin
       (for-each (curry run-insn S*) prog)
       (asserts S S*))))
  (if (unsat? solution) #t
      (begin
        (displayln (evaluate S  solution))
        (displayln (evaluate S* solution))
        solution)))
  
; program synthesizer
(define (synthesize-prog make-state run-insn sketch asserts)
  (define S  (make-state))
  (define S* (make-state))
  (define solution
    (synthesize
     #:forall S
     #:guarantee
     (begin
       (for-each (curry run-insn S*) sketch)
       (asserts S S*))))
  (if (unsat? solution) #f
      (begin
        (evaluate sketch solution))))

(define (optimize-prog max-fuel make-state run-insn sketch-gen asserts)
  (define (worker fuel)
    (printf "fuel ~s~n" fuel)
    (or (synthesize-prog make-state run-insn (sketch-gen fuel) asserts)
        (if (>= fuel max-fuel) #f
            (worker (+ fuel 1)))))
  (worker 0))