( iteration control stacks
  We create two new stacks - a small stack to hold the "current" value
  of the loop, or the "i" stack, and a larger stack to hold any extra 
  state, which we call the "next" stack. Typically the top of the next
  stack would contain an "iterator" that knows how to advance to the next 
  value and how to be cancelled, depending on the needs of the calling code.

  With these two new stacks, we can create a generic loop construct for
  iterating over streaming values. Not only that, but those values can be 
  arbitrarily filtered and transformed simply by pushing a new value onto
  the iter-next stack which calls out to the previous one. )

uservar itop
8 cells userallot
{ userhere @ } const itop-init
uservar nexttop
24 cells userallot
{ userhere @ } const nexttop-init

' task-init :chain 
  >r itop-init itop r@ !far
     nexttop-init nexttop r@ !far <r ;

taskseg task-init drop

{ : :peek 
:ASM ( pixp -- ) >r
  ( i -- v )
  POP AX
  SHL AX 1 #
  MOV BX @[ SS: <r @]
  ADD BX AX
  PUSH @[ SS: BX]
  NEXT ; }

itop    :peek ipeek
nexttop :peek nextpeek

{ : :drop ( pixp -- )
:ASM >r
  MOV BX @[ SS: r@ @]
  INC BX INC BX
  MOV @[ SS: <r @] BX
  NEXT ; 
  : :ndrop ( pixp -- )
:ASM >r ( c -- )
  MOV BX @[ SS: r@ @]
  POP CX
  SHL CX 1 #
  ADD BX CX
  MOV @[ SS: <r @] BX
  NEXT ; }

itop    :drop idrop
itop    :ndrop n-idrop
nexttop :drop nextdrop
nexttop :ndrop n-nextdrop

: iterdrop ( n-next -- ) n-nextdrop idrop ;
: 1-iterdrop nextdrop idrop ; 

{ : :push ( pixp -- )
:ASM >r
  ( v -- )
  POP AX
  MOV BX @[ SS: r@ @]
  DEC BX DEC BX
  MOV @[ SS: <r @] BX
  MOV @[ SS: BX] AX
  NEXT ; }

itop    :push >i
nexttop :push >next

: <i    0 ipeek    idrop ;
: <next 0 nextpeek nextdrop ;
: i 0 ipeek ; : j 1 ipeek ;

:asm n-<next ( n |n| args... -- args... |n| )
  MOV AX SS
  MOV ES AX
  POP CX
  JCXZ 2 @>
  MOV DI @[ SS: nexttop @]
  ( make SP affect the nextstack and DI affect the data stack. )
  STD ( data stack grows down ) 
  XCHG DI SP 
  ( PUSH decrements and then stores; STOSW stores then decrements. )
  SCASW ( pre-decrement )
1 :>
  POP AX
  STOSW
  LOOP 1 <@
  ( fix SP - DI is one word past the end of the stack )
  CLD SCASW XCHG SP DI
  ( update nexttop )
  MOV @[ SS: nexttop @] DI
2 <:
  NEXT

:asm n->next ( args... n |n| -- |n| args... )
  MOV AX SS
  MOV ES AX
  POP CX
  JCXZ 1 @>
  MOV DI @[ SS: nexttop @]
  STD ( next-stack grows down )
  SCASW ( pre-decrement )
0 :>
  POP AX
  STOSW
  LOOP 0 <@
  CLD SCASW ( correct DI - off by one word )
  MOV @[ SS: nexttop @] DI
1 <:
  NEXT

( iterators are pointers to an array containing two function pointers: 
  xt-iter xt-cancel
  The xt-iter word must take care of updating the stacks directly. If 
  there are no more values, it must remove the values from the i-stack, 
  drop itself from the next-stack, and return 0. "finished" and "finish?"
  are useful words to help with this.

  The xt-cancel word should remove all of the iterator's state from the
  iteration stacks and return nothing. "n-nextdrop" and "iterdrop" are useful 
  words to help with this. If the iterator is itself making use of an 
  iterator below it on the stack, the xt-cancel word should call "cancel" to 
  recursively clean that up once it's done.

  Note that all "next" words _must_ be defined in the target Forth!
  This means that any iterator that dereferences near memory, such as "links",
  WILL NOT WORK on the host Forth! )

: iterate 0 nextpeek @ execute ;
: cancel 0 nextpeek cell + @ execute ;

: finished ( -- 0 ) cancel 0 ;
: finish? ( f -- f ) if 1 else finished then ;

: EACH_ <r iterate if cell + else @ then >r ;

{ ( Because we dereference pointers on the return stack, we must run this 
    from the caller's segment. Copy the definition into the host segment. )
  : EACH_ <r iterate if cell + else @ then >r ;
  : each ' EACH_ , here >i 0 , ; immediate
  : continue ' GOTO_ , i cell - , ; immediate
  : next ['] continue here <i ! ; immediate
  :timm each t, EACH_ patchpt >i ;
  : CONTINUE t, GOTO_ i cell - w>t ;
  :timm continue CONTINUE ;
  :timm next CONTINUE <i patch!t ;
  : blankiter, t' finished w>t t' 1-iterdrop w>t ; }

{ : defiter CREATE blankiter, DOES} >next ;
{ : :iter CREATE blankiter, startcolon DOES} 
  >r r@ [ 2 cells lit ] + execute <r >next ;

{ : iter! ( val off -- ) cells latest entry>tcp + !t ;
  : next! 2 iter! ;                 : cancel! 3 iter! ;
  : :next target next! startcolon ; : :cancel target cancel! startcolon ; }

defiter >cancel 
:cancel nextdrop cancel ;

{ : break ' >cancel , ['] continue ; immediate
  :timm break t, >cancel CONTINUE ; }

defiter nothing
' nextdrop cancel!

defiter >single-done
:iter single >i ;
:next nextdrop 1 >single-done ;

:iter times ( n -- ) >i ;
:next <i dup 1- >i finish? ;

defiter >links
:next <i @ dup >i finish? ;
: links ( p -- ) dup if >i >links else nothing then ;

: +for? ( n -- f ) <i + dup >i 1 nextpeek != finish? ;

:iter for ( start lim -- ) >next 1- >i ;
:next 1 +for? ;
:cancel 2 iterdrop ;

:iter for+ ( start lim inc -- ) >next >next 1 nextpeek - >i ;
:next 2 nextpeek +for? ;
:cancel 3 iterdrop ;

:iter pchars ( st -- ) 1- >i ;
:next <i 1+ dup >i b@ finish? ;

: nth ( i -- v ) 0 each 2dup = if drop i break then 1+ next swap drop ;
: count 0 each 1+ next ;

:asm _suspend>args 
( iter |r| yieldpoint -- argcount |r| yieldpoint iter )
  MOV DI @[ -2 @+ BP]
  POP @[ BP]
  INC BP INC BP
  XOR AH AH
  MOV AL @[ DS: DI]
  PUSH AX
  NEXT

:asm 2r>next
  MOV BX @[ SS: nexttop @]
  ( bp grows up. top: BP-2, second: BP-4
    nexttop grows down. top: nexttop, second: nexttop+2 )
  SUB BX 4 # ( make room for 2 more items on next stack )
  MOV AX @[ -4 @+ BP]
  MOV @[ SS: 2 @+ BX] AX
  MOV AX @[ -2 @+ BP]
  MOV @[ SS: BX] AX
  SUB BP 4 # ( pop two off return stack )
  MOV @[ SS: nexttop @] BX
  NEXT

:asm _resume>args ( |n| yieldpoint xt-next -- argcount |n| |r| resumepoint )
  MOV BX @[ SS: nexttop @]
  MOV DI @[ SS: 2 @+ BX]
  ADD BX 4 #
  MOV @[ SS: nexttop @] BX
  XOR AH AH
  MOV AL @[ DS: DI]
  INC DI
  PUSH AX ( argcount )
  MOV @[ BP] DI
  INC BP INC BP
  NEXT

: _resume _resume>args n-<next rswap ;
: _suspend 
  ( args... iter |r| yieldpoint ret -- 1 |r| ret |n| args... yieldpoint iter )
  rswap _suspend>args n->next 2r>next 1 ;
: _cancel _resume>args n-nextdrop rdrop ;

defiter >genstart
:next _resume>args n-<next ;
' _cancel cancel!
: GENSTART_ <r >next >genstart ;

( yielding from a generator has three moving parts:
  suspend: run immediately when the generator yields. takes care of updating
           the i-stack appropriately, as well as potentially saving any
           yield-specific state to the next-stack. After this runs, the 
           generator's extra parameters are pushed to the next-stack, along
           with the generator's "resume" pointer and the yielder's iterator.
  resume: run when the generator is resumed by iterating. Before this runs,
          the yielder will take care of placing the "resume" pointer back
          on the stack and restoring the generator's arguments, so the
          stack environments should match what they were when the suspension
          happened. This deals with cleaning up after the yielder so the
          rest of the generator can run.
  cancel: run when the generator is aborted. Before this runs, the next-stack
          is cleared of the "resume" pointer and all of the generator's
          arguments. This should remove any yield-specific data from the
          iteration stacks, and potentially call "cancel" to clean up the
          iterator underneath if the yielder is intended to map or filter
          another iterator. )

{ var gen-arg-count 
  :timm (( t:| t, GENSTART_ gen-arg-count @ >t ;
  :timm )) t, 0 t|; t, execute 0 gen-arg-count ! ;
  : +arg  1 gen-arg-count !+ ; :timm +arg +arg ;
  : -arg -1 gen-arg-count !+ ; :timm -arg -arg ;
  :timm >arg t, >next +arg ;
  : colonpair ( xt-first xt-second -- xt )
    target <rot gencolon w>t swap w>t t, return ;
  : :yield ( xt-suspend xt-resume xt-cancel -- )
    } t' _cancel swap colonpair swap t' _resume swap colonpair
    target >r w>t w>t gencolon w>t r@ compilenum t, _suspend t, return
    create immediate <r 2 cells + , does> @ w>t gen-arg-count @ >t ; }

: shadow-i <i >next >i ; : 2>i >i >i ; : 2idrop idrop idrop ;
: unmap idrop <next >i ; : mapcancel unmap cancel ;
: unsuspend rdrop rdrop <r 1+ >r ; ( don't yield at all, skip past the yielder )
: suspend? not if unsuspend then ;

(       suspend    resume      cancel )
         ' noop    ' noop      ' noop :yield yield0 [
           ' >i   ' idrop     ' idrop :yield yield
           ' >i      ' <i     ' idrop :yield yield>
          ' 2>i  ' 2idrop         dup :yield yield2
     ' shadow-i   ' unmap ' mapcancel :yield map
         ' noop    ' noop    ' cancel :yield pass
     ' suspend?    ' noop    ' cancel :yield filter

: take ( n -- ) >arg (( each dup if pass else break then 1- next drop )) ;
: readbytes ( -- ) (( each i b@ map next )) ;
: chars ( p -- ) pchars readbytes ;