generators: reduce overhead for take/take-all/yield-from operations.

extra/generators/generators.factor

@@ -1,39 +1,70 @@
 ! Copyright (C) 2023 Keldan Chapman.
 ! See https://factorcode.org/license.txt for BSD license.
-USING: kernel coroutines effects.parser words sequences accessors generalizations
-locals.parser summary combinators.smart math continuations make ;
+USING: accessors combinators.smart continuations coroutines
+effects.parser generalizations kernel locals.parser math
+sequences summary vectors words ;
 IN: generators
 
 TUPLE: generator state ;
+
 ERROR: stop-generator ;
+
 ERROR: has-inputs ;
-M: has-inputs summary drop "Generator quotation cannot require inputs" ;
 
-: assert-no-inputs ( quot -- ) inputs [ has-inputs ] unless-zero ;
-: gen-coroutine ( quot gen -- co ) '[ f _ state<< stop-generator ] compose cocreate ;
-: <generator> ( quot -- gen ) dup assert-no-inputs generator new [ gen-coroutine ] [ state<< ] [ ] tri ;
+M: has-inputs summary
+    drop "Generator quotation cannot require inputs" ;
+
+: assert-no-inputs ( quot -- )
+    inputs [ has-inputs ] unless-zero ;
+
+: gen-coroutine ( quot gen -- co )
+    '[ f _ state<< stop-generator ] compose cocreate ;
+
+: <generator> ( quot -- gen )
+    dup assert-no-inputs generator new
+    [ gen-coroutine ] [ state<< ] [ ] tri ;
+
+: next ( gen -- result )
+    state>> [ *coresume ] [ stop-generator ] if* ;
+
+: next* ( v gen -- result )
+    state>> [ coresume ] [ drop stop-generator ] if* ;
 
-: next ( gen -- result ) state>> [ *coresume ] [ stop-generator ] if* ;
-: next* ( v gen -- result ) state>> [ coresume ] [ drop stop-generator ] if* ;
 ALIAS: yield  coyield*
+
 ALIAS: yield* coyield
 
-: make-gen-quot ( quot effect -- quot ) in>> length [ ncurry <generator> ] 2curry ;
+: make-gen-quot ( quot effect -- quot )
+    in>> length '[ _ _ ncurry <generator> ] ;
 
 SYNTAX: GEN: (:) [ make-gen-quot ] keep define-declared ;
+
 SYNTAX: GEN:: (::) [ make-gen-quot ] keep define-declared ;
 
-! Utilities
 : skip ( gen -- ) next drop ; inline
+
 : skip* ( v gen -- ) next* drop ; inline
 
 : catch-stop-generator ( ..a try: ( ..a -- ..b ) except: ( ..a -- ..b ) -- ..b )
-    [ stop-generator? [ rethrow ] unless ] prepose recover ; inline
-: ?next ( gen -- val/f end? ) [ next f ] [ drop f t ] catch-stop-generator ;
-: ?next* ( v gen -- val/f end? ) [ next* f ] [ 2drop f t ] catch-stop-generator ;
-: take ( gen n -- seq ) [ swap '[ drop _ ?next [ , t ] unless ] all-integers? drop ] { } make ;
-: take-all ( gen -- seq ) '[ _ ?next not ] [ ] produce nip ;
+    [ dup stop-generator? [ drop ] [ rethrow ] if ] prepose recover ; inline
+
+: ?next ( gen -- val/f end? )
+    [ next f ] [ drop f t ] catch-stop-generator ;
+
+: ?next* ( v gen -- val/f end? )
+    [ next* f ] [ 2drop f t ] catch-stop-generator ;
+
+:: take ( gen n -- seq )
+    n <vector> :> accum
+    [ n [ gen next accum push ] times ] [ ] catch-stop-generator
+    accum { } like ;
+
+:: take-all ( gen -- seq )
+    V{ } clone :> accum
+    [ [ gen next accum push t ] loop ] [ ] catch-stop-generator
+    accum { } like ;
 
-: yield-from ( gen -- ) '[ _ ?next [ drop f ] [ yield t ] if ] loop ;
+: yield-from ( gen -- )
+    '[ [ _ next yield t ] loop ] [ ] catch-stop-generator ;
 
 : exhausted? ( gen -- ? ) state>> not ;