From 3bc7f54c1fc2b30e17823d4ec04d6821a52d9e70 Mon Sep 17 00:00:00 2001 From: wiasliaw Date: Tue, 1 Oct 2024 19:54:31 +0800 Subject: [PATCH] move pmtree into zk-kit (#53) ## Description Closes #4. ## Other information This PR moved [pmtree](https://github.com/Rate-Limiting-Nullifier/pmtree) into zk-kit. Some of the information should be check: - [ ] License: There are two different LICENSE in pmtree and need to confirm whether to use the MIT license or follow pmtree's license. - [ ] README: Since zk-kit-pmt has not been released, badges may be unavailable. ## Checklist - [x] My code follows the style guidelines of this project - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have run `make lint` and `cargo build --release` without getting any errors - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes --- Cargo.lock | 371 ++++++++++++++++++++++++++ crates/pmt/Cargo.toml | 16 ++ crates/pmt/LICENSE-APACHE | 201 ++++++++++++++ crates/pmt/LICENSE-MIT | 21 ++ crates/pmt/README.md | 110 ++++++++ crates/pmt/examples/pmt.rs | 84 ++++++ crates/pmt/src/database.rs | 31 +++ crates/pmt/src/hasher.rs | 23 ++ crates/pmt/src/lib.rs | 58 ++++ crates/pmt/src/tree.rs | 429 ++++++++++++++++++++++++++++++ crates/pmt/tests/memory_keccak.rs | 156 +++++++++++ crates/pmt/tests/sled_keccak.rs | 204 ++++++++++++++ 12 files changed, 1704 insertions(+) create mode 100644 crates/pmt/Cargo.toml create mode 100644 crates/pmt/LICENSE-APACHE create mode 100644 crates/pmt/LICENSE-MIT create mode 100644 crates/pmt/README.md create mode 100644 crates/pmt/examples/pmt.rs create mode 100644 crates/pmt/src/database.rs create mode 100644 crates/pmt/src/hasher.rs create mode 100644 crates/pmt/src/lib.rs create mode 100644 crates/pmt/src/tree.rs create mode 100644 crates/pmt/tests/memory_keccak.rs create mode 100644 crates/pmt/tests/sled_keccak.rs diff --git a/Cargo.lock b/Cargo.lock index d1994f7..0c56e9b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,24 +2,177 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "ark-serialize" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d6c2b318ee6e10f8c2853e73a83adc0ccb88995aa978d8a3408d492ab2ee671" +dependencies = [ + "ark-std", + "digest", +] + +[[package]] +name = "ark-std" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1df2c09229cbc5a028b1d70e00fdb2acee28b1055dfb5ca73eea49c5a25c4e7c" +dependencies = [ + "num-traits", + "rand", +] + [[package]] name = "autocfg" version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "crc32fast" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a97769d94ddab943e4510d138150169a2758b5ef3eb191a9ee688de3e23ef7b3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + [[package]] name = "crunchy" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" +[[package]] +name = "digest" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3dd60d1080a57a05ab032377049e0591415d2b31afd7028356dbf3cc6dcb066" +dependencies = [ + "generic-array", +] + +[[package]] +name = "either" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" + +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "fxhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" +dependencies = [ + "byteorder", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "hex" version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hex-literal" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ebdb29d2ea9ed0083cd8cece49bbd968021bd99b0849edb4a9a7ee0fdf6a4e0" + +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "libc" +version = "0.2.158" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" + +[[package]] +name = "lock_api" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" + [[package]] name = "num-bigint" version = "0.4.6" @@ -48,6 +201,152 @@ dependencies = [ "autocfg", ] +[[package]] +name = "parking_lot" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" +dependencies = [ + "instant", + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a2cfe6f0ad2bfc16aefa463b497d5c7a5ecd44a23efa72aa342d90177356dc" +dependencies = [ + "cfg-if", + "instant", + "libc", + "redox_syscall", + "smallvec", + "winapi", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.86" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" + +[[package]] +name = "rayon" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d2df5196e37bcc87abebc0053e20787d73847bb33134a69841207dd0a47f03b" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "sled" +version = "0.34.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f96b4737c2ce5987354855aed3797279def4ebf734436c6aa4552cf8e169935" +dependencies = [ + "crc32fast", + "crossbeam-epoch", + "crossbeam-utils", + "fs2", + "fxhash", + "libc", + "log", + "parking_lot", +] + +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + +[[package]] +name = "syn" +version = "2.0.77" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f35bcdf61fd8e7be6caf75f429fdca8beb3ed76584befb503b1569faee373ed" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "tiny-keccak" version = "2.0.2" @@ -57,6 +356,67 @@ dependencies = [ "crunchy", ] +[[package]] +name = "typenum" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" + +[[package]] +name = "unicode-ident" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "byteorder", + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zk-kit-imt" version = "0.0.6" @@ -65,6 +425,17 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "zk-kit-pmt" +version = "0.1.0" +dependencies = [ + "ark-serialize", + "hex-literal", + "rayon", + "sled", + "tiny-keccak", +] + [[package]] name = "zk-kit-smt" version = "0.0.5" diff --git a/crates/pmt/Cargo.toml b/crates/pmt/Cargo.toml new file mode 100644 index 0000000..a8f6de2 --- /dev/null +++ b/crates/pmt/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "zk-kit-pmt" +version = "0.1.0" +edition.workspace = true +license.workspace = true +publish.workspace = true +description = "Generic for storage Merkle Tree (sparse & fixed-size) in Rust" + +[dependencies] +rayon = { version = "=1.7.0", optional = false } + +[dev-dependencies] +ark-serialize = "=0.3.0" +hex-literal = "=0.3.4" +sled = "=0.34.7" +tiny-keccak = { version = "=2.0.2", features = ["keccak"] } diff --git a/crates/pmt/LICENSE-APACHE b/crates/pmt/LICENSE-APACHE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/crates/pmt/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/crates/pmt/LICENSE-MIT b/crates/pmt/LICENSE-MIT new file mode 100644 index 0000000..c13f991 --- /dev/null +++ b/crates/pmt/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/crates/pmt/README.md b/crates/pmt/README.md new file mode 100644 index 0000000..5539857 --- /dev/null +++ b/crates/pmt/README.md @@ -0,0 +1,110 @@ +

+

+ Persistent Merkle Tree +

+

Persistent Merkle Tree (optimized & sparse & fixed-size) in Rust.

+

+ +
+

+ + 🗣️ Chat & Support + +

+
+ +## References + +- https://github.com/Rate-Limiting-Nullifier/pmtree + +--- + +## 🛠 Install + +```bash +cargo add zk-kit-pmt +``` + +## 📜 Usage + +```rust +struct MemoryDB(HashMap); +struct MyKeccak(()); + +#[derive(Default)] +struct MemoryDBConfig; + +impl Database for MemoryDB { + type Config = MemoryDBConfig; + + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { + Ok(MemoryDB(HashMap::new())) + } + + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )) + } + + fn get(&self, key: DBKey) -> PmtreeResult> { + Ok(self.0.get(&key).cloned()) + } + + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { + self.0.insert(key, value); + + Ok(()) + } + + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { + self.0.extend(subtree); + + Ok(()) + } + + fn close(&mut self) -> PmtreeResult<()> { + Ok(()) + } +} + +impl Hasher for MyKeccak { + type Fr = [u8; 32]; + + fn default_leaf() -> Self::Fr { + [0; 32] + } + + fn serialize(value: Self::Fr) -> Value { + value.to_vec() + } + + fn deserialize(value: Value) -> Self::Fr { + value.try_into().unwrap() + } + + fn hash(input: &[Self::Fr]) -> Self::Fr { + let mut output = [0; 32]; + let mut hasher = Keccak::v256(); + for element in input { + hasher.update(element); + } + hasher.finalize(&mut output); + output + } +} + +fn main() { + let mut mt = MerkleTree::::new(2, MemoryDBConfig).unwrap(); + + assert_eq!(mt.capacity(), 4); + assert_eq!(mt.depth(), 2); + + mt.update_next(hex!( + "c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95" + )) + .unwrap(); + // closes the connection to the database + mt.close().unwrap(); +} +``` diff --git a/crates/pmt/examples/pmt.rs b/crates/pmt/examples/pmt.rs new file mode 100644 index 0000000..7b9a61a --- /dev/null +++ b/crates/pmt/examples/pmt.rs @@ -0,0 +1,84 @@ +use hex_literal::hex; +use std::collections::HashMap; +use tiny_keccak::{Hasher as _, Keccak}; +use zk_kit_pmt::*; + +struct MemoryDB(HashMap); +struct MyKeccak(()); + +#[derive(Default)] +struct MemoryDBConfig; + +impl Database for MemoryDB { + type Config = MemoryDBConfig; + + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { + Ok(MemoryDB(HashMap::new())) + } + + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )) + } + + fn get(&self, key: DBKey) -> PmtreeResult> { + Ok(self.0.get(&key).cloned()) + } + + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { + self.0.insert(key, value); + + Ok(()) + } + + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { + self.0.extend(subtree); + + Ok(()) + } + + fn close(&mut self) -> PmtreeResult<()> { + Ok(()) + } +} + +impl Hasher for MyKeccak { + type Fr = [u8; 32]; + + fn default_leaf() -> Self::Fr { + [0; 32] + } + + fn serialize(value: Self::Fr) -> Value { + value.to_vec() + } + + fn deserialize(value: Value) -> Self::Fr { + value.try_into().unwrap() + } + + fn hash(input: &[Self::Fr]) -> Self::Fr { + let mut output = [0; 32]; + let mut hasher = Keccak::v256(); + for element in input { + hasher.update(element); + } + hasher.finalize(&mut output); + output + } +} + +fn main() { + let mut mt = MerkleTree::::new(2, MemoryDBConfig).unwrap(); + + assert_eq!(mt.capacity(), 4); + assert_eq!(mt.depth(), 2); + + mt.update_next(hex!( + "c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95" + )) + .unwrap(); + // closes the connection to the database + mt.close().unwrap(); +} diff --git a/crates/pmt/src/database.rs b/crates/pmt/src/database.rs new file mode 100644 index 0000000..55f7989 --- /dev/null +++ b/crates/pmt/src/database.rs @@ -0,0 +1,31 @@ +use crate::*; + +use std::collections::HashMap; + +/// Trait that must be implemented for a Database +pub trait Database { + /// Config for database. Default is necessary for a default() pmtree function + type Config: Default; + + /// Creates new instance of db + fn new(config: Self::Config) -> PmtreeResult + where + Self: Sized; + + /// Loades existing db (existence check required) + fn load(config: Self::Config) -> PmtreeResult + where + Self: Sized; + + /// Returns value from db by the key + fn get(&self, key: DBKey) -> PmtreeResult>; + + /// Puts the value to the db by the key + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()>; + + /// Puts the leaves batch to the db + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()>; + + /// Closes the db connection + fn close(&mut self) -> PmtreeResult<()>; +} diff --git a/crates/pmt/src/hasher.rs b/crates/pmt/src/hasher.rs new file mode 100644 index 0000000..3769390 --- /dev/null +++ b/crates/pmt/src/hasher.rs @@ -0,0 +1,23 @@ +use crate::*; + +use std::fmt::Debug; + +/// Trait that must be implemented for Hash Function +pub trait Hasher { + /// Native type for the hash-function + type Fr: Copy + Eq + Default + Sync + Send + Debug; + + /// Serializes Self::Fr + fn serialize(value: Self::Fr) -> Value; + + /// Deserializes Self::Fr + fn deserialize(value: Value) -> Self::Fr; + + /// Outputs the default leaf (Fr::default()) + fn default_leaf() -> Self::Fr { + Self::Fr::default() + } + + /// Calculates hash-function + fn hash(input: &[Self::Fr]) -> Self::Fr; +} diff --git a/crates/pmt/src/lib.rs b/crates/pmt/src/lib.rs new file mode 100644 index 0000000..efbeb16 --- /dev/null +++ b/crates/pmt/src/lib.rs @@ -0,0 +1,58 @@ +pub mod database; +pub mod hasher; +pub mod tree; + +use std::fmt::{Debug, Display}; + +pub use database::*; +pub use hasher::*; +pub use tree::MerkleTree; + +/// Denotes keys in a database +pub type DBKey = [u8; 8]; + +/// Denotes values in a database +pub type Value = Vec; + +/// Denotes pmtree Merkle tree errors +#[derive(Debug)] +pub enum TreeErrorKind { + MerkleTreeIsFull, + InvalidKey, + IndexOutOfBounds, + CustomError(String), +} + +/// Denotes pmtree database errors +#[derive(Debug)] +pub enum DatabaseErrorKind { + CannotLoadDatabase, + DatabaseExists, + CustomError(String), +} + +/// Denotes pmtree errors +#[derive(Debug)] +pub enum PmtreeErrorKind { + /// Error in database + DatabaseError(DatabaseErrorKind), + /// Error in tree + TreeError(TreeErrorKind), + /// Custom error + CustomError(String), +} + +impl Display for PmtreeErrorKind { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + PmtreeErrorKind::DatabaseError(e) => write!(f, "Database error: {e:?}"), + PmtreeErrorKind::TreeError(e) => write!(f, "Tree error: {e:?}"), + PmtreeErrorKind::CustomError(e) => write!(f, "Custom error: {e:?}"), + } + } +} + +impl std::error::Error for PmtreeErrorKind {} + +/// Custom `Result` type with custom `Error` type +pub type PmtreeResult = std::result::Result; diff --git a/crates/pmt/src/tree.rs b/crates/pmt/src/tree.rs new file mode 100644 index 0000000..9c801e1 --- /dev/null +++ b/crates/pmt/src/tree.rs @@ -0,0 +1,429 @@ +use crate::*; + +use std::cmp::{max, min}; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +// db[DEPTH_KEY] = depth +const DEPTH_KEY: DBKey = (u64::MAX - 1).to_be_bytes(); + +// db[NEXT_INDEX_KEY] = next_index; +const NEXT_INDEX_KEY: DBKey = u64::MAX.to_be_bytes(); + +// Default tree depth +const DEFAULT_TREE_DEPTH: usize = 20; + +// Denotes keys (depth, index) in Merkle Tree. Can be converted to DBKey +// TODO! Think about using hashing for that +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Key(usize, usize); +impl From for DBKey { + fn from(key: Key) -> Self { + let cantor_pairing = ((key.0 + key.1) * (key.0 + key.1 + 1) / 2 + key.1) as u64; + cantor_pairing.to_be_bytes() + } +} + +/// The Merkle Tree structure +pub struct MerkleTree +where + D: Database, + H: Hasher, +{ + pub db: D, + depth: usize, + next_index: usize, + cache: Vec, + root: H::Fr, +} + +/// The Merkle proof structure +#[derive(Clone, PartialEq, Eq)] +pub struct MerkleProof(pub Vec<(H::Fr, u8)>); + +impl MerkleTree +where + D: Database, + H: Hasher, +{ + /// Creates tree with specified depth and default "pmtree_db" dbpath. + pub fn default(depth: usize) -> PmtreeResult { + Self::new(depth, D::Config::default()) + } + + /// Creates new `MerkleTree` and store it to the specified path/db + pub fn new(depth: usize, db_config: D::Config) -> PmtreeResult { + // Create new db instance + let mut db = D::new(db_config)?; + + // Insert depth val into db + let depth_val = depth.to_be_bytes().to_vec(); + db.put(DEPTH_KEY, depth_val)?; + + // Insert next_index val into db + let next_index = 0usize; + let next_index_val = next_index.to_be_bytes().to_vec(); + db.put(NEXT_INDEX_KEY, next_index_val)?; + + // Cache nodes + let mut cache = vec![H::default_leaf(); depth + 1]; + + // Initialize one branch of the `Merkle Tree` from bottom to top + cache[depth] = H::default_leaf(); + db.put(Key(depth, 0).into(), H::serialize(cache[depth]))?; + for i in (0..depth).rev() { + cache[i] = H::hash(&[cache[i + 1], cache[i + 1]]); + db.put(Key(i, 0).into(), H::serialize(cache[i]))?; + } + + let root = cache[0]; + + Ok(Self { + db, + depth, + next_index, + cache, + root, + }) + } + + /// Loads existing Merkle Tree from the specified path/db + pub fn load(db_config: D::Config) -> PmtreeResult { + // Load existing db instance + let db = D::load(db_config)?; + + // Load root + let root = match db.get(Key(0, 0).into())? { + Some(root) => H::deserialize(root), + None => H::default_leaf(), + }; + + // Load depth & next_index values from db + let depth = match db.get(DEPTH_KEY)? { + Some(depth) => usize::from_be_bytes(depth.try_into().unwrap()), + None => DEFAULT_TREE_DEPTH, + }; + + let next_index = match db.get(NEXT_INDEX_KEY)? { + Some(next_index) => usize::from_be_bytes(next_index.try_into().unwrap()), + None => 0, + }; + + // Load cache vec + let mut cache = vec![H::default_leaf(); depth + 1]; + cache[depth] = H::default_leaf(); + for i in (0..depth).rev() { + cache[i] = H::hash(&[cache[i + 1], cache[i + 1]]); + } + + Ok(Self { + db, + depth, + next_index, + cache, + root, + }) + } + + /// Closes the db connection + pub fn close(&mut self) -> PmtreeResult<()> { + self.db.close() + } + + /// Sets a leaf at the specified tree index + pub fn set(&mut self, key: usize, leaf: H::Fr) -> PmtreeResult<()> { + if key >= self.capacity() { + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds)); + } + + self.db + .put(Key(self.depth, key).into(), H::serialize(leaf))?; + self.recalculate_from(key)?; + + // Update next_index in memory + self.next_index = max(self.next_index, key + 1); + + // Update next_index in db + let next_index_val = self.next_index.to_be_bytes().to_vec(); + self.db.put(NEXT_INDEX_KEY, next_index_val)?; + + Ok(()) + } + + // Recalculates `Merkle Tree` from the specified key + fn recalculate_from(&mut self, key: usize) -> PmtreeResult<()> { + let mut depth = self.depth; + let mut i = key; + + loop { + let value = self.hash_couple(depth, i)?; + i >>= 1; + depth -= 1; + self.db.put(Key(depth, i).into(), H::serialize(value))?; + + if depth == 0 { + self.root = value; + break; + } + } + + Ok(()) + } + + // Hashes the correct couple for the key + fn hash_couple(&self, depth: usize, key: usize) -> PmtreeResult { + let b = key & !1; + Ok(H::hash(&[ + self.get_elem(Key(depth, b))?, + self.get_elem(Key(depth, b + 1))?, + ])) + } + + // Returns elem by the key + pub fn get_elem(&self, key: Key) -> PmtreeResult { + let res = self + .db + .get(key.into())? + .map_or(self.cache[key.0], |value| H::deserialize(value)); + + Ok(res) + } + + /// Deletes a leaf at the `key` by setting it to its default value + pub fn delete(&mut self, key: usize) -> PmtreeResult<()> { + if key >= self.next_index { + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)); + } + + self.set(key, H::default_leaf())?; + + Ok(()) + } + + /// Inserts a leaf to the next available index + pub fn update_next(&mut self, leaf: H::Fr) -> PmtreeResult<()> { + self.set(self.next_index, leaf)?; + + Ok(()) + } + + /// Batch insertion from starting index + pub fn set_range>( + &mut self, + start: usize, + leaves: I, + ) -> PmtreeResult<()> { + self.batch_insert( + Some(start), + leaves.into_iter().collect::>().as_slice(), + ) + } + + /// Batch insertion, updates the tree in parallel. + pub fn batch_insert(&mut self, start: Option, leaves: &[H::Fr]) -> PmtreeResult<()> { + let start = start.unwrap_or(self.next_index); + let end = start + leaves.len(); + + if end > self.capacity() { + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::MerkleTreeIsFull)); + } + + let mut subtree = HashMap::::new(); + + let root_key = Key(0, 0); + + subtree.insert(root_key, self.root); + self.fill_nodes(root_key, start, end, &mut subtree, leaves, start)?; + + let subtree = Arc::new(RwLock::new(subtree)); + + let root_val = rayon::ThreadPoolBuilder::new() + .num_threads(rayon::current_num_threads()) + .build() + .unwrap() + .install(|| Self::batch_recalculate(root_key, Arc::clone(&subtree), self.depth)); + + let subtree = RwLock::into_inner(Arc::try_unwrap(subtree).unwrap()).unwrap(); + + self.db.put_batch( + subtree + .into_iter() + .map(|(key, value)| (key.into(), H::serialize(value))) + .collect(), + )?; + + // Update next_index value in db + if end > self.next_index { + self.next_index = end; + self.db + .put(NEXT_INDEX_KEY, self.next_index.to_be_bytes().to_vec())?; + } + + // Update root value in memory + self.root = root_val; + + Ok(()) + } + + // Fills hashmap subtree + fn fill_nodes( + &self, + key: Key, + start: usize, + end: usize, + subtree: &mut HashMap, + leaves: &[H::Fr], + from: usize, + ) -> PmtreeResult<()> { + if key.0 == self.depth { + if key.1 >= from { + subtree.insert(key, leaves[key.1 - from]); + } + return Ok(()); + } + + let left = Key(key.0 + 1, key.1 * 2); + let right = Key(key.0 + 1, key.1 * 2 + 1); + + let left_val = self.get_elem(left)?; + let right_val = self.get_elem(right)?; + + subtree.insert(left, left_val); + subtree.insert(right, right_val); + + let half = 1 << (self.depth - key.0 - 1); + + if start < half { + self.fill_nodes(left, start, min(end, half), subtree, leaves, from)?; + } + + if end > half { + self.fill_nodes(right, 0, end - half, subtree, leaves, from)?; + } + + Ok(()) + } + + // Recalculates tree in parallel (in-memory) + fn batch_recalculate( + key: Key, + subtree: Arc>>, + depth: usize, + ) -> H::Fr { + let left_child = Key(key.0 + 1, key.1 * 2); + let right_child = Key(key.0 + 1, key.1 * 2 + 1); + + if key.0 == depth || !subtree.read().unwrap().contains_key(&left_child) { + return *subtree.read().unwrap().get(&key).unwrap(); + } + + let (left, right) = rayon::join( + || Self::batch_recalculate(left_child, Arc::clone(&subtree), depth), + || Self::batch_recalculate(right_child, Arc::clone(&subtree), depth), + ); + + let result = H::hash(&[left, right]); + + subtree.write().unwrap().insert(key, result); + + result + } + + /// Computes a Merkle proof for the leaf at the specified index + pub fn proof(&self, index: usize) -> PmtreeResult> { + if index >= self.capacity() { + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds)); + } + + let mut witness = Vec::with_capacity(self.depth); + + let mut i = index; + let mut depth = self.depth; + while depth != 0 { + i ^= 1; + witness.push(( + self.get_elem(Key(depth, i))?, + (1 - (i & 1)).try_into().unwrap(), + )); + i >>= 1; + depth -= 1; + } + + Ok(MerkleProof(witness)) + } + + /// Verifies a Merkle proof with respect to the input leaf and the tree root + pub fn verify(&self, leaf: &H::Fr, witness: &MerkleProof) -> bool { + let expected_root = witness.compute_root_from(leaf); + + self.root() == expected_root + } + + /// Returns the leaf by the key + pub fn get(&self, key: usize) -> PmtreeResult { + if key >= self.capacity() { + return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds)); + } + + self.get_elem(Key(self.depth, key)) + } + + /// Returns the root of the tree + pub fn root(&self) -> H::Fr { + self.root + } + + /// Returns the total number of leaves set + pub fn leaves_set(&self) -> usize { + self.next_index + } + + /// Returns the capacity of the tree, i.e. the maximum number of leaves + pub fn capacity(&self) -> usize { + 1 << self.depth + } + + /// Returns the depth of the tree + pub fn depth(&self) -> usize { + self.depth + } +} + +impl MerkleProof { + /// Computes the Merkle root by iteratively hashing specified Merkle proof with specified leaf + pub fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr { + let mut acc = *leaf; + for w in self.0.iter() { + if w.1 == 0 { + acc = H::hash(&[acc, w.0]); + } else { + acc = H::hash(&[w.0, acc]); + } + } + + acc + } + + /// Computes the leaf index corresponding to a Merkle proof + pub fn leaf_index(&self) -> usize { + self.get_path_index() + .into_iter() + .rev() + .fold(0, |acc, digit| (acc << 1) + usize::from(digit)) + } + + /// Returns the path indexes forming a Merkle Proof + pub fn get_path_index(&self) -> Vec { + self.0.iter().map(|x| x.1).collect() + } + + /// Returns the path elements forming a Merkle proof + pub fn get_path_elements(&self) -> Vec { + self.0.iter().map(|x| x.0).collect() + } + + /// Returns the length of a Merkle proof + pub fn length(&self) -> usize { + self.0.len() + } +} diff --git a/crates/pmt/tests/memory_keccak.rs b/crates/pmt/tests/memory_keccak.rs new file mode 100644 index 0000000..7d8ea0b --- /dev/null +++ b/crates/pmt/tests/memory_keccak.rs @@ -0,0 +1,156 @@ +use hex_literal::hex; +use std::collections::HashMap; +use tiny_keccak::{Hasher as _, Keccak}; +use zk_kit_pmt::*; + +struct MemoryDB(HashMap); +struct MyKeccak(()); + +#[derive(Default)] +struct MemoryDBConfig; + +impl Database for MemoryDB { + type Config = MemoryDBConfig; + + fn new(_db_config: MemoryDBConfig) -> PmtreeResult { + Ok(MemoryDB(HashMap::new())) + } + + fn load(_db_config: MemoryDBConfig) -> PmtreeResult { + Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )) + } + + fn get(&self, key: DBKey) -> PmtreeResult> { + Ok(self.0.get(&key).cloned()) + } + + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { + self.0.insert(key, value); + + Ok(()) + } + + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { + self.0.extend(subtree); + + Ok(()) + } + + fn close(&mut self) -> PmtreeResult<()> { + Ok(()) + } +} + +impl Hasher for MyKeccak { + type Fr = [u8; 32]; + + fn default_leaf() -> Self::Fr { + [0; 32] + } + + fn serialize(value: Self::Fr) -> Value { + value.to_vec() + } + + fn deserialize(value: Value) -> Self::Fr { + value.try_into().unwrap() + } + + fn hash(input: &[Self::Fr]) -> Self::Fr { + let mut output = [0; 32]; + let mut hasher = Keccak::v256(); + for element in input { + hasher.update(element); + } + hasher.finalize(&mut output); + output + } +} + +#[test] +fn insert_delete() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new(2, MemoryDBConfig)?; + + assert_eq!(mt.capacity(), 4); + assert_eq!(mt.depth(), 2); + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + hex!("0000000000000000000000000000000000000000000000000000000000000003"), + hex!("0000000000000000000000000000000000000000000000000000000000000004"), + ]; + + let default_tree_root = + hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30"); + + assert_eq!(mt.root(), default_tree_root); + + let roots = [ + hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"), + hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b"), + hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c"), + hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"), + ]; + + for i in 0..leaves.len() { + mt.update_next(leaves[i])?; + assert_eq!(mt.root(), roots[i]); + } + + for (i, &leaf) in leaves.iter().enumerate() { + assert!(mt.verify(&leaf, &mt.proof(i)?)); + } + + for i in (0..leaves.len()).rev() { + mt.delete(i)?; + } + + assert_eq!(mt.root(), default_tree_root); + + assert!(mt.update_next(leaves[0]).is_err()); + + Ok(()) +} + +#[test] +fn batch_insertions() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new(2, MemoryDBConfig)?; + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + hex!("0000000000000000000000000000000000000000000000000000000000000003"), + hex!("0000000000000000000000000000000000000000000000000000000000000004"), + ]; + + mt.batch_insert(None, &leaves)?; + + assert_eq!( + mt.root(), + hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36") + ); + + Ok(()) +} + +#[test] +fn set_range() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new(2, MemoryDBConfig)?; + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + ]; + + mt.set_range(2, leaves)?; + + assert_eq!( + mt.root(), + hex!("1e9f6c8d3fd5b7ae3a29792adb094c6d4cc6149d0c81c8c8e57cf06c161a92b8") + ); + + Ok(()) +} diff --git a/crates/pmt/tests/sled_keccak.rs b/crates/pmt/tests/sled_keccak.rs new file mode 100644 index 0000000..4db44af --- /dev/null +++ b/crates/pmt/tests/sled_keccak.rs @@ -0,0 +1,204 @@ +use hex_literal::hex; +use std::collections::HashMap; +use std::fs; +use tiny_keccak::{Hasher as _, Keccak}; +use zk_kit_pmt::*; + +struct MyKeccak(()); +struct MySled(sled::Db); + +#[derive(Default)] +struct SledConfig { + path: String, +} + +impl Database for MySled { + type Config = SledConfig; + + fn new(db_config: SledConfig) -> PmtreeResult { + let db = sled::open(db_config.path).unwrap(); + if db.was_recovered() { + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::DatabaseExists, + )); + } + + Ok(MySled(db)) + } + + fn load(db_config: SledConfig) -> PmtreeResult { + let db = sled::open(&db_config.path).unwrap(); + + if !db.was_recovered() { + fs::remove_dir_all(&db_config.path).expect("Error removing db"); + return Err(PmtreeErrorKind::DatabaseError( + DatabaseErrorKind::CannotLoadDatabase, + )); + } + + Ok(MySled(db)) + } + + fn get(&self, key: DBKey) -> PmtreeResult> { + Ok(self.0.get(key).unwrap().map(|val| val.to_vec())) + } + + fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> { + self.0.insert(key, value).unwrap(); + + self.0.flush().unwrap(); + + Ok(()) + } + + fn put_batch(&mut self, subtree: HashMap) -> PmtreeResult<()> { + let mut batch = sled::Batch::default(); + + for (key, value) in subtree { + batch.insert(&key, value); + } + + self.0.apply_batch(batch).unwrap(); + + Ok(()) + } + + fn close(&mut self) -> PmtreeResult<()> { + self.0.flush().unwrap(); + + Ok(()) + } +} + +impl Hasher for MyKeccak { + type Fr = [u8; 32]; + + fn default_leaf() -> Self::Fr { + [0; 32] + } + + fn serialize(value: Self::Fr) -> Value { + value.to_vec() + } + + fn deserialize(value: Value) -> Self::Fr { + value.to_vec().try_into().unwrap() + } + + fn hash(input: &[Self::Fr]) -> Self::Fr { + let mut output = [0; 32]; + let mut hasher = Keccak::v256(); + for element in input { + hasher.update(element); + } + hasher.finalize(&mut output); + output + } +} + +#[test] +fn insert_delete() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new( + 2, + SledConfig { + path: String::from("abacabas"), + }, + )?; + + assert_eq!(mt.capacity(), 4); + assert_eq!(mt.depth(), 2); + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + hex!("0000000000000000000000000000000000000000000000000000000000000003"), + hex!("0000000000000000000000000000000000000000000000000000000000000004"), + ]; + + let default_tree_root = + hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30"); + + assert_eq!(mt.root(), default_tree_root); + + let roots = [ + hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"), + hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b"), + hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c"), + hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"), + ]; + + for i in 0..leaves.len() { + mt.update_next(leaves[i])?; + assert_eq!(mt.root(), roots[i]); + } + + for (i, &leaf) in leaves.iter().enumerate() { + assert!(mt.verify(&leaf, &mt.proof(i)?)); + } + + for i in (0..leaves.len()).rev() { + mt.delete(i)?; + } + + assert_eq!(mt.root(), default_tree_root); + + assert!(mt.update_next(leaves[0]).is_err()); + + fs::remove_dir_all("abacabas").expect("Error removing db"); + + Ok(()) +} + +#[test] +fn batch_insertions() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new( + 2, + SledConfig { + path: String::from("abacabasa"), + }, + )?; + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + hex!("0000000000000000000000000000000000000000000000000000000000000003"), + hex!("0000000000000000000000000000000000000000000000000000000000000004"), + ]; + + mt.batch_insert(None, &leaves)?; + + assert_eq!( + mt.root(), + hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36") + ); + + fs::remove_dir_all("abacabasa").expect("Error removing db"); + + Ok(()) +} + +#[test] +fn set_range() -> PmtreeResult<()> { + let mut mt = MerkleTree::::new( + 2, + SledConfig { + path: String::from("abacabasab"), + }, + )?; + + let leaves = [ + hex!("0000000000000000000000000000000000000000000000000000000000000001"), + hex!("0000000000000000000000000000000000000000000000000000000000000002"), + ]; + + mt.set_range(2, leaves)?; + + assert_eq!( + mt.root(), + hex!("1e9f6c8d3fd5b7ae3a29792adb094c6d4cc6149d0c81c8c8e57cf06c161a92b8") + ); + + fs::remove_dir_all("abacabasab").expect("Error removing db"); + + Ok(()) +}