#lang rosette/safe

(require rosette/lib/angelic
         rosette/lib/match)
(require "synthesis-base.rkt")
(current-bitwidth #f)

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

#;(define (rotate-right i x)
  (define s (bitvector-size (type-of x)))
  (concat (extract (- i 1) 0 x) (extract (- s 1) i x)))
(define (rotate-left i x)
  (define s (bitvector-size (type-of x)))
  (concat (extract (- s i 1) 0 x) (extract (- s 1) (- s i) x)))

(struct state (reg mem) #:mutable #:transparent)
(define (reg-ref S n) (vector-ref (state-reg S) n))
(define (reg-set! S n v) (vector-set! (state-reg S) n v))
(define (mem-ref S addr) (vector-ref (state-mem S) addr))
(define (mem-set! S addr v) (vector-set! (state-mem S) addr v))

(define-symbolic R0 R1 R2 R3 M0 M1 (bitvector 8))
(define (make-state)
  (state (vector R0 R1 R2 R3)
         (vector M0 M1)))
   
(define (bv8 val) (bv val 8))
(define (choose-imm* w)
  (choose* (bv8 #x00) (bv8 #x01) (bv8 #x03) (bv8 #x07)
           (bv8 #x0f) (bv8 #x1f) (bv8 #x3f) (bv8 #x7f)
           (bv8 #xff) (bv8 #xfe) (bv8 #xfc) (bv8 #xf8)
           (bv8 #xf0) (bv8 #xe0) (bv8 #xc0) (bv8 #x80)))
#;(define (choose-imm* w)
  (define-symbolic* b (bitvector w))
  b)
(define (choose-shift*)
  (define-symbolic* b (bitvector 3))
  b)
#;(define (choose-imm*)
  (choose* (bv8 #x01)))

(struct LSL (a imm) #:transparent)
(struct LSR (a imm) #:transparent)
(struct ASR (a imm) #:transparent)
(struct ROL (a imm) #:transparent)
(struct OR (a b) #:transparent)
(struct ORI (a imm) #:transparent)
(struct AND (a b) #:transparent)
(struct ANDI (a imm) #:transparent)
(struct LD0 (a addr) #:transparent)
(struct ST0 (a addr) #:transparent)
(struct MOV (a b) #:transparent)

(define (format-insn insn)
  (define nat bitvector->natural)
  (match insn
    [(LSL a imm) (format "r~s = r~s lsl #~s" a a (nat imm))]
    [(LSR a imm) (format "r~s = r~s lsr #~s" a a (nat imm))]
    [(ASR a imm) (format "r~s = r~s asr #~s" a a (nat imm))]
    [(ROL a imm) (format "r~s = r~s rol #~s" a a (nat imm))]
    [(OR a b) (format "r~s = r~s or r~s" a a b)]
    [(ORI a imm) (format "r~s = r~s or #~s" a a (nat imm))]
    [(AND a b) (format "r~s = r~s and r~s" a a b)]
    [(ANDI a imm) (format "r~s = r~s and #~s" a a (nat imm))]
    [(LD0 a addr) (format "r~s = [#~s]" a addr)]
    [(ST0 a addr) (format "[#~s] = r~s" addr a)]
    [(MOV a b) (format "r~s = r~s" a b)]
    ))

(define (print-insn insn)
  (displayln (format-insn insn)))

(define (print-prog prog)
  (for-each print-insn prog))

(define (??insn)
  (let ([a (choose* 0 1 2 3)]
        [b (choose* 0 1 2 3)]
        [imm8 (choose-imm* 8)]
        [addr (choose* 0)]
        ;[imm7 (choose-imm* 7)]
        ;[imm6 (choose-imm* 6)]
        [imm3 (choose-shift*)])
    (choose* (LSL a imm3)
             (LSR a imm3)
             ;(ASR a imm3)
             ;(ROL a imm3)
             (OR a b)
             ;(ORI a imm8)
             ;(AND a b)
             (ANDI a imm8)
             (LD0 a addr)
             (ST0 a addr)
             (MOV a b)
             )))
            
(define (run-insn S insn)
  (match insn
    [(LSL a imm)
     (let ([va (reg-ref S a)]
           [shift (zero-extend imm (bitvector 8))])
       (reg-set! S a (bvshl va shift)))]
    [(LSR a imm)
     (let ([va (reg-ref S a)]
           [shift (zero-extend imm (bitvector 8))])
       (reg-set! S a (bvlshr va shift)))]
    [(ASR a imm)
     (let ([va (reg-ref S a)]
           [shift (zero-extend imm (bitvector 8))])
       (reg-set! S a (bvashr va shift)))]
    [(ROL a imm)
     (let ([va (reg-ref S a)]
           [shift (bitvector->natural imm)])
       (reg-set! S a (rotate-left imm va)))]
    [(OR a b)
     (let ([va (reg-ref S a)]
           [vb (reg-ref S b)])
       (reg-set! S a (bvor va vb)))]
    [(ORI a imm)
     (let ([va (reg-ref S a)])
       (reg-set! S a (bvor va imm)))]
    [(AND a b)
     (let ([va (reg-ref S a)]
           [vb (reg-ref S b)])
       (reg-set! S a (bvand va vb)))]
    [(ANDI a imm)
     (let ([va (reg-ref S a)])
       (reg-set! S a (bvand va imm)))]
    [(LD0 a addr)
     (reg-set! S a (mem-ref S addr))]
    [(ST0 a addr)
     (mem-set! S addr (reg-ref S a))]
    [(MOV a b)
     (reg-set! S a (reg-ref S b))]
  ))

(define (assert-preserve S S* . regs)
  (define (assert-preserve-reg n)
    (assert (bveq (reg-ref S n) (reg-ref S* n))))
  (for-each assert-preserve-reg regs))

(define found-prog
  (optimize-prog 12 make-state run-insn
    (lambda (fuel) (??prog ??insn fuel))
    (lambda (S S*)
      (define R  (concat (reg-ref S  2) (reg-ref S  1) (reg-ref S  0)))
      (define R* (concat (reg-ref S* 2) (reg-ref S* 1) (reg-ref S* 0)))
      (assert (bveq (bvshl R (bv 4 24)) R*))
      (assert-preserve S S* 3)
      )))

(if found-prog
    (print-prog found-prog)
    (displayln "No program found!"))