Auto merge of #118024 - notriddle:notriddle/search-speed, r=GuillaumeGomez

rustdoc-search: optimize unifyFunctionTypes

Final profile output:
https://notriddle.com/rustdoc-html-demo-5/profile-4/index.html

This PR contains three commits that improve performance of this hot inner loop: reduces the number of allocations, a fast path for the 1-element basic query case, and reconstructing the multi-element query case to use recursion instead of an explicit `backtracking` array. It also adds new test cases that I found while working on this.

r? `@GuillaumeGomez`

3 files changed, 188 insertions, 150 deletions

diff --git a/src/librustdoc/html/static/js/search.js b/src/librustdoc/html/static/js/search.js
index a3606f4302b..e9dd3c439b3 100644
--- a/src/librustdoc/html/static/js/search.js
+++ b/src/librustdoc/html/static/js/search.js
@@ -1318,7 +1318,7 @@ function initSearch(rawSearchIndex) {
          * then this function will try with a different solution, or bail with false if it
          * runs out of candidates.
          *
-         * @param {Array<FunctionType>} fnTypes - The objects to check.
+         * @param {Array<FunctionType>} fnTypesIn - The objects to check.
          * @param {Array<QueryElement>} queryElems - The elements from the parsed query.
          * @param {[FunctionType]} whereClause - Trait bounds for generic items.
          * @param {Map<number,number>|null} mgensIn
@@ -1329,9 +1329,9 @@ function initSearch(rawSearchIndex) {
          */
         function unifyFunctionTypes(fnTypesIn, queryElems, whereClause, mgensIn, solutionCb) {
             /**
-             * @type Map<integer, integer>
+             * @type Map<integer, integer>|null
              */
-            let mgens = new Map(mgensIn);
+            const mgens = mgensIn === null ? null : new Map(mgensIn);
             if (queryElems.length === 0) {
                 return !solutionCb || solutionCb(mgens);
             }
@@ -1339,169 +1339,170 @@ function initSearch(rawSearchIndex) {
                 return false;
             }
             const ql = queryElems.length;
-            let fl = fnTypesIn.length;
+            const fl = fnTypesIn.length;
+
+            // One element fast path / base case
+            if (ql === 1 && queryElems[0].generics.length === 0) {
+                const queryElem = queryElems[0];
+                for (const fnType of fnTypesIn) {
+                    if (!unifyFunctionTypeIsMatchCandidate(fnType, queryElem, whereClause, mgens)) {
+                        continue;
+                    }
+                    if (fnType.id < 0 && queryElem.id < 0) {
+                        if (mgens && mgens.has(fnType.id) &&
+                            mgens.get(fnType.id) !== queryElem.id) {
+                            continue;
+                        }
+                        const mgensScratch = new Map(mgens);
+                        mgensScratch.set(fnType.id, queryElem.id);
+                        if (!solutionCb || solutionCb(mgensScratch)) {
+                            return true;
+                        }
+                    } else if (!solutionCb || solutionCb(mgens ? new Map(mgens) : null)) {
+                        // unifyFunctionTypeIsMatchCandidate already checks that ids match
+                        return true;
+                    }
+                }
+                for (const fnType of fnTypesIn) {
+                    if (!unifyFunctionTypeIsUnboxCandidate(fnType, queryElem, whereClause, mgens)) {
+                        continue;
+                    }
+                    if (fnType.id < 0) {
+                        if (mgens && mgens.has(fnType.id) &&
+                            mgens.get(fnType.id) !== 0) {
+                            continue;
+                        }
+                        const mgensScratch = new Map(mgens);
+                        mgensScratch.set(fnType.id, 0);
+                        if (unifyFunctionTypes(
+                            whereClause[(-fnType.id) - 1],
+                            queryElems,
+                            whereClause,
+                            mgensScratch,
+                            solutionCb
+                        )) {
+                            return true;
+                        }
+                    } else if (unifyFunctionTypes(
+                        fnType.generics,
+                        queryElems,
+                        whereClause,
+                        mgens ? new Map(mgens) : null,
+                        solutionCb
+                    )) {
+                        return true;
+                    }
+                }
+                return false;
+            }
+
+            // Multiple element recursive case
             /**
              * @type Array<FunctionType>
              */
-            let fnTypes = fnTypesIn.slice();
+            const fnTypes = fnTypesIn.slice();
             /**
-             * loop works by building up a solution set in the working arrays
+             * Algorithm works by building up a solution set in the working arrays
              * fnTypes gets mutated in place to make this work, while queryElems
-             * is left alone
+             * is left alone.
              *
-             *                                  vvvvvvv `i` points here
-             * queryElems = [ good, good, good, unknown, unknown ],
-             * fnTypes    = [ good, good, good, unknown, unknown ],
-             *                ----------------  ^^^^^^^^^^^^^^^^ `j` iterates after `i`,
-             *                |                                   looking for candidates
-             *                everything before `i` is the
-             *                current working solution
+             * It works backwards, because arrays can be cheaply truncated that way.
+             *
+             *                         vvvvvvv `queryElem`
+             * queryElems = [ unknown, unknown, good, good, good ]
+             * fnTypes    = [ unknown, unknown, good, good, good ]
+             *                ^^^^^^^^^^^^^^^^ loop over these elements to find candidates
              *
              * Everything in the current working solution is known to be a good
              * match, but it might not be the match we wind up going with, because
              * there might be more than one candidate match, and we need to try them all
              * before giving up. So, to handle this, it backtracks on failure.
-             *
-             * @type Array<{
-             *     "fnTypesScratch": Array<FunctionType>,
-             *     "queryElemsOffset": integer,
-             *     "fnTypesOffset": integer
-             * }>
              */
-            const backtracking = [];
-            let i = 0;
-            let j = 0;
-            const backtrack = () => {
-                while (backtracking.length !== 0) {
-                    // this session failed, but there are other possible solutions
-                    // to backtrack, reset to (a copy of) the old array, do the swap or unboxing
-                    const {
-                        fnTypesScratch,
-                        mgensScratch,
-                        queryElemsOffset,
-                        fnTypesOffset,
-                        unbox,
-                    } = backtracking.pop();
-                    mgens = new Map(mgensScratch);
-                    const fnType = fnTypesScratch[fnTypesOffset];
-                    const queryElem = queryElems[queryElemsOffset];
-                    if (unbox) {
-                        if (fnType.id < 0) {
-                            if (mgens.has(fnType.id) && mgens.get(fnType.id) !== 0) {
-                                continue;
-                            }
-                            mgens.set(fnType.id, 0);
-                        }
-                        const generics = fnType.id < 0 ?
-                            whereClause[(-fnType.id) - 1] :
-                            fnType.generics;
-                        fnTypes = fnTypesScratch.toSpliced(fnTypesOffset, 1, ...generics);
-                        fl = fnTypes.length;
-                        // re-run the matching algorithm on this item
-                        i = queryElemsOffset - 1;
-                    } else {
-                        if (fnType.id < 0) {
-                            if (mgens.has(fnType.id) && mgens.get(fnType.id) !== queryElem.id) {
-                                continue;
-                            }
-                            mgens.set(fnType.id, queryElem.id);
-                        }
-                        fnTypes = fnTypesScratch.slice();
-                        fl = fnTypes.length;
-                        const tmp = fnTypes[queryElemsOffset];
-                        fnTypes[queryElemsOffset] = fnTypes[fnTypesOffset];
-                        fnTypes[fnTypesOffset] = tmp;
-                        // this is known as a good match; go to the next one
-                        i = queryElemsOffset;
-                    }
-                    return true;
+            const flast = fl - 1;
+            const qlast = ql - 1;
+            const queryElem = queryElems[qlast];
+            let queryElemsTmp = null;
+            for (let i = flast; i >= 0; i -= 1) {
+                const fnType = fnTypes[i];
+                if (!unifyFunctionTypeIsMatchCandidate(fnType, queryElem, whereClause, mgens)) {
+                    continue;
                 }
-                return false;
-            };
-            for (i = 0; i !== ql; ++i) {
-                const queryElem = queryElems[i];
-                /**
-                 * list of potential function types that go with the current query element.
-                 * @type Array<integer>
-                 */
-                const matchCandidates = [];
-                let fnTypesScratch = null;
-                let mgensScratch = null;
-                // don't try anything before `i`, because they've already been
-                // paired off with the other query elements
-                for (j = i; j !== fl; ++j) {
-                    const fnType = fnTypes[j];
-                    if (unifyFunctionTypeIsMatchCandidate(fnType, queryElem, whereClause, mgens)) {
-                        if (!fnTypesScratch) {
-                            fnTypesScratch = fnTypes.slice();
+                let mgensScratch;
+                if (fnType.id < 0) {
+                    mgensScratch = new Map(mgens);
+                    if (mgensScratch.has(fnType.id)
+                        && mgensScratch.get(fnType.id) !== queryElem.id) {
+                        continue;
+                    }
+                    mgensScratch.set(fnType.id, queryElem.id);
+                } else {
+                    mgensScratch = mgens;
+                }
+                // fnTypes[i] is a potential match
+                // fnTypes[flast] is the last item in the list
+                // swap them, and drop the potential match from the list
+                // check if the remaining function types also match
+                fnTypes[i] = fnTypes[flast];
+                fnTypes.length = flast;
+                if (!queryElemsTmp) {
+                    queryElemsTmp = queryElems.slice(0, qlast);
+                }
+                const passesUnification = unifyFunctionTypes(
+                    fnTypes,
+                    queryElemsTmp,
+                    whereClause,
+                    mgensScratch,
+                    mgensScratch => {
+                        if (fnType.generics.length === 0 && queryElem.generics.length === 0) {
+                            return !solutionCb || solutionCb(mgensScratch);
                         }
-                        unifyFunctionTypes(
+                        return unifyFunctionTypes(
                             fnType.generics,
                             queryElem.generics,
                             whereClause,
-                            mgens,
-                            mgensScratch => {
-                                matchCandidates.push({
-                                    fnTypesScratch,
-                                    mgensScratch,
-                                    queryElemsOffset: i,
-                                    fnTypesOffset: j,
-                                    unbox: false,
-                                });
-                                return false; // "reject" all candidates to gather all of them
-                            }
-                        );
-                    }
-                    if (unifyFunctionTypeIsUnboxCandidate(fnType, queryElem, whereClause, mgens)) {
-                        if (!fnTypesScratch) {
-                            fnTypesScratch = fnTypes.slice();
-                        }
-                        if (!mgensScratch) {
-                            mgensScratch = new Map(mgens);
-                        }
-                        backtracking.push({
-                            fnTypesScratch,
                             mgensScratch,
-                            queryElemsOffset: i,
-                            fnTypesOffset: j,
-                            unbox: true,
-                        });
+                            solutionCb
+                        );
                     }
+                );
+                if (passesUnification) {
+                    return true;
                 }
-                if (matchCandidates.length === 0) {
-                    if (backtrack()) {
-                        continue;
-                    } else {
-                        return false;
-                    }
+                // backtrack
+                fnTypes[flast] = fnTypes[i];
+                fnTypes[i] = fnType;
+                fnTypes.length = fl;
+            }
+            for (let i = flast; i >= 0; i -= 1) {
+                const fnType = fnTypes[i];
+                if (!unifyFunctionTypeIsUnboxCandidate(fnType, queryElem, whereClause, mgens)) {
+                    continue;
                 }
-                // use the current candidate
-                const {fnTypesOffset: candidate, mgensScratch: mgensNew} = matchCandidates.pop();
-                if (fnTypes[candidate].id < 0 && queryElems[i].id < 0) {
-                    mgens.set(fnTypes[candidate].id, queryElems[i].id);
-                }
-                for (const [fid, qid] of mgensNew) {
-                    mgens.set(fid, qid);
-                }
-                // `i` and `j` are paired off
-                // `queryElems[i]` is left in place
-                // `fnTypes[j]` is swapped with `fnTypes[i]` to pair them off
-                const tmp = fnTypes[candidate];
-                fnTypes[candidate] = fnTypes[i];
-                fnTypes[i] = tmp;
-                // write other candidates to backtracking queue
-                for (const otherCandidate of matchCandidates) {
-                    backtracking.push(otherCandidate);
-                }
-                // If we're on the last item, check the solution with the callback
-                // backtrack if the callback says its unsuitable
-                while (i === (ql - 1) && solutionCb && !solutionCb(mgens)) {
-                    if (!backtrack()) {
-                        return false;
+                let mgensScratch;
+                if (fnType.id < 0) {
+                    mgensScratch = new Map(mgens);
+                    if (mgensScratch.has(fnType.id) && mgensScratch.get(fnType.id) !== 0) {
+                        continue;
                     }
+                    mgensScratch.set(fnType.id, 0);
+                } else {
+                    mgensScratch = mgens;
+                }
+                const generics = fnType.id < 0 ?
+                    whereClause[(-fnType.id) - 1] :
+                    fnType.generics;
+                const passesUnification = unifyFunctionTypes(
+                    fnTypes.toSpliced(i, 1, ...generics),
+                    queryElems,
+                    whereClause,
+                    mgensScratch,
+                    solutionCb
+                );
+                if (passesUnification) {
+                    return true;
                 }
             }
-            return true;
+            return false;
         }
         function unifyFunctionTypeIsMatchCandidate(fnType, queryElem, whereClause, mgens) {
             // type filters look like `trait:Read` or `enum:Result`
@@ -1514,15 +1515,17 @@ function initSearch(rawSearchIndex) {
             // or, if mgens[fnType.id] = 0, then we've matched this generic with a bare trait
             // and should make that same decision everywhere it appears
             if (fnType.id < 0 && queryElem.id < 0) {
-                if (mgens.has(fnType.id) && mgens.get(fnType.id) !== queryElem.id) {
-                    return false;
-                }
-                for (const [fid, qid] of mgens.entries()) {
-                    if (fnType.id !== fid && queryElem.id === qid) {
+                if (mgens !== null) {
+                    if (mgens.has(fnType.id) && mgens.get(fnType.id) !== queryElem.id) {
                         return false;
                     }
-                    if (fnType.id === fid && queryElem.id !== qid) {
-                        return false;
+                    for (const [fid, qid] of mgens.entries()) {
+                        if (fnType.id !== fid && queryElem.id === qid) {
+                            return false;
+                        }
+                        if (fnType.id === fid && queryElem.id !== qid) {
+                            return false;
+                        }
                     }
                 }
             } else {
@@ -1575,7 +1578,7 @@ function initSearch(rawSearchIndex) {
                 }
                 // mgens[fnType.id] === 0 indicates that we committed to unboxing this generic
                 // mgens[fnType.id] === null indicates that we haven't decided yet
-                if (mgens.has(fnType.id) && mgens.get(fnType.id) !== 0) {
+                if (mgens !== null && mgens.has(fnType.id) && mgens.get(fnType.id) !== 0) {
                     return false;
                 }
                 // This is only a potential unbox if the search query appears in the where clause
diff --git a/tests/rustdoc-js/generics2.js b/tests/rustdoc-js/generics2.js
new file mode 100644
index 00000000000..f08704349a4
--- /dev/null
+++ b/tests/rustdoc-js/generics2.js
@@ -0,0 +1,22 @@
+// exact-check
+
+const EXPECTED = [
+    {
+        'query': 'outside<U>, outside<V> -> outside<W>',
+        'others': [],
+    },
+    {
+        'query': 'outside<V>, outside<U> -> outside<W>',
+        'others': [],
+    },
+    {
+        'query': 'outside<U>, outside<U> -> outside<W>',
+        'others': [],
+    },
+    {
+        'query': 'outside<U>, outside<U> -> outside<U>',
+        'others': [
+            {"path": "generics2", "name": "should_match_3"}
+        ],
+    },
+];
diff --git a/tests/rustdoc-js/generics2.rs b/tests/rustdoc-js/generics2.rs
new file mode 100644
index 00000000000..1177ade6831
--- /dev/null
+++ b/tests/rustdoc-js/generics2.rs
@@ -0,0 +1,13 @@
+pub struct Outside<T>(T);
+
+pub fn no_match<U, V>(a: Outside<U>, b: Outside<V>) -> (Outside<U>, Outside<V>) {
+    unimplemented!();
+}
+
+pub fn no_match_2<U, V>(a: Outside<V>, b: Outside<U>) -> (Outside<U>, Outside<V>) {
+    unimplemented!();
+}
+
+pub fn should_match_3<U>(a: Outside<U>, b: Outside<U>) -> (Outside<U>, Outside<U>) {
+    unimplemented!();
+}