Commit forgotten local changes

Author: Keno

src/stage1/recurse.jl

@@ -456,8 +456,10 @@ function transform!(ci, meth, nargs, sparams, N)
                 return insert_node_rev!(Expr(:call, getfield, Argument(1),
                     i - first(orig_bb_ranges[end]) + 1))
             elseif isa(current_env, BBEnv)
+                bbidx = i - current_env.bb_start_idx + 1
+                @assert bbidx > 0
                 return insert_node_rev!(Expr(:call, getfield, current_env.ctx_obj,
-                    i - current_env.bb_start_idx + 1))
+                    bbidx))
             end
             error()
         end
@@ -489,9 +491,9 @@ function transform!(ci, meth, nargs, sparams, N)
                         insert_node_rev!(Expr(:(=), accumulator, accumed))
                     end
                 end
-                if !isa(stmt, Union{GotoNode, GotoIfNot})
-                    #current_env = BBEnv(access_ctx_map(bb+1),
-                    #    first(ir.cfg.blocks[bb].stmts))
+                if !isa(stmt, Union{GotoNode, GotoIfNot, ReturnNode})
+                    current_env = BBEnv(access_ctx_map(bb+1),
+                        first(ir.cfg.blocks[bb].stmts))
                 end
             end
 

test/runtests.jl

@@ -60,6 +60,15 @@ function myprod(xs)
     return s
 end
 
+function mypow(x, n)
+    r = one(x)
+    while n > 0
+        n -= 1
+        r *= x
+    end
+    return r
+end
+
 function times_three_while(x)
     z = x
     i = 3
@@ -95,10 +104,14 @@ let var"'" = Diffractor.PrimeDerivativeBack
     @test @inferred(complicated_2sin'''(1.0)) == 2sin'''(1.0)
     @test @inferred(complicated_2sin''''(1.0)) == 2sin''''(1.0)
 
+    # Control flow cases
     @test @inferred((x->simple_control_flow(true, x))'(1.0)) == sin'(1.0)
     @test @inferred((x->simple_control_flow(false, x))'(1.0)) == cos'(1.0)
     @test (x->sum(isa_control_flow(Matrix{Float64}, x)))'(Float32[1 2;]) == [1.0 1.0;]
     @test times_three_while'(1.0) == 3.0
+
+    pow5p(x) = (x->mypow(x, 5))'(x)
+    @test pow5p(1.0) == 5.0
 end
 
 # Simple Forward Mode tests