Merge pull request #163 from KristofferC/kc/fix_apply_iterate

fix bug in Core._apply_iterate overdubbing, args[1] also needs to be overdubbed
JuliaLabs · Jan 27, 2020 · 2e0827a · 2e0827a
2 parents 18b5480 + 6d75f47
commit 2e0827a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/src/context.jl b/src/context.jl
@@ -259,7 +259,10 @@ macro context(_Ctx)
 
         @inline Cassette.overdub(ctx::$Ctx, ::typeof(Core._apply), f, args...) = Core._apply(overdub, (ctx, f), args...)
         if VERSION >= v"1.4.0-DEV.304"
-            @inline Cassette.overdub(ctx::$Ctx, ::typeof(Core._apply_iterate), f, args...) = Core._apply_iterate((args...)->overdub(ctx, f, args...), args...)
+            @inline function Cassette.overdub(ctx::$Ctx, ::typeof(Core._apply_iterate), f, args...)
+                new_args = ((_args...) -> overdub(ctx, args[1], _args...), Base.tail(args)...)
+                Core._apply_iterate((args...)->overdub(ctx, f, args...), new_args...)
+            end
         end
 
         # TODO: There are certain non-`Core.Builtin` functions which the compiler often

diff --git a/src/overdub.jl b/src/overdub.jl
@@ -523,7 +523,10 @@ function recurse end
 
 recurse(ctx::Context, ::typeof(Core._apply), f, args...) = Core._apply(recurse, (ctx, f), args...)
 if VERSION >= v"1.4.0-DEV.304"
-    recurse(ctx::Context, ::typeof(Core._apply_iterate), f, args...) = Core._apply_iterate((args...)->recurse(ctx, f, args...), args...)
+    function recurse(ctx::Context, ::typeof(Core._apply_iterate), f, args...)
+        new_args = ((_args...) -> overdub(ctx, args[1], _args...), Base.tail(args)...)
+        Core._apply_iterate((args...)->recurse(ctx, f, args...), new_args...)
+    end
 end
 
 function overdub_definition(line, file)

diff --git a/test/misctests.jl b/test/misctests.jl
@@ -706,3 +706,17 @@ if VERSION >= v"1.4.0-DEV.304"
 else
     @test_broken Cassette.overdub(NukeContext(), launch, Silo()) === ()
 end
+
+if VERSION >= v"1.4.0-DEV.304"
+    Cassette.@context ApplyIterateCtx;
+
+    const instructions = []
+    function Cassette.prehook(ctx::ApplyIterateCtx,
+                              op::Any,
+                              a::T1, b::T2) where {T1, T2}
+        push!(instructions, (op, T1, T2))
+    end
+
+    Cassette.overdub(ApplyIterateCtx(), ()->pi*2.0)
+    @test instructions[end] === (Core.Intrinsics.mul_float, Float64, Float64)
+end