Pass maparg to transform (#178)

In case the custom transform needs to get its argument of a particular
order.

Author: Keno

src/codegen/forward_demand.jl

@@ -254,7 +254,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             return arg
         elseif isa(arg, Argument)
             # TODO: Should we remember whether the callbacks wanted the arg?
-            return transform!(ir, arg, order)
+            return transform!(ir, arg, order, maparg)
         elseif isa(arg, GlobalRef)
             @assert isconst(arg)
             return insert_node!(ir, ssa, NewInstruction(Expr(:call, ZeroBundle{order}, arg), Any))
@@ -277,7 +277,7 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
             continue
         end
         if custom
-            transform!(ir, SSAValue(ssa), order)
+            transform!(ir, SSAValue(ssa), order, maparg)
         else
             inst = ir[SSAValue(ssa)]
             stmt = inst[:inst]

src/stage2/forward.jl

@@ -33,7 +33,7 @@ function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1)
         end
     end
 
-    function transform!(ir::IRCode, ssa::SSAValue, _)
+    function transform!(ir::IRCode, ssa::SSAValue, _, _)
         inst = ir[ssa]
         stmt = inst[:inst]
         if isa(stmt, ReturnNode)
@@ -44,7 +44,7 @@ function dontuse_nth_order_forward_stage2(tt::Type, order::Int=1)
         end
     end
 
-    function transform!(ir::IRCode, arg::Argument, _)
+    function transform!(ir::IRCode, arg::Argument, _, _)
         return insert_node!(ir, SSAValue(1), NewInstruction(Expr(:call, ∂xⁿ{order}(), arg), typeof(∂xⁿ{order}()(1.0))))
     end
 

test/forward_diff_no_inf.jl

@@ -2,11 +2,11 @@
 module forward_diff_no_inf
     using Diffractor, Test
     # this is needed as transform! is *always* called on Arguments regardless of what visit_custom says
-    identity_transform!(ir, ssa::Core.SSAValue, order) = ir[ssa]
-    function identity_transform!(ir, arg::Core.Argument, order)
+    identity_transform!(ir, ssa::Core.SSAValue, order, _) = ir[ssa]
+    function identity_transform!(ir, arg::Core.Argument, order, _)
         return Core.Compiler.insert_node!(ir, Core.SSAValue(1), Core.Compiler.NewInstruction(Expr(:call, Diffractor.ZeroBundle{1}, arg), Any))
     end
-    
+
     @testset "Constructors in forward_diff_no_inf!" begin
         struct Bar148
             v
@@ -47,7 +47,7 @@ module forward_diff_no_inf
             end
             return x - a + b
         end
-    
+
         input_ir = first(only(Base.code_ircode(phi_run, Tuple{Float64})))
         ir = copy(input_ir)
         #Workout where to diff to trigger error
@@ -57,11 +57,11 @@ module forward_diff_no_inf
                 push!(diff_ssa, Core.SSAValue(idx))
             end
         end
-    
+
         Diffractor.forward_diff_no_inf!(ir, diff_ssa .=> 1; transform! = identity_transform!)
         ir2 = Core.Compiler.compact!(ir)
         Core.Compiler.verify_ir(ir2)  # This would error if we were not handling nonconst phi nodes correctly (after https://github.com/JuliaLang/julia/pull/50158)
         f = Core.OpaqueClosure(ir2; do_compile=false)
         @test f(3.5) == 3.5  # this will segfault if we are not handling phi nodes correctly
-    end    
+    end
 end