Merge pull request #148 from JuliaDiff/sf/fix_new_zero_bundle_error

Forward mode handingly for `Expr(:new, )` in `forward_diff_no_inf!`

Author: staticfloat

src/codegen/forward_demand.jl

@@ -183,6 +183,8 @@ function forward_visit!(ir::IRCode, ssa::SSAValue, order::Int, ssa_orders::Vecto
         foreach(recurse, stmt.args)
     elseif isa(stmt, SSAValue)
         recurse(stmt)
+    elseif isexpr(stmt, :code_coverage_effect)
+        return
     elseif !isa(stmt, Expr)
         return
     else
@@ -280,11 +282,12 @@ function forward_diff_no_inf!(ir::IRCode, to_diff::Vector{Pair{SSAValue,Int}};
                 end
                 inst[:inst] = Expr(:call, ∂☆{order}(), newargs...)
                 inst[:type] = Any
-            elseif isexpr(stmt, :call)
+            elseif isexpr(stmt, :call) || isexpr(stmt, :new)
                 newargs = map(stmt.args) do @nospecialize arg
                     maparg(arg, SSAValue(ssa), order)
                 end
-                inst[:inst] = Expr(:call, ∂☆{order}(), newargs...)
+                f = isexpr(stmt, :call) ? ∂☆{order}() : ∂☆new{order}()
+                inst[:inst] = Expr(:call, f, newargs...)
                 inst[:type] = Any
             elseif isa(stmt, PiNode)
                 # TODO: New PiNode that discriminates based on primal?

test/stage2_fwd.jl

@@ -43,4 +43,23 @@ module stage2_fwd
         g(x) = Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(f), Diffractor.TaylorBundle{1}(x, (1.0,)))
         Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(g), Diffractor.TaylorBundle{1}(10f0, (1.0,)))
     end
+
+    @testset "Constructors in forward_diff_no_inf!" begin
+        struct Bar148
+            v
+        end
+        foo_148(x) = Bar148(x)
+
+        # 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)
+            return Core.Compiler.insert_node!(ir, Core.SSAValue(1), Core.Compiler.NewInstruction(Expr(:call, Diffractor.ZeroBundle{1}, arg), Any))
+        end
+
+        ir = first(only(Base.code_ircode(foo_148, Tuple{Float64})))
+        Diffractor.forward_diff_no_inf!(ir, [Core.SSAValue(1) => 1]; transform! = identity_transform!)
+        ir2 = Core.Compiler.compact!(ir)
+        f = Core.OpaqueClosure(ir2; do_compile=false)
+        @test f(1.0) == Bar148(1.0)  # This would error if we were not handling constructors (%new) right
+    end
 end