Optimize DownValues: extract "non-patterns" from Alternatives

Question

Introduction

As we know, when we assign DownValues to a symbol, Mathematica creates a dispatch table for "non-pattern" definitions (i.e. expressions that, when used as patterns, can match only single expression), and can match such definition in constant time (independent of number of definitions).

To match definitions containing "real patterns" (in the sense of expressions that "appear to be able to" match more than one expression) Mathematica tries them one by one, so in worst case Mathematica needs time linear in number of such definitions.

Main question

How can we decompose arbitrary pattern containing Alternatives to collection of separate patterns that when used as DownValues will give exactly the same outputs (and possible side effects) as single pattern, but, if possible, will take advantage of dispatching mechanism for "non-pattern" expressions.

Context

I'm creating a user modifiable predicate: somethingQ. Users will be able to declare arbitrary patterns for which predicate will give True. Expected number of such patterns can grow to the order of $10^6$, but most of them are expected to be "non-pattern" expressions.

In the context of somethingQ it will be "natural" to declare patterns similar to f[g[a | b | c | _h], x]. I could try to document that giving separate patterns: f[g[a], x], f[g[b], x], f[g[c], x], f[g[_h], x] will give somethingQ better performance, but explaining in detail when to decompose pattern and when not to (e.g f[g[a | b], _h] should non be decomposed) is difficult and decomposing patterns by hand can be tedious. Instead of burdening users with understanding DownValues matching and manual decomposition, I'd like to automate this process.

Benchmarks

When we have $10^6$ "non-pattern" expressions "hidden" inside Alternatives we can get noticeable performance gain by assigning "non-pattern" expressions to separate DownValues and gathering "real patterns" to single down value with Alternatives.

To test it let's create $10^6$ patterns with Alternatives, first $9 \times 10^5$ patterns contains only Alternatives of non-pattern expressions rest contains also "real patterns":

ClearAll[a, b, c, g, h, x]
patts = Join[
    {a | Unique[]}, Table[Unique[] | Unique[], 9 10^5],
    {b | _g}, Table[Unique[] | Blank@Unique[], 10^5], {c | _h}
];

Assign them to DownValues of somethingQ:

somethingQ // ClearAll
DownValues@somethingQ = HoldPattern@somethingQ@# :> True & /@ patts;
somethingQ@_ = False;

and test time of calling somethingQ for arguments matching first and last non-pattern and "real pattern" giving True and pattern giving False:

somethingQ@a // RepeatedTiming   (* {2.8*10^-7, True} *)
somethingQ@c // RepeatedTiming   (* {0.14, True} *)
somethingQ@g[] // RepeatedTiming (* {0.13, True} *)
somethingQ@h[] // RepeatedTiming (* {0.14, True} *)
somethingQ@x // RepeatedTiming   (* {0.14, False} *)

Matching expressions from first down value is fast, but matching last down values is slow since Mathematica needs to test all down values one by one.

Putting all patters in one huge Alternatives makes matching last down values three times faster:

somethingQ // ClearAll
somethingQ[Join @@ patts] = True;
somethingQ@_ = False;

somethingQ@a // RepeatedTiming   (* {3.*10^-7, True} *)
somethingQ@c // RepeatedTiming   (* {0.042, True} *)
somethingQ@g[] // RepeatedTiming (* {0.037, True} *)
somethingQ@h[] // RepeatedTiming (* {0.042, True} *)
somethingQ@x // RepeatedTiming   (* {0.042, False} *)

Extracting non-patterns and putting them in separate down values while putting all "real patterns" in one Alternatives gives constant time for all expressions matched by "non-pattern" down values. Worst case time for matching "real patterns" is linear in number of only "real patterns", not in number of all patterns as previously:

somethingQ // ClearAll
DownValues@somethingQ = HoldPattern@somethingQ@# :> True & /@
    Append[#@Symbol, Alternatives @@ #@Blank]& @
    GroupBy[List @@ Join @@ patts, Head];
somethingQ@_ = False;

somethingQ@a // RepeatedTiming   (* {2.*10^-7, True} *)
somethingQ@c // RepeatedTiming   (* {1.8*10^-7, True} *)
somethingQ@g[] // RepeatedTiming (* {4.*10^-7, True} *)
somethingQ@h[] // RepeatedTiming (* {0.0055, True} *)
somethingQ@x // RepeatedTiming   (* {0.0055, False} *)

Of course here extraction of "non-pattern" expressions was trivial since all patterns had known simple structure, but I'm looking for a general way of extracting them from arbitrary patterns.

Answer 1

We can create a recursive top-down parser that, for each sub-expression, of pattern, will call itself on parts of this sub-expression, assemble results, and return pair containing, as first element, list of "non-pattern" expressions and, as second element, "real pattern", that in DownValues will be equivalent to given pattern sub-expression.

Recursion stops on expressions treated as "real patterns" by DownValues and on expressions free of _Alternatives.

When _Alternatives are encountered, we look for first alternative containing _Condition, _PatternTest, or _Pattern, matching this alternative can have side effects (binding variable to a name, done by Pattern, is here also considered a side effect). We leave all alternatives behind it alone, since extracting non-patterns from them could change behavior: f[patt /; sideEffect[] | a] := result is not equivalent to f[a] := result; f[patt /; sideEffect[]] := result. From alternatives before first with possible side-effect we extract non-pattern expressions.

Parser, with behavior sketched above, is implemented in following package.

BeginPackage@"DecomposePattern`"; Unprotect@"`*"; ClearAll@"`*"

DecomposePattern::usage = "DecomposePattern[patt] \
returns {{expr1, expr2, ...}, newPatt} expression, \
where expri are pattern-free expressions matched by patt, \
created by replacing Alternatives, present in patt, \
with their consecutive pattern-free elements. \
newPatt is pattern patt with pattern-free elements of Alternatives \
removed if possible. If expri are only expressions matched by patt, \
then newPatt is Alternatives[].";

Begin@"`Private`"; ClearAll@"`*"

$patternPatt = _Blank | _BlankNullSequence | _BlankSequence | _Condition |
    _Except | _IgnoringInactive | _KeyValuePattern | _Optional |
    _OptionsPattern | _OrderlessPatternSequence |_Pattern | _PatternSequence |
    _PatternTest | _Repeated | _RepeatedNull;

$sideEffectPatt = _Condition | _PatternTest | _Pattern;

hold // Attributes = HoldAllComplete;

parse[alt_Alternatives] := Module[{held, lastBeforeSidEffPos, parsed},
    held = hold @@ Unevaluated@alt;
    (* Since patterns in alternatives are tested in order,
       and in ...Values non-patterns are always tried before patterns,
       then f[patt /; sideEffect[] | a] := ... is not equivalent to
       f[a] := ...; f[patt /; sideEffect[]] := ...,
       so we don't extract non-patterns that are in alternatives after
       first pattern with possible side effects. *)
    lastBeforeSidEffPos = Replace[
        Position[held,
            patt_ /; Not@FreeQ[Unevaluated@patt, $sideEffectPatt],
            {1}, 1, Heads -> False
        ],
        {{} :> Length@held, {{i_}} :> i - 1}
    ];
    parsed = Alternatives @@ parse /@ Take[held, lastBeforeSidEffPos];
    {
        Union @@ parsed[[All, 1]],
        Replace[
            Flatten[
                parsed[[All, 2]] |
                Alternatives @@ hold /@ Drop[held, lastBeforeSidEffPos]
            ],
            Verbatim[Alternatives]@patt_ :> patt
        ]
    }
]
parse[patt : $patternPatt] := {{}, hold@patt}
parse[
    (expr : (h : Except[HoldPattern@Symbol@___, _Symbol] /;
        MemberQ[Attributes@h, Flat | OneIdentity])[___]) |
    (expr_ /; FreeQ[Unevaluated@expr, _Alternatives])
] :=
    If[FreeQ[Unevaluated@expr, $patternPatt],
        {{hold@expr}, Alternatives[]}
    (* else *),
        {{}, hold@expr}
    ]
parse[expr : h_[args___]] := 
    With[{parsed = parse /@ Unevaluated@{h, args}},
        If[TrueQ[Times @@ (Length[#1]&) @@@ parsed <= $maxNonPattProduct],
            {
                #1[##2] & @@@ Tuples@parsed[[All, 1]]
                ,
                Replace[
                    Position[parsed,
                        {_, Except@Verbatim@Alternatives[]},
                        {1}, 2, Heads -> False
                    ],
                    {
                        {} :> Alternatives[],
                        (* If there's only one part with pattern we can
                           replace this part with pattern-only version. *)
                        {{pos_}} :> ReplacePart[hold@expr, {1, pos - 1} ->
                            RuleCondition@parsed[[pos, 2]]],
                        _ :> hold@expr
                    }
                ]
            }
        (* else *),
            {{}, hold@expr}
        ]
    ]
parse // Attributes = HoldAllComplete;

DecomposePattern // Options = {
    "SideEffectFree" -> None,
    "MaxNonPatternProduct" -> 10^6
};
DecomposePattern[patt_, OptionsPattern[]] :=
    Block[
        {
            $sideEffectPatt = Replace[OptionValue@"SideEffectFree",
                {None -> $sideEffectPatt, x_ :> Except[x, $sideEffectPatt]}
            ],
            $maxNonPattProduct = OptionValue@"MaxNonPatternProduct"
        },
        DeleteCases[parse@patt, hold, Infinity, Heads -> True]
    ]

End[]; Protect@"`*"; EndPackage[];

Simple usage examples:

ClearAll[a, b, c, d, e, f, g, h]

f[g[a | b | c | _h], x] // DecomposePattern
(* {{f[g[a], x], f[g[b], x], f[g[c], x]}, f[g[_h], x]} *)

f[g[a | b], _h] // DecomposePattern
(* {{}, f[g[a | b], _h]} *)

f[a | b | c, d | e] // DecomposePattern
(* {{f[a, d], f[a, e], f[b, d], f[b, e], f[c, d], f[c, e]}, Alternatives[]} *)

f[a | b | c, d | _h | e] // DecomposePattern
(* {{f[a, d], f[a, e], f[b, d], f[b, e], f[c, d], f[c, e]}, f[a | b | c, _h]} *)

Complex example with held patterns:

(* Assign some values to test possible evaluation leaks. *)
c := (Print@"evaluation leaked: c"; 0)
f := (Print@"evaluation leaked: f"; 1)
a | __h | b | HoldPattern[c | f[a | _g | b, f[c | d]]] | HoldPattern[f][e] | _?f | g // DecomposePattern
(* {
    {
        a, b, HoldPattern[f][e], HoldPattern[c],
        HoldPattern[f[a, f[c]]], HoldPattern[f[a, f[d]]],
        HoldPattern[f[b, f[c]]], HoldPattern[f[b, f[d]]]
    },
    __h | HoldPattern[f[_g, f[c | d]]] | _?f | g
} *)
ClearAll[c, f]

If we are sure that certain pattern, that potentially could have side effects, is free of them, we can set "SideEffectFree" option:

DecomposePattern[a | _?f | b | _?g | c]
DecomposePattern[a | _?f | b | _?g | c, "SideEffectFree" -> Verbatim[PatternTest][_, f]]
DecomposePattern[a | _?f | b | _?g | c, "SideEffectFree" -> Verbatim[PatternTest][_, f | g]]
(* {{a}, _?f | b | _?g | c} *)
(* {{a, b}, _?f | _?g | c} *)
(* {{a, b, c}, _?f | _?g} *)