A Racket port of some of the code in the paper "Stream Fusion, to Completeness"

stream-fusion-to-completeness.rkt

#lang racket

;; Ported from the paper "Stream Fusion, to Completeness" (POPL 2017)

(struct staged-stream (init-function step-function))

(struct staged-vec-stream-state (index-id vec-id))

(define (staged-vec-stream vec-expr)
  (define (init state-accepting-result-continuation)
    #`(let ([i 0]
            [vec #,vec-expr])
        #,(state-accepting-result-continuation (staged-vec-stream-state #'i #'vec))))
  (define (step state element-accepting-result-continuation)
    (match-define (staged-vec-stream-state index-box-id vec-id) state)
    #`(cond
        [(< #,index-box-id (vector-length #,vec-id))
         (define element (vector-ref #,vec-id #,index-box-id))
         (set! #,index-box-id (add1 #,index-box-id))
         #,(element-accepting-result-continuation #'element)]
        [else
         #,(element-accepting-result-continuation #false)]))
  (staged-stream init step))


(define (staged-map stream f)
  (match-define (staged-stream init step) stream)

  (define (new-step state element-accepting-result-continuation)
    (step state
          (λ (elem)
            (match elem
              [#false (element-accepting-result-continuation #false)]
              [(? syntax? elem-expr) (element-accepting-result-continuation (f elem-expr))]))))

  (staged-stream init new-step))


(define (staged-fold stream f init-fold-value-expr)
  (match-define (staged-stream init step) stream)
  (init
   (λ (state)
     #`(let loop ([fold-val #,init-fold-value-expr])
         #,(step state
                 (λ (elem)
                   (match elem
                     [#false #'fold-val]
                     [(? syntax? elem-expr) #`(loop #,(f #'fold-val elem-expr))])))))))


(staged-fold (staged-vec-stream #'#(1 2 3 4 5))
             (λ (accum x) #`(+ #,accum #,x))
             #'0)

;; outputs the following code:
;;
;; (let ([i 0]
;;       [vec #(1 2 3 4 5)])
;;   (let loop ([fold-val 0])
;;     (cond
;;      [(< i (vector-length vec))
;;       (define element (vector-ref vec i))
;;       (set! i (add1 i))
;;       (loop (+ fold-val element))]
;;      [else fold-val])))