A toy coroutine implementation (implemented premitives: spawn, join, yield, sleep, select)

coroutine.rkt

#lang racket

;;; global states
(define task-queue '())
(define joiners '())
(define curr-task-handle 'main)
(define halt (lambda () (display "not started yet")))


;;; queue-related functions

(define (enqueue! task)
  ;; task-queue := task-queue ++ [task]
  (set! task-queue (append task-queue (list task))))
(define (dequeue!)
  ;; usual dequeue
  (let ([head (car task-queue)])
    (set! task-queue (cdr task-queue))
    head))
(define (filter-queue! f)
  (set! task-queue (filter f task-queue)))

;;; task constructor and access functions
(define (make-task handle cont)
  ;; (handle, cont)
  (cons handle cont))

;; task-handle = \(a,b) -> a
(define task-handle car)
;; task-cont = \(a,b) -> b
(define task-cont cdr)

(define (delete-task! handle)
  (filter-queue! (lambda (task) (not (eq? (task-handle task) handle)))))

;;; joiner constructor and access functions
(define (make-joiner joiner-handle joinee-handle cont)
  (cons joinee-handle (cons cont joiner-handle)))

(define joiner-joinee-handle car)
(define joiner-joiner-handle cddr)
(define joiner-cont cadr)

;;; TODO: make a result map
;;; handle => value

;;; joiner-related functions
(define (find-joiner handle)
  ;; find(joiners, \j -> j.handle == handle)
  (assoc handle joiners))
(define (add-joiner joiner)
  ;; joiners := [joiner] ++ joiners
  (set! joiners (cons joiner joiners)))
(define (del-joiner handle)
  ;; joiners := for j in joiners, j.handle != handle
  (set! joiners
       (filter (lambda (j) (not (eq? handle (joiner-joinee-handle j))))
              joiners)))


;;; coroutine primitives

(define (yield)
  (call/cc
   (lambda (cont)
     (let* ([handle curr-task-handle]
            [task (make-task handle cont)])
       (enqueue! task)
       (run-next-task)))))

(define (run-next-task)
  (if (not (empty? task-queue))
     (begin
       (let* ([task (dequeue!)]
              [cont (task-cont task)]
              [handle (task-handle task)])
         (set! curr-task-handle handle)
         (cont 'resume)))
     (halt)))

(define (spawn name f)
  (let* ([handle (gensym name)]
         [fake-cont (lambda (_arg)
                      (return-value handle (f))
                      (run-next-task))]
         [task (make-task handle fake-cont)])
    (enqueue! task)
    handle))

(define (start handle)
  (call/cc (lambda (cont)
             (set! halt (lambda () (cont 'halt)))
             (join handle))))

(define (return-value handle value)
  (let ([joiner (find-joiner handle)])
    (if joiner
       (begin
         (set! curr-task-handle (joiner-joiner-handle joiner))
         ((joiner-cont joiner) value))
       (begin
         ;; wait till someone joins me
         (yield)
         (return-value handle value)))))

(define (sleep-until t)
  (if (< t (current-milliseconds))
      '()
      (begin
        (yield)
        (sleep-until t))))
(define (sleep ms)
  (sleep-until (+ (current-milliseconds) ms)))

(define (wait-forever)
  (yield)
  (wait-forever))

(define (join handle)
  (let* ([joiner-handle curr-task-handle]
         [retval (call/cc (lambda (cont)
                            (println "joiner")
                            (add-joiner (make-joiner joiner-handle handle cont))
                            (wait-forever)))])
    (delete-task! joiner-handle)
    (del-joiner handle)
    retval))

(define (join-timeout handle ms)
  (let* ([joiner-handle curr-task-handle]
         [retval (call/cc (lambda (cont)
                            (add-joiner (make-joiner joiner-handle handle cont))
                            (sleep ms)
                            'timeout))])
    (delete-task! joiner-handle)
    (del-joiner handle)
    ;; stop the execution of the task
    (delete-task! handle)
    retval))

(define (select handle1 handle2)
  (let* ([joiner-handle curr-task-handle]
         [retval (call/cc (lambda (cont)
                            (add-joiner (make-joiner joiner-handle handle1 cont))
                            (add-joiner (make-joiner joiner-handle handle2 cont))
                            (wait-forever)))])
    (delete-task! joiner-handle)
    (del-joiner handle1)
    (del-joiner handle2)
    (delete-task! handle1)
    (delete-task! handle2)
    retval))

;;; examples

(define handle1 (spawn 'h1 (lambda () (println "1") (yield) (println "2"))))
(define handle2 (spawn 'h2 (lambda () (println "3") (yield) (println "4"))))
(define handle3 (spawn 'h3 (lambda () (println "5") (yield) (println "6"))))

(define sleep1 (spawn 's1 (lambda () (sleep 1000) (println "sleep 1") 1)))
(define sleep2 (spawn 's2 (lambda () (sleep 2000) (println "sleep 2") 2)))

(define main (spawn 'main
                    (lambda ()
                      (join handle3)
                      (println "joined 3")
                      (join handle2)
                      (println "joined 2")
                      (join handle1)
                      (println "joined 1")
                      (println (select sleep1 sleep2))
                      (println "done")
                      'done
                      )))


(start main)
(println task-queue)