extra/generators/generators.factor

   1 ! Copyright (C) 2023 Keldan Chapman.
   2 ! See https://factorcode.org/license.txt for BSD license.
   3 USING: kernel coroutines effects.parser words sequences accessors generalizations
   4 locals.parser summary combinators.smart math continuations make ;
   5 IN: generators
   6
   7 TUPLE: generator state ;
   8 ERROR: stop-generator ;
   9 ERROR: has-inputs ;
  10 M: has-inputs summary drop "Generator quotation cannot require inputs" ;
  11
  12 : assert-no-inputs ( quot -- ) inputs [ has-inputs ] unless-zero ;
  13 : gen-coroutine ( quot gen -- co ) '[ f _ state<< stop-generator ] compose cocreate ;
  14 : <generator> ( quot -- gen ) dup assert-no-inputs generator new [ gen-coroutine ] [ state<< ] [ ] tri ;
  15
  16 : next ( gen -- result ) state>> [ *coresume ] [ stop-generator ] if* ;
  17 : next* ( v gen -- result ) state>> [ coresume ] [ drop stop-generator ] if* ;
  18 ALIAS: yield  coyield*
  19 ALIAS: yield* coyield
  20
  21 : make-gen-quot ( quot effect -- quot ) in>> length [ ncurry <generator> ] 2curry ;
  22
  23 SYNTAX: GEN: (:) [ make-gen-quot ] keep define-declared ;
  24 SYNTAX: GEN:: (::) [ make-gen-quot ] keep define-declared ;
  25
  26 ! Utilities
  27 : skip ( gen -- ) next drop ; inline
  28 : skip* ( v gen -- ) next* drop ; inline
  29
  30 : catch-stop-generator ( ..a try: ( ..a -- ..b ) except: ( ..a -- ..b ) -- ..b )
  31     [ stop-generator? [ rethrow ] unless ] prepose recover ; inline
  32 : ?next ( gen -- val/f end? ) [ next f ] [ drop f t ] catch-stop-generator ;
  33 : ?next* ( v gen -- val/f end? ) [ next* f ] [ 2drop f t ] catch-stop-generator ;
  34 : take ( gen n -- seq ) [ swap '[ drop _ ?next [ , t ] unless ] all-integers? drop ] { } make ;
  35 : take-all ( gen -- seq ) '[ _ ?next not ] [ ] produce nip ;
  36
  37 : yield-from ( gen -- ) '[ _ ?next [ drop f ] [ yield t ] if ] loop ;
  38
  39 : exhausted? ( gen -- ? ) state>> not ;