Definition of `merge(::SortedMultiDict, args...)` causes downstream JET errors

[devmotion]

In 7d2e3c6 (release in 0.19.0) the definition of merge for SortedMultiDict was generalized to merge(::SortedMultiDict, args...) with arbitrary argument types. This causes JET test failures in AdvancedHMC as the return type of merge cannot be inferred as NamedTuple anymore: 7d2e3c6

I'm not completely sure but I tend to think that this is a bug in DataStructures. Maybe similar to other existing methods for AbstractDict in Base it should be restricted to merge(::SortedMultiDict, ::AbstractDict...)?

[StephenVavasis]

The proposed solution of restricting merge(::SortedMultiDict,...) to merge(::SortedMultiDict, ::AbstractDict...) is not appropriate because a SortedMultiDict is not an AbstractDict (since a single key is associated with multiple values in a SortedMultiDict, which is not allowed for an AbstractDict). Can you please elaborate on the downstream type-stability issue? I am not familiar with AdvancedHMC.

[devmotion]

JET (and presumably the compiler) can't infer the return type of simple merge calls anymore due to this function definition. The downstream issues only show up with DataStructures >= 0.19.0 and the inferred return type is Union{NamedTuple,SortedMultiDict}, so 7d2e3c6 seems to be problematic.

[devmotion]

SortedMultiDict is not an AbstractDict (since a single key is associated with multiple values in a SortedMultiDict, which is not allowed for an AbstractDict)

I'm not familiar with SortedMultiDict. Maybe then the method should be restricted to merge(::SortedMultiDict, ::SortedMultiDict...)? In any case, the current definition seems to be problematic.

[StephenVavasis]

Thanks for the additional information. However, I am still puzzled by the problem. If a code invokes merge(s,t) and s is a SortedMultiDict, then the code in sorted_multi_dict.jl is invoked. If s is not a SortedMultiDict, then this code cannot be invoked, so I am not seeing how a type-instability could result unless s itself is already type-unstable.

As for locking down the type of the second argument, ideally the type would be AbstractMultiDict, which is any iterable of pairs of the form (key,value), As far as I know, Base unfortunately does not define such a type.

[devmotion]

It turns out that I'm able to fix this issue downstream: TuringLang/AdvancedHMC.jl#468

The package-wide analysis of JET is solely based on the function definitions, and for untyped arguments this means JET only reasons with Any. Prior to DataStructures 0.19 and its merge definitions, this was sufficient to infer in AdvancedHMC that (iterations = i, ..., tstat..., mass_matrix=metric) is a NamedTuple; with DataStructures >= 0.19, however, it could in principle also be a DataStructures.SortedMultiDict.

So while it doesn't seem to be a bug in DataStructures, it seems somewhat surprising that JET/the compiler was able to infer even with Any arguments that this expression is a NamedTuple prior to DataStructures 0.19 but is not able to infer it anymore with this recent change.