├── .github └── workflows │ └── ci.yml ├── .gitignore ├── CHANGELOG.md ├── Cargo.toml ├── LICENSE ├── README.md ├── benches ├── merk.rs └── ops.rs ├── docs └── algorithms.md ├── merk-dark.svg ├── merk.svg ├── rustfmt.toml ├── scripts └── pgo.sh └── src ├── error.rs ├── lib.rs ├── merk ├── chunks.rs ├── mod.rs ├── restore.rs └── snapshot.rs ├── owner.rs ├── proofs ├── chunk.rs ├── encoding.rs ├── mod.rs ├── query │ ├── map.rs │ └── mod.rs └── tree.rs ├── test_utils ├── crash_merk.rs ├── mod.rs └── temp_merk.rs └── tree ├── commit.rs ├── debug.rs ├── encoding.rs ├── fuzz_tests.rs ├── hash.rs ├── iter.rs ├── kv.rs ├── link.rs ├── mod.rs ├── ops.rs └── walk ├── fetch.rs ├── mod.rs └── ref_walker.rs /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: [master, develop] 6 | pull_request: 7 | branches: [master, develop] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | test-base: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - name: Checkout 17 | uses: actions/checkout@v2 18 | - name: Use Nightly 19 | uses: actions-rs/toolchain@v1 20 | with: 21 | toolchain: nightly-2024-04-25 22 | override: true 23 | - name: Cache 24 | uses: actions/cache@v3 25 | with: 26 | path: | 27 | ~/.cargo/bin/ 28 | ~/.cargo/registry/index/ 29 | ~/.cargo/registry/cache/ 30 | ~/.cargo/git/db/ 31 | ~/.cargo/registry/src/**/librocksdb-sys-* 32 | target/ 33 | key: ${{ runner.os }}-test-base-${{ hashFiles('Cargo.toml') }} 34 | - name: Test 35 | uses: actions-rs/cargo@v1 36 | with: 37 | command: test 38 | args: --verbose 39 | 40 | test-all-features: 41 | runs-on: ubuntu-latest 42 | steps: 43 | - name: Checkout 44 | uses: actions/checkout@v2 45 | - name: Use Nightly 46 | uses: actions-rs/toolchain@v1 47 | with: 48 | toolchain: nightly-2024-04-25 49 | override: true 50 | - name: Cache 51 | uses: actions/cache@v3 52 | with: 53 | path: | 54 | ~/.cargo/bin/ 55 | ~/.cargo/registry/index/ 56 | ~/.cargo/registry/cache/ 57 | ~/.cargo/git/db/ 58 | ~/.cargo/registry/src/**/librocksdb-sys-* 59 | target/ 60 | key: ${{ runner.os }}-test-all-features-${{ hashFiles('Cargo.toml') }} 61 | - name: Test 62 | uses: actions-rs/cargo@v1 63 | with: 64 | command: test 65 | args: --verbose --all-features 66 | 67 | coverage: 68 | runs-on: ubuntu-latest 69 | steps: 70 | - name: Checkout 71 | uses: actions/checkout@v2 72 | - name: Use Nightly 73 | uses: actions-rs/toolchain@v1 74 | with: 75 | toolchain: nightly-2024-04-25 76 | components: llvm-tools-preview 77 | override: true 78 | - name: Cache 79 | uses: actions/cache@v3 80 | with: 81 | path: | 82 | ~/.cargo/bin/ 83 | ~/.cargo/registry/index/ 84 | ~/.cargo/registry/cache/ 85 | ~/.cargo/git/db/ 86 | ~/.cargo/registry/src/**/librocksdb-sys-* 87 | target/ 88 | key: ${{ runner.os }}-coverage-${{ hashFiles('Cargo.toml') }} 89 | - name: Install Coverage Tooling 90 | uses: actions-rs/cargo@v1 91 | with: 92 | command: install 93 | args: cargo-llvm-cov --force 94 | - name: Run Coverage 95 | uses: actions-rs/cargo@v1 96 | with: 97 | command: llvm-cov 98 | args: --all-features --workspace --lcov --output-path lcov.info 99 | - name: Upload to codecov.io 100 | uses: codecov/codecov-action@v1 101 | with: 102 | token: ${{ secrets.CODECOV_TOKEN }} 103 | files: lcov.info 104 | fail_ci_if_error: true 105 | 106 | format: 107 | runs-on: ubuntu-latest 108 | steps: 109 | - name: Checkout 110 | uses: actions/checkout@v2 111 | - name: Use Nightly 112 | uses: actions-rs/toolchain@v1 113 | with: 114 | toolchain: nightly-2024-04-25 115 | components: rustfmt 116 | override: true 117 | - name: Check 118 | uses: actions-rs/cargo@v1 119 | with: 120 | command: fmt 121 | args: --all -- --check 122 | 123 | clippy: 124 | runs-on: ubuntu-latest 125 | steps: 126 | - name: Checkout 127 | uses: actions/checkout@v2 128 | - name: Use Nightly 129 | uses: actions-rs/toolchain@v1 130 | with: 131 | toolchain: nightly-2024-04-25 132 | components: clippy 133 | override: true 134 | - name: Cache 135 | uses: actions/cache@v3 136 | with: 137 | path: | 138 | ~/.cargo/bin/ 139 | ~/.cargo/registry/index/ 140 | ~/.cargo/registry/cache/ 141 | ~/.cargo/git/db/ 142 | ~/.cargo/registry/src/**/librocksdb-sys-* 143 | target/ 144 | key: ${{ runner.os }}-clippy-${{ hashFiles('Cargo.toml') }} 145 | - name: Check 146 | uses: actions-rs/clippy-check@v1 147 | with: 148 | token: ${{ secrets.GITHUB_TOKEN }} 149 | args: --all-features -- -D warnings 150 | 151 | benches: 152 | runs-on: ubuntu-latest 153 | steps: 154 | - name: Checkout 155 | uses: actions/checkout@v2 156 | - name: Use Nightly 157 | uses: actions-rs/toolchain@v1 158 | with: 159 | toolchain: nightly-2024-04-25 160 | override: true 161 | - name: Cache 162 | uses: actions/cache@v3 163 | with: 164 | path: | 165 | ~/.cargo/bin/ 166 | ~/.cargo/registry/index/ 167 | ~/.cargo/registry/cache/ 168 | ~/.cargo/git/db/ 169 | ~/.cargo/registry/src/**/librocksdb-sys-* 170 | target/ 171 | key: ${{ runner.os }}-benches-${{ hashFiles('Cargo.toml') }} 172 | - name: Run Benches 173 | uses: actions-rs/cargo@v1 174 | with: 175 | command: bench 176 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | temp.db 3 | .DS_Store 4 | Cargo.lock 5 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | ## [Unreleased] 4 | 5 | ### Bug Fixes 6 | 7 | - Fixed bug where column families would be non-atomically flushed when one memtable was filled, resulting in inconsistency after a crash. 8 | 9 | [Unreleased]: https://github.com/nomic-io/merk/compare/v1.0.0-alpha.8...HEAD 10 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "merk" 3 | description = "High-performance Merkle key/value store" 4 | version = "2.0.0" 5 | authors = ["Turbofish "] 6 | edition = "2018" 7 | license = "Apache-2.0" 8 | 9 | [dependencies] 10 | thiserror= "1.0.58" 11 | sha2 = "0.10.8" 12 | log = "0.4.21" 13 | 14 | [dependencies.colored] 15 | version = "2.1.0" 16 | optional = true 17 | 18 | [dependencies.num_cpus] 19 | version = "1.16.0" 20 | optional = true 21 | 22 | [dependencies.ed] 23 | version = "0.3.0" 24 | optional = true 25 | 26 | [dependencies.rand] 27 | version = "0.8.5" 28 | features = ["small_rng"] 29 | optional = true 30 | 31 | [dependencies.rocksdb] 32 | version = "0.22.0" 33 | default-features = false 34 | optional = true 35 | 36 | [dependencies.jemallocator] 37 | version = "0.5.4" 38 | features = ["disable_initial_exec_tls"] 39 | optional = true 40 | 41 | [features] 42 | default = ["full", "verify"] 43 | full = [ 44 | "rand", 45 | "rocksdb", 46 | "colored", 47 | "num_cpus", 48 | "ed", 49 | ] 50 | verify = ["ed"] 51 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |

2 | 3 | 4 | 5 | merk 6 | 7 |

8 | 9 | *High-performance Merkle key/value store* 10 | 11 | ![CI](https://github.com/turbofish-org/merk/actions/workflows/ci.yml/badge.svg) 12 | [![codecov](https://codecov.io/gh/turbofish-org/merk/branch/develop/graph/badge.svg?token=TTUTSt2iLz)](https://codecov.io/gh/turbofish-org/merk) 13 | [![Crate](https://img.shields.io/crates/v/merk.svg)](https://crates.io/crates/merk) 14 | [![API](https://docs.rs/merk/badge.svg)](https://docs.rs/merk) 15 | 16 | Merk is a crypto key/value store - more specifically, it's a Merkle AVL tree built on top of RocksDB (Facebook's fork of LevelDB). 17 | 18 | Its priorities are performance and reliability. While Merk was designed to be the state database for blockchains, it can also be used anywhere an auditable key/value store is needed. 19 | 20 | ### Features 21 | - **Fast reads/writes** - Reads have no overhead compared to a normal RocksDB store, and writes are optimized for batch operations (e.g. blocks in a blockchain). 22 | - **Fast proof generation** - Since Merk implements an AVL tree rather than a trie, it is very efficient to create and verify proofs for ranges of keys. 23 | - **Concurrency** - Unlike most other Merkle stores, all operations utilize all available cores - giving huge performance gains and allowing nodes to scale along with Moore's Law. 24 | - **Replication** - The tree is optimized to efficiently build proofs of large chunks, allowing for nodes to download the entire state (e.g. "state syncing"). 25 | - **Checkpointing** - Merk can create checkpoints on disk (an immutable view of the entire store at a certain point in time) without blocking, so there are no delays in availability or liveness. 26 | - **Web-friendly** - Being written in Rust means it is easy to run the proof-verification code in browsers with WebAssembly, allowing for light-clients that can verify data for themselves. 27 | - **Fits any Profile** - Performant on RAM-constrained Raspberry Pi's and beefy validator rigs alike. 28 | 29 | The algorithms are based on AVL, but optimized for batches of operations and random fetches from the backing store. 30 | 31 | ## Usage 32 | 33 | **Install:** 34 | ``` 35 | cargo add merk 36 | ``` 37 | 38 | **Example:** 39 | ```rust 40 | extern crate merk; 41 | use merk::*; 42 | 43 | // load or create a Merk store at the given path 44 | let mut merk = Merk::open("./merk.db").unwrap(); 45 | 46 | // apply some operations 47 | let batch = [ 48 | (b"key", Op::Put(b"value")), 49 | (b"key2", Op::Put(b"value2")), 50 | (b"key3", Op::Put(b"value3")), 51 | (b"key4", Op::Delete) 52 | ]; 53 | merk.apply(&batch).unwrap(); 54 | ``` 55 | Merk is currently used by [Nomic](https://github.com/nomic-io/nomic), a blockchain powering decentralized custody of Bitcoin, built on [Orga](https://github.com/turbofish-org/orga). 56 | 57 | ## Benchmarks 58 | 59 | Benchmarks are measured on a 1M node tree, each node having a key length of 16 bytes and value length of 40 bytes. All tests are single-threaded (not counting RocksDB background threads). 60 | 61 | You can test these yourself by running `cargo bench`. 62 | 63 | ### 2017 Macbook Pro 64 | 65 | *(Using 1 Merk thread and 4 RocksDB compaction threads)* 66 | 67 | **Pruned (no state kept in memory)** 68 | 69 | *RAM usage:* ~20MB average, ~26MB max 70 | 71 | | Test | Ops per second | 72 | | -------- | ------ | 73 | | Random inserts | 23,000 | 74 | | Random updates | 32,000 | 75 | | Random deletes | 26,000 | 76 | | Random reads | 210,000 | 77 | | Random proof generation | 133,000 | 78 | 79 | **Cached (all state kept in memory)** 80 | 81 | *RAM usage:* ~400MB average, ~1.1GB max 82 | 83 | | Test | Ops per second | 84 | | -------- | ------ | 85 | | Random inserts | 58,000 | 86 | | Random updates | 81,000 | 87 | | Random deletes | 72,000 | 88 | | Random reads | 1,565,000 | 89 | | Random proof generation | 311,000 | 90 | 91 | ### i9-9900K Desktop 92 | 93 | *(Using 1 Merk thread and 16 RocksDB compaction threads)* 94 | 95 | **Pruned (no state kept in memory)** 96 | 97 | *RAM usage:* ~20MB average, ~26MB max 98 | 99 | | Test | Ops per second | 100 | | -------- | ------ | 101 | | Random inserts | 40,000 | 102 | | Random updates | 55,000 | 103 | | Random deletes | 45,000 | 104 | | Random reads | 383,000 | 105 | | Random proof generation | 249,000 | 106 | 107 | **Cached (all state kept in memory)** 108 | 109 | *RAM usage:* ~400MB average, ~1.1GB max 110 | 111 | | Test | Ops per second | 112 | | -------- | ------ | 113 | | Random inserts | 93,000 | 114 | | Random updates | 123,000 | 115 | | Random deletes | 111,000 | 116 | | Random reads | 2,370,000 | 117 | | Random proof generation | 497,000 | 118 | 119 | ## Contributing 120 | 121 | Merk is an open-source project spearheaded by Turbofish. Anyone is able to contribute to Merk via GitHub. 122 | 123 | [Contribute to Merk](https://github.com/turbofish-org/merk/contribute) 124 | 125 | ## Security 126 | 127 | ### Security Audits 128 | 129 | | Date | Auditor | Scope | Report | 130 | | ---: | :---: | :--- | :---: | 131 | | October 2024 | Trail of Bits | `orga` `merk` `ed` `abci2` | [📄](https://github.com/trailofbits/publications/blob/master/reviews/2024-11-orgaandmerk-securityreview.pdf) | 132 | 133 | Vulnerabilities should not be reported through public channels, including GitHub Issues. You can report a vulnerability via GitHub's Private Vulnerability Reporting or to Turbofish at `security@turbofish.org`. 134 | 135 | [Report a Vulnerability](https://github.com/turbofish-org/merk/security/advisories/new) 136 | 137 | ## License 138 | 139 | Licensed under the Apache License, Version 2.0 (the "License"); you may not use the files in this repository except in compliance with the License. You may obtain a copy of the License at 140 | 141 | https://www.apache.org/licenses/LICENSE-2.0 142 | 143 | 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. 144 | 145 | --- 146 | 147 | Copyright © 2024 Turbofish, Inc. 148 | -------------------------------------------------------------------------------- /benches/merk.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | 3 | extern crate test; 4 | 5 | use merk::proofs::encode_into as encode_proof_into; 6 | use merk::restore::Restorer; 7 | use merk::test_utils::*; 8 | use merk::{Merk, Result}; 9 | use rand::prelude::*; 10 | use std::thread; 11 | use test::Bencher; 12 | 13 | #[bench] 14 | fn get_1m_rocksdb(b: &mut Bencher) { 15 | let initial_size = 1_000_000; 16 | let batch_size = 2_000; 17 | let num_batches = initial_size / batch_size; 18 | 19 | let path = thread::current().name().unwrap().to_owned(); 20 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 21 | 22 | let mut batches = vec![]; 23 | for i in 0..num_batches { 24 | let batch = make_batch_rand(batch_size, i); 25 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 26 | batches.push(batch); 27 | } 28 | 29 | let mut i = 0; 30 | b.iter(|| { 31 | let batch_index = (i % num_batches) as usize; 32 | let key_index = (i / num_batches) as usize; 33 | 34 | let key = &batches[batch_index][key_index].0; 35 | merk.get(key).expect("get failed"); 36 | 37 | i = (i + 1) % initial_size; 38 | }); 39 | } 40 | 41 | #[bench] 42 | fn insert_1m_2k_seq_rocksdb_noprune(b: &mut Bencher) { 43 | let initial_size = 1_000_000; 44 | let batch_size = 2_000; 45 | 46 | let path = thread::current().name().unwrap().to_owned(); 47 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 48 | 49 | for i in 0..(initial_size / batch_size) { 50 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 51 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 52 | } 53 | 54 | let mut i = initial_size / batch_size; 55 | b.iter(|| { 56 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 57 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 58 | i += 1; 59 | }); 60 | } 61 | 62 | #[bench] 63 | fn insert_1m_2k_rand_rocksdb_noprune(b: &mut Bencher) { 64 | let initial_size = 1_000_000; 65 | let batch_size = 2_000; 66 | 67 | let path = thread::current().name().unwrap().to_owned(); 68 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 69 | 70 | for i in 0..(initial_size / batch_size) { 71 | let batch = make_batch_rand(batch_size, i); 72 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 73 | } 74 | 75 | let mut i = initial_size / batch_size; 76 | b.iter(|| { 77 | let batch = make_batch_rand(batch_size, i); 78 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 79 | i += 1; 80 | }); 81 | } 82 | 83 | #[bench] 84 | fn update_1m_2k_seq_rocksdb_noprune(b: &mut Bencher) { 85 | let initial_size = 1_000_000; 86 | let batch_size = 2_000; 87 | 88 | let path = thread::current().name().unwrap().to_owned(); 89 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 90 | 91 | for i in 0..(initial_size / batch_size) { 92 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 93 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 94 | } 95 | 96 | let mut i = 0; 97 | b.iter(|| { 98 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 99 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 100 | i = (i + 1) % (initial_size / batch_size); 101 | }); 102 | } 103 | 104 | #[bench] 105 | fn update_1m_2k_rand_rocksdb_noprune(b: &mut Bencher) { 106 | let initial_size = 1_000_000; 107 | let batch_size = 2_000; 108 | 109 | let path = thread::current().name().unwrap().to_owned(); 110 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 111 | 112 | for i in 0..(initial_size / batch_size) { 113 | let batch = make_batch_rand(batch_size, i); 114 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 115 | } 116 | 117 | let mut i = 0; 118 | b.iter(|| { 119 | let batch = make_batch_rand(batch_size, i); 120 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 121 | i = (i + 1) % (initial_size / batch_size); 122 | }); 123 | } 124 | 125 | #[bench] 126 | fn delete_1m_2k_rand_rocksdb_noprune(b: &mut Bencher) { 127 | let initial_size = 1_000_000; 128 | let batch_size = 2_000; 129 | 130 | let path = thread::current().name().unwrap().to_owned(); 131 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 132 | 133 | for i in 0..(initial_size / batch_size) { 134 | let batch = make_batch_rand(batch_size, i); 135 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 136 | } 137 | 138 | let mut i = 0; 139 | b.iter(|| { 140 | if i >= (initial_size / batch_size) { 141 | println!("WARNING: too many bench iterations, whole tree deleted"); 142 | return; 143 | } 144 | let batch = make_del_batch_rand(batch_size, i); 145 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 146 | i = (i + 1) % (initial_size / batch_size); 147 | }); 148 | } 149 | 150 | #[bench] 151 | fn prove_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 152 | let initial_size = 1_000_000; 153 | let batch_size = 1_000; 154 | let proof_size = 1; 155 | 156 | let path = thread::current().name().unwrap().to_owned(); 157 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 158 | 159 | for i in 0..(initial_size / batch_size) { 160 | let batch = make_batch_rand(batch_size, i); 161 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 162 | } 163 | 164 | let mut i = 0; 165 | b.iter(|| { 166 | let batch = make_batch_rand(proof_size, i); 167 | let mut keys = Vec::with_capacity(batch.len()); 168 | for (key, _) in batch { 169 | keys.push(merk::proofs::query::QueryItem::Key(key)); 170 | } 171 | merk.prove(keys).expect("prove failed"); 172 | i = (i + 1) % (initial_size / batch_size); 173 | 174 | merk.commit(std::collections::LinkedList::new(), &[]) 175 | .unwrap(); 176 | }); 177 | } 178 | 179 | #[bench] 180 | fn build_trunk_chunk_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 181 | let initial_size = 1_000_000; 182 | let batch_size = 1_000; 183 | 184 | let path = thread::current().name().unwrap().to_owned(); 185 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 186 | 187 | for i in 0..(initial_size / batch_size) { 188 | let batch = make_batch_rand(batch_size, i); 189 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 190 | } 191 | 192 | let mut bytes = vec![]; 193 | 194 | b.iter(|| { 195 | bytes.clear(); 196 | 197 | let (ops, _) = merk.walk(|walker| walker.unwrap().create_trunk_proof().unwrap()); 198 | encode_proof_into(ops.iter(), &mut bytes); 199 | 200 | merk.commit(std::collections::LinkedList::new(), &[]) 201 | .unwrap(); 202 | }); 203 | 204 | b.bytes = bytes.len() as u64; 205 | } 206 | 207 | #[bench] 208 | fn chunkproducer_rand_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 209 | let mut rng = rand::thread_rng(); 210 | 211 | let initial_size = 1_000_000; 212 | let batch_size = 1_000; 213 | 214 | let path = thread::current().name().unwrap().to_owned(); 215 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 216 | 217 | for i in 0..(initial_size / batch_size) { 218 | let batch = make_batch_rand(batch_size, i); 219 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 220 | } 221 | 222 | let mut chunks = merk.chunks().unwrap(); 223 | let mut total_bytes = 0; 224 | let mut i = 0; 225 | 226 | let mut next = || { 227 | let index = rng.gen::() % chunks.len(); 228 | chunks.chunk(index).unwrap() 229 | }; 230 | 231 | b.iter(|| { 232 | let chunk = next(); 233 | total_bytes += chunk.len(); 234 | i += 1; 235 | }); 236 | 237 | b.bytes = (total_bytes / i) as u64; 238 | } 239 | 240 | #[bench] 241 | fn chunk_iter_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 242 | let initial_size = 1_000_000; 243 | let batch_size = 1_000; 244 | 245 | let path = thread::current().name().unwrap().to_owned(); 246 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 247 | 248 | for i in 0..(initial_size / batch_size) { 249 | let batch = make_batch_rand(batch_size, i); 250 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 251 | } 252 | 253 | let mut chunks = merk.chunks().unwrap().into_iter(); 254 | let mut total_bytes = 0; 255 | let mut i = 0; 256 | 257 | let mut next = || match chunks.next() { 258 | Some(chunk) => chunk, 259 | None => { 260 | chunks = merk.chunks().unwrap().into_iter(); 261 | chunks.next().unwrap() 262 | } 263 | }; 264 | 265 | b.iter(|| { 266 | let chunk = next(); 267 | total_bytes += chunk.unwrap().len(); 268 | i += 1; 269 | }); 270 | 271 | b.bytes = (total_bytes / i) as u64; 272 | } 273 | 274 | #[bench] 275 | fn restore_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 276 | let initial_size = 1_000_000; 277 | let batch_size = 1_000; 278 | 279 | let path = thread::current().name().unwrap().to_owned(); 280 | let mut merk = TempMerk::open(path).expect("failed to open merk"); 281 | 282 | for i in 0..(initial_size / batch_size) { 283 | let batch = make_batch_rand(batch_size, i); 284 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 285 | } 286 | 287 | let chunks = merk 288 | .chunks() 289 | .unwrap() 290 | .into_iter() 291 | .collect::>>() 292 | .unwrap(); 293 | 294 | let path = thread::current().name().unwrap().to_owned() + "_restore"; 295 | let mut restorer: Option = None; 296 | 297 | let mut total_bytes = 0; 298 | let mut i = 0; 299 | 300 | b.iter(|| { 301 | if i % chunks.len() == 0 { 302 | if i != 0 { 303 | let restorer_merk = restorer.take().unwrap().finalize(); 304 | drop(restorer_merk); 305 | std::fs::remove_dir_all(&path).unwrap(); 306 | } 307 | 308 | restorer = Some(Merk::restore(&path, merk.root_hash(), chunks.len()).unwrap()); 309 | } 310 | 311 | let restorer = restorer.as_mut().unwrap(); 312 | let chunk = chunks[i % chunks.len()].as_slice(); 313 | restorer.process_chunk(chunk).unwrap(); 314 | 315 | total_bytes += chunk.len(); 316 | i += 1; 317 | }); 318 | 319 | std::fs::remove_dir_all(&path).unwrap(); 320 | 321 | b.bytes = (total_bytes / i) as u64; 322 | } 323 | 324 | #[bench] 325 | fn checkpoint_create_destroy_1m_1_rand_rocksdb_noprune(b: &mut Bencher) { 326 | let initial_size = 1_000_000; 327 | let batch_size = 1_000; 328 | 329 | let path = thread::current().name().unwrap().to_owned(); 330 | let mut merk = TempMerk::open(&path).expect("failed to open merk"); 331 | 332 | for i in 0..(initial_size / batch_size) { 333 | let batch = make_batch_rand(batch_size, i); 334 | unsafe { merk.apply_unchecked(&batch, &[]).expect("apply failed") }; 335 | } 336 | 337 | let path = path + ".checkpoint"; 338 | b.iter(|| { 339 | let checkpoint = merk.checkpoint(&path).unwrap(); 340 | checkpoint.destroy().unwrap(); 341 | }); 342 | } 343 | -------------------------------------------------------------------------------- /benches/ops.rs: -------------------------------------------------------------------------------- 1 | #![feature(test)] 2 | 3 | extern crate test; 4 | 5 | use merk::owner::Owner; 6 | use merk::test_utils::*; 7 | use test::Bencher; 8 | 9 | #[bench] 10 | fn insert_1m_10k_seq_memonly(b: &mut Bencher) { 11 | let initial_size = 1_000_000; 12 | let batch_size = 10_000; 13 | 14 | let mut tree = Owner::new(make_tree_seq(initial_size)); 15 | 16 | let mut i = initial_size / batch_size; 17 | b.iter(|| { 18 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 19 | tree.own(|tree| apply_memonly_unchecked(tree, &batch)); 20 | i += 1; 21 | }); 22 | } 23 | 24 | #[bench] 25 | fn insert_1m_10k_rand_memonly(b: &mut Bencher) { 26 | let initial_size = 1_000_000; 27 | let batch_size = 10_000; 28 | 29 | let mut tree = Owner::new(make_tree_rand(initial_size, batch_size, 0)); 30 | 31 | let mut i = initial_size / batch_size; 32 | b.iter(|| { 33 | let batch = make_batch_rand(batch_size, i); 34 | tree.own(|tree| apply_memonly_unchecked(tree, &batch)); 35 | i += 1; 36 | }); 37 | } 38 | 39 | #[bench] 40 | fn update_1m_10k_seq_memonly(b: &mut Bencher) { 41 | let initial_size = 1_000_000; 42 | let batch_size = 10_000; 43 | 44 | let mut tree = Owner::new(make_tree_seq(initial_size)); 45 | 46 | let mut i = 0; 47 | b.iter(|| { 48 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 49 | tree.own(|tree| apply_memonly_unchecked(tree, &batch)); 50 | i = (i + 1) % (initial_size / batch_size); 51 | }); 52 | } 53 | 54 | #[bench] 55 | fn update_1m_10k_rand_memonly(b: &mut Bencher) { 56 | let initial_size = 1_010_000; 57 | let batch_size = 10_000; 58 | 59 | let mut tree = Owner::new(make_tree_rand(initial_size, batch_size, 0)); 60 | 61 | let mut i = 0; 62 | b.iter(|| { 63 | let batch = make_batch_rand(batch_size, i); 64 | tree.own(|tree| apply_memonly_unchecked(tree, &batch)); 65 | i = (i + 1) % (initial_size / batch_size); 66 | }); 67 | } 68 | -------------------------------------------------------------------------------- /docs/algorithms.md: -------------------------------------------------------------------------------- 1 | # Merk - A High-Performance Merkle AVL Tree 2 | 3 | **Matt Bell ([@mappum](https://twitter.com/mappum))** • [Nomic Hodlings, Inc.](https://nomic.io) 4 | 5 | v0.0.4 - _August 5, 2020_ 6 | 7 | ## Introduction 8 | 9 | Merk is a Merkle AVL tree designed for performance, running on top of a backing key/value store such as RocksDB. Notable features include concurrent operations for higher throughput, an optimized key/value layout for performant usage of the backing store, and efficient proof generation to enable bulk tree replication. 10 | 11 | _Note that this document is meant to be a way to grok how Merk works, rather than an authoritative specification._ 12 | 13 | ## Algorithm Overview 14 | 15 | The Merk tree was inspired by [`tendermint/iavl`](https://github.com/tendermint/iavl) from the [Tendermint](https://tendermint.com) team but makes various fundamental design changes in the name of performance. 16 | 17 | ### Tree Structure 18 | 19 | #### Nodes and Hashing 20 | 21 | In many Merkle tree designs, only leaf nodes contain key/value pairs (inner nodes only contain child hashes). To contrast, every node in a Merk tree contains a key and a value, including inner nodes. 22 | 23 | Each node contains a "kv hash", which is the hash of its key/value pair, in addition to its child hashes. The hash of the node is just the hash of the concatenation of these three hashes: 24 | 25 | ``` 26 | kv_hash = H(key, value) 27 | node_hash = H(kv_hash, left_child_hash, right_child_hash) 28 | ``` 29 | 30 | Note that the `left_child_hash` and/or `right_child_hash` values may be null since it is possible for the node to have no children or only one child. 31 | 32 | In our implementation, the hash function used is SHA512/256 (SHA512 with output truncated to 256 bits) but this choice is trivially swappable. 33 | 34 | #### Database Representation 35 | 36 | In the backing key/value store, nodes are stored using their key/value pair key as the database key, and a binary encoding that contains the fields in the above `Node` structure - minus the `key` field since that is already implied by the database entry. 37 | 38 | Storing nodes by key rather than by hash is an important optimization, and is the reason why inner nodes each have a key/value pair. The implication is that reading a key does not require traversing through the tree structure but only requires a single read in the backing key/value store, meaning there is practically no overhead versus using the backing store without a tree structure. Additionally, we can efficiently iterate through nodes in the tree in their in-order traversal just by iterating by key in the backing store (which RocksDB and LevelDB are optimized for). 39 | 40 | This means we lose the "I" compared to the IAVL library - immutability. Since now we operate on the tree nodes in-place in the backing store, we don't by default have views of past states of the tree. However, **in** our implementation we replicate this functionality with RocksDB's snapshot and checkpoint features which provide a consistent view of the store at a certain point in history - either ephemerally in memory or persistently on disk. 41 | 42 | ### Operations 43 | 44 | Operating on a Merk tree is optimized for batches - in the real world we will only be updating the tree once per block, applying a batch of many changes from many transactions at the same time. 45 | 46 | #### Concurrent Batch Operator 47 | 48 | To mutate the tree, we apply batches of operations, each of which can either be `Put(key, value)` or `Delete(key)`. 49 | 50 | Batches of operations are expected to be sorted by key, with every key appearing only once. Our implementation provides an `apply` method which sorts the batch and checks for duplicate keys, and an `apply_unchecked` method which skips the sorting/checking step for performance reasons when the caller has already ensured the batch is sorted. 51 | 52 | The algorithm to apply these operations to the tree is called recursively on each relevant node. 53 | 54 | _Simplified pseudocode for the operation algorithm:_ 55 | 56 | - Given a node and a batch of operations: 57 | - Binary search for the current node's key in the batch: 58 | - If this node's key is found in the batch at index `i`: 59 | - Apply the operation to this node: 60 | - If operation is `Put`, update its `value` and `kv_hash` 61 | - If the operation is `Delete`, perform a traditional BST node removal 62 | - Split the batch into left and right sub-batches (excluding the operation we just applied): 63 | - Left batch from batch start to index `i` 64 | - Right batch from index `i + 1` to the end of the batch 65 | - If this node's key is not found in the batch, but could be inserted at index `i` maintaining sorted order: 66 | - Split the batch into left and right sub-batches: 67 | - Left batch from batch start to index `i` 68 | - Right batch from index `i` to the end of the batch 69 | - Recurse: 70 | - Apply the left sub-batch to this node's left child 71 | - Apply the right sub-batch to this node's right child 72 | - Balance: 73 | - If after recursing the left and right subtrees are unbalanced (their heights differ by more than 1), perform an AVL tree rotation (possibly more than one) 74 | - Recompute node's hash based on hash of its updated children and `kv_hash`, then return 75 | 76 | This batch application of operations can happen concurrently - recursing into the left and right subtrees of a node are two fully independent operations (operations on one subtree will never involve reading or writing to/from any of the nodes on the other subtree). This means we have an _implicit lock_ - we don't need to coordinate with mutexes but only need to wait for both the left side and right side to finish their operations. 77 | 78 | ### Proofs 79 | 80 | Merk was designed with efficient proofs in mind, both for application queries (e.g. a user checking their account balance) and bulk tree replication (a.k.a. "state syncing") between untrusted nodes. 81 | 82 | #### Structure 83 | 84 | Merk proofs are a list of stack-based operators and node data, with 3 possible operators: `Push(node)`, `Parent`, and `Child`. A stream of these operators can be processed by a verifier in order to reconstruct a sparse representation of part of the tree, in a way where the data can be verified against a known root hash. 85 | 86 | The value of `node` in a `Push` operation can be one of three types: 87 | 88 | - `Hash(hash)` - The hash of a node 89 | - `KVHash(hash)` - The key/value hash of a node 90 | - `KV(key, value)` - The key and value of a node 91 | 92 | This proof format can be encoded in a binary format and has negligible space overhead for efficient transport over the network. 93 | 94 | #### Verification 95 | 96 | A verifier can process a proof by maintaining a stack of connected tree nodes, and executing the operators in order: 97 | 98 | - `Push(node)` - Push some node data onto the stack. 99 | - `Child` - Pop a value from the stack, `child`. Pop another value from the stack, `parent`. Set `child` as the right child of `parent`, and push the combined result back on the stack. 100 | - `Parent` - Pop a value from the stack, `parent`. Pop another value from the stack, `child`. Set `child` as the left child of `parent`, and push the combined result back on the stack. 101 | 102 | Proof verification will fail if e.g. `Child` or `Parent` try to pop a value from the stack but the stack is empty, `Child` or `Parent` try to overwrite an existing child, or the proof does not result in exactly one stack item. 103 | 104 | This proof language can be used to specify any possible set or subset of the tree's data in a way that can be reconstructed efficiently by the verifier. Proofs can contain either an arbitrary set of selected key/value pairs (e.g. in an application query), or contiguous tree chunks (when replicating the tree). After processing an entire proof, the verifier should have derived a root hash which can be compared to the root hash they expect (e.g. the one validators committed to in consensus), and have a set of proven key/value pairs. 105 | 106 | Note that this can be computed in a streaming fashion, e.g. while downloading the proof, which makes the required memory for verification very low even for large proofs. However, the verifier cannot tell if the proof is valid until finishing the entire proof, so very large proofs should be broken up into multiple proofs of smaller size. 107 | 108 | #### Generation 109 | 110 | Efficient proof generation is important since nodes will likely receive a high volume of queries and constantly be serving proofs, essentially providing an API service to end-user application clients, as well as servicing demand for replication when new nodes come onto the network. 111 | 112 | Nodes can generate proofs for a set of keys by traversing through the tree from the root and building up the required proof branches. Much like the batch operator aglorithm, this algorithm takes a batch of sorted, unique keys as input. 113 | 114 | _Simplified pseudocode for proof generation (based on an in-order traversal):_ 115 | 116 | - Given a node and a batch of keys to include in the proof: 117 | - If the batch is empty, append `Push(Hash(node_hash))` to the proof and return 118 | - Binary search the for the current node's key in the batch: 119 | - If this node's key is found in the batch at index `i`: 120 | - Partition the batch into left and right sub-batches at index `i` (excluding index `i`) 121 | - If this node's key is not found in the batch, but could be inserted at index `i` maintaining sorted order: 122 | - Partition the batch into left and right sub-batches at index `i` 123 | - **Recurse left:** If there is a left child: 124 | - If the left sub-batch is not empty, query the left child (appending operators to the proof) 125 | - If the left sub-batch is empty, append `Push(Hash(left_child_hash))` to the proof 126 | - Append proof operator: 127 | - If this node's key is in the batch, or if the left sub-batch was not empty and no left child exists, or if the right sub-batch is not empty and no right child exists,or if the left child's right edge queried a non-existent key, or if the right child's left edge queried a non-existent key, append `Push(KV(key, value))` to the proof 128 | - Otherwise, append `Push(KVHash(kv_hash))` to the proof 129 | - If the left child exists, append `Parent` to the proof 130 | - **Recurse right:** If there is a right child: 131 | - If the right sub-batch is not empty, query the right child (appending operators to the proof) 132 | - If the right sub-batch is empty, append `Push(Hash(left_child_hash))` to the proof 133 | - Append `Child` to the proof 134 | 135 | Since RocksDB allows concurrent reading from a consistent snapshot/checkpoint, nodes can concurrently generate proofs on all cores to service a higher volume of queries, even if our algorithm isn't designed for concurrency. 136 | 137 | #### Binary Format 138 | 139 | We can efficiently encode these proofs by encoding each operator as follows: 140 | 141 | ``` 142 | Push(Hash(hash)) => 0x01 <20-byte hash> 143 | Push(KVHash(hash)) => 0x02 <20-byte hash> 144 | Push(KV(key, value)) => 0x03 <1-byte key length> <2-byte value length> 145 | Parent => 0x10 146 | Child => 0x11 147 | ``` 148 | 149 | This results in a compact binary representation, with a very small space overhead (roughly 2 bytes per node in the proof (1 byte for the `Push` operator type flag, and 1 byte for a `Parent` or `Child` operator), plus 3 bytes per key/value pair (1 byte for the key length, and 2 bytes for the value length)). 150 | 151 | #### Efficient Chunk Proofs for Replication 152 | 153 | An alternate, optimized proof generation can be used when generating proofs for large contiguous subtrees, e.g. chunks for tree replication. This works by iterating sequentially through keys in the backing store (which is much faster than random lookups). 154 | 155 | Based on some early benchmarks, I estimate that typical server hardware should be able to generate this kind of range proof at a rate of hundreds of MB/s, which means the bottleneck for bulk replication will likely be bandwidth rather than CPU. To improve performance further, these proofs can be cached and trivially served by a CDN or a P2P swarm (each node of which can easily verify the chunks they pass around). 156 | 157 | Due to the tree structure we already use, streaming the entries in key-order gives us all the nodes to construct complete contiguous subtrees. For instance, in the diagram below, streaming from keys `1` to `7` will give us a complete subtree. This subtree can be verified to be a part of the full tree as long as we know the hash of `4`. 158 | 159 | ``` 160 | 8 161 | / \ 162 | / ... 163 | 4 164 | / \ 165 | 2 6 166 | / \ / \ 167 | 1 3 5 7 168 | ``` 169 | 170 | Our algorithm builds verifiable chunks by first constructing a chunk of the upper levels of the tree, called the _trunk chunk_, plus each subtree below that (each of which is called a _leaf chunk_). 171 | 172 | The number of levels to include in the trunk can be chosen to control the size of the leaf nodes. For example, a tree of height 10 should have approximately 1,023 nodes. If the trunk contains the top 5 levels, the trunk and the 32 resulting leaf nodes will each contain ~31 nodes. We can even prove to the verifier the trunk size was chosen correctly by also including an approximate tree height proof, by including the branch all the way to the leftmost node of the tree (node `1` in the figure) and using this height as our basis to select the number of trunk levels. 173 | 174 | After the prover builds the trunk by traversing from the root node and making random lookups down to the chosen level, it can generate the leaf nodes extremely efficiently by reading the database keys sequentially as described a few paragraphs above. We can trivially detect when a chunk should end whenever a node at or above the trunk level is encountered (e.g. encountering node `8` signals we have read a complete subtree). 175 | 176 | The generated proofs can be efficiently encoded into the same proof format described above. Verifiers only have the added constraint that none of the data should be abbridged (all nodes contain a key and value, rather than just a hash or kvhash). After first downloading and verifying the trunk, verifiers can also download leaf chunks in parallel and verify that each connects to the trunk by comparing each subtree's root hash. 177 | 178 | Note that this algorithm produces proofs with very little memory requirements, plus little overhead added to the sequential read from disk. In a proof-of-concept benchmark, proof generation was measured to be ~750 MB/s on a modern solid-state drive and processor, meaning a 4GB state tree (the size of the Cosmos Hub state at the time of writing) could be fully proven in ~5 seconds (without considering parallelization). In conjunction with the RocksDB checkpoint feature, this process can happen in the background without blocking the node from executing later blocks. 179 | 180 | _Pseudocode for the range proof generation algorithm:_ 181 | 182 | - Given a tree and a range of keys to prove: 183 | - Create a stack of keys (initially empty) 184 | - **Range iteration:** for every key/value entry within the query range in the backing store: 185 | - Append `Push(KV(key, value))` to the proof 186 | - If the current node has a left child, append `Parent` to the proof 187 | - If the current node has a right child, push the right child's key onto the key stack 188 | - If the current node does not have a right child: 189 | - While the current node's key is greater than or equal to the key at the top of the key stack, append `Child` to the proof and pop from the key stack 190 | 191 | Note that this algorithm produces the proof in a streaming fashion and has very little memory requirements (the only overhead is the key stack, which will be small even for extremely large trees since its length is a maximum of `log N`). 192 | 193 | #### Example Proofs 194 | 195 | Let's walk through a concrete proof example. Consider the following tree: 196 | 197 | ``` 198 | 5 199 | / \ 200 | / \ 201 | 2 9 202 | / \ / \ 203 | 1 4 7 11 204 | / / \ / 205 | 3 6 8 10 206 | ``` 207 | 208 | _Small proof:_ 209 | 210 | First, let's create a proof for a small part of the tree. Let's say the user makes a query for keys `1, 2, 3, 4`. 211 | 212 | If we follow our proof generation algorithm, we should get a proof that looks like this: 213 | 214 | ``` 215 | Push(KV(1, )), 216 | Push(KV(2, )), 217 | Parent, 218 | Push(KV(3, )), 219 | Push(KV(4, )), 220 | Parent, 221 | Child, 222 | Push(KVHash()), 223 | Parent, 224 | Push(Hash()), 225 | Child 226 | ``` 227 | 228 | Let's step through verification to show that this proof works. We'll create a verification stack, which starts out empty, and walk through each operator in the proof, in order: 229 | 230 | ``` 231 | Stack: (empty) 232 | ``` 233 | 234 | We will push a key/value pair on the stack, creating a node. However, note that for verification purposes this node will only need to contain the kv_hash which we will compute at this step. 235 | 236 | ``` 237 | Operator: Push(KV(1, )) 238 | 239 | Stack: 240 | 1 241 | ``` 242 | 243 | ``` 244 | Operator: Push(KV(2, )) 245 | 246 | Stack: 247 | 1 248 | 2 249 | ``` 250 | 251 | Now we connect nodes 1 and 2, with 2 as the parent. 252 | 253 | ``` 254 | Operator: Parent 255 | 256 | Stack: 257 | 2 258 | / 259 | 1 260 | ``` 261 | 262 | ``` 263 | Operator: Push(KV(3, )) 264 | 265 | Stack: 266 | 2 267 | / 268 | 1 269 | 3 270 | ``` 271 | 272 | ``` 273 | Operator: Push(KV(4, )) 274 | 275 | Stack: 276 | 2 277 | / 278 | 1 279 | 3 280 | 4 281 | ``` 282 | 283 | ``` 284 | Operator: Parent 285 | 286 | Stack: 287 | 2 288 | / 289 | 1 290 | 4 291 | / 292 | 3 293 | ``` 294 | 295 | Now connect these two graphs with 4 as the child of 2. 296 | 297 | ``` 298 | Operator: Child 299 | 300 | Stack: 301 | 2 302 | / \ 303 | 1 4 304 | / 305 | 3 306 | ``` 307 | 308 | Since the user isn't querying the data from node 5, we only need its kv_hash. 309 | 310 | ``` 311 | Operator: Push(KVHash()) 312 | 313 | Stack: 314 | 2 315 | / \ 316 | 1 4 317 | / 318 | 3 319 | 5 320 | ``` 321 | 322 | ``` 323 | Operator: Parent 324 | 325 | Stack: 326 | 5 327 | / 328 | 2 329 | / \ 330 | 1 4 331 | / 332 | 3 333 | ``` 334 | 335 | We only need the hash of node 9. 336 | 337 | ``` 338 | Operator: Push(Hash()) 339 | 340 | Stack: 341 | 5 342 | / 343 | 2 344 | / \ 345 | 1 4 346 | / 347 | 3 348 | 9 349 | ``` 350 | 351 | ``` 352 | Operator: Child 353 | 354 | Stack: 355 | 5 356 | / \ 357 | 2 9 358 | / \ 359 | 1 4 360 | / 361 | 3 362 | ``` 363 | 364 | Now after going through all these steps, we have sufficient knowlege of the tree's structure and data to compute node hashes in order to verify. At the end, we will have computed a hash for node 5 (the root), and we verify by comparing this hash to the one we expected. 365 | -------------------------------------------------------------------------------- /merk-dark.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /merk.svg: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | comment_width = 80 2 | wrap_comments = true 3 | 4 | -------------------------------------------------------------------------------- /scripts/pgo.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | default_host_triple="" 4 | default_toolchain="" 5 | IFS=" = " 6 | while read -r name value 7 | do 8 | value="${value//\"/}" 9 | if [ "${name}" == "default_host_triple" ]; then 10 | default_host_triple="${value}" 11 | elif [ "${name}" == "default_toolchain" ]; then 12 | default_toolchain="${value}" 13 | fi 14 | done < ~/.rustup/settings.toml 15 | 16 | echo "default_host_triple=${default_host_triple}" 17 | echo "default_toolchain=${default_toolchain}" 18 | 19 | rustup component add llvm-tools-preview 20 | 21 | rm -rf /tmp/merk-pgo 22 | RUSTFLAGS="-Cprofile-generate=/tmp/merk-pgo" cargo bench rand_rocks 23 | ~/.rustup/toolchains/${default_toolchain}/lib/rustlib/${default_host_triple}/bin/llvm-profdata merge -o /tmp/merk-pgo/merged.profdata /tmp/merk-pgo 24 | RUSTFLAGS="-Cprofile-use=/tmp/merk-pgo/merged.profdata" cargo bench 25 | -------------------------------------------------------------------------------- /src/error.rs: -------------------------------------------------------------------------------- 1 | pub use thiserror::Error; 2 | 3 | #[derive(Error, Debug)] 4 | pub enum Error { 5 | #[error("Attach Error: {0}")] 6 | Attach(String), 7 | #[error("Batch Key Error: {0}")] 8 | BatchKey(String), 9 | #[error("Bound Error: {0}")] 10 | Bound(String), 11 | #[error("Chunk Processing Error: {0}")] 12 | ChunkProcessing(String), 13 | #[error(transparent)] 14 | Ed(#[from] ed::Error), 15 | #[error("Fetch Error: {0}")] 16 | Fetch(String), 17 | #[error("Proof did not match expected hash\n\tExpected: {0:?}\n\tActual: {1:?}")] 18 | HashMismatch([u8; 32], [u8; 32]), 19 | #[error("Index OoB Error: {0}")] 20 | IndexOutOfBounds(String), 21 | #[error("Integer conversion error: {0}")] 22 | IntegerConversionError(#[from] std::num::TryFromIntError), 23 | #[error(transparent)] 24 | IO(#[from] std::io::Error), 25 | #[error("Tried to delete non-existent key {0:?}")] 26 | KeyDelete(Vec), 27 | #[error("Key Error: {0}")] 28 | Key(String), 29 | #[error("Key not found: {0}")] 30 | KeyNotFound(String), 31 | #[error("Proof is missing data for query")] 32 | MissingData, 33 | #[error("Path Error: {0}")] 34 | Path(String), 35 | #[error("Proof Error: {0}")] 36 | Proof(String), 37 | #[cfg(feature = "full")] 38 | #[error(transparent)] 39 | RocksDB(#[from] rocksdb::Error), 40 | #[error("Stack Underflow")] 41 | StackUnderflow, 42 | #[error("Tree Error: {0}")] 43 | Tree(String), 44 | #[error("Unexpected Node Error: {0}")] 45 | UnexpectedNode(String), 46 | #[error("Unknown Error")] 47 | Unknown, 48 | #[error("Version Error: {0}")] 49 | Version(String), 50 | } 51 | 52 | pub type Result = std::result::Result; 53 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A high-performance Merkle key/value store. 2 | //! 3 | //! Merk is a crypto key/value store - more specifically, it's a Merkle AVL tree 4 | //! built on top of RocksDB (Facebook's fork of LevelDB). 5 | //! 6 | //! Its priorities are performance and reliability. While Merk was designed to 7 | //! be the state database for blockchains, it can also be used anywhere an 8 | //! auditable key/value store is needed. 9 | 10 | #![feature(trivial_bounds)] 11 | 12 | #[global_allocator] 13 | #[cfg(feature = "jemallocator")] 14 | static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc; 15 | 16 | #[cfg(feature = "full")] 17 | pub use rocksdb; 18 | 19 | /// Error and Result types. 20 | mod error; 21 | /// The top-level store API. 22 | #[cfg(feature = "full")] 23 | mod merk; 24 | /// Provides a container type that allows temporarily taking ownership of a 25 | /// value. 26 | // TODO: move this into its own crate 27 | pub mod owner; 28 | /// Algorithms for generating and verifying Merkle proofs. 29 | pub mod proofs; 30 | 31 | /// Various helpers useful for tests or benchmarks. 32 | #[cfg(feature = "full")] 33 | pub mod test_utils; 34 | /// The core tree data structure. 35 | pub mod tree; 36 | 37 | #[cfg(feature = "full")] 38 | pub use crate::merk::{chunks, restore, snapshot, Merk, MerkSource, Snapshot}; 39 | 40 | pub use error::{Error, Result}; 41 | pub use tree::{Batch, BatchEntry, Hash, Op, PanicSource, HASH_LENGTH}; 42 | 43 | #[allow(deprecated)] 44 | pub use proofs::query::verify_query; 45 | 46 | pub use proofs::query::verify; 47 | -------------------------------------------------------------------------------- /src/merk/chunks.rs: -------------------------------------------------------------------------------- 1 | //! Provides `ChunkProducer`, which creates chunk proofs for full replication of 2 | //! a Merk. 3 | 4 | use super::Merk; 5 | use crate::proofs::{chunk::get_next_chunk, Node, Op}; 6 | 7 | use crate::{Error, Result}; 8 | use ed::Encode; 9 | use rocksdb::DBRawIterator; 10 | 11 | /// A `ChunkProducer` allows the creation of chunk proofs, used for trustlessly 12 | /// replicating entire Merk trees. 13 | /// 14 | /// Chunks can be generated on the fly in a random order, or iterated in order 15 | /// for slightly better performance. 16 | pub struct ChunkProducer<'a> { 17 | trunk: Vec, 18 | chunk_boundaries: Vec>, 19 | raw_iter: DBRawIterator<'a>, 20 | index: usize, 21 | } 22 | 23 | impl<'a> ChunkProducer<'a> { 24 | /// Creates a new `ChunkProducer` for the given `Merk` instance. In the 25 | /// constructor, the first chunk (the "trunk") will be created. 26 | pub fn new(merk: &'a Merk) -> Result { 27 | let (trunk, has_more) = merk.walk(|maybe_walker| match maybe_walker { 28 | Some(mut walker) => walker.create_trunk_proof(), 29 | None => Ok((vec![], false)), 30 | })?; 31 | 32 | let chunk_boundaries = if has_more { 33 | trunk 34 | .iter() 35 | .filter_map(|op| match op { 36 | Op::Push(Node::KV(key, _)) => Some(key.clone()), 37 | _ => None, 38 | }) 39 | .collect() 40 | } else { 41 | vec![] 42 | }; 43 | 44 | let mut raw_iter = merk.raw_iter(); 45 | raw_iter.seek_to_first(); 46 | 47 | Ok(ChunkProducer { 48 | trunk, 49 | chunk_boundaries, 50 | raw_iter, 51 | index: 0, 52 | }) 53 | } 54 | 55 | /// Gets the chunk with the given index. Errors if the index is out of 56 | /// bounds or the tree is empty - the number of chunks can be checked by 57 | /// calling `producer.len()`. 58 | pub fn chunk(&mut self, index: usize) -> Result> { 59 | if index >= self.len() { 60 | return Err(Error::IndexOutOfBounds("Chunk index out-of-bounds".into())); 61 | } 62 | 63 | self.index = index; 64 | 65 | if index == 0 || index == 1 { 66 | self.raw_iter.seek_to_first(); 67 | } else { 68 | let preceding_key = self.chunk_boundaries.get(index - 2).unwrap(); 69 | self.raw_iter.seek(preceding_key); 70 | self.raw_iter.next(); 71 | } 72 | 73 | self.next_chunk() 74 | } 75 | 76 | /// Returns the total number of chunks for the underlying Merk tree. 77 | #[allow(clippy::len_without_is_empty)] 78 | pub fn len(&self) -> usize { 79 | let boundaries_len = self.chunk_boundaries.len(); 80 | if boundaries_len == 0 { 81 | 1 82 | } else { 83 | boundaries_len + 2 84 | } 85 | } 86 | 87 | /// Gets the next chunk based on the `ChunkProducer`'s internal index state. 88 | /// This is mostly useful for letting `ChunkIter` yield the chunks in order, 89 | /// optimizing throughput compared to random access. 90 | fn next_chunk(&mut self) -> Result> { 91 | if self.index == 0 { 92 | if self.trunk.is_empty() { 93 | return Err(Error::Fetch( 94 | "Attempted to fetch chunk on empty tree".into(), 95 | )); 96 | } 97 | self.index += 1; 98 | return Ok(self.trunk.encode()?); 99 | } 100 | 101 | assert!(self.index < self.len(), "Called next_chunk after end"); 102 | 103 | let end_key = self.chunk_boundaries.get(self.index - 1); 104 | let end_key_slice = end_key.as_ref().map(|k| k.as_slice()); 105 | 106 | self.index += 1; 107 | 108 | let chunk = get_next_chunk(&mut self.raw_iter, end_key_slice)?; 109 | Ok(chunk.encode()?) 110 | } 111 | } 112 | 113 | impl<'a> IntoIterator for ChunkProducer<'a> { 114 | type IntoIter = ChunkIter<'a>; 115 | type Item = as Iterator>::Item; 116 | 117 | fn into_iter(self) -> Self::IntoIter { 118 | ChunkIter(self) 119 | } 120 | } 121 | 122 | /// A `ChunkIter` iterates through all the chunks for the underlying `Merk` 123 | /// instance in order (the first chunk is the "trunk" chunk). Yields `None` 124 | /// after all chunks have been yielded. 125 | pub struct ChunkIter<'a>(ChunkProducer<'a>); 126 | 127 | impl<'a> Iterator for ChunkIter<'a> { 128 | type Item = Result>; 129 | 130 | fn size_hint(&self) -> (usize, Option) { 131 | (self.0.len(), Some(self.0.len())) 132 | } 133 | 134 | fn next(&mut self) -> Option { 135 | if self.0.index >= self.0.len() { 136 | None 137 | } else { 138 | Some(self.0.next_chunk()) 139 | } 140 | } 141 | } 142 | 143 | impl Merk { 144 | /// Creates a `ChunkProducer` which can return chunk proofs for replicating 145 | /// the entire Merk tree. 146 | pub fn chunks(&self) -> Result { 147 | ChunkProducer::new(self) 148 | } 149 | } 150 | 151 | #[cfg(test)] 152 | mod tests { 153 | use super::*; 154 | use crate::{ 155 | proofs::{ 156 | chunk::{verify_leaf, verify_trunk}, 157 | Decoder, 158 | }, 159 | test_utils::*, 160 | }; 161 | 162 | #[test] 163 | fn len_small() { 164 | let mut merk = TempMerk::new().unwrap(); 165 | let batch = make_batch_seq(1..256); 166 | merk.apply(batch.as_slice(), &[]).unwrap(); 167 | 168 | let chunks = merk.chunks().unwrap(); 169 | assert_eq!(chunks.len(), 1); 170 | assert_eq!(chunks.into_iter().size_hint().0, 1); 171 | } 172 | 173 | #[test] 174 | fn len_big() { 175 | let mut merk = TempMerk::new().unwrap(); 176 | let batch = make_batch_seq(1..10_000); 177 | merk.apply(batch.as_slice(), &[]).unwrap(); 178 | 179 | let chunks = merk.chunks().unwrap(); 180 | assert_eq!(chunks.len(), 129); 181 | assert_eq!(chunks.into_iter().size_hint().0, 129); 182 | } 183 | 184 | #[test] 185 | fn generate_and_verify_chunks() -> Result<()> { 186 | let mut merk = TempMerk::new().unwrap(); 187 | let batch = make_batch_seq(1..10_000); 188 | merk.apply(batch.as_slice(), &[]).unwrap(); 189 | 190 | let mut chunks = merk.chunks().unwrap().into_iter().map(Result::unwrap); 191 | 192 | let chunk = chunks.next().unwrap(); 193 | let ops = Decoder::new(chunk.as_slice()); 194 | let (trunk, height) = verify_trunk(ops).unwrap(); 195 | assert_eq!(height, 14); 196 | assert_eq!(trunk.hash()?, merk.root_hash()); 197 | 198 | assert_eq!(trunk.layer(7).count(), 128); 199 | 200 | for (chunk, node) in chunks.zip(trunk.layer(height / 2)) { 201 | let ops = Decoder::new(chunk.as_slice()); 202 | verify_leaf(ops, node.hash()?).unwrap(); 203 | } 204 | Ok(()) 205 | } 206 | 207 | #[test] 208 | fn chunks_from_reopen() { 209 | let time = std::time::SystemTime::now() 210 | .duration_since(std::time::SystemTime::UNIX_EPOCH) 211 | .unwrap() 212 | .as_nanos(); 213 | let path = format!("chunks_from_reopen_{time}.db"); 214 | 215 | let original_chunks = { 216 | let mut merk = Merk::open(&path).unwrap(); 217 | let batch = make_batch_seq(1..10); 218 | merk.apply(batch.as_slice(), &[]).unwrap(); 219 | 220 | merk.chunks() 221 | .unwrap() 222 | .into_iter() 223 | .map(Result::unwrap) 224 | .collect::>() 225 | .into_iter() 226 | }; 227 | 228 | let merk = TempMerk::open(path).unwrap(); 229 | let reopen_chunks = merk.chunks().unwrap().into_iter().map(Result::unwrap); 230 | 231 | for (original, checkpoint) in original_chunks.zip(reopen_chunks) { 232 | assert_eq!(original.len(), checkpoint.len()); 233 | } 234 | } 235 | 236 | #[test] 237 | fn chunks_from_checkpoint() { 238 | let mut merk = TempMerk::new().unwrap(); 239 | let batch = make_batch_seq(1..10); 240 | merk.apply(batch.as_slice(), &[]).unwrap(); 241 | 242 | let path: std::path::PathBuf = "generate_and_verify_chunks_from_checkpoint.db".into(); 243 | if path.exists() { 244 | std::fs::remove_dir_all(&path).unwrap(); 245 | } 246 | let checkpoint = merk.checkpoint(&path).unwrap(); 247 | 248 | let original_chunks = merk.chunks().unwrap().into_iter().map(Result::unwrap); 249 | let checkpoint_chunks = checkpoint.chunks().unwrap().into_iter().map(Result::unwrap); 250 | 251 | for (original, checkpoint) in original_chunks.zip(checkpoint_chunks) { 252 | assert_eq!(original.len(), checkpoint.len()); 253 | } 254 | 255 | std::fs::remove_dir_all(&path).unwrap(); 256 | } 257 | 258 | #[test] 259 | fn random_access_chunks() { 260 | let mut merk = TempMerk::new().unwrap(); 261 | let batch = make_batch_seq(1..111); 262 | merk.apply(batch.as_slice(), &[]).unwrap(); 263 | 264 | let chunks = merk 265 | .chunks() 266 | .unwrap() 267 | .into_iter() 268 | .map(Result::unwrap) 269 | .collect::>(); 270 | 271 | let mut producer = merk.chunks().unwrap(); 272 | for i in 0..chunks.len() * 2 { 273 | let index = i % chunks.len(); 274 | assert_eq!(producer.chunk(index).unwrap(), chunks[index]); 275 | } 276 | } 277 | 278 | #[test] 279 | #[should_panic(expected = "Attempted to fetch chunk on empty tree")] 280 | fn test_chunk_empty() { 281 | let merk = TempMerk::new().unwrap(); 282 | 283 | let _chunks = merk 284 | .chunks() 285 | .unwrap() 286 | .into_iter() 287 | .map(Result::unwrap) 288 | .collect::>(); 289 | } 290 | 291 | #[test] 292 | #[should_panic(expected = "Chunk index out-of-bounds")] 293 | fn test_chunk_index_oob() { 294 | let mut merk = TempMerk::new().unwrap(); 295 | let batch = make_batch_seq(1..42); 296 | merk.apply(batch.as_slice(), &[]).unwrap(); 297 | 298 | let mut producer = merk.chunks().unwrap(); 299 | let _chunk = producer.chunk(50000).unwrap(); 300 | } 301 | 302 | #[test] 303 | fn test_chunk_index_gt_1_access() { 304 | let mut merk = TempMerk::new().unwrap(); 305 | let batch = make_batch_seq(1..513); 306 | merk.apply(batch.as_slice(), &[]).unwrap(); 307 | 308 | let mut producer = merk.chunks().unwrap(); 309 | println!("length: {}", producer.len()); 310 | let chunk = producer.chunk(2).unwrap(); 311 | assert_eq!( 312 | chunk, 313 | vec![ 314 | 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 18, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 315 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 316 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 317 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 318 | 123, 123, 123, 123, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 19, 0, 60, 123, 123, 123, 123, 319 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 320 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 321 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 322 | 123, 123, 123, 123, 123, 123, 123, 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 20, 0, 323 | 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 324 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 325 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 326 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 17, 3, 0, 8, 0, 0, 0, 327 | 0, 0, 0, 0, 21, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 328 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 329 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 330 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 16, 331 | 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 22, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 332 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 333 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 334 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 335 | 123, 123, 123, 123, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 23, 0, 60, 123, 123, 123, 123, 336 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 337 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 338 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 339 | 123, 123, 123, 123, 123, 123, 123, 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 24, 0, 340 | 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 341 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 342 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 343 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 17, 17, 3, 0, 8, 0, 0, 344 | 0, 0, 0, 0, 0, 25, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 345 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 346 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 347 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 348 | 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 26, 0, 60, 123, 123, 123, 123, 123, 123, 349 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 350 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 351 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 352 | 123, 123, 123, 123, 123, 123, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 27, 0, 60, 123, 123, 353 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 354 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 355 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 356 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 357 | 28, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 358 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 359 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 360 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 17, 3, 0, 8, 361 | 0, 0, 0, 0, 0, 0, 0, 29, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 362 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 363 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 364 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 365 | 123, 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 30, 0, 60, 123, 123, 123, 123, 123, 366 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 367 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 368 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 369 | 123, 123, 123, 123, 123, 123, 123, 3, 0, 8, 0, 0, 0, 0, 0, 0, 0, 31, 0, 60, 123, 370 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 371 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 372 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 373 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 16, 3, 0, 8, 0, 0, 0, 0, 0, 374 | 0, 0, 32, 0, 60, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 375 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 376 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 377 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 17, 17, 378 | 17 379 | ] 380 | ); 381 | } 382 | 383 | #[test] 384 | #[should_panic(expected = "Called next_chunk after end")] 385 | fn test_next_chunk_index_oob() { 386 | let mut merk = TempMerk::new().unwrap(); 387 | let batch = make_batch_seq(1..42); 388 | merk.apply(batch.as_slice(), &[]).unwrap(); 389 | 390 | let mut producer = merk.chunks().unwrap(); 391 | let _chunk1 = producer.next_chunk(); 392 | let _chunk2 = producer.next_chunk(); 393 | } 394 | } 395 | -------------------------------------------------------------------------------- /src/merk/restore.rs: -------------------------------------------------------------------------------- 1 | //! Provides `Restorer`, which can create a replica of a Merk instance by 2 | //! receiving chunk proofs. 3 | 4 | use super::Merk; 5 | use crate::{ 6 | merk::MerkSource, 7 | proofs::{ 8 | chunk::{verify_leaf, verify_trunk, MIN_TRUNK_HEIGHT}, 9 | tree::{Child, Tree as ProofTree}, 10 | Decoder, Node, 11 | }, 12 | tree::{Link, RefWalker, Tree}, 13 | Error, Hash, Result, 14 | }; 15 | use rocksdb::WriteBatch; 16 | use std::iter::Peekable; 17 | use std::path::Path; 18 | 19 | /// A `Restorer` handles decoding, verifying, and storing chunk proofs to 20 | /// replicate an entire Merk tree. It expects the chunks to be processed in 21 | /// order, retrying the last chunk if verification fails. 22 | pub struct Restorer { 23 | leaf_hashes: Option>>, 24 | parent_keys: Option>>>, 25 | trunk_height: Option, 26 | merk: Merk, 27 | expected_root_hash: Hash, 28 | stated_length: usize, 29 | } 30 | 31 | impl Restorer { 32 | /// Creates a new `Restorer`, which will initialize a new Merk at the given 33 | /// file path. The first chunk (the "trunk") will be compared against 34 | /// `expected_root_hash`, then each subsequent chunk will be compared 35 | /// against the hashes stored in the trunk, so that the restore process will 36 | /// never allow malicious peers to send more than a single invalid chunk. 37 | /// 38 | /// The `stated_length` should be the number of chunks stated by the peer, 39 | /// which will be verified after processing a valid first chunk to make it 40 | /// easier to download chunks from peers without needing to trust this 41 | /// length. 42 | pub fn new>( 43 | db_path: P, 44 | expected_root_hash: Hash, 45 | stated_length: usize, 46 | ) -> Result { 47 | if db_path.as_ref().exists() { 48 | return Err(Error::Path("The given path already exists".into())); 49 | } 50 | 51 | Ok(Self { 52 | expected_root_hash, 53 | stated_length, 54 | trunk_height: None, 55 | merk: Merk::open(db_path)?, 56 | leaf_hashes: None, 57 | parent_keys: None, 58 | }) 59 | } 60 | 61 | /// Verifies a chunk and writes it to the working RocksDB instance. Expects 62 | /// to be called for each chunk in order. Returns the number of remaining 63 | /// chunks. 64 | /// 65 | /// Once there are no remaining chunks to be processed, `finalize` should 66 | /// be called. 67 | pub fn process_chunk(&mut self, chunk_bytes: &[u8]) -> Result { 68 | let ops = Decoder::new(chunk_bytes); 69 | 70 | match self.leaf_hashes { 71 | None => self.process_trunk(ops), 72 | Some(_) => self.process_leaf(ops), 73 | } 74 | } 75 | 76 | /// Consumes the `Restorer` and returns the newly-created, fully-populated 77 | /// Merk instance. This method will return an error if called before 78 | /// processing all chunks (e.g. `restorer.remaining_chunks()` is not equal 79 | /// to 0). 80 | pub fn finalize(mut self) -> Result { 81 | if self.remaining_chunks().is_none() || self.remaining_chunks().unwrap() != 0 { 82 | return Err(Error::ChunkProcessing( 83 | "Called finalize before all chunks were processed".into(), 84 | )); 85 | } 86 | 87 | if self.trunk_height.unwrap() >= MIN_TRUNK_HEIGHT { 88 | self.rewrite_trunk_child_heights()?; 89 | } 90 | 91 | self.merk.flush()?; 92 | self.merk.load_root()?; 93 | 94 | Ok(self.merk) 95 | } 96 | 97 | /// Returns the number of remaining chunks to be processed. If called before 98 | /// the first chunk is processed, this method will return `None` since we do 99 | /// not yet have enough information to know about the number of chunks. 100 | pub fn remaining_chunks(&self) -> Option { 101 | self.leaf_hashes.as_ref().map(|lh| lh.len()) 102 | } 103 | 104 | /// Writes the data contained in `tree` (extracted from a verified chunk 105 | /// proof) to the RocksDB. 106 | fn write_chunk(&mut self, tree: ProofTree) -> Result<()> { 107 | let mut batch = WriteBatch::default(); 108 | 109 | tree.visit_refs(&mut |proof_node| { 110 | let (key, mut node) = match &proof_node.node { 111 | // TODO: encode tree node without cloning key/value 112 | Node::KV(key, value) => match Tree::new(key.clone(), value.clone()) { 113 | Ok(node) => (key, node), 114 | Err(_) => return, 115 | }, 116 | _ => return, 117 | }; 118 | 119 | *node.slot_mut(true) = proof_node.left.as_ref().map(Child::as_link); 120 | *node.slot_mut(false) = proof_node.right.as_ref().map(Child::as_link); 121 | 122 | let bytes = node.encode(); 123 | batch.put(key, bytes); 124 | }); 125 | 126 | self.merk.write(batch) 127 | } 128 | 129 | /// Verifies the trunk then writes its data to the RocksDB. 130 | /// 131 | /// The trunk contains a height proof which lets us verify the total number 132 | /// of expected chunks is the same as `stated_length` as passed into 133 | /// `Restorer::new()`. We also verify the expected root hash at this step. 134 | fn process_trunk(&mut self, ops: Decoder) -> Result { 135 | let (trunk, height) = verify_trunk(ops)?; 136 | 137 | if trunk.hash()? != self.expected_root_hash { 138 | return Err(Error::HashMismatch(self.expected_root_hash, trunk.hash()?)); 139 | } 140 | 141 | let root_key = trunk.key().to_vec(); 142 | 143 | let trunk_height = height / 2; 144 | self.trunk_height = Some(trunk_height); 145 | 146 | let chunks_remaining = if trunk_height >= MIN_TRUNK_HEIGHT { 147 | let leaf_hashes = trunk 148 | .layer(trunk_height) 149 | .map(|node| node.hash()) 150 | .collect::>>()? 151 | .into_iter() 152 | .peekable(); 153 | self.leaf_hashes = Some(leaf_hashes); 154 | 155 | let parent_keys = trunk 156 | .layer(trunk_height - 1) 157 | .map(|node| node.key().to_vec()) 158 | .collect::>>() 159 | .into_iter() 160 | .peekable(); 161 | self.parent_keys = Some(parent_keys); 162 | assert_eq!( 163 | self.parent_keys.as_ref().unwrap().len(), 164 | self.leaf_hashes.as_ref().unwrap().len() / 2 165 | ); 166 | 167 | let chunks_remaining = (2_usize).pow(trunk_height as u32); 168 | assert_eq!(self.remaining_chunks_unchecked(), chunks_remaining); 169 | chunks_remaining 170 | } else { 171 | self.leaf_hashes = Some(vec![].into_iter().peekable()); 172 | self.parent_keys = Some(vec![].into_iter().peekable()); 173 | 0 174 | }; 175 | 176 | if self.stated_length != chunks_remaining + 1 { 177 | return Err(Error::ChunkProcessing( 178 | "Stated length does not match calculated number of chunks".into(), 179 | )); 180 | } 181 | 182 | // note that these writes don't happen atomically, which is fine here 183 | // because if anything fails during the restore process we will just 184 | // scrap the whole restore and start over 185 | self.write_chunk(trunk)?; 186 | self.merk.set_root_key(root_key)?; 187 | 188 | Ok(chunks_remaining) 189 | } 190 | 191 | /// Verifies a leaf chunk then writes it to the RocksDB. This needs to be 192 | /// called in order, retrying the last chunk for any failed verifications. 193 | fn process_leaf(&mut self, ops: Decoder) -> Result { 194 | let leaf_hashes = self.leaf_hashes.as_mut().unwrap(); 195 | let leaf_hash = leaf_hashes 196 | .peek() 197 | .expect("Received more chunks than expected"); 198 | 199 | let leaf = verify_leaf(ops, *leaf_hash)?; 200 | self.rewrite_parent_link(&leaf)?; 201 | self.write_chunk(leaf)?; 202 | 203 | let leaf_hashes = self.leaf_hashes.as_mut().unwrap(); 204 | leaf_hashes.next(); 205 | 206 | Ok(self.remaining_chunks_unchecked()) 207 | } 208 | 209 | /// The parent of the root node of the leaf does not know the key of its 210 | /// children when it is first written. Now that we have verified this leaf, 211 | /// we can write the key into the parent node's entry. Note that this does 212 | /// not need to recalcuate hashes since it already had the child hash. 213 | fn rewrite_parent_link(&mut self, leaf: &ProofTree) -> Result<()> { 214 | let parent_keys = self.parent_keys.as_mut().unwrap(); 215 | let parent_key = parent_keys.peek().unwrap().clone(); 216 | let mut parent = self 217 | .merk 218 | .fetch_node(parent_key.as_slice())? 219 | .expect("Could not find parent of leaf chunk"); 220 | 221 | let is_left_child = self.remaining_chunks_unchecked() % 2 == 0; 222 | if let Some(Link::Reference { ref mut key, .. }) = parent.link_mut(is_left_child) { 223 | *key = leaf.key().to_vec(); 224 | } else { 225 | panic!("Expected parent links to be type Link::Reference"); 226 | }; 227 | 228 | let parent_bytes = parent.encode(); 229 | self.merk.db.put(parent_key, parent_bytes)?; 230 | 231 | if !is_left_child { 232 | let parent_keys = self.parent_keys.as_mut().unwrap(); 233 | parent_keys.next(); 234 | } 235 | 236 | Ok(()) 237 | } 238 | 239 | fn rewrite_trunk_child_heights(&mut self) -> Result<()> { 240 | fn recurse( 241 | mut node: RefWalker, 242 | remaining_depth: usize, 243 | batch: &mut WriteBatch, 244 | ) -> Result<(u8, u8)> { 245 | if remaining_depth == 0 { 246 | return Ok(node.tree().child_heights()); 247 | } 248 | 249 | let mut cloned_node = 250 | Tree::decode(node.tree().key().to_vec(), node.tree().encode().as_slice()); 251 | 252 | let left_child = node.walk(true)?.unwrap(); 253 | let left_child_heights = recurse(left_child, remaining_depth - 1, batch)?; 254 | let left_height = left_child_heights.0.max(left_child_heights.1) + 1; 255 | *cloned_node.link_mut(true).unwrap().child_heights_mut() = left_child_heights; 256 | 257 | let right_child = node.walk(false)?.unwrap(); 258 | let right_child_heights = recurse(right_child, remaining_depth - 1, batch)?; 259 | let right_height = right_child_heights.0.max(right_child_heights.1) + 1; 260 | *cloned_node.link_mut(false).unwrap().child_heights_mut() = right_child_heights; 261 | 262 | let bytes = cloned_node.encode(); 263 | batch.put(node.tree().key(), bytes); 264 | 265 | Ok((left_height, right_height)) 266 | } 267 | 268 | self.merk.flush()?; 269 | self.merk.load_root()?; 270 | 271 | let mut batch = WriteBatch::default(); 272 | 273 | let depth = self.trunk_height.unwrap(); 274 | self.merk.use_tree_mut(|maybe_tree| { 275 | let tree = maybe_tree.unwrap(); 276 | let walker = RefWalker::new(tree, self.merk.source()); 277 | recurse(walker, depth, &mut batch) 278 | })?; 279 | 280 | self.merk.write(batch)?; 281 | 282 | Ok(()) 283 | } 284 | 285 | /// Returns the number of remaining chunks to be processed. This method will 286 | /// panic if called before processing the first chunk (since that chunk 287 | /// gives us the information to know how many chunks to expect). 288 | pub fn remaining_chunks_unchecked(&self) -> usize { 289 | self.leaf_hashes.as_ref().unwrap().len() 290 | } 291 | } 292 | 293 | impl Merk { 294 | /// Creates a new `Restorer`, which can be used to verify chunk proofs to 295 | /// replicate an entire Merk tree. A new Merk instance will be initialized 296 | /// by creating a RocksDB at `path`. 297 | /// 298 | /// The restoration process will verify integrity by checking that the 299 | /// incoming chunk proofs match `expected_root_hash`. The `stated_length` 300 | /// should be the number of chunks as stated by peers, which will also be 301 | /// verified during the restoration process. 302 | pub fn restore>( 303 | path: P, 304 | expected_root_hash: Hash, 305 | stated_length: usize, 306 | ) -> Result { 307 | Restorer::new(path, expected_root_hash, stated_length) 308 | } 309 | } 310 | 311 | impl ProofTree { 312 | fn child_heights(&self) -> (u8, u8) { 313 | ( 314 | self.left.as_ref().map_or(0, |c| c.tree.height as u8), 315 | self.right.as_ref().map_or(0, |c| c.tree.height as u8), 316 | ) 317 | } 318 | } 319 | 320 | impl Child { 321 | fn as_link(&self) -> Link { 322 | let key = match &self.tree.node { 323 | Node::KV(key, _) => key.as_slice(), 324 | // for the connection between the trunk and leaf chunks, we don't 325 | // have the child key so we must first write in an empty one. once 326 | // the leaf gets verified, we can write in this key to its parent 327 | _ => &[], 328 | }; 329 | 330 | Link::Reference { 331 | hash: self.hash, 332 | child_heights: self.tree.child_heights(), 333 | key: key.to_vec(), 334 | } 335 | } 336 | } 337 | 338 | #[cfg(test)] 339 | mod tests { 340 | use super::*; 341 | use crate::test_utils::*; 342 | use crate::tree::{Batch, Op}; 343 | use std::path::PathBuf; 344 | 345 | fn restore_test(batches: &[&Batch], expected_nodes: usize) { 346 | let mut original = TempMerk::new().unwrap(); 347 | for batch in batches { 348 | original.apply(batch, &[]).unwrap(); 349 | } 350 | original.flush().unwrap(); 351 | 352 | let chunks = original.chunks().unwrap(); 353 | 354 | let path: PathBuf = std::thread::current().name().unwrap().into(); 355 | if path.exists() { 356 | std::fs::remove_dir_all(&path).unwrap(); 357 | } 358 | 359 | let mut restorer = Merk::restore(&path, original.root_hash(), chunks.len()).unwrap(); 360 | 361 | assert_eq!(restorer.remaining_chunks(), None); 362 | 363 | let mut expected_remaining = chunks.len(); 364 | for chunk in chunks { 365 | let chunk = chunk.unwrap(); 366 | let remaining = restorer.process_chunk(chunk.as_slice()).unwrap(); 367 | 368 | expected_remaining -= 1; 369 | assert_eq!(remaining, expected_remaining); 370 | assert_eq!(restorer.remaining_chunks().unwrap(), expected_remaining); 371 | } 372 | assert_eq!(expected_remaining, 0); 373 | 374 | let restored = restorer.finalize().unwrap(); 375 | assert_eq!(restored.root_hash(), original.root_hash()); 376 | assert_raw_db_entries_eq(&restored, &original, expected_nodes); 377 | 378 | std::fs::remove_dir_all(&path).unwrap(); 379 | } 380 | 381 | #[test] 382 | fn restore_10000() { 383 | restore_test(&[&make_batch_seq(0..10_000)], 10_000); 384 | } 385 | 386 | #[test] 387 | fn restore_3() { 388 | restore_test(&[&make_batch_seq(0..3)], 3); 389 | } 390 | 391 | #[test] 392 | fn restore_2_left_heavy() { 393 | restore_test( 394 | &[&[(vec![0], Op::Put(vec![]))], &[(vec![1], Op::Put(vec![]))]], 395 | 2, 396 | ); 397 | } 398 | 399 | #[test] 400 | fn restore_2_right_heavy() { 401 | restore_test( 402 | &[&[(vec![1], Op::Put(vec![]))], &[(vec![0], Op::Put(vec![]))]], 403 | 2, 404 | ); 405 | } 406 | 407 | #[test] 408 | fn restore_1() { 409 | restore_test(&[&make_batch_seq(0..1)], 1); 410 | } 411 | 412 | fn assert_raw_db_entries_eq(restored: &Merk, original: &Merk, length: usize) { 413 | let mut original_entries = original.raw_iter(); 414 | let mut restored_entries = restored.raw_iter(); 415 | original_entries.seek_to_first(); 416 | restored_entries.seek_to_first(); 417 | 418 | let mut i = 0; 419 | loop { 420 | assert_eq!(restored_entries.valid(), original_entries.valid()); 421 | if !restored_entries.valid() { 422 | break; 423 | } 424 | 425 | assert_eq!(restored_entries.key(), original_entries.key()); 426 | assert_eq!(restored_entries.value(), original_entries.value()); 427 | 428 | restored_entries.next(); 429 | original_entries.next(); 430 | 431 | i += 1; 432 | } 433 | 434 | assert_eq!(i, length); 435 | } 436 | } 437 | -------------------------------------------------------------------------------- /src/merk/snapshot.rs: -------------------------------------------------------------------------------- 1 | //! In-memory snapshots of database state. 2 | //! 3 | //! Snapshots are read-only views of the database state at a particular point in 4 | //! time. This can be useful for retaining recent versions of history which can 5 | //! be queried against. Merk snapshots are backed by the similar RocksDB 6 | //! snapshot, but with the added ability to create proofs. 7 | 8 | use std::cell::Cell; 9 | 10 | use crate::{ 11 | proofs::query::QueryItem, 12 | tree::{Fetch, RefWalker, Tree, NULL_HASH}, 13 | Hash, Result, 14 | }; 15 | 16 | /// A read-only view of the database state at a particular point in time. 17 | /// 18 | /// `Snapshot`s are cheap to create since they are just a handle and don't copy 19 | /// any data - they instead just prevent the underlying replaced data from being 20 | /// compacted in RocksDB until they are dropped. They are only held in memory, 21 | /// and will not be persisted after the process exits. 22 | pub struct Snapshot<'a> { 23 | /// The underlying RocksDB snapshot. 24 | ss: Option>, 25 | /// The Merk tree at the time the snapshot was created. 26 | tree: Cell>, 27 | /// Whether the underlying RocksDB snapshot should be dropped when the 28 | /// `Snapshot` is dropped. 29 | should_drop_ss: bool, 30 | } 31 | 32 | impl<'a> Snapshot<'a> { 33 | /// Creates a new `Snapshot` from a RocksDB snapshot and a Merk tree. 34 | /// 35 | /// The RocksDB snapshot will be dropped when the [Snapshot] is dropped. 36 | pub fn new(db: rocksdb::Snapshot<'a>, tree: Option) -> Self { 37 | Snapshot { 38 | ss: Some(db), 39 | tree: Cell::new(tree), 40 | should_drop_ss: true, 41 | } 42 | } 43 | 44 | /// Converts the [Snapshot] into a [StaticSnapshot], an alternative which 45 | /// has easier (but more dangerous) lifetime requirements. 46 | pub fn staticize(mut self) -> StaticSnapshot { 47 | let ss: RocksDBSnapshot = unsafe { std::mem::transmute(self.ss.take().unwrap()) }; 48 | StaticSnapshot { 49 | tree: Cell::new(self.tree.take()), 50 | inner: ss.inner, 51 | should_drop: false, 52 | } 53 | } 54 | 55 | /// Gets the value associated with the given key, from the time the snapshot 56 | /// was created. 57 | pub fn get(&self, key: &[u8]) -> Result>> { 58 | self.use_tree(|maybe_tree| { 59 | maybe_tree 60 | .and_then(|tree| super::get(tree, self.source(), key).transpose()) 61 | .transpose() 62 | }) 63 | } 64 | 65 | /// Gets the root hash of the tree at the time the snapshot was created. 66 | pub fn root_hash(&self) -> Hash { 67 | self.use_tree(|tree| tree.map_or(NULL_HASH, |tree| tree.hash())) 68 | } 69 | 70 | /// Proves the given query against the tree at the time the snapshot was 71 | /// created. 72 | pub fn prove(&self, query: I) -> Result> 73 | where 74 | Q: Into, 75 | I: IntoIterator, 76 | { 77 | self.use_tree_mut(move |maybe_tree| super::prove(maybe_tree, self.source(), query)) 78 | } 79 | 80 | /// Walks the tree at the time the snapshot was created, fetching the child 81 | /// node from the backing store if necessary. 82 | pub fn walk(&self, f: impl FnOnce(Option>) -> T) -> T { 83 | let mut tree = self.tree.take(); 84 | let maybe_walker = tree 85 | .as_mut() 86 | .map(|tree| RefWalker::new(tree, self.source())); 87 | let res = f(maybe_walker); 88 | self.tree.set(tree); 89 | res 90 | } 91 | 92 | /// Returns an iterator over the keys and values in the backing store from 93 | /// the time the snapshot was created. 94 | pub fn raw_iter(&self) -> rocksdb::DBRawIterator { 95 | self.ss.as_ref().unwrap().raw_iterator() 96 | } 97 | 98 | /// A data source which can be used to fetch values from the backing store, 99 | /// from the time the snapshot was created. 100 | fn source(&self) -> SnapshotSource { 101 | SnapshotSource(self.ss.as_ref().unwrap()) 102 | } 103 | 104 | /// Uses the tree, and then puts it back. 105 | fn use_tree(&self, f: impl FnOnce(Option<&Tree>) -> T) -> T { 106 | let tree = self.tree.take(); 107 | let res = f(tree.as_ref()); 108 | self.tree.set(tree); 109 | res 110 | } 111 | 112 | /// Uses the tree mutably, and then puts it back. 113 | fn use_tree_mut(&self, f: impl FnOnce(Option<&mut Tree>) -> T) -> T { 114 | let mut tree = self.tree.take(); 115 | let res = f(tree.as_mut()); 116 | self.tree.set(tree); 117 | res 118 | } 119 | } 120 | 121 | impl<'a> Drop for Snapshot<'a> { 122 | fn drop(&mut self) { 123 | if !self.should_drop_ss { 124 | std::mem::forget(self.ss.take()); 125 | } 126 | } 127 | } 128 | 129 | /// A data source which can be used to fetch values from the backing store, from 130 | /// the time the snapshot was created. 131 | /// 132 | /// This implements [Fetch] and should be used with a type such as [RefWalker]. 133 | #[derive(Clone)] 134 | pub struct SnapshotSource<'a>(&'a rocksdb::Snapshot<'a>); 135 | 136 | impl<'a> Fetch for SnapshotSource<'a> { 137 | fn fetch_by_key(&self, key: &[u8]) -> Result> { 138 | Ok(self 139 | .0 140 | .get(key)? 141 | .map(|bytes| Tree::decode(key.to_vec(), &bytes))) 142 | } 143 | } 144 | 145 | /// A read-only view of the database state at a particular point in time, but 146 | /// with an internal raw pointer to allow for manual lifetime management. 147 | /// 148 | /// This is useful when you would otherwise want a [Snapshot], but you want to 149 | /// use the database while the snapshot is still alive. This is unsafe because 150 | /// it is the caller's responsibility to ensure that the underlying RocksDB 151 | /// snapshot outlives the [StaticSnapshot]. 152 | /// 153 | /// By default, the RocksDB snapshot will not be dropped when the 154 | /// [StaticSnapshot] is dropped, resulting in a memory leak. For correct usage, 155 | /// you must call [StaticSnapshot::drop] to ensure the RocksDB snapshot gets 156 | /// dropped when the [StaticSnapshot] is dropped. 157 | pub struct StaticSnapshot { 158 | /// A Merk tree based on the database state at the time the snapshot was 159 | /// created. 160 | tree: Cell>, 161 | /// A raw pointer to the RocksDB snapshot. 162 | inner: *const (), 163 | /// Used to detect whether the `StaticSnapshot` was set to manually drop 164 | /// before its [Drop::drop] implementation was called. 165 | pub should_drop: bool, 166 | } 167 | 168 | /// An equivalent struct to the [rocksdb::Snapshot] struct within the `rocksdb` 169 | /// crate. This is used to access the private fields of the foreign crate's 170 | /// struct by first transmuting. 171 | /// 172 | /// To guarantee that breaking changes in the `rocksdb` crate do not affect the 173 | /// transmutation into this struct, see the 174 | /// [tests::rocksdb_snapshot_struct_format] test. 175 | struct RocksDBSnapshot<'a> { 176 | /// A reference to the associated RocksDB database. 177 | _db: &'a rocksdb::DB, 178 | /// A raw pointer to the snapshot handle. 179 | inner: *const (), 180 | } 181 | 182 | // We need this because we have a raw pointer to a RocksDB snapshot, but we 183 | // know that our usage of it is thread-safe: 184 | // https://github.com/facebook/rocksdb/blob/main/include/rocksdb/snapshot.h#L15-L16 185 | unsafe impl Send for StaticSnapshot {} 186 | unsafe impl Sync for StaticSnapshot {} 187 | 188 | impl StaticSnapshot { 189 | /// Converts the [StaticSnapshot] to a [Snapshot] by re-associating with the 190 | /// database it was originally created from. 191 | /// 192 | /// # Safety 193 | /// This will cause undefined behavior if a database other than the one 194 | /// originally used to create the snapshot is passed as an argument. 195 | /// 196 | /// This will also cause a memory leak if the underlying RocksDB snapshot is 197 | /// not dropped by calling [StaticSnapshot::drop]. Unlike most uses of 198 | /// [Snapshot], the RocksDB snapshot will not be dropped when the 199 | /// [Snapshot] returned by this method is dropped. 200 | pub unsafe fn with_db<'a>(&self, db: &'a rocksdb::DB) -> Snapshot<'a> { 201 | let db_ss = RocksDBSnapshot { 202 | _db: db, 203 | inner: self.inner, 204 | }; 205 | let db_ss: rocksdb::Snapshot<'a> = std::mem::transmute(db_ss); 206 | 207 | Snapshot { 208 | ss: Some(db_ss), 209 | tree: self.clone_tree(), 210 | should_drop_ss: false, 211 | } 212 | } 213 | 214 | /// Drops the [StaticSnapshot] and the underlying RocksDB snapshot. 215 | /// 216 | /// # Safety 217 | /// This function is unsafe because it results in the RocksDB snapshot being 218 | /// dropped, which could lead to use-after-free bugs if there are still 219 | /// references to the snapshot in other [Snapshot] or [StaticSnapshot] 220 | /// instances. The caller must be sure this is the last remaining reference 221 | /// before calling this method. 222 | pub unsafe fn drop(mut self, db: &rocksdb::DB) { 223 | let mut ss = self.with_db(db); 224 | ss.should_drop_ss = true; 225 | self.should_drop = true; 226 | // the snapshot drop implementation is now called, which includes 227 | // dropping the RocksDB snapshot 228 | } 229 | 230 | /// Clones the root node of the Merk tree into a new [Tree]. 231 | fn clone_tree(&self) -> Cell> { 232 | let tree = self.tree.take().unwrap(); 233 | let tree_clone = Cell::new(Some(Tree::decode( 234 | tree.key().to_vec(), 235 | tree.encode().as_slice(), 236 | ))); 237 | self.tree.set(Some(tree)); 238 | tree_clone 239 | } 240 | } 241 | 242 | impl Drop for StaticSnapshot { 243 | fn drop(&mut self) { 244 | if !self.should_drop { 245 | log::debug!("StaticSnapshot must be manually dropped"); 246 | } 247 | } 248 | } 249 | 250 | impl Clone for StaticSnapshot { 251 | fn clone(&self) -> Self { 252 | Self { 253 | tree: self.clone_tree(), 254 | inner: self.inner, 255 | should_drop: self.should_drop, 256 | } 257 | } 258 | } 259 | 260 | #[cfg(test)] 261 | mod tests { 262 | use std::mem::transmute; 263 | 264 | use super::RocksDBSnapshot; 265 | use crate::test_utils::TempMerk; 266 | 267 | #[test] 268 | fn rocksdb_snapshot_struct_format() { 269 | assert_eq!(std::mem::size_of::(), 16); 270 | 271 | let merk = TempMerk::new().unwrap(); 272 | let exptected_db_ptr = merk.db() as *const _; 273 | 274 | let ss = merk.db().snapshot(); 275 | let ss: RocksDBSnapshot = unsafe { transmute(ss) }; 276 | let db_ptr = ss._db as *const _; 277 | 278 | assert_eq!(exptected_db_ptr, db_ptr); 279 | } 280 | } 281 | -------------------------------------------------------------------------------- /src/owner.rs: -------------------------------------------------------------------------------- 1 | use std::ops::{Deref, DerefMut}; 2 | 3 | /// A container type which holds a value that may be temporarily owned by a 4 | /// consumer. 5 | pub struct Owner { 6 | inner: Option, 7 | } 8 | 9 | impl Owner { 10 | /// Creates a new `Owner` which holds the given value. 11 | pub fn new(value: T) -> Owner { 12 | Owner { inner: Some(value) } 13 | } 14 | 15 | /// Takes temporary ownership of the contained value by passing it to `f`. 16 | /// The function must return a value of the same type (the same value, or a 17 | /// new value to take its place). 18 | /// 19 | /// # Example 20 | /// ``` 21 | /// # use merk::owner::Owner; 22 | /// # struct SomeType(); 23 | /// # impl SomeType { 24 | /// # fn method_which_requires_ownership(self) -> SomeType { self } 25 | /// # } 26 | /// # 27 | /// let mut owner = Owner::new(SomeType()); 28 | /// owner.own(|value| { 29 | /// value.method_which_requires_ownership(); 30 | /// SomeType() // now give back a value of the same type 31 | /// }); 32 | /// ``` 33 | pub fn own T>(&mut self, f: F) { 34 | let old_value = unwrap(self.inner.take()); 35 | let new_value = f(old_value); 36 | self.inner = Some(new_value); 37 | } 38 | 39 | /// Takes temporary ownership of the contained value by passing it to `f`. 40 | /// The function must return a value of the same type (the same value, or a 41 | /// new value to take its place). 42 | /// 43 | /// Like `own`, but uses a tuple return type which allows specifying a value 44 | /// to return from the call to `own_return` for convenience. 45 | /// 46 | /// # Example 47 | /// ``` 48 | /// # use merk::owner::Owner; 49 | /// let mut owner = Owner::new(123); 50 | /// let doubled = owner.own_return(|n| (n, n * 2)); 51 | /// ``` 52 | pub fn own_return(&mut self, f: F) -> R 53 | where 54 | R: Sized, 55 | F: FnOnce(T) -> (T, R), 56 | { 57 | let old_value = unwrap(self.inner.take()); 58 | let (new_value, return_value) = f(old_value); 59 | self.inner = Some(new_value); 60 | return_value 61 | } 62 | 63 | /// Takes temporary ownership of the contained value by passing it to `f`. 64 | /// The function must return a value of the same type (the same value, or a 65 | /// new value to take its place). 66 | /// 67 | /// Like `own`, but with a fallible operation. 68 | /// 69 | /// # Example 70 | /// ``` 71 | /// # use merk::owner::Owner; 72 | /// # use std::convert::TryFrom; 73 | /// let mut owner = Owner::new(123); 74 | /// let converted = owner.own_fallible(|n| u32::try_from(n)); 75 | /// ``` 76 | pub fn own_fallible Result>(&mut self, f: F) -> Result<(), E> { 77 | let old_value = unwrap(self.inner.take()); 78 | let new_value = f(old_value)?; 79 | self.inner = Some(new_value); 80 | Ok(()) 81 | } 82 | 83 | /// Sheds the `Owner` container and returns the value it contained. 84 | pub fn into_inner(mut self) -> T { 85 | unwrap(self.inner.take()) 86 | } 87 | } 88 | 89 | impl Deref for Owner { 90 | type Target = T; 91 | 92 | fn deref(&self) -> &T { 93 | unwrap(self.inner.as_ref()) 94 | } 95 | } 96 | 97 | impl DerefMut for Owner { 98 | fn deref_mut(&mut self) -> &mut T { 99 | unwrap(self.inner.as_mut()) 100 | } 101 | } 102 | 103 | fn unwrap(option: Option) -> T { 104 | match option { 105 | Some(value) => value, 106 | None => unreachable!("value should be Some"), 107 | } 108 | } 109 | 110 | // TODO: unit tests 111 | -------------------------------------------------------------------------------- /src/proofs/chunk.rs: -------------------------------------------------------------------------------- 1 | #[cfg(feature = "full")] 2 | use { 3 | super::tree::{execute, Tree as ProofTree}, 4 | crate::tree::Hash, 5 | crate::tree::Tree, 6 | rocksdb::DBRawIterator, 7 | }; 8 | 9 | use super::{Node, Op}; 10 | use crate::error::{Error, Result}; 11 | use crate::tree::{Fetch, RefWalker}; 12 | 13 | /// The minimum number of layers the trunk will be guaranteed to have before 14 | /// splitting into multiple chunks. 15 | /// 16 | /// If the tree's height is less than double this value, the trunk should be 17 | /// verified as a leaf chunk. 18 | pub const MIN_TRUNK_HEIGHT: usize = 5; 19 | 20 | impl<'a, S> RefWalker<'a, S> 21 | where 22 | S: Fetch + Sized + Send + Clone, 23 | { 24 | /// Generates a trunk proof by traversing the tree. 25 | /// 26 | /// Returns a tuple containing the produced proof, and a boolean indicating 27 | /// whether or not there will be more chunks to follow. If the chunk 28 | /// contains the entire tree, the boolean will be `false`, if the chunk 29 | /// is abdriged and will be connected to leaf chunks, it will be `true`. 30 | pub fn create_trunk_proof(&mut self) -> Result<(Vec, bool)> { 31 | let approx_size = 2usize.pow((self.tree().height() / 2) as u32) * 3; 32 | let mut proof = Vec::with_capacity(approx_size); 33 | 34 | let trunk_height = self.traverse_for_height_proof(&mut proof, 1)?; 35 | 36 | if trunk_height < MIN_TRUNK_HEIGHT { 37 | proof.clear(); 38 | self.traverse_for_trunk(&mut proof, usize::MAX, true)?; 39 | Ok((proof, false)) 40 | } else { 41 | self.traverse_for_trunk(&mut proof, trunk_height, true)?; 42 | Ok((proof, true)) 43 | } 44 | } 45 | 46 | /// Traverses down the left edge of the tree and pushes ops to the proof, to 47 | /// act as a proof of the height of the tree. This is the first step in 48 | /// generating a trunk proof. 49 | fn traverse_for_height_proof(&mut self, proof: &mut Vec, depth: usize) -> Result { 50 | let maybe_left = self.walk(true)?; 51 | let has_left_child = maybe_left.is_some(); 52 | 53 | let trunk_height = if let Some(mut left) = maybe_left { 54 | left.traverse_for_height_proof(proof, depth + 1)? 55 | } else { 56 | depth / 2 57 | }; 58 | 59 | if depth > trunk_height { 60 | proof.push(Op::Push(self.to_kvhash_node())); 61 | 62 | if has_left_child { 63 | proof.push(Op::Parent); 64 | } 65 | 66 | if let Some(right) = self.tree().link(false) { 67 | proof.push(Op::Push(Node::Hash(*right.hash()))); 68 | proof.push(Op::Child); 69 | } 70 | } 71 | 72 | Ok(trunk_height) 73 | } 74 | 75 | /// Traverses down the tree and adds KV push ops for all nodes up to a 76 | /// certain depth. This expects the proof to contain a height proof as 77 | /// generated by `traverse_for_height_proof`. 78 | fn traverse_for_trunk( 79 | &mut self, 80 | proof: &mut Vec, 81 | remaining_depth: usize, 82 | is_leftmost: bool, 83 | ) -> Result<()> { 84 | if remaining_depth == 0 { 85 | // return early if we have reached bottom of trunk 86 | 87 | // for leftmost node, we already have height proof 88 | if is_leftmost { 89 | return Ok(()); 90 | } 91 | 92 | // add this node's hash 93 | proof.push(Op::Push(self.to_hash_node())); 94 | 95 | return Ok(()); 96 | } 97 | 98 | // traverse left 99 | let has_left_child = self.tree().link(true).is_some(); 100 | if has_left_child { 101 | let mut left = self.walk(true)?.unwrap(); 102 | left.traverse_for_trunk(proof, remaining_depth - 1, is_leftmost)?; 103 | } 104 | 105 | // add this node's data 106 | proof.push(Op::Push(self.to_kv_node())); 107 | 108 | if has_left_child { 109 | proof.push(Op::Parent); 110 | } 111 | 112 | // traverse right 113 | if let Some(mut right) = self.walk(false)? { 114 | right.traverse_for_trunk(proof, remaining_depth - 1, false)?; 115 | proof.push(Op::Child); 116 | } 117 | 118 | Ok(()) 119 | } 120 | } 121 | 122 | /// Builds a chunk proof by iterating over values in a RocksDB, ending the chunk 123 | /// when a node with key `end_key` is encountered. 124 | /// 125 | /// Advances the iterator for all nodes in the chunk and the `end_key` (if any). 126 | #[cfg(feature = "full")] 127 | pub(crate) fn get_next_chunk(iter: &mut DBRawIterator, end_key: Option<&[u8]>) -> Result> { 128 | let mut chunk = Vec::with_capacity(512); 129 | let mut stack = Vec::with_capacity(32); 130 | let mut node = Tree::new(vec![], vec![])?; 131 | 132 | while iter.valid() { 133 | let key = iter.key().unwrap(); 134 | 135 | if let Some(end_key) = end_key { 136 | if key == end_key { 137 | break; 138 | } 139 | } 140 | 141 | let encoded_node = iter.value().unwrap(); 142 | Tree::decode_into(&mut node, vec![], encoded_node); 143 | 144 | let kv = Node::KV(key.to_vec(), node.value().to_vec()); 145 | chunk.push(Op::Push(kv)); 146 | 147 | if node.link(true).is_some() { 148 | chunk.push(Op::Parent); 149 | } 150 | 151 | if let Some(child) = node.link(false) { 152 | stack.push(child.key().to_vec()); 153 | } else { 154 | while let Some(top_key) = stack.last() { 155 | if key < top_key.as_slice() { 156 | break; 157 | } 158 | stack.pop(); 159 | chunk.push(Op::Child); 160 | } 161 | } 162 | 163 | iter.next(); 164 | } 165 | 166 | if iter.valid() { 167 | iter.next(); 168 | } 169 | 170 | Ok(chunk) 171 | } 172 | 173 | /// Verifies a leaf chunk proof by executing its operators. Checks that there 174 | /// were no abridged nodes (Hash or KVHash) and the proof hashes to 175 | /// `expected_hash`. 176 | #[cfg(feature = "full")] 177 | pub(crate) fn verify_leaf>>( 178 | ops: I, 179 | expected_hash: Hash, 180 | ) -> Result { 181 | let tree = execute(ops, false, |node| match node { 182 | Node::KV(_, _) => Ok(()), 183 | _ => Err(Error::Tree("Leaf chunks must contain full subtree".into())), 184 | })?; 185 | 186 | if tree.hash()? != expected_hash { 187 | return Err(Error::HashMismatch(expected_hash, tree.hash()?)); 188 | } 189 | 190 | Ok(tree) 191 | } 192 | 193 | /// Verifies a trunk chunk proof by executing its operators. Ensures the 194 | /// resulting tree contains a valid height proof, the trunk is the correct 195 | /// height, and all of its inner nodes are not abridged. Returns the tree and 196 | /// the height given by the height proof. 197 | #[cfg(feature = "full")] 198 | pub(crate) fn verify_trunk>>(ops: I) -> Result<(ProofTree, usize)> { 199 | fn verify_height_proof(tree: &ProofTree) -> Result { 200 | let mut height = 1; 201 | let mut cursor = tree; 202 | while let Some(child) = cursor.child(true) { 203 | if let Node::Hash(_) = child.tree.node { 204 | return Err(Error::UnexpectedNode( 205 | "Expected height proof to only contain KV and KVHash 206 | nodes" 207 | .into(), 208 | )); 209 | } 210 | height += 1; 211 | cursor = &child.tree; 212 | } 213 | Ok(height) 214 | } 215 | 216 | fn verify_completeness(tree: &ProofTree, remaining_depth: usize, leftmost: bool) -> Result<()> { 217 | let recurse = |left, leftmost| { 218 | if let Some(child) = tree.child(left) { 219 | verify_completeness(&child.tree, remaining_depth - 1, left && leftmost)?; 220 | } 221 | Ok(()) 222 | }; 223 | 224 | if remaining_depth > 0 { 225 | match tree.node { 226 | Node::KV(_, _) => {} 227 | _ => { 228 | return Err(Error::UnexpectedNode( 229 | "Expected trunk inner nodes to contain keys and values".into(), 230 | )); 231 | } 232 | } 233 | recurse(true, leftmost)?; 234 | recurse(false, false) 235 | } else if !leftmost { 236 | match tree.node { 237 | Node::Hash(_) => Ok(()), 238 | _ => Err(Error::UnexpectedNode( 239 | "Expected trunk leaves to contain Hash nodes".into(), 240 | )), 241 | } 242 | } else { 243 | match &tree.node { 244 | Node::KVHash(_) => Ok(()), 245 | _ => Err(Error::UnexpectedNode( 246 | "Expected leftmost trunk leaf to contain KVHash node".into(), 247 | )), 248 | } 249 | } 250 | } 251 | 252 | let mut kv_only = true; 253 | let tree = execute(ops, false, |node| { 254 | kv_only &= matches!(node, Node::KV(_, _)); 255 | Ok(()) 256 | })?; 257 | 258 | let height = verify_height_proof(&tree)?; 259 | if height > 64 { 260 | // This is a sanity check to prevent stack overflows in `verify_completeness`, 261 | // but any tree above 64 is probably an error (~3.7e19 nodes). 262 | return Err(Error::Tree("Tree is too large".into())); 263 | } 264 | let trunk_height = height / 2; 265 | 266 | if trunk_height < MIN_TRUNK_HEIGHT { 267 | if !kv_only { 268 | return Err(Error::Tree("Leaf chunks must contain full subtree".into())); 269 | } 270 | } else { 271 | verify_completeness(&tree, trunk_height, true)?; 272 | } 273 | 274 | Ok((tree, height)) 275 | } 276 | 277 | #[cfg(test)] 278 | mod tests { 279 | use super::super::tree::Tree; 280 | use super::*; 281 | use crate::test_utils::*; 282 | use crate::tree::{NoopCommit, PanicSource, Tree as BaseTree}; 283 | use ed::Encode; 284 | 285 | #[derive(Default)] 286 | struct NodeCounts { 287 | hash: usize, 288 | kvhash: usize, 289 | kv: usize, 290 | } 291 | 292 | fn count_node_types(tree: Tree) -> NodeCounts { 293 | let mut counts = NodeCounts::default(); 294 | 295 | tree.visit_nodes(&mut |node| { 296 | match node { 297 | Node::Hash(_) => counts.hash += 1, 298 | Node::KVHash(_) => counts.kvhash += 1, 299 | Node::KV(_, _) => counts.kv += 1, 300 | }; 301 | }); 302 | 303 | counts 304 | } 305 | 306 | #[test] 307 | fn small_trunk_roundtrip() { 308 | let mut tree = make_tree_seq(31); 309 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 310 | 311 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 312 | assert!(!has_more); 313 | 314 | println!("{:?}", &proof); 315 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 316 | 317 | let counts = count_node_types(trunk); 318 | assert_eq!(counts.hash, 0); 319 | assert_eq!(counts.kv, 32); 320 | assert_eq!(counts.kvhash, 0); 321 | } 322 | 323 | #[test] 324 | fn big_trunk_roundtrip() { 325 | let mut tree = make_tree_seq(2u64.pow(MIN_TRUNK_HEIGHT as u32 * 2 + 1) - 1); 326 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 327 | 328 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 329 | assert!(has_more); 330 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 331 | 332 | let counts = count_node_types(trunk); 333 | // are these formulas correct for all values of `MIN_TRUNK_HEIGHT`? 🤔 334 | assert_eq!( 335 | counts.hash, 336 | 2usize.pow(MIN_TRUNK_HEIGHT as u32) + MIN_TRUNK_HEIGHT - 1 337 | ); 338 | assert_eq!(counts.kv, 2usize.pow(MIN_TRUNK_HEIGHT as u32) - 1); 339 | assert_eq!(counts.kvhash, MIN_TRUNK_HEIGHT + 1); 340 | } 341 | 342 | #[test] 343 | fn one_node_tree_trunk_roundtrip() -> Result<()> { 344 | let mut tree = BaseTree::new(vec![0], vec![])?; 345 | tree.commit(&mut NoopCommit {}).unwrap(); 346 | 347 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 348 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 349 | assert!(!has_more); 350 | 351 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 352 | let counts = count_node_types(trunk); 353 | assert_eq!(counts.hash, 0); 354 | assert_eq!(counts.kv, 1); 355 | assert_eq!(counts.kvhash, 0); 356 | Ok(()) 357 | } 358 | 359 | #[test] 360 | fn two_node_right_heavy_tree_trunk_roundtrip() -> Result<()> { 361 | // 0 362 | // \ 363 | // 1 364 | let mut tree = 365 | BaseTree::new(vec![0], vec![])?.attach(false, Some(BaseTree::new(vec![1], vec![])?)); 366 | tree.commit(&mut NoopCommit {}).unwrap(); 367 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 368 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 369 | assert!(!has_more); 370 | 371 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 372 | let counts = count_node_types(trunk); 373 | assert_eq!(counts.hash, 0); 374 | assert_eq!(counts.kv, 2); 375 | assert_eq!(counts.kvhash, 0); 376 | Ok(()) 377 | } 378 | 379 | #[test] 380 | fn two_node_left_heavy_tree_trunk_roundtrip() -> Result<()> { 381 | // 1 382 | // / 383 | // 0 384 | let mut tree = 385 | BaseTree::new(vec![1], vec![])?.attach(true, Some(BaseTree::new(vec![0], vec![])?)); 386 | tree.commit(&mut NoopCommit {}).unwrap(); 387 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 388 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 389 | assert!(!has_more); 390 | 391 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 392 | let counts = count_node_types(trunk); 393 | assert_eq!(counts.hash, 0); 394 | assert_eq!(counts.kv, 2); 395 | assert_eq!(counts.kvhash, 0); 396 | Ok(()) 397 | } 398 | 399 | #[test] 400 | fn three_node_tree_trunk_roundtrip() -> Result<()> { 401 | // 1 402 | // / \ 403 | // 0 2 404 | let mut tree = BaseTree::new(vec![1], vec![])? 405 | .attach(true, Some(BaseTree::new(vec![0], vec![])?)) 406 | .attach(false, Some(BaseTree::new(vec![2], vec![])?)); 407 | tree.commit(&mut NoopCommit {}).unwrap(); 408 | 409 | let mut walker = RefWalker::new(&mut tree, PanicSource {}); 410 | let (proof, has_more) = walker.create_trunk_proof().unwrap(); 411 | assert!(!has_more); 412 | 413 | let (trunk, _) = verify_trunk(proof.into_iter().map(Ok)).unwrap(); 414 | let counts = count_node_types(trunk); 415 | assert_eq!(counts.hash, 0); 416 | assert_eq!(counts.kv, 3); 417 | assert_eq!(counts.kvhash, 0); 418 | Ok(()) 419 | } 420 | 421 | #[test] 422 | fn leaf_chunk_roundtrip() { 423 | let mut merk = TempMerk::new().unwrap(); 424 | let batch = make_batch_seq(0..31); 425 | merk.apply(batch.as_slice(), &[]).unwrap(); 426 | 427 | let root_node = merk.tree.read().unwrap(); 428 | let root_key = root_node.as_ref().unwrap().key().to_vec(); 429 | 430 | // whole tree as 1 leaf 431 | let mut iter = merk.db.raw_iterator(); 432 | iter.seek_to_first(); 433 | let chunk = get_next_chunk(&mut iter, None).unwrap(); 434 | let ops = chunk.into_iter().map(Ok); 435 | let chunk = verify_leaf(ops, merk.root_hash()).unwrap(); 436 | let counts = count_node_types(chunk); 437 | assert_eq!(counts.kv, 31); 438 | assert_eq!(counts.hash, 0); 439 | assert_eq!(counts.kvhash, 0); 440 | drop(iter); 441 | 442 | let mut iter = merk.db.raw_iterator(); 443 | iter.seek_to_first(); 444 | 445 | // left leaf 446 | let chunk = get_next_chunk(&mut iter, Some(root_key.as_slice())).unwrap(); 447 | let ops = chunk.into_iter().map(Ok); 448 | let chunk = verify_leaf( 449 | ops, 450 | [ 451 | 222, 93, 128, 149, 117, 136, 34, 175, 204, 82, 228, 113, 242, 144, 152, 190, 210, 452 | 27, 195, 34, 24, 196, 210, 99, 250, 119, 219, 114, 52, 167, 191, 249, 453 | ], 454 | ) 455 | .unwrap(); 456 | let counts = count_node_types(chunk); 457 | assert_eq!(counts.kv, 15); 458 | assert_eq!(counts.hash, 0); 459 | assert_eq!(counts.kvhash, 0); 460 | 461 | // right leaf 462 | let chunk = get_next_chunk(&mut iter, None).unwrap(); 463 | let ops = chunk.into_iter().map(Ok); 464 | let chunk = verify_leaf( 465 | ops, 466 | [ 467 | 128, 158, 92, 80, 118, 253, 48, 241, 74, 154, 213, 187, 92, 243, 154, 28, 164, 235, 468 | 156, 122, 174, 226, 84, 170, 233, 166, 27, 79, 100, 10, 88, 184, 469 | ], 470 | ) 471 | .unwrap(); 472 | let counts = count_node_types(chunk); 473 | assert_eq!(counts.kv, 15); 474 | assert_eq!(counts.hash, 0); 475 | assert_eq!(counts.kvhash, 0); 476 | } 477 | 478 | #[test] 479 | #[should_panic(expected = "Tree is too large")] 480 | fn test_verify_height_stack_overflow() { 481 | let height = 5_000u32; 482 | let push_op = |i: u32| Op::Push(Node::KV(i.to_be_bytes().to_vec(), vec![])); 483 | let mut ops = Vec::with_capacity((height * 2) as usize); 484 | ops.push(push_op(0)); 485 | for i in 1..height { 486 | ops.push(push_op(i)); 487 | ops.push(Op::Parent) 488 | } 489 | assert!(ops.encoding_length().unwrap() < 50_000); 490 | println!("Len: {}", ops.encoding_length().unwrap()); 491 | let (_, result_height) = verify_trunk(ops.into_iter().map(Ok)).unwrap(); 492 | assert_eq!(height, result_height as u32); 493 | } 494 | } 495 | -------------------------------------------------------------------------------- /src/proofs/encoding.rs: -------------------------------------------------------------------------------- 1 | use std::io::{Read, Write}; 2 | 3 | use ed::{Decode, Encode, Terminated}; 4 | 5 | use super::{Node, Op}; 6 | use crate::error::Result; 7 | use crate::tree::HASH_LENGTH; 8 | 9 | impl Encode for Op { 10 | fn encode_into(&self, dest: &mut W) -> ed::Result<()> { 11 | match self { 12 | Op::Push(Node::Hash(hash)) => { 13 | dest.write_all(&[0x01])?; 14 | dest.write_all(hash)?; 15 | } 16 | Op::Push(Node::KVHash(kv_hash)) => { 17 | dest.write_all(&[0x02])?; 18 | dest.write_all(kv_hash)?; 19 | } 20 | Op::Push(Node::KV(key, value)) => { 21 | debug_assert!(key.len() < 65536); 22 | debug_assert!(value.len() < 65536); 23 | dest.write_all(&[0x03])?; 24 | (key.len() as u16).encode_into(dest)?; 25 | dest.write_all(key)?; 26 | (value.len() as u16).encode_into(dest)?; 27 | dest.write_all(value)?; 28 | } 29 | Op::Parent => dest.write_all(&[0x10])?, 30 | Op::Child => dest.write_all(&[0x11])?, 31 | }; 32 | Ok(()) 33 | } 34 | 35 | fn encoding_length(&self) -> ed::Result { 36 | Ok(match self { 37 | Op::Push(Node::Hash(_)) => 1 + HASH_LENGTH, 38 | Op::Push(Node::KVHash(_)) => 1 + HASH_LENGTH, 39 | Op::Push(Node::KV(key, value)) => 5 + key.len() + value.len(), 40 | Op::Parent => 1, 41 | Op::Child => 1, 42 | }) 43 | } 44 | } 45 | 46 | impl Decode for Op { 47 | fn decode(mut input: R) -> ed::Result { 48 | let variant: u8 = Decode::decode(&mut input)?; 49 | 50 | Ok(match variant { 51 | 0x01 => { 52 | let mut hash = [0; HASH_LENGTH]; 53 | input.read_exact(&mut hash)?; 54 | Op::Push(Node::Hash(hash)) 55 | } 56 | 0x02 => { 57 | let mut hash = [0; HASH_LENGTH]; 58 | input.read_exact(&mut hash)?; 59 | Op::Push(Node::KVHash(hash)) 60 | } 61 | 0x03 => { 62 | let key_len: u16 = Decode::decode(&mut input)?; 63 | let mut key = vec![0; key_len as usize]; 64 | input.read_exact(key.as_mut_slice())?; 65 | 66 | let value_len: u16 = Decode::decode(&mut input)?; 67 | let mut value = vec![0; value_len as usize]; 68 | input.read_exact(value.as_mut_slice())?; 69 | 70 | Op::Push(Node::KV(key, value)) 71 | } 72 | 0x10 => Op::Parent, 73 | 0x11 => Op::Child, 74 | byte => { 75 | return Err(ed::Error::UnexpectedByte(byte)); 76 | } 77 | }) 78 | } 79 | } 80 | 81 | impl Terminated for Op {} 82 | 83 | impl Op { 84 | fn encode_into(&self, dest: &mut W) -> Result<()> { 85 | Ok(Encode::encode_into(self, dest)?) 86 | } 87 | 88 | fn encoding_length(&self) -> usize { 89 | Encode::encoding_length(self).unwrap() 90 | } 91 | 92 | pub fn decode(bytes: &[u8]) -> Result { 93 | Ok(Decode::decode(bytes)?) 94 | } 95 | } 96 | 97 | pub fn encode_into<'a, T: Iterator>(ops: T, output: &mut Vec) { 98 | for op in ops { 99 | op.encode_into(output).unwrap(); 100 | } 101 | } 102 | 103 | pub struct Decoder<'a> { 104 | offset: usize, 105 | bytes: &'a [u8], 106 | } 107 | 108 | impl<'a> Decoder<'a> { 109 | pub fn new(proof_bytes: &'a [u8]) -> Self { 110 | Decoder { 111 | offset: 0, 112 | bytes: proof_bytes, 113 | } 114 | } 115 | } 116 | 117 | impl<'a> Iterator for Decoder<'a> { 118 | type Item = Result; 119 | 120 | fn next(&mut self) -> Option { 121 | if self.offset >= self.bytes.len() { 122 | return None; 123 | } 124 | 125 | Some((|| { 126 | let bytes = &self.bytes[self.offset..]; 127 | let op = Op::decode(bytes)?; 128 | self.offset += op.encoding_length(); 129 | Ok(op) 130 | })()) 131 | } 132 | } 133 | 134 | #[cfg(test)] 135 | mod test { 136 | use super::super::{Node, Op}; 137 | use crate::tree::HASH_LENGTH; 138 | 139 | #[test] 140 | fn encode_push_hash() { 141 | let op = Op::Push(Node::Hash([123; HASH_LENGTH])); 142 | assert_eq!(op.encoding_length(), 1 + HASH_LENGTH); 143 | 144 | let mut bytes = vec![]; 145 | op.encode_into(&mut bytes).unwrap(); 146 | assert_eq!( 147 | bytes, 148 | vec![ 149 | 0x01, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 150 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 151 | 123 152 | ] 153 | ); 154 | } 155 | 156 | #[test] 157 | fn encode_push_kvhash() { 158 | let op = Op::Push(Node::KVHash([123; HASH_LENGTH])); 159 | assert_eq!(op.encoding_length(), 1 + HASH_LENGTH); 160 | 161 | let mut bytes = vec![]; 162 | op.encode_into(&mut bytes).unwrap(); 163 | assert_eq!( 164 | bytes, 165 | vec![ 166 | 0x02, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 167 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 168 | 123 169 | ] 170 | ); 171 | } 172 | 173 | #[test] 174 | fn encode_push_kv() { 175 | let op = Op::Push(Node::KV(vec![1, 2, 3], vec![4, 5, 6])); 176 | assert_eq!(op.encoding_length(), 11); 177 | 178 | let mut bytes = vec![]; 179 | op.encode_into(&mut bytes).unwrap(); 180 | assert_eq!(bytes, vec![0x03, 0, 3, 1, 2, 3, 0, 3, 4, 5, 6]); 181 | } 182 | 183 | #[test] 184 | fn encode_parent() { 185 | let op = Op::Parent; 186 | assert_eq!(op.encoding_length(), 1); 187 | 188 | let mut bytes = vec![]; 189 | op.encode_into(&mut bytes).unwrap(); 190 | assert_eq!(bytes, vec![0x10]); 191 | } 192 | 193 | #[test] 194 | fn encode_child() { 195 | let op = Op::Child; 196 | assert_eq!(op.encoding_length(), 1); 197 | 198 | let mut bytes = vec![]; 199 | op.encode_into(&mut bytes).unwrap(); 200 | assert_eq!(bytes, vec![0x11]); 201 | } 202 | 203 | #[test] 204 | #[should_panic] 205 | fn encode_push_kv_long_key() { 206 | let op = Op::Push(Node::KV(vec![123; 70_000], vec![4, 5, 6])); 207 | let mut bytes = vec![]; 208 | op.encode_into(&mut bytes).unwrap(); 209 | } 210 | 211 | #[test] 212 | fn decode_push_hash() { 213 | let bytes = [ 214 | 0x01, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 215 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 216 | ]; 217 | let op = Op::decode(&bytes[..]).expect("decode failed"); 218 | assert_eq!(op, Op::Push(Node::Hash([123; HASH_LENGTH]))); 219 | } 220 | 221 | #[test] 222 | fn decode_push_kvhash() { 223 | let bytes = [ 224 | 0x02, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 225 | 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 226 | ]; 227 | let op = Op::decode(&bytes[..]).expect("decode failed"); 228 | assert_eq!(op, Op::Push(Node::KVHash([123; HASH_LENGTH]))); 229 | } 230 | 231 | #[test] 232 | fn decode_push_kv() { 233 | let bytes = [0x03, 0, 3, 1, 2, 3, 0, 3, 4, 5, 6]; 234 | let op = Op::decode(&bytes[..]).expect("decode failed"); 235 | assert_eq!(op, Op::Push(Node::KV(vec![1, 2, 3], vec![4, 5, 6]))); 236 | } 237 | 238 | #[test] 239 | fn decode_parent() { 240 | let bytes = [0x10]; 241 | let op = Op::decode(&bytes[..]).expect("decode failed"); 242 | assert_eq!(op, Op::Parent); 243 | } 244 | 245 | #[test] 246 | fn decode_child() { 247 | let bytes = [0x11]; 248 | let op = Op::decode(&bytes[..]).expect("decode failed"); 249 | assert_eq!(op, Op::Child); 250 | } 251 | 252 | #[test] 253 | fn decode_unknown() { 254 | let bytes = [0x88]; 255 | assert!(Op::decode(&bytes[..]).is_err()); 256 | } 257 | } 258 | -------------------------------------------------------------------------------- /src/proofs/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod chunk; 2 | pub mod encoding; 3 | pub mod query; 4 | pub mod tree; 5 | 6 | use crate::tree::Hash; 7 | 8 | pub use encoding::{encode_into, Decoder}; 9 | pub use query::Query; 10 | pub use tree::Tree; 11 | 12 | /// A proof operator, executed to verify the data in a Merkle proof. 13 | #[derive(Debug, PartialEq)] 14 | pub enum Op { 15 | /// Pushes a node on the stack. 16 | Push(Node), 17 | 18 | /// Pops the top stack item as `parent`. Pops the next top stack item as 19 | /// `child`. Attaches `child` as the left child of `parent`. Pushes the 20 | /// updated `parent` back on the stack. 21 | Parent, 22 | 23 | /// Pops the top stack item as `child`. Pops the next top stack item as 24 | /// `parent`. Attaches `child` as the right child of `parent`. Pushes the 25 | /// updated `parent` back on the stack. 26 | Child, 27 | } 28 | 29 | /// A selected piece of data about a single tree node, to be contained in a 30 | /// `Push` operator in a proof. 31 | #[derive(Clone, Debug, PartialEq)] 32 | pub enum Node { 33 | /// Represents the hash of a tree node. 34 | Hash(Hash), 35 | 36 | /// Represents the hash of the key/value pair of a tree node. 37 | KVHash(Hash), 38 | 39 | /// Represents the key and value of a tree node. 40 | KV(Vec, Vec), 41 | } 42 | -------------------------------------------------------------------------------- /src/proofs/tree.rs: -------------------------------------------------------------------------------- 1 | use super::{Node, Op}; 2 | use crate::error::{Error, Result}; 3 | use crate::tree::{kv_hash, node_hash, Hash, Hasher, NULL_HASH}; 4 | 5 | /// Contains a tree's child node and its hash. The hash can always be assumed to 6 | /// be up-to-date. 7 | #[derive(Debug)] 8 | pub struct Child { 9 | /// The child node. 10 | pub tree: Box, 11 | /// The hash of the child node. 12 | pub hash: Hash, 13 | } 14 | 15 | /// A binary tree data structure used to represent a select subset of a tree 16 | /// when verifying Merkle proofs. 17 | #[derive(Debug)] 18 | pub struct Tree { 19 | /// The node at the root of this tree. 20 | pub node: Node, 21 | /// The left child of this tree. 22 | pub left: Option, 23 | /// The right child of this tree. 24 | pub right: Option, 25 | /// The height of this tree. 26 | pub height: usize, 27 | } 28 | 29 | impl From for Tree { 30 | /// Creates a childless tree with the target node as the `node` field. 31 | fn from(node: Node) -> Self { 32 | Tree { 33 | node, 34 | left: None, 35 | right: None, 36 | height: 1, 37 | } 38 | } 39 | } 40 | 41 | impl PartialEq for Tree { 42 | /// Checks equality for the root hashes of the two trees. 43 | fn eq(&self, other: &Self) -> bool { 44 | self.hash() 45 | .and_then(|this_hash| other.hash().map(|other_hash| this_hash == other_hash)) 46 | .unwrap_or_default() 47 | } 48 | } 49 | 50 | impl Tree { 51 | /// Gets or computes the hash for this tree node. 52 | pub fn hash(&self) -> Result { 53 | fn compute_hash(tree: &Tree, kv_hash: Hash) -> Hash { 54 | node_hash::(&kv_hash, &tree.child_hash(true), &tree.child_hash(false)) 55 | } 56 | 57 | match &self.node { 58 | Node::Hash(hash) => Ok(*hash), 59 | Node::KVHash(kv_hash) => Ok(compute_hash(self, *kv_hash)), 60 | Node::KV(key, value) => kv_hash::(key.as_slice(), value.as_slice()) 61 | .map(|kv_hash| compute_hash(self, kv_hash)) 62 | .map_err(Into::into), 63 | } 64 | } 65 | 66 | /// Creates an iterator that yields the in-order traversal of the nodes at 67 | /// the given depth. 68 | pub fn layer(&self, depth: usize) -> LayerIter { 69 | LayerIter::new(self, depth) 70 | } 71 | 72 | /// Consumes the `Tree` and does an in-order traversal over all the nodes in 73 | /// the tree, calling `visit_node` for each. 74 | pub fn visit_nodes(mut self, visit_node: &mut F) { 75 | if let Some(child) = self.left.take() { 76 | child.tree.visit_nodes(visit_node); 77 | } 78 | 79 | let maybe_right_child = self.right.take(); 80 | visit_node(self.node); 81 | 82 | if let Some(child) = maybe_right_child { 83 | child.tree.visit_nodes(visit_node); 84 | } 85 | } 86 | 87 | /// Does an in-order traversal over references to all the nodes in the tree, 88 | /// calling `visit_node` for each. 89 | pub fn visit_refs(&self, visit_node: &mut F) { 90 | if let Some(child) = &self.left { 91 | child.tree.visit_refs(visit_node); 92 | } 93 | 94 | visit_node(self); 95 | 96 | if let Some(child) = &self.right { 97 | child.tree.visit_refs(visit_node); 98 | } 99 | } 100 | 101 | /// Returns an immutable reference to the child on the given side, if any. 102 | pub fn child(&self, left: bool) -> Option<&Child> { 103 | if left { 104 | self.left.as_ref() 105 | } else { 106 | self.right.as_ref() 107 | } 108 | } 109 | 110 | /// Returns a mutable reference to the child on the given side, if any. 111 | pub(crate) fn child_mut(&mut self, left: bool) -> &mut Option { 112 | if left { 113 | &mut self.left 114 | } else { 115 | &mut self.right 116 | } 117 | } 118 | 119 | /// Attaches the child to the `Tree`'s given side. Returns an error if 120 | /// there is already a child attached to this side. 121 | pub(crate) fn attach(&mut self, left: bool, child: Tree) -> Result<()> { 122 | if self.child(left).is_some() { 123 | return Err(Error::Attach( 124 | "Tried to attach to left child, but it is already Some".into(), 125 | )); 126 | } 127 | 128 | if let Node::Hash(_) = self.node { 129 | return Err(Error::Attach("Tried to attach to Hash node".into())); 130 | } 131 | 132 | self.height = self.height.max(child.height + 1); 133 | 134 | let hash = child.hash()?; 135 | let tree = Box::new(child); 136 | *self.child_mut(left) = Some(Child { tree, hash }); 137 | 138 | Ok(()) 139 | } 140 | 141 | /// Returns the already-computed hash for this tree node's child on the 142 | /// given side, if any. If there is no child, returns the null hash 143 | /// (zero-filled). 144 | #[inline] 145 | fn child_hash(&self, left: bool) -> Hash { 146 | self.child(left).map_or(NULL_HASH, |c| c.hash) 147 | } 148 | 149 | /// Consumes the tree node, calculates its hash, and returns a `Node::Hash` 150 | /// variant. 151 | fn try_into_hash(self) -> Result { 152 | self.hash().map(Node::Hash).map(Into::into) 153 | } 154 | 155 | #[cfg(feature = "full")] 156 | pub(crate) fn key(&self) -> &[u8] { 157 | match self.node { 158 | Node::KV(ref key, _) => key, 159 | _ => panic!("Expected node to be type KV"), 160 | } 161 | } 162 | } 163 | 164 | /// `LayerIter` iterates over the nodes in a `Tree` at a given depth. Nodes are 165 | /// visited in order. 166 | pub struct LayerIter<'a> { 167 | stack: Vec<&'a Tree>, 168 | depth: usize, 169 | } 170 | 171 | impl<'a> LayerIter<'a> { 172 | /// Creates a new `LayerIter` that iterates over `tree` at the given depth. 173 | fn new(tree: &'a Tree, depth: usize) -> Self { 174 | let mut iter = LayerIter { 175 | stack: Vec::with_capacity(depth), 176 | depth, 177 | }; 178 | 179 | iter.traverse_to_start(tree, depth); 180 | iter 181 | } 182 | 183 | /// Builds up the stack by traversing through left children to the desired 184 | /// depth. 185 | fn traverse_to_start(&mut self, tree: &'a Tree, remaining_depth: usize) { 186 | self.stack.push(tree); 187 | 188 | if remaining_depth == 0 { 189 | return; 190 | } 191 | 192 | if let Some(child) = tree.child(true) { 193 | self.traverse_to_start(&child.tree, remaining_depth - 1) 194 | } else { 195 | panic!("Could not traverse to given layer") 196 | } 197 | } 198 | } 199 | 200 | impl<'a> Iterator for LayerIter<'a> { 201 | type Item = &'a Tree; 202 | 203 | fn next(&mut self) -> Option { 204 | let item = self.stack.pop(); 205 | let mut popped = item; 206 | 207 | loop { 208 | if self.stack.is_empty() { 209 | return item; 210 | } 211 | 212 | let parent = self.stack.last().unwrap(); 213 | let left_child = parent.child(true).unwrap(); 214 | let right_child = parent.child(false).unwrap(); 215 | 216 | if left_child.tree.as_ref() == popped.unwrap() { 217 | self.stack.push(&right_child.tree); 218 | 219 | while self.stack.len() - 1 < self.depth { 220 | let parent = self.stack.last().unwrap(); 221 | let left_child = parent.child(true).unwrap(); 222 | self.stack.push(&left_child.tree); 223 | } 224 | 225 | return item; 226 | } else { 227 | popped = self.stack.pop(); 228 | } 229 | } 230 | } 231 | } 232 | 233 | /// Executes a proof by stepping through its operators, modifying the 234 | /// verification stack as it goes. The resulting stack item is returned. 235 | /// 236 | /// If the `collapse` option is set to `true`, nodes will be hashed and pruned 237 | /// from memory during execution. This results in the minimum amount of memory 238 | /// usage, and the returned `Tree` will only contain a single node of type 239 | /// `Node::Hash`. If `false`, the returned `Tree` will contain the entire 240 | /// subtree contained in the proof. 241 | /// 242 | /// `visit_node` will be called once for every push operation in the proof, in 243 | /// key-order. If `visit_node` returns an `Err` result, it will halt the 244 | /// execution and `execute` will return the error. 245 | pub(crate) fn execute(ops: I, collapse: bool, mut visit_node: F) -> Result 246 | where 247 | I: IntoIterator>, 248 | F: FnMut(&Node) -> Result<()>, 249 | { 250 | let mut stack: Vec = Vec::with_capacity(32); 251 | let mut maybe_last_key = None; 252 | 253 | fn try_pop(stack: &mut Vec) -> Result { 254 | match stack.pop() { 255 | None => Err(Error::StackUnderflow), 256 | Some(tree) => Ok(tree), 257 | } 258 | } 259 | 260 | for op in ops { 261 | match op? { 262 | Op::Parent => { 263 | let (mut parent, child) = (try_pop(&mut stack)?, try_pop(&mut stack)?); 264 | parent.attach( 265 | true, 266 | if collapse { 267 | child.try_into_hash()? 268 | } else { 269 | child 270 | }, 271 | )?; 272 | stack.push(parent); 273 | } 274 | Op::Child => { 275 | let (child, mut parent) = (try_pop(&mut stack)?, try_pop(&mut stack)?); 276 | parent.attach( 277 | false, 278 | if collapse { 279 | child.try_into_hash()? 280 | } else { 281 | child 282 | }, 283 | )?; 284 | stack.push(parent); 285 | } 286 | Op::Push(node) => { 287 | if let Node::KV(key, _) = &node { 288 | // keys should always increase 289 | if let Some(last_key) = &maybe_last_key { 290 | if key <= last_key { 291 | return Err(Error::Key("Incorrect key ordering".into())); 292 | } 293 | } 294 | 295 | maybe_last_key = Some(key.clone()); 296 | } 297 | 298 | visit_node(&node)?; 299 | 300 | let tree: Tree = node.into(); 301 | stack.push(tree); 302 | } 303 | } 304 | } 305 | 306 | if stack.len() != 1 { 307 | return Err(Error::Proof( 308 | "Expected proof to result in exactly on stack item".into(), 309 | )); 310 | } 311 | 312 | Ok(stack.pop().unwrap()) 313 | } 314 | 315 | #[cfg(test)] 316 | mod test { 317 | use super::super::*; 318 | use super::Tree as ProofTree; 319 | use super::*; 320 | 321 | fn make_7_node_prooftree() -> ProofTree { 322 | let make_node = |i| -> super::super::tree::Tree { Node::KV(vec![i], vec![]).into() }; 323 | 324 | let mut tree = make_node(3); 325 | let mut left = make_node(1); 326 | left.attach(true, make_node(0)).unwrap(); 327 | left.attach(false, make_node(2)).unwrap(); 328 | let mut right = make_node(5); 329 | right.attach(true, make_node(4)).unwrap(); 330 | right.attach(false, make_node(6)).unwrap(); 331 | tree.attach(true, left).unwrap(); 332 | tree.attach(false, right).unwrap(); 333 | 334 | tree 335 | } 336 | 337 | #[test] 338 | fn height_counting() { 339 | fn recurse(tree: &super::Tree, expected_height: usize) { 340 | assert_eq!(tree.height, expected_height); 341 | tree.left 342 | .as_ref() 343 | .into_iter() 344 | .for_each(|l| recurse(&l.tree, expected_height - 1)); 345 | tree.right 346 | .as_ref() 347 | .into_iter() 348 | .for_each(|r| recurse(&r.tree, expected_height - 1)); 349 | } 350 | 351 | let tree = make_7_node_prooftree(); 352 | recurse(&tree, 3); 353 | } 354 | 355 | #[test] 356 | fn layer_iter() { 357 | let tree = make_7_node_prooftree(); 358 | 359 | let assert_node = |node: &Tree, i| match node.node { 360 | Node::KV(ref key, _) => assert_eq!(key[0], i), 361 | _ => unreachable!(), 362 | }; 363 | 364 | let mut iter = tree.layer(0); 365 | assert_node(iter.next().unwrap(), 3); 366 | assert!(iter.next().is_none()); 367 | 368 | let mut iter = tree.layer(1); 369 | assert_node(iter.next().unwrap(), 1); 370 | assert_node(iter.next().unwrap(), 5); 371 | assert!(iter.next().is_none()); 372 | 373 | let mut iter = tree.layer(2); 374 | assert_node(iter.next().unwrap(), 0); 375 | assert_node(iter.next().unwrap(), 2); 376 | assert_node(iter.next().unwrap(), 4); 377 | assert_node(iter.next().unwrap(), 6); 378 | assert!(iter.next().is_none()); 379 | } 380 | 381 | #[test] 382 | fn visit_nodes() { 383 | let tree = make_7_node_prooftree(); 384 | 385 | let assert_node = |node: Node, i| match node { 386 | Node::KV(ref key, _) => assert_eq!(key[0], i), 387 | _ => unreachable!(), 388 | }; 389 | 390 | let mut visited = vec![]; 391 | tree.visit_nodes(&mut |node| visited.push(node)); 392 | 393 | let mut iter = visited.into_iter(); 394 | for i in 0..7 { 395 | assert_node(iter.next().unwrap(), i); 396 | } 397 | assert!(iter.next().is_none()); 398 | } 399 | } 400 | -------------------------------------------------------------------------------- /src/test_utils/crash_merk.rs: -------------------------------------------------------------------------------- 1 | use crate::{Merk, Result}; 2 | use std::fs; 3 | use std::mem::ManuallyDrop; 4 | use std::ops::{Deref, DerefMut}; 5 | use std::path::Path; 6 | 7 | /// Wraps a Merk instance and drops it without flushing once it goes out of 8 | /// scope. 9 | pub struct CrashMerk { 10 | inner: Option>, 11 | path: Box, 12 | } 13 | 14 | impl CrashMerk { 15 | /// Opens a `CrashMerk` at the given file path, creating a new one if it 16 | /// does not exist. 17 | pub fn open>(path: P) -> Result { 18 | let merk = Merk::open(&path)?; 19 | let inner = Some(ManuallyDrop::new(merk)); 20 | Ok(CrashMerk { 21 | inner, 22 | path: path.as_ref().into(), 23 | }) 24 | } 25 | 26 | #[allow(clippy::missing_safety_doc)] 27 | pub unsafe fn crash(&mut self) -> Result<()> { 28 | ManuallyDrop::drop(&mut self.inner.take().unwrap()); 29 | 30 | // rename to invalidate rocksdb's lock 31 | let file_name = format!( 32 | "{}_crashed", 33 | self.path.file_name().unwrap().to_str().unwrap() 34 | ); 35 | let new_path = self.path.with_file_name(file_name); 36 | fs::rename(&self.path, &new_path)?; 37 | 38 | let mut new_merk = CrashMerk::open(&new_path)?; 39 | self.inner = new_merk.inner.take(); 40 | self.path = new_merk.path; 41 | Ok(()) 42 | } 43 | 44 | pub fn into_inner(self) -> Merk { 45 | ManuallyDrop::into_inner(self.inner.unwrap()) 46 | } 47 | 48 | pub fn destroy(self) -> Result<()> { 49 | self.into_inner().destroy() 50 | } 51 | } 52 | 53 | impl Deref for CrashMerk { 54 | type Target = Merk; 55 | 56 | fn deref(&self) -> &Merk { 57 | self.inner.as_ref().unwrap() 58 | } 59 | } 60 | 61 | impl DerefMut for CrashMerk { 62 | fn deref_mut(&mut self) -> &mut Merk { 63 | self.inner.as_mut().unwrap() 64 | } 65 | } 66 | 67 | #[cfg(test)] 68 | mod tests { 69 | use super::CrashMerk; 70 | use crate::Op; 71 | 72 | #[test] 73 | #[ignore] // currently this still works because we enabled the WAL 74 | fn crash() { 75 | let path = std::thread::current().name().unwrap().to_owned(); 76 | 77 | let mut merk = CrashMerk::open(path).expect("failed to open merk"); 78 | merk.apply(&[(vec![1, 2, 3], Op::Put(vec![4, 5, 6]))], &[]) 79 | .expect("apply failed"); 80 | unsafe { 81 | merk.crash().unwrap(); 82 | } 83 | assert_eq!(merk.get(&[1, 2, 3]).expect("failed to get"), None); 84 | merk.into_inner().destroy().unwrap(); 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /src/test_utils/mod.rs: -------------------------------------------------------------------------------- 1 | #![allow(missing_docs)] 2 | 3 | mod crash_merk; 4 | mod temp_merk; 5 | 6 | use crate::tree::{Batch, BatchEntry, NoopCommit, Op, PanicSource, Tree, Walker}; 7 | use rand::prelude::*; 8 | use std::convert::TryInto; 9 | use std::ops::Range; 10 | 11 | pub use crash_merk::CrashMerk; 12 | pub use temp_merk::TempMerk; 13 | 14 | pub fn assert_tree_invariants(tree: &Tree) { 15 | assert!(tree.balance_factor().abs() < 2); 16 | 17 | let maybe_left = tree.link(true); 18 | if let Some(left) = maybe_left { 19 | assert!(left.key() < tree.key()); 20 | assert!(!left.is_modified()); 21 | } 22 | 23 | let maybe_right = tree.link(false); 24 | if let Some(right) = maybe_right { 25 | assert!(right.key() > tree.key()); 26 | assert!(!right.is_modified()); 27 | } 28 | 29 | if let Some(left) = tree.child(true) { 30 | assert_tree_invariants(left); 31 | } 32 | if let Some(right) = tree.child(false) { 33 | assert_tree_invariants(right); 34 | } 35 | } 36 | 37 | pub fn apply_memonly_unchecked(tree: Tree, batch: &Batch) -> Tree { 38 | let walker = Walker::::new(tree, PanicSource {}); 39 | let mut tree = Walker::::apply_to(Some(walker), batch, PanicSource {}) 40 | .expect("apply failed") 41 | .0 42 | .expect("expected tree"); 43 | tree.commit(&mut NoopCommit {}).expect("commit failed"); 44 | tree 45 | } 46 | 47 | pub fn apply_memonly(tree: Tree, batch: &Batch) -> Tree { 48 | let tree = apply_memonly_unchecked(tree, batch); 49 | assert_tree_invariants(&tree); 50 | tree 51 | } 52 | 53 | pub fn apply_to_memonly(maybe_tree: Option, batch: &Batch) -> Option { 54 | let maybe_walker = maybe_tree.map(|tree| Walker::::new(tree, PanicSource {})); 55 | Walker::::apply_to(maybe_walker, batch, PanicSource {}) 56 | .expect("apply failed") 57 | .0 58 | .map(|mut tree| { 59 | tree.commit(&mut NoopCommit {}).expect("commit failed"); 60 | println!("{:?}", &tree); 61 | assert_tree_invariants(&tree); 62 | tree 63 | }) 64 | } 65 | 66 | pub fn seq_key(n: u64) -> Vec { 67 | n.to_be_bytes().to_vec() 68 | } 69 | 70 | pub fn put_entry_value() -> Vec { 71 | vec![123; 60] 72 | } 73 | 74 | pub fn put_entry(n: u64) -> BatchEntry { 75 | (seq_key(n), Op::Put(put_entry_value())) 76 | } 77 | 78 | pub fn del_entry(n: u64) -> BatchEntry { 79 | (seq_key(n), Op::Delete) 80 | } 81 | 82 | pub fn make_batch_seq(range: Range) -> Vec { 83 | let mut batch = Vec::with_capacity((range.end - range.start).try_into().unwrap()); 84 | for n in range { 85 | batch.push(put_entry(n)); 86 | } 87 | batch 88 | } 89 | 90 | pub fn make_del_batch_seq(range: Range) -> Vec { 91 | let mut batch = Vec::with_capacity((range.end - range.start).try_into().unwrap()); 92 | for n in range { 93 | batch.push(del_entry(n)); 94 | } 95 | batch 96 | } 97 | 98 | pub fn make_batch_rand(size: u64, seed: u64) -> Vec { 99 | let mut rng: SmallRng = SeedableRng::seed_from_u64(seed); 100 | let mut batch = Vec::with_capacity(size.try_into().unwrap()); 101 | for _ in 0..size { 102 | let n = rng.gen::(); 103 | batch.push(put_entry(n)); 104 | } 105 | batch.sort_by(|a, b| a.0.cmp(&b.0)); 106 | batch 107 | } 108 | 109 | pub fn make_del_batch_rand(size: u64, seed: u64) -> Vec { 110 | let mut rng: SmallRng = SeedableRng::seed_from_u64(seed); 111 | let mut batch = Vec::with_capacity(size.try_into().unwrap()); 112 | for _ in 0..size { 113 | let n = rng.gen::(); 114 | batch.push(del_entry(n)); 115 | } 116 | batch.sort_by(|a, b| a.0.cmp(&b.0)); 117 | batch 118 | } 119 | 120 | pub fn make_tree_rand(node_count: u64, batch_size: u64, initial_seed: u64) -> Tree { 121 | assert!(node_count >= batch_size); 122 | assert!((node_count % batch_size) == 0); 123 | 124 | let value = vec![123; 60]; 125 | let mut tree = Tree::new(vec![0; 20], value).expect("Tree construction failed"); 126 | 127 | let mut seed = initial_seed; 128 | 129 | let batch_count = node_count / batch_size; 130 | for _ in 0..batch_count { 131 | let batch = make_batch_rand(batch_size, seed); 132 | tree = apply_memonly(tree, &batch); 133 | seed += 1; 134 | } 135 | 136 | tree 137 | } 138 | 139 | pub fn make_tree_seq(node_count: u64) -> Tree { 140 | let batch_size = if node_count >= 10_000 { 141 | assert!(node_count % 10_000 == 0); 142 | 10_000 143 | } else { 144 | node_count 145 | }; 146 | 147 | let value = vec![123; 60]; 148 | let mut tree = Tree::new(vec![0; 20], value).expect("Tree construction failed"); 149 | 150 | let batch_count = node_count / batch_size; 151 | for i in 0..batch_count { 152 | let batch = make_batch_seq((i * batch_size)..((i + 1) * batch_size)); 153 | tree = apply_memonly(tree, &batch); 154 | } 155 | 156 | tree 157 | } 158 | -------------------------------------------------------------------------------- /src/test_utils/temp_merk.rs: -------------------------------------------------------------------------------- 1 | use crate::{Merk, Result}; 2 | use std::env::temp_dir; 3 | use std::ops::{Deref, DerefMut}; 4 | use std::path::{Path, PathBuf}; 5 | use std::time::SystemTime; 6 | 7 | /// Wraps a Merk instance and deletes it from disk it once it goes out of scope. 8 | pub struct TempMerk { 9 | inner: Option, 10 | } 11 | 12 | impl TempMerk { 13 | /// Opens a `TempMerk` at the given file path, creating a new one if it does 14 | /// not exist. 15 | pub fn open>(path: P) -> Result { 16 | let inner = Some(Merk::open(path)?); 17 | Ok(TempMerk { inner }) 18 | } 19 | 20 | /// Opens a `TempMerk` at an autogenerated, temporary file path. 21 | pub fn new() -> Result { 22 | TempMerk::open(Self::create_path()) 23 | } 24 | 25 | pub fn create_path() -> PathBuf { 26 | let time = SystemTime::now() 27 | .duration_since(SystemTime::UNIX_EPOCH) 28 | .unwrap() 29 | .as_nanos(); 30 | let mut path = temp_dir(); 31 | path.push(format!("merk-temp–{time}")); 32 | path 33 | } 34 | } 35 | 36 | impl Drop for TempMerk { 37 | fn drop(&mut self) { 38 | self.inner 39 | .take() 40 | .unwrap() 41 | .destroy() 42 | .expect("failed to delete db"); 43 | } 44 | } 45 | 46 | impl Deref for TempMerk { 47 | type Target = Merk; 48 | 49 | fn deref(&self) -> &Merk { 50 | self.inner.as_ref().unwrap() 51 | } 52 | } 53 | 54 | impl DerefMut for TempMerk { 55 | fn deref_mut(&mut self) -> &mut Merk { 56 | self.inner.as_mut().unwrap() 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /src/tree/commit.rs: -------------------------------------------------------------------------------- 1 | use super::Tree; 2 | use crate::error::Result; 3 | 4 | /// To be used when committing a tree (writing it to a store after applying the 5 | /// changes). 6 | pub trait Commit { 7 | /// Called once per updated node when a finalized tree is to be written to a 8 | /// backing store or cache. 9 | fn write(&mut self, tree: &Tree) -> Result<()>; 10 | 11 | /// Called once per node after writing a node and its children. The returned 12 | /// tuple specifies whether or not to prune the left and right child nodes, 13 | /// respectively. For example, returning `(true, true)` will prune both 14 | /// nodes, removing them from memory. 15 | fn prune(&self, _tree: &Tree) -> (bool, bool) { 16 | (true, true) 17 | } 18 | } 19 | 20 | /// A `Commit` implementation which does not write to a store and does not prune 21 | /// any nodes from the Tree. Useful when only keeping a tree in memory. 22 | pub struct NoopCommit {} 23 | impl Commit for NoopCommit { 24 | fn write(&mut self, _tree: &Tree) -> Result<()> { 25 | Ok(()) 26 | } 27 | 28 | fn prune(&self, _tree: &Tree) -> (bool, bool) { 29 | (false, false) 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /src/tree/debug.rs: -------------------------------------------------------------------------------- 1 | use super::{Link, Tree}; 2 | use colored::Colorize; 3 | use std::fmt::{Debug, Formatter, Result}; 4 | 5 | impl Debug for Tree { 6 | // TODO: unwraps should be results that bubble up 7 | fn fmt(&self, f: &mut Formatter) -> Result { 8 | fn traverse( 9 | f: &mut Formatter, 10 | cursor: &Tree, 11 | stack: &mut Vec<(Vec, Vec)>, 12 | left: bool, 13 | ) { 14 | if let Some(child_link) = cursor.link(true) { 15 | stack.push((child_link.key().to_vec(), cursor.key().to_vec())); 16 | if let Some(child_tree) = child_link.tree() { 17 | traverse(f, child_tree, stack, true); 18 | } else { 19 | traverse_pruned(f, child_link, stack, true); 20 | } 21 | stack.pop(); 22 | } 23 | 24 | let depth = stack.len(); 25 | 26 | if depth > 0 { 27 | // draw ancestor's vertical lines 28 | for (low, high) in stack.iter().take(depth - 1) { 29 | let draw_line = cursor.key() > low && cursor.key() < high; 30 | write!(f, "{}", if draw_line { " │ " } else { " " }.dimmed()).unwrap(); 31 | } 32 | } 33 | 34 | let prefix = if depth == 0 { 35 | "" 36 | } else if left { 37 | " ┌-" 38 | } else { 39 | " └-" 40 | }; 41 | writeln!( 42 | f, 43 | "{}{}", 44 | prefix.dimmed(), 45 | format!("{:?}", cursor.key()).on_bright_black() 46 | ) 47 | .unwrap(); 48 | 49 | if let Some(child_link) = cursor.link(false) { 50 | stack.push((cursor.key().to_vec(), child_link.key().to_vec())); 51 | if let Some(child_tree) = child_link.tree() { 52 | traverse(f, child_tree, stack, false); 53 | } else { 54 | traverse_pruned(f, child_link, stack, false); 55 | } 56 | stack.pop(); 57 | } 58 | } 59 | 60 | fn traverse_pruned( 61 | f: &mut Formatter, 62 | link: &Link, 63 | stack: &mut [(Vec, Vec)], 64 | left: bool, 65 | ) { 66 | let depth = stack.len(); 67 | 68 | if depth > 0 { 69 | // draw ancestor's vertical lines 70 | for (low, high) in stack.iter().take(depth - 1) { 71 | let draw_line = link.key() > low && link.key() < high; 72 | write!(f, "{}", if draw_line { " │ " } else { " " }.dimmed()).unwrap(); 73 | } 74 | } 75 | 76 | let prefix = if depth == 0 { 77 | "" 78 | } else if left { 79 | " ┌-" 80 | } else { 81 | " └-" 82 | }; 83 | writeln!( 84 | f, 85 | "{}{}", 86 | prefix.dimmed(), 87 | format!("{:?}", link.key()).blue() 88 | ) 89 | .unwrap(); 90 | } 91 | 92 | let mut stack = vec![]; 93 | traverse(f, self, &mut stack, false); 94 | writeln!(f) 95 | } 96 | } 97 | -------------------------------------------------------------------------------- /src/tree/encoding.rs: -------------------------------------------------------------------------------- 1 | use std::io::Read; 2 | 3 | use crate::Result; 4 | 5 | use super::{kv::KV, Link, Tree, TreeInner}; 6 | use ed::{Decode, Encode}; 7 | 8 | impl Tree { 9 | #[inline] 10 | pub fn encode(&self) -> Vec { 11 | // operation is infallible so it's ok to unwrap 12 | Encode::encode(self).unwrap() 13 | } 14 | 15 | #[inline] 16 | pub fn encode_into(&self, dest: &mut Vec) { 17 | // operation is infallible so it's ok to unwrap 18 | Encode::encode_into(self, dest).unwrap() 19 | } 20 | 21 | #[inline] 22 | pub fn encoding_length(&self) -> usize { 23 | // operation is infallible so it's ok to unwrap 24 | Encode::encoding_length(self).unwrap() 25 | } 26 | 27 | #[inline] 28 | pub fn decode_into(&mut self, key: Vec, input: &[u8]) { 29 | // operation is infallible so it's ok to unwrap 30 | Decode::decode_into(self, input).unwrap(); 31 | self.inner.kv.key = key; 32 | } 33 | 34 | #[inline] 35 | pub fn decode(key: Vec, input: &[u8]) -> Tree { 36 | // operation is infallible so it's ok to unwrap 37 | let mut tree: Tree = Decode::decode(input).unwrap(); 38 | tree.inner.kv.key = key; 39 | tree 40 | } 41 | 42 | pub fn decode_v0(mut input: R) -> Result { 43 | let mut read_link_v0 = || -> Result> { 44 | let some = bool::decode(&mut input)?; 45 | if some { 46 | let link = Link::decode_v0(&mut input)?; 47 | Ok(Some(link)) 48 | } else { 49 | Ok(None) 50 | } 51 | }; 52 | 53 | let maybe_left = read_link_v0()?; 54 | let maybe_right = read_link_v0()?; 55 | let kv = KV::decode(&mut input)?; 56 | 57 | Ok(Tree { 58 | inner: Box::new(TreeInner { 59 | left: maybe_left, 60 | right: maybe_right, 61 | kv, 62 | }), 63 | }) 64 | } 65 | } 66 | 67 | #[cfg(test)] 68 | mod tests { 69 | use super::super::Link; 70 | use super::*; 71 | use crate::error::Result; 72 | 73 | #[test] 74 | fn encode_leaf_tree() { 75 | let tree = Tree::from_fields(vec![0], vec![1], [55; 32], None, None); 76 | assert_eq!(tree.encoding_length(), 35); 77 | assert_eq!( 78 | tree.encode(), 79 | vec![ 80 | 0, 0, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 81 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 1, 82 | ] 83 | ); 84 | } 85 | 86 | #[test] 87 | #[should_panic] 88 | fn encode_modified_tree() { 89 | let tree = Tree::from_fields( 90 | vec![0], 91 | vec![1], 92 | [55; 32], 93 | Some(Link::Modified { 94 | pending_writes: 1, 95 | child_heights: (123, 124), 96 | tree: Tree::new(vec![2], vec![3]).unwrap(), 97 | }), 98 | None, 99 | ); 100 | tree.encode(); 101 | } 102 | 103 | #[test] 104 | fn encode_loaded_tree() -> Result<()> { 105 | let tree = Tree::from_fields( 106 | vec![0], 107 | vec![1], 108 | [55; 32], 109 | Some(Link::Loaded { 110 | hash: [66; 32], 111 | child_heights: (123, 124), 112 | tree: Tree::new(vec![2], vec![3])?, 113 | }), 114 | None, 115 | ); 116 | assert_eq!( 117 | tree.encode(), 118 | vec![ 119 | 1, 0, 1, 2, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 120 | 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 123, 124, 0, 55, 55, 55, 121 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 122 | 55, 55, 55, 55, 55, 55, 55, 55, 1 123 | ] 124 | ); 125 | Ok(()) 126 | } 127 | 128 | #[test] 129 | fn encode_uncommitted_tree() -> Result<()> { 130 | let tree = Tree::from_fields( 131 | vec![0], 132 | vec![1], 133 | [55; 32], 134 | Some(Link::Uncommitted { 135 | hash: [66; 32], 136 | child_heights: (123, 124), 137 | tree: Tree::new(vec![2], vec![3])?, 138 | }), 139 | None, 140 | ); 141 | assert_eq!( 142 | tree.encode(), 143 | vec![ 144 | 1, 0, 1, 2, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 145 | 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 123, 124, 0, 55, 55, 55, 146 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 147 | 55, 55, 55, 55, 55, 55, 55, 55, 1 148 | ] 149 | ); 150 | Ok(()) 151 | } 152 | 153 | #[test] 154 | fn encode_reference_tree() { 155 | let tree = Tree::from_fields( 156 | vec![0], 157 | vec![1], 158 | [55; 32], 159 | Some(Link::Reference { 160 | hash: [66; 32], 161 | child_heights: (123, 124), 162 | key: vec![2], 163 | }), 164 | None, 165 | ); 166 | assert_eq!(tree.encoding_length(), 71); 167 | assert_eq!( 168 | tree.encode(), 169 | vec![ 170 | 1, 0, 1, 2, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 171 | 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 123, 124, 0, 55, 55, 55, 172 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 173 | 55, 55, 55, 55, 55, 55, 55, 55, 1 174 | ] 175 | ); 176 | } 177 | 178 | #[test] 179 | fn decode_leaf_tree() { 180 | let bytes = vec![ 181 | 0, 0, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 182 | 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 183 | ]; 184 | let tree = Tree::decode(vec![0], bytes.as_slice()); 185 | assert_eq!(tree.key(), &[0]); 186 | assert_eq!(tree.value(), &[1]); 187 | } 188 | 189 | #[test] 190 | fn decode_reference_tree() { 191 | let bytes = vec![ 192 | 1, 0, 1, 2, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 193 | 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 66, 123, 124, 0, 55, 55, 55, 55, 55, 194 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 195 | 55, 55, 55, 55, 55, 1, 196 | ]; 197 | let tree = Tree::decode(vec![0], bytes.as_slice()); 198 | assert_eq!(tree.key(), &[0]); 199 | assert_eq!(tree.value(), &[1]); 200 | if let Some(Link::Reference { 201 | key, 202 | child_heights, 203 | hash, 204 | }) = tree.link(true) 205 | { 206 | assert_eq!(*key, [2]); 207 | assert_eq!(*child_heights, (123_u8, 124_u8)); 208 | assert_eq!(*hash, [66_u8; 32]); 209 | } else { 210 | panic!("Expected Link::Reference"); 211 | } 212 | } 213 | } 214 | -------------------------------------------------------------------------------- /src/tree/fuzz_tests.rs: -------------------------------------------------------------------------------- 1 | #![cfg(test)] 2 | 3 | use crate::test_utils::*; 4 | use crate::tree::*; 5 | use rand::prelude::*; 6 | use std::cell::RefCell; 7 | use std::collections::BTreeMap; 8 | use std::iter::FromIterator; 9 | 10 | const ITERATIONS: usize = 2_000; 11 | type Map = BTreeMap, Vec>; 12 | 13 | #[test] 14 | fn fuzz() { 15 | let mut rng = thread_rng(); 16 | 17 | for _ in 0..ITERATIONS { 18 | let seed = rng.gen::(); 19 | fuzz_case(seed); 20 | } 21 | } 22 | 23 | #[test] 24 | fn fuzz_17391518417409062786() { 25 | fuzz_case(17391518417409062786); 26 | } 27 | 28 | #[test] 29 | fn fuzz_396148930387069749() { 30 | fuzz_case(396148930387069749); 31 | } 32 | 33 | fn fuzz_case(seed: u64) { 34 | let mut rng: SmallRng = SeedableRng::seed_from_u64(seed); 35 | let initial_size = (rng.gen::() % 10) + 1; 36 | let tree = make_tree_rand(initial_size, initial_size, seed); 37 | let mut map = Map::from_iter(tree.iter()); 38 | let mut maybe_tree = Some(tree); 39 | println!("====== MERK FUZZ ======"); 40 | println!("SEED: {}", seed); 41 | println!("{:?}", maybe_tree.as_ref().unwrap()); 42 | 43 | for j in 0..3 { 44 | let batch_size = (rng.gen::() % 3) + 1; 45 | let batch = make_batch(maybe_tree.as_ref(), batch_size, rng.gen::()); 46 | println!("BATCH {}", j); 47 | println!("{:?}", batch); 48 | maybe_tree = apply_to_memonly(maybe_tree, &batch); 49 | apply_to_map(&mut map, &batch); 50 | assert_map(maybe_tree.as_ref(), &map); 51 | if let Some(tree) = &maybe_tree { 52 | println!("{:?}", &tree); 53 | } else { 54 | println!("(Empty tree)"); 55 | } 56 | } 57 | } 58 | 59 | fn make_batch(maybe_tree: Option<&Tree>, size: u64, seed: u64) -> Vec { 60 | let rng: RefCell = RefCell::new(SeedableRng::seed_from_u64(seed)); 61 | let mut batch = Vec::with_capacity(size as usize); 62 | 63 | let get_random_key = || { 64 | let tree = maybe_tree.as_ref().unwrap(); 65 | let entries: Vec<_> = tree.iter().collect(); 66 | let index = rng.borrow_mut().gen::() as usize % entries.len(); 67 | entries[index].0.clone() 68 | }; 69 | 70 | let random_value = |size| { 71 | let mut value = vec![0; size]; 72 | rng.borrow_mut().fill_bytes(&mut value[..]); 73 | value 74 | }; 75 | 76 | let insert = || (random_value(2), Op::Put(random_value(2))); 77 | let update = || { 78 | let key = get_random_key(); 79 | (key.to_vec(), Op::Put(random_value(2))) 80 | }; 81 | let delete = || { 82 | let key = get_random_key(); 83 | (key.to_vec(), Op::Delete) 84 | }; 85 | 86 | for _ in 0..size { 87 | let entry = if maybe_tree.is_some() { 88 | let kind = rng.borrow_mut().gen::() % 3; 89 | if kind == 0 { 90 | insert() 91 | } else if kind == 1 { 92 | update() 93 | } else { 94 | delete() 95 | } 96 | } else { 97 | insert() 98 | }; 99 | batch.push(entry); 100 | } 101 | batch.sort_by(|a, b| a.0.cmp(&b.0)); 102 | 103 | // remove dupes 104 | let mut maybe_prev_key: Option> = None; 105 | let mut deduped_batch = Vec::with_capacity(batch.len()); 106 | for entry in batch { 107 | if let Some(prev_key) = &maybe_prev_key { 108 | if *prev_key == entry.0 { 109 | continue; 110 | } 111 | } 112 | 113 | maybe_prev_key = Some(entry.0.clone()); 114 | deduped_batch.push(entry); 115 | } 116 | deduped_batch 117 | } 118 | 119 | fn apply_to_map(map: &mut Map, batch: &Batch) { 120 | for entry in batch.iter() { 121 | match entry { 122 | (key, Op::Put(value)) => { 123 | map.insert(key.to_vec(), value.to_vec()); 124 | } 125 | (key, Op::Delete) => { 126 | map.remove(key); 127 | } 128 | } 129 | } 130 | } 131 | 132 | fn assert_map(maybe_tree: Option<&Tree>, map: &Map) { 133 | if map.is_empty() { 134 | assert!(maybe_tree.is_none(), "expected tree to be None"); 135 | return; 136 | } 137 | 138 | let tree = maybe_tree.expect("expected tree to be Some"); 139 | 140 | let map_iter = map.iter(); 141 | let tree_iter = tree.iter(); 142 | for (tree_kv, map_kv) in tree_iter.zip(map_iter) { 143 | assert_eq!(tree_kv.0, *map_kv.0); 144 | assert_eq!(tree_kv.1, *map_kv.1); 145 | } 146 | 147 | assert_eq!(tree.iter().count(), map.len()); 148 | } 149 | -------------------------------------------------------------------------------- /src/tree/hash.rs: -------------------------------------------------------------------------------- 1 | use sha2::{Digest, Sha512_256}; 2 | use std::{convert::TryFrom, num::TryFromIntError}; 3 | 4 | /// The hash algorithm used for both KV hashes and node hashes. 5 | pub type Hasher = Sha512_256; 6 | 7 | /// The length of a `Hash` (in bytes). 8 | pub const HASH_LENGTH: usize = 32; 9 | 10 | /// A zero-filled `Hash`. 11 | pub const NULL_HASH: Hash = [0; HASH_LENGTH]; 12 | 13 | /// A cryptographic hash digest. 14 | pub type Hash = [u8; HASH_LENGTH]; 15 | 16 | /// Hashes a key/value pair. 17 | pub fn kv_hash(key: &[u8], value: &[u8]) -> Result { 18 | let mut hasher = D::new(); 19 | hasher.update([0]); 20 | 21 | u32::try_from(key.len()) 22 | .and_then(|key| u32::try_from(value.len()).map(|value| (key, value))) 23 | .map(|(key_length, val_length)| { 24 | hasher.update(key_length.to_le_bytes()); 25 | hasher.update(key); 26 | 27 | hasher.update(val_length.to_le_bytes()); 28 | hasher.update(value); 29 | 30 | let res = hasher.finalize(); 31 | let mut hash: Hash = Default::default(); 32 | hash.copy_from_slice(&res[..]); 33 | hash 34 | }) 35 | } 36 | 37 | /// Hashes a node based on the hash of its left child (if any), its key/value 38 | /// pair, and the hash of its right child (if any). 39 | pub fn node_hash(kv: &Hash, left: &Hash, right: &Hash) -> Hash { 40 | let mut hasher = D::new(); 41 | hasher.update([1]); 42 | hasher.update(left); 43 | hasher.update(kv); 44 | hasher.update(right); 45 | 46 | let res = hasher.finalize(); 47 | let mut hash: Hash = Default::default(); 48 | hash.copy_from_slice(&res[..]); 49 | hash 50 | } 51 | -------------------------------------------------------------------------------- /src/tree/iter.rs: -------------------------------------------------------------------------------- 1 | use super::Tree; 2 | 3 | /// An entry stored on an `Iter`'s stack, containing a reference to a `Tree`, 4 | /// and its traversal state. 5 | /// 6 | /// The `traversed` field represents whether or not the left child, self, and 7 | /// right child have been visited, respectively (`(left, self, right)`). 8 | struct StackItem<'a> { 9 | tree: &'a Tree, 10 | traversed: (bool, bool, bool), 11 | } 12 | 13 | impl<'a> StackItem<'a> { 14 | /// Creates a new `StackItem` for the given tree. The `traversed` state will 15 | /// be `false` since the children and self have not been visited yet, but 16 | /// will default to `true` for sides that do not have a child. 17 | fn new(tree: &'a Tree) -> Self { 18 | StackItem { 19 | tree, 20 | traversed: ( 21 | tree.child(true).is_none(), 22 | false, 23 | tree.child(false).is_none(), 24 | ), 25 | } 26 | } 27 | 28 | /// Gets a tuple to yield from an `Iter`, `(key, value)`. 29 | fn to_entry(&self) -> (Vec, Vec) { 30 | (self.tree.key().to_vec(), self.tree.value().to_vec()) 31 | } 32 | } 33 | 34 | /// An iterator which yields the key/value pairs of the tree, in order, skipping 35 | /// any parts of the tree which are pruned (not currently retained in memory). 36 | pub struct Iter<'a> { 37 | stack: Vec>, 38 | } 39 | 40 | impl<'a> Iter<'a> { 41 | /// Creates a new iterator for the given tree. 42 | pub fn new(tree: &'a Tree) -> Self { 43 | let stack = vec![StackItem::new(tree)]; 44 | Iter { stack } 45 | } 46 | } 47 | 48 | impl<'a> Tree { 49 | /// Creates an iterator which yields `(key, value)` tuples for all of the 50 | /// tree's nodes which are retained in memory (skipping pruned subtrees). 51 | pub fn iter(&'a self) -> Iter<'a> { 52 | Iter::new(self) 53 | } 54 | } 55 | 56 | impl<'a> Iterator for Iter<'a> { 57 | type Item = (Vec, Vec); 58 | 59 | /// Traverses to and yields the next key/value pair, in key order. 60 | fn next(&mut self) -> Option { 61 | if self.stack.is_empty() { 62 | return None; 63 | } 64 | 65 | let last = self.stack.last_mut().unwrap(); 66 | if !last.traversed.0 { 67 | last.traversed.0 = true; 68 | let tree = last.tree.child(true).unwrap(); 69 | self.stack.push(StackItem::new(tree)); 70 | self.next() 71 | } else if !last.traversed.1 { 72 | last.traversed.1 = true; 73 | Some(last.to_entry()) 74 | } else if !last.traversed.2 { 75 | last.traversed.2 = true; 76 | let tree = last.tree.child(false).unwrap(); 77 | self.stack.push(StackItem::new(tree)); 78 | self.next() 79 | } else { 80 | self.stack.pop(); 81 | self.next() 82 | } 83 | } 84 | } 85 | -------------------------------------------------------------------------------- /src/tree/kv.rs: -------------------------------------------------------------------------------- 1 | use super::hash::{kv_hash, Hash, Hasher, HASH_LENGTH, NULL_HASH}; 2 | use ed::{Decode, Encode, Result}; 3 | use std::{ 4 | io::{Read, Write}, 5 | num::TryFromIntError, 6 | }; 7 | 8 | // TODO: maybe use something similar to Vec but without capacity field, 9 | // (should save 16 bytes per entry). also, maybe a shorter length 10 | // field to save even more. also might be possible to combine key 11 | // field and value field. 12 | 13 | /// Contains a key/value pair, and the hash of the key/value pair. 14 | #[derive(Clone, Debug, PartialEq, Eq)] 15 | pub struct KV { 16 | pub(super) key: Vec, 17 | pub(super) value: Vec, 18 | pub(super) hash: Hash, 19 | } 20 | 21 | impl KV { 22 | /// Creates a new `KV` with the given key and value and computes its hash. 23 | #[inline] 24 | pub fn new(key: Vec, value: Vec) -> std::result::Result { 25 | kv_hash::(key.as_slice(), value.as_slice()).map(|hash| KV { key, value, hash }) 26 | } 27 | 28 | /// Creates a new `KV` with the given key, value, and hash. The hash is not 29 | /// checked to be correct for the given key/value. 30 | #[inline] 31 | pub fn from_fields(key: Vec, value: Vec, hash: Hash) -> Self { 32 | KV { key, value, hash } 33 | } 34 | 35 | /// Replaces the `KV`'s value with the given value, updates the hash, and 36 | /// returns the modified `KV`. 37 | #[inline] 38 | pub fn with_value(mut self, value: Vec) -> std::result::Result { 39 | self.value = value; 40 | self.hash = kv_hash::(self.key(), self.value())?; 41 | Ok(self) 42 | } 43 | 44 | /// Returns the key as a slice. 45 | #[inline] 46 | pub fn key(&self) -> &[u8] { 47 | self.key.as_slice() 48 | } 49 | 50 | /// Returns the value as a slice. 51 | #[inline] 52 | pub fn value(&self) -> &[u8] { 53 | self.value.as_slice() 54 | } 55 | 56 | /// Returns the hash. 57 | #[inline] 58 | pub fn hash(&self) -> &Hash { 59 | &self.hash 60 | } 61 | 62 | /// Consumes the `KV` and returns its key without allocating or cloning. 63 | #[inline] 64 | pub fn take_key(self) -> Vec { 65 | self.key 66 | } 67 | } 68 | 69 | impl Encode for KV { 70 | #[inline] 71 | fn encode_into(&self, out: &mut W) -> Result<()> { 72 | out.write_all(&self.hash[..])?; 73 | out.write_all(self.value.as_slice())?; 74 | Ok(()) 75 | } 76 | 77 | #[inline] 78 | fn encoding_length(&self) -> Result { 79 | debug_assert!( 80 | self.key().len() < 65536, 81 | "Key length must be less than 65536" 82 | ); 83 | 84 | Ok(HASH_LENGTH + self.value.len()) 85 | } 86 | } 87 | 88 | impl Decode for KV { 89 | #[inline] 90 | fn decode(input: R) -> Result { 91 | let mut kv = KV { 92 | key: Vec::with_capacity(0), 93 | value: Vec::with_capacity(128), 94 | hash: NULL_HASH, 95 | }; 96 | KV::decode_into(&mut kv, input)?; 97 | Ok(kv) 98 | } 99 | 100 | #[inline] 101 | fn decode_into(&mut self, mut input: R) -> Result<()> { 102 | self.key.clear(); 103 | 104 | input.read_exact(&mut self.hash[..])?; 105 | 106 | self.value.clear(); 107 | input.read_to_end(self.value.as_mut())?; 108 | 109 | Ok(()) 110 | } 111 | } 112 | 113 | #[cfg(test)] 114 | mod test { 115 | use super::*; 116 | 117 | #[test] 118 | fn new_kv() -> std::result::Result<(), TryFromIntError> { 119 | let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6])?; 120 | 121 | assert_eq!(kv.key(), &[1, 2, 3]); 122 | assert_eq!(kv.value(), &[4, 5, 6]); 123 | assert_ne!(kv.hash(), &super::super::hash::NULL_HASH); 124 | Ok(()) 125 | } 126 | 127 | #[test] 128 | fn with_value() -> std::result::Result<(), TryFromIntError> { 129 | let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6])?.with_value(vec![7, 8, 9])?; 130 | 131 | assert_eq!(kv.key(), &[1, 2, 3]); 132 | assert_eq!(kv.value(), &[7, 8, 9]); 133 | assert_ne!(kv.hash(), &super::super::hash::NULL_HASH); 134 | Ok(()) 135 | } 136 | } 137 | -------------------------------------------------------------------------------- /src/tree/link.rs: -------------------------------------------------------------------------------- 1 | use std::cmp::max; 2 | use std::io::{Read, Write}; 3 | 4 | use ed::{Decode, Encode, Result, Terminated}; 5 | 6 | use super::hash::Hash; 7 | use super::Tree; 8 | 9 | // TODO: optimize memory footprint 10 | 11 | /// Represents a reference to a child tree node. Links may or may not contain 12 | /// the child's `Tree` instance (storing its key if not). 13 | #[derive(Debug, Clone, PartialEq, Eq)] 14 | pub enum Link { 15 | /// Represents a child tree node which has been pruned from memory, only 16 | /// retaining a reference to it (its key). The child node can always be 17 | /// fetched from the backing store by this key when necessary. 18 | Reference { 19 | hash: Hash, 20 | child_heights: (u8, u8), 21 | key: Vec, 22 | }, 23 | 24 | /// Represents a tree node which has been modified since the `Tree`'s last 25 | /// hash computation. The child's hash is not stored since it has not yet 26 | /// been recomputed. The child's `Tree` instance is stored in the link. 27 | Modified { 28 | pending_writes: usize, // TODO: rename to `pending_hashes` 29 | child_heights: (u8, u8), 30 | tree: Tree, 31 | }, 32 | 33 | // Represents a tree node which has been modified since the `Tree`'s last 34 | // commit, but which has an up-to-date hash. The child's `Tree` instance is 35 | // stored in the link. 36 | Uncommitted { 37 | hash: Hash, 38 | child_heights: (u8, u8), 39 | tree: Tree, 40 | }, 41 | 42 | /// Represents a tree node which has not been modified, has an up-to-date 43 | /// hash, and which is being retained in memory. 44 | Loaded { 45 | hash: Hash, 46 | child_heights: (u8, u8), 47 | tree: Tree, 48 | }, 49 | } 50 | 51 | impl Link { 52 | /// Creates a `Link::Modified` from the given `Tree`. 53 | #[inline] 54 | pub fn from_modified_tree(tree: Tree) -> Self { 55 | let pending_writes = 1 + tree.child_pending_writes(true) + tree.child_pending_writes(false); 56 | 57 | Link::Modified { 58 | pending_writes, 59 | child_heights: tree.child_heights(), 60 | tree, 61 | } 62 | } 63 | 64 | /// Creates a `Link::Modified` from the given tree, if any. If `None`, 65 | /// returns `None`. 66 | pub fn maybe_from_modified_tree(maybe_tree: Option) -> Option { 67 | maybe_tree.map(Link::from_modified_tree) 68 | } 69 | 70 | /// Returns `true` if the link is of the `Link::Reference` variant. 71 | #[inline] 72 | pub fn is_reference(&self) -> bool { 73 | matches!(self, Link::Reference { .. }) 74 | } 75 | 76 | /// Returns `true` if the link is of the `Link::Modified` variant. 77 | #[inline] 78 | pub fn is_modified(&self) -> bool { 79 | matches!(self, Link::Modified { .. }) 80 | } 81 | 82 | /// Returns `true` if the link is of the `Link::Uncommitted` variant. 83 | #[inline] 84 | pub fn is_uncommitted(&self) -> bool { 85 | matches!(self, Link::Uncommitted { .. }) 86 | } 87 | 88 | /// Returns `true` if the link is of the `Link::Loaded` variant. 89 | #[inline] 90 | pub fn is_stored(&self) -> bool { 91 | matches!(self, Link::Loaded { .. }) 92 | } 93 | 94 | /// Returns the key of the tree referenced by this link, as a slice. 95 | #[inline] 96 | pub fn key(&self) -> &[u8] { 97 | match self { 98 | Link::Reference { key, .. } => key.as_slice(), 99 | Link::Modified { tree, .. } => tree.key(), 100 | Link::Uncommitted { tree, .. } => tree.key(), 101 | Link::Loaded { tree, .. } => tree.key(), 102 | } 103 | } 104 | 105 | /// Returns the `Tree` instance of the tree referenced by the link. If the 106 | /// link is of variant `Link::Reference`, the returned value will be `None`. 107 | #[inline] 108 | pub fn tree(&self) -> Option<&Tree> { 109 | match self { 110 | // TODO: panic for Reference, don't return Option? 111 | Link::Reference { .. } => None, 112 | Link::Modified { tree, .. } => Some(tree), 113 | Link::Uncommitted { tree, .. } => Some(tree), 114 | Link::Loaded { tree, .. } => Some(tree), 115 | } 116 | } 117 | 118 | /// Returns the hash of the tree referenced by the link. Panics if link is 119 | /// of variant `Link::Modified` since we have not yet recomputed the tree's 120 | /// hash. 121 | #[inline] 122 | pub fn hash(&self) -> &Hash { 123 | match self { 124 | Link::Modified { .. } => panic!("Cannot get hash from modified link"), 125 | Link::Reference { hash, .. } => hash, 126 | Link::Uncommitted { hash, .. } => hash, 127 | Link::Loaded { hash, .. } => hash, 128 | } 129 | } 130 | 131 | /// Returns the height of the children of the tree referenced by the link, 132 | /// if any (note: not the height of the referenced tree itself). Return 133 | /// value is `(left_child_height, right_child_height)`. 134 | #[inline] 135 | pub fn height(&self) -> u8 { 136 | let (left_height, right_height) = match self { 137 | Link::Reference { child_heights, .. } => *child_heights, 138 | Link::Modified { child_heights, .. } => *child_heights, 139 | Link::Uncommitted { child_heights, .. } => *child_heights, 140 | Link::Loaded { child_heights, .. } => *child_heights, 141 | }; 142 | 1 + max(left_height, right_height) 143 | } 144 | 145 | /// Returns the balance factor of the tree referenced by the link. 146 | #[inline] 147 | pub fn balance_factor(&self) -> i8 { 148 | let (left_height, right_height) = match self { 149 | Link::Reference { child_heights, .. } => *child_heights, 150 | Link::Modified { child_heights, .. } => *child_heights, 151 | Link::Uncommitted { child_heights, .. } => *child_heights, 152 | Link::Loaded { child_heights, .. } => *child_heights, 153 | }; 154 | right_height as i8 - left_height as i8 155 | } 156 | 157 | /// Consumes the link and converts to variant `Link::Reference`. Panics if 158 | /// the link is of variant `Link::Modified` or `Link::Uncommitted`. 159 | #[inline] 160 | pub fn into_reference(self) -> Self { 161 | match self { 162 | Link::Reference { .. } => self, 163 | Link::Modified { .. } => panic!("Cannot prune Modified tree"), 164 | Link::Uncommitted { .. } => panic!("Cannot prune Uncommitted tree"), 165 | Link::Loaded { 166 | hash, 167 | child_heights, 168 | tree, 169 | } => Link::Reference { 170 | hash, 171 | child_heights, 172 | key: tree.take_key(), 173 | }, 174 | } 175 | } 176 | 177 | #[inline] 178 | #[cfg(feature = "full")] 179 | pub(crate) fn child_heights_mut(&mut self) -> &mut (u8, u8) { 180 | match self { 181 | Link::Reference { 182 | ref mut child_heights, 183 | .. 184 | } => child_heights, 185 | Link::Modified { 186 | ref mut child_heights, 187 | .. 188 | } => child_heights, 189 | Link::Uncommitted { 190 | ref mut child_heights, 191 | .. 192 | } => child_heights, 193 | Link::Loaded { 194 | ref mut child_heights, 195 | .. 196 | } => child_heights, 197 | } 198 | } 199 | } 200 | 201 | impl Encode for Link { 202 | #[inline] 203 | fn encode_into(&self, out: &mut W) -> Result<()> { 204 | let (hash, key, (left_height, right_height)) = match self { 205 | Link::Reference { 206 | hash, 207 | key, 208 | child_heights, 209 | } => (hash, key.as_slice(), child_heights), 210 | Link::Loaded { 211 | hash, 212 | tree, 213 | child_heights, 214 | } => (hash, tree.key(), child_heights), 215 | Link::Uncommitted { 216 | hash, 217 | tree, 218 | child_heights, 219 | } => (hash, tree.key(), child_heights), 220 | 221 | Link::Modified { .. } => panic!("No encoding for Link::Modified"), 222 | }; 223 | 224 | debug_assert!( 225 | self.key().len() < 65536, 226 | "Key length must be less than 65536" 227 | ); 228 | 229 | out.write_all(&(key.len() as u16).to_be_bytes())?; 230 | out.write_all(key)?; 231 | 232 | out.write_all(hash)?; 233 | 234 | out.write_all(&[*left_height, *right_height])?; 235 | 236 | Ok(()) 237 | } 238 | 239 | #[inline] 240 | fn encoding_length(&self) -> Result { 241 | debug_assert!( 242 | self.key().len() < 65536, 243 | "Key length must be less than 65536" 244 | ); 245 | 246 | Ok(match self { 247 | Link::Reference { key, .. } => 1 + key.len() + 32 + 2, 248 | Link::Modified { .. } => panic!("No encoding for Link::Modified"), 249 | Link::Uncommitted { tree, .. } => 1 + tree.key().len() + 32 + 2, 250 | Link::Loaded { tree, .. } => 1 + tree.key().len() + 32 + 2, 251 | }) 252 | } 253 | } 254 | 255 | impl Link { 256 | #[inline] 257 | fn default_reference() -> Self { 258 | Link::Reference { 259 | key: Vec::with_capacity(64), 260 | hash: Default::default(), 261 | child_heights: (0, 0), 262 | } 263 | } 264 | 265 | pub(crate) fn decode_v0(mut input: R) -> Result { 266 | let length = read_u8(&mut input)? as usize; 267 | 268 | let mut key = vec![0; length]; 269 | input.read_exact(&mut key)?; 270 | 271 | let mut hash = [0; 32]; 272 | input.read_exact(&mut hash)?; 273 | 274 | let left_height = read_u8(&mut input)?; 275 | let right_height = read_u8(input)?; 276 | 277 | Ok(Link::Reference { 278 | key, 279 | hash, 280 | child_heights: (left_height, right_height), 281 | }) 282 | } 283 | } 284 | 285 | impl Decode for Link { 286 | #[inline] 287 | fn decode(input: R) -> Result { 288 | let mut link = Link::default_reference(); 289 | Link::decode_into(&mut link, input)?; 290 | Ok(link) 291 | } 292 | 293 | #[inline] 294 | fn decode_into(&mut self, mut input: R) -> Result<()> { 295 | if !self.is_reference() { 296 | // don't create new struct if self is already Link::Reference, 297 | // so we can re-use the key vec 298 | *self = Link::default_reference(); 299 | } 300 | 301 | if let Link::Reference { 302 | ref mut key, 303 | ref mut hash, 304 | ref mut child_heights, 305 | } = self 306 | { 307 | let length = read_u16(&mut input)? as usize; 308 | 309 | key.resize(length, 0); 310 | input.read_exact(key.as_mut())?; 311 | 312 | input.read_exact(&mut hash[..])?; 313 | 314 | child_heights.0 = read_u8(&mut input)?; 315 | child_heights.1 = read_u8(&mut input)?; 316 | } else { 317 | unreachable!() 318 | } 319 | 320 | Ok(()) 321 | } 322 | } 323 | 324 | impl Terminated for Link {} 325 | 326 | #[inline] 327 | fn read_u16(mut input: R) -> Result { 328 | let mut length = [0, 0]; 329 | input.read_exact(length.as_mut())?; 330 | Ok(u16::from_be_bytes(length)) 331 | } 332 | 333 | #[inline] 334 | fn read_u8(mut input: R) -> Result { 335 | let mut length = [0]; 336 | input.read_exact(length.as_mut())?; 337 | Ok(length[0]) 338 | } 339 | 340 | #[cfg(test)] 341 | mod test { 342 | use super::super::hash::NULL_HASH; 343 | use super::super::Tree; 344 | use super::*; 345 | 346 | #[test] 347 | fn from_modified_tree() -> std::result::Result<(), &'static str> { 348 | let tree = Tree::new(vec![0], vec![1]).map_err(|_| "tree construction failed")?; 349 | let link = Link::from_modified_tree(tree); 350 | assert!(link.is_modified()); 351 | assert_eq!(link.height(), 1); 352 | assert_eq!(link.tree().expect("expected tree").key(), &[0]); 353 | if let Link::Modified { pending_writes, .. } = link { 354 | assert_eq!(pending_writes, 1); 355 | Ok(()) 356 | } else { 357 | Err("Expected Link::Modified") 358 | } 359 | } 360 | 361 | #[test] 362 | fn maybe_from_modified_tree() -> std::result::Result<(), crate::error::Error> { 363 | let link = Link::maybe_from_modified_tree(None); 364 | assert!(link.is_none()); 365 | 366 | let tree = Tree::new(vec![0], vec![1])?; 367 | let link = Link::maybe_from_modified_tree(Some(tree)); 368 | assert!(link.expect("expected link").is_modified()); 369 | Ok(()) 370 | } 371 | 372 | #[test] 373 | fn types() -> std::result::Result<(), crate::error::Error> { 374 | let hash = NULL_HASH; 375 | let child_heights = (0, 0); 376 | let pending_writes = 1; 377 | let key = vec![0]; 378 | let tree = || Tree::new(vec![0], vec![1]); 379 | 380 | let reference = Link::Reference { 381 | hash, 382 | child_heights, 383 | key, 384 | }; 385 | let modified = Link::Modified { 386 | pending_writes, 387 | child_heights, 388 | tree: tree()?, 389 | }; 390 | let uncommitted = Link::Uncommitted { 391 | hash, 392 | child_heights, 393 | tree: tree()?, 394 | }; 395 | let loaded = Link::Loaded { 396 | hash, 397 | child_heights, 398 | tree: tree()?, 399 | }; 400 | 401 | assert!(reference.is_reference()); 402 | assert!(!reference.is_modified()); 403 | assert!(!reference.is_uncommitted()); 404 | assert!(!reference.is_stored()); 405 | assert!(reference.tree().is_none()); 406 | assert_eq!(reference.hash(), &[0; 32]); 407 | assert_eq!(reference.height(), 1); 408 | assert!(reference.into_reference().is_reference()); 409 | 410 | assert!(!modified.is_reference()); 411 | assert!(modified.is_modified()); 412 | assert!(!modified.is_uncommitted()); 413 | assert!(!modified.is_stored()); 414 | assert!(modified.tree().is_some()); 415 | assert_eq!(modified.height(), 1); 416 | 417 | assert!(!uncommitted.is_reference()); 418 | assert!(!uncommitted.is_modified()); 419 | assert!(uncommitted.is_uncommitted()); 420 | assert!(!uncommitted.is_stored()); 421 | assert!(uncommitted.tree().is_some()); 422 | assert_eq!(uncommitted.hash(), &[0; 32]); 423 | assert_eq!(uncommitted.height(), 1); 424 | 425 | assert!(!loaded.is_reference()); 426 | assert!(!loaded.is_modified()); 427 | assert!(!loaded.is_uncommitted()); 428 | assert!(loaded.is_stored()); 429 | assert!(loaded.tree().is_some()); 430 | assert_eq!(loaded.hash(), &[0; 32]); 431 | assert_eq!(loaded.height(), 1); 432 | assert!(loaded.into_reference().is_reference()); 433 | Ok(()) 434 | } 435 | 436 | #[test] 437 | #[should_panic(expected = "Cannot get hash from modified link")] 438 | fn modified_hash() { 439 | Tree::new(vec![0], vec![1]) 440 | .map(|tree| Link::Modified { 441 | pending_writes: 1, 442 | child_heights: (1, 1), 443 | tree, 444 | }) 445 | .map(|link| link.hash().to_vec()) 446 | .map(|_| ()) 447 | .unwrap_or_default() 448 | } 449 | 450 | #[test] 451 | #[should_panic] 452 | fn modified_into_reference() { 453 | Link::Modified { 454 | pending_writes: 1, 455 | child_heights: (1, 1), 456 | tree: Tree::new(vec![0], vec![1]).expect("tree construction failed"), 457 | } 458 | .into_reference(); 459 | } 460 | 461 | #[test] 462 | #[should_panic] 463 | fn uncommitted_into_reference() { 464 | Link::Uncommitted { 465 | hash: [1; 32], 466 | child_heights: (1, 1), 467 | tree: Tree::new(vec![0], vec![1]).expect("tree construction failed"), 468 | } 469 | .into_reference(); 470 | } 471 | 472 | #[test] 473 | fn encode_link() { 474 | let link = Link::Reference { 475 | key: vec![1, 2, 3], 476 | child_heights: (123, 124), 477 | hash: [55; 32], 478 | }; 479 | assert_eq!(link.encoding_length().unwrap(), 38); 480 | 481 | let mut bytes = vec![]; 482 | link.encode_into(&mut bytes).unwrap(); 483 | assert_eq!( 484 | bytes, 485 | vec![ 486 | 0, 3, 1, 2, 3, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 487 | 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 55, 123, 124 488 | ] 489 | ); 490 | } 491 | 492 | #[test] 493 | fn encode_link_long_key_valid() { 494 | let link = Link::Reference { 495 | key: vec![123; 60_000], 496 | child_heights: (123, 124), 497 | hash: [55; 32], 498 | }; 499 | let mut bytes = vec![]; 500 | link.encode_into(&mut bytes).unwrap(); 501 | 502 | let decoded = Link::decode(&bytes[..]).unwrap(); 503 | assert_eq!(decoded, link); 504 | } 505 | 506 | #[test] 507 | #[should_panic = "Key length must be less than 65536"] 508 | fn encode_link_long_key_invalid() { 509 | let link = Link::Reference { 510 | key: vec![123; 70_000], 511 | child_heights: (123, 124), 512 | hash: [55; 32], 513 | }; 514 | let mut bytes = vec![]; 515 | link.encode_into(&mut bytes).unwrap(); 516 | } 517 | } 518 | -------------------------------------------------------------------------------- /src/tree/walk/fetch.rs: -------------------------------------------------------------------------------- 1 | use super::super::{Link, Tree}; 2 | use crate::error::{Error, Result}; 3 | 4 | /// A source of data to be used by the tree when encountering a pruned node. 5 | /// 6 | /// This typcially means fetching the tree node from a backing store by its key, 7 | /// but could also implement an in-memory cache for example. 8 | pub trait Fetch { 9 | fn fetch_by_key(&self, key: &[u8]) -> Result>; 10 | 11 | /// Called when the tree needs to fetch a node with the given `Link`. The 12 | /// `link` value will always be a `Link::Reference` variant. 13 | fn fetch(&self, link: &Link) -> Result { 14 | self.fetch_by_key_expect(link.key()) 15 | } 16 | 17 | fn fetch_by_key_expect(&self, key: &[u8]) -> Result { 18 | self.fetch_by_key(key)? 19 | .ok_or_else(|| Error::Key(format!("Key does not exist: {key:?}"))) 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /src/tree/walk/mod.rs: -------------------------------------------------------------------------------- 1 | mod fetch; 2 | mod ref_walker; 3 | 4 | use super::{Link, Tree}; 5 | use crate::error::Result; 6 | use crate::owner::Owner; 7 | pub use fetch::Fetch; 8 | pub use ref_walker::RefWalker; 9 | 10 | /// Allows traversal of a `Tree`, fetching from the given source when traversing 11 | /// to a pruned node, detaching children as they are traversed. 12 | pub struct Walker 13 | where 14 | S: Fetch + Sized + Clone + Send, 15 | { 16 | tree: Owner, 17 | source: S, 18 | } 19 | 20 | impl Walker 21 | where 22 | S: Fetch + Sized + Clone + Send, 23 | { 24 | /// Creates a `Walker` with the given tree and source. 25 | pub fn new(tree: Tree, source: S) -> Self { 26 | Walker { 27 | tree: Owner::new(tree), 28 | source, 29 | } 30 | } 31 | 32 | /// Similar to `Tree#detach`, but yields a `Walker` which fetches from the 33 | /// same source as `self`. Returned tuple is `(updated_self, 34 | /// maybe_child_walker)`. 35 | pub fn detach(mut self, left: bool) -> Result<(Self, Option)> { 36 | let link = match self.tree.link(left) { 37 | None => return Ok((self, None)), 38 | Some(link) => link, 39 | }; 40 | 41 | let child = if link.tree().is_some() { 42 | match self.tree.own_return(|t| t.detach(left)) { 43 | Some(child) => child, 44 | _ => unreachable!("Expected Some"), 45 | } 46 | } else { 47 | let link = self.tree.slot_mut(left).take(); 48 | match link { 49 | Some(Link::Reference { .. }) => (), 50 | _ => unreachable!("Expected Some(Link::Reference)"), 51 | } 52 | self.source.fetch(&link.unwrap())? 53 | }; 54 | 55 | let child = self.wrap(child); 56 | Ok((self, Some(child))) 57 | } 58 | 59 | /// Similar to `Tree#detach_expect`, but yields a `Walker` which fetches 60 | /// from the same source as `self`. Returned tuple is `(updated_self, 61 | /// child_walker)`. 62 | pub fn detach_expect(self, left: bool) -> Result<(Self, Self)> { 63 | let (walker, maybe_child) = self.detach(left)?; 64 | if let Some(child) = maybe_child { 65 | Ok((walker, child)) 66 | } else { 67 | panic!( 68 | "Expected {} child, got None", 69 | if left { "left" } else { "right" } 70 | ); 71 | } 72 | } 73 | 74 | /// Similar to `Tree#walk`, but yields a `Walker` which fetches from the 75 | /// same source as `self`. 76 | pub fn walk(self, left: bool, f: F) -> Result 77 | where 78 | F: FnOnce(Option) -> Result>, 79 | T: Into, 80 | { 81 | let (mut walker, maybe_child) = self.detach(left)?; 82 | let new_child = f(maybe_child)?.map(|t| t.into()); 83 | walker.tree.own(|t| t.attach(left, new_child)); 84 | Ok(walker) 85 | } 86 | 87 | /// Similar to `Tree#walk_expect` but yields a `Walker` which fetches from 88 | /// the same source as `self`. 89 | pub fn walk_expect(self, left: bool, f: F) -> Result 90 | where 91 | F: FnOnce(Self) -> Result>, 92 | T: Into, 93 | { 94 | let (mut walker, child) = self.detach_expect(left)?; 95 | let new_child = f(child)?.map(|t| t.into()); 96 | walker.tree.own(|t| t.attach(left, new_child)); 97 | Ok(walker) 98 | } 99 | 100 | /// Returns an immutable reference to the `Tree` wrapped by this walker. 101 | pub fn tree(&self) -> &Tree { 102 | &self.tree 103 | } 104 | 105 | /// Consumes the `Walker` and returns the `Tree` it wraps. 106 | pub fn into_inner(self) -> Tree { 107 | self.tree.into_inner() 108 | } 109 | 110 | /// Takes a `Tree` and returns a `Walker` which fetches from the same source 111 | /// as `self`. 112 | fn wrap(&self, tree: Tree) -> Self { 113 | Walker::new(tree, self.source.clone()) 114 | } 115 | 116 | /// Returns a clone of this `Walker`'s source. 117 | pub fn clone_source(&self) -> S { 118 | self.source.clone() 119 | } 120 | 121 | /// Similar to `Tree#attach`, but can also take a `Walker` since it 122 | /// implements `Into`. 123 | pub fn attach(mut self, left: bool, maybe_child: Option) -> Self 124 | where 125 | T: Into, 126 | { 127 | self.tree 128 | .own(|t| t.attach(left, maybe_child.map(|t| t.into()))); 129 | self 130 | } 131 | 132 | /// Similar to `Tree#with_value`. 133 | pub fn with_value(mut self, value: Vec) -> Result { 134 | self.tree.own_fallible(|t| t.with_value(value))?; 135 | Ok(self) 136 | } 137 | } 138 | 139 | impl From> for Tree 140 | where 141 | S: Fetch + Sized + Clone + Send, 142 | { 143 | fn from(walker: Walker) -> Tree { 144 | walker.into_inner() 145 | } 146 | } 147 | 148 | #[cfg(test)] 149 | mod test { 150 | use super::super::NoopCommit; 151 | use super::*; 152 | use crate::tree::Tree; 153 | 154 | #[derive(Clone)] 155 | struct MockSource {} 156 | 157 | impl Fetch for MockSource { 158 | fn fetch_by_key(&self, key: &[u8]) -> Result> { 159 | Tree::new(key.to_vec(), b"foo".to_vec()).map(Some) 160 | } 161 | } 162 | 163 | #[test] 164 | fn walk_modified() -> Result<()> { 165 | let tree = Tree::new(b"test".to_vec(), b"abc".to_vec())? 166 | .attach(true, Some(Tree::new(b"foo".to_vec(), b"bar".to_vec())?)); 167 | 168 | let source = MockSource {}; 169 | let walker = Walker::new(tree, source); 170 | 171 | let walker = walker 172 | .walk(true, |child| -> Result> { 173 | assert_eq!(child.expect("should have child").tree().key(), b"foo"); 174 | Ok(None) 175 | }) 176 | .expect("walk failed"); 177 | assert!(walker.into_inner().child(true).is_none()); 178 | Ok(()) 179 | } 180 | 181 | #[test] 182 | fn walk_stored() -> Result<()> { 183 | let mut tree = Tree::new(b"test".to_vec(), b"abc".to_vec())? 184 | .attach(true, Some(Tree::new(b"foo".to_vec(), b"bar".to_vec())?)); 185 | tree.commit(&mut NoopCommit {}).expect("commit failed"); 186 | 187 | let source = MockSource {}; 188 | let walker = Walker::new(tree, source); 189 | 190 | let walker = walker 191 | .walk(true, |child| -> Result> { 192 | assert_eq!(child.expect("should have child").tree().key(), b"foo"); 193 | Ok(None) 194 | }) 195 | .expect("walk failed"); 196 | assert!(walker.into_inner().child(true).is_none()); 197 | Ok(()) 198 | } 199 | 200 | #[test] 201 | fn walk_pruned() { 202 | let tree = Tree::from_fields( 203 | b"test".to_vec(), 204 | b"abc".to_vec(), 205 | Default::default(), 206 | Some(Link::Reference { 207 | hash: Default::default(), 208 | key: b"foo".to_vec(), 209 | child_heights: (0, 0), 210 | }), 211 | None, 212 | ); 213 | 214 | let source = MockSource {}; 215 | let walker = Walker::new(tree, source); 216 | 217 | let walker = walker 218 | .walk_expect(true, |child| -> Result> { 219 | assert_eq!(child.tree().key(), b"foo"); 220 | Ok(None) 221 | }) 222 | .expect("walk failed"); 223 | assert!(walker.into_inner().child(true).is_none()); 224 | } 225 | 226 | #[test] 227 | fn walk_none() -> Result<()> { 228 | let tree = Tree::new(b"test".to_vec(), b"abc".to_vec())?; 229 | 230 | let source = MockSource {}; 231 | let walker = Walker::new(tree, source); 232 | 233 | walker 234 | .walk(true, |child| -> Result> { 235 | assert!(child.is_none()); 236 | Ok(None) 237 | }) 238 | .expect("walk failed"); 239 | Ok(()) 240 | } 241 | } 242 | -------------------------------------------------------------------------------- /src/tree/walk/ref_walker.rs: -------------------------------------------------------------------------------- 1 | use super::super::{Link, Tree}; 2 | use super::Fetch; 3 | use crate::error::Result; 4 | 5 | /// Allows read-only traversal of a `Tree`, fetching from the given source when 6 | /// traversing to a pruned node. 7 | /// 8 | /// The fetched nodes are in memory until they (possibly) get pruned on the next 9 | /// commit. 10 | /// 11 | /// Only finalized trees may be walked (trees which have had `commit` called 12 | /// since the last update). 13 | pub struct RefWalker<'a, S> 14 | where 15 | S: Fetch + Sized + Clone + Send, 16 | { 17 | tree: &'a mut Tree, 18 | source: S, 19 | } 20 | 21 | impl<'a, S> RefWalker<'a, S> 22 | where 23 | S: Fetch + Sized + Clone + Send, 24 | { 25 | /// Creates a `RefWalker` with the given tree and source. 26 | pub fn new(tree: &'a mut Tree, source: S) -> Self { 27 | // TODO: check if tree has modified links, panic if so 28 | RefWalker { tree, source } 29 | } 30 | 31 | /// Gets an immutable reference to the `Tree` wrapped by this `RefWalker`. 32 | pub fn tree(&self) -> &Tree { 33 | self.tree 34 | } 35 | 36 | /// Traverses to the child on the given side (if any), fetching from the 37 | /// source if pruned. When fetching, the link is upgraded from 38 | /// `Link::Reference` to `Link::Loaded`. 39 | pub fn walk(&mut self, left: bool) -> Result>> { 40 | let link = match self.tree.link(left) { 41 | None => return Ok(None), 42 | Some(link) => link, 43 | }; 44 | 45 | match link { 46 | Link::Reference { .. } => { 47 | self.tree.load(left, &self.source)?; 48 | } 49 | Link::Modified { .. } => panic!("Cannot traverse Link::Modified"), 50 | Link::Uncommitted { .. } | Link::Loaded { .. } => {} 51 | } 52 | 53 | let child = self.tree.child_mut(left).unwrap(); 54 | Ok(Some(RefWalker::new(child, self.source.clone()))) 55 | } 56 | } 57 | --------------------------------------------------------------------------------