From be8039e135f47bf0d06aeb0e75e0dca7f90421cb Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Thu, 8 Aug 2024 21:22:00 +0200
Subject: [PATCH 01/23] fix: scoping

---
 docs/src/manual.md |   3 +-
 src/macro.jl       |  34 ++++++++--
 src/transforms.jl  |  81 +++++++++++++++++++-----
 src/utils.jl       | 150 +++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl   |   2 +
 test/test_main.jl  |  35 ++++++-----
 6 files changed, 268 insertions(+), 37 deletions(-)

diff --git a/docs/src/manual.md b/docs/src/manual.md
index a1320ec..70d6429 100644
--- a/docs/src/manual.md
+++ b/docs/src/manual.md
@@ -373,6 +373,7 @@ DocTestSetup = quote
   @resumable function arrays_of_tuples()
     for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
       for i in 1:2
+        local val
         let i=i
           val = [a[i] for a in u]
         end
@@ -413,4 +414,4 @@ DocTestSetup = nothing
 - In a `try` block only top level `@yield` statements are allowed.
 - In a `catch` or a `finally` block a `@yield` statement is not allowed.
 - An anonymous function can not contain a `@yield` statement.
-- If a `FiniteStateMachineIterator` object is used in more than one `for` loop, only the `state` variable is reinitialised. A `@resumable function` that alters its arguments will use the modified values as initial parameters.
\ No newline at end of file
+- If a `FiniteStateMachineIterator` object is used in more than one `for` loop, only the `state` variable is reinitialised. A `@resumable function` that alters its arguments will use the modified values as initial parameters.
diff --git a/src/macro.jl b/src/macro.jl
index 9e37a19..400684b 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -73,6 +73,7 @@ macro resumable(ex::Expr...)
 
   # The function that executes a step of the finite state machine
   func_def = splitdef(expr)
+  @debug func_def[:body]
   rtype = :rtype in keys(func_def) ? func_def[:rtype] : Any
   args, kwargs, arg_dict = get_args(func_def)
   params = ((get_param_name(param) for param in func_def[:whereparams])...,)
@@ -82,7 +83,35 @@ macro resumable(ex::Expr...)
   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])
+  @debug func_def[:body]|>MacroTools.striplines
+  #func_def[:body] = postwalk(x->transform_macro(x), func_def[:body])
+  #@debug func_def[:body]|>MacroTools.striplines
+  #func_def[:body] = postwalk(x->transform_macro_undo(x), func_def[:body])
+  #@debug func_def[:body]|>MacroTools.striplines
+  #func_def[:body] = postwalk(x->transform_let(x), func_def[:body])
+  #@info func_def[:body]|>MacroTools.striplines
+  #func_def[:body] = postwalk(x->transform_local(x), func_def[:body])
+  #@info func_def[:body]|>MacroTools.striplines
+  # Scoping fixes
+
+  # :name is :(fA::A) if it is an overloading call function (fA::A)(...)
+  # ...
+  if func_def[:name] isa Expr
+    @assert func_def[:name].head == :(::)
+    _name = func_def[:name].args[1]
+  else
+    _name = func_def[:name]
+  end
+  
+  scope = ScopeTracker(0, __module__, [Dict(i =>i for i in vcat(args, kwargs, [_name], params...))])
+  #@info func_def[:body]|>MacroTools.striplines
+  func_def[:body] = scoping(copy(func_def[:body]), scope)
+  #@info func_def[:body]|>MacroTools.striplines
+  func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
+  #@info func_def[:body]|>MacroTools.striplines
+
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
+  @debug slots
 
   # check if the resumable function is a callable struct instance (a functional) that is referencing itself
   isfunctional = @capture(func_def[:name], functional_::T_) && inexpr(func_def[:body], functional)
@@ -113,7 +142,6 @@ macro resumable(ex::Expr...)
     fsmi._state = 0x00
     fsmi
   end
-
   # the bare/fallback version of the constructor supplies default slot type parameters
   # we only need to define this if there there are actually slot defaults to be filled
   if !isempty(slot_T)
@@ -139,7 +167,6 @@ macro resumable(ex::Expr...)
     end
   )
   @debug type_expr|>MacroTools.striplines
-
   # The "original" function that now is simply a wrapper around the construction of the finite state machine
   call_def = copy(func_def)
   call_def[:rtype] = nothing
@@ -158,7 +185,7 @@ macro resumable(ex::Expr...)
   end
   call_expr = combinedef(call_def) |> flatten
   @debug call_expr|>MacroTools.striplines
-
+  
   # Finalizing the function stepping through the finite state machine
   if isempty(params)
     func_def[:name] = :((_fsmi::$type_name))
@@ -209,7 +236,6 @@ macro resumable(ex::Expr...)
     call_expr = postwalk(x->x==:(ResumableFunctions.typed_fsmi) ? :(ResumableFunctions.typed_fsmi_fallback) : x, call_expr)
   end
   @debug func_expr|>MacroTools.striplines
-
   # The final expression:
   # - the finite state machine struct
   # - the function stepping through the states
diff --git a/src/transforms.jl b/src/transforms.jl
index 8105a43..9bb988f 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -1,3 +1,20 @@
+function transform_remove_local(ex)
+  ex isa Expr && ex.head === :local && return Expr(:block)
+  return ex
+end
+
+function transform_macro(ex)
+  ex isa Expr || return ex
+  ex.head !== :macrocall && return ex
+  return Expr(:call, :__secret__, ex.args)
+end
+
+function transform_macro_undo(ex)
+  ex isa Expr || return ex
+  (ex.head !== :call || ex.args[1] !== :__secret__) && return ex
+  return Expr(:macrocall, ex.args[2]...)
+end
+
 """
 Function that replaces a variable
 """
@@ -70,22 +87,26 @@ Function that replaces a `for` loop by a corresponding `while` loop saving expli
 """
 function transform_for(expr, ui8::BoxedUInt8)
   @capture(expr, for element_ in iterator_ body_ end) || return expr
+  #@info element
+  localelement = Expr(:local, element)
   ui8.n += one(UInt8)
   next = Symbol("_iteratornext_", ui8.n)
   state = Symbol("_iterstate_", ui8.n)
   iterator_value = Symbol("_iterator_", ui8.n)
   label = Symbol("_iteratorlabel_", ui8.n)
   body = postwalk(x->transform_continue(x, label), :(begin $(body) end))
-  quote
+  res = quote
     $iterator_value = $iterator
     @nosave $next = iterate($iterator_value)
     while $next !== nothing
+      $localelement
       ($element, $state) = $next
       $body
       @label $label
       $next = iterate($iterator_value, $state)
     end
   end
+  res
 end
 
 
@@ -102,7 +123,7 @@ Function that replaces a variable `x` in an expression by `_fsmi.x` where `x` is
 """
 function transform_slots(expr, symbols)
   expr isa Expr || return expr
-  expr.head === :let && return transform_slots_let(expr, symbols)
+  #expr.head === :let && return transform_slots_let(expr, symbols)
   for i in 1:length(expr.args)
     expr.head === :kw && i === 1 && continue
     expr.head === Symbol("quote") && continue
@@ -111,27 +132,53 @@ function transform_slots(expr, symbols)
   expr
 end
 
-"""
-Function that handles `let` block
-"""
-function transform_slots_let(expr::Expr, symbols)
-  @capture(expr, let vars_; body_ end)
-  locals = Set{Symbol}()
-  (isa(vars, Expr) && vars.head==:(=))  || error("@resumable currently supports only single variable declarations in let blocks, i.e. only let blocks exactly of the form `let i=j; ...; end`. If you need multiple variables, please submit an issue on the issue tracker and consider contributing a patch.")
-  sym = vars.args[1].args[2].value
-  push!(locals, sym)
-  vars.args[1] = sym
-  body = postwalk(x->transform_let(x, locals), :(begin $(body) end))
-  :(let $vars; $body end)
+#"""
+#Function that handles `let` block
+#"""
+#function transform_slots_let(expr::Expr, symbols)
+#  @capture(expr, let vars_; body_ end)
+#  locals = Set{Symbol}()
+#  (isa(vars, Expr) && vars.head==:(=))  || error("@resumable currently supports only single variable declarations in let blocks, i.e. only let blocks exactly of the form `let i=j; ...; end`. If you need multiple variables, please submit an issue on the issue tracker and consider contributing a patch.")
+#  sym = vars.args[1].args[2].value
+#  push!(locals, sym)
+#  vars.args[1] = sym
+#  body = postwalk(x->transform_let(x, locals), :(begin $(body) end))
+#  :(let $vars; $body end)
+#end
+
+function transform_let(expr)
+  expr isa Expr || return expr
+  expr.head === :block && return expr
+  #@info "inside transform let"
+  @capture(expr, let arg_; body_; end) || return expr
+  #@info "captured let"
+  #arg |> dump
+  #@info expr
+  #@info arg
+  #error("ASds")
+  res = quote
+    let
+      local $arg
+      $body
+    end
+  end
+  #@info "emitting $res"
+  res
+  #expr.head === :. || return expr
+  #expr = expr.args[2].value in symbols ? :($(expr.args[2].value)) : expr
 end
 
 """
 Function that replaces a variable `_fsmi.x` in an expression by `x` where `x` is a variable declared in a `let` block.
 """
-function transform_let(expr, symbols::Set{Symbol})
+function transform_local(expr)
   expr isa Expr || return expr
-  expr.head === :. || return expr
-  expr = expr.args[2].value in symbols ? :($(expr.args[2].value)) : expr
+  @capture(expr, local arg_ = ex_) || return expr
+  res = quote
+    local $arg
+    $arg = $ex
+  end
+  res
 end
 
 """
diff --git a/src/utils.jl b/src/utils.jl
index 2f09d83..8b132d9 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -208,3 +208,153 @@ end
 function typed_fsmi_fallback(fsmi::Type{T}, fargs...)::T where T
   return T()
 end
+
+mutable struct ScopeTracker
+  i::Int
+  mod::Module
+  scope_stack::Vector
+end
+
+function lookup!(s::Symbol, S::ScopeTracker; new = false)
+  if isdefined(S.mod, s)
+    return s
+  end
+  if !new
+    for D in Iterators.reverse(S.scope_stack)
+      if haskey(D, s)
+        return D[s]
+      end
+    end
+  end
+  D = last(S.scope_stack)
+  new = Symbol(s, Symbol("_$(S.i)"))
+  S.i += 1
+  D[s] = new
+  return new
+end
+
+scoping(e::LineNumberNode, scope) = e
+
+scoping(e::Int, scope) = e
+
+scoping(e::String, scope) = e
+scoping(e::typeof(ResumableFunctions.generate), scope) = e
+scoping(e::typeof(ResumableFunctions.IteratorReturn), scope) = e
+scoping(e::QuoteNode, scope) = e
+scoping(e::Bool, scope) = e
+scoping(e::Nothing, scope) = e
+
+function scoping(s::Symbol, scope; new = false)
+  #@info "scoping $s, $new"
+  return lookup!(s, scope; new = new)
+end
+
+function scoping(expr::Expr, scope)
+  if expr.head === :macrocall
+    for i in 2:length(expr.args)
+      expr.args[i] = scoping(expr.args[i], scope)
+    end
+    return expr
+  end
+  new_stack = false
+  if expr.head === :let
+    # Replace 
+    #   let i, k = 2, j = 1
+    #      [...]
+    #   end
+    #
+    #   by
+    #
+    #   let
+    #     local i_new
+    #     local k_new = 2
+    #     local j_new = 1
+    #   end
+    #
+    #   Caveat:
+    #   let i = i, j = i
+    #
+    #   must be
+    #   new_i = old_i
+    #   new_j = new_i
+    #
+    #   :(
+
+    # defer adding a new scope after the right hand side have been renamed
+    @capture(expr, let arg_; body_ end) || return expr
+    @capture(arg, begin x__ end)
+    replace_rhs = []
+    for i in 1:length(x)
+      y = x[i]
+      fl = @capture(y, k_ = v_)
+      if fl
+        push!(replace_rhs, scoping(v, scope))
+      else
+        # there was no right side
+        push!(replace_rhs, nothing)
+      end
+    end
+    new_stack = true
+    push!(scope.scope_stack, Dict())
+    replace_lhs = []
+    rep = []
+    for i in 1:length(x)
+      y = x[i]
+      fl = @capture(y, k_ = v_)
+      if fl
+        push!(replace_lhs, scoping(k, scope, new = true))
+        push!(rep, quote local $(replace_lhs[i]); $(replace_lhs[i]) = $(replace_rhs[i]) end)
+      else
+        @assert y isa Symbol
+        push!(replace_lhs, scoping(y, scope, new = true))
+        push!(rep, quote local $(replace_lhs[i]) end)
+      end
+    end
+    rep = quote
+      $(rep...)
+    end
+    rep = MacroTools.flatten(rep)
+    expr.args[1] = Expr(:block)
+    pushfirst!(expr.args[2].args, rep)
+
+    # Now continue recursively
+    # but skip the local/dance, since we already replaced them
+    for i in 2:length(expr.args[2].args)
+      a = expr.args[2].args[i]
+      expr.args[2].args[i] = scoping(a, scope)
+    end
+    pop!(scope.scope_stack)
+    return expr
+  end
+
+  if expr.head === :while || expr.head === :let
+    push!(scope.scope_stack, Dict())
+    new_stack = true
+  end
+  if expr.head === :local
+    # this is my local dance
+    # explain and rewrite using @capture
+    if length(expr.args) == 1 && expr.args[1] isa Symbol
+      expr.args[1] = scoping(expr.args[1], scope, new = true)
+    elseif length(expr.args) == 1 && expr.args[1].head === :tuple
+      for i in 1:length(expr.args[1].args)
+        a = expr.args[1].args[i]
+        expr.args[1].args[i] = scoping(a, scope, new = true)
+      end
+    else
+      for i in 1:length(expr.args)
+        a = expr.args[i]
+        expr.args[i] = scoping(a, scope, new = true)
+      end
+    end
+  else
+    for i in 1:length(expr.args)
+      a = expr.args[i]
+      expr.args[i] = scoping(a, scope)
+    end
+  end
+  if new_stack
+    pop!(scope.scope_stack)
+  end
+  return expr
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 10ce0d2..a38ef5f 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -16,6 +16,8 @@ end
 
 macro doset(descr)
     quote
+        @info "====================================="
+        @info $descr
         if doset($descr)
             @safetestset $descr begin
                 include("test_"*$descr*".jl")
diff --git a/test/test_main.jl b/test/test_main.jl
index d169924..722e66e 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -83,6 +83,7 @@ end
 @resumable function test_let()
   for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
     for i in 1:2
+      local val
       let i=i
         val = [a[i] for a in u]
       end
@@ -97,10 +98,10 @@ end
 
 if VERSION >= v"1.8"
 
-q_test_let2 = quote
 @resumable function test_let2()
   for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
     for i in 1:2
+      local val
       let i=i, j=i
         val = [(a[i],a[j]) for a in u]
       end
@@ -108,18 +109,15 @@ q_test_let2 = quote
     end
   end
 end
-end
 
 @testset "test_let2" begin
-#@test collect(test_let2()) == ...
-@test_throws "@resumable currently supports only single" eval(q_test_let2)
-@test_broken false # the test above throws an error when it should not -- it happens because we do not support variables without assignments in let blocks -- see issues #69 and #70
+@test collect(test_let2()) == [[(1, 1), (3, 3)], [(2, 2), (4, 4)], [(5, 5), (7, 7)], [(6, 6), (8, 8)]]
 end
 
-q_test_let_noassignment = quote
 @resumable function test_let_noassignment()
   for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
     for i in 1:2
+      local val
       let i
         val = [a[1] for a in u]
       end
@@ -127,18 +125,15 @@ q_test_let_noassignment = quote
     end
   end
 end
-end
 
 @testset "test_let_noassignment" begin
-#@test collect(test_let_noassignment()) == ...
-@test_throws "@resumable currently supports only single" eval(q_test_let_noassignment)
-@test_broken false # the test above throws an error when it should not -- it happens because we do not support variables without assignments in let blocks -- see issues #69 and #70
+collect(test_let_noassignment()) == [[1, 3], [1, 3], [5, 7], [5, 7]]
 end
 
-q_test_let_multipleargs = quote
 @resumable function test_let_multipleargs()
   for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
     for i in 1:2
+      local val
       let i=i, j
         val = [a[i] for a in u]
       end
@@ -146,12 +141,9 @@ q_test_let_multipleargs = quote
     end
   end
 end
-end
 
 @testset "test_let_multipleargs" begin
-#@test collect(test_let_multipleargs()) == ...
-@test_throws "@resumable currently supports only single" eval(q_test_let_noassignment)
-@test_broken false # the test above throws an error when it should not -- it happens because we do not support variables without assignments in let blocks -- see issues #69 and #70
+@test collect(test_let_multipleargs()) == [[1, 3], [2, 4], [5, 7], [6, 8]]
 end
 
 end # VERSION >= v"1.8"
@@ -228,6 +220,19 @@ end
   @test collect(test_unstable(3)) == ["number 1", "number 2", "number 3"]
 end
 
+@testset "test_scope" begin
+@resumable function test_scope_throws()
+  for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
+    for i in 1:2
+      let i=i, j
+        val = [a[i] for a in u]
+      end
+      @yield val
+    end
+  end
+end
+@test_throws UndefVarError collect(test_scope_throws())
+
 # test length
 
 @testset "test_length" begin

From 9f4cc8faefdb2f17832eca356ab9b0e84536f79b Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Fri, 9 Aug 2024 21:15:24 +0200
Subject: [PATCH 02/23] better

---
 src/macro.jl      |  3 +-
 src/utils.jl      | 93 ++++++++++++++++++++++++++++++++++++++++++++---
 test/runtests.jl  | 12 +++---
 test/test_main.jl | 12 ++++++
 4 files changed, 107 insertions(+), 13 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index 400684b..577dfce 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -91,7 +91,6 @@ macro resumable(ex::Expr...)
   #func_def[:body] = postwalk(x->transform_let(x), func_def[:body])
   #@info func_def[:body]|>MacroTools.striplines
   #func_def[:body] = postwalk(x->transform_local(x), func_def[:body])
-  #@info func_def[:body]|>MacroTools.striplines
   # Scoping fixes
 
   # :name is :(fA::A) if it is an overloading call function (fA::A)(...)
@@ -105,10 +104,12 @@ macro resumable(ex::Expr...)
   
   scope = ScopeTracker(0, __module__, [Dict(i =>i for i in vcat(args, kwargs, [_name], params...))])
   #@info func_def[:body]|>MacroTools.striplines
+  #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = scoping(copy(func_def[:body]), scope)
   #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
   #@info func_def[:body]|>MacroTools.striplines
+  #@info func_def[:body]|>MacroTools.striplines
 
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
   @debug slots
diff --git a/src/utils.jl b/src/utils.jl
index 8b132d9..fa0616a 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -215,7 +215,44 @@ mutable struct ScopeTracker
   scope_stack::Vector
 end
 
+function lookup_lhs!(s::Symbol, S::ScopeTracker; new::Bool = false)
+  if !new
+    for D in Iterators.reverse(S.scope_stack)
+      if haskey(D, s)
+        return D[s]
+      end
+    end
+  end
+  D = last(S.scope_stack)
+  new = Symbol(s, Symbol("_$(S.i)"))
+  S.i += 1
+  D[s] = new
+  return new
+end
+
+function lookup_lhs!(s::Expr, S::ScopeTracker)
+  if s.head === :(.)
+    s.args[1] === lookup_lhs!(s.args[1], S)
+    return s
+  end
+  @assert s.head === :tuple
+  for i in 1:length(s.args)
+    s.args[i] = lookup_rhs!(s.args[i], S)
+  end
+  return s
+end
+
+function lookup_rhs!(s::Symbol, S::ScopeTracker)
+  for D in Iterators.reverse(S.scope_stack)
+    if haskey(D, s)
+      return D[s]
+    end
+  end
+  return s
+end
+
 function lookup!(s::Symbol, S::ScopeTracker; new = false)
+  s == :val && @show s, length(S.scope_stack), new
   if isdefined(S.mod, s)
     return s
   end
@@ -225,6 +262,9 @@ function lookup!(s::Symbol, S::ScopeTracker; new = false)
         return D[s]
       end
     end
+    D = last(S.scope_stack)
+    D[s] = s
+    return s
   end
   D = last(S.scope_stack)
   new = Symbol(s, Symbol("_$(S.i)"))
@@ -246,10 +286,45 @@ scoping(e::Nothing, scope) = e
 
 function scoping(s::Symbol, scope; new = false)
   #@info "scoping $s, $new"
-  return lookup!(s, scope; new = new)
+  return lookup_rhs!(s, scope)
 end
 
 function scoping(expr::Expr, scope)
+  if expr.head === :comprehension
+    # this is again special with respect to scoping
+    if expr.args[1].head === :generator
+      # first the generator case
+      for i in 2:length(expr.args[1].args)
+        @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :(=)
+        expr.args[1].args[i].args[2] = lookup_rhs!(expr.args[1].args[i].args[2], scope)
+      end
+      # now create new scope
+      push!(scope.scope_stack, Dict())
+      for i in 2:length(expr.args[1].args)
+        expr.args[1].args[i].args[1] = lookup_lhs!(expr.args[1].args[i].args[1], scope)
+      end
+
+      expr.args[1].args[1] = scoping(expr.args[1].args[1], scope)
+      pop!(scope.scope_stack)
+      return expr
+    else
+      error("not implemented yet")
+    end
+  end
+
+  if expr.head === :(=)
+    if expr.args[1] isa Symbol
+      expr.args[1] = lookup_lhs!(expr.args[1], scope)
+    else
+      for i in 1:length(expr.args[1].args)
+        expr.args[1].args[i] = lookup_lhs!(expr.args[1].args[i], scope)
+      end
+    end
+    for i in 2:length(expr.args)
+      expr.args[i] = scoping(expr.args[i], scope)
+    end
+    return expr
+  end
   if expr.head === :macrocall
     for i in 2:length(expr.args)
       expr.args[i] = scoping(expr.args[i], scope)
@@ -302,11 +377,11 @@ function scoping(expr::Expr, scope)
       y = x[i]
       fl = @capture(y, k_ = v_)
       if fl
-        push!(replace_lhs, scoping(k, scope, new = true))
+        push!(replace_lhs, lookup_lhs!(k, scope, new = true))
         push!(rep, quote local $(replace_lhs[i]); $(replace_lhs[i]) = $(replace_rhs[i]) end)
       else
         @assert y isa Symbol
-        push!(replace_lhs, scoping(y, scope, new = true))
+        push!(replace_lhs, lookup_lhs!(y, scope, new = true))
         push!(rep, quote local $(replace_lhs[i]) end)
       end
     end
@@ -334,17 +409,23 @@ function scoping(expr::Expr, scope)
   if expr.head === :local
     # this is my local dance
     # explain and rewrite using @capture
+    #
+    # if I see a local x or local x = ...
+    # we always emit a new identifier
     if length(expr.args) == 1 && expr.args[1] isa Symbol
-      expr.args[1] = scoping(expr.args[1], scope, new = true)
+      #expr.args[1] = scoping(expr.args[1], scope, new = true)
+      expr.args[1] = lookup_lhs!(expr.args[1], scope)
     elseif length(expr.args) == 1 && expr.args[1].head === :tuple
       for i in 1:length(expr.args[1].args)
         a = expr.args[1].args[i]
-        expr.args[1].args[i] = scoping(a, scope, new = true)
+        #expr.args[1].args[i] = scoping(a, scope, new = true)
+        expr.args[1].args[i] = lookup_lhs!(a, scope)
       end
     else
       for i in 1:length(expr.args)
         a = expr.args[i]
-        expr.args[i] = scoping(a, scope, new = true)
+        #expr.args[i] = scoping(a, scope, new = true)
+        expr.args[i] = lookup_lhs!(a, scope)
       end
     end
   else
diff --git a/test/runtests.jl b/test/runtests.jl
index a38ef5f..43bcb9c 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -16,13 +16,13 @@ end
 
 macro doset(descr)
     quote
-        @info "====================================="
-        @info $descr
-        if doset($descr)
-            @safetestset $descr begin
+#        @info "====================================="
+#        @info $descr
+#        if doset($descr)
+#            @safetestset $descr begin
                 include("test_"*$descr*".jl")
-            end
-        end
+#            end
+#        end
     end
 end
 
diff --git a/test/test_main.jl b/test/test_main.jl
index 722e66e..08c23d5 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -248,3 +248,15 @@ end
   @test length(collect(test_length(10, 20))) === 10^2 * 20^2
   @test Base.IteratorSize(typeof(test_length(1, 1))) == Base.HasLength()
 end
+
+@testset "test_scope_2" begin
+  @resumable function test_forward()
+    for i in 1:10
+      @yield test_bla(i)
+    end
+  end
+
+  test_bla(i) = i^2
+
+  @test collect(test_forward()) == [i^2 for i in 1:10]
+end

From 1059fd8eb9ee12fae31f8b82c3ce478f7b4e19bf Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Fri, 9 Aug 2024 21:39:08 +0200
Subject: [PATCH 03/23] small fix

---
 src/utils.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/utils.jl b/src/utils.jl
index fa0616a..2f831cc 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -276,6 +276,7 @@ end
 scoping(e::LineNumberNode, scope) = e
 
 scoping(e::Int, scope) = e
+scoping(e::Float64, scope) = e
 
 scoping(e::String, scope) = e
 scoping(e::typeof(ResumableFunctions.generate), scope) = e

From e87515046e5ba767183126ba1fd1d5d93774344c Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Fri, 9 Aug 2024 21:41:09 +0200
Subject: [PATCH 04/23] another

---
 src/utils.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/utils.jl b/src/utils.jl
index 2f831cc..e6da420 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -230,6 +230,10 @@ function lookup_lhs!(s::Symbol, S::ScopeTracker; new::Bool = false)
   return new
 end
 
+function lookup_lhs!(s::QuoteNode, S::ScopeTracker)
+  return s
+end
+
 function lookup_lhs!(s::Expr, S::ScopeTracker)
   if s.head === :(.)
     s.args[1] === lookup_lhs!(s.args[1], S)

From b9497b590efdec9624392f84089e798ae946cccf Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Fri, 9 Aug 2024 22:26:42 +0200
Subject: [PATCH 05/23] kw

---
 src/utils.jl      | 25 +++++++++++++++++++++++++
 test/test_main.jl | 11 +++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/utils.jl b/src/utils.jl
index e6da420..93ef9a0 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -295,6 +295,31 @@ function scoping(s::Symbol, scope; new = false)
 end
 
 function scoping(expr::Expr, scope)
+  if expr.head === :call
+    @info expr
+    # We have to not rename the keyword arguments
+    # Super awkward because of f(x, y = z, w) and f(x; y) is allowed :(
+    # Or even f(x, y = 1, w, z = 2)
+    for i in 2:length(expr.args)
+      if expr.args[i] isa Expr && expr.args[i].head === :kw
+        # this is f(..., x = 2, ...)
+        expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
+      elseif expr.args[i] isa Expr && expr.args[i].head === :paramaters
+        for j in 1:length(expr.args[i].args)
+          if expr.args[i].args[j] isa Symbol
+            # this is f(...; x)
+          else
+            @assert expr.args[i].args[j] isa Expr && expr.args[i].args[j].head === :kw
+            # this is f(...; x = 2)
+            expr.args[i].args[j].args[2] = scoping(expr.args[i].args[j].args[2], scope)
+          end
+        end
+      else
+        expr.args[i] = scoping(expr.args[i], scope)
+      end
+    end
+    return expr
+  end
   if expr.head === :comprehension
     # this is again special with respect to scoping
     if expr.args[1].head === :generator
diff --git a/test/test_main.jl b/test/test_main.jl
index 08c23d5..622bcf7 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -260,3 +260,14 @@ end
 
   @test collect(test_forward()) == [i^2 for i in 1:10]
 end
+
+@testset "test_kw" begin
+  g(x, y; z = 2) = x + y^2 + z^C
+
+  @resumable function f(z)
+    y = 1
+    @yield g(z, z = y, 2)
+  end
+  
+  @test collect(test_kw()) == [8]
+end

From cbc6560deb28f689dfb2ac3e0c6bf3c04591056c Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Fri, 9 Aug 2024 22:50:27 +0200
Subject: [PATCH 06/23] call renaming

---
 src/utils.jl      |  5 ++++-
 test/test_main.jl | 21 ++++++++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index 93ef9a0..19a6529 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -246,6 +246,8 @@ function lookup_lhs!(s::Expr, S::ScopeTracker)
   return s
 end
 
+lookup_rhs!(e::typeof(ResumableFunctions.generate), scope) = e
+
 function lookup_rhs!(s::Symbol, S::ScopeTracker)
   for D in Iterators.reverse(S.scope_stack)
     if haskey(D, s)
@@ -296,7 +298,8 @@ end
 
 function scoping(expr::Expr, scope)
   if expr.head === :call
-    @info expr
+    # Rename the name
+    expr.args[1] = lookup_rhs!(expr.args[1], scope)
     # We have to not rename the keyword arguments
     # Super awkward because of f(x, y = z, w) and f(x; y) is allowed :(
     # Or even f(x, y = 1, w, z = 2)
diff --git a/test/test_main.jl b/test/test_main.jl
index 622bcf7..21aaafc 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -262,12 +262,27 @@ end
 end
 
 @testset "test_kw" begin
-  g(x, y; z = 2) = x + y^2 + z^C
+  g(x, y; z = 2) = x + y^2 + z
 
-  @resumable function f(z)
+  @resumable function test_kw(z)
     y = 1
     @yield g(z, z = y, 2)
   end
   
-  @test collect(test_kw()) == [8]
+  @test collect(test_kw(3)) == [8]
+end
+
+@testset "test_call_renaming" begin
+  g(x) = x^2
+
+  @resumable function test_call_renaming(y)
+    sin = g
+    let g = 3
+      for h in 1:10
+        @yield sin(g + h + y)
+      end
+    end
+  end
+
+  @test collect(test_call_renaming(3)) == [49, 64, 81, 100, 121, 144, 169, 196, 225, 256]
 end

From 827e3d9ba8b3408edc8ee11a13843c28137d513c Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 10 Aug 2024 11:29:18 +0200
Subject: [PATCH 07/23] comprehension

---
 src/macro.jl      |  1 -
 src/utils.jl      | 87 +++++++++++++++++++++++++++++++++--------------
 test/test_main.jl | 47 +++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 26 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index 577dfce..db4e25e 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -109,7 +109,6 @@ macro resumable(ex::Expr...)
   #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
   #@info func_def[:body]|>MacroTools.striplines
-  #@info func_def[:body]|>MacroTools.striplines
 
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
   @debug slots
diff --git a/src/utils.jl b/src/utils.jl
index 19a6529..29a2a48 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -76,8 +76,10 @@ function get_slots(func_def::Dict, args::Dict{Symbol, Any}, mod::Module)
   # eval function
   func_expr = combinedef(func_def) |> flatten
   @eval(mod, @noinline $func_expr)
+  #@info func_def[:body]|>MacroTools.striplines
   # get typed code
   codeinfos = @eval(mod, code_typed($(func_def[:name]), Tuple; optimize=false))
+  #@info codeinfos
   # extract slot names and types
   for codeinfo in codeinfos
     for (name, type) in collect(zip(codeinfo.first.slotnames, codeinfo.first.slottypes))
@@ -234,14 +236,14 @@ function lookup_lhs!(s::QuoteNode, S::ScopeTracker)
   return s
 end
 
-function lookup_lhs!(s::Expr, S::ScopeTracker)
+function lookup_lhs!(s::Expr, S::ScopeTracker; new = false)
   if s.head === :(.)
-    s.args[1] === lookup_lhs!(s.args[1], S)
+    s.args[1] === lookup_lhs!(s.args[1], S; new = new)
     return s
   end
   @assert s.head === :tuple
   for i in 1:length(s.args)
-    s.args[i] = lookup_rhs!(s.args[i], S)
+    s.args[i] = lookup_lhs!(s.args[i], S; new = new)
   end
   return s
 end
@@ -296,10 +298,54 @@ function scoping(s::Symbol, scope; new = false)
   return lookup_rhs!(s, scope)
 end
 
+function scope_generator(expr, scope)
+  @assert expr.head === :generator
+  # first the generator case
+  for i in 2:length(expr.args)
+    @assert expr.args[i] isa Expr && expr.args[i].head === :(=)
+    #expr.args[i].args[2] = lookup_rhs!(expr.args[i].args[2], scope)
+    expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
+  end
+  # now create new scope
+  push!(scope.scope_stack, Dict())
+  for i in 2:length(expr.args)
+    expr.args[i].args[1] = lookup_lhs!(expr.args[i].args[1], scope, new = true)
+  end
+
+  expr.args[1] = scoping(expr.args[1], scope)
+  pop!(scope.scope_stack)
+  return expr
+end
+
 function scoping(expr::Expr, scope)
+  if expr.head === :generator
+    return scope_generator(expr, scope)
+  end
+
+  # Named tuple handling is again pretty awkward
+  # because of the (;a, b) syntax, which we have to expand by hand to
+  # (;a = a, b = b), otherwise we do (;a_1, b_2), and this gets
+  # (;a_1 = a_1, b_2 = b_2) which is nonsensical
+  if expr.head === :tuple && length(expr.args) > 0 && expr.args[1] isa Expr && expr.args[1].head === :parameters
+    # this is a named tuple of the form (;...)
+    # TODO: named tuple recognition not working properly yet
+    # first bring (;...,b,...) in the form (;...,b => b,...)
+    for i in 1:length(expr.args[1].args)
+      if !(expr.args[1].args[i] isa Expr)
+        expr.args[1].args[i] = Expr(:kw, expr.args[1].args[i], expr.args[1].args[i])
+      end
+    end
+    # Now rename the RHS
+    for i in 1:length(expr.args[1].args)
+      @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :kw
+      expr.args[1].args[i].args[2] = scoping(expr.args[1].args[i].args[2], scope)
+    end
+    return expr
+  end
+
   if expr.head === :call
-    # Rename the name
-    expr.args[1] = lookup_rhs!(expr.args[1], scope)
+    # Rename the caller name
+    expr.args[1] = scoping(expr.args[1], scope)
     # We have to not rename the keyword arguments
     # Super awkward because of f(x, y = z, w) and f(x; y) is allowed :(
     # Or even f(x, y = 1, w, z = 2)
@@ -307,16 +353,19 @@ function scoping(expr::Expr, scope)
       if expr.args[i] isa Expr && expr.args[i].head === :kw
         # this is f(..., x = 2, ...)
         expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
-      elseif expr.args[i] isa Expr && expr.args[i].head === :paramaters
+      elseif expr.args[i] isa Expr && expr.args[i].head === :parameters
+        # this is f(...;...)
+        # first normalize f(...;...,x...) to f(...;..., x = x,...)
         for j in 1:length(expr.args[i].args)
-          if expr.args[i].args[j] isa Symbol
-            # this is f(...; x)
-          else
-            @assert expr.args[i].args[j] isa Expr && expr.args[i].args[j].head === :kw
-            # this is f(...; x = 2)
-            expr.args[i].args[j].args[2] = scoping(expr.args[i].args[j].args[2], scope)
+          if !(expr.args[i].args[j] isa Expr)
+            expr.args[i].args[j] = Expr(:kw, expr.args[i].args[j], expr.args[i].args[j])
           end
         end
+        for j in 1:length(expr.args[i].args)
+          @assert expr.args[i].args[j] isa Expr && expr.args[i].args[j].head === :kw
+          # this is f(...; x = 2)
+          expr.args[i].args[j].args[2] = scoping(expr.args[i].args[j].args[2], scope)
+        end
       else
         expr.args[i] = scoping(expr.args[i], scope)
       end
@@ -326,19 +375,7 @@ function scoping(expr::Expr, scope)
   if expr.head === :comprehension
     # this is again special with respect to scoping
     if expr.args[1].head === :generator
-      # first the generator case
-      for i in 2:length(expr.args[1].args)
-        @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :(=)
-        expr.args[1].args[i].args[2] = lookup_rhs!(expr.args[1].args[i].args[2], scope)
-      end
-      # now create new scope
-      push!(scope.scope_stack, Dict())
-      for i in 2:length(expr.args[1].args)
-        expr.args[1].args[i].args[1] = lookup_lhs!(expr.args[1].args[i].args[1], scope)
-      end
-
-      expr.args[1].args[1] = scoping(expr.args[1].args[1], scope)
-      pop!(scope.scope_stack)
+      expr.args[1] = scope_generator(expr.args[1], scope)
       return expr
     else
       error("not implemented yet")
diff --git a/test/test_main.jl b/test/test_main.jl
index 21aaafc..a68d2de 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -270,6 +270,16 @@ end
   end
   
   @test collect(test_kw(3)) == [8]
+
+  g(z; y) = z - y
+
+  @resumable function test_kw_2(x)
+    for y in 1:10
+      @yield g(x; y)
+    end
+  end
+
+  @test collect(test_kw_2(2)) == [1, 0, -1, -2, -3, -4, -5, -6, -7, -8]
 end
 
 @testset "test_call_renaming" begin
@@ -286,3 +296,40 @@ end
 
   @test collect(test_call_renaming(3)) == [49, 64, 81, 100, 121, 144, 169, 196, 225, 256]
 end
+
+@testset "test_quotenode" begin
+  @resumable function test_quotenode(x)
+    @yield x.a^2 + x.b^2
+  end
+
+  @test collect((test_quotenode((a = 3, b = 4)))) == [5^2]
+end
+
+@testset "test_named_tuple" begin
+  @resumable function test_named_tuple(u, v)
+    r = @NamedTuple{a::Int, b::Int}[]
+    for a in u
+      for b in v
+        push!(r, (;a, b))
+        @yield (;a, b)
+      end
+    end
+    @yield r[2]
+  end
+
+  @test collect(test_named_tuple([1, 2], [3, 4])) == [(a = 1, b = 3), (a = 1, b = 4), (a = 2, b = 3), (a = 2, b = 4), (a = 1, b = 4)]
+end
+
+@testset "test_comprehension" begin
+  @resumable function test_comprehension(u)
+    r = Dict{Int, Int}(c =>i for (i, c) in u)
+    s = [u^2 for u in first.(u)]
+    for k in sort(collect(keys(r)))
+      @yield k
+    end
+    for s in s
+      @yield s
+    end
+  end
+  @test collect(test_comprehension([(1, 2), (3, 4)])) == [2,4,1,9]
+end

From 5f087ffe5bcd5d0f87a9081fea1843acf4d41ae5 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 10 Aug 2024 15:30:41 +0200
Subject: [PATCH 08/23] named tuples again and generators

---
 src/macro.jl      |  2 +-
 src/utils.jl      | 69 +++++++++++++++++++++++++++++++----------------
 test/test_main.jl | 35 ++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 24 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index db4e25e..4a236c3 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -108,7 +108,7 @@ macro resumable(ex::Expr...)
   func_def[:body] = scoping(copy(func_def[:body]), scope)
   #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
-  #@info func_def[:body]|>MacroTools.striplines
+  @info func_def[:body]|>MacroTools.striplines
 
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
   @debug slots
diff --git a/src/utils.jl b/src/utils.jl
index 29a2a48..8ab1be6 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -238,14 +238,20 @@ end
 
 function lookup_lhs!(s::Expr, S::ScopeTracker; new = false)
   if s.head === :(.)
-    s.args[1] === lookup_lhs!(s.args[1], S; new = new)
+    s.args[1] = lookup_lhs!(s.args[1], S; new = new)
     return s
   end
-  @assert s.head === :tuple
-  for i in 1:length(s.args)
-    s.args[i] = lookup_lhs!(s.args[i], S; new = new)
+  if s.head  === :tuple
+    for i in 1:length(s.args)
+      s.args[i] = lookup_lhs!(s.args[i], S; new = new)
+    end
+    return s
   end
-  return s
+  if s.head === :ref
+    s = scoping(s, S)
+    return s
+  end
+  error("Not captured")
 end
 
 lookup_rhs!(e::typeof(ResumableFunctions.generate), scope) = e
@@ -326,21 +332,35 @@ function scoping(expr::Expr, scope)
   # because of the (;a, b) syntax, which we have to expand by hand to
   # (;a = a, b = b), otherwise we do (;a_1, b_2), and this gets
   # (;a_1 = a_1, b_2 = b_2) which is nonsensical
-  if expr.head === :tuple && length(expr.args) > 0 && expr.args[1] isa Expr && expr.args[1].head === :parameters
-    # this is a named tuple of the form (;...)
-    # TODO: named tuple recognition not working properly yet
-    # first bring (;...,b,...) in the form (;...,b => b,...)
-    for i in 1:length(expr.args[1].args)
-      if !(expr.args[1].args[i] isa Expr)
-        expr.args[1].args[i] = Expr(:kw, expr.args[1].args[i], expr.args[1].args[i])
+  if expr.head === :tuple
+    if length(expr.args) > 0 && expr.args[1] isa Expr && expr.args[1].head === :parameters
+      # this is a named tuple of the form (;...)
+      # TODO: named tuple recognition not working properly yet
+      # first bring (;...,b,...) in the form (;...,b => b,...)
+      for i in 1:length(expr.args[1].args)
+        if !(expr.args[1].args[i] isa Expr)
+          expr.args[1].args[i] = Expr(:kw, expr.args[1].args[i], expr.args[1].args[i])
+        end
       end
+      # Now rename the RHS
+      for i in 1:length(expr.args[1].args)
+        @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :kw
+        expr.args[1].args[i].args[2] = scoping(expr.args[1].args[i].args[2], scope)
+      end
+      return expr
+    elseif any(a -> a isa Expr && a.head === :(=), expr.args[1:end])
+      # Can be any of (a = 2, b, c = d)
+      # lets first normalize the entries of the form b to b => scoping(b, ...)
+      for i in 2:length(expr.args)
+        if expr.args[i] isa Symbol
+          expr.args[i] = Expr(:(=), expr.args[i], lookup_lhs!(expr.args[i], scope))
+        else
+          @assert expr.args[i].head === :(=)
+          expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
+        end
+      end
+      return expr
     end
-    # Now rename the RHS
-    for i in 1:length(expr.args[1].args)
-      @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :kw
-      expr.args[1].args[i].args[2] = scoping(expr.args[1].args[i].args[2], scope)
-    end
-    return expr
   end
 
   if expr.head === :call
@@ -498,12 +518,15 @@ function scoping(expr::Expr, scope)
         expr.args[i] = lookup_lhs!(a, scope)
       end
     end
-  else
-    for i in 1:length(expr.args)
-      a = expr.args[i]
-      expr.args[i] = scoping(a, scope)
-    end
+    return expr
   end
+
+  # default
+  for i in 1:length(expr.args)
+    a = expr.args[i]
+    expr.args[i] = scoping(a, scope)
+  end
+
   if new_stack
     pop!(scope.scope_stack)
   end
diff --git a/test/test_main.jl b/test/test_main.jl
index a68d2de..714ce19 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -318,6 +318,14 @@ end
   end
 
   @test collect(test_named_tuple([1, 2], [3, 4])) == [(a = 1, b = 3), (a = 1, b = 4), (a = 2, b = 3), (a = 2, b = 4), (a = 1, b = 4)]
+
+  @resumable function test_32()
+    x = 0
+    @yield (x = 1, )
+  end
+
+  @test collect(test_32()) == [(x = 1, )]
+
 end
 
 @testset "test_comprehension" begin
@@ -333,3 +341,30 @@ end
   end
   @test collect(test_comprehension([(1, 2), (3, 4)])) == [2,4,1,9]
 end
+
+@testset "test_ref" begin
+  @resumable function test_ref(x)
+    y = x
+    a = [i^2 for i in 1:3]
+    for i in 1:3
+      y[] = a[i]
+      @yield y[]
+    end
+  end
+  @test collect(test_ref(Ref(1))) == [1,4,9]
+end
+
+@testset "test_getproperty" begin
+  @resumable function test_getproperty()
+    let
+      node = (a = 1, b = 2)
+      v = [[2], node]
+      let node = (a = 2, b = 1)
+        (v[node.a])[node.b] == 3
+      end
+      @yield v
+    end
+  end
+
+  @test collect(test_getproperty()) == [[[2], (a = 1, b = 2)]]
+end

From 27d1d524bf4e7ad7347d6b456cfd2ce4530aeaf2 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 10 Aug 2024 15:31:31 +0200
Subject: [PATCH 09/23] "fix" jet test :)

---
 src/macro.jl     | 2 +-
 test/test_jet.jl | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index 4a236c3..db4e25e 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -108,7 +108,7 @@ macro resumable(ex::Expr...)
   func_def[:body] = scoping(copy(func_def[:body]), scope)
   #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
-  @info func_def[:body]|>MacroTools.striplines
+  #@info func_def[:body]|>MacroTools.striplines
 
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
   @debug slots
diff --git a/test/test_jet.jl b/test/test_jet.jl
index cad1e53..aa9d319 100644
--- a/test/test_jet.jl
+++ b/test/test_jet.jl
@@ -8,7 +8,6 @@ using Test
             Core.Compiler,
         )
     )
-    @show rep
-    @test length(JET.get_reports(rep)) <= 5
+    @test length(JET.get_reports(rep)) <= 7
     @test_broken length(JET.get_reports(rep)) == 0
 end

From e923833f1d2dec8a0ce9cac4595b7e450882f2f5 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 10 Aug 2024 19:29:39 +0200
Subject: [PATCH 10/23] handle all for declarations

---
 src/transforms.jl | 20 +++++++++++++++++++-
 test/test_main.jl | 10 ++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/transforms.jl b/src/transforms.jl
index 9bb988f..59591b2 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -84,10 +84,28 @@ end
 
 """
 Function that replaces a `for` loop by a corresponding `while` loop saving explicitly the *iterator* and its *state*.
+
+For loops of the form `for a, b, c; body; end` are denested.
 """
 function transform_for(expr, ui8::BoxedUInt8)
+  (expr isa Expr && expr.head === :for) || return expr
+  # test for simple for a in b expression
+  expr.args[1].head === :(=) && return transform_for_inner(expr, ui8)
+  # must be a complicated iteration
+  @assert expr.args[1].head === :block
+  body = expr.args[2]
+  # denest, starting at the back
+  for a in reverse(expr.args[1].args)
+    body = Expr(:for, a, body)
+    # turn for into while loop
+    body = transform_for_inner(body, ui8)
+  end
+  return body
+end
+
+function transform_for_inner(expr, ui8::BoxedUInt8)
+  # turning for into while loops
   @capture(expr, for element_ in iterator_ body_ end) || return expr
-  #@info element
   localelement = Expr(:local, element)
   ui8.n += one(UInt8)
   next = Symbol("_iteratornext_", ui8.n)
diff --git a/test/test_main.jl b/test/test_main.jl
index 714ce19..1e73e25 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -368,3 +368,13 @@ end
 
   @test collect(test_getproperty()) == [[[2], (a = 1, b = 2)]]
 end
+
+@testset "test_weird_for" begin
+  @resumable function test_weird_for(n)
+    for i=1:n, j=1:i
+      @yield i, j
+    end
+  end
+
+  @test collect(test_weird_for(3)) == [(1, 1), (2, 1), (2, 2), (3, 1), (3, 2), (3, 3)]
+end

From bbf1f809a8a6f37d4e264cda0876d6589e03538f Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 10 Aug 2024 20:57:06 +0200
Subject: [PATCH 11/23] more named tuple nonsense

---
 src/transforms.jl | 13 +++++++++++++
 src/utils.jl      | 39 +++++++++++++++++++++++++++++++--------
 test/test_main.jl | 11 ++++++++++-
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/src/transforms.jl b/src/transforms.jl
index 59591b2..cc1c9c7 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -142,6 +142,19 @@ Function that replaces a variable `x` in an expression by `_fsmi.x` where `x` is
 function transform_slots(expr, symbols)
   expr isa Expr || return expr
   #expr.head === :let && return transform_slots_let(expr, symbols)
+  
+  # Special handling of named tuples, we are not allowed to replace things in name tuples
+  # note that the scoping pass already turned (;b) into (;b = b) etc
+  if expr.head === :tuple && expr.args[1] isa Expr && expr.args[1].head === :parameters
+    # (;...)
+    for i in 1:length(expr.args[1].args)
+      @assert !(expr.args[1].args[i] isa Symbol)
+      a = expr.args[1].args[i].args[2] # RHS replace
+      a isa Symbol && a in symbols ? :(_fsmi.$(a)) : a
+    end
+    return expr
+  end
+
   for i in 1:length(expr.args)
     expr.head === :kw && i === 1 && continue
     expr.head === Symbol("quote") && continue
diff --git a/src/utils.jl b/src/utils.jl
index 8ab1be6..263217c 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -242,6 +242,10 @@ function lookup_lhs!(s::Expr, S::ScopeTracker; new = false)
     return s
   end
   if s.head  === :tuple
+    # we should never have to treat a (;a) = b here
+    @assert !(s.args[1] isa Expr && s.args[1].head === :parameters)
+
+    # we should have an innocent (a, b, c) = ...
     for i in 1:length(s.args)
       s.args[i] = lookup_lhs!(s.args[i], S; new = new)
     end
@@ -336,7 +340,7 @@ function scoping(expr::Expr, scope)
     if length(expr.args) > 0 && expr.args[1] isa Expr && expr.args[1].head === :parameters
       # this is a named tuple of the form (;...)
       # TODO: named tuple recognition not working properly yet
-      # first bring (;...,b,...) in the form (;...,b => b,...)
+      # first bring (;...,b,...) in the form (;...,b = b,...)
       for i in 1:length(expr.args[1].args)
         if !(expr.args[1].args[i] isa Expr)
           expr.args[1].args[i] = Expr(:kw, expr.args[1].args[i], expr.args[1].args[i])
@@ -351,14 +355,18 @@ function scoping(expr::Expr, scope)
     elseif any(a -> a isa Expr && a.head === :(=), expr.args[1:end])
       # Can be any of (a = 2, b, c = d)
       # lets first normalize the entries of the form b to b => scoping(b, ...)
-      for i in 2:length(expr.args)
+      for i in 1:length(expr.args)
         if expr.args[i] isa Symbol
-          expr.args[i] = Expr(:(=), expr.args[i], lookup_lhs!(expr.args[i], scope))
+          expr.args[i] = Expr(:kw, expr.args[i], lookup_lhs!(expr.args[i], scope))
         else
           @assert expr.args[i].head === :(=)
           expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
+          expr.args[i].head = :kw
         end
       end
+      # Let's normalize to a parameter (;...) form
+      expr.args[1] = Expr(:parameters, expr.args...)
+      expr.args = expr.args[1:1]
       return expr
     end
   end
@@ -403,13 +411,28 @@ function scoping(expr::Expr, scope)
   end
 
   if expr.head === :(=)
-    if expr.args[1] isa Symbol
-      expr.args[1] = lookup_lhs!(expr.args[1], scope)
-    else
-      for i in 1:length(expr.args[1].args)
-        expr.args[1].args[i] = lookup_lhs!(expr.args[1].args[i], scope)
+    # One special case, where we need to have both LHS and RHS at our hands
+    if expr.args[1] isa Expr && expr.args[1].head === :tuple && expr.args[1].args[1] isa Expr && expr.args[1].args[1].head === :parameters
+      # OK, so this (;a, b) = c
+      # lets transform this into
+      # d = c # because c could be an expression itself
+      # a_new = d.a
+      # b_new = d.b
+      d = gensym()
+      res = [quote $(d) = $(expr.args[2]); end]
+      for i in 1:length(expr.args[1].args[1].args)
+        lhs = expr.args[1].args[1].args[i]
+        @assert lhs isa Symbol
+        lhslookup = lookup_lhs!(lhs, scope)
+        push!(res, quote $(lhslookup) = $(d).$(lhs) end)
       end
+      return quote $(res...) end
     end
+
+    # the LHS is a symbol or a tuple of symbols
+    expr.args[1] = lookup_lhs!(expr.args[1], scope)
+
+    # now transform the RHS, this can be anything
     for i in 2:length(expr.args)
       expr.args[i] = scoping(expr.args[i], scope)
     end
diff --git a/test/test_main.jl b/test/test_main.jl
index 1e73e25..9cf1f82 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -326,6 +326,16 @@ end
 
   @test collect(test_32()) == [(x = 1, )]
 
+  @resumable function test_named_tuple2(b)
+    @yield (;a = b.a + 2, b)
+  end
+  @test collect(test_named_tuple2((a = 1, ))) == [(a = 3, b = (a = 1,))]
+
+  @resumable function test_named_tuple2(b)
+    @yield (a = b.a + 2, b)
+  end
+  @test collect(test_named_tuple2((a = 1, ))) == [(a = 3, b = (a = 1,))]
+
 end
 
 @testset "test_comprehension" begin
@@ -375,6 +385,5 @@ end
       @yield i, j
     end
   end
-
   @test collect(test_weird_for(3)) == [(1, 1), (2, 1), (2, 2), (3, 1), (3, 2), (3, 3)]
 end

From 9e9fa2a08454fa7c358fe26f6be0d4e95273451a Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 11 Aug 2024 08:42:46 +0200
Subject: [PATCH 12/23] handle let more correctly

---
 src/utils.jl      | 25 +++++++------------------
 test/test_main.jl | 39 +++++++++++++++++++++++++++++----------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index 263217c..d0bab5b 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -466,36 +466,25 @@ function scoping(expr::Expr, scope)
     #   new_i = old_i
     #   new_j = new_i
     #
-    #   :(
+    #   Thus we add a new scope, resolve the RHS and force the LHS to be new.
+    #   Do this one after the other and everything will be fine.
 
-    # defer adding a new scope after the right hand side have been renamed
     @capture(expr, let arg_; body_ end) || return expr
     @capture(arg, begin x__ end)
-    replace_rhs = []
-    for i in 1:length(x)
-      y = x[i]
-      fl = @capture(y, k_ = v_)
-      if fl
-        push!(replace_rhs, scoping(v, scope))
-      else
-        # there was no right side
-        push!(replace_rhs, nothing)
-      end
-    end
     new_stack = true
     push!(scope.scope_stack, Dict())
-    replace_lhs = []
     rep = []
     for i in 1:length(x)
       y = x[i]
       fl = @capture(y, k_ = v_)
       if fl
-        push!(replace_lhs, lookup_lhs!(k, scope, new = true))
-        push!(rep, quote local $(replace_lhs[i]); $(replace_lhs[i]) = $(replace_rhs[i]) end)
+        replace_rhs = scoping(v, scope)
+        replace_lhs = lookup_lhs!(k, scope, new = true)
+        push!(rep, quote local $(replace_lhs); $(replace_lhs) = $(replace_rhs) end)
       else
         @assert y isa Symbol
-        push!(replace_lhs, lookup_lhs!(y, scope, new = true))
-        push!(rep, quote local $(replace_lhs[i]) end)
+        replace_lhs = lookup_lhs!(y, scope, new = true)
+        push!(rep, quote local $(replace_lhs) end)
       end
     end
     rep = quote
diff --git a/test/test_main.jl b/test/test_main.jl
index 9cf1f82..6c518fb 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -80,20 +80,39 @@ end
 @test collect(test_varargs(1, 2, 3)) == [1, 2, 3]
 end
 
-@resumable function test_let()
-  for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
-    for i in 1:2
-      local val
-      let i=i
-        val = [a[i] for a in u]
+@testset "test_let" begin
+  @resumable function test_let()
+    for u in [[(1,2),(3,4)], [(5,6),(7,8)]]
+      for i in 1:2
+        local val
+        let i=i
+          val = [a[i] for a in u]
+        end
+        @yield val
       end
-      @yield val
     end
   end
-end
+  @test collect(test_let()) == [[1,3],[2,4],[5,7],[6,8]]
 
-@testset "test_let" begin
-@test collect(test_let()) == [[1,3],[2,4],[5,7],[6,8]]
+  @resumable function test_let2()
+    a = 3
+    b = 2
+    let a = b, b = a, c = b
+      @yield a, b, c
+    end
+    @yield a, b
+  end
+  @test collect(test_let2()) == [(2, 2, 2), (3, 2)]
+
+  @resumable function test_let3()
+    a = 3
+    b = 2
+    let a = a, b = a, c = b
+      @yield a, b, c
+    end
+    @yield a, b
+  end
+  @test collect(test_let3()) == [(3, 3, 3), (3, 2)]
 end
 
 if VERSION >= v"1.8"

From e88ff8311d27b931b30f5bb635d448fc6572702b Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 11 Aug 2024 09:12:42 +0200
Subject: [PATCH 13/23] nested comprehensions/generators

---
 src/utils.jl      | 11 +++--------
 test/test_main.jl | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index d0bab5b..a7fd045 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -400,14 +400,9 @@ function scoping(expr::Expr, scope)
     end
     return expr
   end
-  if expr.head === :comprehension
-    # this is again special with respect to scoping
-    if expr.args[1].head === :generator
-      expr.args[1] = scope_generator(expr.args[1], scope)
-      return expr
-    else
-      error("not implemented yet")
-    end
+  if expr.head === :generator
+    expr = scope_generator(expr)
+    return expr
   end
 
   if expr.head === :(=)
diff --git a/test/test_main.jl b/test/test_main.jl
index 6c518fb..2dede87 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -350,10 +350,10 @@ end
   end
   @test collect(test_named_tuple2((a = 1, ))) == [(a = 3, b = (a = 1,))]
 
-  @resumable function test_named_tuple2(b)
+  @resumable function test_named_tuple3(b)
     @yield (a = b.a + 2, b)
   end
-  @test collect(test_named_tuple2((a = 1, ))) == [(a = 3, b = (a = 1,))]
+  @test collect(test_named_tuple3((a = 1, ))) == [(a = 3, b = (a = 1,))]
 
 end
 
@@ -369,6 +369,40 @@ end
     end
   end
   @test collect(test_comprehension([(1, 2), (3, 4)])) == [2,4,1,9]
+
+  @resumable function test_comprehension2()
+    x = 2
+    y = 3
+    @yield [ x + y for x in 1:y, y in 1:x]
+  end
+  @test collect(test_comprehension2()) == [[2 3; 3 4; 4 5]]
+
+  @resumable function test_comprehension3()
+    x = 2
+    y = 3
+    @yield [ x + y for x in 1:x, y in 1:y]
+  end
+  @test collect(test_comprehension3()) == [[2 3 4; 3 4 5]]
+
+  @resumable function test_comprehension4()
+    x = 2
+    y = 3
+    u = [  x + y for x in 1:y for y in 1:x]
+    for k in u
+      @yield k
+    end
+  end
+  @test collect(test_comprehension4()) == [2, 3, 4, 4, 5, 6]
+
+  @resumable function test_comprehension5()
+    x = 2
+    y = 3
+    u = [  x + y for x in 1:x for y in 1:y]
+    for k in u
+      @yield k
+    end
+  end
+  @test collect(test_comprehension5()) == [2, 3, 4, 3, 4, 5]
 end
 
 @testset "test_ref" begin

From 205c4bc7538c95265640cc180fde52002eee94a3 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 11 Aug 2024 09:13:52 +0200
Subject: [PATCH 14/23] simplify transform_slots

---
 src/transforms.jl | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/src/transforms.jl b/src/transforms.jl
index cc1c9c7..a5ec285 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -143,18 +143,7 @@ function transform_slots(expr, symbols)
   expr isa Expr || return expr
   #expr.head === :let && return transform_slots_let(expr, symbols)
   
-  # Special handling of named tuples, we are not allowed to replace things in name tuples
-  # note that the scoping pass already turned (;b) into (;b = b) etc
-  if expr.head === :tuple && expr.args[1] isa Expr && expr.args[1].head === :parameters
-    # (;...)
-    for i in 1:length(expr.args[1].args)
-      @assert !(expr.args[1].args[i] isa Symbol)
-      a = expr.args[1].args[i].args[2] # RHS replace
-      a isa Symbol && a in symbols ? :(_fsmi.$(a)) : a
-    end
-    return expr
-  end
-
+  # "Problematic" expressions all have been transformed into Expr(:kw,...)
   for i in 1:length(expr.args)
     expr.head === :kw && i === 1 && continue
     expr.head === Symbol("quote") && continue

From 824b5b53ea9af081a3f26dea8735d660b1c42bf3 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 11 Aug 2024 10:01:36 +0200
Subject: [PATCH 15/23] better local handling

---
 src/utils.jl | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index a7fd045..9e0558e 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -302,6 +302,7 @@ scoping(e::typeof(ResumableFunctions.IteratorReturn), scope) = e
 scoping(e::QuoteNode, scope) = e
 scoping(e::Bool, scope) = e
 scoping(e::Nothing, scope) = e
+scoping(e::GlobalRef, scope) = e
 
 function scoping(s::Symbol, scope; new = false)
   #@info "scoping $s, $new"
@@ -519,11 +520,16 @@ function scoping(expr::Expr, scope)
         expr.args[1].args[i] = lookup_lhs!(a, scope)
       end
     else
-      for i in 1:length(expr.args)
-        a = expr.args[i]
-        #expr.args[i] = scoping(a, scope, new = true)
-        expr.args[i] = lookup_lhs!(a, scope)
-      end
+      # this is local x = y
+      @assert length(expr.args) == 1 && expr.args[1] isa Expr && expr.args[1].head === :(=)
+      expr.args[1].args[1] = lookup_lhs!(expr.args[1].args[1], scope)
+      expr.args[1].args[2] = scoping(expr.args[1].args[2], scope)
+      #for i in 1:length(expr.args)
+      #  a = expr.args[i]
+      #  #expr.args[i] = scoping(a, scope, new = true)
+      #  expr.args[i] = lookup_lhs!(a, scope)
+      #end
+      expr = quote local $(expr.args[1].args[1]); $(expr.args[1].args[1]) = $(expr.args[1].args[2]); end
     end
     return expr
   end

From 732434558a59c78c52b80f65bfea95c7beeb1ba6 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 11 Aug 2024 10:08:12 +0200
Subject: [PATCH 16/23] fix jet again

---
 test/test_jet.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/test_jet.jl b/test/test_jet.jl
index aa9d319..4974ac4 100644
--- a/test/test_jet.jl
+++ b/test/test_jet.jl
@@ -8,6 +8,7 @@ using Test
             Core.Compiler,
         )
     )
-    @test length(JET.get_reports(rep)) <= 7
+    @show rep
+    @test length(JET.get_reports(rep)) <= 8
     @test_broken length(JET.get_reports(rep)) == 0
 end

From 908a9299e6b38cd45b88071d0054b39324920e17 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sun, 29 Sep 2024 19:25:51 +0200
Subject: [PATCH 17/23] doc

---
 src/macro.jl      |   9 ++--
 src/transforms.jl |   2 +-
 src/utils.jl      | 122 +++++++++++++++++++++++++++++-----------------
 test/test_main.jl |   1 +
 4 files changed, 81 insertions(+), 53 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index db4e25e..8298a14 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -73,7 +73,7 @@ macro resumable(ex::Expr...)
 
   # The function that executes a step of the finite state machine
   func_def = splitdef(expr)
-  @debug func_def[:body]
+  @debug func_def[:body]|>MacroTools.striplines
   rtype = :rtype in keys(func_def) ? func_def[:rtype] : Any
   args, kwargs, arg_dict = get_args(func_def)
   params = ((get_param_name(param) for param in func_def[:whereparams])...,)
@@ -96,19 +96,16 @@ macro resumable(ex::Expr...)
   # :name is :(fA::A) if it is an overloading call function (fA::A)(...)
   # ...
   if func_def[:name] isa Expr
-    @assert func_def[:name].head == :(::)
+    func_def[:name].head != :(::) && error("Unrecognized function name: $(func_def[:name].head)")
     _name = func_def[:name].args[1]
   else
     _name = func_def[:name]
   end
   
   scope = ScopeTracker(0, __module__, [Dict(i =>i for i in vcat(args, kwargs, [_name], params...))])
-  #@info func_def[:body]|>MacroTools.striplines
-  #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = scoping(copy(func_def[:body]), scope)
-  #@info func_def[:body]|>MacroTools.striplines
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
-  #@info func_def[:body]|>MacroTools.striplines
+  @debug func_def[:body]|>MacroTools.striplines
 
   inferfn, slots = get_slots(copy(func_def), arg_dict, __module__)
   @debug slots
diff --git a/src/transforms.jl b/src/transforms.jl
index a5ec285..99ffffa 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -92,7 +92,7 @@ function transform_for(expr, ui8::BoxedUInt8)
   # test for simple for a in b expression
   expr.args[1].head === :(=) && return transform_for_inner(expr, ui8)
   # must be a complicated iteration
-  @assert expr.args[1].head === :block
+  expr.args[1].head !== :block && error("Unrecognized for expression: $(expr.args[1])")
   body = expr.args[2]
   # denest, starting at the back
   for a in reverse(expr.args[1].args)
diff --git a/src/utils.jl b/src/utils.jl
index 9e0558e..b5dd4c7 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -211,6 +211,65 @@ function typed_fsmi_fallback(fsmi::Type{T}, fargs...)::T where T
   return T()
 end
 
+################################################################################
+#
+#  Scoping
+#
+################################################################################
+
+# As every modern programming language, julia has the concept of variable scope,
+# see https://docs.julialang.org/en/v1/manual/variables-and-scoping/.
+#
+# For example, in the following example, there are actual two variables.
+# This can be seen by rewriting this as on the right hand side.
+#
+# a = 1         | a_1 = 1
+# let a = a     | begin
+#   a = 2       |   a_2 = a_1; a_2 = 2
+# end           | end
+#
+# A similar phenomen happens with `local`
+#
+# x = 1           | x_1 = 1
+# begin           | begin
+#   local x = 0   |   x_2 = 0
+# end             | end
+#
+# Simulating a function using a finite state machine (FSM) has a disadvantage
+# that it cannot handle expressions, where identical left-hand sides (RHS)
+# refer to different variables. This applies to both `let` as well as `local`
+# constructions.
+# 
+# We solve this problem by renaming all variables.
+#
+# We use a ScopeTracker type to keep track of the things that are already
+# renamned. It is basically just a Vector{Dict{Symbol, Symbol}},
+# representing a stack, where the top records the renamend variables in the
+# current scope.
+#
+# The renaming is done as follows. If we encounter an assignment of the form
+# x = y
+# there are two cases for x:
+#   1) x has been seen before in some scope. Then we replace x accordingly.
+#   2) x has not been seen before. We give x a new name and store :x => :x_new
+#      in current scope.
+#
+# This is done in lookup_lhs!. Note that some construction, like `let`, create
+# a new variable in a new scope. This is handled by the `new` keyword.
+#
+# For any other symbol y (which is not the left hand side of an assignmetn),
+# there are the following two cases:
+#   1) y has been seen before in some scope. Then we replace y accordinly.
+#   2) y has not been seen before, then we don't rename it.
+#
+# ---
+#
+# Note that we handle local x, by emitting a new variable :x => :x_new
+# inside the current scope.
+#
+# We exploit this when rewriting let and for constructions, see below for
+# examples with let.
+
 mutable struct ScopeTracker
   i::Int
   mod::Module
@@ -243,7 +302,8 @@ function lookup_lhs!(s::Expr, S::ScopeTracker; new = false)
   end
   if s.head  === :tuple
     # we should never have to treat a (;a) = b here
-    @assert !(s.args[1] isa Expr && s.args[1].head === :parameters)
+    (s.args[1] isa Expr && s.args[1].head === :parameters) &&
+        error("Illegal tuple expression in scope lookup: $(s.args[1])")
 
     # we should have an innocent (a, b, c) = ...
     for i in 1:length(s.args)
@@ -269,33 +329,9 @@ function lookup_rhs!(s::Symbol, S::ScopeTracker)
   return s
 end
 
-function lookup!(s::Symbol, S::ScopeTracker; new = false)
-  s == :val && @show s, length(S.scope_stack), new
-  if isdefined(S.mod, s)
-    return s
-  end
-  if !new
-    for D in Iterators.reverse(S.scope_stack)
-      if haskey(D, s)
-        return D[s]
-      end
-    end
-    D = last(S.scope_stack)
-    D[s] = s
-    return s
-  end
-  D = last(S.scope_stack)
-  new = Symbol(s, Symbol("_$(S.i)"))
-  S.i += 1
-  D[s] = new
-  return new
-end
-
 scoping(e::LineNumberNode, scope) = e
-
 scoping(e::Int, scope) = e
 scoping(e::Float64, scope) = e
-
 scoping(e::String, scope) = e
 scoping(e::typeof(ResumableFunctions.generate), scope) = e
 scoping(e::typeof(ResumableFunctions.IteratorReturn), scope) = e
@@ -305,16 +341,14 @@ scoping(e::Nothing, scope) = e
 scoping(e::GlobalRef, scope) = e
 
 function scoping(s::Symbol, scope; new = false)
-  #@info "scoping $s, $new"
   return lookup_rhs!(s, scope)
 end
 
 function scope_generator(expr, scope)
-  @assert expr.head === :generator
+  expr.head !== :generator && error("Illegal generator expression: $(expr)")
   # first the generator case
   for i in 2:length(expr.args)
-    @assert expr.args[i] isa Expr && expr.args[i].head === :(=)
-    #expr.args[i].args[2] = lookup_rhs!(expr.args[i].args[2], scope)
+    !(expr.args[i] isa Expr && expr.args[i].head === :(=)) && error("Illegal expression in generator: $(expr.args[i])")
     expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
   end
   # now create new scope
@@ -349,7 +383,8 @@ function scoping(expr::Expr, scope)
       end
       # Now rename the RHS
       for i in 1:length(expr.args[1].args)
-        @assert expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :kw
+        !(expr.args[1].args[i] isa Expr && expr.args[1].args[i].head === :kw) &&
+          error("Illegal expression in named tuple: $(expr.args[1].args[i])")
         expr.args[1].args[i].args[2] = scoping(expr.args[1].args[i].args[2], scope)
       end
       return expr
@@ -360,7 +395,8 @@ function scoping(expr::Expr, scope)
         if expr.args[i] isa Symbol
           expr.args[i] = Expr(:kw, expr.args[i], lookup_lhs!(expr.args[i], scope))
         else
-          @assert expr.args[i].head === :(=)
+          expr.args[i].head !== :(=) &&
+            error("Uncregonized expression in tuple expression: $(expr.args[i])")
           expr.args[i].args[2] = scoping(expr.args[i].args[2], scope)
           expr.args[i].head = :kw
         end
@@ -391,7 +427,8 @@ function scoping(expr::Expr, scope)
           end
         end
         for j in 1:length(expr.args[i].args)
-          @assert expr.args[i].args[j] isa Expr && expr.args[i].args[j].head === :kw
+          !(expr.args[i].args[j] isa Expr && expr.args[i].args[j].head === :kw) &&
+            error("Unrecognized keyword expression: $(expr.args[i].args[j])")
           # this is f(...; x = 2)
           expr.args[i].args[j].args[2] = scoping(expr.args[i].args[j].args[2], scope)
         end
@@ -418,7 +455,8 @@ function scoping(expr::Expr, scope)
       res = [quote $(d) = $(expr.args[2]); end]
       for i in 1:length(expr.args[1].args[1].args)
         lhs = expr.args[1].args[1].args[i]
-        @assert lhs isa Symbol
+        !(lhs isa Symbol) &&
+          error("Unrecognized expression in named tuple assignment: $(lhs)")
         lhslookup = lookup_lhs!(lhs, scope)
         push!(res, quote $(lhslookup) = $(d).$(lhs) end)
       end
@@ -478,7 +516,8 @@ function scoping(expr::Expr, scope)
         replace_lhs = lookup_lhs!(k, scope, new = true)
         push!(rep, quote local $(replace_lhs); $(replace_lhs) = $(replace_rhs) end)
       else
-        @assert y isa Symbol
+        !(y isa Symbol) &&
+          error("Unrecognized expression in let expression: $(y)")
         replace_lhs = lookup_lhs!(y, scope, new = true)
         push!(rep, quote local $(replace_lhs) end)
       end
@@ -505,30 +544,21 @@ function scoping(expr::Expr, scope)
     new_stack = true
   end
   if expr.head === :local
-    # this is my local dance
-    # explain and rewrite using @capture
-    #
-    # if I see a local x or local x = ...
+    # if we see a local x or local x = ...
     # we always emit a new identifier
     if length(expr.args) == 1 && expr.args[1] isa Symbol
-      #expr.args[1] = scoping(expr.args[1], scope, new = true)
       expr.args[1] = lookup_lhs!(expr.args[1], scope)
     elseif length(expr.args) == 1 && expr.args[1].head === :tuple
       for i in 1:length(expr.args[1].args)
         a = expr.args[1].args[i]
-        #expr.args[1].args[i] = scoping(a, scope, new = true)
         expr.args[1].args[i] = lookup_lhs!(a, scope)
       end
     else
       # this is local x = y
-      @assert length(expr.args) == 1 && expr.args[1] isa Expr && expr.args[1].head === :(=)
+      !(length(expr.args) == 1 && expr.args[1] isa Expr && expr.args[1].head === :(=)) &&
+        error("Illegal local expression: $(expr.args)")
       expr.args[1].args[1] = lookup_lhs!(expr.args[1].args[1], scope)
       expr.args[1].args[2] = scoping(expr.args[1].args[2], scope)
-      #for i in 1:length(expr.args)
-      #  a = expr.args[i]
-      #  #expr.args[i] = scoping(a, scope, new = true)
-      #  expr.args[i] = lookup_lhs!(a, scope)
-      #end
       expr = quote local $(expr.args[1].args[1]); $(expr.args[1].args[1]) = $(expr.args[1].args[2]); end
     end
     return expr
diff --git a/test/test_main.jl b/test/test_main.jl
index 2dede87..3bc6bae 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -440,3 +440,4 @@ end
   end
   @test collect(test_weird_for(3)) == [(1, 1), (2, 1), (2, 2), (3, 1), (3, 2), (3, 3)]
 end
+end

From 88705c466ba7e75078e9bc5ed92ce902c6ee3077 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 26 Oct 2024 23:26:42 +0200
Subject: [PATCH 18/23] remove unused things, document things

---
 src/macro.jl      |  8 +-----
 src/transforms.jl | 64 +++--------------------------------------------
 src/utils.jl      |  2 +-
 test/test_main.jl |  1 -
 4 files changed, 5 insertions(+), 70 deletions(-)

diff --git a/src/macro.jl b/src/macro.jl
index 8298a14..347f369 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -84,13 +84,7 @@ macro resumable(ex::Expr...)
   func_def[:body] = postwalk(transform_yieldfrom, func_def[:body])
   func_def[:body] = postwalk(x->transform_for(x, ui8), func_def[:body])
   @debug func_def[:body]|>MacroTools.striplines
-  #func_def[:body] = postwalk(x->transform_macro(x), func_def[:body])
-  #@debug func_def[:body]|>MacroTools.striplines
-  #func_def[:body] = postwalk(x->transform_macro_undo(x), func_def[:body])
-  #@debug func_def[:body]|>MacroTools.striplines
-  #func_def[:body] = postwalk(x->transform_let(x), func_def[:body])
-  #@info func_def[:body]|>MacroTools.striplines
-  #func_def[:body] = postwalk(x->transform_local(x), func_def[:body])
+
   # Scoping fixes
 
   # :name is :(fA::A) if it is an overloading call function (fA::A)(...)
diff --git a/src/transforms.jl b/src/transforms.jl
index 99ffffa..8b21c9f 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -1,20 +1,11 @@
+"""
+Function that removes `local x` expression.
+"""
 function transform_remove_local(ex)
   ex isa Expr && ex.head === :local && return Expr(:block)
   return ex
 end
 
-function transform_macro(ex)
-  ex isa Expr || return ex
-  ex.head !== :macrocall && return ex
-  return Expr(:call, :__secret__, ex.args)
-end
-
-function transform_macro_undo(ex)
-  ex isa Expr || return ex
-  (ex.head !== :call || ex.args[1] !== :__secret__) && return ex
-  return Expr(:macrocall, ex.args[2]...)
-end
-
 """
 Function that replaces a variable
 """
@@ -152,55 +143,6 @@ function transform_slots(expr, symbols)
   expr
 end
 
-#"""
-#Function that handles `let` block
-#"""
-#function transform_slots_let(expr::Expr, symbols)
-#  @capture(expr, let vars_; body_ end)
-#  locals = Set{Symbol}()
-#  (isa(vars, Expr) && vars.head==:(=))  || error("@resumable currently supports only single variable declarations in let blocks, i.e. only let blocks exactly of the form `let i=j; ...; end`. If you need multiple variables, please submit an issue on the issue tracker and consider contributing a patch.")
-#  sym = vars.args[1].args[2].value
-#  push!(locals, sym)
-#  vars.args[1] = sym
-#  body = postwalk(x->transform_let(x, locals), :(begin $(body) end))
-#  :(let $vars; $body end)
-#end
-
-function transform_let(expr)
-  expr isa Expr || return expr
-  expr.head === :block && return expr
-  #@info "inside transform let"
-  @capture(expr, let arg_; body_; end) || return expr
-  #@info "captured let"
-  #arg |> dump
-  #@info expr
-  #@info arg
-  #error("ASds")
-  res = quote
-    let
-      local $arg
-      $body
-    end
-  end
-  #@info "emitting $res"
-  res
-  #expr.head === :. || return expr
-  #expr = expr.args[2].value in symbols ? :($(expr.args[2].value)) : expr
-end
-
-"""
-Function that replaces a variable `_fsmi.x` in an expression by `x` where `x` is a variable declared in a `let` block.
-"""
-function transform_local(expr)
-  expr isa Expr || return expr
-  @capture(expr, local arg_ = ex_) || return expr
-  res = quote
-    local $arg
-    $arg = $ex
-  end
-  res
-end
-
 """
 Function that replaces a `arg = @yield ret` statement by
 ```julia
diff --git a/src/utils.jl b/src/utils.jl
index b5dd4c7..495414c 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -268,7 +268,7 @@ end
 # inside the current scope.
 #
 # We exploit this when rewriting let and for constructions, see below for
-# examples with let.
+# examples with let. At the end, all `local x` are removed.
 
 mutable struct ScopeTracker
   i::Int
diff --git a/test/test_main.jl b/test/test_main.jl
index 3bc6bae..2dede87 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -440,4 +440,3 @@ end
   end
   @test collect(test_weird_for(3)) == [(1, 1), (2, 1), (2, 2), (3, 1), (3, 2), (3, 3)]
 end
-end

From e96bf99c19ac6992b3a36d64b95296fd6d0ad51d Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 26 Oct 2024 23:33:18 +0200
Subject: [PATCH 19/23] fix test

---
 test/test_main.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/test_main.jl b/test/test_main.jl
index 2dede87..3bc6bae 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -440,3 +440,4 @@ end
   end
   @test collect(test_weird_for(3)) == [(1, 1), (2, 1), (2, 2), (3, 1), (3, 2), (3, 3)]
 end
+end

From 9099f85f27f1df17e6d4614ffd7945e7216a9a4e Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 26 Oct 2024 23:40:24 +0200
Subject: [PATCH 20/23] work around julia regression

---
 test/Project.toml | 1 +
 test/runtests.jl  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/test/Project.toml b/test/Project.toml
index 0628c77..6565403 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -3,6 +3,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/runtests.jl b/test/runtests.jl
index 43bcb9c..79dda8d 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,7 @@
 using ResumableFunctions
 using Test
 using SafeTestsets
+using REPL # work around regression with @doc macro: https://github.com/JuliaLang/julia/issues/54664
 
 function doset(descr)
     if length(ARGS) == 0

From d4d39af17881ff15ae89806fa546c398f39843b2 Mon Sep 17 00:00:00 2001
From: Tommy Hofmann <thofma@gmail.com>
Date: Sat, 26 Oct 2024 23:43:14 +0200
Subject: [PATCH 21/23] fix spelling, thanks spellchecker

---
 src/utils.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/utils.jl b/src/utils.jl
index 495414c..4364895 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -259,7 +259,7 @@ end
 #
 # For any other symbol y (which is not the left hand side of an assignmetn),
 # there are the following two cases:
-#   1) y has been seen before in some scope. Then we replace y accordinly.
+#   1) y has been seen before in some scope. Then we replace y accordingly.
 #   2) y has not been seen before, then we don't rename it.
 #
 # ---

From 49aba4a454b902109f6a93d4595a7f4c0c366a62 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Thu, 21 Nov 2024 11:18:05 -0500
Subject: [PATCH 22/23] minor whitespace cleanup and returning to safetestsets

---
 docs/src/manual.md | 4 ++--
 src/macro.jl       | 4 ++--
 src/transforms.jl  | 2 +-
 src/utils.jl       | 4 ++--
 test/runtests.jl   | 8 ++++----
 test/test_main.jl  | 2 +-
 6 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/docs/src/manual.md b/docs/src/manual.md
index 70d6429..764728e 100644
--- a/docs/src/manual.md
+++ b/docs/src/manual.md
@@ -49,7 +49,7 @@ DocTestSetup = quote
 
   @resumable function basic_example()
     @yield "Initial call"
-    @yield 
+    @yield
     "Final call"
   end
 end
@@ -58,7 +58,7 @@ end
 ```julia
 @resumable function basic_example()
   @yield "Initial call"
-  @yield 
+  @yield
   "Final call"
 end
 ```
diff --git a/src/macro.jl b/src/macro.jl
index 347f369..d7663d9 100755
--- a/src/macro.jl
+++ b/src/macro.jl
@@ -95,7 +95,7 @@ macro resumable(ex::Expr...)
   else
     _name = func_def[:name]
   end
-  
+
   scope = ScopeTracker(0, __module__, [Dict(i =>i for i in vcat(args, kwargs, [_name], params...))])
   func_def[:body] = scoping(copy(func_def[:body]), scope)
   func_def[:body] = postwalk(x->transform_remove_local(x), func_def[:body])
@@ -176,7 +176,7 @@ macro resumable(ex::Expr...)
   end
   call_expr = combinedef(call_def) |> flatten
   @debug call_expr|>MacroTools.striplines
-  
+
   # Finalizing the function stepping through the finite state machine
   if isempty(params)
     func_def[:name] = :((_fsmi::$type_name))
diff --git a/src/transforms.jl b/src/transforms.jl
index 8b21c9f..979986c 100755
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -133,7 +133,7 @@ Function that replaces a variable `x` in an expression by `_fsmi.x` where `x` is
 function transform_slots(expr, symbols)
   expr isa Expr || return expr
   #expr.head === :let && return transform_slots_let(expr, symbols)
-  
+
   # "Problematic" expressions all have been transformed into Expr(:kw,...)
   for i in 1:length(expr.args)
     expr.head === :kw && i === 1 && continue
diff --git a/src/utils.jl b/src/utils.jl
index 4364895..5b97579 100755
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -239,7 +239,7 @@ end
 # that it cannot handle expressions, where identical left-hand sides (RHS)
 # refer to different variables. This applies to both `let` as well as `local`
 # constructions.
-# 
+#
 # We solve this problem by renaming all variables.
 #
 # We use a ScopeTracker type to keep track of the things that are already
@@ -480,7 +480,7 @@ function scoping(expr::Expr, scope)
   end
   new_stack = false
   if expr.head === :let
-    # Replace 
+    # Replace
     #   let i, k = 2, j = 1
     #      [...]
     #   end
diff --git a/test/runtests.jl b/test/runtests.jl
index 79dda8d..774b07e 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,11 +19,11 @@ macro doset(descr)
     quote
 #        @info "====================================="
 #        @info $descr
-#        if doset($descr)
-#            @safetestset $descr begin
+        if doset($descr)
+            @safetestset $descr begin
                 include("test_"*$descr*".jl")
-#            end
-#        end
+            end
+        end
     end
 end
 
diff --git a/test/test_main.jl b/test/test_main.jl
index 3bc6bae..9b601e5 100644
--- a/test/test_main.jl
+++ b/test/test_main.jl
@@ -287,7 +287,7 @@ end
     y = 1
     @yield g(z, z = y, 2)
   end
-  
+
   @test collect(test_kw(3)) == [8]
 
   g(z; y) = z - y

From 0f900951dfe9190ec1819aa35907867c5befafbc Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Thu, 21 Nov 2024 11:19:36 -0500
Subject: [PATCH 23/23] changelog

---
 CHANGELOG.md | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7257787..e3df7f3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # News
 
+## v1.0.0 - dev
+
+- **(breaking)** Support for proper Julia scoping of variables inside of resumable functions.
+
 ## v0.6.10 - 2024-09-08
 
 - Add `length` support by allowing `@resumable length=ex function [...]`
@@ -39,7 +43,7 @@
   * 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 
+* 2023: v0.6.2
   * Julia v1.10 compatibility fix
   * resumable functions can now dispatch on types