From 96c3d55f338c87f5bb95937f636136f2deb62360 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Fri, 18 Sep 2026 16:29:18 +0200
Subject: [PATCH 1/6] Define diff rules as methods and move provider rules to
 extensions

`@define_diffrule` wrote into a global `Dict`, so rules declared outside
DiffRules were lost during precompilation and never reached downstream
packages (#90). It now defines a method of `diffrule` instead, making the
method table the registry: rules declared anywhere are precompiled and
visible, and `diffrules`/`hasdiffrule` read them back by scanning it.

That in turn allows the SpecialFunctions, NaNMath and LogExpFunctions
rules to move into package extensions (#106), which drops those packages
from `[deps]`. Depending on DiffRules now installs 2 packages instead of
22, and no longer pulls in OpenSpecFun_jll. `NaNMath.lgamma` is the one
rule reaching across packages, so it lives in an extension triggered by
both NaNMath and SpecialFunctions.

Rules interpolate the functions they call rather than naming them, so a
rule can be evaluated without its defining package being in scope, and
cannot pick up a shadowed binding.

Lookup by function object and by module is supported alongside the
existing symbol-keyed API:

    diffrule(sin, :x)
    diffrule(Base, :sin, :x)
    diffrule(:Base, :sin, :x)

The module of a rule is derived via `Base.moduleroot`, so functions such
as `Base.sec` that live in `Base.Math` keep reporting `:Base` and stay
splice-able as `M.f` downstream.

`test/baseline.jl` lists every rule key so downstream code generation is
guarded against unintended changes, and `test/precompilation.jl` checks
in a separate process that rules defined in another package survive
precompilation and that the extensions load while a dependent is
precompiled.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 Project.toml                               |  10 +-
 ext/DiffRulesLogExpFunctionsExt.jl         |  34 ++++
 ext/DiffRulesNaNMathExt.jl                 |  36 ++++
 ext/DiffRulesNaNMathSpecialFunctionsExt.jl |  11 ++
 ext/DiffRulesSpecialFunctionsExt.jl        | 152 +++++++++++++++
 src/api.jl                                 | 150 +++++++--------
 src/rules.jl                               | 214 ---------------------
 test/baseline.jl                           | 159 +++++++++++++++
 test/precompilation.jl                     |  48 +++++
 test/registry.jl                           |  61 ++++++
 test/runtests.jl                           |   3 +
 11 files changed, 582 insertions(+), 296 deletions(-)
 create mode 100644 ext/DiffRulesLogExpFunctionsExt.jl
 create mode 100644 ext/DiffRulesNaNMathExt.jl
 create mode 100644 ext/DiffRulesNaNMathSpecialFunctionsExt.jl
 create mode 100644 ext/DiffRulesSpecialFunctionsExt.jl
 create mode 100644 test/baseline.jl
 create mode 100644 test/precompilation.jl
 create mode 100644 test/registry.jl

diff --git a/Project.toml b/Project.toml
index e30e835..3a6f536 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,16 +4,22 @@ version = "1.17.0-dev"
 
 [deps]
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
+
+[weakdeps]
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
+[extensions]
+DiffRulesLogExpFunctionsExt = "LogExpFunctions"
+DiffRulesNaNMathExt = "NaNMath"
+DiffRulesNaNMathSpecialFunctionsExt = ["NaNMath", "SpecialFunctions"]
+DiffRulesSpecialFunctionsExt = "SpecialFunctions"
+
 [compat]
 IrrationalConstants = "0.1.1, 0.2"
 LogExpFunctions = "0.3.2, 1"
 NaNMath = "0.3, 1"
-Random = "1"
 SpecialFunctions = "1.1, 2"
 julia = "1.10"
 
diff --git a/ext/DiffRulesLogExpFunctionsExt.jl b/ext/DiffRulesLogExpFunctionsExt.jl
new file mode 100644
index 0000000..a5d7b3b
--- /dev/null
+++ b/ext/DiffRulesLogExpFunctionsExt.jl
@@ -0,0 +1,34 @@
+module DiffRulesLogExpFunctionsExt
+
+using DiffRules: @define_diffrule
+using LogExpFunctions
+
+###################
+# LogExpFunctions #
+###################
+
+# unary
+@define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x))
+@define_diffrule LogExpFunctions.logistic(x) = :(z = $(LogExpFunctions.logistic)($x); z * (1 - z))
+@define_diffrule LogExpFunctions.logit(x) = :(inv($x * (1 - $x)))
+@define_diffrule LogExpFunctions.log1psq(x) = :(2 * $x / (1 + $x^2))
+@define_diffrule LogExpFunctions.log1pexp(x) = :($(LogExpFunctions.logistic)($x))
+@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - $(LogExpFunctions.log1mexp)($x)))
+@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - $(LogExpFunctions.log2mexp)($x)))
+@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - $(LogExpFunctions.logexpm1)($x)))
+@define_diffrule LogExpFunctions.log1pmx(x) = :(-$x / (1 + $x))
+@define_diffrule LogExpFunctions.logmxp1(x) = :((1 - $x) / $x)
+
+# binary
+@define_diffrule LogExpFunctions.xlogy(x, y) =
+    :(log($y)),
+    :(z = $x / $y; iszero($x) && !isnan($y) ? zero(z) : z)
+@define_diffrule LogExpFunctions.logaddexp(x, y) =
+    :(exp($x - $(LogExpFunctions.logaddexp)($x, $y))), :(exp($y - $(LogExpFunctions.logaddexp)($x, $y)))
+@define_diffrule LogExpFunctions.logsubexp(x, y) =
+    :(z = $(LogExpFunctions.logsubexp)($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
+    :(z = $(LogExpFunctions.logsubexp)($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
+@define_diffrule LogExpFunctions.xlog1py(x, y) =
+    :(log1p($y)),
+    :(z = $x / (1 + $y); iszero($x) && !isnan($y) ? zero(z) : z)
+end # module
diff --git a/ext/DiffRulesNaNMathExt.jl b/ext/DiffRulesNaNMathExt.jl
new file mode 100644
index 0000000..1ac41e0
--- /dev/null
+++ b/ext/DiffRulesNaNMathExt.jl
@@ -0,0 +1,36 @@
+module DiffRulesNaNMathExt
+
+using DiffRules: @define_diffrule
+using IrrationalConstants: logtwo, logten
+using NaNMath
+
+###########
+# NaNMath #
+###########
+
+# unary #
+#-------#
+
+@define_diffrule NaNMath.sqrt(x)   = :(  inv(2 * $(NaNMath.sqrt)($x))                  )
+@define_diffrule NaNMath.sin(x)    = :(  $(NaNMath.cos)($x)                            )
+@define_diffrule NaNMath.cos(x)    = :( -$(NaNMath.sin)($x)                            )
+@define_diffrule NaNMath.tan(x)    = :(  1 + $(NaNMath.pow)($(NaNMath.tan)($x), 2)        )
+@define_diffrule NaNMath.asin(x)   = :(  inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2)))  )
+@define_diffrule NaNMath.acos(x)   = :(  -inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2))) )
+@define_diffrule NaNMath.acosh(x)  = :(  inv($(NaNMath.sqrt)($(NaNMath.pow)($x, 2) - 1))  )
+@define_diffrule NaNMath.atanh(x)  = :(  inv(1 - $(NaNMath.pow)($x, 2))                )
+@define_diffrule NaNMath.log(x)    = :(  inv($x)                                    )
+@define_diffrule NaNMath.log2(x)   = :(  inv($logtwo * $x)                          )
+@define_diffrule NaNMath.log10(x)  = :(  inv($logten * $x)                          )
+@define_diffrule NaNMath.log1p(x)  = :(  inv($x + 1)                                )
+
+
+# binary #
+#--------#
+
+@define_diffrule NaNMath.pow(x, y) = :( $y * $(NaNMath.pow)($x, ($y - 1)) ), :( $(NaNMath.pow)($x, $y) * $(NaNMath.log)($x) )
+@define_diffrule NaNMath.max(x, y) = :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
+                                     :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($y), zero($y))))
+@define_diffrule NaNMath.min(x, y) = :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
+                                     :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($x), zero($x))))
+end # module
diff --git a/ext/DiffRulesNaNMathSpecialFunctionsExt.jl b/ext/DiffRulesNaNMathSpecialFunctionsExt.jl
new file mode 100644
index 0000000..dcda098
--- /dev/null
+++ b/ext/DiffRulesNaNMathSpecialFunctionsExt.jl
@@ -0,0 +1,11 @@
+module DiffRulesNaNMathSpecialFunctionsExt
+
+# `NaNMath.lgamma` is the one rule whose derivative reaches into another package,
+# so it needs both loaded (JuliaDiff/DiffRules.jl#106).
+
+using DiffRules: @define_diffrule
+using NaNMath
+using SpecialFunctions
+
+@define_diffrule NaNMath.lgamma(x) = :(  $(SpecialFunctions.digamma)($x)               )
+end # module
diff --git a/ext/DiffRulesSpecialFunctionsExt.jl b/ext/DiffRulesSpecialFunctionsExt.jl
new file mode 100644
index 0000000..f694cc0
--- /dev/null
+++ b/ext/DiffRulesSpecialFunctionsExt.jl
@@ -0,0 +1,152 @@
+module DiffRulesSpecialFunctionsExt
+
+using DiffRules: @define_diffrule
+using IrrationalConstants: sqrtπ, invsqrtπ
+using SpecialFunctions
+
+@define_diffrule SpecialFunctions.gamma(x) =
+    :(  $(SpecialFunctions.digamma)($x) * $(SpecialFunctions.gamma)($x)  )
+@define_diffrule SpecialFunctions.loggamma(x) =
+    :(  $(SpecialFunctions.digamma)($x)  )
+
+####################
+# SpecialFunctions #
+####################
+
+# unary #
+#-------#
+
+@define_diffrule SpecialFunctions.erf(x)         = :(  2 * ($invsqrtπ * exp(-$x^2))       )
+@define_diffrule SpecialFunctions.erfinv(x)      =
+    :(  ($sqrtπ * exp($(SpecialFunctions.erfinv)($x)^2)) / 2  )
+@define_diffrule SpecialFunctions.erfc(x)        = :( -($invsqrtπ * exp(-$x^2) * 2)       )
+@define_diffrule SpecialFunctions.logerfc(x)     =
+    :( - 2 * ($invsqrtπ * exp(- $x^2 - $(SpecialFunctions.logerfc)($x))) )
+
+@define_diffrule SpecialFunctions.erfcinv(x)     =
+    :( -($sqrtπ * exp($(SpecialFunctions.erfcinv)($x)^2)) / 2  )
+@define_diffrule SpecialFunctions.erfi(x)        = :(  $invsqrtπ * exp($x^2) * 2        )
+@define_diffrule SpecialFunctions.erfcx(x)       =
+    :(  2 * (($x * $(SpecialFunctions.erfcx)($x)) - $invsqrtπ)  )
+@define_diffrule SpecialFunctions.logerfcx(x) =
+    :(  2 * ($x - inv($(SpecialFunctions.erfcx)($x) * $sqrtπ))  )
+
+@define_diffrule SpecialFunctions.dawson(x)      =
+    :(  1 - (2 * $x * $(SpecialFunctions.dawson)($x))  )
+@define_diffrule SpecialFunctions.digamma(x) =
+    :(  $(SpecialFunctions.trigamma)($x)  )
+@define_diffrule SpecialFunctions.invdigamma(x)  =
+    :(  inv($(SpecialFunctions.trigamma)($(SpecialFunctions.invdigamma)($x)))  )
+@define_diffrule SpecialFunctions.trigamma(x)    =
+    :(  $(SpecialFunctions.polygamma)(2, $x)  )
+
+# derivatives for `airybix` and `airybiprimex` are only correct for real inputs
+# `airyaix` and `airyaiprimex` are only defined for positive real inputs
+# `airybix` and `airybiprimex` are unscaled for negative real inputs
+@define_diffrule SpecialFunctions.airyai(x)      =
+    :(  $(SpecialFunctions.airyaiprime)($x)  )
+@define_diffrule SpecialFunctions.airyaiprime(x) =
+    :(  $x * $(SpecialFunctions.airyai)($x)  )
+@define_diffrule SpecialFunctions.airyaix(x)     =
+    :(  $(SpecialFunctions.airyaiprimex)($x) + sqrt($x) * $(SpecialFunctions.airyaix)($x)  )
+@define_diffrule SpecialFunctions.airyaiprimex(x) =
+    :(  $x * $(SpecialFunctions.airyaix)($x) + sqrt($x) * $(SpecialFunctions.airyaiprimex)($x)  )
+@define_diffrule SpecialFunctions.airybi(x)      =
+    :(  $(SpecialFunctions.airybiprime)($x)  )
+@define_diffrule SpecialFunctions.airybiprime(x) =
+    :(  $x * $(SpecialFunctions.airybi)($x)  )
+@define_diffrule SpecialFunctions.airybix(x)     =
+    :(  if $x > zero($x)
+            $(SpecialFunctions.airybiprimex)($x) - sqrt($x) * $(SpecialFunctions.airybix)($x)
+        else
+            $(SpecialFunctions.airybiprimex)($x)
+        end  )
+@define_diffrule SpecialFunctions.airybiprimex(x) =
+    :(  if $x > zero($x)
+            $x * $(SpecialFunctions.airybix)($x) - sqrt($x) * $(SpecialFunctions.airybiprimex)($x)
+        else
+            $x * $(SpecialFunctions.airybix)($x)
+        end  )
+
+@define_diffrule SpecialFunctions.besselj0(x)    =
+    :( -$(SpecialFunctions.besselj1)($x)  )
+@define_diffrule SpecialFunctions.besselj1(x)    =
+    :(  ($(SpecialFunctions.besselj0)($x) - $(SpecialFunctions.besselj)(2, $x)) / 2  )
+@define_diffrule SpecialFunctions.bessely0(x)    =
+    :( -$(SpecialFunctions.bessely1)($x)  )
+@define_diffrule SpecialFunctions.bessely1(x)    =
+    :(  ($(SpecialFunctions.bessely0)($x) - $(SpecialFunctions.bessely)(2, $x)) / 2 )
+
+@define_diffrule SpecialFunctions.sinint(x) = :( sinc($x / π) )
+@define_diffrule SpecialFunctions.cosint(x) = :( cos($x) / $x )
+
+@define_diffrule SpecialFunctions.ellipk(m) =
+    :( ($(SpecialFunctions.ellipe)($m) / (1 - $m) - $(SpecialFunctions.ellipk)($m)) / (2 * $m) )
+@define_diffrule SpecialFunctions.ellipe(m) =
+    :( ($(SpecialFunctions.ellipe)($m) - $(SpecialFunctions.ellipk)($m)) / (2 * $m) )
+
+@define_diffrule SpecialFunctions.expint(x)    = :( -exp(-$x) / $x )
+
+# TODO:
+#
+# eta
+# zeta
+
+# binary #
+#--------#
+
+@define_diffrule SpecialFunctions.erf(x, y) =
+    :(  -2 * ($invsqrtπ * exp(-$x^2))  ), :(  2 * ($invsqrtπ * exp(-$y^2))  )
+
+# derivatives with respect to the order `ν` exist but are not implemented
+# (analogously to the ChainRules definitions in SpecialFunctions)
+
+# derivatives for `besselix`, `besseljx` and `besselyx` are only correct for real inputs
+# see https://github.com/JuliaMath/SpecialFunctions.jl/blob/master/src/chainrules.jl
+# for forward-mode and reverse-mode derivatives for complex inputs
+
+@define_diffrule SpecialFunctions.besselj(ν, x)   =
+    :NaN, :(  ($(SpecialFunctions.besselj)($ν - 1, $x) - $(SpecialFunctions.besselj)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besseljx(ν, x)  =
+    :NaN, :(  ($(SpecialFunctions.besseljx)($ν - 1, $x) - $(SpecialFunctions.besseljx)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besseli(ν, x)   =
+    :NaN, :(  ($(SpecialFunctions.besseli)($ν - 1, $x) + $(SpecialFunctions.besseli)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselix(ν, x)  =
+    :NaN, :(  ($(SpecialFunctions.besselix)($ν - 1, $x) + $(SpecialFunctions.besselix)($ν + 1, $x)) / 2 - sign($x) * $(SpecialFunctions.besselix)($ν, $x)  )
+@define_diffrule SpecialFunctions.bessely(ν, x)   =
+    :NaN, :(  ($(SpecialFunctions.bessely)($ν - 1, $x) - $(SpecialFunctions.bessely)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselyx(ν, x)  =
+    :NaN, :(  ($(SpecialFunctions.besselyx)($ν - 1, $x) - $(SpecialFunctions.besselyx)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselk(ν, x)   =
+    :NaN, :( -($(SpecialFunctions.besselk)($ν - 1, $x) + $(SpecialFunctions.besselk)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselkx(ν, x)  =
+    :NaN, :( -($(SpecialFunctions.besselkx)($ν - 1, $x) + $(SpecialFunctions.besselkx)($ν + 1, $x)) / 2 + $(SpecialFunctions.besselkx)($ν, $x)  )
+@define_diffrule SpecialFunctions.besselh(ν, x)   =
+    :NaN, :(  ($(SpecialFunctions.besselh)($ν - 1, $x) - $(SpecialFunctions.besselh)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselhx(ν, x) =
+    :NaN, :(  ($(SpecialFunctions.besselhx)($ν - 1, $x) - $(SpecialFunctions.besselhx)($ν + 1, $x)) / 2 - im * $(SpecialFunctions.besselhx)($ν, $x)  )
+@define_diffrule SpecialFunctions.hankelh1(ν, x)  =
+    :NaN, :(  ($(SpecialFunctions.hankelh1)($ν - 1, $x) - $(SpecialFunctions.hankelh1)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.hankelh1x(ν, x) =
+    :NaN, :(  ($(SpecialFunctions.hankelh1x)($ν - 1, $x) - $(SpecialFunctions.hankelh1x)($ν + 1, $x)) / 2 - im * $(SpecialFunctions.hankelh1x)($ν, $x)  )
+@define_diffrule SpecialFunctions.hankelh2(ν, x)  =
+    :NaN, :(  ($(SpecialFunctions.hankelh2)($ν - 1, $x) - $(SpecialFunctions.hankelh2)($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.hankelh2x(ν, x) =
+    :NaN, :(  ($(SpecialFunctions.hankelh2x)($ν - 1, $x) - $(SpecialFunctions.hankelh2x)($ν + 1, $x)) / 2 + im * $(SpecialFunctions.hankelh2x)($ν, $x)  )
+
+@define_diffrule SpecialFunctions.polygamma(m, x) =
+    :NaN, :(  $(SpecialFunctions.polygamma)($m + 1, $x)  )
+
+@define_diffrule SpecialFunctions.beta(a, b)      =
+    :( $(SpecialFunctions.beta)($a, $b)*($(SpecialFunctions.digamma)($a) - $(SpecialFunctions.digamma)($a + $b)) ), :(  $(SpecialFunctions.beta)($a, $b)*($(SpecialFunctions.digamma)($b) - $(SpecialFunctions.digamma)($a + $b))     )
+@define_diffrule SpecialFunctions.logbeta(a, b)   =
+    :( $(SpecialFunctions.digamma)($a) - $(SpecialFunctions.digamma)($a + $b)  ), :(  $(SpecialFunctions.digamma)($b) - $(SpecialFunctions.digamma)($a + $b)  )
+
+# derivative wrt to `ν` is not implemented
+@define_diffrule SpecialFunctions.expint(ν, x)    = 
+    :NaN, :( -$(SpecialFunctions.expint)($ν - 1, $x) )
+
+# derivative wrt to `s` is not implemented
+@define_diffrule SpecialFunctions.zeta(s, z)      =
+    :NaN, :( - $s * $(SpecialFunctions.zeta)($s + 1, $z)  )
+end # module
diff --git a/src/api.jl b/src/api.jl
index a045686..7c48754 100644
--- a/src/api.jl
+++ b/src/api.jl
@@ -1,13 +1,10 @@
-
-const DEFINED_DIFFRULES = Dict{Tuple{Union{Expr,Symbol},Symbol,Int},Any}()
-
 """
     @define_diffrule M.f(x) = :(df_dx(\$x))
     @define_diffrule M.f(x, y) = :(df_dx(\$x, \$y)), :(df_dy(\$x, \$y))
     ⋮
 
 Define a new differentiation rule for the function `M.f` and the given arguments, which should
-be treated as bindings to Julia expressions. Return the defined rule's key.
+be treated as bindings to Julia expressions.
 
 The LHS should be a function call with a non-splatted argument list, and the RHS should be
 the derivative expression, or in the `n`-ary case, an `n`-tuple of expressions where the
@@ -16,6 +13,9 @@ interpolated wherever they are used on the RHS.
 
 Note that differentiation rules are purely symbolic, so no type annotations should be used.
 
+Each rule is a method of [`diffrule`](@ref), so rules defined in other packages and in
+package extensions are precompiled and visible like any other method.
+
 # Examples
 
 ```julia
@@ -26,26 +26,19 @@ Note that differentiation rules are purely symbolic, so no type annotations shou
 """
 macro define_diffrule(def)
     @assert isa(def, Expr) && def.head == :(=) "Diff rule expression does not have a left and right side"
-    lhs = def.args[1]
-    rhs = def.args[2]
+    lhs, rhs = def.args
     @assert isa(lhs, Expr) && lhs.head == :call "LHS is not a function call"
-    qualified_f = lhs.args[1]
-    @assert isa(qualified_f, Expr) && qualified_f.head == :(.) "Function is not qualified by module"
-    M = qualified_f.args[1]
-    f = _get_quoted_symbol(qualified_f.args[2])
+    f = lhs.args[1]
+    @assert isa(f, Expr) && f.head == :(.) "Function is not qualified by module"
     args = lhs.args[2:end]
-    rule = Expr(:->, Expr(:tuple, args...), rhs)
-    key = Expr(:tuple, Expr(:quote, M), Expr(:quote, f), length(args))
-    return esc(quote
-        $DiffRules.DEFINED_DIFFRULES[$key] = $rule
-        $key
-    end)
+    return esc(:($DiffRules.diffrule(::typeof($f), $(args...)) = $rhs))
 end
 
 """
-    diffrule(M::Union{Expr,Symbol}, f::Symbol, args...)
+    diffrule(f, args...)
+    diffrule(M::Union{Module,Symbol}, f::Symbol, args...)
 
-Return the derivative expression for `M.f` at the given argument(s), with the argument(s)
+Return the derivative expression for `f` at the given argument(s), with the argument(s)
 interpolated into the returned expression.
 
 In the `n`-ary case, an `n`-tuple of expressions will be returned where the `i`th expression
@@ -54,46 +47,58 @@ is the derivative of `f` w.r.t the `i`th argument.
 # Examples
 
 ```jldoctest
-julia> DiffRules.diffrule(:Base, :sin, 1)
+julia> DiffRules.diffrule(sin, 1)
 :(cos(1))
 
-julia> DiffRules.diffrule(:Base, :sin, :x)
+julia> DiffRules.diffrule(Base, :sin, :x)
 :(cos(x))
 
 julia> DiffRules.diffrule(:Base, :sin, :(x * y^2))
 :(cos(x * y ^ 2))
 ```
 """
-diffrule(M::Union{Expr,Symbol}, f::Symbol, args...) = DEFINED_DIFFRULES[M,f,length(args)](args...)
+function diffrule end
+
+diffrule(M::Module, f::Symbol, args...) = diffrule(getproperty(M, f), args...)
+
+function diffrule(M::Symbol, f::Symbol, args...)
+    fn = _resolve(M, f)
+    fn === nothing && throw(KeyError((M, f, length(args))))
+    return diffrule(fn, args...)
+end
 
 """
-    hasdiffrule(M::Union{Expr,Symbol}, f::Symbol, arity::Int)
+    hasdiffrule(f, arity::Int)
+    hasdiffrule(M::Union{Module,Symbol}, f::Symbol, arity::Int)
 
-Return `true` if a differentiation rule is defined for `M.f` and `arity`, or return `false`
-otherwise.
+Return `true` if a differentiation rule is defined for `f` and `arity`, or return `false`
+otherwise. Here, `arity` refers to the number of arguments accepted by `f`.
 
-Here, `arity` refers to the number of arguments accepted by `f`.
+Rules for a package's functions exist only once that package is loaded, so a query for an
+unloaded package returns `false`.
 
 # Examples
 
 ```jldoctest
-julia> DiffRules.hasdiffrule(:Base, :sin, 1)
+julia> DiffRules.hasdiffrule(sin, 1)
 true
 
-julia> DiffRules.hasdiffrule(:Base, :sin, 2)
+julia> DiffRules.hasdiffrule(sin, 2)
 false
 
-julia> DiffRules.hasdiffrule(:Base, :-, 1)
-true
-
 julia> DiffRules.hasdiffrule(:Base, :-, 2)
 true
-
-julia> DiffRules.hasdiffrule(:Base, :-, 3)
-false
 ```
 """
-hasdiffrule(M::Union{Expr,Symbol}, f::Symbol, arity::Int) = haskey(DEFINED_DIFFRULES, (M, f, arity))
+hasdiffrule(f, arity::Int) = hasmethod(diffrule, Tuple{typeof(f),Vararg{Any,arity}})
+
+hasdiffrule(M::Module, f::Symbol, arity::Int) =
+    isdefined(M, f) && hasdiffrule(getproperty(M, f), arity)
+
+function hasdiffrule(M::Symbol, f::Symbol, arity::Int)
+    fn = _resolve(M, f)
+    return fn !== nothing && hasdiffrule(fn, arity)
+end
 
 # show a deprecation warning if `filter_modules` in `diffrules()` is specified implicitly
 # we use a custom singleton to figure out if the keyword argument was set explicitly
@@ -119,13 +124,14 @@ end
     diffrules(; filter_modules=(:Base, :SpecialFunctions, :NaNMath))
 
 Return a list of keys that can be used to access all defined differentiation rules for
-modules in `filter_modules`.
+modules in `filter_modules`. To include all rules, specify `filter_modules = nothing`.
 
-Each key is of the form `(M::Symbol, f::Symbol, arity::Int)`.
-Here, `arity` refers to the number of arguments accepted by `f` and `M` is one of the
-modules in `filter_modules`.
+Each key is of the form `(M::Symbol, f::Symbol, arity::Int)`, where `M` is the name of the
+package defining `f` and `arity` is the number of arguments accepted by `f`.
 
-To include all rules, specify `filter_modules = nothing`.
+Keys are collected from the method table of [`diffrule`](@ref), so rules defined in other
+packages are included. Rules for a package's functions exist only once that package is
+loaded: querying before `using SpecialFunctions` will not list its rules.
 
 !!! note
     Calling `diffrules()` with the implicit default keyword argument `filter_modules`
@@ -145,55 +151,39 @@ true
 julia> (:Base, :log, 1) in DiffRules.diffrules()
 true
 
-julia> (:Base, :*, 2) in DiffRules.diffrules()
-true
-```
-
-If you call `diffrules()`, only rules for Base, SpecialFunctions, and
-NaNMath are returned but no rules for LogExpFunctions:
-```jldoctest
-julia> any(M === :LogExpFunctions for (M, _, _) in DiffRules.diffrules())
-false
-```
-
-If you set `filter_modules=nothing`, all rules defined in DiffRules are
-returned and in particular also rules for LogExpFunctions:
-```jldoctest
-julia> any(
-           M === :LogExpFunctions
-           for (M, _, _) in DiffRules.diffrules(; filter_modules=nothing)
-       )
-true
-```
-
-If you set `filter_modules=(:Base,)` only rules for functions in Base are
-returned:
-```jldoctest
 julia> all(M === :Base for (M, _, _) in DiffRules.diffrules(; filter_modules=(:Base,)))
 true
 ```
 """
 function diffrules(; filter_modules=DefaultFilterModules())
     modules = deprecated_modules(filter_modules)
-    return if modules === nothing
-        keys(DEFINED_DIFFRULES)
-    else
-        Iterators.filter(keys(DEFINED_DIFFRULES)) do (M, _, _)
-            return M in modules
-        end
-    end
+    rules = [(_pkgname(fn), nameof(fn), arity) for (fn, arity) in _rules()]
+    modules === nothing && return rules
+    return filter(r -> r[1] in modules, rules)
 end
 
-# For v0.6 and v0.7 compatibility, need to support having the diff rule function enter as a
-# `Expr(:quote...)` and a `QuoteNode`. When v0.6 support is dropped, the function will
-# always enter in a `QuoteNode` (#23885).
-function _get_quoted_symbol(ex::Expr)
-    @assert ex.head == :quote
-    @assert length(ex.args) == 1 && isa(ex.args[1], Symbol) "Function not a single symbol"
-    ex.args[1]
+# `parentmodule` reports submodules such as `Base.Math`, which callers cannot splice into
+# `M.f`; the root module is the package name they expect.
+_pkgname(fn) = nameof(Base.moduleroot(parentmodule(fn)))
+
+# Rules are the methods of `diffrule` whose first parameter is a singleton function type.
+# The `Symbol`/`Module` methods above are not, and are skipped.
+function _rules()
+    rules = Tuple{Function,Int}[]
+    for m in methods(diffrule)
+        m.isva && continue
+        params = Base.unwrap_unionall(m.sig).parameters
+        length(params) >= 2 || continue
+        T = params[2]
+        T isa DataType && T <: Function && isdefined(T, :instance) || continue
+        push!(rules, (T.instance, length(params) - 2))
+    end
+    return rules
 end
 
-function _get_quoted_symbol(ex::QuoteNode)
-    @assert isa(ex.value, Symbol) "Function not a single symbol"
-    ex.value
+function _resolve(M::Symbol, f::Symbol)
+    for (fn, _) in _rules()
+        nameof(fn) === f && _pkgname(fn) === M && return fn
+    end
+    return nothing
 end
diff --git a/src/rules.jl b/src/rules.jl
index caa3078..4221866 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -61,11 +61,6 @@
 @define_diffrule Base.deg2rad(x)              = :(   deg2rad(one($x))  )
 @define_diffrule Base.mod2pi(x)               = :(  isinteger($x / $twoπ) ? oftype(float($x), NaN) : one(float($x)) )
 @define_diffrule Base.rad2deg(x)              = :(  rad2deg(one($x))  )
-@define_diffrule SpecialFunctions.gamma(x) =
-    :(  SpecialFunctions.digamma($x) * SpecialFunctions.gamma($x)  )
-@define_diffrule SpecialFunctions.loggamma(x) =
-    :(  SpecialFunctions.digamma($x)  )
-
 @define_diffrule Base.abs(x)                  = :( $(_abs_deriv)($x)            )
 
 # We provide this hook for special number types like `Interval`
@@ -104,212 +99,3 @@ _abs_deriv(x) = signbit(x) ? -one(x) : one(x)
 @define_diffrule Base.ifelse(p, x, y) = false, :($p), :(!$p)
 
 =#
-
-####################
-# SpecialFunctions #
-####################
-
-# unary #
-#-------#
-
-@define_diffrule SpecialFunctions.erf(x)         = :(  2 * ($invsqrtπ * exp(-$x^2))       )
-@define_diffrule SpecialFunctions.erfinv(x)      =
-    :(  ($sqrtπ * exp(SpecialFunctions.erfinv($x)^2)) / 2  )
-@define_diffrule SpecialFunctions.erfc(x)        = :( -($invsqrtπ * exp(-$x^2) * 2)       )
-@define_diffrule SpecialFunctions.logerfc(x)     =
-    :( - 2 * ($invsqrtπ * exp(- $x^2 - SpecialFunctions.logerfc($x))) )
-
-@define_diffrule SpecialFunctions.erfcinv(x)     =
-    :( -($sqrtπ * exp(SpecialFunctions.erfcinv($x)^2)) / 2  )
-@define_diffrule SpecialFunctions.erfi(x)        = :(  $invsqrtπ * exp($x^2) * 2        )
-@define_diffrule SpecialFunctions.erfcx(x)       =
-    :(  2 * (($x * SpecialFunctions.erfcx($x)) - $invsqrtπ)  )
-@define_diffrule SpecialFunctions.logerfcx(x) =
-    :(  2 * ($x - inv(SpecialFunctions.erfcx($x) * $sqrtπ))  )
-
-@define_diffrule SpecialFunctions.dawson(x)      =
-    :(  1 - (2 * $x * SpecialFunctions.dawson($x))  )
-@define_diffrule SpecialFunctions.digamma(x) =
-    :(  SpecialFunctions.trigamma($x)  )
-@define_diffrule SpecialFunctions.invdigamma(x)  =
-    :(  inv(SpecialFunctions.trigamma(SpecialFunctions.invdigamma($x)))  )
-@define_diffrule SpecialFunctions.trigamma(x)    =
-    :(  SpecialFunctions.polygamma(2, $x)  )
-
-# derivatives for `airybix` and `airybiprimex` are only correct for real inputs
-# `airyaix` and `airyaiprimex` are only defined for positive real inputs
-# `airybix` and `airybiprimex` are unscaled for negative real inputs
-@define_diffrule SpecialFunctions.airyai(x)      =
-    :(  SpecialFunctions.airyaiprime($x)  )
-@define_diffrule SpecialFunctions.airyaiprime(x) =
-    :(  $x * SpecialFunctions.airyai($x)  )
-@define_diffrule SpecialFunctions.airyaix(x)     =
-    :(  SpecialFunctions.airyaiprimex($x) + sqrt($x) * SpecialFunctions.airyaix($x)  )
-@define_diffrule SpecialFunctions.airyaiprimex(x) =
-    :(  $x * SpecialFunctions.airyaix($x) + sqrt($x) * SpecialFunctions.airyaiprimex($x)  )
-@define_diffrule SpecialFunctions.airybi(x)      =
-    :(  SpecialFunctions.airybiprime($x)  )
-@define_diffrule SpecialFunctions.airybiprime(x) =
-    :(  $x * SpecialFunctions.airybi($x)  )
-@define_diffrule SpecialFunctions.airybix(x)     =
-    :(  if $x > zero($x)
-            SpecialFunctions.airybiprimex($x) - sqrt($x) * SpecialFunctions.airybix($x)
-        else
-            SpecialFunctions.airybiprimex($x)
-        end  )
-@define_diffrule SpecialFunctions.airybiprimex(x) =
-    :(  if $x > zero($x)
-            $x * SpecialFunctions.airybix($x) - sqrt($x) * SpecialFunctions.airybiprimex($x)
-        else
-            $x * SpecialFunctions.airybix($x)
-        end  )
-
-@define_diffrule SpecialFunctions.besselj0(x)    =
-    :( -SpecialFunctions.besselj1($x)  )
-@define_diffrule SpecialFunctions.besselj1(x)    =
-    :(  (SpecialFunctions.besselj0($x) - SpecialFunctions.besselj(2, $x)) / 2  )
-@define_diffrule SpecialFunctions.bessely0(x)    =
-    :( -SpecialFunctions.bessely1($x)  )
-@define_diffrule SpecialFunctions.bessely1(x)    =
-    :(  (SpecialFunctions.bessely0($x) - SpecialFunctions.bessely(2, $x)) / 2 )
-
-@define_diffrule SpecialFunctions.sinint(x) = :( sinc($x / π) )
-@define_diffrule SpecialFunctions.cosint(x) = :( cos($x) / $x )
-
-@define_diffrule SpecialFunctions.ellipk(m) =
-    :( (SpecialFunctions.ellipe($m) / (1 - $m) - SpecialFunctions.ellipk($m)) / (2 * $m) )
-@define_diffrule SpecialFunctions.ellipe(m) =
-    :( (SpecialFunctions.ellipe($m) - SpecialFunctions.ellipk($m)) / (2 * $m) )
-
-@define_diffrule SpecialFunctions.expint(x)    = :( -exp(-$x) / $x )
-
-# TODO:
-#
-# eta
-# zeta
-
-# binary #
-#--------#
-
-@define_diffrule SpecialFunctions.erf(x, y) =
-    :(  -2 * ($invsqrtπ * exp(-$x^2))  ), :(  2 * ($invsqrtπ * exp(-$y^2))  )
-
-# derivatives with respect to the order `ν` exist but are not implemented
-# (analogously to the ChainRules definitions in SpecialFunctions)
-
-# derivatives for `besselix`, `besseljx` and `besselyx` are only correct for real inputs
-# see https://github.com/JuliaMath/SpecialFunctions.jl/blob/master/src/chainrules.jl
-# for forward-mode and reverse-mode derivatives for complex inputs
-
-@define_diffrule SpecialFunctions.besselj(ν, x)   =
-    :NaN, :(  (SpecialFunctions.besselj($ν - 1, $x) - SpecialFunctions.besselj($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besseljx(ν, x)  =
-    :NaN, :(  (SpecialFunctions.besseljx($ν - 1, $x) - SpecialFunctions.besseljx($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besseli(ν, x)   =
-    :NaN, :(  (SpecialFunctions.besseli($ν - 1, $x) + SpecialFunctions.besseli($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselix(ν, x)  =
-    :NaN, :(  (SpecialFunctions.besselix($ν - 1, $x) + SpecialFunctions.besselix($ν + 1, $x)) / 2 - sign($x) * SpecialFunctions.besselix($ν, $x)  )
-@define_diffrule SpecialFunctions.bessely(ν, x)   =
-    :NaN, :(  (SpecialFunctions.bessely($ν - 1, $x) - SpecialFunctions.bessely($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselyx(ν, x)  =
-    :NaN, :(  (SpecialFunctions.besselyx($ν - 1, $x) - SpecialFunctions.besselyx($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselk(ν, x)   =
-    :NaN, :( -(SpecialFunctions.besselk($ν - 1, $x) + SpecialFunctions.besselk($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselkx(ν, x)  =
-    :NaN, :( -(SpecialFunctions.besselkx($ν - 1, $x) + SpecialFunctions.besselkx($ν + 1, $x)) / 2 + SpecialFunctions.besselkx($ν, $x)  )
-@define_diffrule SpecialFunctions.besselh(ν, x)   =
-    :NaN, :(  (SpecialFunctions.besselh($ν - 1, $x) - SpecialFunctions.besselh($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselhx(ν, x) =
-    :NaN, :(  (SpecialFunctions.besselhx($ν - 1, $x) - SpecialFunctions.besselhx($ν + 1, $x)) / 2 - im * SpecialFunctions.besselhx($ν, $x)  )
-@define_diffrule SpecialFunctions.hankelh1(ν, x)  =
-    :NaN, :(  (SpecialFunctions.hankelh1($ν - 1, $x) - SpecialFunctions.hankelh1($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.hankelh1x(ν, x) =
-    :NaN, :(  (SpecialFunctions.hankelh1x($ν - 1, $x) - SpecialFunctions.hankelh1x($ν + 1, $x)) / 2 - im * SpecialFunctions.hankelh1x($ν, $x)  )
-@define_diffrule SpecialFunctions.hankelh2(ν, x)  =
-    :NaN, :(  (SpecialFunctions.hankelh2($ν - 1, $x) - SpecialFunctions.hankelh2($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.hankelh2x(ν, x) =
-    :NaN, :(  (SpecialFunctions.hankelh2x($ν - 1, $x) - SpecialFunctions.hankelh2x($ν + 1, $x)) / 2 + im * SpecialFunctions.hankelh2x($ν, $x)  )
-
-@define_diffrule SpecialFunctions.polygamma(m, x) =
-    :NaN, :(  SpecialFunctions.polygamma($m + 1, $x)  )
-
-@define_diffrule SpecialFunctions.beta(a, b)      =
-    :( SpecialFunctions.beta($a, $b)*(SpecialFunctions.digamma($a) - SpecialFunctions.digamma($a + $b)) ), :(  SpecialFunctions.beta($a, $b)*(SpecialFunctions.digamma($b) - SpecialFunctions.digamma($a + $b))     )
-@define_diffrule SpecialFunctions.logbeta(a, b)   =
-    :( SpecialFunctions.digamma($a) - SpecialFunctions.digamma($a + $b)  ), :(  SpecialFunctions.digamma($b) - SpecialFunctions.digamma($a + $b)  )
-
-# derivative wrt to `ν` is not implemented
-@define_diffrule SpecialFunctions.expint(ν, x)    = 
-    :NaN, :( -SpecialFunctions.expint($ν - 1, $x) )
-
-# derivative wrt to `s` is not implemented
-@define_diffrule SpecialFunctions.zeta(s, z)      =
-    :NaN, :( - $s * SpecialFunctions.zeta($s + 1, $z)  )
-
-# ternary #
-#---------#
-
-# TODO:
-#
-# besselh
-# besselhx
-
-###########
-# NaNMath #
-###########
-
-# unary #
-#-------#
-
-@define_diffrule NaNMath.sqrt(x)   = :(  inv(2 * NaNMath.sqrt($x))                  )
-@define_diffrule NaNMath.sin(x)    = :(  NaNMath.cos($x)                            )
-@define_diffrule NaNMath.cos(x)    = :( -NaNMath.sin($x)                            )
-@define_diffrule NaNMath.tan(x)    = :(  1 + NaNMath.pow(NaNMath.tan($x), 2)        )
-@define_diffrule NaNMath.asin(x)   = :(  inv(NaNMath.sqrt(1 - NaNMath.pow($x, 2)))  )
-@define_diffrule NaNMath.acos(x)   = :(  -inv(NaNMath.sqrt(1 - NaNMath.pow($x, 2))) )
-@define_diffrule NaNMath.acosh(x)  = :(  inv(NaNMath.sqrt(NaNMath.pow($x, 2) - 1))  )
-@define_diffrule NaNMath.atanh(x)  = :(  inv(1 - NaNMath.pow($x, 2))                )
-@define_diffrule NaNMath.log(x)    = :(  inv($x)                                    )
-@define_diffrule NaNMath.log2(x)   = :(  inv($logtwo * $x)                          )
-@define_diffrule NaNMath.log10(x)  = :(  inv($logten * $x)                          )
-@define_diffrule NaNMath.log1p(x)  = :(  inv($x + 1)                                )
-@define_diffrule NaNMath.lgamma(x) = :(  SpecialFunctions.digamma($x)               )
-
-
-# binary #
-#--------#
-
-@define_diffrule NaNMath.pow(x, y) = :( $y * NaNMath.pow($x, ($y - 1)) ), :( NaNMath.pow($x, $y) * NaNMath.log($x) )
-@define_diffrule NaNMath.max(x, y) = :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
-                                     :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($y), zero($y))))
-@define_diffrule NaNMath.min(x, y) = :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
-                                     :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($x), zero($x))))
-
-###################
-# LogExpFunctions #
-###################
-
-# unary
-@define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x))
-@define_diffrule LogExpFunctions.logistic(x) = :(z = LogExpFunctions.logistic($x); z * (1 - z))
-@define_diffrule LogExpFunctions.logit(x) = :(inv($x * (1 - $x)))
-@define_diffrule LogExpFunctions.log1psq(x) = :(2 * $x / (1 + $x^2))
-@define_diffrule LogExpFunctions.log1pexp(x) = :(LogExpFunctions.logistic($x))
-@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - LogExpFunctions.log1mexp($x)))
-@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - LogExpFunctions.log2mexp($x)))
-@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - LogExpFunctions.logexpm1($x)))
-@define_diffrule LogExpFunctions.log1pmx(x) = :(-$x / (1 + $x))
-@define_diffrule LogExpFunctions.logmxp1(x) = :((1 - $x) / $x)
-
-# binary
-@define_diffrule LogExpFunctions.xlogy(x, y) =
-    :(log($y)),
-    :(z = $x / $y; iszero($x) && !isnan($y) ? zero(z) : z)
-@define_diffrule LogExpFunctions.logaddexp(x, y) =
-    :(exp($x - LogExpFunctions.logaddexp($x, $y))), :(exp($y - LogExpFunctions.logaddexp($x, $y)))
-@define_diffrule LogExpFunctions.logsubexp(x, y) =
-    :(z = LogExpFunctions.logsubexp($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
-    :(z = LogExpFunctions.logsubexp($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
-@define_diffrule LogExpFunctions.xlog1py(x, y) =
-    :(log1p($y)),
-    :(z = $x / (1 + $y); iszero($x) && !isnan($y) ? zero(z) : z)
diff --git a/test/baseline.jl b/test/baseline.jl
new file mode 100644
index 0000000..8549d49
--- /dev/null
+++ b/test/baseline.jl
@@ -0,0 +1,159 @@
+# Every differentiation rule DiffRules defines, as (module, function, arity).
+# Downstream packages generate code by iterating `diffrules()` and splicing `M.f`,
+# so this list guards the keys they rely on against unintended changes.
+
+const BASELINE = [
+    (:Base, :*, 2),
+    (:Base, :+, 1),
+    (:Base, :+, 2),
+    (:Base, :-, 1),
+    (:Base, :-, 2),
+    (:Base, :/, 2),
+    (:Base, :\, 2),
+    (:Base, :^, 2),
+    (:Base, :abs, 1),
+    (:Base, :abs2, 1),
+    (:Base, :acos, 1),
+    (:Base, :acosd, 1),
+    (:Base, :acosh, 1),
+    (:Base, :acot, 1),
+    (:Base, :acotd, 1),
+    (:Base, :acoth, 1),
+    (:Base, :acsc, 1),
+    (:Base, :acscd, 1),
+    (:Base, :acsch, 1),
+    (:Base, :asec, 1),
+    (:Base, :asecd, 1),
+    (:Base, :asech, 1),
+    (:Base, :asin, 1),
+    (:Base, :asind, 1),
+    (:Base, :asinh, 1),
+    (:Base, :atan, 1),
+    (:Base, :atan, 2),
+    (:Base, :atand, 1),
+    (:Base, :atanh, 1),
+    (:Base, :cbrt, 1),
+    (:Base, :cos, 1),
+    (:Base, :cosd, 1),
+    (:Base, :cosh, 1),
+    (:Base, :cospi, 1),
+    (:Base, :cot, 1),
+    (:Base, :cotd, 1),
+    (:Base, :coth, 1),
+    (:Base, :csc, 1),
+    (:Base, :cscd, 1),
+    (:Base, :csch, 1),
+    (:Base, :deg2rad, 1),
+    (:Base, :exp, 1),
+    (:Base, :exp10, 1),
+    (:Base, :exp2, 1),
+    (:Base, :expm1, 1),
+    (:Base, :hypot, 2),
+    (:Base, :inv, 1),
+    (:Base, :ldexp, 2),
+    (:Base, :log, 1),
+    (:Base, :log, 2),
+    (:Base, :log10, 1),
+    (:Base, :log1p, 1),
+    (:Base, :log2, 1),
+    (:Base, :max, 2),
+    (:Base, :min, 2),
+    (:Base, :mod, 2),
+    (:Base, :mod2pi, 1),
+    (:Base, :rad2deg, 1),
+    (:Base, :rem, 2),
+    (:Base, :rem2pi, 2),
+    (:Base, :sec, 1),
+    (:Base, :secd, 1),
+    (:Base, :sech, 1),
+    (:Base, :sin, 1),
+    (:Base, :sinc, 1),
+    (:Base, :sind, 1),
+    (:Base, :sinh, 1),
+    (:Base, :sinpi, 1),
+    (:Base, :sqrt, 1),
+    (:Base, :tan, 1),
+    (:Base, :tand, 1),
+    (:Base, :tanh, 1),
+    (:LogExpFunctions, :log1mexp, 1),
+    (:LogExpFunctions, :log1pexp, 1),
+    (:LogExpFunctions, :log1pmx, 1),
+    (:LogExpFunctions, :log1psq, 1),
+    (:LogExpFunctions, :log2mexp, 1),
+    (:LogExpFunctions, :logaddexp, 2),
+    (:LogExpFunctions, :logexpm1, 1),
+    (:LogExpFunctions, :logistic, 1),
+    (:LogExpFunctions, :logit, 1),
+    (:LogExpFunctions, :logmxp1, 1),
+    (:LogExpFunctions, :logsubexp, 2),
+    (:LogExpFunctions, :xlog1py, 2),
+    (:LogExpFunctions, :xlogx, 1),
+    (:LogExpFunctions, :xlogy, 2),
+    (:NaNMath, :acos, 1),
+    (:NaNMath, :acosh, 1),
+    (:NaNMath, :asin, 1),
+    (:NaNMath, :atanh, 1),
+    (:NaNMath, :cos, 1),
+    (:NaNMath, :lgamma, 1),
+    (:NaNMath, :log, 1),
+    (:NaNMath, :log10, 1),
+    (:NaNMath, :log1p, 1),
+    (:NaNMath, :log2, 1),
+    (:NaNMath, :max, 2),
+    (:NaNMath, :min, 2),
+    (:NaNMath, :pow, 2),
+    (:NaNMath, :sin, 1),
+    (:NaNMath, :sqrt, 1),
+    (:NaNMath, :tan, 1),
+    (:SpecialFunctions, :airyai, 1),
+    (:SpecialFunctions, :airyaiprime, 1),
+    (:SpecialFunctions, :airyaiprimex, 1),
+    (:SpecialFunctions, :airyaix, 1),
+    (:SpecialFunctions, :airybi, 1),
+    (:SpecialFunctions, :airybiprime, 1),
+    (:SpecialFunctions, :airybiprimex, 1),
+    (:SpecialFunctions, :airybix, 1),
+    (:SpecialFunctions, :besselh, 2),
+    (:SpecialFunctions, :besselhx, 2),
+    (:SpecialFunctions, :besseli, 2),
+    (:SpecialFunctions, :besselix, 2),
+    (:SpecialFunctions, :besselj, 2),
+    (:SpecialFunctions, :besselj0, 1),
+    (:SpecialFunctions, :besselj1, 1),
+    (:SpecialFunctions, :besseljx, 2),
+    (:SpecialFunctions, :besselk, 2),
+    (:SpecialFunctions, :besselkx, 2),
+    (:SpecialFunctions, :bessely, 2),
+    (:SpecialFunctions, :bessely0, 1),
+    (:SpecialFunctions, :bessely1, 1),
+    (:SpecialFunctions, :besselyx, 2),
+    (:SpecialFunctions, :beta, 2),
+    (:SpecialFunctions, :cosint, 1),
+    (:SpecialFunctions, :dawson, 1),
+    (:SpecialFunctions, :digamma, 1),
+    (:SpecialFunctions, :ellipe, 1),
+    (:SpecialFunctions, :ellipk, 1),
+    (:SpecialFunctions, :erf, 1),
+    (:SpecialFunctions, :erf, 2),
+    (:SpecialFunctions, :erfc, 1),
+    (:SpecialFunctions, :erfcinv, 1),
+    (:SpecialFunctions, :erfcx, 1),
+    (:SpecialFunctions, :erfi, 1),
+    (:SpecialFunctions, :erfinv, 1),
+    (:SpecialFunctions, :expint, 1),
+    (:SpecialFunctions, :expint, 2),
+    (:SpecialFunctions, :gamma, 1),
+    (:SpecialFunctions, :hankelh1, 2),
+    (:SpecialFunctions, :hankelh1x, 2),
+    (:SpecialFunctions, :hankelh2, 2),
+    (:SpecialFunctions, :hankelh2x, 2),
+    (:SpecialFunctions, :invdigamma, 1),
+    (:SpecialFunctions, :logbeta, 2),
+    (:SpecialFunctions, :logerfc, 1),
+    (:SpecialFunctions, :logerfcx, 1),
+    (:SpecialFunctions, :loggamma, 1),
+    (:SpecialFunctions, :polygamma, 2),
+    (:SpecialFunctions, :sinint, 1),
+    (:SpecialFunctions, :trigamma, 1),
+    (:SpecialFunctions, :zeta, 2),
+]
diff --git a/test/precompilation.jl b/test/precompilation.jl
new file mode 100644
index 0000000..8a5966a
--- /dev/null
+++ b/test/precompilation.jl
@@ -0,0 +1,48 @@
+# Rules are methods, so they survive precompilation of the package defining them, and the
+# extensions are loaded while a package depending on SpecialFunctions precompiles. Neither
+# is observable in-process, so this builds a throwaway package and loads it from a fresh
+# Julia.
+
+@testset "precompilation" begin
+    mktempdir() do dir
+        mkpath(joinpath(dir, "Probe", "src"))
+        write(joinpath(dir, "Probe", "Project.toml"), """
+            name = "Probe"
+            uuid = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01"
+            version = "0.1.0"
+
+            [deps]
+            DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+            SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+            """)
+        write(joinpath(dir, "Probe", "src", "Probe.jl"), """
+            module Probe
+            using DiffRules, SpecialFunctions
+            cube(x) = x^3
+            DiffRules.@define_diffrule Probe.cube(x) = :(3 * \$x^2)
+            const RULES = DiffRules.diffrules(; filter_modules=nothing)
+            end
+            """)
+        write(joinpath(dir, "Project.toml"), """
+            [deps]
+            DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+            Probe = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01"
+
+            [sources]
+            Probe = { path = "Probe" }
+            DiffRules = { path = "$(escape_string(pkgdir(DiffRules)))" }
+            """)
+
+        script = """
+            using Pkg; Pkg.instantiate(; io=devnull)
+            using Probe, DiffRules
+            println((:Probe, :cube, 1) in DiffRules.diffrules(; filter_modules=nothing))
+            println(any(r -> r[1] === :SpecialFunctions, Probe.RULES))
+            """
+        out = read(`$(Base.julia_cmd()) --startup-file=no --project=$dir -e $script`, String)
+        rule_survives, ext_loaded_at_precompile = split(strip(out), '\n')
+
+        @test rule_survives == "true"
+        @test ext_loaded_at_precompile == "true"
+    end
+end
diff --git a/test/registry.jl b/test/registry.jl
new file mode 100644
index 0000000..5c08136
--- /dev/null
+++ b/test/registry.jl
@@ -0,0 +1,61 @@
+include("baseline.jl")
+
+@testset "registry" begin
+    @testset "baseline" begin
+        rules = sort!(collect(DiffRules.diffrules(; filter_modules=nothing));
+                      by = r -> (string(r[1]), string(r[2]), r[3]))
+        @test rules == BASELINE
+        @testset "$M.$f/$n" for (M, f, n) in BASELINE
+            @test DiffRules.hasdiffrule(M, f, n)
+        end
+    end
+
+    # `parentmodule` reports `Base.Math` for these, which would break `@eval $M.$f`.
+    @testset "Base.Math functions report :Base" begin
+        @testset "$f" for f in (:sec, :sind, :hypot, :cot, :deg2rad, :mod2pi, :rem2pi)
+            @test any(r -> r[1] === :Base && r[2] === f,
+                      DiffRules.diffrules(; filter_modules=nothing))
+        end
+    end
+
+    @testset "lookup by function" begin
+        @test DiffRules.diffrule(Base.sin, :x) == :(cos(x))
+        @test DiffRules.diffrule(Base.sec, :x) == :(sec(x) * tan(x))
+        @test DiffRules.hasdiffrule(Base.sin, 1)
+        @test !DiffRules.hasdiffrule(Base.sin, 2)
+        # `log` has a rule at both arities
+        @test DiffRules.hasdiffrule(Base.log, 1) && DiffRules.hasdiffrule(Base.log, 2)
+    end
+
+    @testset "lookup by module" begin
+        @test DiffRules.diffrule(Base, :sin, :x) == :(cos(x))
+        @test DiffRules.hasdiffrule(Base, :sin, 1)
+        @test !DiffRules.hasdiffrule(Base, :sin, 2)
+    end
+
+    # Rules splice in the functions they call, so evaluating one does not require the
+    # defining package to be in scope.
+    @testset "rules evaluate without the package in scope" begin
+        mod = Module(:NoImports)
+        d = DiffRules.diffrule(:SpecialFunctions, :digamma, :x)
+        @test Base.eval(mod, :(let x = 1.5; $d end)) ≈ SpecialFunctions.trigamma(1.5)
+    end
+
+    @testset "unknown rules" begin
+        @test !DiffRules.hasdiffrule(:Base, :nonexistent, 1)
+        @test !DiffRules.hasdiffrule(:NotLoaded, :f, 1)
+        @test_throws Exception DiffRules.diffrule(:Base, :nonexistent, :x)
+    end
+
+    @testset "filter_modules" begin
+        @test all(M === :Base for (M, _, _) in DiffRules.diffrules(; filter_modules=(:Base,)))
+        @test all(M in (:Base, :SpecialFunctions, :NaNMath)
+                  for (M, _, _) in DiffRules.diffrules())
+    end
+
+    # undocumented, but overloaded by IntervalArithmetic and Zygote
+    @testset "_abs_deriv hook" begin
+        @test DiffRules._abs_deriv(-2.0) == -1.0
+        @test DiffRules.diffrule(Base.abs, :x) == :($(DiffRules._abs_deriv)(x))
+    end
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index f1fa041..7b0e9b3 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -12,6 +12,9 @@ Random.seed!(1)
 const finitediff = central_fdm(5, 1, max_range=1e-3)
 
 @testset "DiffRules" begin
+include("registry.jl")
+include("precompilation.jl")
+
 @testset "check rules" begin
 
 non_diffeable_arg_functions = [(:Base, :rem2pi, 2), (:Base, :ldexp, 2), (:Base, :ifelse, 3)]

From e0ee4dbae6bb203ea7429ea43f2290b0f1523e5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Fri, 18 Sep 2026 16:33:28 +0200
Subject: [PATCH 2/6] Document that provider rules live in extensions

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 docs/src/index.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/docs/src/index.md b/docs/src/index.md
index 4d21308..faa9b71 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -13,6 +13,20 @@ Note that DiffRules is *not* a fully-fledged symbolic differentiation tool. It i
 simple global database of common derivative definitions, and was developed with the goal of
 improving derivative coverage in downstream tools.
 
+Rules for SpecialFunctions, NaNMath and LogExpFunctions live in package extensions, so
+DiffRules itself depends on none of them. Load the package you need alongside DiffRules to
+get its rules:
+
+```julia
+using DiffRules, SpecialFunctions
+
+DiffRules.hasdiffrule(SpecialFunctions.erf, 1)  # true
+```
+
+Without `using SpecialFunctions`, that query returns `false` and `diffrules()` does not list
+its rules. Packages generating code from `diffrules()` should therefore load the packages
+whose rules they want, and skip rules for modules they do not have in scope.
+
 ```@docs
 DiffRules.@define_diffrule
 DiffRules.diffrule

From e84823b67e8112ad24888ab5295243f6f34a7074 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Sun, 20 Sep 2026 00:46:46 +0200
Subject: [PATCH 3/6] Tidy the registry lookups

`_resolve` built the full rule list on every symbol lookup and had no early
exit, which downstream packages hit once per rule while precompiling. Scan the
method table directly instead.

Make `diffrule(M::Module, ...)` throw a `KeyError` like its `Symbol`
counterpart rather than an `UndefVarError`, and drop the IrrationalConstants
imports that moved out with the SpecialFunctions rules.

The `@define_diffrule` docstring named `Base.polygamma`, which lives in
SpecialFunctions; harmless when the macro only stored symbols, wrong now that
it evaluates `M.f`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ...pFunctionsExt.jl => LogExpFunctionsExt.jl} |  0
 ext/{DiffRulesNaNMathExt.jl => NaNMathExt.jl} |  0
 ...nsExt.jl => NaNMathSpecialFunctionsExt.jl} |  0
 ...FunctionsExt.jl => SpecialFunctionsExt.jl} |  0
 src/DiffRules.jl                              |  2 +-
 src/api.jl                                    | 43 ++++++++++++++-----
 6 files changed, 33 insertions(+), 12 deletions(-)
 rename ext/{DiffRulesLogExpFunctionsExt.jl => LogExpFunctionsExt.jl} (100%)
 rename ext/{DiffRulesNaNMathExt.jl => NaNMathExt.jl} (100%)
 rename ext/{DiffRulesNaNMathSpecialFunctionsExt.jl => NaNMathSpecialFunctionsExt.jl} (100%)
 rename ext/{DiffRulesSpecialFunctionsExt.jl => SpecialFunctionsExt.jl} (100%)

diff --git a/ext/DiffRulesLogExpFunctionsExt.jl b/ext/LogExpFunctionsExt.jl
similarity index 100%
rename from ext/DiffRulesLogExpFunctionsExt.jl
rename to ext/LogExpFunctionsExt.jl
diff --git a/ext/DiffRulesNaNMathExt.jl b/ext/NaNMathExt.jl
similarity index 100%
rename from ext/DiffRulesNaNMathExt.jl
rename to ext/NaNMathExt.jl
diff --git a/ext/DiffRulesNaNMathSpecialFunctionsExt.jl b/ext/NaNMathSpecialFunctionsExt.jl
similarity index 100%
rename from ext/DiffRulesNaNMathSpecialFunctionsExt.jl
rename to ext/NaNMathSpecialFunctionsExt.jl
diff --git a/ext/DiffRulesSpecialFunctionsExt.jl b/ext/SpecialFunctionsExt.jl
similarity index 100%
rename from ext/DiffRulesSpecialFunctionsExt.jl
rename to ext/SpecialFunctionsExt.jl
diff --git a/src/DiffRules.jl b/src/DiffRules.jl
index 67c76f2..0ce877b 100644
--- a/src/DiffRules.jl
+++ b/src/DiffRules.jl
@@ -2,7 +2,7 @@ __precompile__()
 
 module DiffRules
 
-using IrrationalConstants: logtwo, logten, twoπ, sqrtπ, invsqrtπ
+using IrrationalConstants: logtwo, logten, twoπ
 
 include("api.jl")
 include("rules.jl")
diff --git a/src/api.jl b/src/api.jl
index 7c48754..46019d9 100644
--- a/src/api.jl
+++ b/src/api.jl
@@ -16,12 +16,17 @@ Note that differentiation rules are purely symbolic, so no type annotations shou
 Each rule is a method of [`diffrule`](@ref), so rules defined in other packages and in
 package extensions are precompiled and visible like any other method.
 
+A rule is evaluated by whoever asks for it, which need not have the defining package in
+scope. Interpolate the functions the RHS calls, rather than naming them, so that the
+returned expression stands on its own.
+
 # Examples
 
 ```julia
-@define_diffrule Base.cos(x)          = :(-sin(\$x))
-@define_diffrule Base.:/(x, y)        = :(inv(\$y)), :(-\$x / (\$y^2))
-@define_diffrule Base.polygamma(m, x) = :NaN,       :(polygamma(\$m + 1, \$x))
+@define_diffrule Base.cos(x)      = :(-sin(\$x))
+@define_diffrule Base.:/(x, y)    = :(inv(\$y)), :(-\$x / (\$y^2))
+@define_diffrule Base.ldexp(x, y) = :(exp2(\$y)), :NaN
+@define_diffrule MyPkg.f(x)       = :(\$(MyPkg.g)(\$x))
 ```
 """
 macro define_diffrule(def)
@@ -44,6 +49,9 @@ interpolated into the returned expression.
 In the `n`-ary case, an `n`-tuple of expressions will be returned where the `i`th expression
 is the derivative of `f` w.r.t the `i`th argument.
 
+Throw a `KeyError` if `M.f` cannot be resolved, which for a package's functions is also the
+case as long as that package is not loaded.
+
 # Examples
 
 ```jldoctest
@@ -59,7 +67,10 @@ julia> DiffRules.diffrule(:Base, :sin, :(x * y^2))
 """
 function diffrule end
 
-diffrule(M::Module, f::Symbol, args...) = diffrule(getproperty(M, f), args...)
+function diffrule(M::Module, f::Symbol, args...)
+    isdefined(M, f) || throw(KeyError((nameof(M), f, length(args))))
+    return diffrule(getproperty(M, f), args...)
+end
 
 function diffrule(M::Symbol, f::Symbol, args...)
     fn = _resolve(M, f)
@@ -168,21 +179,31 @@ _pkgname(fn) = nameof(Base.moduleroot(parentmodule(fn)))
 
 # Rules are the methods of `diffrule` whose first parameter is a singleton function type.
 # The `Symbol`/`Module` methods above are not, and are skipped.
+function _rule(m::Method)
+    m.isva && return nothing
+    params = Base.unwrap_unionall(m.sig).parameters
+    length(params) >= 2 || return nothing
+    T = params[2]
+    T isa DataType && T <: Function && isdefined(T, :instance) || return nothing
+    return (T.instance, length(params) - 2)
+end
+
 function _rules()
     rules = Tuple{Function,Int}[]
     for m in methods(diffrule)
-        m.isva && continue
-        params = Base.unwrap_unionall(m.sig).parameters
-        length(params) >= 2 || continue
-        T = params[2]
-        T isa DataType && T <: Function && isdefined(T, :instance) || continue
-        push!(rules, (T.instance, length(params) - 2))
+        rule = _rule(m)
+        rule === nothing || push!(rules, rule)
     end
     return rules
 end
 
+# Downstream packages resolve every rule by name while precompiling, so scan the method table
+# directly instead of building the full list first.
 function _resolve(M::Symbol, f::Symbol)
-    for (fn, _) in _rules()
+    for m in methods(diffrule)
+        rule = _rule(m)
+        rule === nothing && continue
+        fn = rule[1]
         nameof(fn) === f && _pkgname(fn) === M && return fn
     end
     return nothing

From d84359a00bec27e077ec92b117d36d158241547b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Sun, 20 Sep 2026 00:50:00 +0200
Subject: [PATCH 4/6] Drop the DiffRules prefix from the extension names

An extension's `PkgId` is derived from the parent package's UUID, so the prefix
bought nothing.

While here, import the functions the SpecialFunctions and LogExpFunctions rules
call instead of naming them through their module, which leaves the rules
readable. NaNMath keeps the qualified form because its names shadow `Base`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 Project.toml                      |   8 +-
 ext/LogExpFunctionsExt.jl         |  22 ++--
 ext/NaNMathExt.jl                 |  24 ++--
 ext/NaNMathSpecialFunctionsExt.jl |   9 +-
 ext/SpecialFunctionsExt.jl        | 192 ++++++++++++++++++------------
 5 files changed, 146 insertions(+), 109 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3a6f536..872978d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -11,10 +11,10 @@ NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
-DiffRulesLogExpFunctionsExt = "LogExpFunctions"
-DiffRulesNaNMathExt = "NaNMath"
-DiffRulesNaNMathSpecialFunctionsExt = ["NaNMath", "SpecialFunctions"]
-DiffRulesSpecialFunctionsExt = "SpecialFunctions"
+LogExpFunctionsExt = "LogExpFunctions"
+NaNMathExt = "NaNMath"
+NaNMathSpecialFunctionsExt = ["NaNMath", "SpecialFunctions"]
+SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
 IrrationalConstants = "0.1.1, 0.2"
diff --git a/ext/LogExpFunctionsExt.jl b/ext/LogExpFunctionsExt.jl
index a5d7b3b..71a920b 100644
--- a/ext/LogExpFunctionsExt.jl
+++ b/ext/LogExpFunctionsExt.jl
@@ -1,7 +1,8 @@
-module DiffRulesLogExpFunctionsExt
+module LogExpFunctionsExt
 
 using DiffRules: @define_diffrule
-using LogExpFunctions
+using LogExpFunctions:
+    LogExpFunctions, log1mexp, log2mexp, logaddexp, logexpm1, logistic, logsubexp
 
 ###################
 # LogExpFunctions #
@@ -9,13 +10,13 @@ using LogExpFunctions
 
 # unary
 @define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x))
-@define_diffrule LogExpFunctions.logistic(x) = :(z = $(LogExpFunctions.logistic)($x); z * (1 - z))
+@define_diffrule LogExpFunctions.logistic(x) = :(z = $logistic($x); z * (1 - z))
 @define_diffrule LogExpFunctions.logit(x) = :(inv($x * (1 - $x)))
 @define_diffrule LogExpFunctions.log1psq(x) = :(2 * $x / (1 + $x^2))
-@define_diffrule LogExpFunctions.log1pexp(x) = :($(LogExpFunctions.logistic)($x))
-@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - $(LogExpFunctions.log1mexp)($x)))
-@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - $(LogExpFunctions.log2mexp)($x)))
-@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - $(LogExpFunctions.logexpm1)($x)))
+@define_diffrule LogExpFunctions.log1pexp(x) = :($logistic($x))
+@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - $log1mexp($x)))
+@define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - $log2mexp($x)))
+@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - $logexpm1($x)))
 @define_diffrule LogExpFunctions.log1pmx(x) = :(-$x / (1 + $x))
 @define_diffrule LogExpFunctions.logmxp1(x) = :((1 - $x) / $x)
 
@@ -24,11 +25,12 @@ using LogExpFunctions
     :(log($y)),
     :(z = $x / $y; iszero($x) && !isnan($y) ? zero(z) : z)
 @define_diffrule LogExpFunctions.logaddexp(x, y) =
-    :(exp($x - $(LogExpFunctions.logaddexp)($x, $y))), :(exp($y - $(LogExpFunctions.logaddexp)($x, $y)))
+    :(exp($x - $logaddexp($x, $y))), :(exp($y - $logaddexp($x, $y)))
 @define_diffrule LogExpFunctions.logsubexp(x, y) =
-    :(z = $(LogExpFunctions.logsubexp)($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
-    :(z = $(LogExpFunctions.logsubexp)($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
+    :(z = $logsubexp($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
+    :(z = $logsubexp($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
 @define_diffrule LogExpFunctions.xlog1py(x, y) =
     :(log1p($y)),
     :(z = $x / (1 + $y); iszero($x) && !isnan($y) ? zero(z) : z)
+
 end # module
diff --git a/ext/NaNMathExt.jl b/ext/NaNMathExt.jl
index 1ac41e0..c22be76 100644
--- a/ext/NaNMathExt.jl
+++ b/ext/NaNMathExt.jl
@@ -1,8 +1,8 @@
-module DiffRulesNaNMathExt
+module NaNMathExt
 
 using DiffRules: @define_diffrule
 using IrrationalConstants: logtwo, logten
-using NaNMath
+using NaNMath: NaNMath
 
 ###########
 # NaNMath #
@@ -11,19 +11,18 @@ using NaNMath
 # unary #
 #-------#
 
-@define_diffrule NaNMath.sqrt(x)   = :(  inv(2 * $(NaNMath.sqrt)($x))                  )
-@define_diffrule NaNMath.sin(x)    = :(  $(NaNMath.cos)($x)                            )
-@define_diffrule NaNMath.cos(x)    = :( -$(NaNMath.sin)($x)                            )
+@define_diffrule NaNMath.sqrt(x)   = :(  inv(2 * $(NaNMath.sqrt)($x))                     )
+@define_diffrule NaNMath.sin(x)    = :(  $(NaNMath.cos)($x)                               )
+@define_diffrule NaNMath.cos(x)    = :( -$(NaNMath.sin)($x)                               )
 @define_diffrule NaNMath.tan(x)    = :(  1 + $(NaNMath.pow)($(NaNMath.tan)($x), 2)        )
 @define_diffrule NaNMath.asin(x)   = :(  inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2)))  )
-@define_diffrule NaNMath.acos(x)   = :(  -inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2))) )
+@define_diffrule NaNMath.acos(x)   = :( -inv($(NaNMath.sqrt)(1 - $(NaNMath.pow)($x, 2)))  )
 @define_diffrule NaNMath.acosh(x)  = :(  inv($(NaNMath.sqrt)($(NaNMath.pow)($x, 2) - 1))  )
-@define_diffrule NaNMath.atanh(x)  = :(  inv(1 - $(NaNMath.pow)($x, 2))                )
-@define_diffrule NaNMath.log(x)    = :(  inv($x)                                    )
-@define_diffrule NaNMath.log2(x)   = :(  inv($logtwo * $x)                          )
-@define_diffrule NaNMath.log10(x)  = :(  inv($logten * $x)                          )
-@define_diffrule NaNMath.log1p(x)  = :(  inv($x + 1)                                )
-
+@define_diffrule NaNMath.atanh(x)  = :(  inv(1 - $(NaNMath.pow)($x, 2))                   )
+@define_diffrule NaNMath.log(x)    = :(  inv($x)                                          )
+@define_diffrule NaNMath.log2(x)   = :(  inv($logtwo * $x)                                )
+@define_diffrule NaNMath.log10(x)  = :(  inv($logten * $x)                                )
+@define_diffrule NaNMath.log1p(x)  = :(  inv($x + 1)                                      )
 
 # binary #
 #--------#
@@ -33,4 +32,5 @@ using NaNMath
                                      :(ifelse(($y > $x) | (signbit($y) < signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($y), zero($y))))
 @define_diffrule NaNMath.min(x, y) = :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), one($x), zero($x)), ifelse(isnan($x), zero($x), one($x)))),
                                      :(ifelse(($y < $x) | (signbit($y) > signbit($x)), ifelse(isnan($y), zero($y), one($y)), ifelse(isnan($x), one($x), zero($x))))
+
 end # module
diff --git a/ext/NaNMathSpecialFunctionsExt.jl b/ext/NaNMathSpecialFunctionsExt.jl
index dcda098..46425cb 100644
--- a/ext/NaNMathSpecialFunctionsExt.jl
+++ b/ext/NaNMathSpecialFunctionsExt.jl
@@ -1,11 +1,12 @@
-module DiffRulesNaNMathSpecialFunctionsExt
+module NaNMathSpecialFunctionsExt
 
 # `NaNMath.lgamma` is the one rule whose derivative reaches into another package,
 # so it needs both loaded (JuliaDiff/DiffRules.jl#106).
 
 using DiffRules: @define_diffrule
-using NaNMath
-using SpecialFunctions
+using NaNMath: NaNMath
+using SpecialFunctions: digamma
+
+@define_diffrule NaNMath.lgamma(x) = :(  $digamma($x)  )
 
-@define_diffrule NaNMath.lgamma(x) = :(  $(SpecialFunctions.digamma)($x)               )
 end # module
diff --git a/ext/SpecialFunctionsExt.jl b/ext/SpecialFunctionsExt.jl
index f694cc0..34da243 100644
--- a/ext/SpecialFunctionsExt.jl
+++ b/ext/SpecialFunctionsExt.jl
@@ -1,13 +1,50 @@
-module DiffRulesSpecialFunctionsExt
+module SpecialFunctionsExt
 
 using DiffRules: @define_diffrule
 using IrrationalConstants: sqrtπ, invsqrtπ
-using SpecialFunctions
-
-@define_diffrule SpecialFunctions.gamma(x) =
-    :(  $(SpecialFunctions.digamma)($x) * $(SpecialFunctions.gamma)($x)  )
-@define_diffrule SpecialFunctions.loggamma(x) =
-    :(  $(SpecialFunctions.digamma)($x)  )
+using SpecialFunctions:
+    SpecialFunctions,
+    airyai,
+    airyaiprime,
+    airyaiprimex,
+    airyaix,
+    airybi,
+    airybiprime,
+    airybiprimex,
+    airybix,
+    besselh,
+    besselhx,
+    besseli,
+    besselix,
+    besselj,
+    besselj0,
+    besselj1,
+    besseljx,
+    besselk,
+    besselkx,
+    bessely,
+    bessely0,
+    bessely1,
+    besselyx,
+    beta,
+    dawson,
+    digamma,
+    ellipe,
+    ellipk,
+    erfcinv,
+    erfcx,
+    erfinv,
+    expint,
+    gamma,
+    hankelh1,
+    hankelh1x,
+    hankelh2,
+    hankelh2x,
+    invdigamma,
+    logerfc,
+    polygamma,
+    trigamma,
+    zeta
 
 ####################
 # SpecialFunctions #
@@ -16,76 +53,63 @@ using SpecialFunctions
 # unary #
 #-------#
 
+@define_diffrule SpecialFunctions.gamma(x)       = :(  $digamma($x) * $gamma($x)          )
+@define_diffrule SpecialFunctions.loggamma(x)    = :(  $digamma($x)                       )
+
 @define_diffrule SpecialFunctions.erf(x)         = :(  2 * ($invsqrtπ * exp(-$x^2))       )
-@define_diffrule SpecialFunctions.erfinv(x)      =
-    :(  ($sqrtπ * exp($(SpecialFunctions.erfinv)($x)^2)) / 2  )
+@define_diffrule SpecialFunctions.erfinv(x)      = :(  ($sqrtπ * exp($erfinv($x)^2)) / 2  )
 @define_diffrule SpecialFunctions.erfc(x)        = :( -($invsqrtπ * exp(-$x^2) * 2)       )
 @define_diffrule SpecialFunctions.logerfc(x)     =
-    :( - 2 * ($invsqrtπ * exp(- $x^2 - $(SpecialFunctions.logerfc)($x))) )
-
-@define_diffrule SpecialFunctions.erfcinv(x)     =
-    :( -($sqrtπ * exp($(SpecialFunctions.erfcinv)($x)^2)) / 2  )
-@define_diffrule SpecialFunctions.erfi(x)        = :(  $invsqrtπ * exp($x^2) * 2        )
-@define_diffrule SpecialFunctions.erfcx(x)       =
-    :(  2 * (($x * $(SpecialFunctions.erfcx)($x)) - $invsqrtπ)  )
-@define_diffrule SpecialFunctions.logerfcx(x) =
-    :(  2 * ($x - inv($(SpecialFunctions.erfcx)($x) * $sqrtπ))  )
-
-@define_diffrule SpecialFunctions.dawson(x)      =
-    :(  1 - (2 * $x * $(SpecialFunctions.dawson)($x))  )
-@define_diffrule SpecialFunctions.digamma(x) =
-    :(  $(SpecialFunctions.trigamma)($x)  )
-@define_diffrule SpecialFunctions.invdigamma(x)  =
-    :(  inv($(SpecialFunctions.trigamma)($(SpecialFunctions.invdigamma)($x)))  )
-@define_diffrule SpecialFunctions.trigamma(x)    =
-    :(  $(SpecialFunctions.polygamma)(2, $x)  )
+    :( -2 * ($invsqrtπ * exp(-$x^2 - $logerfc($x)))  )
+
+@define_diffrule SpecialFunctions.erfcinv(x)     = :( -($sqrtπ * exp($erfcinv($x)^2)) / 2 )
+@define_diffrule SpecialFunctions.erfi(x)        = :(  $invsqrtπ * exp($x^2) * 2          )
+@define_diffrule SpecialFunctions.erfcx(x)       = :(  2 * (($x * $erfcx($x)) - $invsqrtπ))
+@define_diffrule SpecialFunctions.logerfcx(x)    = :(  2 * ($x - inv($erfcx($x) * $sqrtπ)))
+
+@define_diffrule SpecialFunctions.dawson(x)      = :(  1 - (2 * $x * $dawson($x))         )
+@define_diffrule SpecialFunctions.digamma(x)     = :(  $trigamma($x)                      )
+@define_diffrule SpecialFunctions.invdigamma(x)  = :(  inv($trigamma($invdigamma($x)))    )
+@define_diffrule SpecialFunctions.trigamma(x)    = :(  $polygamma(2, $x)                  )
 
 # derivatives for `airybix` and `airybiprimex` are only correct for real inputs
 # `airyaix` and `airyaiprimex` are only defined for positive real inputs
 # `airybix` and `airybiprimex` are unscaled for negative real inputs
-@define_diffrule SpecialFunctions.airyai(x)      =
-    :(  $(SpecialFunctions.airyaiprime)($x)  )
-@define_diffrule SpecialFunctions.airyaiprime(x) =
-    :(  $x * $(SpecialFunctions.airyai)($x)  )
-@define_diffrule SpecialFunctions.airyaix(x)     =
-    :(  $(SpecialFunctions.airyaiprimex)($x) + sqrt($x) * $(SpecialFunctions.airyaix)($x)  )
+@define_diffrule SpecialFunctions.airyai(x)       = :(  $airyaiprime($x)  )
+@define_diffrule SpecialFunctions.airyaiprime(x)  = :(  $x * $airyai($x)  )
+@define_diffrule SpecialFunctions.airyaix(x)      =
+    :(  $airyaiprimex($x) + sqrt($x) * $airyaix($x)  )
 @define_diffrule SpecialFunctions.airyaiprimex(x) =
-    :(  $x * $(SpecialFunctions.airyaix)($x) + sqrt($x) * $(SpecialFunctions.airyaiprimex)($x)  )
-@define_diffrule SpecialFunctions.airybi(x)      =
-    :(  $(SpecialFunctions.airybiprime)($x)  )
-@define_diffrule SpecialFunctions.airybiprime(x) =
-    :(  $x * $(SpecialFunctions.airybi)($x)  )
-@define_diffrule SpecialFunctions.airybix(x)     =
+    :(  $x * $airyaix($x) + sqrt($x) * $airyaiprimex($x)  )
+@define_diffrule SpecialFunctions.airybi(x)       = :(  $airybiprime($x)  )
+@define_diffrule SpecialFunctions.airybiprime(x)  = :(  $x * $airybi($x)  )
+@define_diffrule SpecialFunctions.airybix(x)      =
     :(  if $x > zero($x)
-            $(SpecialFunctions.airybiprimex)($x) - sqrt($x) * $(SpecialFunctions.airybix)($x)
+            $airybiprimex($x) - sqrt($x) * $airybix($x)
         else
-            $(SpecialFunctions.airybiprimex)($x)
+            $airybiprimex($x)
         end  )
 @define_diffrule SpecialFunctions.airybiprimex(x) =
     :(  if $x > zero($x)
-            $x * $(SpecialFunctions.airybix)($x) - sqrt($x) * $(SpecialFunctions.airybiprimex)($x)
+            $x * $airybix($x) - sqrt($x) * $airybiprimex($x)
         else
-            $x * $(SpecialFunctions.airybix)($x)
+            $x * $airybix($x)
         end  )
 
-@define_diffrule SpecialFunctions.besselj0(x)    =
-    :( -$(SpecialFunctions.besselj1)($x)  )
-@define_diffrule SpecialFunctions.besselj1(x)    =
-    :(  ($(SpecialFunctions.besselj0)($x) - $(SpecialFunctions.besselj)(2, $x)) / 2  )
-@define_diffrule SpecialFunctions.bessely0(x)    =
-    :( -$(SpecialFunctions.bessely1)($x)  )
-@define_diffrule SpecialFunctions.bessely1(x)    =
-    :(  ($(SpecialFunctions.bessely0)($x) - $(SpecialFunctions.bessely)(2, $x)) / 2 )
+@define_diffrule SpecialFunctions.besselj0(x)    = :( -$besselj1($x)                         )
+@define_diffrule SpecialFunctions.besselj1(x)    = :(  ($besselj0($x) - $besselj(2, $x)) / 2 )
+@define_diffrule SpecialFunctions.bessely0(x)    = :( -$bessely1($x)                         )
+@define_diffrule SpecialFunctions.bessely1(x)    = :(  ($bessely0($x) - $bessely(2, $x)) / 2 )
 
-@define_diffrule SpecialFunctions.sinint(x) = :( sinc($x / π) )
-@define_diffrule SpecialFunctions.cosint(x) = :( cos($x) / $x )
+@define_diffrule SpecialFunctions.sinint(x)      = :(  sinc($x / π)  )
+@define_diffrule SpecialFunctions.cosint(x)      = :(  cos($x) / $x  )
 
-@define_diffrule SpecialFunctions.ellipk(m) =
-    :( ($(SpecialFunctions.ellipe)($m) / (1 - $m) - $(SpecialFunctions.ellipk)($m)) / (2 * $m) )
-@define_diffrule SpecialFunctions.ellipe(m) =
-    :( ($(SpecialFunctions.ellipe)($m) - $(SpecialFunctions.ellipk)($m)) / (2 * $m) )
+@define_diffrule SpecialFunctions.ellipk(m)      =
+    :( ($ellipe($m) / (1 - $m) - $ellipk($m)) / (2 * $m) )
+@define_diffrule SpecialFunctions.ellipe(m)      =
+    :( ($ellipe($m) - $ellipk($m)) / (2 * $m) )
 
-@define_diffrule SpecialFunctions.expint(x)    = :( -exp(-$x) / $x )
+@define_diffrule SpecialFunctions.expint(x)      = :( -exp(-$x) / $x )
 
 # TODO:
 #
@@ -106,47 +130,57 @@ using SpecialFunctions
 # for forward-mode and reverse-mode derivatives for complex inputs
 
 @define_diffrule SpecialFunctions.besselj(ν, x)   =
-    :NaN, :(  ($(SpecialFunctions.besselj)($ν - 1, $x) - $(SpecialFunctions.besselj)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($besselj($ν - 1, $x) - $besselj($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besseljx(ν, x)  =
-    :NaN, :(  ($(SpecialFunctions.besseljx)($ν - 1, $x) - $(SpecialFunctions.besseljx)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($besseljx($ν - 1, $x) - $besseljx($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besseli(ν, x)   =
-    :NaN, :(  ($(SpecialFunctions.besseli)($ν - 1, $x) + $(SpecialFunctions.besseli)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($besseli($ν - 1, $x) + $besseli($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besselix(ν, x)  =
-    :NaN, :(  ($(SpecialFunctions.besselix)($ν - 1, $x) + $(SpecialFunctions.besselix)($ν + 1, $x)) / 2 - sign($x) * $(SpecialFunctions.besselix)($ν, $x)  )
+    :NaN, :(  ($besselix($ν - 1, $x) + $besselix($ν + 1, $x)) / 2 - sign($x) * $besselix($ν, $x)  )
 @define_diffrule SpecialFunctions.bessely(ν, x)   =
-    :NaN, :(  ($(SpecialFunctions.bessely)($ν - 1, $x) - $(SpecialFunctions.bessely)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($bessely($ν - 1, $x) - $bessely($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besselyx(ν, x)  =
-    :NaN, :(  ($(SpecialFunctions.besselyx)($ν - 1, $x) - $(SpecialFunctions.besselyx)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($besselyx($ν - 1, $x) - $besselyx($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besselk(ν, x)   =
-    :NaN, :( -($(SpecialFunctions.besselk)($ν - 1, $x) + $(SpecialFunctions.besselk)($ν + 1, $x)) / 2  )
+    :NaN, :( -($besselk($ν - 1, $x) + $besselk($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.besselkx(ν, x)  =
-    :NaN, :( -($(SpecialFunctions.besselkx)($ν - 1, $x) + $(SpecialFunctions.besselkx)($ν + 1, $x)) / 2 + $(SpecialFunctions.besselkx)($ν, $x)  )
+    :NaN, :( -($besselkx($ν - 1, $x) + $besselkx($ν + 1, $x)) / 2 + $besselkx($ν, $x)  )
 @define_diffrule SpecialFunctions.besselh(ν, x)   =
-    :NaN, :(  ($(SpecialFunctions.besselh)($ν - 1, $x) - $(SpecialFunctions.besselh)($ν + 1, $x)) / 2  )
-@define_diffrule SpecialFunctions.besselhx(ν, x) =
-    :NaN, :(  ($(SpecialFunctions.besselhx)($ν - 1, $x) - $(SpecialFunctions.besselhx)($ν + 1, $x)) / 2 - im * $(SpecialFunctions.besselhx)($ν, $x)  )
+    :NaN, :(  ($besselh($ν - 1, $x) - $besselh($ν + 1, $x)) / 2  )
+@define_diffrule SpecialFunctions.besselhx(ν, x)  =
+    :NaN, :(  ($besselhx($ν - 1, $x) - $besselhx($ν + 1, $x)) / 2 - im * $besselhx($ν, $x)  )
 @define_diffrule SpecialFunctions.hankelh1(ν, x)  =
-    :NaN, :(  ($(SpecialFunctions.hankelh1)($ν - 1, $x) - $(SpecialFunctions.hankelh1)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($hankelh1($ν - 1, $x) - $hankelh1($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.hankelh1x(ν, x) =
-    :NaN, :(  ($(SpecialFunctions.hankelh1x)($ν - 1, $x) - $(SpecialFunctions.hankelh1x)($ν + 1, $x)) / 2 - im * $(SpecialFunctions.hankelh1x)($ν, $x)  )
+    :NaN, :(  ($hankelh1x($ν - 1, $x) - $hankelh1x($ν + 1, $x)) / 2 - im * $hankelh1x($ν, $x)  )
 @define_diffrule SpecialFunctions.hankelh2(ν, x)  =
-    :NaN, :(  ($(SpecialFunctions.hankelh2)($ν - 1, $x) - $(SpecialFunctions.hankelh2)($ν + 1, $x)) / 2  )
+    :NaN, :(  ($hankelh2($ν - 1, $x) - $hankelh2($ν + 1, $x)) / 2  )
 @define_diffrule SpecialFunctions.hankelh2x(ν, x) =
-    :NaN, :(  ($(SpecialFunctions.hankelh2x)($ν - 1, $x) - $(SpecialFunctions.hankelh2x)($ν + 1, $x)) / 2 + im * $(SpecialFunctions.hankelh2x)($ν, $x)  )
+    :NaN, :(  ($hankelh2x($ν - 1, $x) - $hankelh2x($ν + 1, $x)) / 2 + im * $hankelh2x($ν, $x)  )
 
 @define_diffrule SpecialFunctions.polygamma(m, x) =
-    :NaN, :(  $(SpecialFunctions.polygamma)($m + 1, $x)  )
+    :NaN, :(  $polygamma($m + 1, $x)  )
 
 @define_diffrule SpecialFunctions.beta(a, b)      =
-    :( $(SpecialFunctions.beta)($a, $b)*($(SpecialFunctions.digamma)($a) - $(SpecialFunctions.digamma)($a + $b)) ), :(  $(SpecialFunctions.beta)($a, $b)*($(SpecialFunctions.digamma)($b) - $(SpecialFunctions.digamma)($a + $b))     )
+    :(  $beta($a, $b) * ($digamma($a) - $digamma($a + $b))  ),
+    :(  $beta($a, $b) * ($digamma($b) - $digamma($a + $b))  )
 @define_diffrule SpecialFunctions.logbeta(a, b)   =
-    :( $(SpecialFunctions.digamma)($a) - $(SpecialFunctions.digamma)($a + $b)  ), :(  $(SpecialFunctions.digamma)($b) - $(SpecialFunctions.digamma)($a + $b)  )
+    :(  $digamma($a) - $digamma($a + $b)  ), :(  $digamma($b) - $digamma($a + $b)  )
 
 # derivative wrt to `ν` is not implemented
-@define_diffrule SpecialFunctions.expint(ν, x)    = 
-    :NaN, :( -$(SpecialFunctions.expint)($ν - 1, $x) )
+@define_diffrule SpecialFunctions.expint(ν, x)    =
+    :NaN, :( -$expint($ν - 1, $x) )
 
 # derivative wrt to `s` is not implemented
 @define_diffrule SpecialFunctions.zeta(s, z)      =
-    :NaN, :( - $s * $(SpecialFunctions.zeta)($s + 1, $z)  )
+    :NaN, :( -$s * $zeta($s + 1, $z) )
+
+# ternary #
+#---------#
+
+# TODO:
+#
+# besselh
+# besselhx
+
 end # module

From e5eb3666d5a20f2308800bb3700a0ba7a01c61c5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Sun, 20 Sep 2026 00:50:34 +0200
Subject: [PATCH 5/6] Isolate and extend the precompilation probes

`Pkg.test` runs the test process with `JULIA_LOAD_PATH="@:<pkgdir>/test"`, which
the spawned probe inherited. That dropped `@stdlib` off its load path, so
`using Pkg` failed and the testset errored. It also stacked our test
environment under the probe's project: a probe with an empty `Project.toml`
could still find SpecialFunctions, so anything the probe asserted about an
unloaded package was vacuous.

Pin the child to its own project plus the stdlibs, and have it print labelled
lines rather than destructuring exactly two lines of stdout.

With the environments actually isolated, probe the two properties the suite
cannot check in-process, since it loads every provider: that no provider rules
exist until the provider is loaded, and that `NaNMath.lgamma` appears only once
NaNMath and SpecialFunctions are both loaded (#106).

Also tighten `unknown rules`, which passed on any exception at all, and assert
that `diffrules()` has no duplicate keys.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 test/precompilation.jl | 132 +++++++++++++++++++++++++++++++----------
 test/registry.jl       |  15 +++--
 test/runtests.jl       |   2 +
 3 files changed, 114 insertions(+), 35 deletions(-)

diff --git a/test/precompilation.jl b/test/precompilation.jl
index 8a5966a..8f2aa18 100644
--- a/test/precompilation.jl
+++ b/test/precompilation.jl
@@ -1,21 +1,59 @@
-# Rules are methods, so they survive precompilation of the package defining them, and the
-# extensions are loaded while a package depending on SpecialFunctions precompiles. Neither
-# is observable in-process, so this builds a throwaway package and loads it from a fresh
-# Julia.
+# Which rules exist depends on what is loaded, which the test suite itself cannot vary: it loads
+# all providers. So probe throwaway environments from fresh Julia processes.
+
+const LOAD_PATH_SEP = Sys.iswindows() ? ';' : ':'
+
+# `Pkg.test` runs us with `JULIA_LOAD_PATH="@:<pkgdir>/test"`. A child inherits that, losing
+# `@stdlib` and gaining our test environment, so pin it to its own project plus the stdlibs.
+function probe(dir, script)
+    cmd = `$(Base.julia_cmd()) --startup-file=no --project=$dir -e $script`
+    return read(
+        addenv(
+            cmd,
+            "JULIA_LOAD_PATH" => "@$(LOAD_PATH_SEP)@stdlib",
+            "JULIA_PROJECT" => nothing,
+        ),
+        String,
+    )
+end
+
+const DIFFRULES_UUID = "b552c78f-8df3-52c6-915a-8e097449b14b"
+const NANMATH_UUID = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+const SPECIALFUNCTIONS_UUID = "276daf66-3868-5448-9aa4-cd146d93841b"
+const PROBE_UUID = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01"
+
+function write_env(dir, deps; probe_pkg=false)
+    mkpath(dir)
+    entries = join(("$name = \"$uuid\"" for (name, uuid) in deps), "\n")
+    sources = "DiffRules = { path = \"$(escape_string(pkgdir(DiffRules)))\" }"
+    if probe_pkg
+        entries *= "\nProbe = \"$PROBE_UUID\""
+        sources *= "\nProbe = { path = \"$(escape_string(joinpath(dir, "Probe")))\" }"
+    end
+    write(joinpath(dir, "Project.toml"), """
+        [deps]
+        $entries
+
+        [sources]
+        $sources
+        """)
+    return dir
+end
 
 @testset "precompilation" begin
-    mktempdir() do dir
-        mkpath(joinpath(dir, "Probe", "src"))
-        write(joinpath(dir, "Probe", "Project.toml"), """
+    mktempdir() do root
+        # `Probe.RULES` records what was visible at `Probe`'s own precompile time.
+        mkpath(joinpath(root, "both", "Probe", "src"))
+        write(joinpath(root, "both", "Probe", "Project.toml"), """
             name = "Probe"
-            uuid = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01"
+            uuid = "$PROBE_UUID"
             version = "0.1.0"
 
             [deps]
-            DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
-            SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+            DiffRules = "$DIFFRULES_UUID"
+            SpecialFunctions = "$SPECIALFUNCTIONS_UUID"
             """)
-        write(joinpath(dir, "Probe", "src", "Probe.jl"), """
+        write(joinpath(root, "both", "Probe", "src", "Probe.jl"), """
             module Probe
             using DiffRules, SpecialFunctions
             cube(x) = x^3
@@ -23,26 +61,60 @@
             const RULES = DiffRules.diffrules(; filter_modules=nothing)
             end
             """)
-        write(joinpath(dir, "Project.toml"), """
-            [deps]
-            DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
-            Probe = "2b9d4c69-3a3d-4f1e-9d84-9f9b5a3f0e01"
 
-            [sources]
-            Probe = { path = "Probe" }
-            DiffRules = { path = "$(escape_string(pkgdir(DiffRules)))" }
-            """)
+        both = write_env(
+            joinpath(root, "both"),
+            ("DiffRules" => DIFFRULES_UUID,
+             "NaNMath" => NANMATH_UUID,
+             "SpecialFunctions" => SPECIALFUNCTIONS_UUID);
+            probe_pkg=true,
+        )
+        bare = write_env(joinpath(root, "bare"), ("DiffRules" => DIFFRULES_UUID,))
+        nanmath = write_env(
+            joinpath(root, "nanmath"),
+            ("DiffRules" => DIFFRULES_UUID, "NaNMath" => NANMATH_UUID),
+        )
+
+        @testset "rules of a dependent package survive its precompilation" begin
+            out = probe(both, """
+                using Pkg; Pkg.instantiate(; io=devnull)
+                using Probe, DiffRules, NaNMath
+                println("rule_survives=",
+                        (:Probe, :cube, 1) in DiffRules.diffrules(; filter_modules=nothing))
+                println("rule_resolves=", DiffRules.diffrule(:Probe, :cube, :x) == :(3 * x^2))
+                println("ext_at_precompile=", any(r -> r[1] === :SpecialFunctions, Probe.RULES))
+                println("lgamma=", DiffRules.hasdiffrule(:NaNMath, :lgamma, 1))
+                """)
+            @test occursin("rule_survives=true", out)
+            @test occursin("rule_resolves=true", out)
+            @test occursin("ext_at_precompile=true", out)
+            @test occursin("lgamma=true", out)
+        end
+
+        @testset "no provider loaded, no provider rules" begin
+            out = probe(bare, """
+                using Pkg; Pkg.instantiate(; io=devnull)
+                using DiffRules
+                println("erf=", DiffRules.hasdiffrule(:SpecialFunctions, :erf, 1))
+                println("nanmath_sin=", DiffRules.hasdiffrule(:NaNMath, :sin, 1))
+                println("base_only=", all(M === :Base for (M, _, _) in
+                                          DiffRules.diffrules(; filter_modules=nothing)))
+                """)
+            @test occursin("erf=false", out)
+            @test occursin("nanmath_sin=false", out)
+            @test occursin("base_only=true", out)
+        end
 
-        script = """
-            using Pkg; Pkg.instantiate(; io=devnull)
-            using Probe, DiffRules
-            println((:Probe, :cube, 1) in DiffRules.diffrules(; filter_modules=nothing))
-            println(any(r -> r[1] === :SpecialFunctions, Probe.RULES))
-            """
-        out = read(`$(Base.julia_cmd()) --startup-file=no --project=$dir -e $script`, String)
-        rule_survives, ext_loaded_at_precompile = split(strip(out), '\n')
-
-        @test rule_survives == "true"
-        @test ext_loaded_at_precompile == "true"
+        # `NaNMath.lgamma` differentiates to `SpecialFunctions.digamma` (#106).
+        @testset "NaNMath alone does not define the lgamma rule" begin
+            out = probe(nanmath, """
+                using Pkg; Pkg.instantiate(; io=devnull)
+                using DiffRules, NaNMath
+                println("nanmath_sin=", DiffRules.hasdiffrule(:NaNMath, :sin, 1))
+                println("lgamma=", DiffRules.hasdiffrule(:NaNMath, :lgamma, 1))
+                """)
+            @test occursin("nanmath_sin=true", out)
+            @test occursin("lgamma=false", out)
+        end
     end
 end
diff --git a/test/registry.jl b/test/registry.jl
index 5c08136..18a8a31 100644
--- a/test/registry.jl
+++ b/test/registry.jl
@@ -1,9 +1,8 @@
-include("baseline.jl")
-
 @testset "registry" begin
     @testset "baseline" begin
-        rules = sort!(collect(DiffRules.diffrules(; filter_modules=nothing));
-                      by = r -> (string(r[1]), string(r[2]), r[3]))
+        rules = DiffRules.diffrules(; filter_modules=nothing)
+        @test allunique(rules)
+        sort!(rules; by = r -> (string(r[1]), string(r[2]), r[3]))
         @test rules == BASELINE
         @testset "$M.$f/$n" for (M, f, n) in BASELINE
             @test DiffRules.hasdiffrule(M, f, n)
@@ -29,6 +28,8 @@ include("baseline.jl")
 
     @testset "lookup by module" begin
         @test DiffRules.diffrule(Base, :sin, :x) == :(cos(x))
+        @test DiffRules.diffrule(SpecialFunctions, :digamma, :x) ==
+              :($(SpecialFunctions.trigamma)(x))
         @test DiffRules.hasdiffrule(Base, :sin, 1)
         @test !DiffRules.hasdiffrule(Base, :sin, 2)
     end
@@ -44,7 +45,11 @@ include("baseline.jl")
     @testset "unknown rules" begin
         @test !DiffRules.hasdiffrule(:Base, :nonexistent, 1)
         @test !DiffRules.hasdiffrule(:NotLoaded, :f, 1)
-        @test_throws Exception DiffRules.diffrule(:Base, :nonexistent, :x)
+        @test !DiffRules.hasdiffrule(Base, :nonexistent, 1)
+        @test !DiffRules.hasdiffrule(:Base, :sum, 1)
+        @test_throws KeyError DiffRules.diffrule(:Base, :nonexistent, :x)
+        @test_throws KeyError DiffRules.diffrule(Base, :nonexistent, :x)
+        @test_throws MethodError DiffRules.diffrule(Base.sin, :x, :y)
     end
 
     @testset "filter_modules" begin
diff --git a/test/runtests.jl b/test/runtests.jl
index 7b0e9b3..d493b57 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,6 +11,8 @@ Random.seed!(1)
 # Set `max_range` to avoid domain errors.
 const finitediff = central_fdm(5, 1, max_range=1e-3)
 
+include("baseline.jl")
+
 @testset "DiffRules" begin
 include("registry.jl")
 include("precompilation.jl")

From abff62c677e208614968b0ab10dc3cd8fe76e0f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Sun, 20 Sep 2026 00:50:49 +0200
Subject: [PATCH 6/6] Test against the lower compat bounds

The rules now evaluate the functions they are defined for, so a provider
version that predates one of them makes the whole extension fail to precompile
instead of registering a key that never resolved. `version: 'min'` pins Julia,
not the packages, so nothing exercised the bounds.

They hold today: at SpecialFunctions 1.1.0, NaNMath 0.3.2 and LogExpFunctions
0.3.2 all four extensions precompile and all 153 rules are present. This guards
the next rule added for a function newer than the bounds.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .github/workflows/ci.yml | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f4887f9..fa06086 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -42,3 +42,19 @@ jobs:
         continue-on-error: true
         with:
           files: lcov.info
+
+  downgrade:
+    name: Downgrade
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: julia-actions/setup-julia@v3
+        with:
+          version: '1'
+      - uses: julia-actions/cache@v3
+      - uses: julia-actions/julia-downgrade-compat@v2
+        with:
+          projects: ., test
+          skip: Random,Test
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1