; Utilities for simulating Rosette programs.
#lang racket/base

(provide simulate-rosette)

(require (only-in rosette bv)
         racket/match
         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.
; - 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.
; - initial-state: The initial state of the module, as generated by Yosys's `write_fuctional_rosette`
;   backend.
(define (simulate-rosette #:function function
                          #:input-helper input-helper
                          #:output-helper output-helper
                          #:initial-state initial-state
                          #:inputs inputs
                          #:outputs outputs)
  (error "TODO: Implement simulate-rosette function"))

; Inputs:
; - config: Association list mapping input name (string) to a configuration value, which is one of the
;   following:
;   - 'exhaustive: The input should be exhaustively tested.
;   - <integer>: The input should be tested with this many random inputs. When the input is not
;     present in the config, it defaults to 'exhaustive.
; -
(define (generate-inputs #:input-helper input-helper
                         #:num-inputs num-inputs
                         #:config config
                         #:inputs inputs)
  ; Fill out missing vallues in the config with 'exhaustive.
  (define config
    (map (λ (input)
           (let ([found (assoc (car input) config)]) (or found (cons (car input) 'exhaustive))))
         inputs))

  ; ; Generate the inputs.
  ; (define generated-inputs
  ;   (map (λ (input)
  ;          (let ([input-name (car input)] [input-bitwidth (cdr input)])
  ;            (cond
  ;              [(equal? 'exhaustive (cdr (assoc input-name config))) (list input-name 'exhaustive)]
  ;              [(number? (cdr (assoc input-name config)))
  ;               (list input-name (make-random-input input-type (cdr (assoc input-name config))))]
  ;              [else (error "Invalid configuration for input" input-name)])))
  ;        inputs))

  (error "TODO"))

; Helper function: for a given input name, bitwidth, and configuration, generate a list of inputs.
; Output: List of Rosette bitvector values for the input.
(define (generate-inputs-for-one input-name bitwidth config)
  (cond
    ; If the configuration is 'exhaustive, or if they request a number of inputs that is greater than
    ; or equal to the number of possible values for the bitwidth, generate all possible inputs.
    [(or (equal? config 'exhaustive) (and (number? config) (>= config (expt 2 bitwidth))))
     (for/list ([n (range (expt 2 bitwidth))])
       (bv n bitwidth))]
    [(and (number? config) (positive? config))
     (map (λ (_) (bv (random (expt 2 bitwidth)) bitwidth)) (range config))]
    [else (error (format "Invalid configuration ~a for input ~a" config input-name))]))

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

(module+ test
  (require rackunit)
  (test-case "generate-inputs-for-one"
    (check-equal? (generate-inputs-for-one "input1" 4 'exhaustive)
                  (list (bv 0 4)
                        (bv 1 4)
                        (bv 2 4)
                        (bv 3 4)
                        (bv 4 4)
                        (bv 5 4)
                        (bv 6 4)
                        (bv 7 4)
                        (bv 8 4)
                        (bv 9 4)
                        (bv 10 4)
                        (bv 11 4)
                        (bv 12 4)
                        (bv 13 4)
                        (bv 14 4)
                        (bv 15 4)))

    ; Requesting fewer inputs than the number of possible values for the bitwidth.
    (check-equal? (length (generate-inputs-for-one "input2" 3 5)) 5)

    ; Requesting more inputs than the number of possible values for the bitwidth.
    (check-equal? (generate-inputs-for-one "input3" 2 5) (list (bv 0 2) (bv 1 2) (bv 2 2) (bv 3 2)))

    ; Requesting equal number of inputs as the number of possible values for the bitwidth.
    (check-equal? (generate-inputs-for-one "input4" 2 4) (list (bv 0 2) (bv 1 2) (bv 2 2) (bv 3 2)))

    ; Requesting invalid configuration should raise an error.
    (check-exn exn:fail? (λ () (generate-inputs-for-one "input5" 2 'invalid-config)))
    (check-exn exn:fail? (λ () (generate-inputs-for-one "input5" 2 bytes->string/latin-1)))))