handle Expr(:isdefined,) in forwards mode

Author: oxinabox

src/codegen/forward.jl

@@ -10,8 +10,9 @@ function fwd_transform!(ci, mi, nargs, N)
     ssa_mapping = Int[]
     loc_mapping = Int[]
 
-    function emit!(@nospecialize stmt)
-        (isexpr(stmt, :call) || isexpr(stmt, :(=)) || isexpr(stmt, :new)) || return stmt
+    emit!(@nospecialize stmt) = stmt
+    function emit!(stmt::Expr)
+        stmt.head ∈ (:call, :(=), :new, :isdefined) || return stmt
         push!(new_code, stmt)
         push!(new_codelocs, isempty(new_codelocs) ? 0 : new_codelocs[end])
         return SSAValue(length(new_code))
@@ -58,7 +59,8 @@ function fwd_transform!(ci, mi, nargs, N)
             # Can't trust that meta annotations are still valid in the AD'd
             # version.
             return nothing
-        
+        elseif isexpr(stmt, :isdefined)
+            return Expr(:call, ZeroBundle{N}, emit!(stmt))
         # Always disable `@inbounds`, as we don't actually know if the AD'd
         # code is truly `@inbounds` or not.
         elseif isexpr(stmt, :boundscheck)

test/forward.jl

@@ -147,6 +147,19 @@ end
     end
 end
 
+@testset "isdefined" begin
+    function foo_isdefined(x)
+        if (@noinline rand()) < 2
+            a=1  # always happens
+        end
+
+        2*@isdefined(a)*x + @isdefined(_thing_that_is_not_defined)
+    end
+    let var"'" = Diffractor.PrimeDerivativeFwd
+        @test foo_isdefined'(100.0) == 2.0
+    end
+end
+
 
 @testset "taylor_compatible" begin
     taylor_compatible = Diffractor.taylor_compatible