Add debug mode for bad inplace

Author: oxinabox

docs/src/debug_mode.md

@@ -11,3 +11,8 @@ To enable, redefine the [`ChainRulesCore.debug_mode`](@ref) function to return `
 ```julia
 ChainRulesCore.debug_mode() = true
 ```
+
+## Features of Debug Mode:
+
+ - If you add a `Composite` to a primal value, and it was unable to construct a new primal values, then a better error message will be displayed detailing what overloads need to be written to fix this.
+ - during [`add!!`](@ref), if an `InplaceThunk` is used, and it runs the code that is supposed to run in place, but the return result is not the input (with updated values), then an error is thrown. Rather than silently using what ever values were returned.

src/accumulation.jl

@@ -14,7 +14,11 @@ The specialization of `add!!` for [`InplaceableThunk`](@ref) promises to only ca
 """
 function add!!(x, t::InplaceableThunk)
     return if is_inplaceable_destination(x)
-        t.add!(x)
+        if !debug_mode()
+            t.add!(x)
+        else
+            return debug_add!(x, t)
+        end
     else
         x + t
     end
@@ -51,3 +55,28 @@ function is_inplaceable_destination(x::AbstractArray)
     # processing and so are mutable if their `parent` is.
     return is_inplaceable_destination(p)
 end
+
+function debug_add!(accumuland, t::InplaceableThunk)
+    returned_value = t.add!(accumuland)
+    if returned_value !== accumuland
+        throw(BadInplaceException(t, accumuland, returned_value))
+    end
+    return returned_value
+end
+struct BadInplaceException <: Exception
+    t::InplaceableThunk
+    accumuland
+    returned_value
+end
+function Base.showerror(io::IO, err::BadInplaceException)
+    println(io, "add! of $(err.t) did not return an updated accumuland.")
+    println(io, "accumuland = $(err.accumuland)")
+    println(io, "returned_value = $(err.returned_value)")
+
+    if err.accumuland == err.returned_value
+        println(
+            io,
+            "Which in this case happenned to be equal. But they are not the same object."
+        )
+    end
+end

test/accumulation.jl

@@ -106,5 +106,30 @@
                 @test accumuland == [1.0 2.0; 3.0 4.0]  # must not have mutated
             end
         end
+
+        @testset "not actually inplace but said it was" begin
+            ithunk = InplaceableThunk(
+                @thunk(@assert false),  # this should never be used in this test
+                x -> 77*ones(2, 2)  # not actually inplace (also wrong)
+            )
+            accumuland = ones(2, 2)
+            @assert ChainRulesCore.debug_mode() == false
+            # without debug being enabled should return the result, not error
+            @test 77*ones(2, 2) == add!!(accumuland, ithunk)
+
+            ChainRulesCore.debug_mode() = true  # enable debug mode
+            # with debug being enabled should error
+            @test_throws ChainRulesCore.BadInplaceException add!!(accumuland, ithunk)
+            ChainRulesCore.debug_mode() = false  # disable it again
+        end
+    end
+    @testset "showerror BadInplaceException" begin
+        BadInplaceException = ChainRulesCore.BadInplaceException
+        ithunk = InplaceableThunk(@thunk(@assert false), x̄->nothing)
+        msg = sprint(showerror, BadInplaceException(ithunk, [22], [23]))
+        @test occursin("22", msg)
+
+        msg_equal = sprint(showerror, BadInplaceException(ithunk, [22], [22]))
+        @test occursin("equal", msg_equal)
     end
 end