Add method for checking rules added to DataType etc

Author: oxinabox

src/ChainRulesTestUtils.jl

@@ -15,6 +15,7 @@ export TestIterator
 export test_approx, test_scalar, test_frule, test_rrule, generate_well_conditioned_matrix
 export ⊢, rand_tangent
 export @maybe_inferred
+export test_method_tables_sensibility
 
 __init__() = init_test_inferred_setting!()
 
@@ -31,4 +32,6 @@ include("check_result.jl")
 include("rule_config.jl")
 include("finite_difference_calls.jl")
 include("testers.jl")
+
+include("global_checks.jl")
 end # module

src/global_checks.jl

@@ -0,0 +1,56 @@
+"""
+    _parameters(type)
+Extracts the type-parameters of the `type`.
+e.g. `_parameters(Foo{A, B, C}) == [A, B, C]`
+"""
+_parameters(sig::UnionAll) = _parameters(sig.body)
+_parameters(sig::DataType) = sig.parameters
+_parameters(sig::Union) = Base.uniontypes(sig)
+
+"""
+    test_method_tables_sensibility()
+
+Checks that the method tables for `rrule` and `frule` are sensible.
+This in future may carry out a number of checks, but presently just checks to make sure that
+no rules have been added to the very general `DataType`, `Union` or `UnionAll` types.
+This is easy to do when writing rules for constructors.
+It happens if you writeg. `rrule(::typeof(Foo), x)` rather than `rrule(::Type{<:Foo}, x)`:
+This would then actually define `rrule(::DataType, x)`. (or `UnionAll` if `Foo`
+was parametric, or `Union` if `Foo` was a type alias for a `Union`)
+"""
+function test_method_tables_sensibility()
+    @testset "Make sure methods haven't been added to DataType/UnionAll/Union" begin
+        # if someone wrote e.g. `rrule(::typeof(Foo), x)` rather than
+        # `rrule(::Type{<:Foo}, x)` then that would actually define `rrule(::DataType, x)`
+        # which would be bad. This test checks for that and fails if such a method exists.
+        for method in methods(rrule)
+            function_type = if method.sig <: Tuple{Any, RuleConfig, Type, Vararg}
+                _parameters(method.sig)[3]
+            elseif method.sig <: Tuple{Any, Type, Vararg}
+                _parameters(method.sig)[2]
+            else
+                nothing
+            end
+            
+            if function_type ∈ (DataType, UnionAll Union)
+                @error "Bad constructor rrule. typeof(T)` not `Type{T}`" method
+                @test false
+            end
+        end
+
+        # frule
+        for method in methods(frule)
+            function_type = if method.sig <: Tuple{Any, RuleConfig, Any, Type, Vararg}
+                _parameters(method.sig)[4]
+            elseif method.sig <: Tuple{Any, Any, Type, Vararg}
+                @show _parameters(method.sig)[3]
+            else
+                nothing
+            end
+            
+            if function_type ∈ (DataType, UnionAll Union)
+                @error "Bad constructor frule. typeof(T)` not `Type{T}`" method
+                @test false
+            end
+        end
+end

test/global_checks.jl

@@ -0,0 +1,4 @@
+@testset "global_checks.jl" begin
+    # Just check it doesn't error.
+    test_method_tables_sensibility()
+end

test/runtests.jl

@@ -16,4 +16,6 @@ ChainRulesTestUtils.TEST_INFERRED[] = true
     include("testers.jl")
     include("data_generation.jl")
     include("rand_tangent.jl")
+
+    include("global_checks.jl")
 end