; Utilities for simulating Rosette programs.
;
; Tests can be run with `raco test <this file>`.
#lang racket/base

(provide simulate-rosette)

(require (only-in rosette bv)
         racket/list)

; Inputs:
; - function: The function for the module to simulate. This should be a Rosette function generated by
;   Yosys's `write_fuctional_rosette` backend.
; - initial-state: The initial state of the module, as generated by Yosys's `write_fuctional_rosette`
;   backend.
; - inputs: A list of association lists. The function will be called with each association list as
;   inputs, and the state will be threaded through each call.
;
; Outputs:
; - A list of outputs, one for each cycle. The outputs are a list of the output objects generated by
;   `function`.
(define (simulate-rosette #:function function #:initial-state initial-state #:inputs inputs)

  (define outputs-and-states
    (drop (reverse (foldl (lambda (input acc)
                            (let* ([outputs (function input (cdr (car acc)))]) (cons outputs acc)))
                          (list (cons 'unused initial-state))
                          inputs))
          1))

  (define outputs (map car outputs-and-states))

  outputs)

; Inputs:
; - inputs: association list mapping string name to bitwidth.
; - num-inputs: number of inputs to generate.
; TODO(@gussmith23): If `num-inputs` is more than the number of possible values, just enumerate.
(define (generate-inputs #:inputs inputs #:num-inputs num-inputs)
  (define (generate-random-input inputs)
    (map (lambda (pair) (cons (car pair) (bv (random (expt 2 (cdr pair))) (cdr pair)))) inputs))
  (for/list ([_ (range num-inputs)])
    (generate-random-input inputs)))

; Generates a clock signal for the given inputs.
;
; Given a string of inputs, one per clock cycle, this function generates a clock signal alongside the
; inputs. It does so by alternating the clock signal between 0 and 1 for each cycle, starting with 0.
; For example, if the inputs are (list inputs1 inputs2 inputs3), the output will be (list (cons (cons
; "clk" (bv 0 1)) inputs1) (cons (cons "clk" (bv 1 1)) inputs1) (cons (cons "clk" (bv 0 1)) inputs2)
; (cons (cons "clk" (bv 1 1)) inputs2) ... ).
;
; Inputs:
; - clock-name: The name of the clock signal.
; - inputs: A list of inputs in association list form, as output by `generate-inputs`.
;
; Outputs:
; - A list of association lists, each containing a new clock signal. Will be twice the length of the
;   inputs list.
(define (generate-clock #:clock-name clock-name #:inputs inputs)
  (apply append
         (for/list ([this-cycle-inputs inputs])
           (list (cons (cons clock-name (bv 0 1)) this-cycle-inputs)
                 (cons (cons clock-name (bv 1 1)) this-cycle-inputs)))))

; This is what gets executed when the script is run.
(module main racket/base
  (require racket/cmdline)

  ; - input-helper, output-helper: association-list-based helpers for input and output struct, generated
  ;   by Yosys's `write_fuctional_rosette` backend with `-assoc-list-helpers` enabled.
  )

(module+ test
  (require rackunit
           (only-in rosette bv bvadd))
  (test-case "generate-inputs"
    (check-equal? (length (generate-inputs #:inputs (list (cons "input1" 4)) #:num-inputs 10)) 10)
    ; Check that this call generates a list of one-length lists, each containing a single association
    ; list with the key "input1" and a random value.
    (check-true (foldl (lambda (input acc)
                         (and acc (equal? (length input) 1) (equal? (car (first input)) "input1")))
                       #t
                       (generate-inputs #:inputs (list (cons "input1" 4)) #:num-inputs 10))))

  (test-case "generate-clock"
    (define inputs
      (list (list (cons "input1" (bv 4 4)) (cons "input2" (bv 3 3)) (cons "input3" (bv 2 2)))
            (list (cons "input1" (bv 3 4)) (cons "input2" (bv 4 3)) (cons "input3" (bv 1 2)))
            (list (cons "input1" (bv 2 4)) (cons "input2" (bv 5 3)) (cons "input3" (bv 0 2)))))

    (check-equal? (length (generate-clock #:clock-name "clk" #:inputs inputs)) 6)

    (check-equal?
     (generate-clock #:clock-name "clk" #:inputs inputs)
     (list (cons (cons "clk" (bv 0 1))
                 (list (cons "input1" (bv 4 4)) (cons "input2" (bv 3 3)) (cons "input3" (bv 2 2))))
           (cons (cons "clk" (bv 1 1))
                 (list (cons "input1" (bv 4 4)) (cons "input2" (bv 3 3)) (cons "input3" (bv 2 2))))
           (cons (cons "clk" (bv 0 1))
                 (list (cons "input1" (bv 3 4)) (cons "input2" (bv 4 3)) (cons "input3" (bv 1 2))))
           (cons (cons "clk" (bv 1 1))
                 (list (cons "input1" (bv 3 4)) (cons "input2" (bv 4 3)) (cons "input3" (bv 1 2))))
           (cons (cons "clk" (bv 0 1))
                 (list (cons "input1" (bv 2 4)) (cons "input2" (bv 5 3)) (cons "input3" (bv 0 2))))
           (cons (cons "clk" (bv 1 1))
                 (list (cons "input1" (bv 2 4)) (cons "input2" (bv 5 3)) (cons "input3" (bv 0 2)))))))
  (test-case "simulate-rosette"

    ; This function will take association lists as inputs, so the helper function is simply identity.
    ; This is not generally true of Yosys-generated code. Similarly, this function uses an association
    ; list for state, which is not what Yosys generates, but it's easier for testing.
    ;
    ; A one-stage adder. Inputs are registered in one clock cycle, and the output is the sum of the
    ; two registered inputs.
    (define (module-function inputs state)
      (let* ([a (cdr (assoc "a" inputs))]
             [b (cdr (assoc "b" inputs))]
             [clk (cdr (assoc "clk" inputs))]
             [old-clk (cdr (assoc "clk" state))]
             [prev-a (cdr (assoc "prev-a" state))]
             [prev-b (cdr (assoc "prev-b" state))]
             [a-reg (cdr (assoc "a-reg" state))]
             [b-reg (cdr (assoc "b-reg" state))]
             [clk-ticked (and (equal? clk (bv 1 1)) (equal? old-clk (bv 0 1)))]
             [new-a-reg (if clk-ticked prev-a a-reg)]
             [new-b-reg (if clk-ticked prev-b b-reg)]
             [out (list (cons "o" (bvadd new-a-reg new-b-reg)))]
             [new-state (list (cons "prev-a" a)
                              (cons "a-reg" new-a-reg)
                              (cons "prev-b" b)
                              (cons "b-reg" new-b-reg)
                              (cons "clk" clk))])
        (cons out new-state)))

    (define outputs
      (simulate-rosette #:function module-function
                        #:initial-state (list (cons "a-reg" (bv 0 4))
                                              (cons "b-reg" (bv 0 4))
                                              (cons "prev-a" (bv 0 4))
                                              (cons "prev-b" (bv 0 4))
                                              (cons "clk" (bv 0 1)))
                        #:inputs
                        (list (list (cons "clk" (bv 0 1)) (cons "a" (bv 4 4)) (cons "b" (bv 4 4)))
                              (list (cons "clk" (bv 1 1)) (cons "a" (bv 3 4)) (cons "b" (bv 0 4)))
                              (list (cons "clk" (bv 0 1)) (cons "a" (bv 10 4)) (cons "b" (bv 9 4)))
                              (list (cons "clk" (bv 1 1)) (cons "a" (bv 2 4)) (cons "b" (bv -1 4)))
                              (list (cons "clk" (bv 0 1)) (cons "a" (bv 4 4)) (cons "b" (bv -15 4)))
                              (list (cons "clk" (bv 1 1)) (cons "a" (bv 0 4)) (cons "b" (bv 0 4))))))

    (check-equal? outputs
                  (list (list (cons "o" (bv 0 4)))
                        (list (cons "o" (bv 8 4)))
                        (list (cons "o" (bv 8 4)))
                        (list (cons "o" (bv 3 4)))
                        (list (cons "o" (bv 3 4)))
                        (list (cons "o" (bv -11 4)))))))