Add support for delegation and return values with @yieldfrom (#57)

---------

Signed-off-by: Stefan Krastanov <stefan@krastanov.org>
Co-authored-by: Stefan Krastanov <stefan@krastanov.org>
Co-authored-by: Stefan Krastanov <github.acc@krastanov.org>

.github/workflows/ci.yml

@@ -12,7 +12,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.2'
           - '1.6'
           - '1'
         os:

Project.toml

@@ -5,11 +5,11 @@ license = "MIT"
 desc = "C# sharp style generators a.k.a. semi-coroutines for Julia."
 authors = ["Ben Lauwens <ben.lauwens@gmail.com>"]
 repo = "https://github.com/BenLauwens/ResumableFunctions.jl.git"
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]
 MacroTools = "0.5"
-julia = "1.2"
+julia = "1.6"

README.md

@@ -133,6 +133,11 @@ Iteration protocol:
 
 ## Release notes
 
+* 2023: v0.6.3
+  * Julia 1.6 or newer is required
+  * introduction of `@yieldfrom` to delegate to another resumable function or iterator (similar to [Python's `yield from`](https://peps.python.org/pep-0380/))
+  * resumable functions are now allowed to return values, so that `r = @yieldfrom f` also stores the return value of `f` in `r`
+
 * 2023: v0.6.2 
   * Julia v1.10 compatibility fix
   * resumable functions can now dispatch on types

src/ResumableFunctions.jl

@@ -6,7 +6,7 @@ module ResumableFunctions
   using MacroTools
   using MacroTools: combinedef, combinearg, flatten, postwalk
 
-  export @resumable, @yield, @nosave
+  export @resumable, @yield, @nosave, @yieldfrom 
 
   include("types.jl")
   include("transforms.jl")

src/macro.jl

@@ -5,6 +5,13 @@ macro yield(expr=nothing)
   error("@yield macro outside a @resumable function!")
 end
 
+"""
+Macro if used in a `@resumable function` that delegates to `expr` otherwise throws an error.
+"""
+macro yieldfrom(expr=nothing)
+  error("@yieldfrom macro outside a @resumable function!")
+end
+
 """
 Macro if used in a `@resumable function` that creates a not saved variable otherwise throws an error.
 """
@@ -29,7 +36,8 @@ macro resumable(expr::Expr)
   args, kwargs, arg_dict = get_args(func_def)
   params = ((get_param_name(param) for param in func_def[:whereparams])...,)
   ui8 = BoxedUInt8(zero(UInt8))
-  postwalk(x->hasreturnvalue(x), func_def[:body])
+  func_def[:body] = postwalk(transform_arg_yieldfrom, func_def[:body])
+  func_def[:body] = postwalk(transform_yieldfrom, func_def[:body])
   func_def[:body] = postwalk(x->transform_for(x, ui8), func_def[:body])
   slots = get_slots(copy(func_def), arg_dict, __module__)
   type_name = gensym()

src/transforms.jl

@@ -7,6 +7,64 @@ function transform_nosave(expr, nosaves::Set{Symbol})
   :($var = $body)
 end
 
+"""
+Function that replaces a `@yieldfrom iter` statement with
+```julia
+  _other_ = iter...
+  _ret_ = generate(_other_, nothing)
+  while !(_ret_ isa IteratorReturn)
+    _value_, _state_ = _ret_
+    _newvalue_ = @yield _value_
+    _ret_ = generate(_other_, _newvalue_, _state_)
+  end
+  _
+```
+"""
+function transform_yieldfrom(expr)
+  _is_yieldfrom(expr) || return expr
+  iter = expr.args[3:end]
+  quote
+    _other_ = $(iter...)
+    _ret_ = $generate(_other_, nothing)
+    while !(_ret_ isa $IteratorReturn)
+      _value_, _state_ = _ret_
+      _newvalue_ = @yield _value_
+      _ret_ = $generate(_other_, _newvalue_, _state_)
+    end
+  end
+end
+
+"""
+Function that replaces an `arg = @yieldfrom iter` statement by
+```julia
+  @yieldfrom iter
+  arg = _ret_.value
+```
+"""
+function transform_arg_yieldfrom(expr)
+  @capture(expr, arg_ = ex_) || return expr
+  _is_yieldfrom(ex) || return expr
+  iter = ex.args[3:end]
+  quote
+    @yieldfrom $(iter...)
+    $arg = _ret_.value
+  end
+end
+
+"""
+Function returning whether an expression is a `@yieldfrom` macro
+"""
+_is_yieldfrom(ex) = false
+
+function _is_yieldfrom(ex::Expr)
+  is_ = ex.head === :macrocall && ex.args[1] === Symbol("@yieldfrom")
+  if is_ && length(ex.args) < 3
+    error("@yieldfrom without arguments!")
+  end
+  return is_
+end
+
+
 """
 Function that replaces a `for` loop by a corresponding `while` loop saving explicitely the *iterator* and its *state*.
 """

src/types.jl

@@ -20,9 +20,14 @@ Implements the `eltype` method of the *iterator* interface for a subtype of `Fin
 """
 Base.eltype(::Type{T}) where T<:FiniteStateMachineIterator{R} where R = R
 
-function Base.iterate(fsm_iter::T, state=nothing) where T<:FiniteStateMachineIterator
-  #fsm_iter._state = state
-  result = fsm_iter()
-  fsm_iter._state === 0xff && return nothing
-  result, nothing
-end
+function Base.iterate(fsm_iter::FiniteStateMachineIterator)
+  ret = generate(fsm_iter, nothing)
+  ret isa IteratorReturn && return nothing
+  ret
+end
+
+function Base.iterate(fsm_iter::FiniteStateMachineIterator, state)
+  ret = generate(fsm_iter, nothing, state)
+  ret isa IteratorReturn && return nothing
+  ret
+end

src/utils.jl

@@ -84,12 +84,26 @@ function make_arg_any(expr, slots::Dict{Symbol, Any})
   expr
 end
 
-"""
-Function checking the use of a return statement with value
-"""
-function hasreturnvalue(expr)
-  @capture(expr, return val_) || return expr
-  (val === :nothing || val === nothing) && return expr
-  @warn "@resumable function contains return statement with value!"
-  expr
+struct IteratorReturn{T}
+  value :: T
+  IteratorReturn(value) = new{typeof(value)}(value)
+end
+
+@inline function generate(fsm_iter::FiniteStateMachineIterator, send, state=nothing)
+  #fsm_iter._state = state
+  result = fsm_iter(send)
+  fsm_iter._state === 0xff && return IteratorReturn(result)
+  result, nothing
+end
+
+@inline function generate(iter, _)
+  ret = iterate(iter)
+  isnothing(ret) && return IteratorReturn(nothing)
+  ret
+end
+
+@inline function generate(iter, _, state)
+  ret = iterate(iter, state)
+  isnothing(ret) && return IteratorReturn(nothing)
+  ret
 end

test/runtests.jl

@@ -27,6 +27,7 @@ end
 println("Starting tests with $(Threads.nthreads()) threads out of `Sys.CPU_THREADS = $(Sys.CPU_THREADS)`...")
 
 @doset "main"
+@doset "yieldfrom"
 @doset "typeparams"
 VERSION >= v"1.8" && @doset "doctests"
 VERSION >= v"1.8" && @doset "aqua"

test/test_yieldfrom.jl

@@ -0,0 +1,55 @@
+using Test
+using ResumableFunctions
+
+@resumable function test_yieldfrom_inner(n)
+  for i in 1:n
+    @yield i^2
+  end
+  42, n
+end
+
+@resumable function test_yieldfrom(n)
+  @yieldfrom [42, 314]   
+  m, n = @yieldfrom test_yieldfrom_inner(n)
+  @test m == 42
+  @yield n
+  @yieldfrom test_yieldfrom_inner(n+1)
+end
+
+@testset "test_yieldfrom" begin
+@test collect(test_yieldfrom(4)) == [42, 314, 1, 4, 9, 16, 4, 1, 4, 9, 16, 25]
+end
+
+@resumable function test_echo()
+  x = 0
+  while true
+    x = @yield x
+  end
+  return "Done"
+end
+
+@resumable function test_forward()
+  ret = @yieldfrom test_echo()
+  @test ret == "Done"
+  @yieldfrom test_echo()
+end
+
+@testset "test_yieldfrom_twoway" begin
+  forward = test_forward()
+  @test forward() == 0
+  for i in 1:5
+    @test forward(i) == i
+  end
+end
+
+@testset "test_yieldfrom_nonresumable" begin
+  @test_throws LoadError eval(:(@yieldfrom [1,2,3]))
+end
+
+
+@testset "test_yieldfrom_noargs" begin
+  test_f = :(@resumable function test_yieldfrom_noargs()
+    @yieldfrom
+  end)
+  @test_throws LoadError eval(test_f)
+end