├── .github └── workflows │ ├── coverage.yml │ └── test.yml ├── .gitignore ├── Cargo.toml ├── LICENSE-APACHE ├── LICENSE-MIT ├── Readme.md ├── benches └── benchmark.rs ├── codecov.yml ├── examples └── random_access.rs └── src ├── fmt.rs ├── fuzzing ├── basic.rs ├── mod.rs ├── set_ops.rs └── traversals.rs ├── inner.rs ├── lib.rs ├── map ├── entry.rs ├── iter.rs └── mod.rs ├── prefix.rs ├── serde.rs ├── set.rs ├── set_backup ├── test.rs └── trieview ├── difference.rs ├── intersection.rs ├── mod.rs └── union.rs /.github/workflows/coverage.yml: -------------------------------------------------------------------------------- 1 | name: Coverage 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | RUSTFLAGS: "-Cinstrument-coverage" 12 | RUSTDOCFLAGS: "-Cinstrument-coverage -Z unstable-options --persist-doctests target/debug/doctestbins" 13 | 14 | jobs: 15 | build: 16 | 17 | runs-on: ubuntu-latest 18 | 19 | steps: 20 | - uses: actions/checkout@v3 21 | - name: Install latest nightly 22 | uses: actions-rs/toolchain@v1 23 | with: 24 | toolchain: nightly 25 | override: true 26 | components: rustfmt, clippy, llvm-tools-preview 27 | - name: install grcov 28 | run: cargo install grcov 29 | - name: run tests 30 | run: LLVM_PROFILE_FILE="your_name-%p-%m.profraw" cargo test --verbose 31 | - name: Generate coverage report 32 | run: grcov . --binary-path ./target/debug/ -s . -t lcov --branch --ignore-not-existing --ignore "examples" --ignore "/*" -o lcov.info 33 | - name: Upload coverage 34 | uses: codecov/codecov-action@v5 35 | env: 36 | CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} 37 | with: 38 | files: lcov.info 39 | fail_ci_if_error: false 40 | verbose: true 41 | 42 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Rust 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | test: 14 | runs-on: ubuntu-latest 15 | strategy: 16 | matrix: 17 | toolchain: 18 | - stable 19 | - beta 20 | - nightly 21 | steps: 22 | - uses: actions/checkout@v3 23 | - name: Install rustup 24 | uses: actions-rs/toolchain@v1 25 | with: 26 | toolchain: ${{ matrix.toolchain }} 27 | override: true 28 | components: rustfmt, clippy 29 | - name: Install cargo-hack 30 | run: cargo install cargo-hack 31 | 32 | - name: Run clippy 33 | run: cargo hack --feature-powerset clippy --verbose 34 | - name: Run rustfmt 35 | run: cargo fmt --check 36 | - name: Run tests 37 | run: cargo hack --feature-powerset test --verbose --lib --release 38 | - name: Run doctests 39 | run: cargo hack --feature-powerset test --verbose --doc 40 | 41 | msrv: 42 | runs-on: ubuntu-latest 43 | steps: 44 | - uses: actions/checkout@v3 45 | - name: Install rustup 46 | uses: actions-rs/toolchain@v1 47 | with: 48 | toolchain: stable 49 | override: true 50 | - name: Install cargo-hack 51 | run: cargo install cargo-msrv 52 | - name: Check rust version 53 | run: cargo msrv verify 54 | 55 | miri: 56 | runs-on: ubuntu-latest 57 | steps: 58 | - uses: actions/checkout@v3 59 | - name: Install rustup 60 | uses: actions-rs/toolchain@v1 61 | with: 62 | toolchain: nightly 63 | override: true 64 | components: miri 65 | - name: Install nextest 66 | run: cargo install cargo-nextest 67 | - name: Run miri on libtests 68 | run: cargo miri nextest run -j4 69 | - name: Run miri on doctests 70 | run: cargo miri test --doc 71 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | Cargo.lock 2 | target 3 | coverage 4 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "prefix-trie" 3 | version = "0.7.0" 4 | edition = "2021" 5 | license = "MIT OR Apache-2.0" 6 | repository = "https://github.com/tiborschneider/prefix-trie" 7 | homepage = "https://github.com/tiborschneider/prefix-trie" 8 | documentation = "https://docs.rs/prefix-trie" 9 | readme = "Readme.md" 10 | categories = ["network-programming"] 11 | keywords = ["IP", "prefix", "trie", "tree", "collection"] 12 | description = "Prefix trie (tree) datastructure (both a set and a map) that provides exact and longest-prefix matches." 13 | rust-version = "1.65.0" 14 | 15 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 16 | 17 | [[bench]] 18 | name = "benchmark" 19 | harness = false 20 | required-features = ["ipnet"] 21 | 22 | [[example]] 23 | name = "random_access" 24 | required-features = ["ipnet"] 25 | 26 | [features] 27 | default = ["ipnet"] 28 | serde = ["dep:serde"] 29 | ipnetwork = ["dep:ipnetwork"] 30 | ipnet = ["dep:ipnet"] 31 | cidr = ["dep:cidr"] 32 | 33 | [dependencies] 34 | ipnet = { version = "2", optional = true} 35 | ipnetwork = { version = "0.20", optional = true } 36 | cidr = { version = "0.3", optional = true } 37 | num-traits = "0.2" 38 | serde = { version = "1", optional = true} 39 | 40 | [dev-dependencies] 41 | criterion = { version = "0.5.1", features = ["html_reports"] } 42 | generic-tests = "0.1.2" 43 | ip_network_table-deps-treebitmap = "0.5.0" 44 | itertools = "0.14.0" 45 | paste = "1.0.10" 46 | pretty_assertions = "1.4.0" 47 | quickcheck = "1.0.3" 48 | quickcheck_macros = "1.0.0" 49 | rand = "0.8.5" 50 | -------------------------------------------------------------------------------- /LICENSE-APACHE: -------------------------------------------------------------------------------- 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 2017 Juniper Networks, Inc. 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 | -------------------------------------------------------------------------------- /LICENSE-MIT: -------------------------------------------------------------------------------- 1 | Copyright 2022 Tibor Schneider 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 8 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | [![CI test](https://img.shields.io/github/actions/workflow/status/tiborschneider/prefix-trie/test.yml)](https://github.com/tiborschneider/prefix-trie/actions) 2 | [![codecov](https://codecov.io/gh/tiborschneider/prefix-trie/branch/main/graph/badge.svg?token=EEJXNNURMW)](https://codecov.io/gh/tiborschneider/prefix-trie) 3 | [![version](https://img.shields.io/crates/v/prefix-trie)](https://crates.io/crates/prefix-trie) 4 | [![downloads](https://img.shields.io/crates/d/prefix-trie)](https://crates.io/crates/prefix-trie) 5 | [![docs.rs](https://img.shields.io/docsrs/prefix-trie/latest)](https://docs.rs/prefix-trie/latest/prefix_trie/) 6 | [![license](https://img.shields.io/crates/l/prefix-trie/0.2.2)](https://crates.io/crates/prefix-trie) 7 | [![Crates.io MSRV](https://img.shields.io/crates/msrv/prefix-trie)](https://crates.io/crates/prefix-trie) 8 | 9 | # Prefix-Trie 10 | 11 | This crate provides a simple prefix tree for IP prefixes. Any lookup performs longest-prefix 12 | match. This crate supports both IPv4 and IPv6 (from either [ipnet](https://docs.rs/ipnet/2.10.0) 13 | or [ipnetwork](https://crates.io/crates/ipnetwork) or [cidr](https://crates.io/crates/cidr)). 14 | It also supports any tuple `(R, u8)`, where `R` is any unsigned primitive integer (`u8`, `u16`, 15 | `u32`, `u64`, `u128`, or `usize`). 16 | 17 | ## Comparison with related projects 18 | 19 | [`ip_network_table-deps-treebitmap`](https://crates.io/crates/ip_network_table-deps-treebitmap) 20 | provides an IP lookup table, similar to `PrefixMap`. 21 | 22 | The following compares the two approaches in the case of *dense* or *sparse* maps. Each test case 23 | performs 100'000 modifications or lookups. However, the dense cases randomly pick any IPv4 24 | address, while the sparse case only picks 20 different IPv4 addresses. See `benches/benchmark.rs` 25 | for more details. 26 | 27 | | Operation | Mode | `PrefixMap` | `treebitmap` | factor | 28 | |-----------------|--------|-------------|--------------|--------| 29 | | Insert & Remove | dense | **31.78ms** | 47.52ms | ~1.5x | 30 | | Lookup | dense | 32.36ms | **8.409ms** | ~0.25x | 31 | | Insert & Remove | sparse | **6.645ms** | 7.329ms | ~1.1x | 32 | | Lookup | sparse | **8.394ms** | 12.30ms | ~1.5x | 33 | 34 | In addition, `prefix-trie` includes a `PrefixSet` analogous to `std::collections::HashSet`, 35 | including union, intersection and difference operations that are implemented as simultaneous 36 | tree traversals. Further, `prefix-trie` has an interface similar to `std::collections`, and 37 | offers a general longest-prefix match that is not limited to individual addresses. Finally, 38 | `prefix-trie` allows you to (mutably) borrow a sub-trie using views. 39 | 40 | ## Description of the Tree 41 | 42 | The tree is structured as follows: Each node consists of a prefix, a container for a potential 43 | value (`Option`), and two optional children. Adding a new child, or traversing into the tree is 44 | done as follows: we look at the most significant bit that is **not** part of the prefix 45 | itself. If it is not set, then we take the left branch, and otherwise, we take the right one. 46 | 47 | # Traversals 48 | 49 | Any iteration over all elements in the tree is implemented as a graph traversal that will yield 50 | elements in lexicographic order. 51 | 52 | The library offers set operations of different maps or sets. We implement a union, intersection, 53 | difference, and covering_difference. These iterators are implemented using simultaneous tree 54 | traversals. They will yield elements in lexicographic order. Whenever appropriate, the yielded 55 | items will also include the longest prefix match. 56 | 57 | # `TrieView` and `TrieViewMut` 58 | 59 | You can create a view of a (sub)-trie. Such a view has an any node as its root. Any operations 60 | on that view will only traverse that node and all its children. You can iterate over all 61 | children, search in that sub-trie, and perform set operations (union, intersection, difference, 62 | or the covering difference) on them. 63 | 64 | A view can point to one of three possible nodes: 65 | - A node in the tree that is actually present in the map, 66 | - A branching node that does not exist in the map, but is needed for the tree structure (or that 67 | was deleted using the function `remove_keep_tree`) 68 | - A virtual node that does not exist as a node in the tree. This is only the case if you call 69 | `PrefixView::find` or `AsView::view_at` with a node that is not present in the tree, but 70 | that contains elements present in the tree. Virtual nodes are treated as if they are actually 71 | present in the tree as branching nodes. 72 | 73 | # Operations on the tree 74 | 75 | There are several operations one can do on the tree. Regular inserts are handled using the 76 | `Entry` structure. An `Entry` is a pointer to a location in the tree to either insert a value or 77 | modify an existing one. Removals however are different. 78 | 79 | The following are the computational complexities of the functions, where `n` is the number of 80 | elements in the tree. 81 | 82 | | Operation | Complexity | 83 | |--------------------------------------------|------------| 84 | | `entry`, `insert` | `O(log n)` | 85 | | `remove`, `remove_keep_tree` | `O(log n)` | 86 | | `remove_children` (calling `drop` on `T`) | `O(n)` | 87 | | `get`, `get_lpm`, `get_mut` | `O(log n)` | 88 | | `retain` | `O(n)` | 89 | | `clear` (calling `drop` on `T`) | `O(n)` | 90 | | Operations on [`map::Entry`] | `O(1)` | 91 | | `len` and `is_empty` | `O(1)` | 92 | | `union`, `intersection`, `difference`, ... | `O(n)` | 93 | 94 | There are three kinds of removals you! can do: 95 | 96 | - `PrefixMap::remove` will remove an entry from the tree and modify the tree structure as if 97 | the value was never inserted before. `PrefixMap::remove` will always exactly revert the 98 | operation of `PrefixMap::insert`. When only calling this function to remove elements, you 99 | are guaranteed that the tree structure is indistinguishable to a different tree where you 100 | only inserted elements. 101 | - `PrefixMap::remove_children` will remove all entries that are contained within the given 102 | prefix. This operation will search for the node with the shortest prefix length that is 103 | contained within the given prefix and remove it, including all of its children. 104 | - `PrefixMap::remove_keep_tree` will not change anything in the tree structure. It will only 105 | remove a value from a node. As soon as you call `remove_keep_tree` once on a tree structure, 106 | the tree will no longer be optimal. 107 | 108 | ## TODO 109 | 110 | Migrate to a TreeBitMap, described by [W. Eatherton, Z. Dittia, G. Varghes](https://doi.org/10.1145/997150.997160). 111 | -------------------------------------------------------------------------------- /benches/benchmark.rs: -------------------------------------------------------------------------------- 1 | use criterion::{criterion_group, criterion_main, Criterion}; 2 | use ip_network_table_deps_treebitmap::IpLookupTable; 3 | use ipnet::Ipv4Net; 4 | use prefix_trie::*; 5 | use rand::prelude::*; 6 | use std::collections::HashSet; 7 | use std::net::Ipv4Addr; 8 | 9 | const ITERS: usize = 100_000; 10 | const NUM_SPARSE_ADDR: usize = 20; 11 | 12 | enum Insn { 13 | Insert(Ipv4Addr, u8, u32), 14 | Remove(Ipv4Addr, u8), 15 | ExactMatch(Ipv4Addr, u8), 16 | LongestPrefixMatch(Ipv4Addr, u8), 17 | } 18 | 19 | fn min_prefix_len(addr: u32) -> u8 { 20 | let mut bit: u32 = 0x00000001; 21 | let mut len: u8 = 32; 22 | while len > 0 && bit & addr == 0 { 23 | len = len.saturating_sub(1); 24 | (bit, _) = bit.overflowing_shl(1); 25 | } 26 | len 27 | } 28 | 29 | fn random_addr(rng: &mut ThreadRng) -> (Ipv4Addr, u8) { 30 | let addr: u32 = rng.gen::(); 31 | let max_len = 32; 32 | let min_len = min_prefix_len(addr); 33 | let len = rng.gen_range(min_len..=max_len); 34 | (addr.into(), len) 35 | } 36 | 37 | fn generate_random_mods_dense() -> (Vec, HashSet<(Ipv4Addr, u8)>) { 38 | let mut rng = thread_rng(); 39 | let mut result = Vec::new(); 40 | 41 | let mut addresses = HashSet::new(); 42 | 43 | for _ in 0..ITERS { 44 | if addresses.is_empty() || rng.gen_bool(0.8) { 45 | let (addr, len) = random_addr(&mut rng); 46 | let val = rng.gen::(); 47 | result.push(Insn::Insert(addr, len, val)); 48 | addresses.insert((addr, len)); 49 | } else { 50 | let (addr, len) = addresses 51 | .iter() 52 | .choose(&mut rng) 53 | .map(|(addr, len)| (*addr, *len)) 54 | .unwrap(); 55 | addresses.remove(&(addr, len)); 56 | result.push(Insn::Remove(addr, len)); 57 | } 58 | } 59 | (result, addresses) 60 | } 61 | 62 | fn generate_random_lookups_dense(addresses: &HashSet<(Ipv4Addr, u8)>) -> Vec { 63 | let mut rng = thread_rng(); 64 | let mut result = Vec::new(); 65 | 66 | for _ in 0..ITERS { 67 | if rng.gen_bool(0.5) { 68 | let (addr, len) = if addresses.is_empty() || rng.gen_bool(0.5) { 69 | random_addr(&mut rng) 70 | } else { 71 | addresses 72 | .iter() 73 | .choose(&mut rng) 74 | .map(|(addr, len)| (*addr, *len)) 75 | .unwrap() 76 | }; 77 | result.push(Insn::ExactMatch(addr, len)); 78 | } else { 79 | let (addr, len) = random_addr(&mut rng); 80 | result.push(Insn::LongestPrefixMatch(addr, len)); 81 | } 82 | } 83 | result 84 | } 85 | 86 | fn sparse_addresses() -> Vec<(Ipv4Addr, u8)> { 87 | let mut rng = thread_rng(); 88 | (0..NUM_SPARSE_ADDR) 89 | .map(|_| random_addr(&mut rng)) 90 | .collect() 91 | } 92 | 93 | fn generate_random_mods_sparse(addresses: &[(Ipv4Addr, u8)]) -> Vec { 94 | let mut rng = thread_rng(); 95 | (0..ITERS) 96 | .map(|_| { 97 | let (addr, len) = addresses.iter().choose(&mut rng).unwrap(); 98 | if rng.gen_bool(0.7) { 99 | let val = rng.gen::(); 100 | Insn::Insert(*addr, *len, val) 101 | } else { 102 | Insn::Remove(*addr, *len) 103 | } 104 | }) 105 | .collect() 106 | } 107 | 108 | fn generate_random_lookups_sparse(addresses: &[(Ipv4Addr, u8)]) -> Vec { 109 | let mut rng = thread_rng(); 110 | (0..ITERS) 111 | .map(|_| { 112 | let (addr, len) = addresses.iter().choose(&mut rng).unwrap(); 113 | if rng.gen_bool(0.5) { 114 | Insn::ExactMatch(*addr, *len) 115 | } else { 116 | Insn::LongestPrefixMatch(*addr, *len) 117 | } 118 | }) 119 | .collect() 120 | } 121 | 122 | fn execute_prefix_map(map: &mut PrefixMap, insns: &Vec) { 123 | for insn in insns { 124 | criterion::black_box(match insn { 125 | Insn::Insert(addr, len, val) => map.insert(Ipv4Net::new(*addr, *len).unwrap(), *val), 126 | Insn::Remove(addr, len) => map.remove(&Ipv4Net::new(*addr, *len).unwrap()), 127 | Insn::ExactMatch(addr, len) => map.get(&Ipv4Net::new(*addr, *len).unwrap()).copied(), 128 | Insn::LongestPrefixMatch(addr, len) => map 129 | .get_lpm(&Ipv4Net::new(*addr, *len).unwrap()) 130 | .map(|(_, x)| *x), 131 | }); 132 | } 133 | } 134 | 135 | fn lookup_prefix_map(map: &PrefixMap, insns: &Vec) { 136 | for insn in insns { 137 | criterion::black_box(match insn { 138 | Insn::Insert(_, _, _) => unreachable!(), 139 | Insn::Remove(_, _) => unreachable!(), 140 | Insn::ExactMatch(addr, len) => map.get(&Ipv4Net::new(*addr, *len).unwrap()).copied(), 141 | Insn::LongestPrefixMatch(addr, len) => map 142 | .get_lpm(&Ipv4Net::new(*addr, *len).unwrap()) 143 | .map(|(_, x)| *x), 144 | }); 145 | } 146 | } 147 | 148 | fn execute_treebitmap(map: &mut IpLookupTable, insns: &Vec) { 149 | for insn in insns { 150 | criterion::black_box(match insn { 151 | Insn::Insert(addr, len, val) => map.insert(*addr, *len as u32, *val), 152 | Insn::Remove(addr, len) => map.remove(*addr, *len as u32), 153 | Insn::ExactMatch(addr, len) => map.exact_match(*addr, *len as u32).copied(), 154 | Insn::LongestPrefixMatch(addr, _) => { 155 | let mut octets = addr.octets(); 156 | octets[3] += 1; 157 | let addr = octets.into(); 158 | map.longest_match(addr).map(|(_, _, x)| *x) 159 | } 160 | }); 161 | } 162 | } 163 | 164 | fn lookup_treebitmap(map: &IpLookupTable, insns: &Vec) { 165 | for insn in insns { 166 | criterion::black_box(match insn { 167 | Insn::Insert(_, _, _) => unreachable!(), 168 | Insn::Remove(_, _) => unreachable!(), 169 | Insn::ExactMatch(addr, len) => map.exact_match(*addr, *len as u32).copied(), 170 | Insn::LongestPrefixMatch(addr, _) => { 171 | let mut octets = addr.octets(); 172 | octets[3] += 1; 173 | let addr = octets.into(); 174 | map.longest_match(addr).map(|(_, _, x)| *x) 175 | } 176 | }); 177 | } 178 | } 179 | 180 | pub fn dense_mods(c: &mut Criterion) { 181 | let mut group = c.benchmark_group("dense modification"); 182 | 183 | let (insn, _) = generate_random_mods_dense(); 184 | 185 | group.bench_function("PrefixMap", |b| { 186 | b.iter(|| { 187 | let mut map = PrefixMap::new(); 188 | execute_prefix_map(&mut map, &insn); 189 | }) 190 | }); 191 | group.bench_function("TreeBitMap dense", |b| { 192 | b.iter(|| { 193 | let mut map = IpLookupTable::new(); 194 | execute_treebitmap(&mut map, &insn); 195 | }) 196 | }); 197 | 198 | group.finish(); 199 | } 200 | 201 | pub fn dense_lookup(c: &mut Criterion) { 202 | let (mods, addrs) = generate_random_mods_dense(); 203 | let lookups = generate_random_lookups_dense(&addrs); 204 | 205 | let mut prefix_map = PrefixMap::new(); 206 | let mut treebitmap = IpLookupTable::new(); 207 | execute_prefix_map(&mut prefix_map, &mods); 208 | execute_treebitmap(&mut treebitmap, &mods); 209 | 210 | let mut group = c.benchmark_group("dense lookups"); 211 | 212 | group.bench_function("PrefixMap", |b| { 213 | b.iter(|| { 214 | lookup_prefix_map(&prefix_map, &lookups); 215 | }) 216 | }); 217 | group.bench_function("TreeBitMap", |b| { 218 | b.iter(|| { 219 | lookup_treebitmap(&treebitmap, &lookups); 220 | }) 221 | }); 222 | 223 | group.finish(); 224 | } 225 | 226 | pub fn sparse_mods(c: &mut Criterion) { 227 | let mut group = c.benchmark_group("sparse modification"); 228 | 229 | let addrs = sparse_addresses(); 230 | let insn = generate_random_mods_sparse(&addrs); 231 | 232 | group.bench_function("PrefixMap", |b| { 233 | b.iter(|| { 234 | let mut map = PrefixMap::new(); 235 | execute_prefix_map(&mut map, &insn); 236 | }) 237 | }); 238 | group.bench_function("TreeBitMap sparse", |b| { 239 | b.iter(|| { 240 | let mut map = IpLookupTable::new(); 241 | execute_treebitmap(&mut map, &insn); 242 | }) 243 | }); 244 | 245 | group.finish(); 246 | } 247 | 248 | pub fn sparse_lookup(c: &mut Criterion) { 249 | let mut group = c.benchmark_group("sparse lookups"); 250 | 251 | let addrs = sparse_addresses(); 252 | let mods = generate_random_mods_sparse(&addrs); 253 | let lookups = generate_random_lookups_sparse(&addrs); 254 | 255 | let mut prefix_map = PrefixMap::new(); 256 | let mut treebitmap = IpLookupTable::new(); 257 | execute_prefix_map(&mut prefix_map, &mods); 258 | execute_treebitmap(&mut treebitmap, &mods); 259 | 260 | group.bench_function("PrefixMap", |b| { 261 | b.iter(|| { 262 | lookup_prefix_map(&prefix_map, &lookups); 263 | }) 264 | }); 265 | group.bench_function("TreeBitMap", |b| { 266 | b.iter(|| { 267 | lookup_treebitmap(&treebitmap, &lookups); 268 | }) 269 | }); 270 | 271 | group.finish(); 272 | } 273 | 274 | criterion_group!( 275 | benches, 276 | dense_lookup, 277 | dense_mods, 278 | sparse_lookup, 279 | sparse_mods 280 | ); 281 | criterion_main!(benches); 282 | -------------------------------------------------------------------------------- /codecov.yml: -------------------------------------------------------------------------------- 1 | ignore: 2 | - "src/fuzzing" 3 | - "src/test.rs" 4 | - "examples" 5 | -------------------------------------------------------------------------------- /examples/random_access.rs: -------------------------------------------------------------------------------- 1 | use std::net::Ipv4Addr; 2 | 3 | use prefix_trie::*; 4 | 5 | use ipnet::Ipv4Net; 6 | use rand::prelude::*; 7 | 8 | fn main() { 9 | let mut pm = PrefixMap::::new(); 10 | 11 | let mut rng = thread_rng(); 12 | 13 | for _ in 0..1_000_000 { 14 | let prefix = Ipv4Net::new(Ipv4Addr::new(rng.gen(), 0, 0, 0), rng.gen_range(1..=8)).unwrap(); 15 | let prefix = Ipv4Net::new(prefix.mask().into(), prefix.prefix_len()).unwrap(); 16 | 17 | if rng.gen_bool(0.7) { 18 | let value: u32 = rng.gen::() as u32; 19 | pm.insert(prefix, value); 20 | } else if rng.gen_bool(0.1) { 21 | // remove all children of that prefix 22 | pm.remove_children(&prefix); 23 | } else { 24 | pm.remove(&prefix); 25 | } 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /src/fmt.rs: -------------------------------------------------------------------------------- 1 | //! Formatting implementation for the PrefixMap 2 | 3 | use std::fmt::{Debug, Formatter, Result}; 4 | 5 | use super::*; 6 | 7 | impl Debug for PrefixMap { 8 | fn fmt(&self, f: &mut Formatter<'_>) -> Result { 9 | DebugPrefixMap(self, 0).fmt(f) 10 | } 11 | } 12 | 13 | struct DebugPrefixMap<'a, P, T>(&'a PrefixMap, usize); 14 | 15 | impl Debug for DebugPrefixMap<'_, P, T> { 16 | fn fmt(&self, f: &mut Formatter<'_>) -> Result { 17 | let map = self.0; 18 | let idx = self.1; 19 | let node = &map.table[idx]; 20 | match (node.value.as_ref(), node.left, node.right) { 21 | (None, None, None) => node.prefix.fmt(f), 22 | (None, None, Some(child)) | (None, Some(child), None) => f 23 | .debug_map() 24 | .entry(&node.prefix, &Self(map, child)) 25 | .finish(), 26 | (None, Some(left), Some(right)) => f 27 | .debug_map() 28 | .entry(&node.prefix, &(Self(map, left), Self(map, right))) 29 | .finish(), 30 | (Some(v), None, None) => f.debug_map().entry(&node.prefix, v).finish(), 31 | (Some(v), None, Some(child)) | (Some(v), Some(child), None) => f 32 | .debug_map() 33 | .entry(&node.prefix, &(v, Self(map, child))) 34 | .finish(), 35 | (Some(v), Some(left), Some(right)) => f 36 | .debug_map() 37 | .entry(&node.prefix, &(v, Self(map, left), Self(map, right))) 38 | .finish(), 39 | } 40 | } 41 | } 42 | 43 | impl Debug for PrefixSet

{ 44 | fn fmt(&self, f: &mut Formatter<'_>) -> Result { 45 | DebugPrefixMap(&self.0, 0).fmt(f) 46 | } 47 | } 48 | -------------------------------------------------------------------------------- /src/fuzzing/basic.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use super::*; 4 | use itertools::Itertools; 5 | 6 | qc!(new, _new); 7 | fn _new(list: Vec<(TestPrefix, i32)>) -> bool { 8 | let mut pmap = PrefixMap::new(); 9 | let mut hmap = HashMap::new(); 10 | 11 | for (p, t) in list { 12 | pmap.insert(p, t); 13 | hmap.insert(p, t); 14 | } 15 | 16 | // assert that the iterator of both is the same 17 | pmap.into_iter().eq(hmap.into_iter().sorted()) 18 | } 19 | 20 | qc!(new_mods, _new_mods); 21 | fn _new_mods(list: Vec>) -> bool { 22 | let mut pmap = PrefixMap::new(); 23 | let mut hmap = HashMap::new(); 24 | 25 | for op in list { 26 | match op { 27 | Operation::Add(p, t) => { 28 | pmap.insert(p, t); 29 | hmap.insert(p, t); 30 | } 31 | Operation::Remove(p) => { 32 | pmap.remove(&p); 33 | hmap.remove(&p); 34 | } 35 | } 36 | } 37 | 38 | // assert that the iterator of both is the same 39 | pmap.into_iter().eq(hmap.into_iter().sorted()) 40 | } 41 | 42 | qc!(new_mods_entry, _new_mods_entry); 43 | fn _new_mods_entry(list: Vec>) -> bool { 44 | let mut pmap = PrefixMap::new(); 45 | let mut hmap = HashMap::new(); 46 | 47 | for op in list { 48 | match op { 49 | Operation::Add(p, t) => { 50 | let _ = pmap.entry(p).insert(t); 51 | hmap.insert(p, t); 52 | } 53 | Operation::Remove(p) => { 54 | pmap.remove(&p); 55 | hmap.remove(&p); 56 | } 57 | } 58 | } 59 | 60 | // assert that the iterator of both is the same 61 | pmap.into_iter().eq(hmap.into_iter().sorted()) 62 | } 63 | 64 | qc!(equality, _equality); 65 | fn _equality(list: Vec>) -> bool { 66 | let mut map = PrefixMap::default(); 67 | 68 | for op in list { 69 | match op { 70 | Operation::Add(p, t) => { 71 | map.insert(p, t); 72 | } 73 | Operation::Remove(p) => { 74 | map.remove(&p); 75 | } 76 | } 77 | } 78 | 79 | let clone = map.clone().into_iter().collect::>(); 80 | 81 | // assert that the iterator of both is the same 82 | map == clone && map.len() == clone.len() && map.is_empty() == clone.is_empty() 83 | } 84 | 85 | qc!(equality_keep_tree, _equality_keep_tree); 86 | fn _equality_keep_tree(list: Vec>) -> bool { 87 | let mut map = PrefixMap::default(); 88 | 89 | for op in list { 90 | match op { 91 | Operation::Add(p, t) => { 92 | map.insert(p, t); 93 | } 94 | Operation::Remove(p) => { 95 | map.remove_keep_tree(&p); 96 | } 97 | } 98 | } 99 | 100 | let clone = map.clone().into_iter().collect::>(); 101 | 102 | // assert that the iterator of both is the same 103 | map == clone && map.len() == clone.len() && map.is_empty() == clone.is_empty() 104 | } 105 | 106 | qc!(equality_set, _equality_set); 107 | fn _equality_set(list: Vec>) -> bool { 108 | let mut set = PrefixSet::default(); 109 | 110 | for op in list { 111 | match op { 112 | Operation::Add(p, _) => { 113 | set.insert(p); 114 | } 115 | Operation::Remove(p) => { 116 | set.remove(&p); 117 | } 118 | } 119 | } 120 | 121 | let clone = set.iter().copied().collect::>(); 122 | 123 | // assert that the iterator of both is the same 124 | set == clone && set.len() == clone.len() && set.is_empty() == clone.is_empty() 125 | } 126 | 127 | qc!(remove_children, _remove_children); 128 | fn _remove_children((mut map, root): (PrefixMap, TestPrefix)) -> bool { 129 | let want = select(&map, |p, _| !root.contains(p)); 130 | map.remove_children(&root); 131 | map.len() == want.len() && map.into_iter().eq(want) 132 | } 133 | 134 | qc!(retain, _retain); 135 | fn _retain((mut map, root): (PrefixMap, TestPrefix)) -> bool { 136 | let want = select(&map, |p, _| !(root.contains(p) && p.1 >= root.1 + 2)); 137 | map.retain(|p, _| !(root.contains(p) && p.1 >= root.1 + 2)); 138 | map.into_iter().eq(want) 139 | } 140 | 141 | qc!(view_at, _view_at); 142 | fn _view_at((map, root): (PrefixMap, TestPrefix)) -> bool { 143 | let value = map.get(&root).cloned(); 144 | let root_has_nodes = map.iter().any(|(p, _)| root.contains(p)); 145 | match map.view_at(root) { 146 | None => !root_has_nodes, 147 | Some(view) => view.value() == value.as_ref(), 148 | } 149 | } 150 | qc!(view_mut_at, _view_mut_at); 151 | fn _view_mut_at((mut map, root): (PrefixMap, TestPrefix)) -> bool { 152 | let value = map.get(&root).cloned(); 153 | let root_has_nodes = map.iter().any(|(p, _)| root.contains(p)); 154 | match map.view_mut_at(root) { 155 | None => !root_has_nodes, 156 | Some(view) => view.value() == value.as_ref(), 157 | } 158 | } 159 | 160 | qc!(view_left, _view_left); 161 | fn _view_left((map, root): (PrefixMap, TestPrefix)) -> bool { 162 | let left_prefix_has_nodes = map.iter().any(|(p, _)| root.left().contains(p)); 163 | map.view_at(root).and_then(|v| v.left()).is_some() == left_prefix_has_nodes 164 | } 165 | 166 | qc!(view_right, _view_right); 167 | fn _view_right((map, root): (PrefixMap, TestPrefix)) -> bool { 168 | let right_prefix_has_nodes = map.iter().any(|(p, _)| root.right().contains(p)); 169 | map.view_at(root).and_then(|v| v.right()).is_some() == right_prefix_has_nodes 170 | } 171 | 172 | qc!(view_mut_left, _view_mut_left); 173 | fn _view_mut_left((mut m, root): (PrefixMap, TestPrefix)) -> bool { 174 | let left_prefix_has_nodes = m.iter().any(|(p, _)| root.left().contains(p)); 175 | let c1 = m.view_mut_at(root).and_then(|v| v.left().ok()).is_some() == left_prefix_has_nodes; 176 | let c2 = m.view_mut_at(root).map(|v| v.has_left()).unwrap_or(false) == left_prefix_has_nodes; 177 | c1 && c2 178 | } 179 | 180 | qc!(view_mut_right, _view_mut_right); 181 | fn _view_mut_right((mut m, root): (PrefixMap, TestPrefix)) -> bool { 182 | let right_prefix_has_nodes = m.iter().any(|(p, _)| root.right().contains(p)); 183 | let c1 = m.view_mut_at(root).and_then(|v| v.right().ok()).is_some() == right_prefix_has_nodes; 184 | let c2 = m.view_mut_at(root).map(|v| v.has_right()).unwrap_or(false) == right_prefix_has_nodes; 185 | c1 && c2 186 | } 187 | 188 | qc!(view_mut_split, _view_mut_split); 189 | fn _view_mut_split((mut map, root): (PrefixMap, TestPrefix)) -> bool { 190 | let left_prefix_has_nodes = map.iter().any(|(p, _)| root.left().contains(p)); 191 | let right_prefix_has_nodes = map.iter().any(|(p, _)| root.right().contains(p)); 192 | map.view_mut_at(root) 193 | .map(|view| view.split()) 194 | .map(|(l, r)| (l.is_some(), r.is_some())) 195 | .unwrap_or((false, false)) 196 | == (left_prefix_has_nodes, right_prefix_has_nodes) 197 | } 198 | -------------------------------------------------------------------------------- /src/fuzzing/mod.rs: -------------------------------------------------------------------------------- 1 | //! Module for testing using fuzzing (quickcheck) 2 | #![allow(clippy::type_complexity)] 3 | 4 | use std::fmt::Debug; 5 | 6 | use crate::*; 7 | use quickcheck::Arbitrary; 8 | 9 | mod basic; 10 | mod set_ops; 11 | mod traversals; 12 | 13 | #[derive(Debug, PartialEq, Clone, Copy)] 14 | enum Operation { 15 | Add(P, T), 16 | Remove(P), 17 | } 18 | 19 | #[cfg(miri)] 20 | const DEFAULT_NUM_TESTS: usize = 10; 21 | #[cfg(not(miri))] 22 | const DEFAULT_NUM_TESTS: usize = 10000; 23 | const DEFAULT_GEN_SIZE: usize = 100; 24 | 25 | fn proptest_runner bool>(f: F) { 26 | let num_tests: usize = std::env::var("QUICKCHECK_TESTS") 27 | .ok() 28 | .and_then(|x| x.parse::().ok()) 29 | .unwrap_or(DEFAULT_NUM_TESTS); 30 | 31 | let gen_size: usize = std::env::var("QUICKCHECK_GENERATOR_SIZE") 32 | .ok() 33 | .and_then(|x| x.parse::().ok()) 34 | .unwrap_or(DEFAULT_GEN_SIZE); 35 | 36 | let mut gen = quickcheck::Gen::new(gen_size); 37 | 38 | // sample all inputs 39 | for _ in 0..num_tests { 40 | let input = A::arbitrary(&mut gen); 41 | let input_c = input.clone(); 42 | let success = f(input_c); 43 | if !success { 44 | shrink_failure(f, input) 45 | } 46 | } 47 | } 48 | 49 | fn shrink_failure bool>(f: F, input: A) -> ! { 50 | for i in input.shrink() { 51 | let i_c = i.clone(); 52 | let success = f(i_c); 53 | if !success { 54 | shrink_failure(f, i) 55 | } 56 | } 57 | // if we reach this point, then all shrunken inputs work. Therefore, `inputs` is the minimal 58 | // input 59 | panic!( 60 | "[QUICKCHECK] Test case failed!\n Minimal input:\n {:?}", 61 | input 62 | ); 63 | } 64 | 65 | #[allow(missing_docs)] 66 | #[macro_export] 67 | macro_rules! qc { 68 | ($name:ident, $f:ident) => { 69 | #[test] 70 | fn $name() { 71 | proptest_runner($f) 72 | } 73 | }; 74 | } 75 | 76 | fn select bool>(map: &PrefixMap, f: F) -> Vec<(P, T)> { 77 | map.iter() 78 | .map(|(p, t)| (p.clone(), t.clone())) 79 | .filter(|(p, t)| f(p, t)) 80 | .collect() 81 | } 82 | 83 | fn select_ref bool>(map: &PrefixMap, f: F) -> Vec<(&P, &T)> { 84 | map.iter().filter(|(p, t)| f(*p, *t)).collect() 85 | } 86 | 87 | fn select_keys bool>(map: &PrefixMap, f: F) -> Vec<&P> { 88 | map.iter() 89 | .filter(|(p, t)| f(*p, *t)) 90 | .map(|(p, _)| p) 91 | .collect() 92 | } 93 | 94 | fn select_values bool>(map: &PrefixMap, f: F) -> Vec<&T> { 95 | map.iter() 96 | .filter(|(p, t)| f(*p, *t)) 97 | .map(|(_, t)| t) 98 | .collect() 99 | } 100 | 101 | impl Arbitrary for PrefixMap { 102 | fn arbitrary(g: &mut quickcheck::Gen) -> Self { 103 | as Arbitrary>::arbitrary(g) 104 | .into_iter() 105 | .collect() 106 | } 107 | 108 | fn shrink(&self) -> Box> { 109 | let elems = self.clone().into_iter().collect::>(); 110 | let shrinked = elems.shrink(); 111 | Box::new(shrinked.map(PrefixMap::from_iter)) 112 | } 113 | } 114 | 115 | impl Arbitrary for Operation { 116 | fn arbitrary(g: &mut quickcheck::Gen) -> Self { 117 | let p = P::arbitrary(g); 118 | if g.choose(&[ 119 | true, true, true, true, true, true, true, false, false, false, 120 | ]) 121 | .copied() 122 | .unwrap_or_default() 123 | { 124 | let t = T::arbitrary(g); 125 | Self::Add(p, t) 126 | } else { 127 | Self::Remove(p) 128 | } 129 | } 130 | 131 | fn shrink(&self) -> Box> { 132 | match self { 133 | Operation::Add(p, t) => { 134 | let t = t.clone(); 135 | Box::new( 136 | p.clone() 137 | .shrink() 138 | .map(move |p| Operation::Add(p, t.clone())), 139 | ) 140 | } 141 | Operation::Remove(p) => Box::new(p.clone().shrink().map(|p| Operation::Remove(p))), 142 | } 143 | } 144 | } 145 | 146 | #[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)] 147 | struct TestPrefix(u32, u8); 148 | 149 | impl TestPrefix { 150 | fn left(self) -> Self { 151 | TestPrefix(self.0, self.1 + 1) 152 | } 153 | 154 | fn right(self) -> Self { 155 | TestPrefix(self.0 + (1 << (31 - self.1)), self.1 + 1) 156 | } 157 | } 158 | 159 | impl Debug for TestPrefix { 160 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 161 | let addr = format!("{:032b}", self.0)[..10].to_string(); 162 | write!(f, "0b{addr}/{}", self.1) 163 | } 164 | } 165 | 166 | impl Arbitrary for TestPrefix { 167 | fn arbitrary(g: &mut quickcheck::Gen) -> Self { 168 | #[rustfmt::skip] 169 | let len: u8 = *g 170 | .choose(&[ 171 | 0, 172 | 1, 1, 173 | 2, 2, 2, 174 | 3, 3, 3, 3, 175 | 4, 4, 4, 4, 4, 176 | 5, 5, 5, 5, 5, 5, 177 | 6, 6, 6, 6, 6, 6, 6, 178 | 7, 7, 7, 7, 7, 7, 7, 7, 179 | 8, 8, 8, 8, 8, 8, 8, 8, 8, 180 | 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 181 | ]) 182 | .unwrap(); 183 | let x = u32::arbitrary(g); 184 | Self::from_repr_len(x, len) 185 | } 186 | 187 | fn shrink(&self) -> Box> { 188 | if self.1 == 0 { 189 | quickcheck::empty_shrinker() 190 | } else { 191 | let len = self.1 - 1; 192 | let x = Self::from_repr_len(self.0, len); 193 | quickcheck::single_shrinker(x) 194 | } 195 | } 196 | } 197 | 198 | impl Prefix for TestPrefix { 199 | type R = u32; 200 | 201 | fn repr(&self) -> Self::R { 202 | self.0 203 | } 204 | 205 | fn prefix_len(&self) -> u8 { 206 | self.1 207 | } 208 | 209 | fn from_repr_len(repr: Self::R, len: u8) -> Self { 210 | let x = Prefix::mask(&(repr, len)); 211 | Self(x, len) 212 | } 213 | } 214 | -------------------------------------------------------------------------------- /src/fuzzing/set_ops.rs: -------------------------------------------------------------------------------- 1 | use std::collections::HashMap; 2 | 3 | use super::*; 4 | use itertools::Itertools; 5 | use trieview::UnionItem; 6 | 7 | qc!(union, _union); 8 | fn _union((a, b): (PrefixMap, PrefixMap)) -> bool { 9 | let mut union_set: HashMap, Option)> = HashMap::new(); 10 | for (p, a) in a.iter() { 11 | let _ = union_set.entry(*p).or_default().0.insert(*a); 12 | } 13 | for (p, b) in b.iter() { 14 | let _ = union_set.entry(*p).or_default().1.insert(*b); 15 | } 16 | 17 | let want = union_set.into_iter().sorted().collect::>(); 18 | let got = a 19 | .view() 20 | .union(&b) 21 | .map(|x| match x { 22 | UnionItem::Left { prefix, left, .. } => (*prefix, (Some(*left), None)), 23 | UnionItem::Right { prefix, right, .. } => (*prefix, (None, Some(*right))), 24 | UnionItem::Both { 25 | prefix, 26 | left, 27 | right, 28 | } => (*prefix, (Some(*left), Some(*right))), 29 | }) 30 | .collect::>(); 31 | 32 | want == got 33 | } 34 | 35 | qc!(union_mut, _union_mut); 36 | fn _union_mut((mut a, mut b): (PrefixMap, PrefixMap)) -> bool { 37 | let mut union_set: HashMap, Option)> = HashMap::new(); 38 | for (p, a) in a.iter() { 39 | let _ = union_set.entry(*p).or_default().0.insert(*a); 40 | } 41 | for (p, b) in b.iter() { 42 | let _ = union_set.entry(*p).or_default().1.insert(*b); 43 | } 44 | 45 | let want = union_set.into_iter().sorted().collect::>(); 46 | let got = a 47 | .view_mut() 48 | .union_mut(&mut b) 49 | .map(|(p, a, b)| (*p, (a.copied(), b.copied()))) 50 | .collect::>(); 51 | 52 | want == got 53 | } 54 | 55 | qc!(union_lpm, _union_lpm); 56 | fn _union_lpm((a, b): (PrefixMap, PrefixMap)) -> bool { 57 | let mut union_set: HashMap, Option<(TestPrefix, i32)>)> = 58 | HashMap::new(); 59 | for (p, a) in a.iter() { 60 | union_set 61 | .entry(*p) 62 | .or_insert_with(|| (Some((*p, *a)), b.get_lpm(p).copied())); 63 | } 64 | for (p, b) in b.iter() { 65 | union_set 66 | .entry(*p) 67 | .or_insert_with(|| (a.get_lpm(p).copied(), Some((*p, *b)))); 68 | } 69 | 70 | let want = union_set.into_iter().sorted().collect::>(); 71 | let got = a 72 | .view() 73 | .union(&b) 74 | .map(|x| (*x.prefix(), (x.left().copied(), x.right().copied()))) 75 | .collect::>(); 76 | 77 | want == got 78 | } 79 | 80 | qc!(intersection, _intersection); 81 | fn _intersection((a, b): (PrefixMap, PrefixMap)) -> bool { 82 | let mut intersection_set: HashMap = HashMap::new(); 83 | for (p, a) in a.iter() { 84 | if let Some(b) = b.get(p) { 85 | intersection_set.insert(*p, (*a, *b)); 86 | } 87 | } 88 | 89 | let want = intersection_set.into_iter().sorted().collect::>(); 90 | let got = a 91 | .view() 92 | .intersection(&b) 93 | .map(|(p, a, b)| (*p, (*a, *b))) 94 | .collect::>(); 95 | 96 | want == got 97 | } 98 | 99 | qc!(intersection_mut, _intersection_mut); 100 | fn _intersection_mut( 101 | (mut a, mut b): (PrefixMap, PrefixMap), 102 | ) -> bool { 103 | let mut intersection_set: HashMap = HashMap::new(); 104 | for (p, a) in a.iter() { 105 | if let Some(b) = b.get(p) { 106 | intersection_set.insert(*p, (*a, *b)); 107 | } 108 | } 109 | 110 | let want = intersection_set.into_iter().sorted().collect::>(); 111 | let got = a 112 | .view_mut() 113 | .intersection_mut(&mut b) 114 | .map(|(p, a, b)| (*p, (*a, *b))) 115 | .collect::>(); 116 | 117 | want == got 118 | } 119 | 120 | qc!(difference, _difference); 121 | fn _difference((a, b): (PrefixMap, PrefixMap)) -> bool { 122 | let mut difference_set: HashMap)> = HashMap::new(); 123 | for (p, a) in a.iter() { 124 | let b = b.get_lpm(p).copied(); 125 | if b.map(|(lpm, _)| &lpm == p).unwrap_or(false) { 126 | // same prefix! ignore 127 | } else { 128 | difference_set.insert(*p, (*a, b)); 129 | } 130 | } 131 | 132 | let want = difference_set.into_iter().sorted().collect::>(); 133 | let got = a 134 | .view() 135 | .difference(&b) 136 | .map(|i| (*i.prefix, (*i.value, i.right.copied()))) 137 | .collect::>(); 138 | 139 | want == got 140 | } 141 | 142 | qc!(difference_mut, _difference_mut); 143 | fn _difference_mut((mut a, b): (PrefixMap, PrefixMap)) -> bool { 144 | let mut difference_set: HashMap)> = HashMap::new(); 145 | for (p, a) in a.iter() { 146 | let b = b.get_lpm(p).copied(); 147 | if b.map(|(lpm, _)| &lpm == p).unwrap_or(false) { 148 | // same prefix! ignore 149 | } else { 150 | difference_set.insert(*p, (*a, b)); 151 | } 152 | } 153 | 154 | let want = difference_set.into_iter().sorted().collect::>(); 155 | let got = a 156 | .view_mut() 157 | .difference_mut(&b) 158 | .map(|i| (*i.prefix, (*i.value, i.right.copied()))) 159 | .collect::>(); 160 | 161 | want == got 162 | } 163 | 164 | qc!(covering_difference, _covering_difference); 165 | fn _covering_difference((a, b): (PrefixMap, PrefixMap)) -> bool { 166 | let mut difference_set: HashMap = HashMap::new(); 167 | for (p, a) in a.iter() { 168 | if b.get_lpm(p).is_none() { 169 | difference_set.insert(*p, *a); 170 | } 171 | } 172 | 173 | let want = difference_set.into_iter().sorted().collect::>(); 174 | let got = a 175 | .view() 176 | .covering_difference(&b) 177 | .map(|(p, t)| (*p, *t)) 178 | .collect::>(); 179 | 180 | want == got 181 | } 182 | 183 | qc!(covering_difference_mut, _covering_difference_mut); 184 | fn _covering_difference_mut( 185 | (mut a, b): (PrefixMap, PrefixMap), 186 | ) -> bool { 187 | let mut difference_set: HashMap = HashMap::new(); 188 | for (p, a) in a.iter() { 189 | if b.get_lpm(p).is_none() { 190 | difference_set.insert(*p, *a); 191 | } 192 | } 193 | 194 | let want = difference_set.into_iter().sorted().collect::>(); 195 | let got = a 196 | .view_mut() 197 | .covering_difference_mut(&b) 198 | .map(|(p, t)| (*p, *t)) 199 | .collect::>(); 200 | 201 | want == got 202 | } 203 | 204 | trait MyCopy { 205 | type Out; 206 | 207 | fn copied(&self) -> Self::Out; 208 | } 209 | 210 | impl MyCopy for Option<(&P, &T)> { 211 | type Out = Option<(P, T)>; 212 | 213 | fn copied(&self) -> Self::Out { 214 | self.map(|(p, t)| (*p, *t)) 215 | } 216 | } 217 | 218 | qc!(self_union_mut, _self_union_mut); 219 | fn _self_union_mut(mut map: PrefixMap) -> bool { 220 | let original_map = map.clone(); 221 | // go to the first split 222 | let mut view = map.view_mut(); 223 | let (mut left, right) = loop { 224 | match view.split() { 225 | (None, None) => return true, 226 | (None, Some(v)) | (Some(v), None) => view = v, 227 | (Some(left), Some(right)) => break (left, right), 228 | } 229 | }; 230 | 231 | let left_prefix = left.prefix(); 232 | let right_prefix = right.prefix(); 233 | 234 | let want = original_map 235 | .iter() 236 | .map(|(p, t)| { 237 | ( 238 | *p, 239 | if left_prefix.contains(p) || right_prefix.contains(p) { 240 | t.saturating_mul(2) 241 | } else { 242 | *t 243 | }, 244 | ) 245 | }) 246 | .sorted() 247 | .collect::>(); 248 | 249 | // take the union of both left and right and multiply each entry by 2 250 | for (_, l, r) in left.union_mut(right) { 251 | match (l, r) { 252 | (None, None) | (Some(_), Some(_)) => return false, 253 | (None, Some(t)) | (Some(t), None) => *t = t.saturating_mul(2), 254 | } 255 | } 256 | 257 | let got = map.into_iter().collect::>(); 258 | want == got 259 | } 260 | 261 | qc!(self_intersection_mut, _self_intersection_mut); 262 | fn _self_intersection_mut(mut map: PrefixMap) -> bool { 263 | // go to the first split 264 | let mut view = map.view_mut(); 265 | let (mut left, right) = loop { 266 | match view.split() { 267 | (None, None) => return true, 268 | (None, Some(v)) | (Some(v), None) => view = v, 269 | (Some(left), Some(right)) => break (left, right), 270 | } 271 | }; 272 | 273 | // take the union of both left and right and multiply each entry by 2 274 | left.intersection_mut(right).count() == 0 275 | } 276 | -------------------------------------------------------------------------------- /src/fuzzing/traversals.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | sync::{Arc, Mutex}, 3 | thread, 4 | }; 5 | 6 | use itertools::Itertools; 7 | 8 | use super::*; 9 | 10 | qc!(children, _children); 11 | fn _children((map, start): (PrefixMap, TestPrefix)) -> bool { 12 | let want = select_ref(&map, |p, _| start.contains(p)); 13 | map.children(&start).eq(want) 14 | } 15 | 16 | qc!(children_trieview, _children_trieview); 17 | fn _children_trieview((map, start): (PrefixMap, TestPrefix)) -> bool { 18 | let want = select_ref(&map, |p, _| start.contains(p)); 19 | if let Some(view) = map.view_at(start) { 20 | view.iter().eq(want) 21 | } else { 22 | want.is_empty() 23 | } 24 | } 25 | 26 | qc!(children_trieview_mut, _children_trieview_mut); 27 | fn _children_trieview_mut((mut map, start): (PrefixMap, TestPrefix)) -> bool { 28 | let want = select(&map, |p, _| start.contains(p)); 29 | if let Some(view) = map.view_mut_at(start) { 30 | view.into_iter().map(|(p, t)| (*p, *t)).eq(want) 31 | } else { 32 | want.is_empty() 33 | } 34 | } 35 | 36 | qc!(children_keys_trieview, _children_keys_trieview); 37 | fn _children_keys_trieview((map, start): (PrefixMap, TestPrefix)) -> bool { 38 | let want = select_keys(&map, |p, _| start.contains(p)); 39 | if let Some(view) = map.view_at(start) { 40 | view.keys().eq(want) 41 | } else { 42 | want.is_empty() 43 | } 44 | } 45 | 46 | qc!(children_values_trieview, _children_values_trieview); 47 | fn _children_values_trieview((map, start): (PrefixMap, TestPrefix)) -> bool { 48 | let want = select_values(&map, |p, _| start.contains(p)); 49 | if let Some(view) = map.view_at(start) { 50 | view.values().eq(want) 51 | } else { 52 | want.is_empty() 53 | } 54 | } 55 | 56 | qc!(simultaneous_iter_mut, _simultaneous_iter_mut); 57 | fn _simultaneous_iter_mut(mut map: PrefixMap) -> bool { 58 | // go to the first split 59 | let mut view = map.view_mut(); 60 | let (left, right) = loop { 61 | match view.split() { 62 | (None, None) => return true, 63 | (None, Some(v)) | (Some(v), None) => view = v, 64 | (Some(left), Some(right)) => break (left, right), 65 | } 66 | }; 67 | 68 | let want = left 69 | .view() 70 | .into_iter() 71 | .chain(right.view()) 72 | .map(|(p, t)| (*p, t.saturating_mul(4))) 73 | .sorted() 74 | .collect::>(); 75 | 76 | // simultaneously traverse left and right 77 | let result = Arc::new(Mutex::new(Vec::new())); 78 | let res_left = result.clone(); 79 | let res_right = result.clone(); 80 | 81 | fn foo<'a>( 82 | view: TrieViewMut<'a, TestPrefix, i32>, 83 | res: Arc>>, 84 | ) { 85 | let mut references = Vec::new(); 86 | for (p, t) in view { 87 | *t = t.saturating_mul(2); 88 | references.push((p, t)); 89 | } 90 | res.lock().unwrap().extend(references); 91 | } 92 | 93 | thread::scope(move |s| { 94 | let thread_left = s.spawn(move || foo(left, res_left)); 95 | let thread_right = s.spawn(move || foo(right, res_right)); 96 | thread_left.join().unwrap(); 97 | thread_right.join().unwrap(); 98 | }); 99 | 100 | // take the result and multiply all values by 2 again 101 | let mut result = result.lock().unwrap(); 102 | let mut got = Vec::new(); 103 | for (p, t) in result.iter_mut() { 104 | **t = t.saturating_mul(2); 105 | got.push((**p, **t)); 106 | } 107 | got.sort(); 108 | 109 | want == got 110 | } 111 | -------------------------------------------------------------------------------- /src/inner.rs: -------------------------------------------------------------------------------- 1 | //! The inner datastructure of a PrefixTrie that offers interior mutability. 2 | 3 | use std::{ 4 | cell::UnsafeCell, 5 | ops::{Index, IndexMut}, 6 | }; 7 | 8 | use crate::{to_right, Prefix}; 9 | 10 | #[derive(Clone)] 11 | pub(crate) struct Node { 12 | pub(crate) prefix: P, 13 | pub(crate) value: Option, 14 | pub(crate) left: Option, 15 | pub(crate) right: Option, 16 | } 17 | 18 | impl Node { 19 | /// get the tuple of prefix and value. 20 | pub(crate) fn prefix_value(&self) -> Option<(&P, &T)> { 21 | self.value.as_ref().map(|v| (&self.prefix, v)) 22 | } 23 | 24 | /// get the tuple of prefix and value. 25 | pub(crate) fn prefix_value_mut(&mut self) -> Option<(&P, &mut T)> { 26 | self.value.as_mut().map(|v| (&self.prefix, v)) 27 | } 28 | } 29 | 30 | /// A table to the prefix-trie that offers interior mutability. 31 | /// 32 | /// # Safety 33 | /// Owning a mutable reference to the Table implies that you can safely get a mutable reference to 34 | /// the inner data. If, however, you own an immutable reference, then you must guarantee that there 35 | /// is no other reference to the Table that potentially accesses the same node mutably. This interior 36 | /// mutability is only ever provided in `get_mut`. 37 | pub(crate) struct Table(UnsafeCell>>); 38 | 39 | // Safety: 40 | // - Sending a PrefixMap over thread boundary is fine. No-one besides us can have the raw pointer, 41 | // otherwise, the map would be borrowed. 42 | // - Sending a reference of PrefixMap over thread boundaries (i.e., TrieView is Send) is safe, 43 | // because we ensure that the existence of a TrieView on a sub-tree implies the absence of a 44 | // TrieViewMut that overlaps with that sub-tree. 45 | // - Sending a mutable reference of PrefixMap over thread boundaries (i.e., TrieView is Send) is 46 | // safe, because we ensure that the existence of a TrieViewMut on a sub-tree implies the absence 47 | // of any other TrieView or TrieViewMut that overlaps with that sub-tree. 48 | // The same argument holds for Sync. 49 | unsafe impl Send for Table {} 50 | unsafe impl Sync for Table {} 51 | 52 | impl AsRef>> for Table { 53 | fn as_ref(&self) -> &Vec> { 54 | // Safety: We own an immutable reference to the table. 55 | unsafe { self.0.get().as_ref().unwrap() } 56 | } 57 | } 58 | 59 | impl AsMut>> for Table { 60 | fn as_mut(&mut self) -> &mut Vec> { 61 | self.0.get_mut() 62 | } 63 | } 64 | 65 | impl Index for Table { 66 | type Output = Node; 67 | 68 | fn index(&self, index: usize) -> &Self::Output { 69 | &self.as_ref()[index] 70 | } 71 | } 72 | 73 | impl IndexMut for Table { 74 | fn index_mut(&mut self, index: usize) -> &mut Self::Output { 75 | &mut self.as_mut()[index] 76 | } 77 | } 78 | 79 | impl Clone for Table { 80 | fn clone(&self) -> Self { 81 | Self(UnsafeCell::new(self.as_ref().clone())) 82 | } 83 | } 84 | 85 | impl Default for Table 86 | where 87 | P: Prefix, 88 | { 89 | fn default() -> Self { 90 | Self(UnsafeCell::new(vec![Node { 91 | prefix: P::zero(), 92 | value: None, 93 | left: None, 94 | right: None, 95 | }])) 96 | } 97 | } 98 | 99 | pub(crate) enum Direction { 100 | /// The prefix is already reached. 101 | Reached, 102 | /// Enter the next index and search again. 103 | Enter { next: usize, right: bool }, 104 | /// The node was not found. 105 | Missing, 106 | } 107 | 108 | pub(crate) enum DirectionForInsert

{ 109 | /// The prefix is already reached. 110 | Reached, 111 | /// Enter the next index and search again. 112 | Enter { next: usize, right: bool }, 113 | /// Insert a new child at the given position as a leaf. 114 | NewLeaf { right: bool }, 115 | /// Insert a new child at the given position, moving all old children to be a child of the new 116 | /// prefix. `parent_right` tells where to insert the new node, while `child_right` tells where 117 | /// to insert the old child (on the right or the left of the new node). 118 | NewChild { right: bool, child_right: bool }, 119 | /// Insert a new branch at the parent with the given prefix. `parent_right` tells where to 120 | /// insert the branch, while `prefix_right` tells where to insert the new node at the 121 | /// branch. The old child of the parent should be inserted at `!prefix_right` of the branch. 122 | NewBranch { 123 | branch_prefix: P, 124 | right: bool, 125 | prefix_right: bool, 126 | }, 127 | } 128 | 129 | impl Table { 130 | pub(crate) fn into_inner(self) -> Vec> { 131 | self.0.into_inner() 132 | } 133 | 134 | /// *Safety*: You must ensure for the lifetime of 'a, that you will never construct a second 135 | /// reference to that node (neither mutable nor immutable). 136 | #[allow(clippy::mut_from_ref)] 137 | pub(crate) unsafe fn get_mut(&self, idx: usize) -> &mut Node { 138 | // old implementation that caused issues with Miri: 139 | // unsafe { &mut self.0.get().as_mut().unwrap()[idx] } 140 | 141 | // new implementation based on manually offsetting the pointers: 142 | unsafe { 143 | // do the bounds check 144 | let len = self.0.get().as_ref().unwrap().len(); 145 | if idx >= len { 146 | panic!("index out of bounds: the len is {len} but the index is {idx}"); 147 | } 148 | let ptr_to_slice = self.0.get().as_ref().unwrap().as_ptr(); 149 | let ptr_to_elem = ptr_to_slice.add(idx); 150 | (ptr_to_elem as *mut Node).as_mut().unwrap() 151 | } 152 | } 153 | } 154 | 155 | impl Table { 156 | /// Get the child of a node, either to the left or the right 157 | #[inline(always)] 158 | pub(crate) fn get_child(&self, idx: usize, right: bool) -> Option { 159 | if right { 160 | self[idx].right 161 | } else { 162 | self[idx].left 163 | } 164 | } 165 | 166 | /// set the child of a node (either to the left or the right), and return the index of the old child. 167 | #[inline(always)] 168 | pub(crate) fn set_child(&mut self, idx: usize, child: usize, right: bool) -> Option { 169 | if right { 170 | self[idx].right.replace(child) 171 | } else { 172 | self[idx].left.replace(child) 173 | } 174 | } 175 | 176 | /// remove a child from a node (just the reference). 177 | #[inline(always)] 178 | pub(crate) fn clear_child(&mut self, idx: usize, right: bool) -> Option { 179 | if right { 180 | self[idx].right.take() 181 | } else { 182 | self[idx].left.take() 183 | } 184 | } 185 | 186 | /// Get the directions from some node `idx` to get to `prefix`. 187 | #[inline(always)] 188 | pub(crate) fn get_direction(&self, cur: usize, prefix: &P) -> Direction { 189 | let cur_p = &self[cur].prefix; 190 | if cur_p.eq(prefix) { 191 | Direction::Reached 192 | } else { 193 | let right = to_right(cur_p, prefix); 194 | match self.get_child(cur, right) { 195 | Some(child) if self[child].prefix.contains(prefix) => { 196 | Direction::Enter { next: child, right } 197 | } 198 | _ => Direction::Missing, 199 | } 200 | } 201 | } 202 | 203 | /// Get the directions from some node `idx` to get to `prefix`. 204 | #[inline(always)] 205 | pub(crate) fn get_direction_for_insert(&self, cur: usize, prefix: &P) -> DirectionForInsert

{ 206 | let cur_p = &self[cur].prefix; 207 | if cur_p.eq(prefix) { 208 | DirectionForInsert::Reached 209 | } else { 210 | let right = to_right(cur_p, prefix); 211 | if let Some(child) = self.get_child(cur, right) { 212 | let child_p = &self[child].prefix; 213 | if child_p.contains(prefix) { 214 | DirectionForInsert::Enter { next: child, right } 215 | } else if prefix.contains(child_p) { 216 | DirectionForInsert::NewChild { 217 | right, 218 | child_right: to_right(prefix, child_p), 219 | } 220 | } else { 221 | let branch_prefix = prefix.longest_common_prefix(child_p); 222 | let prefix_right = to_right(&branch_prefix, prefix); 223 | DirectionForInsert::NewBranch { 224 | branch_prefix, 225 | right, 226 | prefix_right, 227 | } 228 | } 229 | } else { 230 | DirectionForInsert::NewLeaf { right } 231 | } 232 | } 233 | } 234 | } 235 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! This crate provides a simple prefix tree for IP prefixes. Any lookup performs longest-prefix 2 | //! match. This crate supports both IPv4 and IPv6 (from either [ipnet](https://docs.rs/ipnet/2.10.0) 3 | //! or [ipnetwork](https://crates.io/crates/ipnetwork) or [cidr](https://crates.io/crates/cidr)). 4 | //! It also supports any tuple `(R, u8)`, where `R` is any unsigned primitive integer (`u8`, `u16`, 5 | //! `u32`, `u64`, `u128`, or `usize`). 6 | //! 7 | //! # Comparison with related projects 8 | //! 9 | //! [`ip_network_table-deps-treebitmap`](https://crates.io/crates/ip_network_table-deps-treebitmap) 10 | //! provides an IP lookup table, similar to [`PrefixMap`]. 11 | //! 12 | //! The following compares the two approaches in the case of *dense* or *sparse* maps. Each test 13 | //! case performs 100'000 modifications or lookups. However, the dense cases randomly picks any IPv4 14 | //! address, while the sparse case only pick 20 different IPv4 addresses. See `benches/benchmark.rs` 15 | //! for more details. 16 | //! 17 | //! | Operation | Mode | `PrefixMap` | `treebitmap` | factor | 18 | //! |-----------------|--------|-------------|--------------|--------| 19 | //! | Insert & Remove | dense | **31.78ms** | 47.52ms | ~1.5x | 20 | //! | Lookup | dense | 32.36ms | **8.409ms** | ~0.25x | 21 | //! | Insert & Remove | sparse | **6.645ms** | 7.329ms | ~1.1x | 22 | //! | Lookup | sparse | **8.394ms** | 12.30ms | ~1.5x | 23 | //! 24 | //! 25 | //! In addition, `prefix-trie` includes a [`PrefixSet`] analogous to `std::collections::HashSet`, 26 | //! including union, intersection and difference operations that are implemented as simultaneous 27 | //! tree traversals. Further, `prefix-trie` has an interface similar to `std::collections`, and 28 | //! offers a general longest-prefix match that is not limited to individual addresses. Finally, 29 | //! `prefix-trie` allows you to (mutably) borrow a sub-trie using views. 30 | //! 31 | //! # Description of the Tree 32 | //! 33 | //! The tree is structured as follows: Each node consists of a prefix, a container for a potential 34 | //! value (`Option`), and two optional children. Adding a new child, or traversing into the tree is 35 | //! done as follows: we look at the most significant bit that is **not** part of the prefix 36 | //! itself. If it is not set, then we take the left branch, and otherwise, we take the right one. 37 | //! 38 | //! # Traversals 39 | //! 40 | //! Any iteration over all elements in the tree is implemented as a graph traversal that will yield 41 | //! elements in lexicographic order. 42 | //! 43 | //! The library offers set operations of different maps or sets. We implement a union, intersection, 44 | //! difference, and covering_difference. These iterators are implemented using simultaneous tree 45 | //! traversals. They will yield elements in lexicographic order. Whenever appropriate, the yielded 46 | //! items will also include the longest prefix match. 47 | //! 48 | //! # [`TrieView`] and [`TrieViewMut`] 49 | //! 50 | //! You can create a view of a (sub)-trie. Such a view has an any node as its root. Any operations 51 | //! on that view will only traverse that node and all its children. You can iterate over all 52 | //! children, search in that sub-trie, and perform set operations (union, intersection, difference, 53 | //! or the covering difference) on them. 54 | //! 55 | //! A view can point to one of three possible nodes: 56 | //! - A node in the tree that is actually present in the map, 57 | //! - A branching node that does not exist in the map, but is needed for the tree structure (or that 58 | //! was deleted using the function `remove_keep_tree`) 59 | //! - A virtual node that does not exist as a node in the tree. This is only the case if you call 60 | //! [`TrieView::find`] or [`AsView::view_at`] with a node that is not present in the tree, but 61 | //! that contains elements present in the tree. Virtual nodes are treated as if they are actually 62 | //! present in the tree as branching nodes. 63 | //! 64 | //! # Operations on the tree 65 | //! 66 | //! There are several operations one can do on the tree. Regular inserts are handled using the 67 | //! `Entry` structure. An `Entry` is a pointer to a location in the tree to either insert a value or 68 | //! modify an existing one. Removals however are different. 69 | //! 70 | //! The following are the computational complexities of the functions, where `n` is the number of 71 | //! elements in the tree. 72 | //! 73 | //! | Operation | Complexity | 74 | //! |--------------------------------------------|------------| 75 | //! | `entry`, `insert` | `O(log n)` | 76 | //! | `remove`, `remove_keep_tree` | `O(log n)` | 77 | //! | `remove_children` (calling `drop` on `T`) | `O(n)` | 78 | //! | `get`, `get_lpm`, `get_mut` | `O(log n)` | 79 | //! | `retain` | `O(n)` | 80 | //! | `clear` (calling `drop` on `T`) | `O(n)` | 81 | //! | Operations on [`map::Entry`] | `O(1)` | 82 | //! | `len` and `is_empty` | `O(1)` | 83 | //! | `union`, `intersection`, `difference`, ... | `O(n)` | 84 | //! 85 | //! There are three kinds of removals you! can do: 86 | //! 87 | //! - [`PrefixMap::remove`] will remove an entry from the tree and modify the tree structure as if 88 | //! the value was never inserted before. [`PrefixMap::remove`] will always exactly revert the 89 | //! operation of [`PrefixMap::insert`]. When only calling this function to remove elements, you 90 | //! are guaranteed that the tree structure is indistinguishable to a different tree where you 91 | //! only inserted elements. 92 | //! - [`PrefixMap::remove_children`] will remove all entries that are contained within the given 93 | //! prefix. This operation will search for the node with the shortest prefix length that is 94 | //! contained within the given prefix and remove it, including all of its children. 95 | //! - [`PrefixMap::remove_keep_tree`] will not change anything in the tree structure. It will only 96 | //! remove a value from a node. As soon as you call `remove_keep_tree` once on a tree structure, 97 | //! the tree will no longer be optimal. 98 | //! 99 | //! # TODO 100 | //! 101 | //! Migrate to a TreeBitMap, described by 102 | //! [W. Eatherton, Z. Dittia, G. Varghes](https://doi.org/10.1145/997150.997160). 103 | 104 | #![allow(clippy::collapsible_else_if)] 105 | #![deny(missing_docs)] 106 | 107 | mod fmt; 108 | #[cfg(test)] 109 | mod fuzzing; 110 | pub(crate) mod inner; 111 | mod prefix; 112 | #[cfg(feature = "serde")] 113 | mod serde; 114 | #[cfg(feature = "ipnet")] 115 | #[cfg(test)] 116 | mod test; 117 | 118 | pub mod map; 119 | pub mod set; 120 | pub mod trieview; 121 | 122 | pub use map::PrefixMap; 123 | pub use prefix::Prefix; 124 | pub use set::PrefixSet; 125 | pub use trieview::{AsView, AsViewMut, TrieView, TrieViewMut}; 126 | 127 | #[inline(always)] 128 | pub(crate) fn to_right(branch_p: &P, child_p: &P) -> bool { 129 | child_p.is_bit_set(branch_p.prefix_len()) 130 | } 131 | -------------------------------------------------------------------------------- /src/map/entry.rs: -------------------------------------------------------------------------------- 1 | //! Code for inserting elements and the entry pattern. 2 | 3 | use super::*; 4 | 5 | /// A mutable view into a single entry in a map, which may either be vacant or occupied. 6 | pub enum Entry<'a, P, T> { 7 | /// The entry is not present in the tree. 8 | Vacant(VacantEntry<'a, P, T>), 9 | /// The entry is already present in the tree. 10 | Occupied(OccupiedEntry<'a, P, T>), 11 | } 12 | 13 | /// A mutable view into a missing entry. The information within this structure describes a path 14 | /// towards that missing node, and how to insert it. 15 | pub struct VacantEntry<'a, P, T> { 16 | pub(super) map: &'a mut PrefixMap, 17 | pub(super) prefix: P, 18 | pub(super) idx: usize, 19 | pub(super) direction: DirectionForInsert

, 20 | } 21 | 22 | /// A mutable view into an occupied entry. An occupied entry represents a node that is already 23 | /// present on the tree. 24 | pub struct OccupiedEntry<'a, P, T> { 25 | pub(super) node: &'a mut Node, 26 | pub(super) prefix: P, // needed to replace the prefix on the thing if we perform insert. 27 | } 28 | 29 | impl Entry<'_, P, T> { 30 | /// Get the value if it exists 31 | /// 32 | /// ``` 33 | /// # use prefix_trie::*; 34 | /// # #[cfg(feature = "ipnet")] 35 | /// # fn main() -> Result<(), Box> { 36 | /// let mut pm: PrefixMap = PrefixMap::new(); 37 | /// pm.insert("192.168.1.0/24".parse()?, 1); 38 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).get(), Some(&1)); 39 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).get(), None); 40 | /// # Ok(()) 41 | /// # } 42 | /// # #[cfg(not(feature = "ipnet"))] 43 | /// # fn main() {} 44 | /// ``` 45 | pub fn get(&self) -> Option<&T> { 46 | match self { 47 | Entry::Vacant(_) => None, 48 | Entry::Occupied(e) => e.node.value.as_ref(), 49 | } 50 | } 51 | 52 | /// Get the value if it exists 53 | /// 54 | /// ``` 55 | /// # use prefix_trie::*; 56 | /// # #[cfg(feature = "ipnet")] 57 | /// # fn main() -> Result<(), Box> { 58 | /// let mut pm: PrefixMap = PrefixMap::new(); 59 | /// pm.insert("192.168.1.0/24".parse()?, 1); 60 | /// pm.entry("192.168.1.0/24".parse()?).get_mut().map(|x| *x += 1); 61 | /// pm.entry("192.168.2.0/24".parse()?).get_mut().map(|x| *x += 1); 62 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&2)); 63 | /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), None); 64 | /// # Ok(()) 65 | /// # } 66 | /// # #[cfg(not(feature = "ipnet"))] 67 | /// # fn main() {} 68 | /// ``` 69 | pub fn get_mut(&mut self) -> Option<&mut T> { 70 | match self { 71 | Entry::Vacant(_) => None, 72 | Entry::Occupied(e) => e.node.value.as_mut(), 73 | } 74 | } 75 | 76 | /// get the key of the current entry 77 | /// 78 | /// ``` 79 | /// # use prefix_trie::*; 80 | /// # #[cfg(feature = "ipnet")] 81 | /// # fn main() -> Result<(), Box> { 82 | /// let mut pm: PrefixMap = PrefixMap::new(); 83 | /// pm.insert("192.168.1.0/24".parse()?, 1); 84 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).key(), &"192.168.1.0/24".parse()?); 85 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).key(), &"192.168.2.0/24".parse()?); 86 | /// # Ok(()) 87 | /// # } 88 | /// # #[cfg(not(feature = "ipnet"))] 89 | /// # fn main() {} 90 | /// ``` 91 | pub fn key(&self) -> &P { 92 | match self { 93 | Entry::Vacant(e) => &e.prefix, 94 | Entry::Occupied(e) => &e.node.prefix, 95 | } 96 | } 97 | } 98 | 99 | impl<'a, P, T> Entry<'a, P, T> 100 | where 101 | P: Prefix, 102 | { 103 | /// Replace the current entry, and return the entry that was stored before. This will also 104 | /// replace the key with the one provided to the `entry` function. 105 | /// 106 | /// ``` 107 | /// # use prefix_trie::*; 108 | /// # #[cfg(feature = "ipnet")] 109 | /// # fn main() -> Result<(), Box> { 110 | /// let mut pm: PrefixMap = PrefixMap::new(); 111 | /// pm.insert("192.168.1.0/24".parse()?, 1); 112 | /// 113 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).insert(10), Some(1)); 114 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).insert(20), None); 115 | /// 116 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&10)); 117 | /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), Some(&20)); 118 | /// # Ok(()) 119 | /// # } 120 | /// # #[cfg(not(feature = "ipnet"))] 121 | /// # fn main() {} 122 | /// ``` 123 | #[inline(always)] 124 | pub fn insert(self, v: T) -> Option { 125 | match self { 126 | Entry::Vacant(e) => { 127 | e._insert(v); 128 | None 129 | } 130 | Entry::Occupied(e) => Some(e.insert(v)), 131 | } 132 | } 133 | 134 | /// Ensures a value is in the entry by inserting the default if empty, and returns a mutable 135 | /// reference to the value in the entry. 136 | /// 137 | /// ``` 138 | /// # use prefix_trie::*; 139 | /// # #[cfg(feature = "ipnet")] 140 | /// # fn main() -> Result<(), Box> { 141 | /// let mut pm: PrefixMap = PrefixMap::new(); 142 | /// pm.insert("192.168.1.0/24".parse()?, 1); 143 | /// 144 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).or_insert(10), &1); 145 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).or_insert(20), &20); 146 | /// 147 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&1)); 148 | /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), Some(&20)); 149 | /// # Ok(()) 150 | /// # } 151 | /// # #[cfg(not(feature = "ipnet"))] 152 | /// # fn main() {} 153 | /// ``` 154 | #[inline(always)] 155 | pub fn or_insert(self, default: T) -> &'a mut T { 156 | match self { 157 | Entry::Vacant(e) => e._insert(default).value.as_mut().unwrap(), 158 | Entry::Occupied(e) => e.node.value.get_or_insert(default), 159 | } 160 | } 161 | 162 | /// Ensures a value is in the entry by inserting the result of the default function if empty, 163 | /// and returns a mutable reference to the value in the entry. 164 | /// 165 | /// ``` 166 | /// # use prefix_trie::*; 167 | /// # #[cfg(feature = "ipnet")] 168 | /// # fn main() -> Result<(), Box> { 169 | /// let mut pm: PrefixMap = PrefixMap::new(); 170 | /// pm.insert("192.168.1.0/24".parse()?, 1); 171 | /// 172 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).or_insert_with(|| 10), &1); 173 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).or_insert_with(|| 20), &20); 174 | /// 175 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&1)); 176 | /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), Some(&20)); 177 | /// # Ok(()) 178 | /// # } 179 | /// # #[cfg(not(feature = "ipnet"))] 180 | /// # fn main() {} 181 | /// ``` 182 | #[inline(always)] 183 | pub fn or_insert_with T>(self, default: F) -> &'a mut T { 184 | match self { 185 | Entry::Vacant(e) => e._insert(default()).value.as_mut().unwrap(), 186 | Entry::Occupied(e) => e.node.value.get_or_insert_with(default), 187 | } 188 | } 189 | 190 | /// Provides in-place mutable access to an occupied entry before any potential inserts into the 191 | /// map. 192 | /// 193 | /// ``` 194 | /// # use prefix_trie::*; 195 | /// # #[cfg(feature = "ipnet")] 196 | /// # fn main() -> Result<(), Box> { 197 | /// let mut pm: PrefixMap = PrefixMap::new(); 198 | /// pm.insert("192.168.1.0/24".parse()?, 1); 199 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).and_modify(|x| *x += 1).get(), Some(&2)); 200 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).and_modify(|x| *x += 1).get(), None); 201 | /// # Ok(()) 202 | /// # } 203 | /// # #[cfg(not(feature = "ipnet"))] 204 | /// # fn main() {} 205 | /// ``` 206 | #[inline(always)] 207 | pub fn and_modify(self, f: F) -> Self { 208 | match self { 209 | Entry::Vacant(e) => Entry::Vacant(e), 210 | Entry::Occupied(e) => { 211 | e.node.value.as_mut().map(f); 212 | Entry::Occupied(e) 213 | } 214 | } 215 | } 216 | } 217 | 218 | impl<'a, P, T> Entry<'a, P, T> 219 | where 220 | P: Prefix, 221 | T: Default, 222 | { 223 | /// Ensures a value is in the entry by inserting the default value if empty, and returns a 224 | /// mutable reference to the value in the entry. 225 | /// 226 | /// ``` 227 | /// # use prefix_trie::*; 228 | /// # #[cfg(feature = "ipnet")] 229 | /// # fn main() -> Result<(), Box> { 230 | /// let mut pm: PrefixMap = PrefixMap::new(); 231 | /// pm.insert("192.168.1.0/24".parse()?, 1); 232 | /// 233 | /// assert_eq!(pm.entry("192.168.1.0/24".parse()?).or_default(), &1); 234 | /// assert_eq!(pm.entry("192.168.2.0/24".parse()?).or_default(), &0); 235 | /// 236 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&1)); 237 | /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), Some(&0)); 238 | /// # Ok(()) 239 | /// # } 240 | /// # #[cfg(not(feature = "ipnet"))] 241 | /// # fn main() {} 242 | /// ``` 243 | #[allow(clippy::unwrap_or_default)] 244 | #[inline(always)] 245 | pub fn or_default(self) -> &'a mut T { 246 | self.or_insert_with(Default::default) 247 | } 248 | } 249 | 250 | impl<'a, P, T> VacantEntry<'a, P, T> 251 | where 252 | P: Prefix, 253 | { 254 | fn _insert(self, v: T) -> &'a mut Node { 255 | match self.direction { 256 | DirectionForInsert::Reached => { 257 | // increment the count, as node.value will be `None`. We do it here as we borrow 258 | // `map` mutably in the next line. 259 | self.map.count += 1; 260 | let node = &mut self.map.table[self.idx]; 261 | node.prefix = self.prefix; 262 | debug_assert!(node.value.is_none()); 263 | node.value = Some(v); 264 | node 265 | } 266 | DirectionForInsert::NewLeaf { right } => { 267 | let new = self.map.new_node(self.prefix, Some(v)); 268 | self.map.table.set_child(self.idx, new, right); 269 | &mut self.map.table[new] 270 | } 271 | DirectionForInsert::NewChild { right, child_right } => { 272 | let new = self.map.new_node(self.prefix, Some(v)); 273 | let child = self.map.table.set_child(self.idx, new, right).unwrap(); 274 | self.map.table.set_child(new, child, child_right); 275 | &mut self.map.table[new] 276 | } 277 | DirectionForInsert::NewBranch { 278 | branch_prefix, 279 | right, 280 | prefix_right, 281 | } => { 282 | let branch = self.map.new_node(branch_prefix, None); 283 | let new = self.map.new_node(self.prefix, Some(v)); 284 | let child = self.map.table.set_child(self.idx, branch, right).unwrap(); 285 | self.map.table.set_child(branch, new, prefix_right); 286 | self.map.table.set_child(branch, child, !prefix_right); 287 | &mut self.map.table[new] 288 | } 289 | DirectionForInsert::Enter { .. } => unreachable!(), 290 | } 291 | } 292 | } 293 | 294 | impl OccupiedEntry<'_, P, T> { 295 | /// Gets a reference to the key in the entry. This is the key that is currently stored, and not 296 | /// the key that was used in the insert. 297 | /// 298 | /// ``` 299 | /// # use prefix_trie::*; 300 | /// use prefix_trie::map::Entry; 301 | /// # #[cfg(feature = "ipnet")] 302 | /// # fn main() -> Result<(), Box> { 303 | /// let mut pm: PrefixMap = PrefixMap::new(); 304 | /// pm.insert("192.168.1.0/24".parse()?, 1); 305 | /// match pm.entry("192.168.1.1/24".parse()?) { 306 | /// Entry::Occupied(e) => assert_eq!(e.key(), &"192.168.1.0/24".parse()?), 307 | /// Entry::Vacant(_) => unreachable!(), 308 | /// } 309 | /// # Ok(()) 310 | /// # } 311 | /// # #[cfg(not(feature = "ipnet"))] 312 | /// # fn main() {} 313 | /// ``` 314 | pub fn key(&self) -> &P { 315 | &self.node.prefix 316 | } 317 | 318 | /// Gets a reference to the value in the entry. 319 | /// 320 | /// ``` 321 | /// # use prefix_trie::*; 322 | /// use prefix_trie::map::Entry; 323 | /// # #[cfg(feature = "ipnet")] 324 | /// # fn main() -> Result<(), Box> { 325 | /// 326 | /// let mut pm: PrefixMap = PrefixMap::new(); 327 | /// pm.insert("192.168.1.0/24".parse()?, 1); 328 | /// match pm.entry("192.168.1.0/24".parse()?) { 329 | /// Entry::Occupied(e) => assert_eq!(e.get(), &1), 330 | /// Entry::Vacant(_) => unreachable!(), 331 | /// } 332 | /// # Ok(()) 333 | /// # } 334 | /// # #[cfg(not(feature = "ipnet"))] 335 | /// # fn main() {} 336 | /// ``` 337 | pub fn get(&self) -> &T { 338 | self.node.value.as_ref().unwrap() 339 | } 340 | 341 | /// Gets a mutable reference to the value in the entry. 342 | /// 343 | /// This call will not modify the prefix stored in the tree. In case the prefix used to create 344 | /// the entry is different from the stored one (has additional information in the host part), 345 | /// and you wish that prefix to be overwritten, use `insert`. 346 | /// 347 | /// ``` 348 | /// # use prefix_trie::*; 349 | /// use prefix_trie::map::Entry; 350 | /// # #[cfg(feature = "ipnet")] 351 | /// # fn main() -> Result<(), Box> { 352 | /// 353 | /// let mut pm: PrefixMap = PrefixMap::new(); 354 | /// pm.insert("192.168.1.0/24".parse()?, 1); 355 | /// match pm.entry("192.168.1.0/24".parse()?) { 356 | /// Entry::Occupied(mut e) => *e.get_mut() += 1, 357 | /// Entry::Vacant(_) => unreachable!(), 358 | /// } 359 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&2)); 360 | /// # Ok(()) 361 | /// # } 362 | /// # #[cfg(not(feature = "ipnet"))] 363 | /// # fn main() {} 364 | /// ``` 365 | pub fn get_mut(&mut self) -> &mut T { 366 | self.node.value.as_mut().unwrap() 367 | } 368 | 369 | /// Insert a new value into the entry, returning the old value. This operation will also replace 370 | /// the prefix with the provided one. 371 | /// 372 | /// ``` 373 | /// # use prefix_trie::*; 374 | /// use prefix_trie::map::Entry; 375 | /// # #[cfg(feature = "ipnet")] 376 | /// # fn main() -> Result<(), Box> { 377 | /// 378 | /// let mut pm: PrefixMap = PrefixMap::new(); 379 | /// pm.insert("192.168.1.0/24".parse()?, 1); 380 | /// match pm.entry("192.168.1.0/24".parse()?) { 381 | /// Entry::Occupied(mut e) => assert_eq!(e.insert(10), 1), 382 | /// Entry::Vacant(_) => unreachable!(), 383 | /// } 384 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&10)); 385 | /// # Ok(()) 386 | /// # } 387 | /// # #[cfg(not(feature = "ipnet"))] 388 | /// # fn main() {} 389 | /// ``` 390 | pub fn insert(self, value: T) -> T { 391 | self.node.prefix = self.prefix; 392 | self.node.value.replace(value).unwrap() 393 | } 394 | 395 | /// Remove the current value and return it. The tree will not be modified (the same effect as 396 | /// `PrefixMap::remove_keep_tree`). 397 | /// 398 | /// ``` 399 | /// # use prefix_trie::*; 400 | /// use prefix_trie::map::Entry; 401 | /// # #[cfg(feature = "ipnet")] 402 | /// # fn main() -> Result<(), Box> { 403 | /// 404 | /// let mut pm: PrefixMap = PrefixMap::new(); 405 | /// pm.insert("192.168.1.0/24".parse()?, 1); 406 | /// match pm.entry("192.168.1.0/24".parse()?) { 407 | /// Entry::Occupied(mut e) => assert_eq!(e.remove(), 1), 408 | /// Entry::Vacant(_) => unreachable!(), 409 | /// } 410 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), None); 411 | /// # Ok(()) 412 | /// # } 413 | /// # #[cfg(not(feature = "ipnet"))] 414 | /// # fn main() {} 415 | /// ``` 416 | pub fn remove(&mut self) -> T { 417 | self.node.value.take().unwrap() 418 | } 419 | } 420 | 421 | impl VacantEntry<'_, P, T> { 422 | /// Gets a reference to the key in the entry. 423 | /// 424 | /// ``` 425 | /// # use prefix_trie::*; 426 | /// use prefix_trie::map::Entry; 427 | /// # #[cfg(feature = "ipnet")] 428 | /// # fn main() -> Result<(), Box> { 429 | /// let mut pm: PrefixMap = PrefixMap::new(); 430 | /// match pm.entry("192.168.1.0/24".parse()?) { 431 | /// Entry::Vacant(e) => assert_eq!(e.key(), &"192.168.1.0/24".parse()?), 432 | /// Entry::Occupied(_) => unreachable!(), 433 | /// } 434 | /// # Ok(()) 435 | /// # } 436 | /// # #[cfg(not(feature = "ipnet"))] 437 | /// # fn main() {} 438 | /// ``` 439 | pub fn key(&self) -> &P { 440 | &self.prefix 441 | } 442 | } 443 | 444 | impl<'a, P, T> VacantEntry<'a, P, T> 445 | where 446 | P: Prefix, 447 | { 448 | /// Get a mutable reference to the value. If the value is yet empty, set it to the given default 449 | /// value. 450 | /// 451 | /// ``` 452 | /// # use prefix_trie::*; 453 | /// use prefix_trie::map::Entry; 454 | /// # #[cfg(feature = "ipnet")] 455 | /// # fn main() -> Result<(), Box> { 456 | /// let mut pm: PrefixMap = PrefixMap::new(); 457 | /// match pm.entry("192.168.1.0/24".parse()?) { 458 | /// Entry::Vacant(mut e) => assert_eq!(e.insert(10), &10), 459 | /// Entry::Occupied(_) => unreachable!(), 460 | /// } 461 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&10)); 462 | /// # Ok(()) 463 | /// # } 464 | /// # #[cfg(not(feature = "ipnet"))] 465 | /// # fn main() {} 466 | /// ``` 467 | pub fn insert(self, default: T) -> &'a mut T { 468 | let node = self._insert(default); 469 | node.value.as_mut().unwrap() 470 | } 471 | 472 | /// Get a mutable reference to the value. If the value is yet empty, set it to the return value 473 | /// from the given function. 474 | /// 475 | /// ``` 476 | /// # use prefix_trie::*; 477 | /// use prefix_trie::map::Entry; 478 | /// # #[cfg(feature = "ipnet")] 479 | /// # fn main() -> Result<(), Box> { 480 | /// let mut pm: PrefixMap = PrefixMap::new(); 481 | /// match pm.entry("192.168.1.0/24".parse()?) { 482 | /// Entry::Vacant(mut e) => assert_eq!(e.insert_with(|| 10), &10), 483 | /// Entry::Occupied(_) => unreachable!(), 484 | /// } 485 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&10)); 486 | /// # Ok(()) 487 | /// # } 488 | /// # #[cfg(not(feature = "ipnet"))] 489 | /// # fn main() {} 490 | /// ``` 491 | pub fn insert_with T>(self, default: F) -> &'a mut T { 492 | let node = self._insert(default()); 493 | node.value.as_mut().unwrap() 494 | } 495 | } 496 | 497 | impl<'a, P, T> VacantEntry<'a, P, T> 498 | where 499 | P: Prefix, 500 | T: Default, 501 | { 502 | /// Get a mutable reference to the value. If the value is yet empty, set it to the default value 503 | /// using `Default::default()`. 504 | /// 505 | /// ``` 506 | /// # use prefix_trie::*; 507 | /// use prefix_trie::map::Entry; 508 | /// # #[cfg(feature = "ipnet")] 509 | /// # fn main() -> Result<(), Box> { 510 | /// let mut pm: PrefixMap = PrefixMap::new(); 511 | /// match pm.entry("192.168.1.0/24".parse()?) { 512 | /// Entry::Vacant(e) => assert_eq!(e.default(), &0), 513 | /// Entry::Occupied(_) => unreachable!(), 514 | /// } 515 | /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&0)); 516 | /// # Ok(()) 517 | /// # } 518 | /// # #[cfg(not(feature = "ipnet"))] 519 | /// # fn main() {} 520 | /// ``` 521 | pub fn default(self) -> &'a mut T { 522 | let node = self._insert(Default::default()); 523 | node.value.as_mut().unwrap() 524 | } 525 | } 526 | -------------------------------------------------------------------------------- /src/map/iter.rs: -------------------------------------------------------------------------------- 1 | //! Module that contains the implementation for the iterators 2 | 3 | use map::Table; 4 | 5 | use crate::*; 6 | 7 | use super::Node; 8 | 9 | /// An iterator over all entries of a [`PrefixMap`] in lexicographic order. 10 | #[derive(Clone)] 11 | pub struct Iter<'a, P, T> { 12 | table: Option<&'a Table>, 13 | nodes: Vec, 14 | } 15 | 16 | impl Default for Iter<'_, P, T> { 17 | fn default() -> Self { 18 | Self { 19 | table: None, 20 | nodes: Vec::new(), 21 | } 22 | } 23 | } 24 | 25 | impl<'a, P, T> Iter<'a, P, T> { 26 | pub(crate) fn new(table: &'a Table, nodes: Vec) -> Self { 27 | Self { 28 | table: Some(table), 29 | nodes, 30 | } 31 | } 32 | } 33 | 34 | impl<'a, P, T> Iterator for Iter<'a, P, T> { 35 | type Item = (&'a P, &'a T); 36 | 37 | fn next(&mut self) -> Option<(&'a P, &'a T)> { 38 | while let Some(cur) = self.nodes.pop() { 39 | let node = &self.table.as_ref()?[cur]; 40 | if let Some(right) = node.right { 41 | self.nodes.push(right); 42 | } 43 | if let Some(left) = node.left { 44 | self.nodes.push(left); 45 | } 46 | if let Some(v) = &node.value { 47 | return Some((&node.prefix, v)); 48 | } 49 | } 50 | None 51 | } 52 | } 53 | 54 | /// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order. 55 | #[derive(Clone, Default)] 56 | pub struct Keys<'a, P, T> { 57 | pub(crate) inner: Iter<'a, P, T>, 58 | } 59 | 60 | impl<'a, P, T> Iterator for Keys<'a, P, T> { 61 | type Item = &'a P; 62 | 63 | fn next(&mut self) -> Option<&'a P> { 64 | self.inner.next().map(|(k, _)| k) 65 | } 66 | } 67 | 68 | /// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated 69 | /// prefixes. 70 | #[derive(Clone, Default)] 71 | pub struct Values<'a, P, T> { 72 | pub(crate) inner: Iter<'a, P, T>, 73 | } 74 | 75 | impl<'a, P, T> Iterator for Values<'a, P, T> { 76 | type Item = &'a T; 77 | 78 | fn next(&mut self) -> Option<&'a T> { 79 | self.inner.next().map(|(_, v)| v) 80 | } 81 | } 82 | 83 | /// An iterator over all owned entries of a [`PrefixMap`] in lexicographic order. 84 | #[derive(Clone)] 85 | pub struct IntoIter { 86 | table: Vec>, 87 | nodes: Vec, 88 | } 89 | 90 | impl Iterator for IntoIter { 91 | type Item = (P, T); 92 | 93 | fn next(&mut self) -> Option<(P, T)> { 94 | while let Some(cur) = self.nodes.pop() { 95 | let node = &mut self.table[cur]; 96 | if let Some(right) = node.right { 97 | self.nodes.push(right); 98 | } 99 | if let Some(left) = node.left { 100 | self.nodes.push(left); 101 | } 102 | if let Some(v) = node.value.take() { 103 | return Some((std::mem::replace(&mut node.prefix, P::zero()), v)); 104 | } 105 | } 106 | None 107 | } 108 | } 109 | 110 | /// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order. 111 | #[derive(Clone)] 112 | pub struct IntoKeys { 113 | inner: IntoIter, 114 | } 115 | 116 | impl Iterator for IntoKeys { 117 | type Item = P; 118 | 119 | fn next(&mut self) -> Option

{ 120 | self.inner.next().map(|(k, _)| k) 121 | } 122 | } 123 | 124 | /// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated 125 | /// prefix. 126 | #[derive(Clone)] 127 | pub struct IntoValues { 128 | inner: IntoIter, 129 | } 130 | 131 | impl Iterator for IntoValues { 132 | type Item = T; 133 | 134 | fn next(&mut self) -> Option { 135 | self.inner.next().map(|(_, v)| v) 136 | } 137 | } 138 | 139 | impl IntoIterator for PrefixMap { 140 | type Item = (P, T); 141 | 142 | type IntoIter = IntoIter; 143 | 144 | fn into_iter(self) -> Self::IntoIter { 145 | IntoIter { 146 | table: self.table.into_inner(), 147 | nodes: vec![0], 148 | } 149 | } 150 | } 151 | 152 | impl<'a, P, T> IntoIterator for &'a PrefixMap { 153 | type Item = (&'a P, &'a T); 154 | 155 | type IntoIter = Iter<'a, P, T>; 156 | 157 | fn into_iter(self) -> Self::IntoIter { 158 | // Safety: we own an immutable reference, and `Iter` will only ever read the table. 159 | Iter::new(&self.table, vec![0]) 160 | } 161 | } 162 | 163 | /// A mutable iterator over a [`PrefixMap`]. This iterator yields elements in lexicographic order of 164 | /// their associated prefix. 165 | pub struct IterMut<'a, P, T> { 166 | table: Option<&'a Table>, 167 | nodes: Vec, 168 | } 169 | 170 | impl Default for IterMut<'_, P, T> { 171 | fn default() -> Self { 172 | Self { 173 | table: None, 174 | nodes: Vec::new(), 175 | } 176 | } 177 | } 178 | 179 | impl<'a, P, T> IterMut<'a, P, T> { 180 | /// # Safety 181 | /// - First, you must ensure that 'a is tied to a mutable reference of the original table. 182 | /// - Second, you are allowed to create mutiple such `IterMut`s, as long as none of the root 183 | /// nodes is the parent of another root node (of any of the iterators). This also applies if 184 | /// you only create a single iterator with multiple root nodes. 185 | /// 186 | /// The iterator will only ever access its roots or its children. 187 | pub(crate) unsafe fn new(table: &'a Table, nodes: Vec) -> Self { 188 | Self { 189 | table: Some(table), 190 | nodes, 191 | } 192 | } 193 | } 194 | 195 | impl<'a, P, T> Iterator for IterMut<'a, P, T> { 196 | type Item = (&'a P, &'a mut T); 197 | 198 | fn next(&mut self) -> Option { 199 | while let Some(cur) = self.nodes.pop() { 200 | // Safety: 201 | // In the following, we assume that there are no two iterators that may reach the same 202 | // sub-tree (see the safety comment above). 203 | // 204 | // The iterator borrows from &'a mut PrefixMap, see `PrefixMap::iter_mut` where 'a is 205 | // linked to a mutable reference. Then, we must ensure that we only ever construct a 206 | // mutable reference to each element exactly once. We ensure this by the fact that we 207 | // iterate over a tree. Thus, each node is visited exactly once. 208 | let node: &'a mut Node = unsafe { self.table.as_ref()?.get_mut(cur) }; 209 | 210 | if let Some(right) = node.right { 211 | self.nodes.push(right); 212 | } 213 | if let Some(left) = node.left { 214 | self.nodes.push(left); 215 | } 216 | if node.value.is_some() { 217 | let v = node.value.as_mut().unwrap(); 218 | return Some((&node.prefix, v)); 219 | } 220 | } 221 | None 222 | } 223 | } 224 | 225 | /// A mutable iterator over values of [`PrefixMap`]. This iterator yields elements in lexicographic 226 | /// order. 227 | #[derive(Default)] 228 | pub struct ValuesMut<'a, P, T> { 229 | // # Safety 230 | // You must ensure that there only ever exists one such iterator for each tree. You may create 231 | // multiple such iterators for the same tree if you start with distinct starting nodes! This 232 | // ensures that any one iteration will never yield elements of the other iterator. 233 | pub(crate) inner: IterMut<'a, P, T>, 234 | } 235 | 236 | impl<'a, P, T> Iterator for ValuesMut<'a, P, T> { 237 | type Item = &'a mut T; 238 | 239 | fn next(&mut self) -> Option { 240 | self.inner.next().map(|(_, v)| v) 241 | } 242 | } 243 | 244 | impl PrefixMap { 245 | /// An iterator visiting all key-value pairs in lexicographic order. The iterator element type 246 | /// is `(&P, &T)`. 247 | /// 248 | /// ``` 249 | /// # use prefix_trie::*; 250 | /// # #[cfg(feature = "ipnet")] 251 | /// # fn main() -> Result<(), Box> { 252 | /// let mut pm: PrefixMap = PrefixMap::new(); 253 | /// pm.insert("192.168.0.0/22".parse()?, 1); 254 | /// pm.insert("192.168.0.0/23".parse()?, 2); 255 | /// pm.insert("192.168.2.0/23".parse()?, 3); 256 | /// pm.insert("192.168.0.0/24".parse()?, 4); 257 | /// pm.insert("192.168.2.0/24".parse()?, 5); 258 | /// assert_eq!( 259 | /// pm.iter().collect::>(), 260 | /// vec![ 261 | /// (&"192.168.0.0/22".parse()?, &1), 262 | /// (&"192.168.0.0/23".parse()?, &2), 263 | /// (&"192.168.0.0/24".parse()?, &4), 264 | /// (&"192.168.2.0/23".parse()?, &3), 265 | /// (&"192.168.2.0/24".parse()?, &5), 266 | /// ] 267 | /// ); 268 | /// # Ok(()) 269 | /// # } 270 | /// # #[cfg(not(feature = "ipnet"))] 271 | /// # fn main() {} 272 | /// ``` 273 | #[inline(always)] 274 | pub fn iter(&self) -> Iter<'_, P, T> { 275 | self.into_iter() 276 | } 277 | 278 | /// Get a mutable iterator over all key-value pairs. The order of this iterator is lexicographic. 279 | pub fn iter_mut(&mut self) -> IterMut<'_, P, T> { 280 | // Safety: We get the pointer to the table by and construct the `IterMut`. Its lifetime is 281 | // now tied to the mutable borrow of `self`, so we are allowed to access elements of that 282 | // table mutably. 283 | unsafe { IterMut::new(&self.table, vec![0]) } 284 | } 285 | 286 | /// An iterator visiting all keys in lexicographic order. The iterator element type is `&P`. 287 | /// 288 | /// ``` 289 | /// # use prefix_trie::*; 290 | /// # #[cfg(feature = "ipnet")] 291 | /// # fn main() -> Result<(), Box> { 292 | /// let mut pm: PrefixMap = PrefixMap::new(); 293 | /// pm.insert("192.168.0.0/22".parse()?, 1); 294 | /// pm.insert("192.168.0.0/23".parse()?, 2); 295 | /// pm.insert("192.168.2.0/23".parse()?, 3); 296 | /// pm.insert("192.168.0.0/24".parse()?, 4); 297 | /// pm.insert("192.168.2.0/24".parse()?, 5); 298 | /// assert_eq!( 299 | /// pm.keys().collect::>(), 300 | /// vec![ 301 | /// &"192.168.0.0/22".parse()?, 302 | /// &"192.168.0.0/23".parse()?, 303 | /// &"192.168.0.0/24".parse()?, 304 | /// &"192.168.2.0/23".parse()?, 305 | /// &"192.168.2.0/24".parse()?, 306 | /// ] 307 | /// ); 308 | /// # Ok(()) 309 | /// # } 310 | /// # #[cfg(not(feature = "ipnet"))] 311 | /// # fn main() {} 312 | /// ``` 313 | #[inline(always)] 314 | pub fn keys(&self) -> Keys<'_, P, T> { 315 | Keys { inner: self.iter() } 316 | } 317 | 318 | /// Creates a consuming iterator visiting all keys in lexicographic order. The iterator element 319 | /// type is `P`. 320 | #[inline(always)] 321 | pub fn into_keys(self) -> IntoKeys { 322 | IntoKeys { 323 | inner: IntoIter { 324 | table: self.table.into_inner(), 325 | nodes: vec![0], 326 | }, 327 | } 328 | } 329 | 330 | /// An iterator visiting all values in lexicographic order. The iterator element type is `&P`. 331 | /// 332 | /// ``` 333 | /// # use prefix_trie::*; 334 | /// # #[cfg(feature = "ipnet")] 335 | /// # fn main() -> Result<(), Box> { 336 | /// let mut pm: PrefixMap = PrefixMap::new(); 337 | /// pm.insert("192.168.0.0/22".parse()?, 1); 338 | /// pm.insert("192.168.0.0/23".parse()?, 2); 339 | /// pm.insert("192.168.2.0/23".parse()?, 3); 340 | /// pm.insert("192.168.0.0/24".parse()?, 4); 341 | /// pm.insert("192.168.2.0/24".parse()?, 5); 342 | /// assert_eq!(pm.values().collect::>(), vec![&1, &2, &4, &3, &5]); 343 | /// # Ok(()) 344 | /// # } 345 | /// # #[cfg(not(feature = "ipnet"))] 346 | /// # fn main() {} 347 | /// ``` 348 | #[inline(always)] 349 | pub fn values(&self) -> Values<'_, P, T> { 350 | Values { inner: self.iter() } 351 | } 352 | 353 | /// Creates a consuming iterator visiting all values in lexicographic order. The iterator 354 | /// element type is `P`. 355 | #[inline(always)] 356 | pub fn into_values(self) -> IntoValues { 357 | IntoValues { 358 | inner: IntoIter { 359 | table: self.table.into_inner(), 360 | nodes: vec![0], 361 | }, 362 | } 363 | } 364 | 365 | /// Get a mutable iterator over all values. The order of this iterator is lexicographic. 366 | pub fn values_mut(&mut self) -> ValuesMut<'_, P, T> { 367 | ValuesMut { 368 | inner: self.iter_mut(), 369 | } 370 | } 371 | } 372 | 373 | impl PrefixMap 374 | where 375 | P: Prefix, 376 | { 377 | /// Get an iterator over the node itself and all children. All elements returned have a prefix 378 | /// that is contained within `prefix` itself (or are the same). The iterator yields references 379 | /// to both keys and values, i.e., type `(&'a P, &'a T)`. The iterator yields elements in 380 | /// lexicographic order. 381 | /// 382 | /// **Note**: Consider using [`AsView::view_at`] as an alternative. 383 | /// 384 | /// ``` 385 | /// # use prefix_trie::*; 386 | /// # #[cfg(feature = "ipnet")] 387 | /// # fn main() -> Result<(), Box> { 388 | /// let mut pm: PrefixMap = PrefixMap::new(); 389 | /// pm.insert("192.168.0.0/22".parse()?, 1); 390 | /// pm.insert("192.168.0.0/23".parse()?, 2); 391 | /// pm.insert("192.168.2.0/23".parse()?, 3); 392 | /// pm.insert("192.168.0.0/24".parse()?, 4); 393 | /// pm.insert("192.168.2.0/24".parse()?, 5); 394 | /// assert_eq!( 395 | /// pm.children(&"192.168.0.0/23".parse()?).collect::>(), 396 | /// vec![ 397 | /// (&"192.168.0.0/23".parse()?, &2), 398 | /// (&"192.168.0.0/24".parse()?, &4), 399 | /// ] 400 | /// ); 401 | /// # Ok(()) 402 | /// # } 403 | /// # #[cfg(not(feature = "ipnet"))] 404 | /// # fn main() {} 405 | /// ``` 406 | pub fn children<'a>(&'a self, prefix: &P) -> Iter<'a, P, T> { 407 | let nodes = lpm_children_iter_start(&self.table, prefix); 408 | Iter { 409 | table: Some(&self.table), 410 | nodes, 411 | } 412 | } 413 | 414 | /// Get an iterator of mutable references of the node itself and all its children. All elements 415 | /// returned have a prefix that is contained within `prefix` itself (or are the same). The 416 | /// iterator yields references to the keys, and mutable references to the values, i.e., type 417 | /// `(&'a P, &'a mut T)`. The iterator yields elements in lexicographic order. 418 | /// 419 | /// **Note**: Consider using [`AsViewMut::view_mut_at`] as an alternative. 420 | /// 421 | /// ``` 422 | /// # use prefix_trie::*; 423 | /// # #[cfg(feature = "ipnet")] 424 | /// # fn main() -> Result<(), Box> { 425 | /// let mut pm: PrefixMap = PrefixMap::new(); 426 | /// pm.insert("192.168.0.0/22".parse()?, 1); 427 | /// pm.insert("192.168.0.0/23".parse()?, 2); 428 | /// pm.insert("192.168.0.0/24".parse()?, 3); 429 | /// pm.insert("192.168.2.0/23".parse()?, 4); 430 | /// pm.insert("192.168.2.0/24".parse()?, 5); 431 | /// pm.children_mut(&"192.168.0.0/23".parse()?).for_each(|(_, x)| *x *= 10); 432 | /// assert_eq!( 433 | /// pm.into_iter().collect::>(), 434 | /// vec![ 435 | /// ("192.168.0.0/22".parse()?, 1), 436 | /// ("192.168.0.0/23".parse()?, 20), 437 | /// ("192.168.0.0/24".parse()?, 30), 438 | /// ("192.168.2.0/23".parse()?, 4), 439 | /// ("192.168.2.0/24".parse()?, 5), 440 | /// ] 441 | /// ); 442 | /// # Ok(()) 443 | /// # } 444 | /// # #[cfg(not(feature = "ipnet"))] 445 | /// # fn main() {} 446 | /// ``` 447 | pub fn children_mut<'a>(&'a mut self, prefix: &P) -> IterMut<'a, P, T> { 448 | let nodes = lpm_children_iter_start(&self.table, prefix); 449 | IterMut { 450 | table: Some(&self.table), 451 | nodes, 452 | } 453 | } 454 | 455 | /// Get an iterator over the node itself and all children with a value. All elements returned 456 | /// have a prefix that is contained within `prefix` itself (or are the same). This function will 457 | /// consume `self`, returning an iterator over all owned children. 458 | /// 459 | /// ``` 460 | /// # use prefix_trie::*; 461 | /// # #[cfg(feature = "ipnet")] 462 | /// # fn main() -> Result<(), Box> { 463 | /// let mut pm: PrefixMap = PrefixMap::new(); 464 | /// pm.insert("192.168.0.0/22".parse()?, 1); 465 | /// pm.insert("192.168.0.0/23".parse()?, 2); 466 | /// pm.insert("192.168.2.0/23".parse()?, 3); 467 | /// pm.insert("192.168.0.0/24".parse()?, 4); 468 | /// pm.insert("192.168.2.0/24".parse()?, 5); 469 | /// assert_eq!( 470 | /// pm.into_children(&"192.168.0.0/23".parse()?).collect::>(), 471 | /// vec![ 472 | /// ("192.168.0.0/23".parse()?, 2), 473 | /// ("192.168.0.0/24".parse()?, 4), 474 | /// ] 475 | /// ); 476 | /// # Ok(()) 477 | /// # } 478 | /// # #[cfg(not(feature = "ipnet"))] 479 | /// # fn main() {} 480 | /// ``` 481 | pub fn into_children(self, prefix: &P) -> IntoIter { 482 | let nodes = lpm_children_iter_start(&self.table, prefix); 483 | IntoIter { 484 | table: self.table.into_inner(), 485 | nodes, 486 | } 487 | } 488 | } 489 | 490 | fn lpm_children_iter_start(table: &Table, prefix: &P) -> Vec { 491 | let mut idx = 0; 492 | let mut cur_p = &table[idx].prefix; 493 | 494 | loop { 495 | if cur_p.eq(prefix) { 496 | break vec![idx]; 497 | } 498 | let right = to_right(cur_p, prefix); 499 | match table.get_child(idx, right) { 500 | Some(c) => { 501 | cur_p = &table[c].prefix; 502 | if cur_p.contains(prefix) { 503 | // continue traversal 504 | idx = c; 505 | } else if prefix.contains(cur_p) { 506 | break vec![c]; 507 | } else { 508 | break vec![]; 509 | } 510 | } 511 | None => break vec![], 512 | } 513 | } 514 | } 515 | 516 | impl FromIterator<(P, T)> for PrefixMap 517 | where 518 | P: Prefix, 519 | { 520 | fn from_iter>(iter: I) -> Self { 521 | let mut map = Self::new(); 522 | iter.into_iter().for_each(|(p, v)| { 523 | map.insert(p, v); 524 | }); 525 | map 526 | } 527 | } 528 | 529 | /// An iterator that yields all items in a `PrefixMap` that covers a given prefix (including the 530 | /// prefix itself if preseint). See [`PrefixMap::cover`] for how to create this iterator. 531 | pub struct Cover<'a, 'p, P, T> { 532 | pub(super) table: &'a Table, 533 | pub(super) idx: Option, 534 | pub(super) prefix: &'p P, 535 | } 536 | 537 | impl<'a, P, T> Iterator for Cover<'a, '_, P, T> 538 | where 539 | P: Prefix, 540 | { 541 | type Item = (&'a P, &'a T); 542 | 543 | fn next(&mut self) -> Option { 544 | // check if self.idx is None. If so, then check if the first branch is present in the map 545 | if self.idx.is_none() { 546 | self.idx = Some(0); 547 | let entry = &self.table[0]; 548 | if let Some(value) = entry.value.as_ref() { 549 | return Some((&entry.prefix, value)); 550 | } 551 | } 552 | 553 | // if we reach here, then self.idx is not None! 554 | 555 | loop { 556 | let map::Direction::Enter { next, .. } = 557 | self.table.get_direction(self.idx.unwrap(), self.prefix) 558 | else { 559 | return None; 560 | }; 561 | self.idx = Some(next); 562 | let entry = &self.table[next]; 563 | if let Some(value) = entry.value.as_ref() { 564 | return Some((&entry.prefix, value)); 565 | } 566 | } 567 | } 568 | } 569 | 570 | /// An iterator that yields all keys (prefixes) in a `PrefixMap` that covers a given prefix 571 | /// (including the prefix itself if preseint). See [`PrefixMap::cover_keys`] for how to create this 572 | /// iterator. 573 | pub struct CoverKeys<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>); 574 | 575 | impl<'a, P, T> Iterator for CoverKeys<'a, '_, P, T> 576 | where 577 | P: Prefix, 578 | { 579 | type Item = &'a P; 580 | 581 | fn next(&mut self) -> Option { 582 | self.0.next().map(|(p, _)| p) 583 | } 584 | } 585 | 586 | /// An iterator that yields all values in a `PrefixMap` that covers a given prefix (including the 587 | /// prefix itself if preseint). See [`PrefixMap::cover_values`] for how to create this iterator. 588 | pub struct CoverValues<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>); 589 | 590 | impl<'a, P, T> Iterator for CoverValues<'a, '_, P, T> 591 | where 592 | P: Prefix, 593 | { 594 | type Item = &'a T; 595 | 596 | fn next(&mut self) -> Option { 597 | self.0.next().map(|(_, t)| t) 598 | } 599 | } 600 | -------------------------------------------------------------------------------- /src/prefix.rs: -------------------------------------------------------------------------------- 1 | //! Description of the generic type `Prefix`. 2 | 3 | #[cfg(feature = "cidr")] 4 | use cidr::{Ipv4Cidr, Ipv4Inet, Ipv6Cidr, Ipv6Inet}; 5 | #[cfg(feature = "ipnet")] 6 | use ipnet::{Ipv4Net, Ipv6Net}; 7 | #[cfg(feature = "ipnetwork")] 8 | use ipnetwork::{Ipv4Network, Ipv6Network}; 9 | use num_traits::{CheckedShr, PrimInt, Unsigned, Zero}; 10 | 11 | /// Trait for defining prefixes. 12 | pub trait Prefix: Sized + std::fmt::Debug { 13 | /// How can the prefix be represented. This must be one of `u8`, `u16`, `u32`, `u64`, or `u128`. 14 | type R: Unsigned + PrimInt + Zero + CheckedShr; 15 | 16 | /// Get raw representation of the address, ignoring the prefix length. This function must return 17 | /// the representation with the mask already applied. 18 | fn repr(&self) -> Self::R; 19 | 20 | /// Prefix length 21 | fn prefix_len(&self) -> u8; 22 | 23 | /// Create a new prefix from the representation and the prefix pength. 24 | fn from_repr_len(repr: Self::R, len: u8) -> Self; 25 | 26 | /// mask `self.repr()` using `self.len()`. If you can guarantee that `repr` is already masked, 27 | /// them simply re-implement this function for your type. 28 | fn mask(&self) -> Self::R { 29 | self.repr() & mask_from_prefix_len(self.prefix_len()) 30 | } 31 | 32 | /// Create a prefix that matches everything 33 | fn zero() -> Self { 34 | Self::from_repr_len(Self::R::zero(), 0) 35 | } 36 | 37 | /// longest common prefix 38 | fn longest_common_prefix(&self, other: &Self) -> Self { 39 | let a = self.mask(); 40 | let b = other.mask(); 41 | let len = ((a ^ b).leading_zeros() as u8) 42 | .min(self.prefix_len()) 43 | .min(other.prefix_len()); 44 | let repr = a & mask_from_prefix_len(len); 45 | Self::from_repr_len(repr, len) 46 | } 47 | 48 | /// Check if `self` contains `other` in its prefix range. This function also returns `True` if 49 | /// `self` is identical to `other`. 50 | fn contains(&self, other: &Self) -> bool { 51 | if self.prefix_len() > other.prefix_len() { 52 | return false; 53 | } 54 | other.repr() & mask_from_prefix_len(self.prefix_len()) == self.mask() 55 | } 56 | 57 | /// Check if a specific bit is set (counted from the left, where 0 is the first bit from the 58 | /// left). 59 | fn is_bit_set(&self, bit: u8) -> bool { 60 | let mask = (!Self::R::zero()) 61 | .checked_shr(bit as u32) 62 | .unwrap_or_else(Self::R::zero) 63 | ^ (!Self::R::zero()) 64 | .checked_shr(1u32 + bit as u32) 65 | .unwrap_or_else(Self::R::zero); 66 | mask & self.mask() != Self::R::zero() 67 | } 68 | 69 | /// Compare two prefixes together 70 | fn eq(&self, other: &Self) -> bool { 71 | self.mask() == other.mask() && self.prefix_len() == other.prefix_len() 72 | } 73 | } 74 | 75 | pub(crate) fn mask_from_prefix_len(len: u8) -> R 76 | where 77 | R: PrimInt + Zero, 78 | { 79 | if len as u32 == R::zero().count_zeros() { 80 | !R::zero() 81 | } else if len == 0 { 82 | R::zero() 83 | } else { 84 | !((!R::zero()) >> len as usize) 85 | } 86 | } 87 | 88 | #[cfg(feature = "ipnet")] 89 | impl Prefix for Ipv4Net { 90 | type R = u32; 91 | 92 | fn repr(&self) -> u32 { 93 | self.addr().into() 94 | } 95 | 96 | fn prefix_len(&self) -> u8 { 97 | self.prefix_len() 98 | } 99 | 100 | fn from_repr_len(repr: u32, len: u8) -> Self { 101 | Ipv4Net::new(repr.into(), len).unwrap() 102 | } 103 | 104 | fn mask(&self) -> u32 { 105 | self.network().into() 106 | } 107 | 108 | fn zero() -> Self { 109 | Default::default() 110 | } 111 | 112 | fn longest_common_prefix(&self, other: &Self) -> Self { 113 | let a = self.repr(); 114 | let b = other.repr(); 115 | let len = ((a ^ b).leading_zeros() as u8) 116 | .min(self.prefix_len()) 117 | .min(other.prefix_len()); 118 | let repr = a & mask_from_prefix_len::(len); 119 | Ipv4Net::new(repr.into(), len).unwrap() 120 | } 121 | 122 | fn contains(&self, other: &Self) -> bool { 123 | self.contains(other) 124 | } 125 | } 126 | 127 | #[cfg(feature = "ipnet")] 128 | impl Prefix for Ipv6Net { 129 | type R = u128; 130 | 131 | fn repr(&self) -> u128 { 132 | self.addr().into() 133 | } 134 | 135 | fn prefix_len(&self) -> u8 { 136 | self.prefix_len() 137 | } 138 | 139 | fn from_repr_len(repr: u128, len: u8) -> Self { 140 | Ipv6Net::new(repr.into(), len).unwrap() 141 | } 142 | 143 | fn mask(&self) -> u128 { 144 | self.network().into() 145 | } 146 | 147 | fn zero() -> Self { 148 | Default::default() 149 | } 150 | 151 | fn longest_common_prefix(&self, other: &Self) -> Self { 152 | let a = self.repr(); 153 | let b = other.repr(); 154 | let len = ((a ^ b).leading_zeros() as u8) 155 | .min(self.prefix_len()) 156 | .min(other.prefix_len()); 157 | let repr = a & mask_from_prefix_len::(len); 158 | Ipv6Net::new(repr.into(), len).unwrap() 159 | } 160 | 161 | fn contains(&self, other: &Self) -> bool { 162 | self.contains(other) 163 | } 164 | } 165 | 166 | #[cfg(feature = "ipnetwork")] 167 | impl Prefix for Ipv4Network { 168 | type R = u32; 169 | 170 | fn repr(&self) -> u32 { 171 | self.ip().into() 172 | } 173 | 174 | fn prefix_len(&self) -> u8 { 175 | self.prefix() 176 | } 177 | 178 | fn from_repr_len(repr: u32, len: u8) -> Self { 179 | Ipv4Network::new(repr.into(), len).unwrap() 180 | } 181 | 182 | fn mask(&self) -> u32 { 183 | self.network().into() 184 | } 185 | } 186 | 187 | #[cfg(feature = "ipnetwork")] 188 | impl Prefix for Ipv6Network { 189 | type R = u128; 190 | 191 | fn repr(&self) -> u128 { 192 | self.ip().into() 193 | } 194 | 195 | fn prefix_len(&self) -> u8 { 196 | self.prefix() 197 | } 198 | 199 | fn from_repr_len(repr: u128, len: u8) -> Self { 200 | Ipv6Network::new(repr.into(), len).unwrap() 201 | } 202 | 203 | fn mask(&self) -> u128 { 204 | self.network().into() 205 | } 206 | } 207 | 208 | #[cfg(feature = "cidr")] 209 | impl Prefix for Ipv4Cidr { 210 | type R = u32; 211 | 212 | fn repr(&self) -> Self::R { 213 | self.first_address().into() 214 | } 215 | 216 | fn prefix_len(&self) -> u8 { 217 | self.network_length() 218 | } 219 | 220 | fn from_repr_len(repr: Self::R, len: u8) -> Self { 221 | let repr = repr & mask_from_prefix_len::(len); 222 | Self::new(repr.into(), len).unwrap() 223 | } 224 | 225 | fn mask(&self) -> Self::R { 226 | self.first_address().into() 227 | } 228 | } 229 | 230 | #[cfg(feature = "cidr")] 231 | impl Prefix for Ipv6Cidr { 232 | type R = u128; 233 | 234 | fn repr(&self) -> Self::R { 235 | self.first_address().into() 236 | } 237 | 238 | fn prefix_len(&self) -> u8 { 239 | self.network_length() 240 | } 241 | 242 | fn from_repr_len(repr: Self::R, len: u8) -> Self { 243 | let repr = repr & mask_from_prefix_len::(len); 244 | Self::new(repr.into(), len).unwrap() 245 | } 246 | 247 | fn mask(&self) -> Self::R { 248 | self.first_address().into() 249 | } 250 | } 251 | 252 | #[cfg(feature = "cidr")] 253 | impl Prefix for Ipv4Inet { 254 | type R = u32; 255 | 256 | fn repr(&self) -> Self::R { 257 | self.address().into() 258 | } 259 | 260 | fn prefix_len(&self) -> u8 { 261 | self.network_length() 262 | } 263 | 264 | fn from_repr_len(repr: Self::R, len: u8) -> Self { 265 | Self::new(repr.into(), len).unwrap() 266 | } 267 | 268 | fn mask(&self) -> Self::R { 269 | self.network().first_address().into() 270 | } 271 | } 272 | 273 | #[cfg(feature = "cidr")] 274 | impl Prefix for Ipv6Inet { 275 | type R = u128; 276 | 277 | fn repr(&self) -> Self::R { 278 | self.address().into() 279 | } 280 | 281 | fn prefix_len(&self) -> u8 { 282 | self.network_length() 283 | } 284 | 285 | fn from_repr_len(repr: Self::R, len: u8) -> Self { 286 | Self::new(repr.into(), len).unwrap() 287 | } 288 | 289 | fn mask(&self) -> Self::R { 290 | self.network().first_address().into() 291 | } 292 | } 293 | 294 | impl Prefix for (R, u8) 295 | where 296 | R: Unsigned + PrimInt + Zero + CheckedShr + std::fmt::Debug, 297 | { 298 | type R = R; 299 | 300 | fn repr(&self) -> R { 301 | self.0 302 | } 303 | 304 | fn prefix_len(&self) -> u8 { 305 | self.1 306 | } 307 | 308 | fn from_repr_len(repr: R, len: u8) -> Self { 309 | (repr, len) 310 | } 311 | } 312 | 313 | #[cfg(test)] 314 | #[cfg(feature = "ipnet")] 315 | mod test { 316 | use super::*; 317 | 318 | macro_rules! pfx { 319 | ($p:literal) => { 320 | $p.parse::().unwrap() 321 | }; 322 | } 323 | 324 | #[test] 325 | fn mask_from_len() { 326 | assert_eq!(mask_from_prefix_len::(3), 0b11100000); 327 | assert_eq!(mask_from_prefix_len::(5), 0b11111000); 328 | assert_eq!(mask_from_prefix_len::(8), 0b11111111); 329 | assert_eq!(mask_from_prefix_len::(0), 0b00000000); 330 | 331 | assert_eq!(mask_from_prefix_len::(0), 0x00000000); 332 | assert_eq!(mask_from_prefix_len::(8), 0xff000000); 333 | assert_eq!(mask_from_prefix_len::(16), 0xffff0000); 334 | assert_eq!(mask_from_prefix_len::(24), 0xffffff00); 335 | assert_eq!(mask_from_prefix_len::(32), 0xffffffff); 336 | } 337 | 338 | #[test] 339 | fn prefix_mask() { 340 | let addr = pfx!("10.1.0.0/8"); 341 | assert_eq!(Prefix::prefix_len(&addr), 8); 342 | assert_eq!(Prefix::repr(&addr), (10 << 24) + (1 << 16)); 343 | assert_eq!(Prefix::mask(&addr), 10u32 << 24); 344 | } 345 | 346 | #[test] 347 | fn contains() { 348 | let larger = pfx!("10.128.0.0/9"); 349 | let smaller = pfx!("10.0.0.0/8"); 350 | let larger_c = pfx!("10.130.2.5/9"); 351 | let smaller_c = pfx!("10.25.2.8/8"); 352 | assert!(smaller.contains(&larger)); 353 | assert!(smaller.contains(&larger_c)); 354 | assert!(smaller_c.contains(&larger)); 355 | assert!(smaller_c.contains(&larger_c)); 356 | assert!(!larger.contains(&smaller)); 357 | assert!(!larger.contains(&smaller_c)); 358 | assert!(!larger_c.contains(&smaller)); 359 | assert!(!larger_c.contains(&smaller_c)); 360 | assert!(smaller.contains(&smaller)); 361 | assert!(smaller.contains(&smaller_c)); 362 | assert!(smaller_c.contains(&smaller)); 363 | assert!(smaller_c.contains(&smaller_c)); 364 | } 365 | 366 | #[test] 367 | fn longest_common_prefix() { 368 | macro_rules! assert_lcp { 369 | ($a:literal, $b:literal, $c:literal) => { 370 | assert_eq!(pfx!($a).longest_common_prefix(&pfx!($b)), pfx!($c)); 371 | assert_eq!(pfx!($b).longest_common_prefix(&pfx!($a)), pfx!($c)); 372 | }; 373 | } 374 | assert_lcp!("1.2.3.4/24", "1.3.3.4/24", "1.2.0.0/15"); 375 | assert_lcp!("1.2.3.4/24", "1.1.3.4/24", "1.0.0.0/14"); 376 | assert_lcp!("1.2.3.4/24", "1.2.3.4/30", "1.2.3.0/24"); 377 | } 378 | 379 | #[test] 380 | fn is_bit_set() { 381 | assert!(pfx!("255.0.0.0/8").is_bit_set(0)); 382 | assert!(pfx!("255.0.0.0/8").is_bit_set(7)); 383 | assert!(!pfx!("255.0.0.0/8").is_bit_set(8)); 384 | assert!(!pfx!("255.255.0.0/8").is_bit_set(8)); 385 | } 386 | 387 | #[generic_tests::define] 388 | mod t { 389 | use num_traits::NumCast; 390 | 391 | use super::*; 392 | 393 | fn new(repr: u32, len: u8) -> P { 394 | let repr = <

::R as NumCast>::from(repr).unwrap(); 395 | let num_zeros = <

::R as Zero>::zero().count_zeros() as u8; 396 | let len = len + (num_zeros - 32); 397 | P::from_repr_len(repr, len) 398 | } 399 | 400 | #[test] 401 | fn repr_len() { 402 | for x in [0x01000000u32, 0x010f0000u32, 0xffff0000u32] { 403 | let repr = <

::R as NumCast>::from(x).unwrap(); 404 | let num_zeros = <

::R as Zero>::zero().count_zeros() as u8; 405 | let len = 16 + (num_zeros - 32); 406 | let prefix = P::from_repr_len(repr, len); 407 | assert!(prefix.repr() == repr); 408 | assert!(prefix.prefix_len() == len); 409 | } 410 | } 411 | 412 | #[test] 413 | fn keep_host_addr() { 414 | #[allow(unused_mut)] 415 | #[allow(unused_assignments)] 416 | let mut prefix_is_masked = false; 417 | #[cfg(feature = "cidr")] 418 | { 419 | let p_id = std::any::TypeId::of::

(); 420 | // Ipv4Cidr and Ipv6Cidr addresses are always masked. 421 | prefix_is_masked = p_id == std::any::TypeId::of::() 422 | || p_id == std::any::TypeId::of::(); 423 | } 424 | let mask = 0xffff0000u32; 425 | for mut x in [0x01001234u32, 0x010fabcdu32, 0xffff5678u32] { 426 | let prefix: P = new(x, 16); 427 | if prefix_is_masked { 428 | x &= mask; 429 | } 430 | assert_eq!(::from(prefix.repr()), Some(x)); 431 | } 432 | } 433 | 434 | #[test] 435 | fn mask() { 436 | let mask = 0xffff0000u32; 437 | for x in [0x01001234u32, 0x010fabcdu32, 0xffff5678u32] { 438 | let prefix: P = new(x, 16); 439 | assert_eq!(::from(prefix.mask()), Some(x & mask)); 440 | } 441 | } 442 | 443 | #[test] 444 | fn zero() { 445 | let prefix = P::from_repr_len(P::R::zero(), 0); 446 | assert!(P::zero().eq(&prefix)); 447 | } 448 | 449 | #[test] 450 | fn longest_common_prefix() { 451 | for ((a, al), (b, bl), (c, cl)) in [ 452 | ((0x01020304, 24), (0x01030304, 24), (0x01020000, 15)), 453 | ((0x12345678, 24), (0x12345678, 16), (0x12340000, 16)), 454 | ] { 455 | let a: P = new(a, al); 456 | let b: P = new(b, bl); 457 | let c: P = new(c, cl); 458 | let lcp = a.longest_common_prefix(&b); 459 | assert!(lcp.repr() == c.repr()); 460 | assert!(lcp.prefix_len() == c.prefix_len()); 461 | } 462 | } 463 | 464 | #[test] 465 | fn contains() { 466 | assert!(new::

(0x01020000, 16).contains(&new(0x0102ffff, 24))); 467 | assert!(new::

(0x01020304, 16).contains(&new(0x0102ffff, 24))); 468 | assert!(new::

(0x01020304, 16).contains(&new(0x0102ffff, 16))); 469 | assert!(!new::

(0x01020304, 24).contains(&new(0x0102ffff, 16))); 470 | } 471 | 472 | #[test] 473 | fn is_bit_set() { 474 | let x = 0x12345678u32; 475 | let num_zeros = <

::R as Zero>::zero().count_zeros() as u8; 476 | let offset = num_zeros - 32; 477 | let p: P = new(x, 16); 478 | for i in 0..64 { 479 | let j = i + offset; 480 | if i >= 16 { 481 | assert!(!p.is_bit_set(j)) 482 | } else { 483 | let mask = 0x80000000u32 >> i; 484 | assert_eq!(p.is_bit_set(j), x & mask != 0) 485 | } 486 | } 487 | } 488 | 489 | #[instantiate_tests()] 490 | mod ipv4net {} 491 | 492 | #[instantiate_tests()] 493 | mod ipv6net {} 494 | 495 | #[cfg(feature = "ipnetwork")] 496 | #[instantiate_tests()] 497 | mod ipv4network {} 498 | 499 | #[cfg(feature = "ipnetwork")] 500 | #[instantiate_tests()] 501 | mod ipv6network {} 502 | 503 | #[cfg(feature = "cidr")] 504 | #[instantiate_tests()] 505 | mod ipv4cidr {} 506 | 507 | #[cfg(feature = "cidr")] 508 | #[instantiate_tests()] 509 | mod ipv4inet {} 510 | 511 | #[cfg(feature = "cidr")] 512 | #[instantiate_tests()] 513 | mod ipv6cidr {} 514 | 515 | #[cfg(feature = "cidr")] 516 | #[instantiate_tests()] 517 | mod ipv6inet {} 518 | 519 | #[instantiate_tests(<(u32, u8)>)] 520 | mod u32_u8 {} 521 | 522 | #[instantiate_tests(<(u64, u8)>)] 523 | mod u64_u8 {} 524 | } 525 | } 526 | -------------------------------------------------------------------------------- /src/serde.rs: -------------------------------------------------------------------------------- 1 | //! Serialization and Deserialization implementation 2 | 3 | use std::collections::{HashMap, HashSet}; 4 | 5 | use ::serde::{Deserialize, Deserializer, Serialize, Serializer}; 6 | 7 | use super::*; 8 | 9 | impl Serialize for PrefixMap { 10 | fn serialize(&self, serializer: S) -> Result 11 | where 12 | S: Serializer, 13 | { 14 | let map: HashMap<&P, &T> = HashMap::from_iter(self); 15 | map.serialize(serializer) 16 | } 17 | } 18 | 19 | impl Serialize for PrefixSet

{ 20 | fn serialize(&self, serializer: S) -> Result 21 | where 22 | S: Serializer, 23 | { 24 | let set: HashSet<&P> = HashSet::from_iter(self); 25 | set.serialize(serializer) 26 | } 27 | } 28 | 29 | impl<'de, P: Prefix + Deserialize<'de> + Eq + std::hash::Hash, T: Deserialize<'de>> Deserialize<'de> 30 | for PrefixMap 31 | { 32 | fn deserialize(deserializer: D) -> Result 33 | where 34 | D: Deserializer<'de>, 35 | { 36 | let map: HashMap = HashMap::deserialize(deserializer)?; 37 | Ok(Self::from_iter(map)) 38 | } 39 | } 40 | 41 | impl<'de, P: Prefix + Deserialize<'de> + Eq + std::hash::Hash> Deserialize<'de> for PrefixSet

{ 42 | fn deserialize(deserializer: D) -> Result 43 | where 44 | D: Deserializer<'de>, 45 | { 46 | let set: HashSet

= HashSet::deserialize(deserializer)?; 47 | Ok(Self::from_iter(set)) 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /src/set.rs: -------------------------------------------------------------------------------- 1 | //! PrefixSet, that is implemened as a simple binary tree, based on the [`PrefixMap`]. 2 | 3 | use crate::{map::CoverKeys, Prefix, PrefixMap}; 4 | 5 | /// Set of prefixes, organized in a tree. This strucutre gives efficient access to the longest 6 | /// prefix in the set that contains another prefix. 7 | /// 8 | /// You can perform union, intersection, and (covering) difference operations by first creating a 9 | /// view over the map using [`crate::AsView`] or [`crate::AsViewMut`]. 10 | #[derive(Clone)] 11 | pub struct PrefixSet

(pub(crate) PrefixMap); 12 | 13 | impl PrefixSet

{ 14 | /// Create a new, empty prefixset. 15 | pub fn new() -> Self { 16 | Self(Default::default()) 17 | } 18 | 19 | /// Returns the number of elements stored in `self`. 20 | #[inline(always)] 21 | pub fn len(&self) -> usize { 22 | self.0.len() 23 | } 24 | 25 | /// Returns `true` if the set contains no elements. 26 | #[inline(always)] 27 | pub fn is_empty(&self) -> bool { 28 | self.0.is_empty() 29 | } 30 | 31 | /// Check wether some prefix is present in the set, without using longest prefix match. 32 | /// 33 | /// ``` 34 | /// # use prefix_trie::*; 35 | /// # #[cfg(feature = "ipnet")] 36 | /// # fn main() -> Result<(), Box> { 37 | /// let mut set: PrefixSet = PrefixSet::new(); 38 | /// set.insert("192.168.1.0/24".parse()?); 39 | /// assert!(set.contains(&"192.168.1.0/24".parse()?)); 40 | /// assert!(!set.contains(&"192.168.2.0/24".parse()?)); 41 | /// assert!(!set.contains(&"192.168.0.0/23".parse()?)); 42 | /// assert!(!set.contains(&"192.168.1.128/25".parse()?)); 43 | /// # Ok(()) 44 | /// # } 45 | /// # #[cfg(not(feature = "ipnet"))] 46 | /// # fn main() {} 47 | /// ``` 48 | pub fn contains(&self, prefix: &P) -> bool { 49 | self.0.contains_key(prefix) 50 | } 51 | 52 | /// Get a reference to the stored prefix. This function allows you to retrieve the host part of 53 | /// the prefix. The returned prefix will always have the same network address and prefix length. 54 | /// 55 | /// ``` 56 | /// # use prefix_trie::*; 57 | /// # #[cfg(feature = "ipnet")] 58 | /// # fn main() -> Result<(), Box> { 59 | /// let mut set: PrefixSet = PrefixSet::new(); 60 | /// set.insert("192.168.0.254/24".parse()?); 61 | /// assert_eq!(set.get(&"192.168.0.0/24".parse()?), Some(&"192.168.0.254/24".parse()?)); 62 | /// # Ok(()) 63 | /// # } 64 | /// # #[cfg(not(feature = "ipnet"))] 65 | /// # fn main() {} 66 | /// ``` 67 | pub fn get<'a>(&'a self, prefix: &P) -> Option<&'a P> { 68 | self.0.get_key_value(prefix).map(|(p, _)| p) 69 | } 70 | 71 | /// Get the longest prefix in the set that contains the given preifx. 72 | /// 73 | /// ``` 74 | /// # use prefix_trie::*; 75 | /// # #[cfg(feature = "ipnet")] 76 | /// # fn main() -> Result<(), Box> { 77 | /// let mut set: PrefixSet = PrefixSet::new(); 78 | /// set.insert("192.168.1.0/24".parse()?); 79 | /// set.insert("192.168.0.0/23".parse()?); 80 | /// assert_eq!(set.get_lpm(&"192.168.1.1/32".parse()?), Some(&"192.168.1.0/24".parse()?)); 81 | /// assert_eq!(set.get_lpm(&"192.168.1.0/24".parse()?), Some(&"192.168.1.0/24".parse()?)); 82 | /// assert_eq!(set.get_lpm(&"192.168.0.0/24".parse()?), Some(&"192.168.0.0/23".parse()?)); 83 | /// assert_eq!(set.get_lpm(&"192.168.2.0/24".parse()?), None); 84 | /// # Ok(()) 85 | /// # } 86 | /// # #[cfg(not(feature = "ipnet"))] 87 | /// # fn main() {} 88 | /// ``` 89 | pub fn get_lpm<'a>(&'a self, prefix: &P) -> Option<&'a P> { 90 | self.0.get_lpm(prefix).map(|(p, _)| p) 91 | } 92 | 93 | /// Get the shortest prefix in the set that contains the given preifx. 94 | /// 95 | /// ``` 96 | /// # use prefix_trie::*; 97 | /// # #[cfg(feature = "ipnet")] 98 | /// # fn main() -> Result<(), Box> { 99 | /// let mut set: PrefixSet = PrefixSet::new(); 100 | /// set.insert("192.168.1.0/24".parse()?); 101 | /// set.insert("192.168.0.0/23".parse()?); 102 | /// assert_eq!(set.get_spm(&"192.168.1.1/32".parse()?), Some(&"192.168.0.0/23".parse()?)); 103 | /// assert_eq!(set.get_spm(&"192.168.1.0/24".parse()?), Some(&"192.168.0.0/23".parse()?)); 104 | /// assert_eq!(set.get_spm(&"192.168.0.0/23".parse()?), Some(&"192.168.0.0/23".parse()?)); 105 | /// assert_eq!(set.get_spm(&"192.168.2.0/24".parse()?), None); 106 | /// # Ok(()) 107 | /// # } 108 | /// # #[cfg(not(feature = "ipnet"))] 109 | /// # fn main() {} 110 | /// ``` 111 | pub fn get_spm<'a>(&'a self, prefix: &P) -> Option<&'a P> { 112 | self.0.get_spm_prefix(prefix) 113 | } 114 | 115 | /// Adds a value to the set. 116 | /// 117 | /// Returns whether the value was newly inserted. That is: 118 | /// - If the set did not previously contain this value, `true` is returned. 119 | /// - If the set already contained this value, `false` is returned. 120 | /// 121 | /// This operation will always replace the currently stored prefix. This allows you to store 122 | /// additional information in the host aprt of the prefix. 123 | /// 124 | /// ``` 125 | /// # use prefix_trie::*; 126 | /// # #[cfg(feature = "ipnet")] 127 | /// # fn main() -> Result<(), Box> { 128 | /// let mut set: PrefixSet = PrefixSet::new(); 129 | /// assert!(set.insert("192.168.0.0/23".parse()?)); 130 | /// assert!(set.insert("192.168.1.0/24".parse()?)); 131 | /// assert!(!set.insert("192.168.1.0/24".parse()?)); 132 | /// # Ok(()) 133 | /// # } 134 | /// # #[cfg(not(feature = "ipnet"))] 135 | /// # fn main() {} 136 | /// ``` 137 | pub fn insert(&mut self, prefix: P) -> bool { 138 | self.0.insert(prefix, ()).is_none() 139 | } 140 | 141 | /// Removes a value from the set. Returns whether the value was present in the set. 142 | /// 143 | /// ``` 144 | /// # use prefix_trie::*; 145 | /// # #[cfg(feature = "ipnet")] 146 | /// # fn main() -> Result<(), Box> { 147 | /// let mut set: PrefixSet = PrefixSet::new(); 148 | /// let prefix = "192.168.1.0/24".parse()?; 149 | /// set.insert(prefix); 150 | /// assert!(set.contains(&prefix)); 151 | /// assert!(set.remove(&prefix)); 152 | /// assert!(!set.contains(&prefix)); 153 | /// # Ok(()) 154 | /// # } 155 | /// # #[cfg(not(feature = "ipnet"))] 156 | /// # fn main() {} 157 | /// ``` 158 | pub fn remove(&mut self, prefix: &P) -> bool { 159 | self.0.remove(prefix).is_some() 160 | } 161 | 162 | /// Removes a prefix from the set, returning wether the prefix was present or not. In contrast 163 | /// to [`Self::remove`], his operation will keep the tree structure as is, but only remove the 164 | /// element from it. This allows any future `insert` on the same prefix to be faster. However 165 | /// future reads from the tree might be a bit slower because they need to traverse more nodes. 166 | /// 167 | /// ``` 168 | /// # use prefix_trie::*; 169 | /// # #[cfg(feature = "ipnet")] 170 | /// # fn main() -> Result<(), Box> { 171 | /// let mut set: PrefixSet = PrefixSet::new(); 172 | /// let prefix = "192.168.1.0/24".parse()?; 173 | /// set.insert(prefix); 174 | /// assert!(set.contains(&prefix)); 175 | /// assert!(set.remove_keep_tree(&prefix)); 176 | /// assert!(!set.contains(&prefix)); 177 | /// 178 | /// // future inserts of the same key are now faster! 179 | /// set.insert(prefix); 180 | /// # Ok(()) 181 | /// # } 182 | /// # #[cfg(not(feature = "ipnet"))] 183 | /// # fn main() {} 184 | /// ``` 185 | pub fn remove_keep_tree(&mut self, prefix: &P) -> bool { 186 | self.0.remove_keep_tree(prefix).is_some() 187 | } 188 | 189 | /// Remove all elements that are contained within `prefix`. This will change the tree 190 | /// structure. This operation is `O(n)`, as the entries must be freed up one-by-one. 191 | /// 192 | /// ``` 193 | /// # use prefix_trie::*; 194 | /// # #[cfg(feature = "ipnet")] 195 | /// # fn main() -> Result<(), Box> { 196 | /// let mut set: PrefixSet = PrefixSet::new(); 197 | /// set.insert("192.168.0.0/22".parse()?); 198 | /// set.insert("192.168.0.0/23".parse()?); 199 | /// set.insert("192.168.0.0/24".parse()?); 200 | /// set.insert("192.168.2.0/23".parse()?); 201 | /// set.insert("192.168.2.0/24".parse()?); 202 | /// set.remove_children(&"192.168.0.0/23".parse()?); 203 | /// assert!(!set.contains(&"192.168.0.0/23".parse()?)); 204 | /// assert!(!set.contains(&"192.168.0.0/24".parse()?)); 205 | /// assert!(set.contains(&"192.168.2.0/23".parse()?)); 206 | /// assert!(set.contains(&"192.168.2.0/24".parse()?)); 207 | /// # Ok(()) 208 | /// # } 209 | /// # #[cfg(not(feature = "ipnet"))] 210 | /// # fn main() {} 211 | /// ``` 212 | pub fn remove_children(&mut self, prefix: &P) { 213 | self.0.remove_children(prefix) 214 | } 215 | 216 | /// Clear the set but keep the allocated memory. 217 | /// 218 | /// ``` 219 | /// # use prefix_trie::*; 220 | /// # #[cfg(feature = "ipnet")] 221 | /// # fn main() -> Result<(), Box> { 222 | /// let mut set: PrefixSet = PrefixSet::new(); 223 | /// set.insert("192.168.0.0/24".parse()?); 224 | /// set.insert("192.168.1.0/24".parse()?); 225 | /// set.clear(); 226 | /// assert!(!set.contains(&"192.168.0.0/24".parse()?)); 227 | /// assert!(!set.contains(&"192.168.1.0/24".parse()?)); 228 | /// # Ok(()) 229 | /// # } 230 | /// # #[cfg(not(feature = "ipnet"))] 231 | /// # fn main() {} 232 | /// ``` 233 | pub fn clear(&mut self) { 234 | self.0.clear() 235 | } 236 | 237 | /// Iterate over all prefixes in the set 238 | pub fn iter(&self) -> Iter<'_, P> { 239 | self.into_iter() 240 | } 241 | 242 | /// Keep only the elements in the map that satisfy the given condition `f`. 243 | /// 244 | /// ``` 245 | /// # use prefix_trie::*; 246 | /// # #[cfg(feature = "ipnet")] 247 | /// # fn main() -> Result<(), Box> { 248 | /// let mut set: PrefixSet = PrefixSet::new(); 249 | /// set.insert("192.168.0.0/24".parse()?); 250 | /// set.insert("192.168.1.0/24".parse()?); 251 | /// set.insert("192.168.2.0/24".parse()?); 252 | /// set.insert("192.168.2.0/25".parse()?); 253 | /// set.retain(|p| p.prefix_len() == 24); 254 | /// assert!(set.contains(&"192.168.0.0/24".parse()?)); 255 | /// assert!(set.contains(&"192.168.1.0/24".parse()?)); 256 | /// assert!(set.contains(&"192.168.2.0/24".parse()?)); 257 | /// assert!(!set.contains(&"192.168.2.0/25".parse()?)); 258 | /// # Ok(()) 259 | /// # } 260 | /// # #[cfg(not(feature = "ipnet"))] 261 | /// # fn main() {} 262 | /// ``` 263 | pub fn retain(&mut self, mut f: F) 264 | where 265 | F: FnMut(&P) -> bool, 266 | { 267 | let _ = self.0._retain(0, None, false, None, false, |p, _| f(p)); 268 | } 269 | 270 | /// Get an iterator over the node itself and all children. All elements returned have a prefix 271 | /// that is contained within `prefix` itself (or are the same). The iterator yields elements in 272 | /// lexicographic order. 273 | /// 274 | /// **Info**: Use the [`crate::trieview::TrieView`] abstraction that provides more flexibility. 275 | /// 276 | /// ``` 277 | /// # use prefix_trie::*; 278 | /// # #[cfg(feature = "ipnet")] 279 | /// # fn main() -> Result<(), Box> { 280 | /// let mut set: PrefixSet = PrefixSet::new(); 281 | /// set.insert("192.168.0.0/22".parse()?); 282 | /// set.insert("192.168.0.0/23".parse()?); 283 | /// set.insert("192.168.2.0/23".parse()?); 284 | /// set.insert("192.168.0.0/24".parse()?); 285 | /// set.insert("192.168.2.0/24".parse()?); 286 | /// assert_eq!( 287 | /// set.children(&"192.168.0.0/23".parse()?).collect::>(), 288 | /// vec![ 289 | /// &"192.168.0.0/23".parse()?, 290 | /// &"192.168.0.0/24".parse()?, 291 | /// ] 292 | /// ); 293 | /// # Ok(()) 294 | /// # } 295 | /// # #[cfg(not(feature = "ipnet"))] 296 | /// # fn main() {} 297 | /// ``` 298 | pub fn children<'a>(&'a self, prefix: &P) -> Iter<'a, P> { 299 | Iter(self.0.children(prefix)) 300 | } 301 | 302 | /// Iterate over all prefixes in the set that covers the given `prefix` (including `prefix` 303 | /// itself if that is present in the set). The returned iterator yields `&'a P`. 304 | /// 305 | /// The iterator will always yield elements ordered by their prefix length, i.e., their depth in 306 | /// the tree. 307 | /// 308 | /// ``` 309 | /// # use prefix_trie::*; 310 | /// # #[cfg(feature = "ipnet")] 311 | /// # fn main() -> Result<(), Box> { 312 | /// let mut set: PrefixSet = PrefixSet::new(); 313 | /// let p0 = "10.0.0.0/8".parse()?; 314 | /// let p1 = "10.1.0.0/16".parse()?; 315 | /// let p2 = "10.1.1.0/24".parse()?; 316 | /// set.insert(p0); 317 | /// set.insert(p1); 318 | /// set.insert(p2); 319 | /// set.insert("10.1.2.0/24".parse()?); // disjoint prefixes are not covered 320 | /// set.insert("10.1.1.0/25".parse()?); // more specific prefixes are not covered 321 | /// set.insert("11.0.0.0/8".parse()?); // Branch points that don't contain values are skipped 322 | /// assert_eq!(set.cover(&p2).collect::>(), vec![&p0, &p1, &p2]); 323 | /// # Ok(()) 324 | /// # } 325 | /// # #[cfg(not(feature = "ipnet"))] 326 | /// # fn main() {} 327 | /// ``` 328 | pub fn cover<'a, 'p>(&'a self, prefix: &'p P) -> CoverKeys<'a, 'p, P, ()> { 329 | self.0.cover_keys(prefix) 330 | } 331 | } 332 | 333 | impl Default for PrefixSet

{ 334 | fn default() -> Self { 335 | Self::new() 336 | } 337 | } 338 | 339 | impl

PartialEq for PrefixSet

340 | where 341 | P: Prefix + PartialEq, 342 | { 343 | fn eq(&self, other: &Self) -> bool { 344 | self.iter().zip(other.iter()).all(|(a, b)| a == b) 345 | } 346 | } 347 | 348 | impl

Eq for PrefixSet

where P: Prefix + Eq {} 349 | 350 | #[derive(Clone, Default)] 351 | /// An iterator over all entries of a [`PrefixSet`] in lexicographic order. 352 | pub struct Iter<'a, P>(crate::map::Iter<'a, P, ()>); 353 | 354 | impl<'a, P: Prefix> Iterator for Iter<'a, P> { 355 | type Item = &'a P; 356 | 357 | fn next(&mut self) -> Option { 358 | self.0.next().map(|(p, _)| p) 359 | } 360 | } 361 | 362 | #[derive(Clone)] 363 | /// A consuming iterator over all entries of a [`PrefixSet`] in lexicographic order. 364 | pub struct IntoIter

(crate::map::IntoIter); 365 | 366 | impl Iterator for IntoIter

{ 367 | type Item = P; 368 | 369 | fn next(&mut self) -> Option { 370 | self.0.next().map(|(p, _)| p) 371 | } 372 | } 373 | 374 | impl IntoIterator for PrefixSet

{ 375 | type Item = P; 376 | 377 | type IntoIter = IntoIter

; 378 | 379 | fn into_iter(self) -> Self::IntoIter { 380 | IntoIter(self.0.into_iter()) 381 | } 382 | } 383 | 384 | impl<'a, P: Prefix> IntoIterator for &'a PrefixSet

{ 385 | type Item = &'a P; 386 | 387 | type IntoIter = Iter<'a, P>; 388 | 389 | fn into_iter(self) -> Self::IntoIter { 390 | Iter(self.0.iter()) 391 | } 392 | } 393 | 394 | impl FromIterator

for PrefixSet

{ 395 | fn from_iter>(iter: I) -> Self { 396 | let mut set = Self::new(); 397 | for p in iter { 398 | set.insert(p); 399 | } 400 | set 401 | } 402 | } 403 | -------------------------------------------------------------------------------- /src/test.rs: -------------------------------------------------------------------------------- 1 | use ipnet::Ipv4Net; 2 | use num_traits::{NumCast, PrimInt}; 3 | 4 | use super::inner::Node; 5 | use super::*; 6 | 7 | type Map

= PrefixMap; 8 | 9 | fn ip(s: &str) -> P { 10 | let ip: Ipv4Net = s.parse().unwrap(); 11 | let r = ip.addr().to_bits(); 12 | let len = ip.prefix_len(); 13 | 14 | let type_len = P::zero().repr().count_zeros() as usize; 15 | assert!(type_len >= 32); 16 | 17 | let r:

::R = <

::R as NumCast>::from(r).unwrap() << (type_len - 32); 18 | P::from_repr_len(r, len) 19 | } 20 | 21 | struct TestNode { 22 | prefix: P, 23 | value: Option, 24 | left: Option>>, 25 | right: Option>>, 26 | } 27 | 28 | impl TestNode

{ 29 | pub fn create(self) -> PrefixMap { 30 | assert_eq!(self.prefix.prefix_len(), 0); 31 | let mut map = PrefixMap::new(); 32 | self.build(&mut map); 33 | map 34 | } 35 | 36 | pub fn build(mut self, map: &mut PrefixMap) -> usize { 37 | let idx = if self.prefix.prefix_len() == 0 { 38 | 0 39 | } else { 40 | map.table.as_ref().len() 41 | }; 42 | map.table.as_mut().push(Node { 43 | prefix: self.prefix, 44 | value: self.value, 45 | left: None, 46 | right: None, 47 | }); 48 | if let Some(left) = self.left.take() { 49 | let left = left.build(map); 50 | map.table[idx].left = Some(left); 51 | } 52 | if let Some(right) = self.right.take() { 53 | let right = right.build(map); 54 | map.table[idx].right = Some(right); 55 | } 56 | idx 57 | } 58 | } 59 | 60 | macro_rules! map { 61 | ($($args:tt),* $(,)?) => { 62 | _map!($($args),*).create() 63 | } 64 | } 65 | 66 | macro_rules! _map { 67 | ($ip:literal $(,)?) => { 68 | TestNode::

{ 69 | prefix: ip($ip), 70 | value: None, 71 | left: None, 72 | right: None, 73 | } 74 | }; 75 | ($ip:literal, $val:literal $(,)?) => { 76 | TestNode::

{ 77 | prefix: ip($ip), 78 | value: Some($val), 79 | left: None, 80 | right: None, 81 | } 82 | }; 83 | ($ip:literal, (), ($($args:tt),+) $(,)?) => { 84 | TestNode::

{ 85 | prefix: ip($ip), 86 | value: None, 87 | left: None 88 | right: Some(Box::new(_map!($($args),+))), 89 | } 90 | }; 91 | ($ip:literal, ($($args:tt),+), () $(,)?) => { 92 | TestNode::

{ 93 | prefix: ip($ip), 94 | value: None, 95 | left: Some(Box::new(_map!($($args),+))), 96 | right: None 97 | } 98 | }; 99 | ($ip:literal, $val:literal, (), ($($args:tt),+) $(,)?) => { 100 | TestNode::

{ 101 | prefix: ip($ip), 102 | value: Some($val), 103 | left: None, 104 | right: Some(Box::new(_map!($($args),+))), 105 | } 106 | }; 107 | ($ip:literal, $val:literal, ($($args:tt),+), () $(,)?) => { 108 | TestNode::

{ 109 | prefix: ip($ip), 110 | value: Some($val), 111 | left: Some(Box::new(_map!($($args),+))), 112 | right: None, 113 | } 114 | }; 115 | ($ip:literal, ($($left:tt),+), ($($right:tt),+) $(,)?) => { 116 | TestNode::

{ 117 | prefix: ip($ip), 118 | value: None, 119 | left: Some(Box::new(_map!($($left),+))), 120 | right: Some(Box::new(_map!($($right),+))), 121 | } 122 | }; 123 | ($ip:literal, $val:literal, ($($left:tt),+), ($($right:tt),+) $(,)?) => { 124 | TestNode::

{ 125 | prefix: ip($ip), 126 | value: Some($val), 127 | left: Some(Box::new(_map!($($left),+))), 128 | right: Some(Box::new(_map!($($right),+))), 129 | } 130 | }; 131 | } 132 | 133 | macro_rules! assert_map { 134 | ($exp:expr, ($($acq:tt),+)) => { 135 | pretty_assertions::assert_eq!(format!("{:#?}", $exp), format!("{:#?}", map!($($acq),+))) 136 | }; 137 | } 138 | 139 | macro_rules! assert_get_exact { 140 | ($map:expr, $ip:expr, $val:expr) => { 141 | let map_repr = format!("{:#?}", $map); 142 | if $val.is_some() { 143 | assert!( 144 | $map.contains_key(&ip($ip)), 145 | "Missing prefix {}\n{}", 146 | $ip, 147 | map_repr, 148 | ); 149 | } 150 | assert_eq!( 151 | $map.get(&ip($ip)).cloned(), 152 | $val, 153 | "Invalid content for prefix {}\n{}", 154 | $ip, 155 | map_repr, 156 | ); 157 | assert_eq!( 158 | $map.get_key_value(&ip($ip)).map(|(p, v)| (*p, *v)), 159 | $val.map(|v| (ip($ip), v)), 160 | "Invalid content for prefix {}\n{}", 161 | $ip, 162 | map_repr, 163 | ); 164 | assert_eq!( 165 | $map.get_mut(&ip($ip)).cloned(), 166 | $val, 167 | "Invalid content for prefix {}\n{}", 168 | $ip, 169 | map_repr, 170 | ); 171 | if let Some(val) = $val { 172 | *$map.get_mut(&ip($ip)).unwrap() += 1; 173 | assert_eq!( 174 | $map.get_mut(&ip($ip)).cloned(), 175 | Some(val + 1), 176 | "Invalid content for prefix {} after increment\n{}", 177 | $ip, 178 | map_repr, 179 | ); 180 | *$map.get_mut(&ip($ip)).unwrap() -= 1; 181 | } 182 | }; 183 | } 184 | 185 | macro_rules! assert_get_lpm { 186 | ($map:expr, $ip:expr, $lpm:expr, $val:expr) => { 187 | let map_repr = format!("{:#?}", $map); 188 | assert_eq!( 189 | $map.get_lpm_prefix(&ip($ip)).cloned(), 190 | $val.map(|_| ip($lpm)), 191 | "Invalid LPM prefix for prefix {}\n{}", 192 | $ip, 193 | map_repr, 194 | ); 195 | assert_eq!( 196 | $map.get_lpm(&ip($ip)).map(|(p, v)| (*p, *v)), 197 | $val.map(|v| (ip($lpm), v)), 198 | "Invalid LPM match for prefix {}\n{}", 199 | $ip, 200 | map_repr, 201 | ); 202 | }; 203 | } 204 | 205 | macro_rules! assert_get { 206 | ($map:expr, $ip:literal, $val:literal) => { 207 | assert_get_exact!($map, $ip, Some($val)); 208 | assert_get_lpm!($map, $ip, $ip, Some($val)); 209 | }; 210 | } 211 | 212 | macro_rules! assert_remove { 213 | ($map:expr, $ip:literal, $val:literal) => { 214 | assert_get!($map, $ip, $val); 215 | assert_eq!($map.remove_keep_tree(&ip($ip)), Some($val)); 216 | assert_get_exact!($map, $ip, None::); 217 | }; 218 | } 219 | 220 | macro_rules! assert_iter { 221 | ($map:expr) => { 222 | assert_iter!($map,) 223 | }; 224 | ($map:expr, $(($ip:literal, $val:literal)),*) => { 225 | let exp: Vec<(_, u32)> = vec![$((ip($ip), $val)),*]; 226 | assert_iter!($map, exp); 227 | }; 228 | ($map:expr, $exp:expr) => { 229 | let own_i: Vec<(_, u32)> = $exp.clone(); 230 | let ref_i: Vec<(&_, &u32)> = own_i.iter().map(|(p, v)| (p, v)).collect(); 231 | let own_p: Vec<_> = $exp.iter().map(|(p, _)| *p).collect(); 232 | let ref_p: Vec<&P> = own_p.iter().collect(); 233 | let own_v: Vec = $exp.iter().map(|(_, v)| *v).collect(); 234 | let ref_v: Vec<&u32> = own_v.iter().collect(); 235 | 236 | let mut duplicate_a = $map.clone(); 237 | let mut duplicate_b = $map.clone(); 238 | let double: Vec<(_, u32)> = $exp.iter().map(|(p, v)| (*p, v * 2)).collect(); 239 | 240 | pretty_assertions::assert_eq!($map.iter().collect::>(), ref_i); 241 | pretty_assertions::assert_eq!($map.keys().collect::>(), ref_p); 242 | pretty_assertions::assert_eq!($map.values().collect::>(), ref_v); 243 | pretty_assertions::assert_eq!($map.clone().into_iter().collect::>(), own_i); 244 | pretty_assertions::assert_eq!($map.clone().into_keys().collect::>(), own_p); 245 | pretty_assertions::assert_eq!($map.clone().into_values().collect::>(), own_v); 246 | 247 | duplicate_a.iter_mut().for_each(|(_, v)| *v *= 2); 248 | duplicate_b.values_mut().for_each(|v| *v *= 2); 249 | pretty_assertions::assert_eq!(duplicate_a.into_iter().collect::>(), double); 250 | pretty_assertions::assert_eq!(duplicate_b.into_iter().collect::>(), double); 251 | }; 252 | } 253 | 254 | #[generic_tests::define] 255 | mod t { 256 | use super::*; 257 | 258 | #[test] 259 | fn child() { 260 | let mut pm = Map::

::new(); 261 | pm.insert(ip("1.0.0.0/8"), 1); 262 | assert_map!(pm, ("0.0.0.0/0", ("1.0.0.0/8", 1), ())); 263 | assert_iter!(pm, ("1.0.0.0/8", 1)); 264 | assert_get!(pm, "1.0.0.0/8", 1); 265 | assert_eq!(pm.len(), 1); 266 | } 267 | 268 | #[test] 269 | fn chain() { 270 | let mut pm = Map::new(); 271 | pm.insert(ip("1.0.0.0/8"), 1); 272 | pm.insert(ip("1.2.0.0/16"), 2); 273 | pm.insert(ip("1.2.3.0/24"), 3); 274 | assert_eq!(pm.len(), 3); 275 | 276 | assert_map!( 277 | pm, 278 | ( 279 | "0.0.0.0/0", 280 | ("1.0.0.0/8", 1, ("1.2.0.0/16", 2, ("1.2.3.0/24", 3), ()), ()), 281 | () 282 | ) 283 | ); 284 | 285 | assert_iter!(pm, ("1.0.0.0/8", 1), ("1.2.0.0/16", 2), ("1.2.3.0/24", 3)); 286 | 287 | assert_remove!(pm, "1.0.0.0/8", 1); 288 | assert_remove!(pm, "1.2.0.0/16", 2); 289 | assert_remove!(pm, "1.2.3.0/24", 3); 290 | 291 | assert_eq!(pm.len(), 0); 292 | 293 | assert_map!( 294 | pm, 295 | ( 296 | "0.0.0.0/0", 297 | ("1.0.0.0/8", ("1.2.0.0/16", ("1.2.3.0/24"), ()), ()), 298 | () 299 | ) 300 | ); 301 | 302 | assert_iter!(pm); 303 | } 304 | 305 | #[test] 306 | fn chain_reverse() { 307 | let mut pm = Map::new(); 308 | pm.insert(ip("1.2.3.0/24"), 3); 309 | pm.insert(ip("1.2.0.0/16"), 2); 310 | pm.insert(ip("1.0.0.0/8"), 1); 311 | 312 | assert_eq!(pm.len(), 3); 313 | 314 | assert_map!( 315 | pm, 316 | ( 317 | "0.0.0.0/0", 318 | ("1.0.0.0/8", 1, ("1.2.0.0/16", 2, ("1.2.3.0/24", 3), ()), ()), 319 | () 320 | ) 321 | ); 322 | 323 | assert_iter!(pm, ("1.0.0.0/8", 1), ("1.2.0.0/16", 2), ("1.2.3.0/24", 3)); 324 | 325 | assert_remove!(pm, "1.0.0.0/8", 1); 326 | assert_remove!(pm, "1.2.0.0/16", 2); 327 | assert_remove!(pm, "1.2.3.0/24", 3); 328 | 329 | assert_eq!(pm.len(), 0); 330 | 331 | assert_map!( 332 | pm, 333 | ( 334 | "0.0.0.0/0", 335 | ("1.0.0.0/8", ("1.2.0.0/16", ("1.2.3.0/24"), ()), ()), 336 | () 337 | ) 338 | ); 339 | assert_iter!(pm); 340 | } 341 | 342 | #[test] 343 | fn branch_direct() { 344 | let mut pm = Map::new(); 345 | pm.insert(ip("0.0.0.0/7"), 1); 346 | pm.insert(ip("0.0.0.0/8"), 2); 347 | pm.insert(ip("1.0.0.0/8"), 3); 348 | 349 | assert_eq!(pm.len(), 3); 350 | 351 | assert_map!( 352 | pm, 353 | ( 354 | "0.0.0.0/0", 355 | ("0.0.0.0/7", 1, ("0.0.0.0/8", 2), ("1.0.0.0/8", 3)), 356 | () 357 | ) 358 | ); 359 | 360 | assert_iter!(pm, ("0.0.0.0/7", 1), ("0.0.0.0/8", 2), ("1.0.0.0/8", 3)); 361 | 362 | assert_remove!(pm, "0.0.0.0/7", 1); 363 | assert_remove!(pm, "0.0.0.0/8", 2); 364 | assert_remove!(pm, "1.0.0.0/8", 3); 365 | 366 | assert_eq!(pm.len(), 0); 367 | 368 | assert_map!( 369 | pm, 370 | ("0.0.0.0/0", ("0.0.0.0/7", ("0.0.0.0/8"), ("1.0.0.0/8")), ()) 371 | ); 372 | assert_iter!(pm); 373 | } 374 | 375 | #[test] 376 | fn branch_indirect() { 377 | let mut pm = Map::new(); 378 | pm.insert(ip("0.0.0.0/8"), 1); 379 | pm.insert(ip("1.0.0.0/8"), 2); 380 | 381 | assert_eq!(pm.len(), 2); 382 | 383 | assert_map!( 384 | pm, 385 | ( 386 | "0.0.0.0/0", 387 | ("0.0.0.0/7", ("0.0.0.0/8", 1), ("1.0.0.0/8", 2)), 388 | () 389 | ) 390 | ); 391 | 392 | assert_iter!(pm, ("0.0.0.0/8", 1), ("1.0.0.0/8", 2)); 393 | 394 | assert_remove!(pm, "0.0.0.0/8", 1); 395 | assert_remove!(pm, "1.0.0.0/8", 2); 396 | assert_map!( 397 | pm, 398 | ("0.0.0.0/0", ("0.0.0.0/7", ("0.0.0.0/8"), ("1.0.0.0/8")), ()) 399 | ); 400 | assert_eq!(pm.len(), 0); 401 | assert_iter!(pm); 402 | } 403 | 404 | #[test] 405 | fn branch_indirect_child() { 406 | let mut pm = Map::new(); 407 | pm.insert(ip("0.0.0.0/8"), 1); 408 | pm.insert(ip("4.0.0.0/8"), 2); 409 | assert_eq!(pm.len(), 2); 410 | assert_map!( 411 | pm, 412 | ( 413 | "0.0.0.0/0", 414 | ("0.0.0.0/5", ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)), 415 | () 416 | ) 417 | ); 418 | 419 | assert_iter!(pm, ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)); 420 | 421 | assert_remove!(pm, "0.0.0.0/8", 1); 422 | assert_remove!(pm, "4.0.0.0/8", 2); 423 | assert_eq!(pm.len(), 0); 424 | assert_map!( 425 | pm, 426 | ("0.0.0.0/0", ("0.0.0.0/5", ("0.0.0.0/8"), ("4.0.0.0/8")), ()) 427 | ); 428 | assert_iter!(pm); 429 | } 430 | 431 | #[test] 432 | fn branch_indirect_with_value() { 433 | let mut pm = Map::new(); 434 | pm.insert(ip("0.0.0.0/8"), 1); 435 | pm.insert(ip("4.0.0.0/8"), 2); 436 | pm.insert(ip("0.0.0.0/5"), 3); 437 | assert_eq!(pm.len(), 3); 438 | assert_map!( 439 | pm, 440 | ( 441 | "0.0.0.0/0", 442 | ("0.0.0.0/5", 3, ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)), 443 | () 444 | ) 445 | ); 446 | 447 | assert_iter!(pm, ("0.0.0.0/5", 3), ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)); 448 | 449 | assert_remove!(pm, "0.0.0.0/8", 1); 450 | assert_remove!(pm, "4.0.0.0/8", 2); 451 | assert_remove!(pm, "0.0.0.0/5", 3); 452 | assert_eq!(pm.len(), 0); 453 | assert_map!( 454 | pm, 455 | ("0.0.0.0/0", ("0.0.0.0/5", ("0.0.0.0/8"), ("4.0.0.0/8")), ()) 456 | ); 457 | assert_iter!(pm); 458 | } 459 | 460 | #[test] 461 | fn branch_indirect_twice() { 462 | let mut pm = Map::new(); 463 | pm.insert(ip("0.0.0.0/8"), 1); 464 | pm.insert(ip("4.0.0.0/8"), 2); 465 | pm.insert(ip("8.0.0.0/8"), 3); 466 | assert_eq!(pm.len(), 3); 467 | assert_map!( 468 | pm, 469 | ( 470 | "0.0.0.0/0", 471 | ( 472 | "0.0.0.0/4", 473 | ("0.0.0.0/5", ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)), 474 | ("8.0.0.0/8", 3) 475 | ), 476 | () 477 | ) 478 | ); 479 | 480 | assert_iter!(pm, ("0.0.0.0/8", 1), ("4.0.0.0/8", 2), ("8.0.0.0/8", 3)); 481 | 482 | assert_remove!(pm, "0.0.0.0/8", 1); 483 | assert_remove!(pm, "4.0.0.0/8", 2); 484 | assert_remove!(pm, "8.0.0.0/8", 3); 485 | assert_eq!(pm.len(), 0); 486 | assert_map!( 487 | pm, 488 | ( 489 | "0.0.0.0/0", 490 | ( 491 | "0.0.0.0/4", 492 | ("0.0.0.0/5", ("0.0.0.0/8"), ("4.0.0.0/8")), 493 | ("8.0.0.0/8") 494 | ), 495 | () 496 | ) 497 | ); 498 | assert_iter!(pm); 499 | } 500 | 501 | #[test] 502 | fn get_exact() { 503 | let mut pm = Map::new(); 504 | pm.insert(ip("0.0.0.0/8"), 1); 505 | pm.insert(ip("4.0.0.0/8"), 2); 506 | pm.insert(ip("8.0.0.0/8"), 3); 507 | pm.insert(ip("0.0.0.0/4"), 4); 508 | assert_eq!(pm.len(), 4); 509 | assert_map!( 510 | pm, 511 | ( 512 | "0.0.0.0/0", 513 | ( 514 | "0.0.0.0/4", 515 | 4, 516 | ("0.0.0.0/5", ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)), 517 | ("8.0.0.0/8", 3) 518 | ), 519 | () 520 | ) 521 | ); 522 | eprintln!("FOO"); 523 | 524 | assert_iter!( 525 | pm, 526 | ("0.0.0.0/4", 4), 527 | ("0.0.0.0/8", 1), 528 | ("4.0.0.0/8", 2), 529 | ("8.0.0.0/8", 3) 530 | ); 531 | 532 | assert_get_exact!(pm, "0.0.0.0/0", None::); 533 | assert_get_exact!(pm, "0.0.0.0/1", None::); 534 | assert_get_exact!(pm, "0.0.0.0/2", None::); 535 | assert_get_exact!(pm, "0.0.0.0/3", None::); 536 | assert_get_exact!(pm, "0.0.0.0/4", Some(4)); 537 | assert_get_exact!(pm, "0.0.0.0/5", None::); 538 | assert_get_exact!(pm, "0.0.0.0/6", None::); 539 | assert_get_exact!(pm, "0.0.0.0/7", None::); 540 | assert_get_exact!(pm, "0.0.0.0/8", Some(1)); 541 | assert_get_exact!(pm, "0.0.0.0/5", None::); 542 | assert_get_exact!(pm, "4.0.0.0/6", None::); 543 | assert_get_exact!(pm, "4.0.0.0/7", None::); 544 | assert_get_exact!(pm, "4.0.0.0/8", Some(2)); 545 | assert_get_exact!(pm, "8.0.0.0/5", None::); 546 | assert_get_exact!(pm, "8.0.0.0/6", None::); 547 | assert_get_exact!(pm, "8.0.0.0/7", None::); 548 | assert_get_exact!(pm, "8.0.0.0/8", Some(3)); 549 | } 550 | 551 | #[test] 552 | fn get_lpm() { 553 | let mut pm = Map::

::new(); 554 | pm.insert(ip("0.0.0.0/8"), 1); 555 | pm.insert(ip("4.0.0.0/8"), 2); 556 | pm.insert(ip("8.0.0.0/8"), 3); 557 | pm.insert(ip("0.0.0.0/4"), 4); 558 | assert_eq!(pm.len(), 4); 559 | assert_map!( 560 | pm, 561 | ( 562 | "0.0.0.0/0", 563 | ( 564 | "0.0.0.0/4", 565 | 4, 566 | ("0.0.0.0/5", ("0.0.0.0/8", 1), ("4.0.0.0/8", 2)), 567 | ("8.0.0.0/8", 3) 568 | ), 569 | () 570 | ) 571 | ); 572 | assert_get_lpm!(pm, "0.0.0.0/0", "0.0.0.0/0", None::); 573 | assert_get_lpm!(pm, "0.0.0.0/1", "0.0.0.0/0", None::); 574 | assert_get_lpm!(pm, "0.0.0.0/2", "0.0.0.0/0", None::); 575 | assert_get_lpm!(pm, "0.0.0.0/3", "0.0.0.0/0", None::); 576 | assert_get_lpm!(pm, "0.0.0.0/4", "0.0.0.0/4", Some(4)); 577 | 578 | assert_get_lpm!(pm, "0.0.0.0/5", "0.0.0.0/4", Some(4)); 579 | assert_get_lpm!(pm, "0.0.0.0/6", "0.0.0.0/4", Some(4)); 580 | assert_get_lpm!(pm, "0.0.0.0/7", "0.0.0.0/4", Some(4)); 581 | assert_get_lpm!(pm, "0.0.0.0/8", "0.0.0.0/8", Some(1)); 582 | assert_get_lpm!(pm, "0.0.0.0/9", "0.0.0.0/8", Some(1)); 583 | 584 | assert_get_lpm!(pm, "0.0.0.0/5", "0.0.0.0/4", Some(4)); 585 | assert_get_lpm!(pm, "4.0.0.0/6", "0.0.0.0/4", Some(4)); 586 | assert_get_lpm!(pm, "4.0.0.0/7", "0.0.0.0/4", Some(4)); 587 | assert_get_lpm!(pm, "4.0.0.0/8", "4.0.0.0/8", Some(2)); 588 | assert_get_lpm!(pm, "4.0.0.0/9", "4.0.0.0/8", Some(2)); 589 | 590 | assert_get_lpm!(pm, "8.0.0.0/5", "0.0.0.0/4", Some(4)); 591 | assert_get_lpm!(pm, "8.0.0.0/6", "0.0.0.0/4", Some(4)); 592 | assert_get_lpm!(pm, "8.0.0.0/7", "0.0.0.0/4", Some(4)); 593 | assert_get_lpm!(pm, "8.0.0.0/8", "8.0.0.0/8", Some(3)); 594 | assert_get_lpm!(pm, "8.0.0.0/9", "8.0.0.0/8", Some(3)); 595 | 596 | assert_get_lpm!(pm, "9.0.0.0/8", "0.0.0.0/4", Some(4)); 597 | assert_get_lpm!(pm, "10.0.0.0/8", "0.0.0.0/4", Some(4)); 598 | } 599 | 600 | #[test] 601 | fn test_spm_vs_lpm_all_routes() { 602 | let prefix_set = PrefixSet::

::from_iter(vec![ 603 | ip("0.0.0.0/0"), 604 | ip("192.168.0.0/23"), 605 | ip("192.168.0.0/24"), 606 | ]); 607 | 608 | let prefix = ip("192.168.0.1/32"); 609 | 610 | assert_eq!(prefix_set.get_spm(&prefix).unwrap(), &ip("0.0.0.0/0")); 611 | assert_eq!(prefix_set.get_lpm(&prefix).unwrap(), &ip("192.168.0.0/24")); 612 | } 613 | 614 | #[test] 615 | fn test_remove_iter_mut() { 616 | let mut map: PrefixMap = PrefixMap::from_iter([ 617 | (ip("192.168.0.0/20"), 1), 618 | (ip("192.168.0.0/22"), 2), 619 | (ip("192.168.0.0/24"), 3), 620 | (ip("192.168.2.0/23"), 4), 621 | (ip("192.168.4.0/22"), 5), 622 | ]); 623 | let mut view = map.view_mut_at(ip("192.168.0.0/22")).unwrap(); 624 | 625 | view.remove(); 626 | 627 | let x = view.into_iter().collect::>(); 628 | println!("{:?}", *x[0].0); 629 | println!("{:?}", *x[0].1); 630 | println!("{:?}", *x[1].0); 631 | println!("{:?}", *x[1].1); 632 | } 633 | 634 | #[test] 635 | fn insert_with_host_part() { 636 | let mut set = PrefixSet::

::new(); 637 | set.insert(ip("192.168.0.254/24")); 638 | set.insert(ip("192.168.1.254/24")); 639 | // Now, we have a branching node at 192.168.0.0/23 with that address. 640 | set.insert(ip("192.168.0.254/23")); 641 | // This will not overwrite the address 642 | assert_eq!( 643 | Vec::from_iter(set), 644 | vec![ 645 | ip("192.168.0.254/23"), 646 | ip("192.168.0.254/24"), 647 | ip("192.168.1.254/24"), 648 | ] 649 | ); 650 | } 651 | 652 | #[test] 653 | fn replace_with_host_part() { 654 | let mut set = PrefixSet::

::new(); 655 | set.insert(ip("192.168.0.0/24")); 656 | set.insert(ip("192.168.0.1/24")); 657 | // This will not overwrite the address 658 | assert_eq!(Vec::from_iter(set), vec![ip("192.168.0.1/24"),]); 659 | } 660 | 661 | #[test] 662 | fn entry_with_host_part() { 663 | let mut set = PrefixMap::::new(); 664 | set.insert(ip("192.168.0.254/24"), 1); 665 | set.insert(ip("192.168.1.254/24"), 2); 666 | println!("{set:?}"); 667 | // Now, we have a branching node at 192.168.0.0/23 with that address. 668 | set.entry(ip("192.168.0.254/23")).or_insert(3); 669 | println!("{set:?}"); 670 | // This will not overwrite the address 671 | assert_eq!( 672 | Vec::from_iter(set), 673 | vec![ 674 | (ip("192.168.0.254/23"), 3), 675 | (ip("192.168.0.254/24"), 1), 676 | (ip("192.168.1.254/24"), 2), 677 | ] 678 | ); 679 | } 680 | 681 | #[test] 682 | fn replace_entry_with_host_part() { 683 | let mut map = PrefixMap::::new(); 684 | map.insert(ip("192.168.0.0/24"), 0); 685 | map.entry(ip("192.168.0.1/24")).insert(0); 686 | // This will not overwrite the address 687 | assert_eq!(Vec::from_iter(map), vec![(ip("192.168.0.1/24"), 0)]); 688 | } 689 | 690 | #[test] 691 | fn from_iter_with_host_part() { 692 | let set = PrefixSet::

::from_iter([ 693 | ip("192.168.0.254/24"), 694 | ip("192.168.1.0/24"), 695 | ip("192.168.0.254/23"), 696 | ip("192.168.1.254/24"), 697 | ]); 698 | assert_eq!( 699 | Vec::from_iter(set), 700 | vec![ 701 | ip("192.168.0.254/23"), 702 | ip("192.168.0.254/24"), 703 | ip("192.168.1.254/24"), 704 | ] 705 | ); 706 | } 707 | 708 | #[instantiate_tests(<(u32, u8)>)] 709 | mod raw32 {} 710 | 711 | #[instantiate_tests(<(u64, u8)>)] 712 | mod raw64 {} 713 | 714 | #[instantiate_tests(<(u128, u8)>)] 715 | mod raw128 {} 716 | 717 | #[cfg(feature = "ipnet")] 718 | #[instantiate_tests()] 719 | mod ipv4net {} 720 | 721 | #[cfg(feature = "ipnet")] 722 | #[instantiate_tests()] 723 | mod ipv6net {} 724 | 725 | #[cfg(feature = "ipnetwork")] 726 | #[instantiate_tests()] 727 | mod ipv4network {} 728 | 729 | #[cfg(feature = "ipnetwork")] 730 | #[instantiate_tests()] 731 | mod ipv6network {} 732 | 733 | #[cfg(feature = "cidr")] 734 | #[instantiate_tests()] 735 | mod ipv4cidr {} 736 | 737 | #[cfg(feature = "cidr")] 738 | #[instantiate_tests()] 739 | mod ipv6cidr {} 740 | 741 | #[cfg(feature = "cidr")] 742 | #[instantiate_tests()] 743 | mod ipv4inet {} 744 | 745 | #[cfg(feature = "cidr")] 746 | #[instantiate_tests()] 747 | mod ipv6inet {} 748 | } 749 | -------------------------------------------------------------------------------- /src/trieview/intersection.rs: -------------------------------------------------------------------------------- 1 | use crate::to_right; 2 | 3 | use super::*; 4 | 5 | /// an iterator over the intersection of two [`crate::PrefixSet`]s in lexicographic order. 6 | pub struct Intersection<'a, P, L, R> { 7 | pub(super) table_l: &'a Table, 8 | pub(super) table_r: &'a Table, 9 | pub(super) nodes: Vec, 10 | } 11 | 12 | /// an iterator over the intersection of two [`crate::PrefixSet`]s in lexicographic order, yielding 13 | /// mutable references to all elements. 14 | pub struct IntersectionMut<'a, P, L, R> { 15 | pub(super) table_l: &'a Table, 16 | pub(super) table_r: &'a Table, 17 | pub(super) nodes: Vec, 18 | } 19 | 20 | impl<'a, P, L, R> IntersectionMut<'a, P, L, R> { 21 | /// Safety: 22 | /// 1. Table_l must come from a `TrieViewMut` and satisfy the conditions in `TrieViewMut::new`. 23 | /// 2. Table_r must come from a `TrieViewMut` and satisfy the conditions in `TrieViewMut::new`. 24 | /// 3. Table_l and Table_r must be distinct. This is implicitly given by the two rules above. 25 | unsafe fn new( 26 | table_l: &'a Table, 27 | table_r: &'a Table, 28 | nodes: Vec, 29 | ) -> Self { 30 | Self { 31 | table_l, 32 | table_r, 33 | nodes, 34 | } 35 | } 36 | } 37 | 38 | pub(super) enum IntersectionIndex { 39 | Both(usize, usize), 40 | FirstA(usize, usize), 41 | FirstB(usize, usize), 42 | } 43 | 44 | impl<'a, P, L> TrieView<'a, P, L> 45 | where 46 | P: Prefix, 47 | { 48 | /// Iterate over the union of both Views. Each element will yield a reference to the prefix and 49 | /// the value stored in `self` and `other` (if the prefix is in both views). 50 | /// 51 | /// ``` 52 | /// # use prefix_trie::*; 53 | /// # #[cfg(feature = "ipnet")] 54 | /// macro_rules! net { ($x:literal) => {$x.parse::().unwrap()}; } 55 | /// 56 | /// # #[cfg(feature = "ipnet")] 57 | /// # { 58 | /// let mut map_a: PrefixMap = PrefixMap::from_iter([ 59 | /// (net!("192.168.0.0/20"), 1), 60 | /// (net!("192.168.0.0/22"), 2), 61 | /// (net!("192.168.0.0/24"), 3), 62 | /// (net!("192.168.2.0/23"), 4), 63 | /// ]); 64 | /// let mut map_b: PrefixMap = PrefixMap::from_iter([ 65 | /// (net!("192.168.0.0/20"), "a"), 66 | /// (net!("192.168.0.0/22"), "b"), 67 | /// (net!("192.168.0.0/23"), "c"), 68 | /// (net!("192.168.0.0/24"), "d"), 69 | /// (net!("192.168.2.0/24"), "e"), 70 | /// ]); 71 | /// let sub_a = map_a.view_at(net!("192.168.0.0/22")).unwrap(); 72 | /// let sub_b = map_b.view_at(net!("192.168.0.0/22")).unwrap(); 73 | /// assert_eq!( 74 | /// sub_a.intersection(sub_b).collect::>(), 75 | /// vec![ 76 | /// (&net!("192.168.0.0/22"), &2, &"b"), 77 | /// (&net!("192.168.0.0/24"), &3, &"d"), 78 | /// ] 79 | /// ); 80 | /// # } 81 | /// ``` 82 | pub fn intersection(&self, other: impl AsView<'a, P, R>) -> Intersection<'a, P, L, R> { 83 | let other = other.view(); 84 | Intersection { 85 | table_l: self.table, 86 | table_r: other.table, 87 | nodes: Vec::from_iter(next_indices( 88 | self.table, 89 | other.table, 90 | Some(self.loc.idx()), 91 | Some(other.loc.idx()), 92 | )), 93 | } 94 | } 95 | } 96 | 97 | impl TrieViewMut<'_, P, L> 98 | where 99 | P: Prefix, 100 | { 101 | /// Iterate over the union of both Views. Each element will yield a reference to the prefix and 102 | /// mutable references to the values stored in `self` and `other` (if the prefix is in both 103 | /// views). 104 | /// 105 | /// ``` 106 | /// # use prefix_trie::*; 107 | /// # use prefix_trie::trieview::UnionItem; 108 | /// # #[cfg(feature = "ipnet")] 109 | /// macro_rules! net { ($x:literal) => {$x.parse::().unwrap()}; } 110 | /// 111 | /// # #[cfg(feature = "ipnet")] 112 | /// # { 113 | /// let mut map_a: PrefixMap = PrefixMap::from_iter([ 114 | /// (net!("192.168.0.0/20"), 1), 115 | /// (net!("192.168.0.0/22"), 2), 116 | /// (net!("192.168.0.0/24"), 3), 117 | /// (net!("192.168.2.0/23"), 4), 118 | /// ]); 119 | /// let mut map_b: PrefixMap = PrefixMap::from_iter([ 120 | /// (net!("192.168.0.0/22"), 10), 121 | /// (net!("192.168.0.0/23"), 20), 122 | /// (net!("192.168.2.0/24"), 30), 123 | /// ]); 124 | /// 125 | /// // Modify the two maps by adding their values for elements of the same prefix. 126 | /// for (_, l, r) in map_a.view_mut().intersection_mut(&mut map_b) { 127 | /// *l += *r; 128 | /// *r = *l; 129 | /// } 130 | /// 131 | /// assert_eq!( 132 | /// map_a.into_iter().collect::>(), 133 | /// vec![ 134 | /// (net!("192.168.0.0/20"), 1), 135 | /// (net!("192.168.0.0/22"), 12), 136 | /// (net!("192.168.0.0/24"), 3), 137 | /// (net!("192.168.2.0/23"), 4), 138 | /// ], 139 | /// ); 140 | /// assert_eq!( 141 | /// map_b.into_iter().collect::>(), 142 | /// vec![ 143 | /// (net!("192.168.0.0/22"), 12), 144 | /// (net!("192.168.0.0/23"), 20), 145 | /// (net!("192.168.2.0/24"), 30), 146 | /// ], 147 | /// ); 148 | /// # } 149 | /// ``` 150 | pub fn intersection_mut<'b, R>( 151 | &'b mut self, 152 | other: impl AsViewMut<'b, P, R>, 153 | ) -> IntersectionMut<'b, P, L, R> { 154 | let other = other.view_mut(); 155 | let nodes = Vec::from_iter(next_indices( 156 | self.table, 157 | other.table, 158 | Some(self.loc.idx()), 159 | Some(other.loc.idx()), 160 | )); 161 | // Safety: Both `self` and `other` are `TrieViewMut`s, and must adhere to the safety 162 | // constraints in `TrieViewMut::new`. 163 | unsafe { IntersectionMut::new(self.table, other.table, nodes) } 164 | } 165 | } 166 | 167 | impl<'a, P: Prefix, L, R> Iterator for Intersection<'a, P, L, R> { 168 | type Item = (&'a P, &'a L, &'a R); 169 | 170 | fn next(&mut self) -> Option { 171 | while let Some(cur) = self.nodes.pop() { 172 | match cur { 173 | IntersectionIndex::Both(l, r) => { 174 | let node_l = &self.table_l[l]; 175 | let node_r = &self.table_r[r]; 176 | self.nodes.extend(next_indices( 177 | self.table_l, 178 | self.table_r, 179 | node_l.right, 180 | node_r.right, 181 | )); 182 | self.nodes.extend(next_indices( 183 | self.table_l, 184 | self.table_r, 185 | node_l.left, 186 | node_r.left, 187 | )); 188 | if let (Some(left), Some(right)) = 189 | (node_l.value.as_ref(), node_r.value.as_ref()) 190 | { 191 | return Some((&node_l.prefix, left, right)); 192 | } 193 | } 194 | IntersectionIndex::FirstA(l, r) => { 195 | let node_l = &self.table_l[l]; 196 | self.nodes.extend(next_indices_first_a( 197 | self.table_l, 198 | self.table_r, 199 | l, 200 | node_l.left, 201 | node_l.right, 202 | r, 203 | )); 204 | } 205 | IntersectionIndex::FirstB(l, r) => { 206 | let node_r = &self.table_r[r]; 207 | self.nodes.extend(next_indices_first_b( 208 | self.table_l, 209 | self.table_r, 210 | l, 211 | r, 212 | node_r.left, 213 | node_r.right, 214 | )); 215 | } 216 | } 217 | } 218 | None 219 | } 220 | } 221 | 222 | impl<'a, P: Prefix, L, R> Iterator for IntersectionMut<'a, P, L, R> { 223 | type Item = (&'a P, &'a mut L, &'a mut R); 224 | 225 | fn next(&mut self) -> Option { 226 | // safety: map is a tree. Every node is visited exactly once during the iteration 227 | // (self.nodes is not public). Therefore, each in each iteration of this loop (also between 228 | // multiple calls to `next`), the index `cur` is different to any of the earlier 229 | // iterations. It is therefore safe to extend the lifetime of the elements to 'a (which is 230 | // the lifetime for which `self` has an exclusive reference over the map). 231 | while let Some(cur) = self.nodes.pop() { 232 | match cur { 233 | IntersectionIndex::Both(l, r) => { 234 | let node_l = &self.table_l[l]; 235 | let node_r = &self.table_r[r]; 236 | self.nodes.extend(next_indices( 237 | self.table_l, 238 | self.table_r, 239 | node_l.right, 240 | node_r.right, 241 | )); 242 | self.nodes.extend(next_indices( 243 | self.table_l, 244 | self.table_r, 245 | node_l.left, 246 | node_r.left, 247 | )); 248 | let node_l = unsafe { self.table_l.get_mut(l) }; 249 | let node_r = unsafe { self.table_r.get_mut(r) }; 250 | if let (Some(left), Some(right)) = 251 | (node_l.value.as_mut(), node_r.value.as_mut()) 252 | { 253 | return Some((&node_l.prefix, left, right)); 254 | } 255 | } 256 | IntersectionIndex::FirstA(l, r) => { 257 | let node_l = &self.table_l[l]; 258 | self.nodes.extend(next_indices_first_a( 259 | self.table_l, 260 | self.table_r, 261 | l, 262 | node_l.left, 263 | node_l.right, 264 | r, 265 | )); 266 | } 267 | IntersectionIndex::FirstB(l, r) => { 268 | let node_r = &self.table_r[r]; 269 | self.nodes.extend(next_indices_first_b( 270 | self.table_l, 271 | self.table_r, 272 | l, 273 | r, 274 | node_r.left, 275 | node_r.right, 276 | )); 277 | } 278 | } 279 | } 280 | None 281 | } 282 | } 283 | 284 | fn next_indices<'a, P: Prefix, L, R>( 285 | table_l: &'a Table, 286 | table_r: &'a Table, 287 | node_l: Option, 288 | node_r: Option, 289 | ) -> Option { 290 | match (node_l, node_r) { 291 | (None, Some(_)) => None, 292 | (Some(_), None) => None, 293 | (Some(a), Some(b)) => { 294 | let p_a = &table_l[a].prefix; 295 | let p_b = &table_r[b].prefix; 296 | if p_a.prefix_len() == p_b.prefix_len() { 297 | match p_a.mask().cmp(&p_b.mask()) { 298 | std::cmp::Ordering::Equal => Some(IntersectionIndex::Both(a, b)), 299 | _ => None, 300 | } 301 | } else if p_a.contains(p_b) { 302 | Some(IntersectionIndex::FirstA(a, b)) 303 | } else if p_b.contains(p_a) { 304 | Some(IntersectionIndex::FirstB(a, b)) 305 | } else { 306 | None 307 | } 308 | } 309 | _ => None, 310 | } 311 | } 312 | 313 | fn next_indices_first_a<'a, P: Prefix, L, R>( 314 | table_l: &'a Table, 315 | table_r: &'a Table, 316 | l: usize, 317 | ll: Option, 318 | lr: Option, 319 | r: usize, 320 | ) -> Option { 321 | match (ll, lr) { 322 | (None, None) => None, 323 | (None, Some(lr)) => next_indices(table_l, table_r, Some(lr), Some(r)), 324 | (Some(ll), None) => next_indices(table_l, table_r, Some(ll), Some(r)), 325 | (Some(ll), Some(lr)) => { 326 | if to_right(&table_l[l].prefix, &table_r[r].prefix) { 327 | next_indices(table_l, table_r, Some(lr), Some(r)) 328 | } else { 329 | next_indices(table_l, table_r, Some(ll), Some(r)) 330 | } 331 | } 332 | } 333 | } 334 | 335 | fn next_indices_first_b<'a, P: Prefix, L, R>( 336 | table_l: &'a Table, 337 | table_r: &'a Table, 338 | l: usize, 339 | r: usize, 340 | rl: Option, 341 | rr: Option, 342 | ) -> Option { 343 | match (rl, rr) { 344 | (None, None) => None, 345 | (None, Some(rr)) => next_indices(table_l, table_r, Some(l), Some(rr)), 346 | (Some(rl), None) => next_indices(table_l, table_r, Some(l), Some(rl)), 347 | (Some(rl), Some(rr)) => { 348 | if to_right(&table_r[r].prefix, &table_l[l].prefix) { 349 | next_indices(table_l, table_r, Some(l), Some(rr)) 350 | } else { 351 | next_indices(table_l, table_r, Some(l), Some(rl)) 352 | } 353 | } 354 | } 355 | } 356 | -------------------------------------------------------------------------------- /src/trieview/union.rs: -------------------------------------------------------------------------------- 1 | use crate::to_right; 2 | 3 | use super::*; 4 | 5 | /// An iterator over the union of two TrieViews that always yields either the exact value or the 6 | /// longest prefix match of both of them. 7 | pub struct Union<'a, P, L, R> { 8 | table_l: &'a Table, 9 | table_r: &'a Table, 10 | nodes: Vec>, 11 | } 12 | 13 | /// An iterator over the union of two TrieViews that always yields either the exact value or the 14 | /// longest prefix match of both of them. 15 | pub struct UnionMut<'a, P, L, R> { 16 | // Safety: table_l must be distinct from table_r 17 | table_l: &'a Table, 18 | // Safety: table_l must be distinct from table_r 19 | table_r: &'a Table, 20 | nodes: Vec, 21 | } 22 | 23 | impl<'a, P, L, R> UnionMut<'a, P, L, R> { 24 | /// Safety: 25 | /// 1. Table_l must come from a `TrieViewMut` and satisfy the conditions in `TrieViewMut::new`. 26 | /// 2. Table_r must come from a `TrieViewMut` and satisfy the conditions in `TrieViewMut::new`. 27 | /// 3. Table_l and Table_r must be distinct. This is implicitly given by the two rules above. 28 | unsafe fn new( 29 | table_l: &'a Table, 30 | table_r: &'a Table, 31 | nodes: Vec, 32 | ) -> Self { 33 | Self { 34 | table_l, 35 | table_r, 36 | nodes, 37 | } 38 | } 39 | } 40 | 41 | type Lpm<'a, P, T> = Option<(&'a P, &'a T)>; 42 | type Node<'a, P, L, R> = (UnionIndex, Lpm<'a, P, L>, Lpm<'a, P, R>); 43 | 44 | /// An item of the [`Union`] iterator. 45 | #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] 46 | pub enum UnionItem<'a, P, L, R> { 47 | /// The prefix is only present in the left TrieView (`self`). 48 | Left { 49 | /// The prefix of the element. 50 | prefix: &'a P, 51 | /// The value of the element in the left TrieView (`self`). 52 | left: &'a L, 53 | /// The longest prefix match in the right TrieView (`other`). 54 | right: Option<(&'a P, &'a R)>, 55 | }, 56 | /// The prefix is only present in the right TrieView (`other`). 57 | Right { 58 | /// The prefix of the element. 59 | prefix: &'a P, 60 | /// The longest prefix match in the left TrieView (`self`). 61 | left: Option<(&'a P, &'a L)>, 62 | /// The value of the element in the right TrieView (`other`). 63 | right: &'a R, 64 | }, 65 | /// The prefix is only present in the right TrieView (`other`). 66 | Both { 67 | /// The prefix of the element. 68 | prefix: &'a P, 69 | /// The value of the element in the left TrieView (`self`). 70 | left: &'a L, 71 | /// The value of the element in the right TrieView (`other`). 72 | right: &'a R, 73 | }, 74 | } 75 | 76 | impl<'a, P, L, R> UnionItem<'a, P, L, R> { 77 | /// Get the prefix of the current element (in the exact match). 78 | pub fn prefix(&self) -> &'a P { 79 | match self { 80 | UnionItem::Left { prefix, .. } 81 | | UnionItem::Right { prefix, .. } 82 | | UnionItem::Both { prefix, .. } => prefix, 83 | } 84 | } 85 | 86 | /// Get the element in both the left and right map, but only if they are both present. By 87 | /// doing `a.union(b).filter_map(|x| x.both)`, you get an iterator that yields only elements 88 | /// that are present in both tries. 89 | pub fn both(&self) -> Option<(&'a P, &'a L, &'a R)> { 90 | match self { 91 | UnionItem::Left { .. } | UnionItem::Right { .. } => None, 92 | UnionItem::Both { 93 | prefix, 94 | left, 95 | right, 96 | } => Some((prefix, left, right)), 97 | } 98 | } 99 | 100 | /// Get the value of the left item (`self`). This either returns the exact match or the 101 | /// longest-prefix match. 102 | pub fn left(&self) -> Option<(&'a P, &'a L)> { 103 | match self { 104 | UnionItem::Right { left, .. } => *left, 105 | UnionItem::Left { prefix, left, .. } | UnionItem::Both { prefix, left, .. } => { 106 | Some((prefix, left)) 107 | } 108 | } 109 | } 110 | 111 | /// Get the value of the right item (`other`). This either returns the exact match or the 112 | /// longest-prefix match. 113 | pub fn right(&self) -> Option<(&'a P, &'a R)> { 114 | match self { 115 | UnionItem::Left { right, .. } => *right, 116 | UnionItem::Right { prefix, right, .. } | UnionItem::Both { prefix, right, .. } => { 117 | Some((prefix, right)) 118 | } 119 | } 120 | } 121 | } 122 | 123 | enum UnionIndex { 124 | Both(usize, usize), 125 | FirstL(usize, usize), 126 | FirstR(usize, usize), 127 | OnlyL(usize), 128 | OnlyR(usize), 129 | } 130 | 131 | impl<'a, P, L> TrieView<'a, P, L> 132 | where 133 | P: Prefix, 134 | { 135 | /// Iterate over the union of two views. If a prefix is present in both trees, the iterator 136 | /// will yield both elements. Otherwise, the iterator will yield the element of one TrieView 137 | /// together with the longest prefix match in the other TrieView. Elements are of type 138 | /// [`UnionItem`]. 139 | /// 140 | /// **Warning**: The iterator will only yield elements of the given TrieViews. If either of the 141 | /// two TrieViews is pointing to a branching or a virtual node, then the longest prefix match 142 | /// returned may be `None`, even though it exists in the larger tree. 143 | /// 144 | /// ``` 145 | /// # use prefix_trie::*; 146 | /// # use prefix_trie::trieview::UnionItem; 147 | /// # #[cfg(feature = "ipnet")] 148 | /// macro_rules! net { ($x:literal) => {$x.parse::().unwrap()}; } 149 | /// 150 | /// # #[cfg(feature = "ipnet")] 151 | /// # { 152 | /// let mut map_a: PrefixMap = PrefixMap::from_iter([ 153 | /// (net!("192.168.0.0/20"), 1), 154 | /// (net!("192.168.0.0/22"), 2), 155 | /// (net!("192.168.0.0/24"), 3), 156 | /// ]); 157 | /// let mut map_b: PrefixMap = PrefixMap::from_iter([ 158 | /// (net!("192.168.0.0/22"), "a"), 159 | /// (net!("192.168.0.0/23"), "b"), 160 | /// ]); 161 | /// assert_eq!( 162 | /// map_a.view().union(&map_b).collect::>(), 163 | /// vec![ 164 | /// UnionItem::Left{ 165 | /// prefix: &net!("192.168.0.0/20"), 166 | /// left: &1, 167 | /// right: None, 168 | /// }, 169 | /// UnionItem::Both{ 170 | /// prefix: &net!("192.168.0.0/22"), 171 | /// left: &2, 172 | /// right: &"a", 173 | /// }, 174 | /// UnionItem::Right{ 175 | /// prefix: &net!("192.168.0.0/23"), 176 | /// left: Some((&net!("192.168.0.0/22"), &2)), 177 | /// right: &"b", 178 | /// }, 179 | /// UnionItem::Left{ 180 | /// prefix: &net!("192.168.0.0/24"), 181 | /// left: &3, 182 | /// right: Some((&net!("192.168.0.0/23"), &"b")), 183 | /// }, 184 | /// ] 185 | /// ); 186 | /// # } 187 | /// ``` 188 | pub fn union(&self, other: impl AsView<'a, P, R>) -> Union<'a, P, L, R> { 189 | let other = other.view(); 190 | Union { 191 | table_l: self.table, 192 | table_r: other.table, 193 | nodes: extend_lpm( 194 | self.table, 195 | other.table, 196 | self.table[self.loc.idx()].prefix_value(), 197 | other.table[other.loc.idx()].prefix_value(), 198 | next_indices( 199 | self.table, 200 | other.table, 201 | Some(self.loc.idx()), 202 | Some(other.loc.idx()), 203 | ), 204 | ) 205 | .collect(), 206 | } 207 | } 208 | } 209 | 210 | impl TrieViewMut<'_, P, L> 211 | where 212 | P: Prefix, 213 | { 214 | /// Iterate over the union of two views. If a prefix is present in both trees, the iterator 215 | /// will yield mutable references to both elements. Longest prefix matches are not returned. 216 | /// 217 | /// ``` 218 | /// # use prefix_trie::*; 219 | /// # use prefix_trie::trieview::UnionItem; 220 | /// # #[cfg(feature = "ipnet")] 221 | /// macro_rules! net { ($x:literal) => {$x.parse::().unwrap()}; } 222 | /// 223 | /// # #[cfg(feature = "ipnet")] 224 | /// # { 225 | /// let mut map_a: PrefixMap = PrefixMap::from_iter([ 226 | /// (net!("192.168.0.0/20"), 1), 227 | /// (net!("192.168.0.0/22"), 2), 228 | /// (net!("192.168.0.0/24"), 3), 229 | /// (net!("192.168.2.0/23"), 4), 230 | /// ]); 231 | /// let mut map_b: PrefixMap = PrefixMap::from_iter([ 232 | /// (net!("192.168.0.0/22"), 10), 233 | /// (net!("192.168.0.0/23"), 20), 234 | /// (net!("192.168.2.0/24"), 30), 235 | /// ]); 236 | /// 237 | /// // Modify the two maps by adding their values for elements of the same prefix. 238 | /// for (_, l, r) in map_a.view_mut().union_mut(&mut map_b) { 239 | /// if let (Some(l), Some(r)) = (l, r) { 240 | /// *l += *r; 241 | /// *r = *l; 242 | /// } 243 | /// } 244 | /// 245 | /// assert_eq!( 246 | /// map_a.into_iter().collect::>(), 247 | /// vec![ 248 | /// (net!("192.168.0.0/20"), 1), 249 | /// (net!("192.168.0.0/22"), 12), 250 | /// (net!("192.168.0.0/24"), 3), 251 | /// (net!("192.168.2.0/23"), 4), 252 | /// ], 253 | /// ); 254 | /// assert_eq!( 255 | /// map_b.into_iter().collect::>(), 256 | /// vec![ 257 | /// (net!("192.168.0.0/22"), 12), 258 | /// (net!("192.168.0.0/23"), 20), 259 | /// (net!("192.168.2.0/24"), 30), 260 | /// ], 261 | /// ); 262 | /// # } 263 | /// ``` 264 | pub fn union_mut<'b, R>( 265 | &'b mut self, 266 | other: impl AsViewMut<'b, P, R>, 267 | ) -> UnionMut<'b, P, L, R> { 268 | let other = other.view_mut(); 269 | let nodes = next_indices( 270 | self.table, 271 | other.table, 272 | Some(self.loc.idx()), 273 | Some(other.loc.idx()), 274 | ); 275 | // Safety: We take the reference to the table from two TrieViewMut. Since they both have to 276 | // be created using TrieViewMut::new, we satisfy the conditions in `UnionMut::new`. 277 | unsafe { UnionMut::new(self.table, other.table, nodes) } 278 | } 279 | } 280 | 281 | impl<'a, P: Prefix, L, R> Union<'a, P, L, R> { 282 | fn extend( 283 | &mut self, 284 | indices: impl IntoIterator + 'static, 285 | lpm_l: Lpm<'a, P, L>, 286 | lpm_r: Lpm<'a, P, R>, 287 | ) { 288 | self.nodes.extend(extend_lpm( 289 | self.table_l, 290 | self.table_r, 291 | lpm_l, 292 | lpm_r, 293 | indices, 294 | )); 295 | } 296 | 297 | fn get_next( 298 | &self, 299 | prefix: &'a P, 300 | l: Option<&'a L>, 301 | r: Option<&'a R>, 302 | lpm_l: Lpm<'a, P, L>, 303 | lpm_r: Lpm<'a, P, R>, 304 | ) -> Option> { 305 | match (l, r) { 306 | (None, None) => None, 307 | (None, Some(right)) => Some(UnionItem::Right { 308 | prefix, 309 | left: lpm_l, 310 | right, 311 | }), 312 | (Some(left), None) => Some(UnionItem::Left { 313 | prefix, 314 | left, 315 | right: lpm_r, 316 | }), 317 | (Some(left), Some(right)) => Some(UnionItem::Both { 318 | prefix, 319 | left, 320 | right, 321 | }), 322 | } 323 | } 324 | } 325 | 326 | impl<'a, P: Prefix, L, R> Iterator for Union<'a, P, L, R> { 327 | type Item = UnionItem<'a, P, L, R>; 328 | 329 | fn next(&mut self) -> Option { 330 | while let Some((cur, lpm_l, lpm_r)) = self.nodes.pop() { 331 | match cur { 332 | UnionIndex::Both(l, r) => { 333 | let node_l = &self.table_l[l]; 334 | let node_r = &self.table_r[r]; 335 | self.extend( 336 | next_indices(self.table_l, self.table_r, node_l.right, node_r.right), 337 | lpm_l, 338 | lpm_r, 339 | ); 340 | self.extend( 341 | next_indices(self.table_l, self.table_r, node_l.left, node_r.left), 342 | lpm_l, 343 | lpm_r, 344 | ); 345 | if let Some(x) = self.get_next( 346 | &node_l.prefix, 347 | node_l.value.as_ref(), 348 | node_r.value.as_ref(), 349 | lpm_l, 350 | lpm_r, 351 | ) { 352 | return Some(x); 353 | } 354 | } 355 | UnionIndex::FirstL(l, r) => { 356 | let node_l = &self.table_l[l]; 357 | self.extend( 358 | next_indices_first_l( 359 | self.table_l, 360 | self.table_r, 361 | l, 362 | node_l.left, 363 | node_l.right, 364 | r, 365 | ), 366 | lpm_l, 367 | lpm_r, 368 | ); 369 | if let Some(x) = 370 | self.get_next(&node_l.prefix, node_l.value.as_ref(), None, lpm_l, lpm_r) 371 | { 372 | return Some(x); 373 | } 374 | } 375 | UnionIndex::FirstR(l, r) => { 376 | let node_r = &self.table_r[r]; 377 | self.extend( 378 | next_indices_first_r( 379 | self.table_l, 380 | self.table_r, 381 | l, 382 | r, 383 | node_r.left, 384 | node_r.right, 385 | ), 386 | lpm_l, 387 | lpm_r, 388 | ); 389 | if let Some(x) = 390 | self.get_next(&node_r.prefix, None, node_r.value.as_ref(), lpm_l, lpm_r) 391 | { 392 | return Some(x); 393 | } 394 | } 395 | UnionIndex::OnlyL(l) => { 396 | let node_l = &self.table_l[l]; 397 | if let Some(right) = node_l.right { 398 | self.extend([UnionIndex::OnlyL(right)], lpm_l, lpm_r); 399 | } 400 | if let Some(left) = node_l.left { 401 | self.extend([UnionIndex::OnlyL(left)], lpm_l, lpm_r); 402 | } 403 | if let Some(x) = 404 | self.get_next(&node_l.prefix, node_l.value.as_ref(), None, lpm_l, lpm_r) 405 | { 406 | return Some(x); 407 | } 408 | } 409 | UnionIndex::OnlyR(r) => { 410 | let node_r = &self.table_r[r]; 411 | if let Some(right) = node_r.right { 412 | self.extend([UnionIndex::OnlyR(right)], lpm_l, lpm_r); 413 | } 414 | if let Some(left) = node_r.left { 415 | self.extend([UnionIndex::OnlyR(left)], lpm_l, lpm_r); 416 | } 417 | if let Some(x) = 418 | self.get_next(&node_r.prefix, None, node_r.value.as_ref(), lpm_l, lpm_r) 419 | { 420 | return Some(x); 421 | } 422 | } 423 | } 424 | } 425 | None 426 | } 427 | } 428 | 429 | impl<'a, P: Prefix, L, R> Iterator for UnionMut<'a, P, L, R> { 430 | type Item = (&'a P, Option<&'a mut L>, Option<&'a mut R>); 431 | 432 | fn next(&mut self) -> Option { 433 | while let Some(cur) = self.nodes.pop() { 434 | // safety: map is a tree. Every node is visited exactly once during the iteration 435 | // (self.nodes is not public). Therefore, each in each iteration of this loop (also 436 | // between multiple calls to `next`), the index `cur` is different to any of the earlier 437 | // iterations. It is therefore safe to extend the lifetime of the elements to 'a (which 438 | // is the lifetime for which `self` has an exclusive reference over the map). 439 | match cur { 440 | UnionIndex::Both(l, r) => { 441 | let node_l = &self.table_l[l]; 442 | let node_r = &self.table_r[r]; 443 | self.nodes.extend(next_indices( 444 | self.table_l, 445 | self.table_r, 446 | node_l.right, 447 | node_r.right, 448 | )); 449 | self.nodes.extend(next_indices( 450 | self.table_l, 451 | self.table_r, 452 | node_l.left, 453 | node_r.left, 454 | )); 455 | let node_l = unsafe { self.table_l.get_mut(l) }; 456 | let node_r = unsafe { self.table_r.get_mut(r) }; 457 | if node_l.value.is_some() || node_r.value.is_some() { 458 | return Some(( 459 | &node_l.prefix, 460 | node_l.value.as_mut(), 461 | node_r.value.as_mut(), 462 | )); 463 | } 464 | } 465 | UnionIndex::FirstL(l, r) => { 466 | let node_l = &self.table_l[l]; 467 | self.nodes.extend(next_indices_first_l( 468 | self.table_l, 469 | self.table_r, 470 | l, 471 | node_l.left, 472 | node_l.right, 473 | r, 474 | )); 475 | let node_l = unsafe { self.table_l.get_mut(l) }; 476 | if node_l.value.is_some() { 477 | return Some((&node_l.prefix, node_l.value.as_mut(), None)); 478 | } 479 | } 480 | UnionIndex::FirstR(l, r) => { 481 | let node_r = &self.table_r[r]; 482 | self.nodes.extend(next_indices_first_r( 483 | self.table_l, 484 | self.table_r, 485 | l, 486 | r, 487 | node_r.left, 488 | node_r.right, 489 | )); 490 | let node_r = unsafe { self.table_r.get_mut(r) }; 491 | if node_r.value.is_some() { 492 | return Some((&node_r.prefix, None, node_r.value.as_mut())); 493 | } 494 | } 495 | UnionIndex::OnlyL(l) => { 496 | let node_l = unsafe { self.table_l.get_mut(l) }; 497 | if let Some(right) = node_l.right { 498 | self.nodes.push(UnionIndex::OnlyL(right)); 499 | } 500 | if let Some(left) = node_l.left { 501 | self.nodes.push(UnionIndex::OnlyL(left)); 502 | } 503 | if node_l.value.is_some() { 504 | return Some((&node_l.prefix, node_l.value.as_mut(), None)); 505 | } 506 | } 507 | UnionIndex::OnlyR(r) => { 508 | let node_r = unsafe { self.table_r.get_mut(r) }; 509 | if let Some(right) = node_r.right { 510 | self.nodes.push(UnionIndex::OnlyR(right)); 511 | } 512 | if let Some(left) = node_r.left { 513 | self.nodes.push(UnionIndex::OnlyR(left)); 514 | } 515 | if node_r.value.is_some() { 516 | return Some((&node_r.prefix, None, node_r.value.as_mut())); 517 | } 518 | } 519 | } 520 | } 521 | None 522 | } 523 | } 524 | 525 | fn next_indices<'a, P: Prefix, L, R>( 526 | table_l: &'a Table, 527 | table_r: &'a Table, 528 | node_l: Option, 529 | node_r: Option, 530 | ) -> Vec { 531 | match (node_l, node_r) { 532 | (None, Some(b)) => vec![UnionIndex::OnlyR(b)], 533 | (Some(a), None) => vec![UnionIndex::OnlyL(a)], 534 | (Some(a), Some(b)) => { 535 | let p_a = &table_l[a].prefix; 536 | let p_b = &table_r[b].prefix; 537 | if p_a.prefix_len() == p_b.prefix_len() { 538 | match p_a.mask().cmp(&p_b.mask()) { 539 | std::cmp::Ordering::Less => { 540 | vec![UnionIndex::OnlyR(b), UnionIndex::OnlyL(a)] 541 | } 542 | std::cmp::Ordering::Equal => { 543 | vec![UnionIndex::Both(a, b)] 544 | } 545 | std::cmp::Ordering::Greater => { 546 | vec![UnionIndex::OnlyL(a), UnionIndex::OnlyR(b)] 547 | } 548 | } 549 | } else if p_a.contains(p_b) { 550 | vec![UnionIndex::FirstL(a, b)] 551 | } else if p_b.contains(p_a) { 552 | vec![UnionIndex::FirstR(a, b)] 553 | } else { 554 | if p_a.mask() < p_b.mask() { 555 | vec![UnionIndex::OnlyR(b), UnionIndex::OnlyL(a)] 556 | } else { 557 | vec![UnionIndex::OnlyL(a), UnionIndex::OnlyR(b)] 558 | } 559 | } 560 | } 561 | _ => vec![], 562 | } 563 | } 564 | 565 | fn next_indices_first_l<'a, P: Prefix, L, R>( 566 | table_l: &'a Table, 567 | table_r: &'a Table, 568 | l: usize, 569 | ll: Option, 570 | lr: Option, 571 | r: usize, 572 | ) -> Vec { 573 | match (ll, lr) { 574 | (None, None) => vec![UnionIndex::OnlyR(r)], 575 | (None, Some(lr)) => next_indices(table_l, table_r, Some(lr), Some(r)), 576 | (Some(ll), None) => next_indices(table_l, table_r, Some(ll), Some(r)), 577 | (Some(ll), Some(lr)) => { 578 | if to_right(&table_l[l].prefix, &table_r[r].prefix) { 579 | let mut idxes = next_indices(table_l, table_r, Some(lr), Some(r)); 580 | idxes.push(UnionIndex::OnlyL(ll)); 581 | idxes 582 | } else { 583 | let mut idxes = next_indices(table_l, table_r, Some(ll), Some(r)); 584 | idxes.insert(0, UnionIndex::OnlyL(lr)); 585 | idxes 586 | } 587 | } 588 | } 589 | } 590 | 591 | fn next_indices_first_r<'a, P: Prefix, L, R>( 592 | table_l: &'a Table, 593 | table_r: &'a Table, 594 | l: usize, 595 | r: usize, 596 | rl: Option, 597 | rr: Option, 598 | ) -> Vec { 599 | match (rl, rr) { 600 | (None, None) => vec![UnionIndex::OnlyL(l)], 601 | (None, Some(rr)) => next_indices(table_l, table_r, Some(l), Some(rr)), 602 | (Some(rl), None) => next_indices(table_l, table_r, Some(l), Some(rl)), 603 | (Some(rl), Some(rr)) => { 604 | if to_right(&table_r[r].prefix, &table_l[l].prefix) { 605 | let mut idxes = next_indices(table_l, table_r, Some(l), Some(rr)); 606 | idxes.push(UnionIndex::OnlyR(rl)); 607 | idxes 608 | } else { 609 | let mut idxes = next_indices(table_l, table_r, Some(l), Some(rl)); 610 | idxes.insert(0, UnionIndex::OnlyR(rr)); 611 | idxes 612 | } 613 | } 614 | } 615 | } 616 | 617 | fn extend_lpm<'a, P: Prefix, L, R>( 618 | table_l: &'a Table, 619 | table_r: &'a Table, 620 | lpm_l: Lpm<'a, P, L>, 621 | lpm_r: Lpm<'a, P, R>, 622 | indices: impl IntoIterator + 'static, 623 | ) -> impl Iterator> + 'a { 624 | let get_lpm_l = move |l: usize| table_l[l].prefix_value().or(lpm_l); 625 | let get_lpm_r = move |r: usize| table_r[r].prefix_value().or(lpm_r); 626 | indices.into_iter().map(move |x| match x { 627 | UnionIndex::Both(l, r) => (x, get_lpm_l(l), get_lpm_r(r)), 628 | UnionIndex::FirstL(l, _) | UnionIndex::OnlyL(l) => (x, get_lpm_l(l), lpm_r), 629 | UnionIndex::FirstR(_, r) | UnionIndex::OnlyR(r) => (x, lpm_l, get_lpm_r(r)), 630 | }) 631 | } 632 | --------------------------------------------------------------------------------