; Copyright (c) 2024-2025 Diogo Behrens
;
; Permission to use, copy, modify, and/or distribute this software for any
; purpose with or without fee is hereby granted.
;
; THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
; REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
; AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
; INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
; LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
; OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
; PERFORMANCE OF THIS SOFTWARE.

(define-syntax print
  (syntax-rules ()
    ((_ ARG ...)
     (begin
       (display ARG) ...
       (newline)))))

;; ----------------------------------------------------------------------------
;; DSL
;; ----------------------------------------------------------------------------

;;; Cooperative yield point.
;;; Suspends the current task and resumes another.
(define (YIELD)
  (yield))

;;; Spawn a new task to run `thunk` (a 0-argument procedure).
(define (SPAWN thunk)
  (spawn thunk))

;;; Explore all executions (subject to pruning).
(define (EXPLORE)
  (explore))

;;; Non-deterministically pick one of `options` (used to model choices).
(define (CHOOSE . options)
  (when (not (*finisher*))
    (error "CHOOSE: scheduler not running"))
  (when (null? options)
    (error "CHOOSE: empty option list"))
  (let* ((i (pick-next options)))
    (list-ref options i)))

;; -----------------------------------------------------------------------------
;; taskqueue
;; -----------------------------------------------------------------------------
(define *tasks* '())

(define (no-tasks?)
  (null? *tasks*))

(define (push-task task)
  (when (memq task *tasks*)
    (error "push-task: duplicate continuation"))
  (set! *tasks* (cons task *tasks*)))

(define (take-task idx)
  (cond ((not idx) #f)
        ((or (negative? idx)
             (>= idx (length *tasks*))) #f)
        (else (let ((task (list-ref *tasks* idx)))
                (set! *tasks* (delete task *tasks*))
                task))))

(define (clear-taskqueue!)
  (set! *tasks* '()))

;; Return a snapshot of the scheduler state
(define (snapshot-taskqueue)
  *tasks*)

;; Restore the scheduler state with a snapshot
(define (restore-taskqueue! tasks)
  (set! *tasks* tasks))

;; -----------------------------------------------------------------------------
;; scheduler
;; -----------------------------------------------------------------------------

(define *finisher* (make-parameter #f))
(define (make-finisher cc)
  (lambda (_) (clear-taskqueue!) (cc #f)))

;; Spawn a new task given a thunk.
(define (spawn thunk)
  (push-task (lambda (_)
               (thunk)
               (resume)
               (error "unreachable"))))

;; Yield current task and resume another.
(define (yield)
  (when (not (*finisher*))
    (error "scheduler not running"))
  (call/cc (lambda (cc)
             (push-task cc)
             (resume)
             (error "unreachable"))))

(define (resume)
  (if (no-tasks?)
      ((*finisher*) #f)
      (let* ((idx (pick-next *tasks*))
             (cc (take-task idx)))
        (cc #f))))

(define (run)
  (call/cc (lambda (cc)
             (parameterize ((*finisher* (make-finisher cc)))
               (resume)))))

;; -----------------------------------------------------------------------------
;; history
;; -----------------------------------------------------------------------------

;; history of selections form (cons idx max-idx)
(define *history* '())

;; returns true if in any scheduling point of the future,
;; there still a different option to take
(define (has-future? hist)
  (cond
    ((null? hist) #f)
    ((< (caar hist) (cdar hist)) #t)
    (else (has-future? (cdr hist)))))

;; add index selection to history
(define (add-history idx max-idx)
  (let ((item (cons idx max-idx)))
    (set! *history* (cons item *history*))))

(define (extract-history)
  (reverse *history*))

(define (swap-history!)
  (let ((prev-hist (extract-history)))
    (set! *history* '())
    prev-hist))

;; -----------------------------------------------------------------------------
;; systematic exploration
;; -----------------------------------------------------------------------------

(define (pick/dfs hist)
  (cond
    ((null? hist)
     (values 0 '()))

    ; some future change point has an option
    ((has-future? (cdr hist))
     (let ((idx (caar hist)))
       (values idx (cdr hist))))

    ; there is still some option in the current change point
    ((< (caar hist) (cdar hist))
     (let ((idx (+ (caar hist) 1)))
       (values idx '())))

    (else (error "should never happen"))))

(define *prev-hist* '())
(define (pick-next lst)
  (let-values (((nxt hist) (pick/dfs *prev-hist*)))
    (set! *prev-hist* hist)
    (add-history nxt (- (length lst) 1))
    nxt))

(define (explore)
  (let ((initial-taskqueue (snapshot-taskqueue)))
    (let loop ((execution-count 0))
      (if (or (zero? execution-count)
              (has-future? (extract-history)))
          (let ((prev-hist (swap-history!)))
            (print "\n=== EXECUTION " (+ 1 execution-count) " ===")
            (set! *prev-hist* prev-hist)
            (run)
            (restore-taskqueue! initial-taskqueue)
            (print "== HISTORY: " (extract-history))
            (loop (+ execution-count 1)))
          execution-count))))