refactor: Path compression for DiscrTree

[nomeata]

Trie data structures tend to be very wasteful if one doesn’t apply path compression: If you have a sequence of nodes of degree 1, better put them in a flat array, instead of nested .nodes with singleton arrays.

If applied to the DiscrTree, this would reduce the size of the DiscrTree used by mathlib’s library_search from 217 MB to 41 MB, and the maximum depth from 2596 to 27. The maximum depths was causing stack overflow in mathlib PRs.

[ghost]

• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-24 21:01:27) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-24 21:20:48) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-25 07:23:33) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-25 08:51:18) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-25 08:59:29) View Log
• 🟡 Mathlib branch lean-pr-testing-2577 build against this PR was cancelled. (2023-09-25 09:37:59) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-09-25 09:40:14) View Log
• ✅ Mathlib branch lean-pr-testing-2577 has successfully built against this PR. (2023-09-25 11:13:24) View Log
• 🟡 Mathlib branch lean-pr-testing-2577 build this PR didn't complete normally. (2023-10-14 14:25:15) View Log
• 🟡 Mathlib branch lean-pr-testing-2577 build this PR didn't complete normally. (2023-10-16 21:26:44) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-10-16 22:14:22) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-10-17 21:15:05) View Log
• 🟡 Mathlib branch lean-pr-testing-2577 build this PR didn't complete normally. (2023-10-18 01:41:24) View Log
• 💥 Mathlib branch lean-pr-testing-2577 build failed against this PR. (2023-10-18 01:47:35) View Log
• ✅ Mathlib branch lean-pr-testing-2577 has successfully built against this PR. (2023-10-18 18:53:27) View Log

[nomeata]

Could someone !bench this?

[Kha]

!bench

[leanprover-bot]

Here are the benchmark results for commit 83e2599.
There were no significant changes against commit 2ac782c.

[nomeata]

There were no significant changes against commit 2ac782c.

That’s a bit disappointing. I’ll make mathlib compile, maybe there is something there.

[nomeata]

From 237MB to:

$ ls -sh ./build/lib/MathlibExtras/LibrarySearch.extra
58M ./build/lib/MathlibExtras/LibrarySearch.extra

That's at least pretty nice.

[nomeata]

Results are in: http://speed.lean-fro.org/mathlib4/compare/e7b27246-a3e6-496a-b552-ff4b45c7236e/to/ee9e4caf-7a54-4dc1-a6d9-19101bc7d2d9?hash1=a477df946a0648231b66e5a7f0f04839546ebf25

Mild improvements across the board, it seems?

[nomeata]

awaiting-review

Guess we have the data to weigh this idea. The code certainly gets more complex, but the size gains at least for library_search (didn't check the size of other in instances used) are significant, and a 3.5% mathlib compile time improvement is at least nice.

Jovan Gerbscheid pointed out that due to the well-nested structure of key sequences, values are only ever found at the leaves. One could therefore remove the .empty constructor and the recursive argument to .values, and make insertion panic when the key sequence is invalid, for a bit code reduction. OTOH, the current code is a bit more robust, for example when people mess with the key sequences as in leanprover-community/mathlib4#7345, so maybe better like this for now.

[nomeata]

I updated the branches, and created a new comparison report:
http://speed.lean-fro.org/mathlib4/compare/63a60bc6-4997-4dbe-a363-9ceff22def1c/to/1c8d5e96-3fad-4da2-b00c-4b1a63cdc75b

I’m not sure how to best get a good signal from the report.
This time build wall-clock has improved by 0.9% (probably noise). Instruction counts go down through the bank, but less than 1%. All not very exiciting, not sure if this is worth pursuing, it seems only the size improvement to the library_search and rw? cache still look good.

[kim-em]

Oh, I think you're underselling. A 1% improvement on Mathlib is still great news, and the size improvements for the caches are a significant QOL win, because they are usually the last to finish when running lake exe cache get.

[kim-em]

@nomeata, I'm marking this back to awaiting-author now that there are conflicts.

[nomeata]

Absolutely! Although I think I’ll see leanprover-community/batteries#285 (which I also have to update now) through first, so that refactorings like the present only affect one downstream repo, and not three.

[nomeata]

I’ll close this. Others are working on DiscrTree more actively these days, have a better sense if this is useful, and if it is, can probably cherry-pick the ideas easily.

[JovanGerb]

@nomeata, do you have interest in reviving this PR in some form? As mathlib and Lean continue to grow, the number of simp lemmas keeps growing, and building the simp discrimination tree (at import time) is taking a significant amount of time when importing all of mathlib. The refactor in this PR would be one way to mitigate this inefficiency.

[nomeata]

Do we have evidence that this PR is helping with import times?

[JovanGerb]

I measured that for import Mathlib, the simp discrimination tree takes about 0.5s and the instances discrimination tree takes about 0.1s. When building mathlib, the imports are typically smaller, so it will take less time per file, but with over 8000 files this will add up. (see #mathlib4 > Performance cost of environment extensions @ 💬)

And I recall that when this PR was benchmarked against mathlib there was a significant speedup. So I figured that this speedup is from the reduced import times, and that the effect would be even more significant now that mathlib has grown more.

[nomeata]

Improving startup time for import Mathlib would be nice. I’ll ask claude to update this PR and we can get some measurement :-)

[JovanGerb]

An alternative approach would be to make use of a Thunk data structure so as to never compute the entire discrimination tree. However when I tried this, I realized that the implementation would require some more flexibility from the environment extension framework: at import time we would have to first construct the root node of the discrimination tree, and after that is created, map over it with some Thunk.mk. Otherwise we would need to use a Thunk.map for each imported simp lemma, which I somehow imagine to not be efficient.

[nomeata]

Closing this PR in favor of #12838

Some form of lazy loading may possibly also be interesting, of course.