;;; -*- Gerbil -*-
;;; (C) vyzo
;;; Iterations and comprehensions

(import :gerbil/gambit
        :std/error
        :std/contract
        :std/generic
        :std/coroutine)
(export
  (struct-out iterator)
  iter :iter iter-end iter-end? iter-next! iter-fini!
  for for* for/collect for/fold
  in-iota in-range in-naturals in-hash in-hash-keys in-hash-values
  in-input-port in-input-lines in-input-chars in-input-bytes
  in-coroutine in-cothread
  yield
  )

(defstruct iterator (e next fini)
  constructor: :init! unchecked: #t final: #t)

(defmethod {:init! iterator}
  (lambda (self e next (fini #f))
    (struct-instance-init! self e next fini)))

(defstruct :iter-end ())
(def iter-end
  (make-:iter-end))
(def (iter-end? obj)
  (eq? iter-end obj))

(def (iter obj)
  (if (iterator? obj) obj
      (:iter obj)))

(defgeneric :iter)
(defmethod (:iter (it iterator))
  it)
(defmethod (:iter (obj <pair>))
  (iter-list obj))
(defmethod (:iter (obj <null>))
  (iter-null))
(defmethod (:iter (obj <vector>))
  (iter-vector obj))
(defmethod (:iter (obj <string>))
  (iter-string obj))
(defmethod (:iter (obj <hash-table>))
  (iter-hash-table obj))
(defmethod (:iter (obj <procedure>))
  (iter-coroutine obj))
(defmethod (:iter (obj <port>))
  (if (input-port? obj)
    (iter-input-port obj)
    (raise-bad-argument :iter "input-port" obj)))
(defmethod (:iter (obj <object>))
  {:iter obj})

(def (iter-null)
  (make-iterator iter-end iterator-e))

(def (iter-list lst)
  (def (next it)
    (using (it :- iterator)
      (match it.e
        ([hd . rest]
         (set! it.e rest)
         hd)
        (else iter-end))))
  (make-iterator lst next))

(defrules defiter-vector ()
  ((_ iter-vector length-e ref-e)
   (def (iter-vector vec)
     (declare (not safe))
     (def (next it)
       (using (it :- iterator)
         (with ([vec . index] it.e)
           (if (fx< index (length-e vec))
             (let (v (ref-e vec index))
               (set! (cdr it.e) (fx1+ index))
               v)
             iter-end))))
     (make-iterator (cons vec 0) next))))

(defiter-vector iter-vector ##vector-length ##vector-ref)
(defiter-vector iter-string ##string-length ##string-ref)

(def (iter-hash-table ht)
  (def (iterate)
    (hash-for-each yield ht))
  (iter-coroutine iterate))

(def (iter-coroutine proc)
  (def (next it)
    (using (it :- iterator)
      (continue it.e)))
  (let (cort (coroutine (lambda () (proc) iter-end)))
    (make-iterator cort next)))

(def (iter-cothread proc)
  (def (next it)
    (using (it :- iterator)
      (continue it.e)))
  (def (fini it)
    (using (it :- iterator)
      (when it.e
        (cothread-stop! it.e)
        (set! it.e #f))))
  (let* ((cothr (cothread (lambda () (proc) iter-end)))
         (it (make-iterator cothr next fini)))
    (make-will it fini)
    it))

(def (iter-input-port port (read-e read))
  (declare (not safe))
  (def (next it)
    (using (it :- iterator)
      (let (val (read-e it.e))
        (if (eof-object? val)
          iter-end
          val))))
  (make-iterator port next))

(def (iter-in-iota start count step)
  (declare (not safe))
  (def (next it)
    (using (it :- iterator)
      (with ([value . limit] it.e)
        (if (fx> limit 0)
          (begin
            (set! (car it.e) (+ value step))
            (set! (cdr it.e) (fx1- limit))
            value)
          iter-end))))
  (unless (and (number? start) (fixnum? count) (number? step))
    (raise-bad-argument in-iota "iota parameters: (count:fixnum start:number step:number)"
      count start step))
  (make-iterator (cons start count) next))

(def* in-iota
  ((count) (iter-in-iota 0 count 1))
  ((count start) (iter-in-iota start count 1))
  ((count start step) (iter-in-iota start count step)))

(defrules defiter-in-range ()
  ((_ iter-in-range cmp)
   (def (iter-in-range start end step)
     (declare (not safe))
     (def (next it)
       (using (it :- iterator)
         (if (cmp it.e end)
           (let (e it.e)
             (set! it.e (+ e step))
             e)
           iter-end)))
     (unless (and (real? start) (real? end) (real? step))
       (raise-bad-argument in-range "range parameters: real numbers" start end step))
     (make-iterator start next))))

(defiter-in-range iter-in-range< <)
(defiter-in-range iter-in-range> >)

(def* in-range
  ((end) (iter-in-range< 0 end 1))
  ((start end)
   (if (> start end)
     (iter-in-range> start end -1)
     (iter-in-range< start end  1)))
  ((start end step)
   (if (negative? step)
     (iter-in-range> start end step)
     (iter-in-range< start end step))))

(def (in-naturals (start 0) (step 1))
  (declare (not safe))
  (def (next it)
    (using (it :- iterator)
      (let (value it.e)
        (let (value+step (+ value step))
          (set! it.e value+step)
          value))))
  (make-iterator start next))

(def (in-hash ht)
  (iter-hash-table ht))

(def (in-hash-keys ht)
  (def (iterate)
    (hash-for-each (lambda (k v) (yield k)) ht))
  (iter-coroutine iterate))

(def (in-hash-values ht)
  (def (iterate)
    (hash-for-each (lambda (k v) (yield v)) ht))
  (iter-coroutine iterate))

(def (in-input-port obj (read-e read))
  (iter-input-port obj read-e))

(def (in-input-lines obj)
  (iter-input-port obj read-line))

(def (in-input-chars obj)
  (iter-input-port obj read-char))

(def (in-input-bytes obj)
  (iter-input-port obj read-u8))

(def (in-coroutine proc . args)
  (iter-coroutine (if (null? args) proc (cut apply proc args))))

(def (in-cothread proc . args)
  (iter-cothread (if (null? args) proc (cut apply proc args))))

(def (iter-next! it)
  (using (it :- iterator)
    (declare (not safe))
    (it.next it)))

(def (iter-fini! it)
  (declare (not safe))
  (@iter-fini! it))

(defrules @iter-fini! ()
  ((_ it)
   (cond
    ((&iterator-fini it) => (cut <> it)))))

(def (iter-filter pred it)
  (def (iterate)
    (for (val it)
      (when (pred val)
        (yield val))))
  (iter-coroutine iterate))

(begin-syntax
  (def (for-binding? bind)
    (syntax-case bind (when unless)
      ((pat expr) (match-pattern? #'pat))
      ((pat expr when filter-expr) (match-pattern? #'pat))
      ((pat expr unless filter-expr) (match-pattern? #'pat))
      (_ #f)))

  (def (for-binding-expr binding)
    (syntax-case binding (when unless)
      ((bind bind-e) #'bind-e)
      ((bind bind-e when filter-e)
       #'(iter-filter (match <> (bind filter-e)) bind-e))
      ((bind bind-e unless filter-e)
       #'(iter-filter (match <> (bind (not filter-e))) bind-e))))

  (def (for-binding-bind binding)
    (syntax-case binding ()
      ((bind bind-e . _) #'bind)))

  (def (for-iota-args iter-e)
    (syntax-case iter-e ()
      ((_ n) #'(n 0 1))
      ((_ n start) #'(n start 1))
      ((_ n start step) #'(n start step))))

  (def (for-range-args iter-e)
    (syntax-case iter-e ()
      ((_ end) #'(0 end 1))
      ((_ start end) #'(start end (if (> $start $end) -1 1)))
      ((_ start end step) #'(start end step))))

  (def (for-naturals-args iter-e)
    (syntax-case iter-e ()
      ((_) #'(0 1))
      ((_ start) #'(start 1))
      ((_ start step) #'(start step)))))

(defsyntax (for stx)
  (def (generate-for bindings filter body)
    (if (fx= (length bindings) 1)
      (generate-for1 (car bindings) filter body)
      (generate-for* bindings filter body)))

  (def (generate-for1 bind filter body)
    (let ((iter-e (for-binding-expr bind))
          (bind-e (for-binding-bind bind)))
      (syntax-case iter-e (in-iota in-range in-naturals)
        ((in-iota . _)
         (generate-for1-iota iter-e bind-e filter body))
        ((in-range . _)
         (generate-for1-range iter-e bind-e filter body))
        ((in-naturals . _)
         (generate-for1-naturals iter-e bind-e filter body))
        (_
         (with-syntax*
             ((iter-e iter-e)
              (bind-e bind-e)
              ((body ...) body)
              (iter-do-e
               (if filter
                 (with-syntax ((filter-e filter))
                   #'(lambda (val)
                       (with ((bind-e val))
                         (when filter-e
                           body ...))))
                 #'(lambda (val)
                     (with ((bind-e val))
                       body ...)))))
           #'(let ((iterable iter-e)
                   (iter-do iter-do-e))
               (cond
                ;; speculatively inline list iteration
                ((pair? iterable)
                 (for-each iter-do iterable))
                ((null? iterable) (void))
                (else
                 ;; full iteration protocol
                 (let* ((it (iter iterable))
                        (next! (&iterator-next it)))
                   (declare (not safe))
                   (let lp ()
                     (let (val (next! it))
                       (unless (eq? iter-end val)
                         (iter-do val)
                         (lp))))
                   (@iter-fini! it)
                   (void))))))))))

  (def (generate-for1-iota iter-e bind-e filter body)
    (with-syntax* (((count start step) (for-iota-args iter-e))
                   (bind-e bind-e)
                   ((body ...) body)
                   (iter-do-e
                    (if filter
                      (with-syntax ((filter-e filter))
                        #'(lambda (n)
                            (let ((bind-e n))
                              (when filter-e
                                body ...))))
                      #'(lambda (n)
                          (let ((bind-e n))
                            body ...)))))
      #'(let ((iter-do iter-do-e)
              ($count count)
              ($start start)
              ($step step))
          (let lp ((i 0) (val $start))
            (when (fx< i $count)
              (iter-do val)
              (lp (fx1+ i) (+ val $step)))))))

  (def (generate-for1-range iter-e bind-e filter body)
    (with-syntax* (((start end step) (for-range-args iter-e))
                   (bind-e bind-e)
                   ((body ...) body)
                   (iter-do-e
                    (if filter
                      (with-syntax ((filter-e filter))
                        #'(lambda (n)
                            (let ((bind-e n))
                              (when filter-e
                                body ...))))
                      #'(lambda (n)
                          (let ((bind-e n))
                            body ...)))))
      #'(let* ((iter-do iter-do-e)
               ($start start)
               ($end end)
               ($step step))
          (if (negative? $step)
            (let lp ((val $start))
              (when (> val $end)
                (iter-do val)
                (lp (+ val $step))))
            (let lp ((val $start))
              (when (< val $end)
                (iter-do val)
                (lp (+ val $step))))))))

  (def (generate-for1-naturals iter-e bind-e filter body)
    (with-syntax* (((start step) (for-naturals-args iter-e))
                   (bind-e bind-e)
                   ((body ...) body)
                   (iter-do-e
                    (if filter
                      (with-syntax ((filter-e filter))
                        #'(lambda (n)
                            (let ((bind-e n))
                              (when filter-e
                                body ...))))
                      #'(lambda (n)
                          (let ((bind-e n))
                            body ...)))))
      #'(let ((iter-do iter-do-e)
              ($start start)
              ($step step))
          (let lp ((val $start))
            (iter-do val)
            (lp (+ val $step))))))

  (def (generate-for* bindings filter body)
    (with-syntax*
        (((it ...)
          (gentemps bindings))
         ((next! ...)
          (gentemps bindings))
         ((iter-e ...)
          (stx-map for-binding-expr bindings))
         ((bind-id ...)
          (gentemps bindings))
         ((bind-e ...)
          (stx-map for-binding-bind bindings))
         ((body ...) body)
         (iter-do-e
          (if filter
            (with-syntax ((filter-e filter))
              #'(lambda (bind-id ...)
                  (with ((bind-e bind-id) ...)
                    (when filter-e
                      body ...))))
            #'(lambda (bind-id ...)
                (with ((bind-e bind-id) ...)
                  body ...)))))
      #'(let* ((it (iter iter-e)) ...
               (next! (&iterator-next it)) ...
               (iter-do iter-do-e))
          (declare (not safe))
          (let lp ()
            (let ((bind-id (next! it)) ...)
              (unless (or (eq? iter-end bind-id) ...)
                (iter-do bind-id ...)
                (lp))))
          (@iter-fini! it) ...
          (void))))

  (syntax-case stx (when unless)
    ((_ bind body ...)
     (for-binding? #'bind)
     (generate-for [#'bind] #f #'(body ...)))
    ((_ (bind ...) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #f #'(body ...)))
    ((_ (bind ... when expr) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #'expr #'(body ...)))
    ((_ (bind ... unless expr) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #'(not expr) #'(body ...)))))

(defrules for* (when unless)
  ((recur (bind when expr . rest) body ...)
   (for-binding? #'bind)
   (for (bind when expr) (recur rest body ...)))
  ((recur (bind unless expr . rest) body ...)
   (for-binding? #'bind)
   (for (bind unless expr) (recur rest body ...)))
  ((recur (bind . rest) body ...)
   (for-binding? #'bind)
   (for (bind) (recur rest body ...)))
  ((_ () body ...)
   (begin body ...)))

(defsyntax (for/collect stx)
  (def (generate-for bindings filter body)
    (if (fx= (length bindings) 1)
      (generate-for1 (car bindings) filter body)
      (generate-for* bindings filter body)))

  (def (generate-for1 bind filter body)
    (let ((iter-e (for-binding-expr bind))
          (bind-e (for-binding-bind bind)))
      (syntax-case iter-e (in-iota in-range)
        ((in-iota . _)
         (generate-for1-iota iter-e bind-e filter body))
        ((in-range . _)
         (generate-for1-range iter-e bind-e filter body))
        (_
         (with-syntax
             ((iter-e iter-e)
              (bind-e bind-e)
              ((body ...) body))
           (if filter
             (with-syntax ((filter-e filter))
               #'(let ((iterable iter-e)
                       (iter-do
                        (lambda (val)
                          (with ((bind-e val))
                            body ...)))
                       (iter-test
                        (lambda (val)
                          (with ((bind-e val))
                            filter-e))))
                   ;; full iteration protocol
                   (let* ((it (iter iterable))
                          (next! (&iterator-next it)))
                     (declare (not safe))
                     (let (root [#f])
                       (let lp ((tl root))
                         (let (val (next! it))
                           (if (eq? iter-end val)
                             (begin
                               (@iter-fini! it)
                               (cdr root))
                             (if (iter-test val)
                               (let* ((xval (iter-do val))
                                      (tl* [xval]))
                                 (set! (cdr tl) tl*)
                                 (lp tl*))
                               (lp tl)))))))))
             #'(let ((iterable iter-e)
                     (iter-do
                      (lambda (val)
                        (with ((bind-e val))
                          body ...))))
                 (cond
                  ;; speculatively inline list iteration
                  ((pair? iterable)
                   (map iter-do iterable))
                  ((null? iterable) [])
                  (else
                   ;; full iteration protocol
                   (let* ((it (iter iterable))
                          (next! (&iterator-next it)))
                     (declare (not safe))
                     (let (root [#f])
                       (let lp ((tl root))
                         (let (val (next! it))
                           (if (eq? iter-end val)
                             (begin
                               (@iter-fini! it)
                               (cdr root))
                             (let* ((xval (iter-do val))
                                    (tl* [xval]))
                               (set! (cdr tl) tl*)
                               (lp tl*))))))))))))))))

  (def (generate-for1-iota iter-e bind-e filter body)
    (with-syntax (((count start step) (for-iota-args iter-e))
                  (bind-e bind-e)
                  ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let ((iter-do (lambda (n) (let ((bind-e n)) body ...)))
                  (iter-test (lambda (n) (let ((bind-e n)) filter-e)))
                  ($count count)
                  ($start start)
                  ($step step))
              (let (root [#f])
                (let lp ((i 0) (val $start) (tl root))
                  (if (fx< i $count)
                    (lp (##fx+ i 1)
                        (+ val $step)
                        (if (iter-test val)
                          (let (tl* [(iter-do val)])
                            (##set-cdr! tl tl*)
                            tl*)
                          tl))
                    (##cdr root))))))
        #'(let ((iter-do (lambda (n) (let ((bind-e n)) body ...)))
                ($count count)
                ($start start)
                ($step step))
            (let (root [#f])
              (let lp ((i 0) (val $start) (tl root))
                (if (fx< i $count)
                  (lp (##fx+ i 1)
                      (+ val $step)
                      (let (tl* [(iter-do val)])
                        (##set-cdr! tl tl*)
                        tl*))
                  (##cdr root))))))))

  (def (generate-for1-range iter-e bind-e filter body)
    (with-syntax (((start end step) (for-range-args iter-e))
                  (bind-e bind-e)
                  ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let* ((iter-do (lambda (n) (let ((bind-e n)) body ...)))
                   (iter-test (lambda (n) (let ((bind-e n)) filter-e)))
                   ($start start)
                   ($end end)
                   ($step step))
              (let (root [#f])
                (if (negative? $step)
                  (let lp ((val $start) (tl root))
                    (if (> val $end)
                      (lp (+ val $step)
                          (if (iter-test val)
                            (let (tl* [(iter-do val)])
                              (##set-cdr! tl tl*)
                              tl*)
                            tl))
                      (##cdr root)))
                  (let lp ((val $start) (tl root))
                    (if (< val $end)
                      (lp (+ val $step)
                          (if (iter-test val)
                            (let (tl* [(iter-do val)])
                              (##set-cdr! tl tl*)
                              tl*)
                            tl))
                      (##cdr root)))))))
        #'(let* ((iter-do (lambda (n) (let ((bind-e n)) body ...)))
                 ($start start)
                 ($end end)
                 ($step step))
            (let (root [#f])
              (if (negative? $step)
                (let lp ((val $start) (tl root))
                  (if (> val $end)
                    (lp (+ val $step)
                        (let (tl* [(iter-do val)])
                          (##set-cdr! tl tl*)
                          tl*))
                    (##cdr root)))
                (let lp ((val $start) (tl root))
                  (if (< val $end)
                    (lp (+ val $step)
                        (let (tl* [(iter-do val)])
                          (##set-cdr! tl tl*)
                          tl*))
                    (##cdr root)))))))))

  (def (generate-for* bindings filter body)
    (with-syntax
        ((value  (genident 'value))
         (rvalue (genident 'rvalue))
         ((it ...)
          (gentemps bindings))
         ((next! ...)
          (gentemps bindings))
         ((iter-e ...)
          (stx-map for-binding-expr bindings))
         ((bind-id ...)
          (gentemps bindings))
         ((bind-e ...)
          (stx-map for-binding-bind bindings))
         ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let* ((it (iter iter-e)) ...
                   (next! (&iterator-next it)) ...
                   (iter-do
                    (lambda (bind-id ...)
                      (with ((bind-e bind-id) ...)
                        body ...)))
                   (iter-test
                    (lambda (bind-id ...)
                      (with ((bind-e bind-id) ...)
                        filter-e))))
              (declare (not safe))
              (let (root [#f])
                (let lp ((tl root))
                  (let ((bind-id (next! it)) ...)
                    (if (or (eq? iter-end bind-id) ...)
                      (begin (@iter-fini! it) ...
                             (cdr root))
                      (if (iter-test bind-id ...)
                        (let* ((value (iter-do bind-id ...))
                               (tl* [value]))
                          (set! (cdr tl) tl*)
                          (lp tl*))
                        (lp tl))))))))
        #'(let* ((it (iter iter-e)) ...
                 (next! (&iterator-next it)) ...
                 (iter-do
                  (lambda (bind-id ...)
                    (with ((bind-e bind-id) ...)
                      body ...))))
            (declare (not safe))
            (let (root [#f])
              (let lp ((tl root))
                (let ((bind-id (next! it)) ...)
                  (if (or (eq? iter-end bind-id) ...)
                    (begin (@iter-fini! it) ...
                           (cdr root))
                    (let* ((value (iter-do bind-id ...))
                           (tl* [value]))
                      (set! (cdr tl) tl*)
                      (lp tl*))))))))))

  (syntax-case stx (when unless)
    ((_ bind body ...)
     (for-binding? #'bind)
     (generate-for [#'bind] #f #'(body ...)))
    ((_ (bind ...) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #f #'(body ...)))
    ((_ (bind ... when expr) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #'expr #'(body ...)))
    ((_ (bind ... unless expr) body ...)
     (stx-andmap for-binding? #'(bind ...))
     (generate-for #'(bind ...) #'(not expr) #'(body ...)))))

(defsyntax (for/fold stx)
  (def (for/fold-bind? bind)
    (syntax-case bind ()
      ((id expr) (identifier? #'id))
      (else #f)))

  (def (generate-for fold-bind bindings filter body)
    (if (fx= (length bindings) 1)
      (generate-for1 fold-bind (car bindings) filter body)
      (generate-for* fold-bind bindings filter body)))

  (def (generate-for1 fold-bind bind filter body)
    (let ((iter-e (for-binding-expr bind))
          (bind-e (for-binding-bind bind)))
      (syntax-case iter-e (in-iota in-range)
        ((in-iota . _)
         (generate-for1-iota iter-e bind-e fold-bind filter body))
        ((in-range . _)
         (generate-for1-range iter-e bind-e fold-bind filter body))
        (_
         (with-syntax
             ((iter-e iter-e)
              (bind-e bind-e)
              ((fold-iv fold-e) fold-bind)
              ((body ...) body))
           (if filter
             (with-syntax ((filter-e filter))
               #'(let ((iterable iter-e)
                       (iter-do
                        (lambda (val fold-iv)
                          (with ((bind-e val))
                            body ...)))
                       (iter-test
                        (lambda (val fold-iv)
                          (with ((bind-e val))
                            filter-e)))
                       (fold-iv fold-e))
                   ;; full iteration protocol
                   (let* ((it (iter iterable))
                          (next! (&iterator-next it)))
                     (declare (not safe))
                     (let lp ((rval fold-e))
                       (let (val (next! it))
                         (if (eq? iter-end val)
                           (begin
                             (@iter-fini! it)
                             rval)
                           (if (iter-test val rval)
                             (let (xval (iter-do val rval))
                               (lp xval))
                             (lp rval))))))))
             #'(let ((iterable iter-e)
                     (iter-do
                      (lambda (val fold-iv)
                        (with ((bind-e val))
                          body ...)))
                     (fold-iv fold-e))
                 (cond
                  ;; speculatively inline list iteration
                  ((pair? iterable)
                   (foldl iter-do fold-iv iterable))
                  ((null? iterable) fold-iv)
                  (else
                   ;; full iteration protocol
                   (let* ((it (iter iterable))
                          (next! (&iterator-next it)))
                     (declare (not safe))
                     (let lp ((rval fold-e))
                       (let (val (next! it))
                         (if (eq? iter-end val)
                           (begin
                             (@iter-fini! it)
                             rval)
                           (let (xval (iter-do val rval))
                             (lp xval)))))))))))))))

  (def (generate-for1-iota iter-e bind-e fold-bind filter body)
    (with-syntax (((count start step) (for-iota-args iter-e))
                  (bind-e bind-e)
                  ((fold-iv fold-e) fold-bind)
                  ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let ((iter-do (lambda (n fold-iv) (let ((bind-e n)) body ...)))
                  (iter-test (lambda (n fold-iv) (let ((bind-e n)) filter-e)))
                  ($count count)
                  ($start start)
                  ($step step))
              (let lp ((i 0) (val $start) (r fold-e))
                (if (fx< i $count)
                  (lp (fx1+ i)
                      (+ val $step)
                      (if (iter-test val r)
                        (iter-do val r)
                        r))
                  r))))
        #'(let ((iter-do (lambda (n fold-iv) (let ((bind-e n)) body ...)))
                ($count count)
                ($start start)
                ($step step))
            (let lp ((i 0) (val $start) (r fold-e))
              (if (fx< i $count)
                (lp (fx1+ i)
                    (+ val $step)
                    (iter-do val r))
                r))))))

  (def (generate-for1-range iter-e bind-e fold-bind filter body)
    (with-syntax (((start end step) (for-range-args iter-e))
                  (bind-e bind-e)
                  ((fold-iv fold-e) fold-bind)
                  ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let* ((iter-do (lambda (n fold-iv) (let ((bind-e n)) body ...)))
                   (iter-test (lambda (n fold-iv) (let ((bind-e n)) filter-e)))
                   ($start start)
                   ($end end)
                   ($step step))
              (if (negative? $step)
                (let lp ((val $start) (r fold-e))
                  (if (> val $end)
                    (lp (+ val $step)
                        (if (iter-test val r)
                          (iter-do val r)
                          r))
                    r))
                (let lp ((val $start) (r fold-e))
                  (if (< val $end)
                    (lp (+ val $step)
                        (if (iter-test val r)
                          (iter-do val r)
                          r))
                    r)))))
        #'(let* ((iter-do (lambda (n fold-iv) (let ((bind-e n)) body ...)))
                 ($start start)
                 ($end end)
                 ($step step))
            (if (negative? $step)
              (let lp ((val $start) (r fold-e))
                (if (> val $end)
                  (lp (+ val $step)
                      (iter-do val r))
                  r))
              (let lp ((val $start) (r fold-e))
                (if (< val $end)
                  (lp (+ val $step)
                      (iter-do val r))
                  r)))))))

  (def (generate-for* fold-bind bindings filter body)
    (with-syntax
        ((value  (genident 'value))
         ((loop-id loop-e)
          fold-bind)
         ((it ...)
          (gentemps bindings))
         ((next! ...)
          (gentemps bindings))
         ((iter-e ...)
          (stx-map for-binding-expr bindings))
         ((bind-id ...)
          (gentemps bindings))
         ((bind-e ...)
          (stx-map for-binding-bind bindings))
         ((body ...) body))
      (if filter
        (with-syntax ((filter-e filter))
          #'(let* ((it (iter iter-e)) ...
                   (next! (&iterator-next it)) ...
                   (iter-do
                    (lambda (loop-id bind-id ...)
                      (with ((bind-e bind-id) ...)
                        body ...)))
                   (iter-test
                    (lambda (loop-id bind-id ...)
                      (with ((bind-e bind-id) ...)
                        filter-e))))
              (declare (not safe))
              (let lp ((loop-id loop-e))
                (let ((bind-id (next! it)) ...)
                  (if (or (eq? iter-end bind-id) ...)
                    (begin (@iter-fini! it) ...
                           loop-id)
                    (if (iter-test loop-id bind-id ...)
                      (let (value (iter-do loop-id bind-id ...))
                        (lp value))
                      (lp loop-id)))))))
        #'(let* ((it (iter iter-e)) ...
                 (next! (&iterator-next it)) ...
                 (iter-do
                  (lambda (loop-id bind-id ...)
                    (with ((bind-e bind-id) ...)
                      body ...))))
            (declare (not safe))
            (let lp ((loop-id loop-e))
              (let ((bind-id (next! it)) ...)
                (if (or (eq? iter-end bind-id) ...)
                  (begin (@iter-fini! it) ...
                         loop-id)
                  (let (value (iter-do loop-id bind-id ...))
                    (lp value)))))))))

  (syntax-case stx (when unless)
    ((_ fold-bind bind body ...)
     (and (for/fold-bind? #'fold-bind)
          (for-binding? #'bind))
     (generate-for #'fold-bind [#'bind] #f #'(body ...)))
    ((_ fold-bind (bind ...) body ...)
     (and (for/fold-bind? #'fold-bind)
          (stx-andmap for-binding? #'(bind ...)))
     (generate-for #'fold-bind #'(bind ...) #f #'(body ...)))
    ((_ fold-bind (bind ... when expr) body ...)
     (and (for/fold-bind? #'fold-bind)
          (stx-andmap for-binding? #'(bind ...)))
     (generate-for #'fold-bind #'(bind ...) #'expr #'(body ...)))
    ((_ fold-bind (bind ... unless expr) body ...)
     (and (for/fold-bind? #'fold-bind)
          (stx-andmap for-binding? #'(bind ...)))
     (generate-for #'fold-bind #'(bind ...) #'(not expr) #'(body ...)))))
