inference: Add a hook for users to be able to specify custom recursion relations

When infering recursive functions, we try to detect cases where
the recursion can be shown to be terminating and in such cases
allow significantly more inference among the recursive functions
than in cases where we cannot prove termination. The rationale for
this is to allow inference of functions that are recursive over
structures, while not spending useless (or infinite) compile time
chasing down recursions that build up infinitely large types (those
get expensive *fast*). Our built-in recursion relation here allows
recursion to proceed if argument types are syntactic subsets, if
they are decreasing integers, tuples of decreasing length, and
a few other cases. It is worth noting that similar considerations
come in in significantly more static system, where non-termination
is (statically) disallowed. E.g. in Coq, for non-syntactic recursion,
a proof needs to be provided that a recursive function does indeed
terminate, before one is allowed to define it [1]. More modern languages
like Dafny allow the user to specify a predicate over the incoming
values that is shown to decrease over subsequent function calls [2].

My motivation comes from Diffractor, where we get call chains like:

∂⃖{1}(sin'', ...) -> ∂⃖{2}(sin', ...) -> ∂⃖{3}(sin, ...) -> ∂⃖{2}(rrule, sin, ...) ->
∂⃖{1}(rrule, sin)

In this example, the first two calls are both the same method, as are the
last two. Unfortunately, particularly for the first two calls, there isn't
really a good way to express the recursion rule here in a way that is
generic.

Thus, to address these cases, this PR adds a per-method field, similar
to a generator that allows packages to provide arbitrary recursion
relations that take advantage of the (known) special semantics of those
methods to expand the allowed set of recursions.

Originally I had hoped to use this hook in place of the existing
`type_more_complex` check. However, our code currently requires
transitivity of the `type_more_complex` check for soundness of
the termination analysis.

This runs into problems in the specified use case, because we may
have interleaved chains of calls, that are both the same method,
but are not actually part of a cycle as such because their ultimate
underlying methods are different (in particular this happens when
chaining two-Cassette like generated functions). We do not
currently express enough about the semantics of these Cassette-like
methods in order for inference to reasonably compute whether
two instances are part of the same cycle or not (we have
`method_for_inference_heuristics` of course, which takes
care of one level of this, but does not take care of the nested
case).

By having this hook, but not requring transitivity, it is legal
for the hook to compute whether the ultimate underlying method
is the same (by using its knowledge of what the methods actually
do) and answering accordingly. In the long run, I would like to
bring these Cassette-like capabilities more closely into the
compiler, at which point inference itself may have enough
information to compute the cycles and we'd be able to get away
with requiring transitivity.

All that said, this mechanism is quite simple and achieves its
goal. I don't think it is particularly pretty and should
definitely be considered unstable. I'm not providing any
user-facing APIs for this, so those in the know will have to
manually poke the methods. I do think a more general
language-level framework for proving termination could be
useful, particularly as part of more rigurous definitions
of when various constant propagation happens, but this is
not that, yet.

I've tried this in Diffractor and with appropriate definitions
of the recursion relation for the relevant functions, Diffractor
becomes nicely inferable:

```
julia> using Diffractor: var"'", ∂⃖

julia> Base.return_types(sin''', Tuple{Float64})
1-element Vector{Any}:
 Float64

julia> Base.return_types(sin'''', Tuple{Float64})
1-element Vector{Any}:
 Float64

julia> Base.return_types(sin''''', Tuple{Float64})
1-element Vector{Any}:
 Float64

julia> Base.return_types(sin'''''', Tuple{Float64})
1-element Vector{Any}:
 Float64
```

Diffractor's Phase 1 design goal was to infer fine at 3rd and
4th order - which this meets. The fact that it also infers
at higher orders is nice, but inference times also increase
to impractical levels for real-world functions, e.g. 5th
order above takes a few seconds to infer just `sin` and
6th order takes about 20s or so. Of course that is still
better than Zyogte, even at second order, but fixing this
properly will be part of Diffractor Phase 2.

With this (plus some additional tweaks to constprop
heuristics for OpaqueClosure that I'll be putting up
separately), we do also generate very nice code:

```
julia> @code_typed sin'''(1.0)
CodeInfo(
1 ─ %1  = invoke ChainRules.sincos(_2::Float64)::Tuple{Float64, Float64}
│   %2  = Base.getfield(%1, 1)::Float64
│   %3  = Base.getfield(%1, 2)::Float64
│   %4  = Diffractor.getfield(%1, 1)::Float64
│   %5  = Diffractor.getfield(%1, 2)::Float64
│   %6  = Diffractor.getfield(%1, 2)::Float64
│   %7  = Base.mul_float(%6, 1.0)::Float64
│   %8  = Base.mul_float(0.0, %7)::Float64
│   %9  = Base.neg_float(%4)::Float64
│   %10 = Base.mul_float(%9, 1.0)::Float64
│   %11 = Base.mul_float(1.0, %8)::Float64
│   %12 = Base.mul_float(%5, 1.0)::Float64
│   %13 = Base.mul_float(%12, %7)::Float64
│   %14 = Base.mul_float(0.0, %12)::Float64
│   %15 = Base.mul_float(0.0, %13)::Float64
│   %16 = Base.mul_float(%14, 1.0)::Float64
│   %17 = Base.mul_float(%6, %14)::Float64
│   %18 = Base.mul_float(%10, 1.0)::Float64
│   %19 = Base.mul_float(1.0, %10)::Float64
│   %20 = Base.add_float(%17, %18)::Float64
│   %21 = Base.mul_float(0.0, %20)::Float64
│   %22 = Base.mul_float(%21, 1.0)::Float64
│   %23 = Base.mul_float(%6, %21)::Float64
│   %24 = Base.add_float(%22, %16)::Float64
│   %25 = Base.add_float(%23, %19)::Float64
│   %26 = Base.mul_float(0.0, %25)::Float64
│   %27 = Base.add_float(%15, %26)::Float64
│   %28 = Base.add_float(%24, %11)::Float64
│   %29 = Base.add_float(%27, -1.0)::Float64
│   %30 = Base.neg_float(%2)::Float64
│   %31 = Base.mul_float(%3, %29)::Float64
│   %32 = Base.muladd_float(%30, %28, %31)::Float64
└──       return %32
) => Float64
```

[1] http://adam.chlipala.net/cpdt/html/GeneralRec.html
[2] https://rise4fun.com/Dafny/tutorial/Termination

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -337,6 +337,41 @@ function abstract_call_method(interp::AbstractInterpreter, method::Method, @nosp
         return parent.linfo.def === sv.linfo.def && sv_method2 === parent_method2
     end
 
+    function edge_matches_sv(frame::InferenceState)
+        inf_method2 = frame.src.method_for_inference_limit_heuristics # limit only if user token match
+        inf_method2 isa Method || (inf_method2 = nothing) # Union{Method, Nothing}
+        if callee_method2 !== inf_method2
+            return false
+        end
+        if !hardlimit
+            # if this is a soft limit,
+            # also inspect the parent of this edge,
+            # to see if they are the same Method as sv
+            # in which case we'll need to ensure it is convergent
+            # otherwise, we don't
+
+            # check in the cycle list first
+            # all items in here are mutual parents of all others
+            if !_any(matches_sv, frame.callers_in_cycle)
+                let parent = frame.parent
+                    parent !== nothing || return false
+                    parent = parent::InferenceState
+                    (parent.cached || parent.parent !== nothing) || return false
+                    matches_sv(parent) || return false
+                end
+            end
+
+            # If the method defines a recursion relation, give it a chance
+            # to tell us that this recursion is actually ok.
+            if isdefined(method, :recursion_relation)
+                if Core._apply_pure(method.recursion_relation, Any[method, callee_method2, sig, frame.linfo.specTypes])
+                    return false
+                end
+            end
+        end
+        return true
+    end
+
     for infstate in InfStackUnwind(sv)
         if method === infstate.linfo.def
             if infstate.linfo.specTypes == sig
@@ -355,40 +390,19 @@ function abstract_call_method(interp::AbstractInterpreter, method::Method, @nosp
                 break
             end
             topmost === nothing || continue
-            inf_method2 = infstate.src.method_for_inference_limit_heuristics # limit only if user token match
-            inf_method2 isa Method || (inf_method2 = nothing) # Union{Method, Nothing}
-            if callee_method2 === inf_method2
-                if !hardlimit
-                    # if this is a soft limit,
-                    # also inspect the parent of this edge,
-                    # to see if they are the same Method as sv
-                    # in which case we'll need to ensure it is convergent
-                    # otherwise, we don't
-
-                    # check in the cycle list first
-                    # all items in here are mutual parents of all others
-                    if !_any(matches_sv, infstate.callers_in_cycle)
-                        let parent = infstate.parent
-                            parent !== nothing || continue
-                            parent = parent::InferenceState
-                            (parent.cached || parent.parent !== nothing) || continue
-                            matches_sv(parent) || continue
-                        end
-                    end
-                end
-
+            if edge_matches_sv(infstate)
                 topmost = infstate
                 edgecycle = true
             end
         end
     end
 
-    if !(topmost === nothing)
-        topmost = topmost::InferenceState
+    if topmost !== nothing
         sigtuple = unwrap_unionall(sig)::DataType
         msig = unwrap_unionall(method.sig)::DataType
         spec_len = length(msig.parameters) + 1
         ls = length(sigtuple.parameters)
+
         if method === sv.linfo.def
             # Under direct self-recursion, permit much greater use of reducers.
             # here we assume that complexity(specTypes) :>= complexity(sig)
@@ -398,6 +412,13 @@ function abstract_call_method(interp::AbstractInterpreter, method::Method, @nosp
         else
             comparison = method.sig
         end
+
+        if isdefined(method, :recursion_relation)
+            # We don't recquire the recursion_relation to be transitive, so
+            # apply a hard limit
+            hardlimit = true
+        end
+
         # see if the type is actually too big (relative to the caller), and limit it if required
         newsig = limit_type_size(sig, comparison, hardlimit ? comparison : sv.linfo.specTypes, InferenceParams(interp).TUPLE_COMPLEXITY_LIMIT_DEPTH, spec_len)
 

src/dump.c

@@ -660,6 +660,7 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_li
         jl_serialize_value(s, (jl_value_t*)m->unspecialized);
         jl_serialize_value(s, (jl_value_t*)m->generator);
         jl_serialize_value(s, (jl_value_t*)m->invokes);
+        jl_serialize_value(s, (jl_value_t*)m->recursion_relation);
     }
     else if (jl_is_method_instance(v)) {
         jl_method_instance_t *mi = (jl_method_instance_t*)v;
@@ -1503,6 +1504,9 @@ static jl_value_t *jl_deserialize_value_method(jl_serializer_state *s, jl_value_
         jl_gc_wb(m, m->generator);
     m->invokes = jl_deserialize_value(s, (jl_value_t**)&m->invokes);
     jl_gc_wb(m, m->invokes);
+    m->recursion_relation = jl_deserialize_value(s, (jl_value_t**)&m->recursion_relation);
+    if (m->recursion_relation)
+        jl_gc_wb(m, m->recursion_relation);
     JL_MUTEX_INIT(&m->writelock);
     return (jl_value_t*)m;
 }

src/jltypes.c

@@ -2200,7 +2200,7 @@ void jl_init_types(void) JL_GC_DISABLED
     jl_method_type =
         jl_new_datatype(jl_symbol("Method"), core,
                         jl_any_type, jl_emptysvec,
-                        jl_perm_symsvec(24,
+                        jl_perm_symsvec(25,
                             "name",
                             "module",
                             "file",
@@ -2217,6 +2217,7 @@ void jl_init_types(void) JL_GC_DISABLED
                             "roots",
                             "ccallable",
                             "invokes",
+                            "recursion_relation",
                             "nargs",
                             "called",
                             "nospecialize",
@@ -2225,7 +2226,7 @@ void jl_init_types(void) JL_GC_DISABLED
                             "pure",
                             "is_for_opaque_closure",
                             "aggressive_constprop"),
-                        jl_svec(24,
+                        jl_svec(25,
                             jl_symbol_type,
                             jl_module_type,
                             jl_symbol_type,
@@ -2242,6 +2243,7 @@ void jl_init_types(void) JL_GC_DISABLED
                             jl_array_any_type,
                             jl_simplevector_type,
                             jl_any_type,
+                            jl_any_type,
                             jl_int32_type,
                             jl_int32_type,
                             jl_int32_type,

src/julia.h

@@ -322,6 +322,12 @@ typedef struct _jl_method_t {
     // the most specific for the argument types.
     jl_typemap_t *invokes;
 
+    // A function that compares two specializations of this method, returning
+    // `true` if the first signature is to be considered "smaller" than the
+    // second for purposes of recursion analysis. Set to NULL to use
+    // the default recusion relation.
+    jl_value_t *recursion_relation;
+
     int32_t nargs;
     int32_t called;        // bit flags: whether each of the first 8 arguments is called
     int32_t nospecialize;  // bit flags: which arguments should not be specialized

src/method.c

@@ -662,6 +662,7 @@ JL_DLLEXPORT jl_method_t *jl_new_method_uninit(jl_module_t *module)
     m->nospecialize = module->nospecialize;
     m->nkw = 0;
     m->invokes = NULL;
+    m->recursion_relation = NULL;
     m->isva = 0;
     m->nargs = 0;
     m->primary_world = 1;

stdlib/Serialization/src/Serialization.jl

@@ -79,7 +79,7 @@ const TAGS = Any[
 
 @assert length(TAGS) == 255
 
-const ser_version = 14 # do not make changes without bumping the version #!
+const ser_version = 15 # do not make changes without bumping the version #!
 
 format_version(::AbstractSerializer) = ser_version
 format_version(s::Serializer) = s.version
@@ -429,6 +429,11 @@ function serialize(s::AbstractSerializer, meth::Method)
     else
         serialize(s, nothing)
     end
+    if isdefined(meth, :recursion_relation)
+        serialize(s, method.recursion_relation)
+    else
+        serialize(s, nothing)
+    end
     nothing
 end
 
@@ -1007,6 +1012,10 @@ function deserialize(s::AbstractSerializer, ::Type{Method})
         template = template_or_is_opaque
     end
     generator = deserialize(s)
+    recursion_relation = nothing
+    if format_version(s) >= 15
+        recursion_relation = deserialize(s)
+    end
     if makenew
         meth.module = mod
         meth.name = name
@@ -1033,6 +1042,9 @@ function deserialize(s::AbstractSerializer, ::Type{Method})
             linfo.def = meth
             meth.generator = linfo
         end
+        if recursion_relation !== nothing
+            meth.recursion_relation = recursion_relation
+        end
         if !is_for_opaque_closure
             mt = ccall(:jl_method_table_for, Any, (Any,), sig)
             if mt !== nothing && nothing === ccall(:jl_methtable_lookup, Any, (Any, Any, UInt), mt, sig, typemax(UInt))