├── .copywrite.hcl ├── .github ├── dependabot.yaml ├── release.yaml └── workflows │ └── ci.yaml ├── .gitignore ├── .golangci.yaml ├── CODEOWNERS ├── LICENSE ├── README.md ├── benchmark_test.go ├── collection.go ├── collection_test.go ├── examples_hashset_test.go ├── examples_set_test.go ├── examples_treeset_test.go ├── go.mod ├── go.sum ├── hashset.go ├── hashset_test.go ├── serialization.go ├── serialization_test.go ├── set.go ├── set_test.go ├── stack.go ├── stack_test.go ├── treeset.go └── treeset_test.go /.copywrite.hcl: -------------------------------------------------------------------------------- 1 | schema_version = 1 2 | 3 | project { 4 | license = "MPL-2.0" 5 | copyright_year = 2022 6 | header_ignore = [ 7 | ".golangci.yaml", 8 | ".copywrite.hcl", 9 | ".github/**", 10 | ] 11 | } 12 | 13 | -------------------------------------------------------------------------------- /.github/dependabot.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) HashiCorp, Inc. 2 | # SPDX-License-Identifier: MPL-2.0 3 | 4 | version: 2 5 | updates: 6 | - package-ecosystem: gomod 7 | directory: "/" 8 | schedule: 9 | interval: "monthly" 10 | labels: 11 | - "dependabot" 12 | -------------------------------------------------------------------------------- /.github/release.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) HashiCorp, Inc. 2 | # SPDX-License-Identifier: MPL-2.0 3 | 4 | changelog: 5 | categories: 6 | - title: Changes 7 | labels: 8 | - "*" 9 | 10 | -------------------------------------------------------------------------------- /.github/workflows/ci.yaml: -------------------------------------------------------------------------------- 1 | name: Run CI Tests 2 | on: 3 | push: 4 | branches: 5 | - main 6 | paths-ignore: 7 | - README.md 8 | - .gitignore 9 | pull_request: 10 | paths-ignore: 11 | - README.md 12 | - .gitignore 13 | jobs: 14 | run-copywrite: 15 | runs-on: ubuntu-24.04 16 | steps: 17 | - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0 18 | - uses: hashicorp/setup-copywrite@v1.1.3 19 | - name: verify copyright 20 | run: | 21 | copywrite headers --plan 22 | run-lint: 23 | timeout-minutes: 10 24 | runs-on: ubuntu-24.04 25 | steps: 26 | - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0 27 | - uses: hashicorp/setup-golang@v3 28 | with: 29 | version-file: go.mod 30 | - uses: golangci/golangci-lint-action@3a919529898de77ec3da873e3063ca4b10e7f5cc # v3.7.0 31 | with: 32 | version: v1.60.1 33 | skip-cache: true 34 | run-tests: 35 | timeout-minutes: 10 36 | runs-on: ubuntu-24.04 37 | steps: 38 | - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0 39 | - uses: hashicorp/setup-golang@v3 40 | with: 41 | version-file: go.mod 42 | - name: Run Go Vet 43 | run: | 44 | go vet ./... 45 | - name: Run Go Fmt 46 | run: | 47 | files=$(go fmt ./...) 48 | if [ -n "$files" ]; then 49 | echo "Please run gofmt on these files ..." 50 | echo "$files" 51 | exit 1 52 | fi 53 | - name: Run Go Test 54 | run: | 55 | go test -race -v ./... 56 | 57 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # Test binary, built with `go test -c` 9 | *.test 10 | 11 | # Output of the go coverage tool, specifically when used with LiteIDE 12 | *.out 13 | 14 | # Dependency directories (remove the comment below to include it) 15 | # vendor/ 16 | 17 | # Go work files 18 | go.work 19 | go.work.sum 20 | 21 | # IDE files 22 | .idea 23 | .fleet 24 | -------------------------------------------------------------------------------- /.golangci.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) HashiCorp, Inc. 2 | # SPDX-License-Identifier: MPL-2.0 3 | 4 | run: 5 | timeout: 5m 6 | linters: 7 | enable: 8 | - gofmt 9 | - errcheck 10 | - errname 11 | - errorlint 12 | - bodyclose 13 | - durationcheck 14 | - exhaustive 15 | - goconst 16 | - whitespace 17 | 18 | -------------------------------------------------------------------------------- /CODEOWNERS: -------------------------------------------------------------------------------- 1 | # codeowner default 2 | * @hashicorp/github-nomad-core @hashicorp/nomad-eng 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2022 HashiCorp, Inc. 2 | 3 | Mozilla Public License Version 2.0 4 | ================================== 5 | 6 | 1. Definitions 7 | -------------- 8 | 9 | 1.1. "Contributor" 10 | means each individual or legal entity that creates, contributes to 11 | the creation of, or owns Covered Software. 12 | 13 | 1.2. "Contributor Version" 14 | means the combination of the Contributions of others (if any) used 15 | by a Contributor and that particular Contributor's Contribution. 16 | 17 | 1.3. "Contribution" 18 | means Covered Software of a particular Contributor. 19 | 20 | 1.4. "Covered Software" 21 | means Source Code Form to which the initial Contributor has attached 22 | the notice in Exhibit A, the Executable Form of such Source Code 23 | Form, and Modifications of such Source Code Form, in each case 24 | including portions thereof. 25 | 26 | 1.5. "Incompatible With Secondary Licenses" 27 | means 28 | 29 | (a) that the initial Contributor has attached the notice described 30 | in Exhibit B to the Covered Software; or 31 | 32 | (b) that the Covered Software was made available under the terms of 33 | version 1.1 or earlier of the License, but not also under the 34 | terms of a Secondary License. 35 | 36 | 1.6. "Executable Form" 37 | means any form of the work other than Source Code Form. 38 | 39 | 1.7. "Larger Work" 40 | means a work that combines Covered Software with other material, in 41 | a separate file or files, that is not Covered Software. 42 | 43 | 1.8. "License" 44 | means this document. 45 | 46 | 1.9. "Licensable" 47 | means having the right to grant, to the maximum extent possible, 48 | whether at the time of the initial grant or subsequently, any and 49 | all of the rights conveyed by this License. 50 | 51 | 1.10. "Modifications" 52 | means any of the following: 53 | 54 | (a) any file in Source Code Form that results from an addition to, 55 | deletion from, or modification of the contents of Covered 56 | Software; or 57 | 58 | (b) any new file in Source Code Form that contains any Covered 59 | Software. 60 | 61 | 1.11. "Patent Claims" of a Contributor 62 | means any patent claim(s), including without limitation, method, 63 | process, and apparatus claims, in any patent Licensable by such 64 | Contributor that would be infringed, but for the grant of the 65 | License, by the making, using, selling, offering for sale, having 66 | made, import, or transfer of either its Contributions or its 67 | Contributor Version. 68 | 69 | 1.12. "Secondary License" 70 | means either the GNU General Public License, Version 2.0, the GNU 71 | Lesser General Public License, Version 2.1, the GNU Affero General 72 | Public License, Version 3.0, or any later versions of those 73 | licenses. 74 | 75 | 1.13. "Source Code Form" 76 | means the form of the work preferred for making modifications. 77 | 78 | 1.14. "You" (or "Your") 79 | means an individual or a legal entity exercising rights under this 80 | License. For legal entities, "You" includes any entity that 81 | controls, is controlled by, or is under common control with You. For 82 | purposes of this definition, "control" means (a) the power, direct 83 | or indirect, to cause the direction or management of such entity, 84 | whether by contract or otherwise, or (b) ownership of more than 85 | fifty percent (50%) of the outstanding shares or beneficial 86 | ownership of such entity. 87 | 88 | 2. License Grants and Conditions 89 | -------------------------------- 90 | 91 | 2.1. Grants 92 | 93 | Each Contributor hereby grants You a world-wide, royalty-free, 94 | non-exclusive license: 95 | 96 | (a) under intellectual property rights (other than patent or trademark) 97 | Licensable by such Contributor to use, reproduce, make available, 98 | modify, display, perform, distribute, and otherwise exploit its 99 | Contributions, either on an unmodified basis, with Modifications, or 100 | as part of a Larger Work; and 101 | 102 | (b) under Patent Claims of such Contributor to make, use, sell, offer 103 | for sale, have made, import, and otherwise transfer either its 104 | Contributions or its Contributor Version. 105 | 106 | 2.2. Effective Date 107 | 108 | The licenses granted in Section 2.1 with respect to any Contribution 109 | become effective for each Contribution on the date the Contributor first 110 | distributes such Contribution. 111 | 112 | 2.3. Limitations on Grant Scope 113 | 114 | The licenses granted in this Section 2 are the only rights granted under 115 | this License. No additional rights or licenses will be implied from the 116 | distribution or licensing of Covered Software under this License. 117 | Notwithstanding Section 2.1(b) above, no patent license is granted by a 118 | Contributor: 119 | 120 | (a) for any code that a Contributor has removed from Covered Software; 121 | or 122 | 123 | (b) for infringements caused by: (i) Your and any other third party's 124 | modifications of Covered Software, or (ii) the combination of its 125 | Contributions with other software (except as part of its Contributor 126 | Version); or 127 | 128 | (c) under Patent Claims infringed by Covered Software in the absence of 129 | its Contributions. 130 | 131 | This License does not grant any rights in the trademarks, service marks, 132 | or logos of any Contributor (except as may be necessary to comply with 133 | the notice requirements in Section 3.4). 134 | 135 | 2.4. Subsequent Licenses 136 | 137 | No Contributor makes additional grants as a result of Your choice to 138 | distribute the Covered Software under a subsequent version of this 139 | License (see Section 10.2) or under the terms of a Secondary License (if 140 | permitted under the terms of Section 3.3). 141 | 142 | 2.5. Representation 143 | 144 | Each Contributor represents that the Contributor believes its 145 | Contributions are its original creation(s) or it has sufficient rights 146 | to grant the rights to its Contributions conveyed by this License. 147 | 148 | 2.6. Fair Use 149 | 150 | This License is not intended to limit any rights You have under 151 | applicable copyright doctrines of fair use, fair dealing, or other 152 | equivalents. 153 | 154 | 2.7. Conditions 155 | 156 | Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted 157 | in Section 2.1. 158 | 159 | 3. Responsibilities 160 | ------------------- 161 | 162 | 3.1. Distribution of Source Form 163 | 164 | All distribution of Covered Software in Source Code Form, including any 165 | Modifications that You create or to which You contribute, must be under 166 | the terms of this License. You must inform recipients that the Source 167 | Code Form of the Covered Software is governed by the terms of this 168 | License, and how they can obtain a copy of this License. You may not 169 | attempt to alter or restrict the recipients' rights in the Source Code 170 | Form. 171 | 172 | 3.2. Distribution of Executable Form 173 | 174 | If You distribute Covered Software in Executable Form then: 175 | 176 | (a) such Covered Software must also be made available in Source Code 177 | Form, as described in Section 3.1, and You must inform recipients of 178 | the Executable Form how they can obtain a copy of such Source Code 179 | Form by reasonable means in a timely manner, at a charge no more 180 | than the cost of distribution to the recipient; and 181 | 182 | (b) You may distribute such Executable Form under the terms of this 183 | License, or sublicense it under different terms, provided that the 184 | license for the Executable Form does not attempt to limit or alter 185 | the recipients' rights in the Source Code Form under this License. 186 | 187 | 3.3. Distribution of a Larger Work 188 | 189 | You may create and distribute a Larger Work under terms of Your choice, 190 | provided that You also comply with the requirements of this License for 191 | the Covered Software. If the Larger Work is a combination of Covered 192 | Software with a work governed by one or more Secondary Licenses, and the 193 | Covered Software is not Incompatible With Secondary Licenses, this 194 | License permits You to additionally distribute such Covered Software 195 | under the terms of such Secondary License(s), so that the recipient of 196 | the Larger Work may, at their option, further distribute the Covered 197 | Software under the terms of either this License or such Secondary 198 | License(s). 199 | 200 | 3.4. Notices 201 | 202 | You may not remove or alter the substance of any license notices 203 | (including copyright notices, patent notices, disclaimers of warranty, 204 | or limitations of liability) contained within the Source Code Form of 205 | the Covered Software, except that You may alter any license notices to 206 | the extent required to remedy known factual inaccuracies. 207 | 208 | 3.5. Application of Additional Terms 209 | 210 | You may choose to offer, and to charge a fee for, warranty, support, 211 | indemnity or liability obligations to one or more recipients of Covered 212 | Software. However, You may do so only on Your own behalf, and not on 213 | behalf of any Contributor. You must make it absolutely clear that any 214 | such warranty, support, indemnity, or liability obligation is offered by 215 | You alone, and You hereby agree to indemnify every Contributor for any 216 | liability incurred by such Contributor as a result of warranty, support, 217 | indemnity or liability terms You offer. You may include additional 218 | disclaimers of warranty and limitations of liability specific to any 219 | jurisdiction. 220 | 221 | 4. Inability to Comply Due to Statute or Regulation 222 | --------------------------------------------------- 223 | 224 | If it is impossible for You to comply with any of the terms of this 225 | License with respect to some or all of the Covered Software due to 226 | statute, judicial order, or regulation then You must: (a) comply with 227 | the terms of this License to the maximum extent possible; and (b) 228 | describe the limitations and the code they affect. Such description must 229 | be placed in a text file included with all distributions of the Covered 230 | Software under this License. Except to the extent prohibited by statute 231 | or regulation, such description must be sufficiently detailed for a 232 | recipient of ordinary skill to be able to understand it. 233 | 234 | 5. Termination 235 | -------------- 236 | 237 | 5.1. The rights granted under this License will terminate automatically 238 | if You fail to comply with any of its terms. However, if You become 239 | compliant, then the rights granted under this License from a particular 240 | Contributor are reinstated (a) provisionally, unless and until such 241 | Contributor explicitly and finally terminates Your grants, and (b) on an 242 | ongoing basis, if such Contributor fails to notify You of the 243 | non-compliance by some reasonable means prior to 60 days after You have 244 | come back into compliance. Moreover, Your grants from a particular 245 | Contributor are reinstated on an ongoing basis if such Contributor 246 | notifies You of the non-compliance by some reasonable means, this is the 247 | first time You have received notice of non-compliance with this License 248 | from such Contributor, and You become compliant prior to 30 days after 249 | Your receipt of the notice. 250 | 251 | 5.2. If You initiate litigation against any entity by asserting a patent 252 | infringement claim (excluding declaratory judgment actions, 253 | counter-claims, and cross-claims) alleging that a Contributor Version 254 | directly or indirectly infringes any patent, then the rights granted to 255 | You by any and all Contributors for the Covered Software under Section 256 | 2.1 of this License shall terminate. 257 | 258 | 5.3. In the event of termination under Sections 5.1 or 5.2 above, all 259 | end user license agreements (excluding distributors and resellers) which 260 | have been validly granted by You or Your distributors under this License 261 | prior to termination shall survive termination. 262 | 263 | ************************************************************************ 264 | * * 265 | * 6. Disclaimer of Warranty * 266 | * ------------------------- * 267 | * * 268 | * Covered Software is provided under this License on an "as is" * 269 | * basis, without warranty of any kind, either expressed, implied, or * 270 | * statutory, including, without limitation, warranties that the * 271 | * Covered Software is free of defects, merchantable, fit for a * 272 | * particular purpose or non-infringing. The entire risk as to the * 273 | * quality and performance of the Covered Software is with You. * 274 | * Should any Covered Software prove defective in any respect, You * 275 | * (not any Contributor) assume the cost of any necessary servicing, * 276 | * repair, or correction. This disclaimer of warranty constitutes an * 277 | * essential part of this License. No use of any Covered Software is * 278 | * authorized under this License except under this disclaimer. * 279 | * * 280 | ************************************************************************ 281 | 282 | ************************************************************************ 283 | * * 284 | * 7. Limitation of Liability * 285 | * -------------------------- * 286 | * * 287 | * Under no circumstances and under no legal theory, whether tort * 288 | * (including negligence), contract, or otherwise, shall any * 289 | * Contributor, or anyone who distributes Covered Software as * 290 | * permitted above, be liable to You for any direct, indirect, * 291 | * special, incidental, or consequential damages of any character * 292 | * including, without limitation, damages for lost profits, loss of * 293 | * goodwill, work stoppage, computer failure or malfunction, or any * 294 | * and all other commercial damages or losses, even if such party * 295 | * shall have been informed of the possibility of such damages. This * 296 | * limitation of liability shall not apply to liability for death or * 297 | * personal injury resulting from such party's negligence to the * 298 | * extent applicable law prohibits such limitation. Some * 299 | * jurisdictions do not allow the exclusion or limitation of * 300 | * incidental or consequential damages, so this exclusion and * 301 | * limitation may not apply to You. * 302 | * * 303 | ************************************************************************ 304 | 305 | 8. Litigation 306 | ------------- 307 | 308 | Any litigation relating to this License may be brought only in the 309 | courts of a jurisdiction where the defendant maintains its principal 310 | place of business and such litigation shall be governed by laws of that 311 | jurisdiction, without reference to its conflict-of-law provisions. 312 | Nothing in this Section shall prevent a party's ability to bring 313 | cross-claims or counter-claims. 314 | 315 | 9. Miscellaneous 316 | ---------------- 317 | 318 | This License represents the complete agreement concerning the subject 319 | matter hereof. If any provision of this License is held to be 320 | unenforceable, such provision shall be reformed only to the extent 321 | necessary to make it enforceable. Any law or regulation which provides 322 | that the language of a contract shall be construed against the drafter 323 | shall not be used to construe this License against a Contributor. 324 | 325 | 10. Versions of the License 326 | --------------------------- 327 | 328 | 10.1. New Versions 329 | 330 | Mozilla Foundation is the license steward. Except as provided in Section 331 | 10.3, no one other than the license steward has the right to modify or 332 | publish new versions of this License. Each version will be given a 333 | distinguishing version number. 334 | 335 | 10.2. Effect of New Versions 336 | 337 | You may distribute the Covered Software under the terms of the version 338 | of the License under which You originally received the Covered Software, 339 | or under the terms of any subsequent version published by the license 340 | steward. 341 | 342 | 10.3. Modified Versions 343 | 344 | If you create software not governed by this License, and you want to 345 | create a new license for such software, you may create and use a 346 | modified version of this License if you rename the license and remove 347 | any references to the name of the license steward (except to note that 348 | such modified license differs from this License). 349 | 350 | 10.4. Distributing Source Code Form that is Incompatible With Secondary 351 | Licenses 352 | 353 | If You choose to distribute Source Code Form that is Incompatible With 354 | Secondary Licenses under the terms of this version of the License, the 355 | notice described in Exhibit B of this License must be attached. 356 | 357 | Exhibit A - Source Code Form License Notice 358 | ------------------------------------------- 359 | 360 | This Source Code Form is subject to the terms of the Mozilla Public 361 | License, v. 2.0. If a copy of the MPL was not distributed with this 362 | file, You can obtain one at http://mozilla.org/MPL/2.0/. 363 | 364 | If it is not possible or desirable to put the notice in a particular 365 | file, then You may include the notice in a location (such as a LICENSE 366 | file in a relevant directory) where a recipient would be likely to look 367 | for such a notice. 368 | 369 | You may add additional accurate notices of copyright ownership. 370 | 371 | Exhibit B - "Incompatible With Secondary Licenses" Notice 372 | --------------------------------------------------------- 373 | 374 | This Source Code Form is "Incompatible With Secondary Licenses", as 375 | defined by the Mozilla Public License, v. 2.0. 376 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # go-set 2 | 3 | [![Go Reference](https://pkg.go.dev/badge/github.com/hashicorp/go-set.svg)](https://pkg.go.dev/github.com/hashicorp/go-set/v3) 4 | [![Run CI Tests](https://github.com/hashicorp/go-set/actions/workflows/ci.yaml/badge.svg)](https://github.com/hashicorp/go-set/actions/workflows/ci.yaml) 5 | [![GitHub](https://img.shields.io/github/license/hashicorp/go-set)](LICENSE) 6 | 7 | The `go-set` repository provides a `set` package containing a few 8 | generic [Set](https://en.wikipedia.org/wiki/Set) implementations for Go. 9 | 10 | --- 11 | 12 | **PSA** August 2024 - The **v3** version of this package has been published, 13 | starting at tag version `v3.0.0`. A description of the changes including 14 | backwards incompatibilities can be found in https://github.com/hashicorp/go-set/issues/90 15 | 16 | _Requires `go1.23` or later._ 17 | 18 | --- 19 | 20 | **PSA** October 2023 - The **v2** version of this package has been published, 21 | starting at tag version `v2.1.0`. A description of the changes including 22 | backwards incompatibilities can be found in https://github.com/hashicorp/go-set/issues/73 23 | 24 | --- 25 | 26 | Each implementation is optimal for a particular use case. 27 | 28 | **Set[T]** is ideal for `comparable` types. 29 | - backed by `map` builtin 30 | - commonly used with `string`, `int`, simple `struct` types, etc. 31 | 32 | **HashSet[T]** is useful for types that implement a `Hash()` function. 33 | - backed by `map` builtin 34 | - commonly used with complex structs 35 | - also works with custom `HashFunc[T]` implementations 36 | 37 | **TreeSet[T]** is useful for comparable data (via `CompareFunc[T]`) 38 | - backed by Red-Black Binary Search Tree 39 | - commonly used with complex structs with extrinsic order 40 | - efficient iteration in sort order 41 | - additional methods `Min` / `Max` / `TopK` / `BottomK` 42 | 43 | This package is not thread-safe. 44 | 45 | --- 46 | 47 | # Documentation 48 | 49 | The full `go-set` package reference is available on [pkg.go.dev](https://pkg.go.dev/github.com/hashicorp/go-set/v3). 50 | 51 | # Install 52 | 53 | ```shell 54 | go get github.com/hashicorp/go-set/v3@latest 55 | ``` 56 | 57 | ```shell 58 | import "github.com/hashicorp/go-set/v3" 59 | ``` 60 | 61 | # Motivation 62 | 63 | Package `set` helps reduce the boiler plate of using a `map[]struct{}` as a set. 64 | 65 | Say we want to de-duplicate a slice of strings 66 | ```go 67 | items := []string{"mitchell", "armon", "jack", "dave", "armon", "dave"} 68 | ``` 69 | 70 | A typical example of the classic way using `map` built-in: 71 | ```go 72 | m := make(map[string]struct{}) 73 | for _, item := range items { 74 | m[item] = struct{}{} 75 | } 76 | list := make([]string, 0, len(items)) 77 | for k := range m { 78 | list = append(list, k) 79 | } 80 | ``` 81 | 82 | The same result, but in one line using package `go-set`. 83 | ```go 84 | list := set.From[string](items).Slice() 85 | ``` 86 | 87 | # Set 88 | 89 | The `go-set` package includes `Set` for types that satisfy the `comparable` constraint. 90 | Uniqueness of a set elements is guaranteed via shallow comparison (result of == operator). 91 | 92 | Note: if pointers or structs with pointer fields are stored in the `Set`, they will 93 | be compared in the sense of pointer addresses, not in the sense of referenced values. 94 | Due to this fact the `Set` type is recommended to be used with builtin types like 95 | `string`, `int`, or simple struct types with no pointers. `Set` usage with pointers or 96 | structs with pointer is also possible if shallow equality is acceptable. 97 | 98 | # HashSet 99 | 100 | The `go-set` package includes `HashSet` for types that implement a `Hash()` function. 101 | The custom type must satisfy `HashFunc[H Hash]` - essentially any `Hash()` function 102 | that returns a `string` or `integer`. This enables types to use string-y hash 103 | functions like `md5`, `sha1`, or even `GoString()`, but also enables types to 104 | implement an efficient hash function using a hash code based on prime multiples. 105 | 106 | # TreeSet 107 | 108 | The `go-set` package includes `TreeSet` for creating sorted sets. A `TreeSet` may 109 | be used with any type `T` as the comparison between elements is provided by implementing 110 | `CompareFunc[T]`. The standard library `cmp.Compare` function provides a convenient 111 | implementation of `CompareFunc` for `cmp.Ordered` types like `string` or `int`. A 112 | `TreeSet` is backed by an underlying balanced binary search tree, making operations 113 | like in-order traversal efficient, in addition to enabling functions like `Min()`, 114 | `Max()`, `TopK()`, and `BottomK()`. 115 | 116 | # Collection[T] 117 | 118 | The `Collection[T]` interface is implemented by each of `Set`, `HashSet`, and `TreeSet`. 119 | 120 | It serves as a useful abstraction over the common methods implemented by each set type. 121 | 122 | ### Iteration 123 | 124 | Starting with `v3` each of `Set`, `HashSet`, and `TreeSet` implement an `Items` 125 | method. It can be used with the `range` keyword for iterating through each 126 | element in the set. 127 | 128 | ```go 129 | // e.g. print each element in the set 130 | for _, item := range s.Items() { 131 | fmt.Println(item) 132 | } 133 | ``` 134 | 135 | # Set Examples 136 | 137 | Below are simple example usages of `Set` 138 | 139 | ```go 140 | s := set.New[int](10) 141 | s.Insert(1) 142 | s.InsertSlice([]int{2, 3, 4}) 143 | s.Size() 144 | ``` 145 | 146 | ```go 147 | s := set.From[string]([]string{"one", "two", "three"}) 148 | s.Contains("three") 149 | s.Remove("one") 150 | ``` 151 | 152 | 153 | ```go 154 | a := set.From[int]([]int{2, 4, 6, 8}) 155 | b := set.From[int]([]int{4, 5, 6}) 156 | a.Intersect(b) 157 | ``` 158 | 159 | # HashSet Examples 160 | 161 | Below are simple example usages of `HashSet` 162 | 163 | (using a hash code) 164 | ```go 165 | type inventory struct { 166 | item int 167 | serial int 168 | } 169 | 170 | func (i *inventory) Hash() int { 171 | code := 3 * item * 5 * serial 172 | return code 173 | } 174 | 175 | i1 := &inventory{item: 42, serial: 101} 176 | 177 | s := set.NewHashSet[*inventory, int](10) 178 | s.Insert(i1) 179 | ``` 180 | 181 | (using a string hash) 182 | ```go 183 | type employee struct { 184 | name string 185 | id int 186 | } 187 | 188 | func (e *employee) Hash() string { 189 | return fmt.Sprintf("%s:%d", e.name, e.id) 190 | } 191 | 192 | e1 := &employee{name: "armon", id: 2} 193 | 194 | s := set.NewHashSet[*employee, string](10) 195 | s.Insert(e1) 196 | ``` 197 | 198 | # TreeSet Examples 199 | 200 | Below are simple example usages of `TreeSet` 201 | 202 | ```go 203 | ts := NewTreeSet[int](Compare[int]) 204 | ts.Insert(5) 205 | ``` 206 | 207 | ```go 208 | type waypoint struct { 209 | distance int 210 | name string 211 | } 212 | 213 | // compare implements CompareFunc 214 | compare := func(w1, w2 *waypoint) int { 215 | return w1.distance - w2.distance 216 | } 217 | 218 | ts := NewTreeSet[*waypoint](compare) 219 | ts.Insert(&waypoint{distance: 42, name: "tango"}) 220 | ts.Insert(&waypoint{distance: 13, name: "alpha"}) 221 | ts.Insert(&waypoint{distance: 71, name: "xray"}) 222 | ``` 223 | 224 | -------------------------------------------------------------------------------- /benchmark_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "cmp" 8 | "math/rand" 9 | "sort" 10 | "testing" 11 | ) 12 | 13 | type test struct { 14 | size int 15 | name string 16 | } 17 | 18 | var cases []test = []test{ 19 | {size: 10, name: "10"}, 20 | {size: 1_000, name: "1000"}, 21 | {size: 100_000, name: "100000"}, 22 | {size: 1_000_000, name: "1000000"}, 23 | } 24 | 25 | func random[I ~int](n int) []I { 26 | result := make([]I, n) 27 | for i := 0; i < n; i++ { 28 | result[i] = I(rand.Int()) 29 | } 30 | return result 31 | } 32 | 33 | func unsort[I ~int](s []I) { 34 | s[0], s[len(s)-1] = s[len(s)-1], s[0] 35 | } 36 | 37 | type hashint int 38 | 39 | func (i hashint) Hash() int { 40 | return int(i) 41 | } 42 | 43 | func BenchmarkSlice_Insert(b *testing.B) { 44 | for _, tc := range cases { 45 | s := random[int](tc.size) 46 | b.Run(tc.name, func(b *testing.B) { 47 | for i := 0; i < b.N; i++ { 48 | s = append(s, i) 49 | } 50 | }) 51 | } 52 | } 53 | 54 | func BenchmarkSet_Insert(b *testing.B) { 55 | for _, tc := range cases { 56 | s := From(random[int](tc.size)) 57 | b.Run(tc.name, func(b *testing.B) { 58 | for i := 0; i < b.N; i++ { 59 | s.Insert(i) 60 | } 61 | }) 62 | } 63 | } 64 | 65 | func BenchmarkHashSet_Insert(b *testing.B) { 66 | for _, tc := range cases { 67 | hs := HashSetFrom[hashint, int](random[hashint](tc.size)) 68 | b.Run(tc.name, func(b *testing.B) { 69 | for i := 0; i < b.N; i++ { 70 | hs.Insert(hashint(i)) 71 | } 72 | }) 73 | } 74 | } 75 | 76 | func BenchmarkTreeSet_Insert(b *testing.B) { 77 | for _, tc := range cases { 78 | ts := TreeSetFrom[int](random[int](tc.size), cmp.Compare[int]) 79 | b.Run(tc.name, func(b *testing.B) { 80 | for i := 0; i < b.N; i++ { 81 | ts.Insert(i) 82 | } 83 | }) 84 | } 85 | } 86 | 87 | func BenchmarkSlice_Minimum(b *testing.B) { 88 | for _, tc := range cases { 89 | slice := random[int](tc.size) 90 | b.Run(tc.name, func(b *testing.B) { 91 | for i := 0; i < b.N; i++ { 92 | sort.Ints(slice) 93 | _ = slice[0] 94 | unsort(slice) 95 | } 96 | }) 97 | } 98 | } 99 | 100 | func BenchmarkSet_Minimum(b *testing.B) { 101 | for _, tc := range cases { 102 | s := From(random[int](tc.size)) 103 | b.Run(tc.name, func(b *testing.B) { 104 | for i := 0; i < b.N; i++ { 105 | values := s.Slice() 106 | sort.Ints(values) 107 | _ = values[0] 108 | } 109 | }) 110 | } 111 | } 112 | 113 | func BenchmarkHashSet_Minimum(b *testing.B) { 114 | for _, tc := range cases { 115 | hs := HashSetFrom[hashint, int](random[hashint](tc.size)) 116 | b.Run(tc.name, func(b *testing.B) { 117 | for i := 0; i < b.N; i++ { 118 | values := hs.Slice() 119 | sort.Slice(values, func(a, b int) bool { return values[a] < values[b] }) 120 | _ = values[0] 121 | unsort(values) 122 | } 123 | }) 124 | } 125 | } 126 | 127 | func BenchmarkTreeSet_Minimum(b *testing.B) { 128 | for _, tc := range cases { 129 | ts := TreeSetFrom[int](random[int](tc.size), cmp.Compare[int]) 130 | b.Run(tc.name, func(b *testing.B) { 131 | for i := 0; i < b.N; i++ { 132 | _ = ts.Min() 133 | } 134 | }) 135 | } 136 | } 137 | 138 | func BenchmarkSlice_Contains(b *testing.B) { 139 | contains := func(s []int, target int) bool { 140 | for i := 0; i < len(s); i++ { 141 | if s[i] == target { 142 | return true 143 | } 144 | } 145 | return false 146 | } 147 | 148 | for _, tc := range cases { 149 | slice := random[int](tc.size) 150 | b.Run(tc.name, func(b *testing.B) { 151 | for i := 0; i < b.N; i++ { 152 | _ = contains(slice, i) 153 | } 154 | }) 155 | } 156 | } 157 | 158 | func BenchmarkSet_Contains(b *testing.B) { 159 | for _, tc := range cases { 160 | s := From(random[int](tc.size)) 161 | b.Run(tc.name, func(b *testing.B) { 162 | for i := 0; i < b.N; i++ { 163 | _ = s.Contains(i) 164 | } 165 | }) 166 | } 167 | } 168 | 169 | func BenchmarkHashSet_Contains(b *testing.B) { 170 | for _, tc := range cases { 171 | hs := HashSetFrom[hashint, int](random[hashint](tc.size)) 172 | b.Run(tc.name, func(b *testing.B) { 173 | for i := 0; i < b.N; i++ { 174 | _ = hs.Contains(hashint(i)) 175 | } 176 | }) 177 | } 178 | } 179 | 180 | func BenchmarkTreeSet_Contains(b *testing.B) { 181 | for _, tc := range cases { 182 | ts := TreeSetFrom[int](random[int](tc.size), cmp.Compare[int]) 183 | b.Run(tc.name, func(b *testing.B) { 184 | for i := 0; i < b.N; i++ { 185 | _ = ts.Contains(i) 186 | } 187 | }) 188 | } 189 | } 190 | -------------------------------------------------------------------------------- /collection.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import "iter" 7 | 8 | // Collection is a minimal common interface that all sets implement. 9 | 10 | // Fundamental set operations and familiar utility methods are part of this 11 | // interface. Each of Set, HashSet, and TreeSet may also provide implementation 12 | // specific methods not part of this interface. 13 | type Collection[T any] interface { 14 | 15 | // Insert an element into the set. 16 | // 17 | // Returns true if the set is modified as a result. 18 | Insert(T) bool 19 | 20 | // InsertSlice will insert each element of a given slice into the set. 21 | // 22 | // Returns true if the set was modified as a result. 23 | InsertSlice([]T) bool 24 | 25 | // InsertSet will insert each element of a given Collection into the set. 26 | // 27 | // Returns true if the set was modified as a result. 28 | InsertSet(Collection[T]) bool 29 | 30 | // Remove will remove the given element from the set, if present. 31 | // 32 | // Returns true if the set was modified as a result of the operation. 33 | Remove(T) bool 34 | 35 | // RemoveSlice will remove each element of a slice from the set, if present. 36 | // 37 | // Returns true if the set was modified as a result of the operation. 38 | RemoveSlice([]T) bool 39 | 40 | // RemoveSet will remove each element of a Collection from the set. 41 | // 42 | // Returns true if the set was modified as a result of the operation. 43 | RemoveSet(Collection[T]) bool 44 | 45 | // RemoveFunc will remove each element from the set that satisfies the given predicate. 46 | // 47 | // Returns true if the set was modified as a result of the opearation. 48 | RemoveFunc(func(T) bool) bool 49 | 50 | // Contains returns whether an element is present in the set. 51 | Contains(T) bool 52 | 53 | // ContainsSlice returns whether the set contains the same set of elements as 54 | // the given slice. The elements of the slice may contain duplicates. 55 | ContainsSlice([]T) bool 56 | 57 | // Subset returns whether the given Collection is a subset of the set. 58 | Subset(Collection[T]) bool 59 | 60 | // ProperSubset returns whether the given Collection is a proper subset of the set. 61 | ProperSubset(Collection[T]) bool 62 | 63 | // Size returns the number of elements in the set. 64 | Size() int 65 | 66 | // Empty returns whether the set contains no elements. 67 | Empty() bool 68 | 69 | // Union returns a new set containing the unique elements of both this set 70 | // and a given Collection. 71 | // 72 | // https://en.wikipedia.org/wiki/Union_(set_theory) 73 | Union(Collection[T]) Collection[T] 74 | 75 | // Difference returns a new set that contains elements this set that are not 76 | // in a given Collection. 77 | // 78 | // https://en.wikipedia.org/wiki/Complement_(set_theory) 79 | Difference(Collection[T]) Collection[T] 80 | 81 | // Intersect returns a new set that contains only the elements present in 82 | // both this and a given Collection. 83 | // 84 | // https://en.wikipedia.org/wiki/Intersection_(set_theory) 85 | Intersect(Collection[T]) Collection[T] 86 | 87 | // Slice returns a slice of all elements in the set. 88 | // 89 | // For iterating elements, consider using Items() instead. 90 | // 91 | // Note: order of elements depends on the underlying implementation. 92 | Slice() []T 93 | 94 | // String creates a string representation of this set. 95 | // 96 | // Note: string representation depends on underlying implementation. 97 | String() string 98 | 99 | // StringFunc creates a string representation of this set, using the given 100 | // function to transform each element into a string. 101 | // 102 | // Note: string representation depends on underlying implementation. 103 | StringFunc(func(T) string) string 104 | 105 | // EqualSet returns whether this set and a given Collection contain the same 106 | // elements. 107 | EqualSet(Collection[T]) bool 108 | 109 | // EqualSlice returns whether this set and a given slice contain the same 110 | // elements, where the slice may contain duplicates. 111 | EqualSlice([]T) bool 112 | 113 | // EqualSliceSet returns whether this set and a given slice contain exactly 114 | // the same elements, where the slice must not contain duplicates. 115 | EqualSliceSet([]T) bool 116 | 117 | // Items returns a generator function for use with the range keyword 118 | // enabling iteration of each element in the set. 119 | // 120 | // Note: iteration order depends on the underlying implementation. 121 | // 122 | // for element := range s.Items() { ... } 123 | Items() iter.Seq[T] 124 | } 125 | 126 | // InsertSliceFunc inserts all elements from items into col, applying the transform 127 | // function to each element before insertion. 128 | // 129 | // Returns true if col was modified as a result of the operation. 130 | func InsertSliceFunc[T, E any](col Collection[T], items []E, transform func(element E) T) bool { 131 | modified := false 132 | for _, item := range items { 133 | if col.Insert(transform(item)) { 134 | modified = true 135 | } 136 | } 137 | return modified 138 | } 139 | 140 | // InsertSetFunc inserts the elements of a into b, applying the transform function 141 | // to each element before insertion. 142 | // 143 | // Returns true if b was modified as a result. 144 | func InsertSetFunc[T, E any](a Collection[T], b Collection[E], transform func(T) E) bool { 145 | modified := false 146 | for item := range a.Items() { 147 | if b.Insert(transform(item)) { 148 | modified = true 149 | } 150 | } 151 | return modified 152 | } 153 | 154 | // SliceFunc produces a slice of the elements in s, applying the transform 155 | // function to each element first. 156 | func SliceFunc[T, E any](s Collection[T], transform func(T) E) []E { 157 | slice := make([]E, 0, s.Size()) 158 | for item := range s.Items() { 159 | slice = append(slice, transform(item)) 160 | } 161 | return slice 162 | } 163 | 164 | func insert[T any](destination, col Collection[T]) { 165 | for item := range col.Items() { 166 | destination.Insert(item) 167 | } 168 | } 169 | 170 | func intersect[T any](destination, a, b Collection[T]) { 171 | var ( 172 | big Collection[T] = a 173 | small Collection[T] = b 174 | ) 175 | if a.Size() < b.Size() { 176 | big, small = b, a 177 | } 178 | for item := range small.Items() { 179 | if big.Contains(item) { 180 | destination.Insert(item) 181 | } 182 | } 183 | } 184 | 185 | func containsSlice[T any](col Collection[T], items []T) bool { 186 | for _, item := range items { 187 | if !col.Contains(item) { 188 | return false 189 | } 190 | } 191 | return true 192 | } 193 | 194 | func equalSet[T any](a, b Collection[T]) bool { 195 | // fast paths: sets are empty or different sizes 196 | sizeA, sizeB := a.Size(), b.Size() 197 | if sizeA == 0 && sizeB == 0 { 198 | return true 199 | } 200 | if sizeA != sizeB { 201 | return false 202 | } 203 | 204 | // look for any missing element 205 | for item := range a.Items() { 206 | if !b.Contains(item) { 207 | return false 208 | } 209 | } 210 | return true 211 | } 212 | 213 | func removeSet[T any](s, col Collection[T]) bool { 214 | modified := false 215 | for item := range col.Items() { 216 | if s.Remove(item) { 217 | modified = true 218 | } 219 | } 220 | return modified 221 | } 222 | 223 | func removeFunc[T any](col Collection[T], predicate func(T) bool) bool { 224 | remove := make([]T, 0) 225 | for item := range col.Items() { 226 | if predicate(item) { 227 | remove = append(remove, item) 228 | } 229 | } 230 | return col.RemoveSlice(remove) 231 | } 232 | 233 | func subset[T any](a, b Collection[T]) bool { 234 | if b.Size() > a.Size() { 235 | return false 236 | } 237 | 238 | for item := range b.Items() { 239 | if !a.Contains(item) { 240 | return false 241 | } 242 | } 243 | 244 | return true 245 | } 246 | -------------------------------------------------------------------------------- /collection_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "cmp" 8 | "sort" 9 | "strconv" 10 | "testing" 11 | 12 | "github.com/shoenig/test/must" 13 | ) 14 | 15 | func TestInsertSliceFunc(t *testing.T) { 16 | numbers := ints(3) 17 | 18 | t.Run("set", func(t *testing.T) { 19 | s := New[string](10) 20 | transform := func(element int) string { return strconv.Itoa(element) } 21 | InsertSliceFunc[string](s, numbers, transform) 22 | slices := s.Slice() 23 | sort.Strings(slices) 24 | must.SliceEqFunc(t, slices, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 25 | }) 26 | 27 | t.Run("hashset", func(t *testing.T) { 28 | s := NewHashSet[*company, string](10) 29 | InsertSliceFunc[*company](s, numbers, func(element int) *company { 30 | return &company{ 31 | address: "InsertSliceFunc", 32 | floor: element, 33 | } 34 | }) 35 | must.MapContainsKeys(t, s.items, []string{ 36 | "InsertSliceFunc:1", "InsertSliceFunc:2", "InsertSliceFunc:3", 37 | }) 38 | }) 39 | 40 | t.Run("treeSet", func(t *testing.T) { 41 | s := NewTreeSet[string](cmp.Compare[string]) 42 | InsertSliceFunc[string](s, numbers, func(element int) string { 43 | return strconv.Itoa(element) 44 | }) 45 | invariants(t, s, cmp.Compare[string]) 46 | must.SliceEqFunc(t, s.Slice(), []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 47 | }) 48 | } 49 | 50 | func TestSliceFunc(t *testing.T) { 51 | t.Run("set", func(t *testing.T) { 52 | s := From(ints(3)) 53 | slice := SliceFunc[int](s, func(element int) string { 54 | return strconv.Itoa(element) 55 | }) 56 | sort.Strings(slice) 57 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 58 | }) 59 | 60 | t.Run("hashset", func(t *testing.T) { 61 | s := NewHashSet[*company, string](10) 62 | s.InsertSlice([]*company{c1, c2, c3}) 63 | slice := SliceFunc[*company](s, func(element *company) string { 64 | return element.Hash() 65 | }) 66 | sort.Strings(slice) 67 | must.SliceEqFunc(t, slice, []string{"street:1", "street:2", "street:3"}, func(a, b string) bool { return a == b }) 68 | }) 69 | 70 | t.Run("treeSet", func(t *testing.T) { 71 | s := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 72 | slice := SliceFunc[int](s, func(element int) string { 73 | return strconv.Itoa(element) 74 | }) 75 | sort.Strings(slice) 76 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 77 | }) 78 | } 79 | 80 | func TestInsertSetFunc(t *testing.T) { 81 | t.Run("set", func(t *testing.T) { 82 | a := From(ints(3)) 83 | t.Run("set -> set", func(t *testing.T) { 84 | b := New[string](3) 85 | modified := InsertSetFunc[int, string](a, b, func(element int) string { 86 | return strconv.Itoa(element) 87 | }) 88 | must.True(t, modified) 89 | slice := b.Slice() 90 | sort.Strings(slice) 91 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 92 | }) 93 | 94 | t.Run("set -> hashset", func(t *testing.T) { 95 | b := NewHashSet[*company, string](10) 96 | modified := InsertSetFunc[int, *company](a, b, func(element int) *company { 97 | return &company{ 98 | address: "street", 99 | floor: element, 100 | } 101 | }) 102 | must.True(t, modified) 103 | must.MapContainsKeys(t, b.items, []string{ 104 | "street:1", "street:2", "street:3", 105 | }) 106 | }) 107 | 108 | t.Run("set -> treeSet", func(t *testing.T) { 109 | b := NewTreeSet[string](cmp.Compare[string]) 110 | modified := InsertSetFunc[int, string](a, b, func(element int) string { 111 | return strconv.Itoa(element) 112 | }) 113 | must.True(t, modified) 114 | slice := b.Slice() 115 | sort.Strings(slice) 116 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 117 | }) 118 | 119 | t.Run("not modified", func(t *testing.T) { 120 | b := a.Copy() 121 | modified := InsertSetFunc[int, int](a, b, func(element int) int { 122 | return element 123 | }) 124 | must.False(t, modified) 125 | }) 126 | }) 127 | 128 | t.Run("hashSet", func(t *testing.T) { 129 | a := NewHashSet[*company, string](10) 130 | a.InsertSlice([]*company{c1, c2, c3}) 131 | 132 | t.Run("hashSet -> set", func(t *testing.T) { 133 | b := New[int](3) 134 | modified := InsertSetFunc[*company, int](a, b, func(element *company) int { 135 | return element.floor 136 | }) 137 | must.True(t, modified) 138 | slice := b.Slice() 139 | sort.Ints(slice) 140 | must.SliceEqFunc(t, slice, []int{1, 2, 3}, func(a, b int) bool { return a == b }) 141 | }) 142 | 143 | t.Run("hashSet -> hashSet", func(t *testing.T) { 144 | b := NewHashSet[*company, string](10) 145 | modified := InsertSetFunc[*company, *company](a, b, func(element *company) *company { 146 | return &company{ 147 | address: element.address, 148 | floor: element.floor * 5, 149 | } 150 | }) 151 | must.True(t, modified) 152 | must.MapContainsKeys(t, b.items, []string{ 153 | "street:5", "street:10", "street:15", 154 | }) 155 | }) 156 | 157 | t.Run("hashSet -> treeSet", func(t *testing.T) { 158 | b := NewTreeSet[int](cmp.Compare[int]) 159 | modified := InsertSetFunc[*company, int](a, b, func(element *company) int { 160 | return element.floor 161 | }) 162 | must.True(t, modified) 163 | slice := b.Slice() 164 | sort.Ints(slice) 165 | must.SliceEqFunc(t, slice, []int{1, 2, 3}, func(a, b int) bool { return a == b }) 166 | }) 167 | 168 | t.Run("not modified", func(t *testing.T) { 169 | b := a.Copy() 170 | modified := InsertSetFunc[*company, *company](a, b, func(element *company) *company { 171 | return element 172 | }) 173 | must.False(t, modified) 174 | }) 175 | }) 176 | 177 | t.Run("treeSet", func(t *testing.T) { 178 | a := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 179 | 180 | t.Run("treeSet -> set", func(t *testing.T) { 181 | b := New[string](3) 182 | modified := InsertSetFunc[int, string](a, b, func(element int) string { 183 | return strconv.Itoa(element) 184 | }) 185 | must.True(t, modified) 186 | slice := b.Slice() 187 | sort.Strings(slice) 188 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 189 | }) 190 | 191 | t.Run("treeSet -> hashSet", func(t *testing.T) { 192 | b := NewHashSet[*company, string](10) 193 | modified := InsertSetFunc[int, *company](a, b, func(element int) *company { 194 | return &company{ 195 | address: "street", 196 | floor: element, 197 | } 198 | }) 199 | must.True(t, modified) 200 | must.MapContainsKeys(t, b.items, []string{ 201 | "street:1", "street:2", "street:3", 202 | }) 203 | }) 204 | 205 | t.Run("treeSet -> treeSet", func(t *testing.T) { 206 | b := NewTreeSet[string](cmp.Compare[string]) 207 | modified := InsertSetFunc[int, string](a, b, func(element int) string { 208 | return strconv.Itoa(element) 209 | }) 210 | must.True(t, modified) 211 | slice := b.Slice() 212 | sort.Strings(slice) 213 | must.SliceEqFunc(t, slice, []string{"1", "2", "3"}, func(a, b string) bool { return a == b }) 214 | }) 215 | 216 | t.Run("not modified", func(t *testing.T) { 217 | b := a.Copy() 218 | modified := InsertSetFunc[int, int](a, b, func(element int) int { 219 | return element 220 | }) 221 | must.False(t, modified) 222 | }) 223 | }) 224 | } 225 | 226 | func TestEqualSet(t *testing.T) { 227 | t.Run("equal ok", func(t *testing.T) { 228 | a := From(ints(3)) 229 | b := From(ints(3)) 230 | must.True(t, a.EqualSet(b)) 231 | }) 232 | t.Run("none equal none", func(t *testing.T) { 233 | a := New[int](0) 234 | b := New[int](0) 235 | must.True(t, a.EqualSet(b)) 236 | }) 237 | t.Run("size not equal", func(t *testing.T) { 238 | a := From(ints(3)) 239 | b := From(ints(4)) 240 | must.False(t, a.EqualSet(b)) 241 | }) 242 | t.Run("items not equal", func(t *testing.T) { 243 | a := From(ints(3)) 244 | b := From([]int{1, 2, 4}) 245 | must.False(t, a.EqualSet(b)) 246 | }) 247 | } 248 | -------------------------------------------------------------------------------- /examples_hashset_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "encoding/json" 8 | "fmt" 9 | "sort" 10 | ) 11 | 12 | type person struct { 13 | Name string 14 | ID int 15 | } 16 | 17 | func (p *person) Hash() string { 18 | return fmt.Sprintf("%s:%d", p.Name, p.ID) 19 | } 20 | 21 | func (p *person) String() string { 22 | return p.Name 23 | } 24 | 25 | func ExampleHashSet_Insert() { 26 | s := NewHashSet[*person, string](10) 27 | s.Insert(&person{Name: "dave", ID: 108}) 28 | s.Insert(&person{Name: "armon", ID: 101}) 29 | s.Insert(&person{Name: "mitchell", ID: 100}) 30 | s.Insert(&person{Name: "armon", ID: 101}) 31 | 32 | fmt.Println(s) 33 | 34 | // Output: 35 | // [armon dave mitchell] 36 | } 37 | 38 | func ExampleHashSet_InsertSlice() { 39 | s := NewHashSet[*person, string](10) 40 | s.InsertSlice([]*person{ 41 | {Name: "dave", ID: 108}, 42 | {Name: "mitchell", ID: 100}, 43 | {Name: "dave", ID: 108}, 44 | {Name: "armon", ID: 101}, 45 | }) 46 | 47 | fmt.Println(s) 48 | 49 | // Output: 50 | // [armon dave mitchell] 51 | } 52 | 53 | func ExampleHashSet_InsertSet() { 54 | anna := &person{Name: "anna", ID: 94} 55 | bill := &person{Name: "bill", ID: 50} 56 | carl := &person{Name: "carl", ID: 10} 57 | dave := &person{Name: "dave", ID: 32} 58 | s1 := HashSetFrom[*person, string]([]*person{anna, carl}) 59 | s2 := HashSetFrom[*person, string]([]*person{carl, dave, bill}) 60 | s2.InsertSet(s1) 61 | 62 | fmt.Println(s1) 63 | fmt.Println(s2) 64 | 65 | // Output: 66 | // [anna carl] 67 | // [anna bill carl dave] 68 | } 69 | 70 | func ExampleHashSet_Remove() { 71 | anna := &person{Name: "anna", ID: 94} 72 | bill := &person{Name: "bill", ID: 50} 73 | carl := &person{Name: "carl", ID: 10} 74 | dave := &person{Name: "dave", ID: 32} 75 | s := HashSetFrom[*person, string]([]*person{anna, carl, dave, bill}) 76 | 77 | fmt.Println(s) 78 | 79 | s.Remove(carl) 80 | 81 | fmt.Println(s) 82 | 83 | // Output: 84 | // [anna bill carl dave] 85 | // [anna bill dave] 86 | } 87 | 88 | func ExampleHashSet_RemoveSlice() { 89 | anna := &person{Name: "anna", ID: 94} 90 | bill := &person{Name: "bill", ID: 50} 91 | carl := &person{Name: "carl", ID: 10} 92 | dave := &person{Name: "dave", ID: 32} 93 | s := HashSetFrom[*person, string]([]*person{anna, carl, dave, bill}) 94 | 95 | fmt.Println(s) 96 | 97 | s.RemoveSlice([]*person{anna, carl}) 98 | 99 | fmt.Println(s) 100 | 101 | // Output: 102 | // [anna bill carl dave] 103 | // [bill dave] 104 | } 105 | 106 | func ExampleHashSet_RemoveSet() { 107 | anna := &person{Name: "anna", ID: 94} 108 | bill := &person{Name: "bill", ID: 50} 109 | carl := &person{Name: "carl", ID: 10} 110 | dave := &person{Name: "dave", ID: 32} 111 | s := HashSetFrom[*person, string]([]*person{carl, dave, bill}) 112 | r := HashSetFrom[*person, string]([]*person{anna, carl}) 113 | 114 | fmt.Println(s) 115 | 116 | s.RemoveSet(r) 117 | 118 | fmt.Println(s) 119 | 120 | // Output: 121 | // [bill carl dave] 122 | // [bill dave] 123 | } 124 | 125 | func ExampleHashSet_RemoveFunc() { 126 | anna := &person{Name: "anna", ID: 94} 127 | bill := &person{Name: "bill", ID: 50} 128 | carl := &person{Name: "carl", ID: 10} 129 | dave := &person{Name: "dave", ID: 32} 130 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl, dave}) 131 | 132 | idAbove50 := func(p *person) bool { 133 | return p.ID >= 50 134 | } 135 | 136 | fmt.Println(s) 137 | 138 | s.RemoveFunc(idAbove50) 139 | 140 | fmt.Println(s) 141 | 142 | // Output: 143 | // [anna bill carl dave] 144 | // [carl dave] 145 | } 146 | 147 | func ExampleHashSet_Contains() { 148 | anna := &person{Name: "anna", ID: 94} 149 | bill := &person{Name: "bill", ID: 50} 150 | carl := &person{Name: "carl", ID: 10} 151 | dave := &person{Name: "dave", ID: 32} 152 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 153 | 154 | fmt.Println(s.Contains(anna)) 155 | fmt.Println(s.Contains(dave)) 156 | 157 | // Output: 158 | // true 159 | // false 160 | } 161 | 162 | func ExampleHashSet_ContainsSlice() { 163 | anna := &person{Name: "anna", ID: 94} 164 | bill := &person{Name: "bill", ID: 50} 165 | carl := &person{Name: "carl", ID: 10} 166 | dave := &person{Name: "dave", ID: 32} 167 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 168 | 169 | fmt.Println(s.ContainsSlice([]*person{anna, bill})) 170 | fmt.Println(s.ContainsSlice([]*person{anna, bill, carl})) 171 | fmt.Println(s.ContainsSlice([]*person{carl, dave})) 172 | 173 | // Output: 174 | // false 175 | // true 176 | // false 177 | } 178 | 179 | func ExampleHashSet_Subset() { 180 | anna := &person{Name: "anna", ID: 94} 181 | bill := &person{Name: "bill", ID: 50} 182 | carl := &person{Name: "carl", ID: 10} 183 | dave := &person{Name: "dave", ID: 32} 184 | s1 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 185 | s2 := HashSetFrom[*person, string]([]*person{anna, bill}) 186 | s3 := HashSetFrom[*person, string]([]*person{bill, carl, dave}) 187 | 188 | fmt.Println(s1.Subset(s2)) 189 | fmt.Println(s1.Subset(s3)) 190 | 191 | // Output: 192 | // true 193 | // false 194 | } 195 | 196 | func ExampleHashSet_Size() { 197 | anna := &person{Name: "anna", ID: 94} 198 | bill := &person{Name: "bill", ID: 50} 199 | carl := &person{Name: "carl", ID: 10} 200 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 201 | 202 | fmt.Println(s.Size()) 203 | 204 | // Output: 205 | // 3 206 | } 207 | 208 | func ExampleHashSet_Empty() { 209 | s := NewHashSet[*person, string](0) 210 | 211 | fmt.Println(s.Empty()) 212 | 213 | // Output: 214 | // true 215 | } 216 | 217 | func ExampleHashSet_Union() { 218 | anna := &person{Name: "anna", ID: 94} 219 | bill := &person{Name: "bill", ID: 50} 220 | carl := &person{Name: "carl", ID: 10} 221 | dave := &person{Name: "dave", ID: 32} 222 | s1 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 223 | s2 := HashSetFrom[*person, string]([]*person{anna, bill, dave}) 224 | union := s1.Union(s2) 225 | 226 | fmt.Println(s1) 227 | fmt.Println(s2) 228 | fmt.Println(union) 229 | 230 | // Output: 231 | // [anna bill carl] 232 | // [anna bill dave] 233 | // [anna bill carl dave] 234 | } 235 | 236 | func ExampleHashSet_Difference() { 237 | anna := &person{Name: "anna", ID: 94} 238 | bill := &person{Name: "bill", ID: 50} 239 | carl := &person{Name: "carl", ID: 10} 240 | dave := &person{Name: "dave", ID: 32} 241 | s1 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 242 | s2 := HashSetFrom[*person, string]([]*person{anna, bill, dave}) 243 | difference := s1.Difference(s2) 244 | 245 | fmt.Println(s1) 246 | fmt.Println(s2) 247 | fmt.Println(difference) 248 | 249 | // Output: 250 | // [anna bill carl] 251 | // [anna bill dave] 252 | // [carl] 253 | } 254 | 255 | func ExampleHashSet_Intersect() { 256 | anna := &person{Name: "anna", ID: 94} 257 | bill := &person{Name: "bill", ID: 50} 258 | carl := &person{Name: "carl", ID: 10} 259 | dave := &person{Name: "dave", ID: 32} 260 | s1 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 261 | s2 := HashSetFrom[*person, string]([]*person{anna, bill, dave}) 262 | intersect := s1.Intersect(s2) 263 | 264 | fmt.Println(s1) 265 | fmt.Println(s2) 266 | fmt.Println(intersect) 267 | 268 | // Output: 269 | // [anna bill carl] 270 | // [anna bill dave] 271 | // [anna bill] 272 | } 273 | 274 | func ExampleHashSet_Equal() { 275 | anna := &person{Name: "anna", ID: 94} 276 | bill := &person{Name: "bill", ID: 50} 277 | carl := &person{Name: "carl", ID: 10} 278 | dave := &person{Name: "dave", ID: 32} 279 | 280 | s1 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 281 | s2 := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 282 | s3 := HashSetFrom[*person, string]([]*person{anna, bill, dave}) 283 | 284 | fmt.Println(s1.Equal(s2)) 285 | fmt.Println(s1.Equal(s3)) 286 | 287 | // Output: 288 | // true 289 | // false 290 | } 291 | 292 | func ExampleHashSet_EqualSlice() { 293 | anna := &person{Name: "anna", ID: 94} 294 | bill := &person{Name: "bill", ID: 50} 295 | carl := &person{Name: "carl", ID: 10} 296 | dave := &person{Name: "dave", ID: 32} 297 | 298 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 299 | 300 | fmt.Println(s.EqualSlice([]*person{bill, anna, carl})) 301 | fmt.Println(s.EqualSlice([]*person{anna, anna, bill, carl})) 302 | fmt.Println(s.EqualSlice([]*person{dave, bill, carl})) 303 | 304 | // Output: 305 | // true 306 | // true 307 | // false 308 | } 309 | 310 | func ExampleHashSet_EqualSliceSet() { 311 | anna := &person{Name: "anna", ID: 94} 312 | bill := &person{Name: "bill", ID: 50} 313 | carl := &person{Name: "carl", ID: 10} 314 | dave := &person{Name: "dave", ID: 32} 315 | 316 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 317 | 318 | fmt.Println(s.EqualSliceSet([]*person{bill, anna, carl})) 319 | fmt.Println(s.EqualSliceSet([]*person{anna, anna, bill, carl})) 320 | fmt.Println(s.EqualSliceSet([]*person{dave, bill, carl})) 321 | 322 | // Output: 323 | // true 324 | // false 325 | // false 326 | } 327 | 328 | func ExampleHashSet_Copy() { 329 | anna := &person{Name: "anna", ID: 94} 330 | bill := &person{Name: "bill", ID: 50} 331 | carl := &person{Name: "carl", ID: 10} 332 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 333 | c := s.Copy() 334 | 335 | fmt.Println(c) 336 | 337 | // Output: 338 | // [anna bill carl] 339 | } 340 | 341 | func ExampleHashSet_Slice() { 342 | anna := &person{Name: "anna", ID: 94} 343 | bill := &person{Name: "bill", ID: 50} 344 | carl := &person{Name: "carl", ID: 10} 345 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 346 | 347 | slice := s.Slice() 348 | sort.Slice(slice, func(a, b int) bool { 349 | return slice[a].ID < slice[b].ID 350 | }) 351 | 352 | fmt.Println(slice) 353 | 354 | // Output: 355 | // [carl bill anna] 356 | } 357 | 358 | func ExampleHashSet_String() { 359 | anna := &person{Name: "anna", ID: 94} 360 | bill := &person{Name: "bill", ID: 50} 361 | carl := &person{Name: "carl", ID: 10} 362 | s := HashSetFrom[*person, string]([]*person{anna, bill, carl}) 363 | 364 | fmt.Println(s.String()) 365 | 366 | // Output: 367 | // [anna bill carl] 368 | } 369 | 370 | // TODO: will not work as long as [HashFunc] cannot be derived from the type parameters. 371 | func ExampleHashSet_UnmarshalJSON() { 372 | // type Foo struct { 373 | // Persons *HashSet[*person, string] `json:"persons"` 374 | // } 375 | 376 | // in := `{"persons":[{"Name":"anna","ID":94},{"Name":"bill","ID":50},{"Name":"bill","ID":50},{"Name":"carl","ID":10}]}` 377 | // var out Foo 378 | 379 | // _ = json.Unmarshal([]byte(in), &out) 380 | 381 | // fmt.Println(out.Persons) 382 | 383 | // Output: 384 | } 385 | 386 | func ExampleHashSet_MarshalJSON() { 387 | type Foo struct { 388 | Persons *HashSet[*person, string] `json:"persons"` 389 | } 390 | 391 | f := Foo{ 392 | Persons: HashSetFrom([]*person{ 393 | {Name: "anna", ID: 94}, 394 | {Name: "bill", ID: 50}, 395 | {Name: "carl", ID: 10}, 396 | }), 397 | } 398 | 399 | b, _ := json.Marshal(f) 400 | 401 | fmt.Println(string(b)) 402 | 403 | // Output: 404 | // {"persons":[{"Name":"anna","ID":94},{"Name":"bill","ID":50},{"Name":"carl","ID":10}]} 405 | } 406 | -------------------------------------------------------------------------------- /examples_set_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "encoding/json" 8 | "fmt" 9 | "sort" 10 | ) 11 | 12 | func ExampleSet_Insert() { 13 | s := New[int](10) 14 | s.Insert(1) 15 | s.Insert(1) 16 | s.Insert(2) 17 | s.Insert(3) 18 | s.Insert(2) 19 | 20 | fmt.Println(s) 21 | 22 | // Output: 23 | // [1 2 3] 24 | } 25 | 26 | func ExampleSet_InsertSlice() { 27 | s := New[int](10) 28 | s.InsertSlice([]int{1, 1, 2, 3, 2}) 29 | 30 | fmt.Println(s) 31 | 32 | // Output: 33 | // [1 2 3] 34 | } 35 | 36 | func ExampleSet_InsertSet() { 37 | s := New[int](10) 38 | s.InsertSet(From([]int{1, 1, 2, 3, 2})) 39 | 40 | fmt.Println(s) 41 | 42 | // Output: 43 | // [1 2 3] 44 | } 45 | 46 | func ExampleSet_Remove() { 47 | s := New[int](10) 48 | s.InsertSlice([]int{1, 1, 2, 3, 2}) 49 | s.Remove(2) 50 | 51 | fmt.Println(s) 52 | 53 | // Output: 54 | // [1 3] 55 | } 56 | 57 | func ExampleSet_RemoveSlice() { 58 | s := New[int](10) 59 | s.InsertSlice([]int{1, 1, 2, 3, 2}) 60 | s.RemoveSlice([]int{2, 3}) 61 | 62 | fmt.Println(s) 63 | 64 | // Output: 65 | // [1] 66 | } 67 | 68 | func ExampleSet_RemoveSet() { 69 | s := New[int](10) 70 | s.InsertSlice([]int{1, 1, 2, 3, 2}) 71 | s.RemoveSet(From([]int{2, 3})) 72 | 73 | fmt.Println(s) 74 | 75 | // Output: 76 | // [1] 77 | } 78 | 79 | func ExampleSet_RemoveFunc() { 80 | s := From([]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}) 81 | even := func(i int) bool { 82 | return i%2 == 0 83 | } 84 | s.RemoveFunc(even) 85 | 86 | fmt.Println(s) 87 | 88 | // Output: 89 | // [1 3 5 7 9] 90 | } 91 | 92 | func ExampleSet_Contains() { 93 | s := From([]string{"red", "green", "blue"}) 94 | 95 | fmt.Println(s.Contains("red")) 96 | fmt.Println(s.Contains("orange")) 97 | 98 | // Output: 99 | // true 100 | // false 101 | } 102 | 103 | func ExampleSet_ContainsSlice() { 104 | s := From([]string{"red", "green", "blue"}) 105 | 106 | fmt.Println(s.ContainsSlice([]string{"red", "blue"})) 107 | fmt.Println(s.ContainsSlice([]string{"red", "blue", "orange"})) 108 | fmt.Println(s.ContainsSlice([]string{"red", "blue", "green"})) 109 | 110 | // Output: 111 | // true 112 | // false 113 | // true 114 | } 115 | 116 | func ExampleSet_Subset() { 117 | t1 := From([]string{"red", "green", "blue"}) 118 | t2 := From([]string{"red", "blue"}) 119 | t3 := From([]string{"red", "orange"}) 120 | 121 | fmt.Println(t1.Subset(t2)) 122 | fmt.Println(t1.Subset(t3)) 123 | 124 | // Output: 125 | // true 126 | // false 127 | } 128 | 129 | func ExampleSet_Size() { 130 | s := From([]string{"red", "green", "blue"}) 131 | 132 | fmt.Println(s.Size()) 133 | 134 | // Output: 135 | // 3 136 | } 137 | 138 | func ExampleSet_Empty() { 139 | s := New[string](10) 140 | 141 | fmt.Println(s.Empty()) 142 | 143 | // Output: 144 | // true 145 | } 146 | 147 | func ExampleSet_Union() { 148 | t1 := From([]string{"red", "green", "blue"}) 149 | t2 := From([]string{"red", "blue"}) 150 | t3 := From([]string{"red", "orange"}) 151 | 152 | fmt.Println(t1.Union(t2)) 153 | fmt.Println(t1.Union(t3)) 154 | 155 | // Output: 156 | // [blue green red] 157 | // [blue green orange red] 158 | } 159 | 160 | func ExampleSet_Difference() { 161 | t1 := From([]string{"red", "green", "blue"}) 162 | t2 := From([]string{"red", "blue"}) 163 | t3 := From([]string{"red", "orange"}) 164 | 165 | fmt.Println(t1.Difference(t2)) 166 | fmt.Println(t1.Difference(t3)) 167 | 168 | // Output: 169 | // [green] 170 | // [blue green] 171 | } 172 | 173 | func ExampleSet_Intersect() { 174 | t1 := From([]string{"red", "green", "blue"}) 175 | t2 := From([]string{"red", "blue"}) 176 | t3 := From([]string{"red", "orange"}) 177 | t4 := From([]string{"yellow"}) 178 | 179 | fmt.Println(t1.Intersect(t2)) 180 | fmt.Println(t1.Intersect(t3)) 181 | fmt.Println(t1.Intersect(t4)) 182 | 183 | // Output: 184 | // [blue red] 185 | // [red] 186 | // [] 187 | } 188 | 189 | func ExampleSet_Equal() { 190 | t1 := From([]string{"red", "green", "blue"}) 191 | t2 := From([]string{"red", "blue"}) 192 | t3 := From([]string{"red", "green", "yellow"}) 193 | t4 := From([]string{"red", "green", "blue"}) 194 | 195 | fmt.Println(t1.Equal(t2)) 196 | fmt.Println(t1.Equal(t3)) 197 | fmt.Println(t1.Equal(t4)) 198 | 199 | // Output: 200 | // false 201 | // false 202 | // true 203 | } 204 | 205 | func ExampleSet_EqualSlice() { 206 | s := From([]string{"red", "green", "blue"}) 207 | 208 | fmt.Println(s.EqualSlice([]string{"red", "blue", "green"})) 209 | fmt.Println(s.EqualSlice([]string{"red", "red", "blue", "green"})) 210 | fmt.Println(s.EqualSlice([]string{"yellow", "blue", "green"})) 211 | 212 | // Output: 213 | // true 214 | // true 215 | // false 216 | } 217 | 218 | func ExampleSet_EqualSliceSet() { 219 | s := From([]string{"red", "green", "blue"}) 220 | 221 | fmt.Println(s.EqualSliceSet([]string{"red", "blue", "green"})) 222 | fmt.Println(s.EqualSliceSet([]string{"red", "red", "blue", "green"})) 223 | fmt.Println(s.EqualSliceSet([]string{"yellow", "blue", "green"})) 224 | 225 | // Output: 226 | // true 227 | // false 228 | // false 229 | } 230 | 231 | func ExampleSet_Copy() { 232 | s := From([]string{"red", "green", "blue"}) 233 | t := s.Copy() 234 | 235 | fmt.Println(t) 236 | 237 | // Output: 238 | // [blue green red] 239 | } 240 | 241 | func ExampleSet_Slice() { 242 | s := From([]string{"red", "green", "blue"}) 243 | t := s.Slice() 244 | 245 | sort.Strings(t) 246 | fmt.Println(t) 247 | 248 | // Output: 249 | // [blue green red] 250 | } 251 | 252 | func ExampleSet_String() { 253 | s := From([]string{"red", "green", "blue"}) 254 | 255 | fmt.Println(s.String()) 256 | 257 | // Output: 258 | // [blue green red] 259 | } 260 | 261 | func ExampleSet_UnmarshalJSON() { 262 | type Foo struct { 263 | Colors *Set[string] `json:"colors"` 264 | } 265 | 266 | in := `{"colors":["red","green","green","blue"]}` 267 | var out Foo 268 | 269 | _ = json.Unmarshal([]byte(in), &out) 270 | 271 | fmt.Println(out.Colors) 272 | 273 | // Output: 274 | // [blue green red] 275 | } 276 | 277 | func ExampleSet_MarshalJSON() { 278 | type Foo struct { 279 | Colors *Set[string] `json:"colors"` 280 | } 281 | 282 | f := Foo{ 283 | Colors: From([]string{"red", "green", "blue"}), 284 | } 285 | 286 | b, _ := json.Marshal(f) 287 | 288 | fmt.Println(string(b)) 289 | 290 | // Output: 291 | // {"colors":["red","green","blue"]} 292 | } 293 | -------------------------------------------------------------------------------- /examples_treeset_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "cmp" 8 | "encoding/json" 9 | "fmt" 10 | ) 11 | 12 | func ExampleCompare_contestant() { 13 | type contestant struct { 14 | name string 15 | score int 16 | } 17 | 18 | compare := func(a, b contestant) int { 19 | return a.score - b.score 20 | } 21 | 22 | s := NewTreeSet[contestant](compare) 23 | s.Insert(contestant{name: "alice", score: 80}) 24 | s.Insert(contestant{name: "dave", score: 90}) 25 | s.Insert(contestant{name: "bob", score: 70}) 26 | 27 | fmt.Println(s) 28 | 29 | // Output: 30 | // [{bob 70} {alice 80} {dave 90}] 31 | } 32 | 33 | func ExampleCmp_strings() { 34 | s := NewTreeSet[string](cmp.Compare[string]) 35 | s.Insert("red") 36 | s.Insert("green") 37 | s.Insert("blue") 38 | 39 | fmt.Println(s) 40 | fmt.Println("min:", s.Min()) 41 | fmt.Println("max:", s.Max()) 42 | 43 | // Output: 44 | // [blue green red] 45 | // min: blue 46 | // max: red 47 | } 48 | 49 | func ExampleCmp_ints() { 50 | s := NewTreeSet[int](cmp.Compare[int]) 51 | s.Insert(50) 52 | s.Insert(42) 53 | s.Insert(100) 54 | 55 | fmt.Println(s) 56 | fmt.Println("min:", s.Min()) 57 | fmt.Println("max:", s.Max()) 58 | 59 | // Output: 60 | // [42 50 100] 61 | // min: 42 62 | // max: 100 63 | } 64 | 65 | func ExampleTreeSet_Insert() { 66 | s := TreeSetFrom[string]([]string{}, cmp.Compare[string]) 67 | 68 | fmt.Println(s) 69 | 70 | s.Insert("red") 71 | s.Insert("green") 72 | s.Insert("blue") 73 | 74 | fmt.Println(s) 75 | 76 | // [] 77 | // [blue green red] 78 | } 79 | 80 | func ExampleTreeSet_InsertSlice() { 81 | s := TreeSetFrom[string]([]string{}, cmp.Compare[string]) 82 | 83 | fmt.Println(s) 84 | 85 | s.InsertSlice([]string{"red", "green", "blue"}) 86 | 87 | fmt.Println(s) 88 | 89 | // [] 90 | // [blue green red] 91 | } 92 | 93 | func ExampleTreeSet_InsertSet() { 94 | s1 := TreeSetFrom[string]([]string{"red", "green"}, cmp.Compare[string]) 95 | s2 := TreeSetFrom[string]([]string{"green", "blue"}, cmp.Compare[string]) 96 | 97 | fmt.Println(s1) 98 | fmt.Println(s2) 99 | 100 | s1.InsertSet(s2) 101 | 102 | fmt.Println(s1) 103 | 104 | // Output: 105 | // [green red] 106 | // [blue green] 107 | // [blue green red] 108 | } 109 | 110 | func ExampleTreeSet_Remove() { 111 | s := TreeSetFrom[string]([]string{"red", "green", "blue"}, cmp.Compare[string]) 112 | 113 | fmt.Println(s) 114 | 115 | fmt.Println(s.Remove("green")) 116 | fmt.Println(s.Remove("orange")) 117 | 118 | fmt.Println(s) 119 | 120 | // Output: 121 | // [blue green red] 122 | // true 123 | // false 124 | // [blue red] 125 | } 126 | 127 | func ExampleTreeSet_RemoveSlice() { 128 | s := TreeSetFrom[string]([]string{"red", "green", "blue"}, cmp.Compare[string]) 129 | 130 | fmt.Println(s) 131 | 132 | fmt.Println(s.RemoveSlice([]string{"red", "blue"})) 133 | fmt.Println(s.RemoveSlice([]string{"orange", "white"})) 134 | 135 | fmt.Println(s) 136 | 137 | // Output: 138 | // [blue green red] 139 | // true 140 | // false 141 | // [green] 142 | } 143 | 144 | func ExampleTreeSet_RemoveSet() { 145 | s1 := TreeSetFrom[string]([]string{"a", "b", "c", "d", "e", "f"}, cmp.Compare[string]) 146 | s2 := TreeSetFrom[string]([]string{"e", "z", "a"}, cmp.Compare[string]) 147 | 148 | fmt.Println(s1) 149 | fmt.Println(s2) 150 | 151 | s1.RemoveSet(s2) 152 | 153 | fmt.Println(s1) 154 | 155 | // Output: 156 | // [a b c d e f] 157 | // [a e z] 158 | // [b c d f] 159 | } 160 | 161 | func ExampleTreeSet_RemoveFunc() { 162 | s := TreeSetFrom[int](ints(20), cmp.Compare[int]) 163 | 164 | fmt.Println(s) 165 | 166 | even := func(i int) bool { 167 | return i%3 != 0 168 | } 169 | s.RemoveFunc(even) 170 | 171 | fmt.Println(s) 172 | 173 | // Output: 174 | // [1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20] 175 | // [3 6 9 12 15 18] 176 | } 177 | 178 | func ExampleTreeSet_Contains() { 179 | s := TreeSetFrom[string]([]string{"red", "green", "blue"}, cmp.Compare[string]) 180 | 181 | fmt.Println(s.Contains("green")) 182 | fmt.Println(s.Contains("orange")) 183 | 184 | // Output: 185 | // true 186 | // false 187 | } 188 | 189 | func ExampleTreeSet_ContainsSlice() { 190 | s := TreeSetFrom[string]([]string{"red", "green", "blue"}, cmp.Compare[string]) 191 | 192 | fmt.Println(s.ContainsSlice([]string{"red", "green"})) 193 | fmt.Println(s.ContainsSlice([]string{"red", "orange"})) 194 | 195 | // Output: 196 | // true 197 | // false 198 | } 199 | 200 | func ExampleTreeSet_Subset() { 201 | s1 := TreeSetFrom[string]([]string{"a", "b", "c", "d", "e"}, cmp.Compare[string]) 202 | s2 := TreeSetFrom[string]([]string{"b", "d"}, cmp.Compare[string]) 203 | s3 := TreeSetFrom[string]([]string{"a", "z"}, cmp.Compare[string]) 204 | 205 | fmt.Println(s1.Subset(s2)) 206 | fmt.Println(s1.Subset(s3)) 207 | 208 | // Output: 209 | // true 210 | // false 211 | } 212 | 213 | func ExampleTreeSet_Size() { 214 | s := TreeSetFrom[string]([]string{"red", "green", "blue"}, cmp.Compare[string]) 215 | 216 | fmt.Println(s.Size()) 217 | 218 | // Output: 219 | // 3 220 | } 221 | 222 | func ExampleTreeSet_Empty() { 223 | s := TreeSetFrom[string]([]string{}, cmp.Compare[string]) 224 | 225 | fmt.Println(s.Empty()) 226 | 227 | s.InsertSlice([]string{"red", "green", "blue"}) 228 | 229 | fmt.Println(s.Empty()) 230 | 231 | // Output: 232 | // true 233 | // false 234 | } 235 | 236 | func ExampleTreeSet_Union() { 237 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 238 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 239 | f := TreeSetFrom[int]([]int{1, 3, 5, 7, 9}, cmp.Compare[int]) 240 | 241 | fmt.Println(s.Union(t)) 242 | fmt.Println(s.Union(f)) 243 | 244 | // Output: 245 | // [1 2 3 4 5] 246 | // [1 2 3 4 5 7 9] 247 | } 248 | 249 | func ExampleTreeSet_Difference() { 250 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 251 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 252 | f := TreeSetFrom[int]([]int{1, 3, 5, 7, 9}, cmp.Compare[int]) 253 | 254 | fmt.Println(s.Difference(t)) 255 | fmt.Println(s.Difference(f)) 256 | 257 | // Output: 258 | // [] 259 | // [2 4] 260 | } 261 | 262 | func ExampleTreeSet_Intersect() { 263 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 264 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 265 | f := TreeSetFrom[int]([]int{1, 3, 5, 7, 9}, cmp.Compare[int]) 266 | 267 | fmt.Println(s.Intersect(t)) 268 | fmt.Println(s.Intersect(f)) 269 | 270 | // Output: 271 | // [1 2 3 4 5] 272 | // [1 3 5] 273 | } 274 | 275 | func ExampleTreeSet_Equal() { 276 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 277 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 278 | f := TreeSetFrom[int]([]int{1, 3, 5, 7, 9}, cmp.Compare[int]) 279 | 280 | fmt.Println(s.Equal(t)) 281 | fmt.Println(s.Equal(f)) 282 | 283 | // Output: 284 | // true 285 | // false 286 | } 287 | 288 | func ExampleTreeSet_EqualSlice() { 289 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 290 | 291 | fmt.Println(t.EqualSlice([]int{1, 2, 3, 4, 5})) 292 | fmt.Println(t.EqualSlice([]int{1, 1, 2, 3, 4, 5})) 293 | fmt.Println(t.EqualSlice([]int{0, 2, 3, 4, 5})) 294 | 295 | // Output: 296 | // true 297 | // true 298 | // false 299 | } 300 | 301 | func ExampleTreeSet_EqualSliceSet() { 302 | t := TreeSetFrom[int]([]int{5, 4, 3, 2, 1}, cmp.Compare[int]) 303 | 304 | fmt.Println(t.EqualSliceSet([]int{1, 2, 3, 4, 5})) 305 | fmt.Println(t.EqualSliceSet([]int{1, 1, 2, 3, 4, 5})) 306 | fmt.Println(t.EqualSliceSet([]int{0, 2, 3, 4, 5})) 307 | 308 | // Output: 309 | // true 310 | // false 311 | // false 312 | } 313 | 314 | func ExampleTreeSet_Copy() { 315 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 316 | c := s.Copy() 317 | s.Remove(2) 318 | s.Remove(4) 319 | 320 | fmt.Println(s) 321 | fmt.Println(c) 322 | 323 | // Output: 324 | // [1 3 5] 325 | // [1 2 3 4 5] 326 | } 327 | 328 | func ExampleTreeSet_Slice() { 329 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 330 | slice := s.Slice() 331 | 332 | fmt.Println(slice) 333 | fmt.Println(len(slice)) 334 | 335 | // Output: 336 | // [1 2 3 4 5] 337 | // 5 338 | } 339 | 340 | func ExampleTreeSet_String() { 341 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 342 | 343 | fmt.Println(s.String() == "[1 2 3 4 5]") 344 | 345 | // Output: 346 | // true 347 | } 348 | 349 | func ExampleTreeSet_Min() { 350 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 351 | r := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, func(a int, b int) int { 352 | return b - a 353 | }) 354 | 355 | fmt.Println("asc:", s.Min()) 356 | fmt.Println("desc:", r.Min()) 357 | 358 | // Output: 359 | // asc: 1 360 | // desc: 5 361 | } 362 | 363 | func ExampleTreeSet_Max() { 364 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 365 | r := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, func(a int, b int) int { 366 | return b - a 367 | }) 368 | 369 | fmt.Println("asc:", s.Max()) 370 | fmt.Println("desc:", r.Max()) 371 | 372 | // Output: 373 | // asc: 5 374 | // desc: 1 375 | } 376 | 377 | func ExampleTreeSet_TopK() { 378 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 379 | 380 | fmt.Println(s.TopK(0)) 381 | fmt.Println(s.TopK(1)) 382 | fmt.Println(s.TopK(3)) 383 | fmt.Println(s.TopK(5)) 384 | 385 | // Output: 386 | // [] 387 | // [1] 388 | // [1 2 3] 389 | // [1 2 3 4 5] 390 | } 391 | 392 | func ExampleTreeSet_BottomK() { 393 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 394 | 395 | fmt.Println(s.BottomK(0)) 396 | fmt.Println(s.BottomK(1)) 397 | fmt.Println(s.BottomK(3)) 398 | fmt.Println(s.BottomK(5)) 399 | 400 | // Output: 401 | // [] 402 | // [5] 403 | // [5 4 3] 404 | // [5 4 3 2 1] 405 | } 406 | 407 | func ExampleTreeSet_FirstAbove() { 408 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 409 | 410 | fmt.Println(s.FirstAbove(3)) 411 | fmt.Println(s.FirstAbove(5)) 412 | fmt.Println(s.FirstAbove(10)) 413 | 414 | // Output: 415 | // 4 true 416 | // 0 false 417 | // 0 false 418 | } 419 | 420 | func ExampleTreeSet_FirstAboveEqual() { 421 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 422 | 423 | fmt.Println(s.FirstAboveEqual(3)) 424 | fmt.Println(s.FirstAboveEqual(5)) 425 | fmt.Println(s.FirstAboveEqual(10)) 426 | 427 | // Output: 428 | // 3 true 429 | // 5 true 430 | // 0 false 431 | } 432 | 433 | func ExampleTreeSet_Above() { 434 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 435 | 436 | fmt.Println(s.Above(3)) 437 | fmt.Println(s.Above(5)) 438 | fmt.Println(s.Above(10)) 439 | 440 | // Output: 441 | // [4 5] 442 | // [] 443 | // [] 444 | } 445 | 446 | func ExampleTreeSet_AboveEqual() { 447 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 448 | 449 | fmt.Println(s.AboveEqual(3)) 450 | fmt.Println(s.AboveEqual(5)) 451 | fmt.Println(s.AboveEqual(10)) 452 | 453 | // Output: 454 | // [3 4 5] 455 | // [5] 456 | // [] 457 | } 458 | 459 | func ExampleTreeSet_FirstBelow() { 460 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 461 | 462 | fmt.Println(s.FirstBelow(1)) 463 | fmt.Println(s.FirstBelow(3)) 464 | fmt.Println(s.FirstBelow(10)) 465 | 466 | // Output: 467 | // 0 false 468 | // 2 true 469 | // 5 true 470 | } 471 | 472 | func ExampleTreeSet_FirstBelowEqual() { 473 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 474 | 475 | fmt.Println(s.FirstBelowEqual(1)) 476 | fmt.Println(s.FirstBelowEqual(3)) 477 | fmt.Println(s.FirstBelowEqual(10)) 478 | 479 | // Output: 480 | // 1 true 481 | // 3 true 482 | // 5 true 483 | } 484 | 485 | func ExampleTreeSet_Below() { 486 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 487 | 488 | fmt.Println(s.Below(1)) 489 | fmt.Println(s.Below(3)) 490 | fmt.Println(s.Below(10)) 491 | 492 | // Output: 493 | // [] 494 | // [1 2] 495 | // [1 2 3 4 5] 496 | } 497 | 498 | func ExampleTreeSet_BelowEqual() { 499 | s := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 500 | 501 | fmt.Println(s.BelowEqual(1)) 502 | fmt.Println(s.BelowEqual(3)) 503 | fmt.Println(s.BelowEqual(10)) 504 | 505 | // Output: 506 | // [1] 507 | // [1 2 3] 508 | // [1 2 3 4 5] 509 | } 510 | 511 | // TODO: will not work as long as [CompareFunc] cannot be derived from the type parameters. 512 | func ExampleTreeSet_UnmarshalJSON() { 513 | // type Foo struct { 514 | // Colors *TreeSet[string] `json:"colors"` 515 | // } 516 | 517 | // in := `{"colors":["red","green","green","blue"]}` 518 | // var out Foo 519 | 520 | // _ = json.Unmarshal([]byte(in), &out) 521 | 522 | // fmt.Println(out.Colors) 523 | 524 | // Output: 525 | } 526 | 527 | func ExampleTreeSet_MarshalJSON() { 528 | type Foo struct { 529 | Colors *TreeSet[string] `json:"colors"` 530 | } 531 | 532 | f := Foo{ 533 | Colors: TreeSetFrom([]string{"red", "green", "blue"}, cmp.Compare[string]), 534 | } 535 | 536 | b, _ := json.Marshal(f) 537 | 538 | fmt.Println(string(b)) 539 | 540 | // Output: 541 | // {"colors":["blue","green","red"]} 542 | } 543 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/hashicorp/go-set/v3 2 | 3 | go 1.23 4 | 5 | require github.com/shoenig/test v1.12.1 6 | 7 | require github.com/google/go-cmp v0.6.0 // indirect 8 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= 2 | github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= 3 | github.com/shoenig/test v1.12.1 h1:mLHfnMv7gmhhP44WrvT+nKSxKkPDiNkIuHGdIGI9RLU= 4 | github.com/shoenig/test v1.12.1/go.mod h1:UxJ6u/x2v/TNs/LoLxBNJRV9DiwBBKYxXSyczsBHFoI= 5 | -------------------------------------------------------------------------------- /hashset.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "fmt" 8 | "iter" 9 | "sort" 10 | ) 11 | 12 | // Hash represents the output type of a Hash() function defined on a type. 13 | // 14 | // A Hash could be string-like or int-like. A string hash could be something like 15 | // and md5, sha1, or GoString() representation of a type. An int hash could be 16 | // something like the prime multiple hash code of a type. 17 | type Hash interface { 18 | ~string | ~int | ~uint | ~int64 | ~uint64 | ~int32 | ~uint32 | ~int16 | ~uint16 | ~int8 | ~uint8 19 | } 20 | 21 | // Hasher represents a type that implements a Hash() method. Types that wish to 22 | // cache a hash value with an internal field should implement Hash accordingly. 23 | type Hasher[H Hash] interface { 24 | Hash() H 25 | } 26 | 27 | // HasherFunc creates a closure around the T.Hash function so that the type can 28 | // be used as the HashFunc for a HashSet. 29 | func HasherFunc[T Hasher[H], H Hash]() HashFunc[T, H] { 30 | return func(t T) H { 31 | return t.Hash() 32 | } 33 | } 34 | 35 | // HashFunc represents a function that that produces a hash value when applied 36 | // to a given T. Typically this will be implemented as T.Hash but by separating 37 | // HashFunc a HashSet can be made to make use of any hash implementation. 38 | type HashFunc[T any, H Hash] func(T) H 39 | 40 | // HashSet is a generic implementation of the mathematical data structure, oriented 41 | // around the use of a HashFunc to make hash values from other types. 42 | type HashSet[T any, H Hash] struct { 43 | fn HashFunc[T, H] 44 | items map[H]T 45 | } 46 | 47 | // NewHashSet creates a HashSet with underlying capacity of size and will compute 48 | // hash values from the T.Hash method. 49 | func NewHashSet[T Hasher[H], H Hash](size int) *HashSet[T, H] { 50 | return NewHashSetFunc[T, H](size, HasherFunc[T, H]()) 51 | } 52 | 53 | // NewHashSetFunc creates a HashSet with underlying capacity of size and uses 54 | // the given hashing function to compute hashes on elements. 55 | // 56 | // A HashSet will automatically grow or shrink its capacity as items are added 57 | // or removed. 58 | func NewHashSetFunc[T any, H Hash](size int, fn HashFunc[T, H]) *HashSet[T, H] { 59 | return &HashSet[T, H]{ 60 | fn: fn, 61 | items: make(map[H]T, max(0, size)), 62 | } 63 | } 64 | 65 | // HashSetFrom creates a new HashSet containing each element in items. 66 | // 67 | // T must implement HashFunc[H], where H is of type Hash. This allows custom types 68 | // that include non-comparable fields to provide their own hash algorithm. 69 | func HashSetFrom[T Hasher[H], H Hash](items []T) *HashSet[T, H] { 70 | s := NewHashSet[T, H](len(items)) 71 | s.InsertSlice(items) 72 | return s 73 | } 74 | 75 | // NewHashSetFromFunc creates a new HashSet containing each element in items. 76 | func HashSetFromFunc[T any, H Hash](items []T, hash HashFunc[T, H]) *HashSet[T, H] { 77 | s := NewHashSetFunc[T, H](len(items), hash) 78 | s.InsertSlice(items) 79 | return s 80 | } 81 | 82 | // Insert item into s. 83 | // 84 | // Return true if s was modified (item was not already in s), false otherwise. 85 | func (s *HashSet[T, H]) Insert(item T) bool { 86 | key := s.fn(item) 87 | if _, exists := s.items[key]; exists { 88 | return false 89 | } 90 | s.items[key] = item 91 | return true 92 | } 93 | 94 | // InsertSlice will insert each item in items into s. 95 | // 96 | // Return true if s was modified (at least one item was not already in s), false otherwise. 97 | func (s *HashSet[T, H]) InsertSlice(items []T) bool { 98 | modified := false 99 | for _, item := range items { 100 | if s.Insert(item) { 101 | modified = true 102 | } 103 | } 104 | return modified 105 | } 106 | 107 | // InsertSet will insert each element of col into s. 108 | // 109 | // Return true if s was modified (at least one item of col was not already in s), false otherwise. 110 | func (s *HashSet[T, H]) InsertSet(col Collection[T]) bool { 111 | modified := false 112 | for item := range col.Items() { 113 | if s.Insert(item) { 114 | modified = true 115 | } 116 | } 117 | return modified 118 | } 119 | 120 | // Remove will remove item from s. 121 | // 122 | // Return true if s was modified (item was present), false otherwise. 123 | func (s *HashSet[T, H]) Remove(item T) bool { 124 | key := s.fn(item) 125 | if _, exists := s.items[key]; !exists { 126 | return false 127 | } 128 | delete(s.items, key) 129 | return true 130 | } 131 | 132 | // RemoveSlice will remove each item in items from s. 133 | // 134 | // Return true if s was modified (any item was present), false otherwise. 135 | func (s *HashSet[T, H]) RemoveSlice(items []T) bool { 136 | modified := false 137 | for _, item := range items { 138 | if s.Remove(item) { 139 | modified = true 140 | } 141 | } 142 | return modified 143 | } 144 | 145 | // RemoveSet will remove each element of col from s. 146 | // 147 | // Return true if s was modified (any item of col was present in s), false otherwise. 148 | func (s *HashSet[T, H]) RemoveSet(col Collection[T]) bool { 149 | return removeSet(s, col) 150 | } 151 | 152 | // RemoveFunc will remove each element from s that satisfies condition f. 153 | // 154 | // Return true if s was modified, false otherwise. 155 | func (s *HashSet[T, H]) RemoveFunc(f func(item T) bool) bool { 156 | return removeFunc(s, f) 157 | } 158 | 159 | // Contains returns whether item is present in s. 160 | func (s *HashSet[T, H]) Contains(item T) bool { 161 | hash := s.fn(item) 162 | _, exists := s.items[hash] 163 | return exists 164 | } 165 | 166 | // ContainsSlice returns whether s contains the same set of of elements 167 | // that are in items. The elements of items may contain duplicates. 168 | // 169 | // If the slice is known to be set-like (no duplicates), EqualSlice provides 170 | // a more efficient implementation. 171 | func (s *HashSet[T, H]) ContainsSlice(items []T) bool { 172 | return s.Equal(HashSetFromFunc[T, H](items, s.fn)) 173 | } 174 | 175 | // Subset returns whether col is a subset of s. 176 | func (s *HashSet[T, H]) Subset(col Collection[T]) bool { 177 | return subset(s, col) 178 | } 179 | 180 | // ProperSubset returns whether col is a proper subset of s. 181 | func (s *HashSet[T, H]) ProperSubset(col Collection[T]) bool { 182 | if len(s.items) <= col.Size() { 183 | return false 184 | } 185 | return s.Subset(col) 186 | } 187 | 188 | // Size returns the cardinality of s. 189 | func (s *HashSet[T, H]) Size() int { 190 | return len(s.items) 191 | } 192 | 193 | // Empty returns true if s contains no elements, false otherwise. 194 | func (s *HashSet[T, H]) Empty() bool { 195 | return s.Size() == 0 196 | } 197 | 198 | // Union returns a set that contains all elements of s and col combined. 199 | // 200 | // Elements in s take priority in the event of colliding hash values. 201 | func (s *HashSet[T, H]) Union(col Collection[T]) Collection[T] { 202 | result := NewHashSetFunc[T, H](s.Size(), s.fn) 203 | insert(result, s) 204 | insert(result, col) 205 | return result 206 | } 207 | 208 | // Difference returns a set that contains elements of s that are not in col. 209 | func (s *HashSet[T, H]) Difference(col Collection[T]) Collection[T] { 210 | result := NewHashSetFunc[T, H](max(0, s.Size()-col.Size()), s.fn) 211 | for item := range s.Items() { 212 | if !col.Contains(item) { 213 | result.Insert(item) 214 | } 215 | } 216 | return result 217 | } 218 | 219 | // Intersect returns a set that contains elements that are present in both s and col. 220 | func (s *HashSet[T, H]) Intersect(col Collection[T]) Collection[T] { 221 | result := NewHashSetFunc[T, H](0, s.fn) 222 | intersect(result, s, col) 223 | return result 224 | } 225 | 226 | // Copy creates a shallow copy of s. 227 | func (s *HashSet[T, H]) Copy() *HashSet[T, H] { 228 | result := NewHashSetFunc[T, H](s.Size(), s.fn) 229 | for key, item := range s.items { 230 | result.items[key] = item 231 | } 232 | return result 233 | } 234 | 235 | // Slice creates a copy of s as a slice. 236 | // 237 | // The result is not ordered. 238 | func (s *HashSet[T, H]) Slice() []T { 239 | result := make([]T, 0, s.Size()) 240 | for _, item := range s.items { 241 | result = append(result, item) 242 | } 243 | return result 244 | } 245 | 246 | // String creates a string representation of s, using "%v" printf formatting to transform 247 | // each element into a string. The result contains elements sorted by their lexical 248 | // string order. 249 | func (s *HashSet[T, H]) String() string { 250 | return s.StringFunc(func(element T) string { 251 | return fmt.Sprintf("%v", element) 252 | }) 253 | } 254 | 255 | // StringFunc creates a string representation of s, using f to transform each element 256 | // into a string. The result contains elements sorted by their string order. 257 | func (s *HashSet[T, H]) StringFunc(f func(element T) string) string { 258 | l := make([]string, 0, s.Size()) 259 | for _, item := range s.items { 260 | l = append(l, f(item)) 261 | } 262 | sort.Strings(l) 263 | return fmt.Sprintf("%s", l) 264 | } 265 | 266 | // Equal returns whether s and o contain the same elements. 267 | func (s *HashSet[T, H]) Equal(o *HashSet[T, H]) bool { 268 | if len(s.items) != len(o.items) { 269 | return false 270 | } 271 | for _, item := range s.items { 272 | if !o.Contains(item) { 273 | return false 274 | } 275 | } 276 | return true 277 | } 278 | 279 | // EqualSet returns whether s and col contain the same elements. 280 | func (s *HashSet[T, H]) EqualSet(col Collection[T]) bool { 281 | return equalSet(s, col) 282 | } 283 | 284 | // EqualSlice returns whether s and items contain the same elements. 285 | // 286 | // The items slice may contain duplicates. 287 | // 288 | // If the items slice is known to contain no duplicates, EqualSliceSet can be 289 | // used instead as a faster implementation. 290 | // 291 | // To detect if a slice is a subset of s, use ContainsSlice. 292 | func (s *HashSet[T, H]) EqualSlice(items []T) bool { 293 | other := HashSetFromFunc[T, H](items, s.fn) 294 | return s.Equal(other) 295 | } 296 | 297 | // EqualSliceSet returns whether s and items contain exactly the same elements. 298 | // 299 | // If items contains duplicates EqualSliceSet will return false. The elements of 300 | // items are assumed to be set-like. For comparing s to a slice that may contain 301 | // duplicate elements, use EqualSlice instead. 302 | // 303 | // To detect if a slice is a subset of s, use ContainsSlice. 304 | func (s *HashSet[T, H]) EqualSliceSet(items []T) bool { 305 | if len(items) != s.Size() { 306 | return false 307 | } 308 | for _, item := range items { 309 | if !s.Contains(item) { 310 | return false 311 | } 312 | } 313 | return true 314 | } 315 | 316 | // MarshalJSON implements the json.Marshaler interface. 317 | func (s *HashSet[T, H]) MarshalJSON() ([]byte, error) { 318 | return marshalJSON[T](s) 319 | } 320 | 321 | // UnmarshalJSON implements the json.Unmarshaler interface. 322 | func (s *HashSet[T, H]) UnmarshalJSON(data []byte) error { 323 | return unmarshalJSON[T](s, data) 324 | } 325 | 326 | // Items returns a generator function for iterating each element in s by using 327 | // the range keyword. 328 | // 329 | // for element := range s.Items() { ... } 330 | func (s *HashSet[T, H]) Items() iter.Seq[T] { 331 | return func(yield func(T) bool) { 332 | for _, item := range s.items { 333 | if !yield(item) { 334 | return 335 | } 336 | } 337 | } 338 | } 339 | -------------------------------------------------------------------------------- /hashset_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "fmt" 8 | "strconv" 9 | "strings" 10 | "testing" 11 | 12 | "github.com/shoenig/test/must" 13 | ) 14 | 15 | // assertion that HashSet[T] implements Collection[T] 16 | var _ Collection[*company] = (*HashSet[*company, string])(nil) 17 | 18 | // company is an example type that is not comparable, and implements Hash() string 19 | type company struct { 20 | _ func() // not comparable 21 | address string 22 | floor int 23 | } 24 | 25 | func (c *company) Equal(o *company) bool { 26 | return c.address == o.address && c.floor == o.floor 27 | } 28 | 29 | func (c *company) Hash() string { 30 | return fmt.Sprintf("%s:%d", c.address, c.floor) 31 | } 32 | 33 | func (c *company) String() string { 34 | return fmt.Sprintf("<%s %d>", c.address, c.floor) 35 | } 36 | func (c *company) MarshalJSON() ([]byte, error) { 37 | return []byte(fmt.Sprintf("{\"%s\":%d}", c.address, c.floor)), nil 38 | } 39 | 40 | func (c *company) UnmarshalJSON(data []byte) error { 41 | s := strings.TrimLeft(string(data), "{") 42 | s = strings.TrimRight(s, "}") 43 | 44 | splitSlice := strings.Split(s, ":") 45 | address := strings.TrimLeft(splitSlice[0], "\"") 46 | address = strings.TrimRight(address, "\"") 47 | c.address = address 48 | c.floor, _ = strconv.Atoi(splitSlice[1]) 49 | return nil 50 | } 51 | 52 | var ( 53 | c1 = &company{address: "street", floor: 1} 54 | c2 = &company{address: "street", floor: 2} 55 | c3 = &company{address: "street", floor: 3} 56 | c4 = &company{address: "street", floor: 4} 57 | c5 = &company{address: "street", floor: 5} 58 | c6 = &company{address: "street", floor: 6} 59 | c7 = &company{address: "street", floor: 7} 60 | c8 = &company{address: "street", floor: 8} 61 | c10 = &company{address: "street", floor: 10} 62 | ) 63 | 64 | // coded is an example type that maintains its own hash code, implementing Hash() int 65 | type coded struct { 66 | i int // internal hash code 67 | } 68 | 69 | func (c *coded) Hash() int { 70 | return c.i 71 | } 72 | 73 | var ( 74 | s1 = &coded{i: 1} 75 | s2 = &coded{i: 2} 76 | s3 = &coded{i: 3} 77 | ) 78 | 79 | func TestHashSet_New(t *testing.T) { 80 | t.Run("positive", func(t *testing.T) { 81 | s := NewHashSet[*company, string](1) 82 | must.MapEmpty(t, s.items) 83 | }) 84 | 85 | t.Run("zero", func(t *testing.T) { 86 | s := NewHashSet[*company, string](0) 87 | must.MapEmpty(t, s.items) 88 | }) 89 | 90 | t.Run("negative", func(t *testing.T) { 91 | s := NewHashSet[*company, string](-1) 92 | must.MapEmpty(t, s.items) 93 | }) 94 | } 95 | 96 | func TestHashSet_Insert(t *testing.T) { 97 | t.Run("one", func(t *testing.T) { 98 | s := NewHashSet[*company, string](1) 99 | must.True(t, s.Insert(c1)) 100 | must.MapContainsKeys(t, s.items, []string{"street:1"}) 101 | }) 102 | 103 | t.Run("re-insert", func(t *testing.T) { 104 | s := NewHashSet[*company, string](1) 105 | must.True(t, s.Insert(c1)) 106 | must.False(t, s.Insert(c1)) 107 | must.MapContainsKeys(t, s.items, []string{"street:1"}) 108 | }) 109 | 110 | t.Run("insert several", func(t *testing.T) { 111 | s := NewHashSet[*company, string](3) 112 | must.True(t, s.Insert(c1)) 113 | must.True(t, s.Insert(c2)) 114 | must.True(t, s.Insert(c3)) 115 | must.MapContainsKeys(t, s.items, []string{ 116 | "street:1", "street:2", "street:3", 117 | }) 118 | }) 119 | } 120 | 121 | func TestHashSet_InsertSlice(t *testing.T) { 122 | t.Run("insert none", func(t *testing.T) { 123 | empty := NewHashSet[*company, string](0) 124 | must.False(t, empty.InsertSlice(nil)) 125 | must.MapEmpty(t, empty.items) 126 | }) 127 | 128 | t.Run("insert some", func(t *testing.T) { 129 | s := NewHashSet[*company, string](0) 130 | must.True(t, s.InsertSlice([]*company{c1, c2, c3})) 131 | must.MapContainsKeys(t, s.items, []string{ 132 | "street:1", "street:2", "street:3", 133 | }) 134 | }) 135 | } 136 | 137 | func TestHashSet_InsertSet(t *testing.T) { 138 | t.Run("insert empty", func(t *testing.T) { 139 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 140 | b := NewHashSet[*company, string](0) 141 | must.False(t, a.InsertSet(b)) 142 | must.MapContainsKeys(t, a.items, []string{ 143 | "street:1", "street:2", "street:3", 144 | }) 145 | }) 146 | 147 | t.Run("insert some", func(t *testing.T) { 148 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 149 | b := HashSetFrom[*company, string]([]*company{c3, c4, c5}) 150 | a.InsertSet(b) 151 | must.MapContainsKeys(t, a.items, []string{ 152 | "street:1", "street:2", "street:3", "street:4", "street:5", 153 | }) 154 | }) 155 | } 156 | 157 | func TestHashSet_Remove(t *testing.T) { 158 | t.Run("empty remove item", func(t *testing.T) { 159 | s := NewHashSet[*company, string](10) 160 | must.False(t, s.Remove(c1)) 161 | must.MapEmpty(t, s.items) 162 | }) 163 | 164 | t.Run("set remove item", func(t *testing.T) { 165 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 166 | must.True(t, s.Remove(c2)) 167 | must.MapContainsKeys(t, s.items, []string{ 168 | "street:1", "street:3", 169 | }) 170 | }) 171 | 172 | t.Run("set remove missing", func(t *testing.T) { 173 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 174 | must.False(t, s.Remove(c4)) 175 | must.MapContainsKeys(t, s.items, []string{ 176 | "street:1", "street:2", "street:3", 177 | }) 178 | }) 179 | } 180 | 181 | func TestHashSet_RemoveSlice(t *testing.T) { 182 | t.Run("empty remove all", func(t *testing.T) { 183 | s := NewHashSet[*company, string](0) 184 | must.False(t, s.RemoveSlice([]*company{c1, c2, c3})) 185 | must.MapEmpty(t, s.items) 186 | }) 187 | 188 | t.Run("set remove nothing", func(t *testing.T) { 189 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 190 | must.False(t, s.RemoveSlice([]*company{c4, c5})) 191 | must.MapContainsKeys(t, s.items, []string{ 192 | "street:1", "street:2", "street:3", 193 | }) 194 | }) 195 | 196 | t.Run("set remove some", func(t *testing.T) { 197 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 198 | must.True(t, s.RemoveSlice([]*company{c4, c2})) 199 | must.MapContainsKeys(t, s.items, []string{ 200 | "street:1", "street:3", "street:5", 201 | }) 202 | }) 203 | } 204 | 205 | func TestHashSet_RemoveSet(t *testing.T) { 206 | t.Run("empty remove empty", func(t *testing.T) { 207 | a := NewHashSet[*company, string](0) 208 | b := NewHashSet[*company, string](0) 209 | must.False(t, a.RemoveSet(b)) 210 | must.MapEmpty(t, a.items) 211 | }) 212 | 213 | t.Run("empty remove some", func(t *testing.T) { 214 | a := NewHashSet[*company, string](0) 215 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 216 | must.False(t, a.RemoveSet(b)) 217 | must.MapEmpty(t, a.items) 218 | }) 219 | 220 | t.Run("set remove some", func(t *testing.T) { 221 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 222 | b := HashSetFrom[*company, string]([]*company{c2, c3}) 223 | must.True(t, a.RemoveSet(b)) 224 | must.MapContainsKeys(t, a.items, []string{ 225 | "street:1", "street:4", "street:5", 226 | }) 227 | }) 228 | } 229 | 230 | func TestHashSet_RemoveFunc(t *testing.T) { 231 | t.Run("empty", func(t *testing.T) { 232 | s := NewHashSet[*company, string](10) 233 | modified := s.RemoveFunc(func(c *company) bool { 234 | return c.floor > 3 235 | }) 236 | must.Empty(t, s) 237 | must.False(t, modified) 238 | }) 239 | 240 | t.Run("none match", func(t *testing.T) { 241 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 242 | modified := s.RemoveFunc(func(c *company) bool { 243 | return c.floor > 3 244 | }) 245 | must.False(t, modified) 246 | must.Size(t, 3, s) 247 | }) 248 | 249 | t.Run("some match", func(t *testing.T) { 250 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5, c6}) 251 | modified := s.RemoveFunc(func(c *company) bool { 252 | return c.floor > 3 253 | }) 254 | must.True(t, modified) 255 | must.Size(t, 3, s) 256 | must.Contains[*company](t, c1, s) 257 | must.Contains[*company](t, c2, s) 258 | must.Contains[*company](t, c3, s) 259 | }) 260 | 261 | t.Run("all match", func(t *testing.T) { 262 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5, c6}) 263 | modified := s.RemoveFunc(func(c *company) bool { 264 | return c.floor >= 0 265 | }) 266 | must.True(t, modified) 267 | must.Empty(t, s) 268 | }) 269 | } 270 | 271 | func TestHashSet_Contains(t *testing.T) { 272 | t.Run("empty contains", func(t *testing.T) { 273 | a := NewHashSet[*company, string](0) 274 | must.False(t, a.Contains(c1)) 275 | }) 276 | 277 | t.Run("not contains", func(t *testing.T) { 278 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 279 | must.False(t, s.Contains(c4)) 280 | }) 281 | 282 | t.Run("does contain", func(t *testing.T) { 283 | s := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 284 | must.True(t, s.Contains(c1)) 285 | }) 286 | } 287 | 288 | func TestHashSet_ContainsSlice(t *testing.T) { 289 | t.Run("empty empty", func(t *testing.T) { 290 | a := NewHashSet[*company, string](0) 291 | b := make([]*company, 0) 292 | must.True(t, a.ContainsSlice(b)) 293 | }) 294 | 295 | t.Run("empty some", func(t *testing.T) { 296 | a := NewHashSet[*company, string](0) 297 | b := []*company{c1, c2, c3} 298 | must.False(t, a.ContainsSlice(b)) 299 | }) 300 | 301 | t.Run("some empty", func(t *testing.T) { 302 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 303 | b := make([]*company, 0) 304 | must.False(t, a.ContainsSlice(b)) 305 | }) 306 | 307 | t.Run("equal", func(t *testing.T) { 308 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 309 | b := []*company{c3, c2, c1} 310 | must.True(t, a.ContainsSlice(b)) 311 | }) 312 | 313 | t.Run("not equal", func(t *testing.T) { 314 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 315 | b := []*company{c2, c3, c4} 316 | must.False(t, a.ContainsSlice(b)) 317 | }) 318 | 319 | t.Run("subset", func(t *testing.T) { 320 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 321 | b := []*company{c2, c3, c4} 322 | must.False(t, a.ContainsSlice(b)) 323 | }) 324 | 325 | t.Run("superset", func(t *testing.T) { 326 | a := HashSetFrom[*company, string]([]*company{c2, c3, c4}) 327 | b := []*company{c1, c2, c3, c4, c5} 328 | must.False(t, a.ContainsSlice(b)) 329 | }) 330 | 331 | t.Run("duplicates", func(t *testing.T) { 332 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 333 | b := []*company{c1, c2, c2, c3, c3, c4, c5} 334 | must.True(t, a.ContainsSlice(b)) 335 | }) 336 | } 337 | 338 | func TestHashSet_Subset(t *testing.T) { 339 | t.Run("empty empty", func(t *testing.T) { 340 | a := NewHashSet[*company, string](0) 341 | b := NewHashSet[*company, string](0) 342 | must.True(t, a.Subset(b)) 343 | }) 344 | 345 | t.Run("empty some", func(t *testing.T) { 346 | a := NewHashSet[*company, string](0) 347 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 348 | must.False(t, a.Subset(b)) 349 | }) 350 | 351 | t.Run("some empty", func(t *testing.T) { 352 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 353 | b := NewHashSet[*company, string](0) 354 | must.True(t, a.Subset(b)) 355 | }) 356 | 357 | t.Run("equal", func(t *testing.T) { 358 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 359 | b := HashSetFrom[*company, string]([]*company{c2, c3, c1}) 360 | must.True(t, a.Subset(b)) 361 | }) 362 | 363 | t.Run("subset", func(t *testing.T) { 364 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 365 | b := HashSetFrom[*company, string]([]*company{c3, c1}) 366 | must.True(t, a.Subset(b)) 367 | }) 368 | 369 | t.Run("superset", func(t *testing.T) { 370 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 371 | b := HashSetFrom[*company, string]([]*company{c3, c1, c2, c4}) 372 | must.False(t, a.Subset(b)) 373 | }) 374 | } 375 | 376 | func TestHashSet_ProperSubset(t *testing.T) { 377 | t.Run("empty empty", func(t *testing.T) { 378 | a := NewHashSet[*company, string](0) 379 | b := NewHashSet[*company, string](0) 380 | must.False(t, a.ProperSubset(b)) 381 | }) 382 | 383 | t.Run("empty some", func(t *testing.T) { 384 | a := NewHashSet[*company, string](0) 385 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 386 | must.False(t, a.ProperSubset(b)) 387 | }) 388 | 389 | t.Run("some empty", func(t *testing.T) { 390 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 391 | b := NewHashSet[*company, string](0) 392 | must.True(t, a.ProperSubset(b)) 393 | }) 394 | 395 | t.Run("equal", func(t *testing.T) { 396 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 397 | b := HashSetFrom[*company, string]([]*company{c2, c3, c1}) 398 | must.False(t, a.ProperSubset(b)) 399 | }) 400 | 401 | t.Run("subset", func(t *testing.T) { 402 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 403 | b := HashSetFrom[*company, string]([]*company{c3, c1}) 404 | must.True(t, a.ProperSubset(b)) 405 | }) 406 | 407 | t.Run("superset", func(t *testing.T) { 408 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 409 | b := HashSetFrom[*company, string]([]*company{c3, c1, c2, c4}) 410 | must.False(t, a.ProperSubset(b)) 411 | }) 412 | } 413 | 414 | func TestHashSet_Size(t *testing.T) { 415 | t.Run("size empty", func(t *testing.T) { 416 | s := NewHashSet[*company, string](10) 417 | must.Zero(t, s.Size()) 418 | }) 419 | 420 | t.Run("size not empty", func(t *testing.T) { 421 | s := NewHashSet[*company, string](10) 422 | must.True(t, s.Insert(c1)) 423 | must.True(t, s.Insert(c2)) 424 | must.Eq(t, 2, s.Size()) 425 | }) 426 | } 427 | 428 | func TestHashSet_Empty(t *testing.T) { 429 | t.Run("is empty", func(t *testing.T) { 430 | s := NewHashSet[*company, string](10) 431 | must.Empty(t, s) 432 | }) 433 | 434 | t.Run("is not empty", func(t *testing.T) { 435 | s := NewHashSet[*company, string](10) 436 | must.True(t, s.Insert(c1)) 437 | must.True(t, s.Insert(c2)) 438 | must.NotEmpty(t, s) 439 | }) 440 | } 441 | 442 | func TestHashSet_Difference(t *testing.T) { 443 | t.Run("empty \\ empty", func(t *testing.T) { 444 | a := NewHashSet[*company, string](10) 445 | b := NewHashSet[*company, string](10) 446 | diff := a.Difference(b).(*HashSet[*company, string]) 447 | must.MapEmpty(t, diff.items) 448 | }) 449 | 450 | t.Run("empty \\ set", func(t *testing.T) { 451 | a := NewHashSet[*company, string](10) 452 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 453 | diff := a.Difference(b).(*HashSet[*company, string]) 454 | must.MapEmpty(t, diff.items) 455 | }) 456 | 457 | t.Run("set \\ empty", func(t *testing.T) { 458 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 459 | b := NewHashSet[*company, string](10) 460 | diff := a.Difference(b).(*HashSet[*company, string]) 461 | must.MapContainsKeys(t, diff.items, []string{ 462 | "street:1", "street:2", "street:3", "street:4", "street:5", 463 | }) 464 | }) 465 | 466 | t.Run("set \\ other", func(t *testing.T) { 467 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5, c6, c7, c8}) 468 | b := HashSetFrom[*company, string]([]*company{c2, c4, c6, c8, c10, c10}) 469 | diff := a.Difference(b).(*HashSet[*company, string]) 470 | must.MapContainsKeys(t, diff.items, []string{ 471 | "street:1", "street:3", "street:5", "street:7", 472 | }) 473 | }) 474 | } 475 | 476 | func TestHashSet_Intersect(t *testing.T) { 477 | t.Run("empty ∩ empty", func(t *testing.T) { 478 | a := NewHashSet[*company, string](10) 479 | b := NewHashSet[*company, string](10) 480 | intersect := a.Intersect(b).(*HashSet[*company, string]) 481 | must.MapEmpty(t, intersect.items) 482 | }) 483 | 484 | t.Run("set ∩ empty", func(t *testing.T) { 485 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 486 | b := NewHashSet[*company, string](10) 487 | intersect := a.Intersect(b).(*HashSet[*company, string]) 488 | must.MapEmpty(t, intersect.items) 489 | }) 490 | 491 | t.Run("empty ∩ set", func(t *testing.T) { 492 | a := NewHashSet[*company, string](10) 493 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 494 | intersect := a.Intersect(b).(*HashSet[*company, string]) 495 | must.MapEmpty(t, intersect.items) 496 | }) 497 | 498 | t.Run("big ∩ small", func(t *testing.T) { 499 | a := HashSetFrom[*company, string]([]*company{c2, c3, c4, c6, c8}) 500 | b := HashSetFrom[*company, string]([]*company{c4, c5, c6, c7}) 501 | intersect := a.Intersect(b).(*HashSet[*company, string]) 502 | must.MapContainsKeys(t, intersect.items, []string{ 503 | "street:4", "street:6", 504 | }) 505 | }) 506 | 507 | t.Run("small ∩ big", func(t *testing.T) { 508 | a := HashSetFrom[*company, string]([]*company{c4, c5, c6, c7}) 509 | b := HashSetFrom[*company, string]([]*company{c2, c3, c4, c6, c8}) 510 | intersect := a.Intersect(b).(*HashSet[*company, string]) 511 | must.MapContainsKeys(t, intersect.items, []string{ 512 | "street:4", "street:6", 513 | }) 514 | }) 515 | } 516 | 517 | type special struct { 518 | hash string 519 | version int // not part of the hash 520 | } 521 | 522 | func (s *special) Hash() string { 523 | return s.hash 524 | } 525 | 526 | func TestHashSet_Union(t *testing.T) { 527 | t.Run("empty ∪ empty", func(t *testing.T) { 528 | a := NewHashSet[*company, string](10) 529 | b := NewHashSet[*company, string](10) 530 | union := a.Union(b).(*HashSet[*company, string]) 531 | must.MapEmpty(t, union.items) 532 | }) 533 | 534 | t.Run("set ∪ empty", func(t *testing.T) { 535 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 536 | b := NewHashSet[*company, string](10) 537 | union := a.Union(b).(*HashSet[*company, string]) 538 | must.MapContainsKeys(t, union.items, []string{ 539 | "street:1", "street:2", "street:3", 540 | }) 541 | }) 542 | 543 | t.Run("empty ∪ set", func(t *testing.T) { 544 | a := NewHashSet[*company, string](10) 545 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 546 | union := a.Union(b).(*HashSet[*company, string]) 547 | must.MapContainsKeys(t, union.items, []string{ 548 | "street:1", "street:2", "street:3", 549 | }) 550 | }) 551 | 552 | t.Run("big ∪ small", func(t *testing.T) { 553 | a := HashSetFrom[*company, string]([]*company{c4, c5, c6, c7}) 554 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 555 | union := a.Union(b).(*HashSet[*company, string]) 556 | must.MapContainsKeys(t, union.items, []string{ 557 | "street:1", "street:2", "street:3", "street:4", "street:5", "street:6", "street:7", 558 | }) 559 | }) 560 | 561 | t.Run("overlap", func(t *testing.T) { 562 | a := HashSetFrom[*company, string]([]*company{c4, c5, c6, c7}) 563 | b := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 564 | union := a.Union(b).(*HashSet[*company, string]) 565 | must.MapContainsKeys(t, union.items, []string{ 566 | "street:1", "street:2", "street:3", "street:4", "street:5", "street:6", "street:7", 567 | }) 568 | }) 569 | 570 | t.Run("overlap priority", func(t *testing.T) { 571 | x1 := &special{hash: "x", version: 1} 572 | x2 := &special{hash: "x", version: 2} 573 | 574 | a := HashSetFrom[*special, string]([]*special{x1}) 575 | b := HashSetFrom[*special, string]([]*special{x2}) 576 | 577 | // receiver elements take priority over argument elements 578 | union1 := a.Union(b).(*HashSet[*special, string]) 579 | must.MapContainsValues(t, union1.items, []*special{ 580 | x1, 581 | }) 582 | 583 | // checking in the other direction 584 | union2 := b.Union(a).(*HashSet[*special, string]) 585 | must.MapContainsValues(t, union2.items, []*special{ 586 | x2, 587 | }) 588 | }) 589 | } 590 | 591 | func TestHashSet_Copy(t *testing.T) { 592 | t.Run("copy empty", func(t *testing.T) { 593 | a := NewHashSet[*company, string](0) 594 | b := a.Copy() 595 | must.MapEmpty(t, b.items) 596 | must.True(t, b.Insert(c3)) 597 | must.MapEmpty(t, a.items) 598 | must.MapContainsKeys(t, b.items, []string{"street:3"}) 599 | }) 600 | 601 | t.Run("copy some", func(t *testing.T) { 602 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4}) 603 | b := a.Copy() 604 | must.MapContainsKeys(t, b.items, []string{ 605 | "street:1", "street:2", "street:3", "street:4", 606 | }) 607 | must.True(t, b.RemoveSlice([]*company{c1, c3})) 608 | must.MapContainsKeys(t, b.items, []string{"street:2", "street:4"}) 609 | must.MapContainsKeys(t, a.items, []string{ 610 | "street:1", "street:2", "street:3", "street:4", 611 | }) 612 | }) 613 | } 614 | 615 | func TestHashSet_Slice(t *testing.T) { 616 | t.Run("slice empty", func(t *testing.T) { 617 | a := NewHashSet[*company, string](10) 618 | l := a.Slice() 619 | must.SliceEmpty(t, l) 620 | }) 621 | 622 | t.Run("slice set", func(t *testing.T) { 623 | a := HashSetFrom[*company, string]([]*company{c1, c2}) 624 | l := a.Slice() 625 | must.Len(t, 2, l) 626 | must.SliceContainsEqual(t, l, c1) 627 | must.SliceContainsEqual(t, l, c2) 628 | }) 629 | } 630 | 631 | func TestHashSet_String(t *testing.T) { 632 | a := HashSetFrom[*company, string]([]*company{c2, c1}) 633 | result := a.String() 634 | must.Eq(t, "[ ]", result) 635 | } 636 | 637 | func TestHashSet_StringFunc(t *testing.T) { 638 | t.Run("empty", func(t *testing.T) { 639 | a := NewHashSet[*company, string](10) 640 | s := a.StringFunc(nil) 641 | must.Eq(t, "[]", s) 642 | }) 643 | 644 | t.Run("some", func(t *testing.T) { 645 | a := HashSetFrom[*company, string]([]*company{c1, c2}) 646 | s := a.StringFunc(func(c *company) string { 647 | return fmt.Sprintf("(%s %d)", c.address, c.floor) 648 | }) 649 | must.Eq(t, "[(street 1) (street 2)]", s) 650 | }) 651 | } 652 | 653 | func TestHashSet_EqualSet(t *testing.T) { 654 | t.Run("empty empty", func(t *testing.T) { 655 | a := NewHashSet[*company, string](0) 656 | b := NewHashSet[*company, string](0) 657 | must.True(t, a.EqualSet(b)) 658 | }) 659 | 660 | t.Run("different", func(t *testing.T) { 661 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 662 | b := HashSetFrom[*company, string]([]*company{c1, c2, c4}) 663 | must.False(t, a.EqualSet(b)) 664 | must.False(t, b.EqualSet(a)) 665 | }) 666 | } 667 | 668 | func TestHashSet_EqualSlice(t *testing.T) { 669 | t.Run("empty empty", func(t *testing.T) { 670 | a := NewHashSet[*company, string](0) 671 | b := make([]*company, 0) 672 | must.True(t, a.EqualSlice(b)) 673 | }) 674 | 675 | t.Run("empty some", func(t *testing.T) { 676 | a := NewHashSet[*company, string](0) 677 | b := []*company{c1, c2, c3} 678 | must.False(t, a.EqualSlice(b)) 679 | }) 680 | 681 | t.Run("some empty", func(t *testing.T) { 682 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 683 | b := make([]*company, 0) 684 | must.False(t, a.EqualSlice(b)) 685 | }) 686 | 687 | t.Run("equal", func(t *testing.T) { 688 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 689 | b := []*company{c3, c2, c1} 690 | must.True(t, a.EqualSlice(b)) 691 | }) 692 | 693 | t.Run("not equal", func(t *testing.T) { 694 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3}) 695 | b := []*company{c2, c3, c4} 696 | must.False(t, a.EqualSlice(b)) 697 | }) 698 | 699 | t.Run("subset", func(t *testing.T) { 700 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 701 | b := []*company{c2, c3, c4} 702 | must.False(t, a.EqualSlice(b)) 703 | }) 704 | 705 | t.Run("superset", func(t *testing.T) { 706 | a := HashSetFrom[*company, string]([]*company{c2, c3, c4}) 707 | b := []*company{c1, c2, c3, c4, c5} 708 | must.False(t, a.EqualSlice(b)) 709 | }) 710 | 711 | t.Run("duplicates", func(t *testing.T) { 712 | a := HashSetFrom[*company, string]([]*company{c1, c2, c3, c4, c5}) 713 | b := []*company{c1, c2, c2, c3, c3, c4, c5} 714 | must.True(t, a.EqualSlice(b)) 715 | }) 716 | } 717 | 718 | func TestHashSet_HashCode(t *testing.T) { 719 | a := NewHashSet[*coded, int](0) 720 | a.Insert(s1) 721 | a.Insert(s2) 722 | must.MapContainsKeys(t, a.items, []int{1, 2}) 723 | must.True(t, a.Contains(s1)) 724 | must.True(t, a.Contains(s2)) 725 | must.False(t, a.Contains(s3)) 726 | } 727 | 728 | func TestHashSet_Items(t *testing.T) { 729 | a := NewHashSet[*coded, int](0) 730 | a.Insert(s1) 731 | a.Insert(s2) 732 | a.Insert(s3) 733 | 734 | sum := 0 735 | for element := range a.Items() { 736 | sum += element.i 737 | } 738 | 739 | must.Eq(t, 6, sum) 740 | } 741 | -------------------------------------------------------------------------------- /serialization.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import "encoding/json" 7 | 8 | // marshalJSON will serialize a Serializable[T] into a json byte array 9 | func marshalJSON[T any](s Collection[T]) ([]byte, error) { 10 | return json.Marshal(s.Slice()) 11 | } 12 | 13 | // unmarshalJSON will deserialize a json byte array into a Serializable[T] 14 | func unmarshalJSON[T any](s Collection[T], data []byte) error { 15 | slice := make([]T, 0) 16 | err := json.Unmarshal(data, &slice) 17 | if err != nil { 18 | return err 19 | } 20 | s.InsertSlice(slice) 21 | return nil 22 | } 23 | -------------------------------------------------------------------------------- /serialization_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "cmp" 8 | "encoding/json" 9 | "testing" 10 | 11 | "github.com/shoenig/test/must" 12 | ) 13 | 14 | func TestSerialization(t *testing.T) { 15 | t.Run("Set", func(t *testing.T) { 16 | set := New[int](3) 17 | set.InsertSlice([]int{1, 2, 3}) 18 | bs, err := json.Marshal(set) 19 | must.NoError(t, err) 20 | must.StrContains(t, string(bs), "1") 21 | must.StrContains(t, string(bs), "2") 22 | must.StrContains(t, string(bs), "3") 23 | 24 | dstSet := New[int](3) 25 | err = json.Unmarshal(bs, dstSet) 26 | must.NoError(t, err) 27 | must.MapEq(t, dstSet.items, set.items) 28 | }) 29 | 30 | t.Run("HashSet", func(t *testing.T) { 31 | set := NewHashSet[*company, string](10) 32 | set.InsertSlice([]*company{c1, c2, c3}) 33 | bs, err := json.Marshal(set) 34 | must.NoError(t, err) 35 | must.StrContains(t, string(bs), `"street":1`) 36 | must.StrContains(t, string(bs), `"street":2`) 37 | must.StrContains(t, string(bs), `"street":3`) 38 | 39 | dstSet := NewHashSet[*company, string](10) 40 | err = json.Unmarshal(bs, dstSet) 41 | must.NoError(t, err) 42 | must.MapEqual(t, dstSet.items, set.items) 43 | }) 44 | 45 | t.Run("TreeSet", func(t *testing.T) { 46 | set := NewTreeSet[int](cmp.Compare[int]) 47 | set.InsertSlice([]int{10, 3, 13}) 48 | bs, err := json.Marshal(set) 49 | must.NoError(t, err) 50 | must.StrContains(t, string(bs), "10") 51 | must.StrContains(t, string(bs), "3") 52 | must.StrContains(t, string(bs), "13") 53 | 54 | dstSet := NewTreeSet[int](cmp.Compare[int]) 55 | err = json.Unmarshal(bs, dstSet) 56 | must.NoError(t, err) 57 | must.Eq(t, set.Slice(), dstSet.Slice()) 58 | }) 59 | } 60 | -------------------------------------------------------------------------------- /set.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | // Package set provides a basic generic set implementation. 5 | // 6 | // https://en.wikipedia.org/wiki/Set_(mathematics) 7 | package set 8 | 9 | import ( 10 | "fmt" 11 | "iter" 12 | "sort" 13 | ) 14 | 15 | type nothing struct{} 16 | 17 | var sentinel = nothing{} 18 | 19 | // New creates a new Set with initial underlying capacity of size. 20 | // 21 | // A Set will automatically grow or shrink its capacity as items are added or 22 | // removed. 23 | // 24 | // T may be any comparable type. Keep in mind that pointer types or structs 25 | // containing pointer fields will be compared using shallow equality. For deep 26 | // equality use HashSet instead. 27 | func New[T comparable](size int) *Set[T] { 28 | return &Set[T]{ 29 | items: make(map[T]nothing, max(0, size)), 30 | } 31 | } 32 | 33 | // From creates a new Set containing each item in items. 34 | // 35 | // T may be any comparable type. Keep in mind that pointer types or structs 36 | // containing pointer fields will be compared using shallow equality. For deep 37 | // equality use HashSet instead. 38 | func From[T comparable](items []T) *Set[T] { 39 | s := New[T](len(items)) 40 | s.InsertSlice(items) 41 | return s 42 | } 43 | 44 | // FromFunc creates a new Set containing a conversion of each item in items. 45 | // 46 | // T may be any comparable type. Keep in mind that pointer types or structs 47 | // containing pointer fields will be compared using shallow equality. For deep 48 | // equality use HashSet instead. 49 | func FromFunc[A any, T comparable](items []A, conversion func(A) T) *Set[T] { 50 | s := New[T](len(items)) 51 | for _, item := range items { 52 | s.Insert(conversion(item)) 53 | } 54 | return s 55 | } 56 | 57 | // Set is a simple, generic implementation of the set mathematical data structure. 58 | // It is optimized for correctness and convenience, as a replacement for the use 59 | // of map[interface{}]struct{}. 60 | type Set[T comparable] struct { 61 | items map[T]nothing 62 | } 63 | 64 | // Insert item into s. 65 | // 66 | // Return true if s was modified (item was not already in s), false otherwise. 67 | func (s *Set[T]) Insert(item T) bool { 68 | if _, exists := s.items[item]; exists { 69 | return false 70 | } 71 | if s.items == nil { 72 | s.items = make(map[T]nothing) 73 | } 74 | s.items[item] = sentinel 75 | return true 76 | } 77 | 78 | // InsertSlice will insert each item in items into s. 79 | // 80 | // Return true if s was modified (at least one item was not already in s), false otherwise. 81 | func (s *Set[T]) InsertSlice(items []T) bool { 82 | modified := false 83 | for _, item := range items { 84 | if s.Insert(item) { 85 | modified = true 86 | } 87 | } 88 | return modified 89 | } 90 | 91 | // InsertSet will insert each element of col into s. 92 | // 93 | // Return true if s was modified (at least one item of col was not already in s), false otherwise. 94 | func (s *Set[T]) InsertSet(col Collection[T]) bool { 95 | modified := false 96 | for item := range col.Items() { 97 | if s.Insert(item) { 98 | modified = true 99 | } 100 | } 101 | return modified 102 | } 103 | 104 | // Remove will remove item from s. 105 | // 106 | // Return true if s was modified (item was present), false otherwise. 107 | func (s *Set[T]) Remove(item T) bool { 108 | if _, exists := s.items[item]; !exists { 109 | return false 110 | } 111 | delete(s.items, item) 112 | return true 113 | } 114 | 115 | // RemoveSlice will remove each item in items from s. 116 | // 117 | // Return true if s was modified (any item was present), false otherwise. 118 | func (s *Set[T]) RemoveSlice(items []T) bool { 119 | modified := false 120 | for _, item := range items { 121 | if s.Remove(item) { 122 | modified = true 123 | } 124 | } 125 | return modified 126 | } 127 | 128 | // RemoveSet will remove each element of col from s. 129 | // 130 | // Return true if s was modified (any item of o was present in s), false otherwise. 131 | func (s *Set[T]) RemoveSet(col Collection[T]) bool { 132 | return removeSet(s, col) 133 | } 134 | 135 | // RemoveFunc will remove each element from s that satisfies condition f. 136 | // 137 | // Return true if s was modified, false otherwise. 138 | func (s *Set[T]) RemoveFunc(f func(T) bool) bool { 139 | return removeFunc(s, f) 140 | } 141 | 142 | // Contains returns whether item is present in s. 143 | func (s *Set[T]) Contains(item T) bool { 144 | _, exists := s.items[item] 145 | return exists 146 | } 147 | 148 | // ContainsSlice returns whether all elements in items are present in s. 149 | func (s *Set[T]) ContainsSlice(items []T) bool { 150 | return containsSlice(s, items) 151 | } 152 | 153 | // Subset returns whether col is a subset of s. 154 | func (s *Set[T]) Subset(col Collection[T]) bool { 155 | return subset(s, col) 156 | } 157 | 158 | // Subset returns whether col is a proper subset of s. 159 | func (s *Set[T]) ProperSubset(col Collection[T]) bool { 160 | if len(s.items) <= col.Size() { 161 | return false 162 | } 163 | return s.Subset(col) 164 | } 165 | 166 | // Size returns the cardinality of s. 167 | func (s *Set[T]) Size() int { 168 | return len(s.items) 169 | } 170 | 171 | // Empty returns true if s contains no elements, false otherwise. 172 | func (s *Set[T]) Empty() bool { 173 | return s.Size() == 0 174 | } 175 | 176 | // Union returns a set that contains all elements of s and col combined. 177 | func (s *Set[T]) Union(col Collection[T]) Collection[T] { 178 | size := max(s.Size(), col.Size()) 179 | result := New[T](size) 180 | insert(result, s) 181 | insert(result, col) 182 | return result 183 | } 184 | 185 | // Difference returns a set that contains elements of s that are not in col. 186 | func (s *Set[T]) Difference(col Collection[T]) Collection[T] { 187 | result := New[T](max(0, s.Size()-col.Size())) 188 | for item := range s.items { 189 | if !col.Contains(item) { 190 | result.items[item] = sentinel 191 | } 192 | } 193 | return result 194 | } 195 | 196 | // Intersect returns a set that contains elements that are present in both s and col. 197 | func (s *Set[T]) Intersect(col Collection[T]) Collection[T] { 198 | result := New[T](0) 199 | intersect(result, s, col) 200 | return result 201 | } 202 | 203 | // Copy creates a copy of s. 204 | func (s *Set[T]) Copy() *Set[T] { 205 | result := New[T](s.Size()) 206 | for item := range s.items { 207 | result.items[item] = sentinel 208 | } 209 | return result 210 | } 211 | 212 | // Slice creates a copy of s as a slice. Elements are in no particular order. 213 | func (s *Set[T]) Slice() []T { 214 | result := make([]T, 0, s.Size()) 215 | for item := range s.items { 216 | result = append(result, item) 217 | } 218 | return result 219 | } 220 | 221 | // String creates a string representation of s, using "%v" printf formating to transform 222 | // each element into a string. The result contains elements sorted by their lexical 223 | // string order. 224 | func (s *Set[T]) String() string { 225 | return s.StringFunc(func(element T) string { 226 | return fmt.Sprintf("%v", element) 227 | }) 228 | } 229 | 230 | // StringFunc creates a string representation of s, using f to transform each element 231 | // into a string. The result contains elements sorted by their lexical string order. 232 | func (s *Set[T]) StringFunc(f func(element T) string) string { 233 | l := make([]string, 0, s.Size()) 234 | for item := range s.items { 235 | l = append(l, f(item)) 236 | } 237 | sort.Strings(l) 238 | return fmt.Sprintf("%s", l) 239 | } 240 | 241 | // Equal returns whether s and o contain the same elements. 242 | func (s *Set[T]) Equal(o *Set[T]) bool { 243 | if len(s.items) != len(o.items) { 244 | return false 245 | } 246 | for item := range s.items { 247 | if !o.Contains(item) { 248 | return false 249 | } 250 | } 251 | return true 252 | } 253 | 254 | // EqualSet returns whether s and col contain the same elements. 255 | func (s *Set[T]) EqualSet(col Collection[T]) bool { 256 | return equalSet(s, col) 257 | } 258 | 259 | // EqualSlice returns whether s and items contain the same elements. 260 | // 261 | // The items slice may contain duplicates. 262 | // 263 | // If the items slice is known to contain no duplicates, EqualSliceSet may be 264 | // used instead as a faster implementation. 265 | // 266 | // To detect if a slice is a subset of s, use ContainsSlice. 267 | func (s *Set[T]) EqualSlice(items []T) bool { 268 | other := From[T](items) 269 | return s.Equal(other) 270 | } 271 | 272 | // EqualSliceSet returns whether s and items contain exactly the same elements. 273 | // 274 | // If items contains duplicates EqualSliceSet will return false. The elements of 275 | // items are assumed to be set-like. For comparing s to a slice that may contain 276 | // duplicate elements, use EqualSlice instead. 277 | // 278 | // To detect if a slice is a subset of s, use ContainsSlice. 279 | func (s *Set[T]) EqualSliceSet(items []T) bool { 280 | if len(items) != s.Size() { 281 | return false 282 | } 283 | for _, item := range items { 284 | if !s.Contains(item) { 285 | return false 286 | } 287 | } 288 | return true 289 | } 290 | 291 | // MarshalJSON implements the json.Marshaler interface. 292 | func (s *Set[T]) MarshalJSON() ([]byte, error) { 293 | return marshalJSON[T](s) 294 | } 295 | 296 | // UnmarshalJSON implements the json.Unmarshaler interface. 297 | func (s *Set[T]) UnmarshalJSON(data []byte) error { 298 | return unmarshalJSON[T](s, data) 299 | } 300 | 301 | // Items returns a generator function for iterating each element in s by using 302 | // the range keyword. 303 | // 304 | // for element := range s.Items() { ... } 305 | func (s *Set[T]) Items() iter.Seq[T] { 306 | return func(yield func(T) bool) { 307 | for item := range s.items { 308 | if !yield(item) { 309 | return 310 | } 311 | } 312 | } 313 | } 314 | -------------------------------------------------------------------------------- /set_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "fmt" 8 | "testing" 9 | 10 | "github.com/shoenig/test/must" 11 | ) 12 | 13 | // assertion that Set[T] implements Collection[T] 14 | var _ Collection[nothing] = (*Set[nothing])(nil) 15 | 16 | type employee struct { 17 | name string 18 | id int 19 | } 20 | 21 | func (e *employee) String() string { 22 | return fmt.Sprintf("(%d %s)", e.id, e.name) 23 | } 24 | 25 | func TestSet_New(t *testing.T) { 26 | t.Run("positive", func(t *testing.T) { 27 | s := New[float64](1) 28 | must.MapEmpty(t, s.items) 29 | }) 30 | 31 | t.Run("zero", func(t *testing.T) { 32 | s := New[int](0) 33 | must.MapEmpty(t, s.items) 34 | }) 35 | 36 | t.Run("negative", func(t *testing.T) { 37 | s := New[string](-1) // assume zero 38 | must.MapEmpty(t, s.items) 39 | }) 40 | } 41 | 42 | func TestSet_From(t *testing.T) { 43 | t.Run("from nil", func(t *testing.T) { 44 | s := From[string](nil) 45 | must.MapEmpty(t, s.items) 46 | }) 47 | 48 | t.Run("from some", func(t *testing.T) { 49 | s := From[string]([]string{"apple", "banana", "cherry"}) 50 | must.MapContainsKeys(t, s.items, []string{"apple", "banana", "cherry"}) 51 | }) 52 | } 53 | 54 | func TestSet_FromFunc(t *testing.T) { 55 | employees := []employee{ 56 | {"alice", 1}, {"bob", 2}, {"bob", 2}, {"carol", 3}, {"dave", 4}, 57 | } 58 | s := FromFunc(employees, func(e employee) string { 59 | return e.name 60 | }) 61 | must.MapContainsKeys(t, s.items, []string{"alice", "bob", "carol", "dave"}) 62 | } 63 | 64 | func TestSet_Insert(t *testing.T) { 65 | t.Run("one int", func(t *testing.T) { 66 | s := New[int](10) 67 | must.True(t, s.Insert(1)) 68 | must.MapContainsKeys(t, s.items, []int{1}) 69 | }) 70 | 71 | t.Run("one string", func(t *testing.T) { 72 | s := New[string](10) 73 | must.True(t, s.Insert("apple")) 74 | must.MapContainsKeys(t, s.items, []string{"apple"}) 75 | }) 76 | 77 | t.Run("re-insert", func(t *testing.T) { 78 | s := New[int](10) 79 | must.True(t, s.Insert(2)) 80 | must.False(t, s.Insert(2)) 81 | must.MapContainsKeys(t, s.items, []int{2}) 82 | }) 83 | 84 | t.Run("insert several", func(t *testing.T) { 85 | s := New[int](10) 86 | must.True(t, s.Insert(1)) 87 | must.True(t, s.Insert(2)) 88 | must.True(t, s.Insert(3)) 89 | must.True(t, s.Insert(4)) 90 | must.True(t, s.Insert(5)) 91 | must.MapContainsKeys(t, s.items, []int{1, 2, 3, 4, 5}) 92 | }) 93 | 94 | t.Run("insert custom", func(t *testing.T) { 95 | s := New[employee](10) 96 | must.True(t, s.Insert(employee{"mitchell", 1})) 97 | must.True(t, s.Insert(employee{"armon", 2})) 98 | must.True(t, s.Insert(employee{"jack", 3})) 99 | must.False(t, s.Insert(employee{"jack", 3})) 100 | must.False(t, s.Insert(employee{"armon", 2})) 101 | must.False(t, s.Insert(employee{"mitchell", 1})) 102 | must.MapContainsKeys(t, s.items, []employee{ 103 | {"mitchell", 1}, {"armon", 2}, {"jack", 3}, 104 | }) 105 | }) 106 | } 107 | 108 | func TestSet_InsertSlice(t *testing.T) { 109 | t.Run("insert none", func(t *testing.T) { 110 | empty := New[int](0) 111 | must.False(t, empty.InsertSlice(nil)) 112 | must.MapEmpty(t, empty.items) 113 | }) 114 | 115 | t.Run("insert some", func(t *testing.T) { 116 | s := New[string](0) 117 | must.True(t, s.InsertSlice([]string{"apple", "banana", "cherry"})) 118 | must.MapContainsKeys(t, s.items, []string{"apple", "banana", "cherry"}) 119 | }) 120 | 121 | t.Run("insert duplicates", func(t *testing.T) { 122 | s := New[int](0) 123 | must.True(t, s.InsertSlice([]int{2, 4, 6, 8})) 124 | must.True(t, s.InsertSlice([]int{4, 5, 6})) 125 | must.MapContainsKeys(t, s.items, []int{2, 4, 5, 6, 8}) 126 | }) 127 | } 128 | 129 | func TestSet_InsertSet(t *testing.T) { 130 | t.Run("insert empty", func(t *testing.T) { 131 | a := From[int]([]int{1, 2, 3, 4}) 132 | b := New[int](0) 133 | must.False(t, a.InsertSet(b)) 134 | must.MapContainsKeys(t, a.items, []int{1, 2, 3, 4}) 135 | }) 136 | 137 | t.Run("insert some", func(t *testing.T) { 138 | a := From[int]([]int{1, 2, 3, 4}) 139 | b := From[int]([]int{3, 4, 5, 6, 7}) 140 | must.True(t, a.InsertSet(b)) 141 | must.MapContainsKeys(t, a.items, []int{1, 2, 3, 4, 5, 6, 7}) 142 | }) 143 | } 144 | 145 | func TestSet_Contains(t *testing.T) { 146 | t.Run("contains string item", func(t *testing.T) { 147 | s := New[string](10) 148 | must.True(t, s.InsertSlice([]string{"apple", "banana", "chery"})) 149 | must.True(t, s.Contains("apple")) 150 | must.True(t, s.Contains("banana")) 151 | must.True(t, s.Contains("chery")) 152 | must.False(t, s.Contains("zucchini")) 153 | }) 154 | 155 | t.Run("contains custom item", func(t *testing.T) { 156 | s := New[employee](10) 157 | must.True(t, s.Insert(employee{"mitchell", 1})) 158 | must.True(t, s.Insert(employee{"armon", 2})) 159 | must.True(t, s.Insert(employee{"jack", 3})) 160 | must.Contains[employee](t, employee{"mitchell", 1}, s) 161 | must.Contains[employee](t, employee{"armon", 2}, s) 162 | must.Contains[employee](t, employee{"jack", 3}, s) 163 | must.NotContains[employee](t, employee{"dave", 27}, s) 164 | }) 165 | } 166 | 167 | func TestSet_ContainsSlice(t *testing.T) { 168 | t.Run("empty empty", func(t *testing.T) { 169 | a := New[int](0) 170 | b := make([]int, 0) 171 | must.True(t, a.ContainsSlice(b)) 172 | }) 173 | 174 | t.Run("empty some", func(t *testing.T) { 175 | a := New[int](0) 176 | b := []int{1, 2, 3} 177 | must.False(t, a.ContainsSlice(b)) 178 | }) 179 | 180 | t.Run("some empty", func(t *testing.T) { 181 | a := From[int]([]int{1, 2, 3}) 182 | b := make([]int, 0) 183 | must.True(t, a.ContainsSlice(b)) 184 | }) 185 | 186 | t.Run("equal", func(t *testing.T) { 187 | a := From[int]([]int{1, 2, 3}) 188 | b := []int{3, 2, 1} 189 | must.True(t, a.ContainsSlice(b)) 190 | }) 191 | 192 | t.Run("not equal", func(t *testing.T) { 193 | a := From[int]([]int{1, 2, 3}) 194 | b := []int{2, 3, 4} 195 | must.False(t, a.ContainsSlice(b)) 196 | }) 197 | 198 | t.Run("subset", func(t *testing.T) { 199 | a := From[int]([]int{1, 2, 3, 4, 5}) 200 | b := []int{2, 3, 4} 201 | must.True(t, a.ContainsSlice(b)) 202 | }) 203 | 204 | t.Run("superset", func(t *testing.T) { 205 | a := From[int]([]int{2, 3, 4}) 206 | b := []int{1, 2, 3, 4, 5} 207 | must.False(t, a.ContainsSlice(b)) 208 | }) 209 | 210 | t.Run("duplicates", func(t *testing.T) { 211 | a := From[int]([]int{1, 2, 3, 4, 5}) 212 | b := []int{1, 2, 2, 3, 3, 4, 5} 213 | must.True(t, a.ContainsSlice(b)) 214 | }) 215 | } 216 | 217 | func TestSet_Size(t *testing.T) { 218 | t.Run("size empty", func(t *testing.T) { 219 | s := New[int](10) 220 | must.Zero(t, s.Size()) 221 | }) 222 | 223 | t.Run("size not empty", func(t *testing.T) { 224 | s := New[int](10) 225 | must.True(t, s.Insert(1)) 226 | must.True(t, s.Insert(2)) 227 | must.Eq(t, 2, s.Size()) 228 | }) 229 | } 230 | 231 | func TestSet_Empty(t *testing.T) { 232 | t.Run("is empty", func(t *testing.T) { 233 | s := New[int](10) 234 | must.Empty(t, s) 235 | }) 236 | 237 | t.Run("is not empty", func(t *testing.T) { 238 | s := From[int]([]int{1, 2, 3, 4}) 239 | must.NotEmpty(t, s) 240 | }) 241 | } 242 | 243 | func TestSet_Union(t *testing.T) { 244 | t.Run("empty ∪ empty", func(t *testing.T) { 245 | a := New[int](0) 246 | b := New[int](10) 247 | union := a.Union(b).(*Set[int]) 248 | must.MapEmpty(t, union.items) 249 | }) 250 | 251 | t.Run("empty ∪ set", func(t *testing.T) { 252 | a := New[int](10) 253 | b := New[int](10) 254 | b.InsertSlice([]int{1, 2, 3, 4, 5}) 255 | union := a.Union(b).(*Set[int]) 256 | must.MapContainsKeys(t, union.items, []int{1, 2, 3, 4, 5}) 257 | }) 258 | 259 | t.Run("set ∪ empty", func(t *testing.T) { 260 | a := New[int](10) 261 | a.InsertSlice([]int{1, 2, 3, 4, 5}) 262 | b := New[int](10) 263 | union := a.Union(b).(*Set[int]) 264 | must.MapContainsKeys(t, union.items, []int{1, 2, 3, 4, 5}) 265 | }) 266 | 267 | t.Run("set ∪ other", func(t *testing.T) { 268 | a := New[int](10) 269 | must.True(t, a.InsertSlice([]int{2, 4, 6, 8})) 270 | b := New[int](10) 271 | must.True(t, b.InsertSlice([]int{4, 5, 6})) 272 | union := a.Union(b).(*Set[int]) 273 | must.MapContainsKeys(t, union.items, []int{2, 4, 5, 6, 8}) 274 | }) 275 | } 276 | 277 | func TestSet_Difference(t *testing.T) { 278 | t.Run("empty \\ empty", func(t *testing.T) { 279 | a := New[int](10) 280 | b := New[int](10) 281 | diff := a.Difference(b).(*Set[int]) 282 | must.MapEmpty(t, diff.items) 283 | }) 284 | 285 | t.Run("empty \\ set", func(t *testing.T) { 286 | a := New[int](10) 287 | b := From([]int{1, 2, 3, 4, 5}) 288 | diff := a.Difference(b).(*Set[int]) 289 | must.MapEmpty(t, diff.items) 290 | }) 291 | 292 | t.Run("set \\ empty", func(t *testing.T) { 293 | a := From([]int{1, 2, 3, 4, 5}) 294 | b := New[int](10) 295 | diff := a.Difference(b).(*Set[int]) 296 | must.MapContainsKeys(t, diff.items, []int{1, 2, 3, 4, 5}) 297 | }) 298 | 299 | t.Run("set \\ other", func(t *testing.T) { 300 | a := From([]int{1, 2, 3, 4, 5, 6, 7, 8}) 301 | b := From([]int{2, 4, 6, 8, 10, 12}) 302 | diff := a.Difference(b).(*Set[int]) 303 | must.MapContainsKeys(t, diff.items, []int{1, 3, 5, 7}) 304 | }) 305 | } 306 | 307 | func TestSet_Intersect(t *testing.T) { 308 | t.Run("empty ∩ empty", func(t *testing.T) { 309 | a := New[int](10) 310 | b := New[int](10) 311 | intersect := a.Intersect(b).(*Set[int]) 312 | must.MapEmpty(t, intersect.items) 313 | }) 314 | 315 | t.Run("set ∩ empty", func(t *testing.T) { 316 | a := From[int]([]int{1, 2, 3}) 317 | b := New[int](10) 318 | intersect := a.Intersect(b).(*Set[int]) 319 | must.MapEmpty(t, intersect.items) 320 | }) 321 | 322 | t.Run("empty ∩ set", func(t *testing.T) { 323 | a := New[int](10) 324 | b := From[int]([]int{1, 2, 3}) 325 | intersect := a.Intersect(b).(*Set[int]) 326 | must.MapEmpty(t, intersect.items) 327 | }) 328 | 329 | t.Run("big ∩ small", func(t *testing.T) { 330 | a := From[int]([]int{2, 3, 4, 6, 8}) 331 | b := From[int]([]int{4, 5, 6, 7}) 332 | intersect := a.Intersect(b).(*Set[int]) 333 | must.MapContainsKeys(t, intersect.items, []int{4, 6}) 334 | }) 335 | 336 | t.Run("small ∩ big", func(t *testing.T) { 337 | a := From[int]([]int{4, 5, 6, 7}) 338 | b := From[int]([]int{2, 3, 4, 6, 8}) 339 | intersect := a.Intersect(b).(*Set[int]) 340 | must.MapContainsKeys(t, intersect.items, []int{4, 6}) 341 | }) 342 | } 343 | 344 | func TestSet_Remove(t *testing.T) { 345 | t.Run("empty remove item", func(t *testing.T) { 346 | s := New[int](10) 347 | must.False(t, s.Remove(32)) 348 | must.MapEmpty(t, s.items) 349 | }) 350 | 351 | t.Run("set remove item", func(t *testing.T) { 352 | s := From[string]([]string{"apple", "banana", "cherry"}) 353 | must.True(t, s.Remove("banana")) 354 | must.MapContainsKeys(t, s.items, []string{"apple", "cherry"}) 355 | }) 356 | 357 | t.Run("set remove missing", func(t *testing.T) { 358 | s := From[string]([]string{"apple", "banana", "cherry"}) 359 | must.False(t, s.Remove("zucchini")) 360 | must.MapContainsKeys(t, s.items, []string{"apple", "banana", "cherry"}) 361 | }) 362 | } 363 | 364 | func TestSet_RemoveSlice(t *testing.T) { 365 | t.Run("empty remove all", func(t *testing.T) { 366 | s := New[int](10) 367 | must.False(t, s.RemoveSlice([]int{1, 2, 3})) 368 | must.MapEmpty(t, s.items) 369 | }) 370 | 371 | t.Run("set remove nothing", func(t *testing.T) { 372 | s := From[int]([]int{1, 2, 3}) 373 | must.False(t, s.RemoveSlice(nil)) 374 | must.MapContainsKeys(t, s.items, []int{1, 2, 3}) 375 | }) 376 | 377 | t.Run("set remove some", func(t *testing.T) { 378 | s := From[int]([]int{1, 2, 3, 4, 5, 6}) 379 | must.True(t, s.RemoveSlice([]int{5, 6, 7, 8, 9})) 380 | must.MapContainsKeys(t, s.items, []int{1, 2, 3, 4}) 381 | }) 382 | } 383 | 384 | func TestSet_RemoveSet(t *testing.T) { 385 | t.Run("empty remove empty", func(t *testing.T) { 386 | a := New[int](0) 387 | b := New[int](0) 388 | must.False(t, a.RemoveSet(b)) 389 | must.MapEmpty(t, a.items) 390 | }) 391 | 392 | t.Run("empty remove some", func(t *testing.T) { 393 | a := New[int](0) 394 | b := From[int]([]int{1, 2, 3, 4}) 395 | must.False(t, a.RemoveSet(b)) 396 | must.MapEmpty(t, a.items) 397 | }) 398 | 399 | t.Run("set remove some", func(t *testing.T) { 400 | a := From[int]([]int{1, 2, 3, 4, 5, 6, 7, 8}) 401 | b := From[int]([]int{2, 4, 6, 8}) 402 | must.True(t, a.RemoveSet(b)) 403 | must.MapContainsKeys(t, a.items, []int{1, 3, 5, 7}) 404 | }) 405 | } 406 | 407 | func TestSet_RemoveFunc(t *testing.T) { 408 | t.Run("empty", func(t *testing.T) { 409 | a := New[int](10) 410 | modified := a.RemoveFunc(func(i int) bool { 411 | return i%2 == 0 412 | }) 413 | must.Empty(t, a) 414 | must.False(t, modified) 415 | }) 416 | 417 | t.Run("none match", func(t *testing.T) { 418 | a := From[int]([]int{1, 3, 5, 7, 9}) 419 | modified := a.RemoveFunc(func(i int) bool { 420 | return i%2 == 0 421 | }) 422 | must.True(t, a.ContainsSlice([]int{1, 3, 5, 7, 9})) 423 | must.False(t, modified) 424 | }) 425 | 426 | t.Run("some match", func(t *testing.T) { 427 | a := From[int]([]int{1, 2, 3, 4, 5, 6, 7, 8, 9}) 428 | modified := a.RemoveFunc(func(i int) bool { 429 | return i%2 == 0 430 | }) 431 | must.True(t, a.ContainsSlice([]int{1, 3, 5, 7, 9})) 432 | must.True(t, modified) 433 | }) 434 | 435 | t.Run("all match", func(t *testing.T) { 436 | a := From[int]([]int{1, 3, 5, 7, 9}) 437 | modified := a.RemoveFunc(func(i int) bool { 438 | return i%2 != 0 439 | }) 440 | must.Empty(t, a) 441 | must.True(t, modified) 442 | }) 443 | } 444 | 445 | func TestSet_Copy(t *testing.T) { 446 | t.Run("copy empty", func(t *testing.T) { 447 | a := New[int](0) 448 | b := a.Copy() 449 | must.MapEmpty(t, b.items) 450 | must.True(t, b.Insert(3)) 451 | must.MapEmpty(t, a.items) 452 | must.MapContainsKeys(t, b.items, []int{3}) 453 | }) 454 | 455 | t.Run("copy some", func(t *testing.T) { 456 | a := From[int]([]int{1, 2, 3, 4}) 457 | b := a.Copy() 458 | must.MapContainsKeys(t, b.items, []int{1, 2, 3, 4}) 459 | must.True(t, b.RemoveSlice([]int{1, 3})) 460 | must.MapContainsKeys(t, b.items, []int{2, 4}) 461 | must.MapContainsKeys(t, a.items, []int{1, 2, 3, 4}) 462 | }) 463 | } 464 | 465 | func TestSet_Slice(t *testing.T) { 466 | t.Run("slice empty", func(t *testing.T) { 467 | a := New[string](10) 468 | l := a.Slice() 469 | must.SliceEmpty(t, l) 470 | }) 471 | 472 | t.Run("slice set", func(t *testing.T) { 473 | a := From([]string{"apple", "banana", "cherry"}) 474 | l := a.Slice() 475 | must.Len(t, 3, l) 476 | must.SliceContains(t, l, "apple") 477 | must.SliceContains(t, l, "banana") 478 | must.SliceContains(t, l, "cherry") 479 | }) 480 | } 481 | 482 | func TestSet_String(t *testing.T) { 483 | t.Run("ints", func(t *testing.T) { 484 | a := From([]int{1, 2, 3}) 485 | result := a.String() 486 | must.Eq(t, "[1 2 3]", result) 487 | }) 488 | 489 | t.Run("custom", func(t *testing.T) { 490 | a := From([]*employee{{"bob", 2}, {"alice", 1}, {"carl", 3}}) 491 | result := a.String() 492 | must.Eq(t, "[(1 alice) (2 bob) (3 carl)]", result) 493 | }) 494 | } 495 | 496 | func TestSet_StringFunc(t *testing.T) { 497 | t.Run("empty", func(t *testing.T) { 498 | a := New[string](10) 499 | s := a.StringFunc(nil) 500 | must.Eq(t, "[]", s) 501 | }) 502 | 503 | t.Run("int", func(t *testing.T) { 504 | a := From([]int{5, 2, 5, 1, 3}) 505 | s := a.StringFunc(func(i int) string { 506 | return fmt.Sprintf("%d", i) 507 | }) 508 | must.Eq(t, "[1 2 3 5]", s) 509 | }) 510 | 511 | t.Run("custom", func(t *testing.T) { 512 | a := From([]employee{ 513 | {"mitchell", 1}, 514 | {"jack", 3}, 515 | {"armon", 2}, 516 | }) 517 | s := a.StringFunc(func(e employee) string { 518 | return fmt.Sprintf("(%d %s)", e.id, e.name) 519 | }) 520 | must.Eq(t, "[(1 mitchell) (2 armon) (3 jack)]", s) 521 | }) 522 | } 523 | 524 | func TestSet_Equal(t *testing.T) { 525 | t.Run("empty empty", func(t *testing.T) { 526 | a := New[int](0) 527 | b := New[int](10) 528 | must.True(t, a.Equal(b)) 529 | }) 530 | 531 | t.Run("empty some", func(t *testing.T) { 532 | a := New[int](0) 533 | b := From[int]([]int{1, 2, 3}) 534 | must.False(t, a.Equal(b)) 535 | }) 536 | 537 | t.Run("same", func(t *testing.T) { 538 | a := From[int]([]int{3, 2, 1}) 539 | b := From[int]([]int{1, 2, 3}) 540 | must.True(t, a.Equal(b)) 541 | }) 542 | 543 | t.Run("subset", func(t *testing.T) { 544 | a := From[int]([]int{2, 3}) 545 | b := From[int]([]int{1, 2, 3}) 546 | must.False(t, a.Equal(b)) 547 | must.False(t, b.Equal(a)) 548 | }) 549 | } 550 | 551 | func TestSet_Subset(t *testing.T) { 552 | t.Run("empty empty", func(t *testing.T) { 553 | a := New[int](0) 554 | b := New[int](0) 555 | must.True(t, a.Subset(b)) 556 | }) 557 | 558 | t.Run("empty some", func(t *testing.T) { 559 | a := New[int](0) 560 | b := From[int]([]int{1, 2, 3}) 561 | must.False(t, a.Subset(b)) 562 | }) 563 | 564 | t.Run("some empty", func(t *testing.T) { 565 | a := From[int]([]int{1, 2, 3}) 566 | b := New[int](0) 567 | must.True(t, a.Subset(b)) 568 | }) 569 | 570 | t.Run("equal", func(t *testing.T) { 571 | a := From[int]([]int{1, 2, 3}) 572 | b := From[int]([]int{2, 3, 1}) 573 | must.True(t, a.Subset(b)) 574 | }) 575 | 576 | t.Run("subset", func(t *testing.T) { 577 | a := From[int]([]int{1, 2, 3}) 578 | b := From[int]([]int{3, 1}) 579 | must.True(t, a.Subset(b)) 580 | }) 581 | 582 | t.Run("superset", func(t *testing.T) { 583 | a := From[int]([]int{1, 2, 3}) 584 | b := From[int]([]int{3, 1, 2, 4}) 585 | must.False(t, a.Subset(b)) 586 | }) 587 | } 588 | 589 | func TestSet_ProperSubset(t *testing.T) { 590 | t.Run("empty empty", func(t *testing.T) { 591 | a := New[int](0) 592 | b := New[int](0) 593 | must.True(t, a.Subset(b)) 594 | }) 595 | 596 | t.Run("empty some", func(t *testing.T) { 597 | a := New[int](0) 598 | b := From[int]([]int{1, 2, 3}) 599 | must.False(t, a.ProperSubset(b)) 600 | }) 601 | 602 | t.Run("some empty", func(t *testing.T) { 603 | a := From[int]([]int{1, 2, 3}) 604 | b := New[int](0) 605 | must.True(t, a.ProperSubset(b)) 606 | }) 607 | 608 | t.Run("equal", func(t *testing.T) { 609 | a := From[int]([]int{1, 2, 3}) 610 | b := From[int]([]int{1, 2, 3}) 611 | must.False(t, a.ProperSubset(b)) 612 | must.False(t, b.ProperSubset(a)) 613 | }) 614 | 615 | t.Run("proper subset", func(t *testing.T) { 616 | a := From[int]([]int{1, 2, 3}) 617 | b := From[int]([]int{1, 2}) 618 | must.True(t, a.ProperSubset(b)) 619 | must.False(t, b.ProperSubset(a)) 620 | }) 621 | } 622 | 623 | func TestSet_EqualSet(t *testing.T) { 624 | t.Run("empty empty", func(t *testing.T) { 625 | a := New[int](0) 626 | b := New[int](0) 627 | must.True(t, a.EqualSet(b)) 628 | }) 629 | 630 | t.Run("empty some", func(t *testing.T) { 631 | a := New[int](0) 632 | b := From[int]([]int{1, 2, 3}) 633 | must.False(t, a.EqualSet(b)) 634 | must.False(t, b.EqualSet(a)) 635 | }) 636 | 637 | t.Run("same", func(t *testing.T) { 638 | a := From[int]([]int{1, 2, 3}) 639 | b := From[int]([]int{3, 2, 1}) 640 | must.True(t, a.EqualSet(b)) 641 | must.True(t, b.EqualSet(a)) 642 | }) 643 | 644 | t.Run("different", func(t *testing.T) { 645 | a := From[int]([]int{1, 2, 3}) 646 | b := From[int]([]int{1, 2, 4}) 647 | must.False(t, a.EqualSet(b)) 648 | must.False(t, b.EqualSet(a)) 649 | }) 650 | } 651 | 652 | func TestSet_EqualSlice(t *testing.T) { 653 | t.Run("empty empty", func(t *testing.T) { 654 | a := New[int](0) 655 | b := make([]int, 0) 656 | must.True(t, a.EqualSlice(b)) 657 | }) 658 | 659 | t.Run("empty some", func(t *testing.T) { 660 | a := New[int](0) 661 | b := []int{1, 2, 3} 662 | must.False(t, a.EqualSlice(b)) 663 | }) 664 | 665 | t.Run("some empty", func(t *testing.T) { 666 | a := From[int]([]int{1, 2, 3}) 667 | b := make([]int, 0) 668 | must.False(t, a.EqualSlice(b)) 669 | }) 670 | 671 | t.Run("equal", func(t *testing.T) { 672 | a := From[int]([]int{1, 2, 3}) 673 | b := []int{3, 2, 1} 674 | must.True(t, a.EqualSlice(b)) 675 | }) 676 | 677 | t.Run("not equal", func(t *testing.T) { 678 | a := From[int]([]int{1, 2, 3}) 679 | b := []int{2, 3, 4} 680 | must.False(t, a.EqualSlice(b)) 681 | }) 682 | 683 | t.Run("subset", func(t *testing.T) { 684 | a := From[int]([]int{1, 2, 3, 4, 5}) 685 | b := []int{2, 3, 4} 686 | must.False(t, a.EqualSlice(b)) 687 | }) 688 | 689 | t.Run("superset", func(t *testing.T) { 690 | a := From[int]([]int{2, 3, 4}) 691 | b := []int{1, 2, 3, 4, 5} 692 | must.False(t, a.EqualSlice(b)) 693 | }) 694 | 695 | t.Run("duplicates", func(t *testing.T) { 696 | a := From[int]([]int{1, 2, 3, 4, 5}) 697 | b := []int{1, 2, 2, 3, 3, 4, 5} 698 | must.True(t, a.EqualSlice(b)) 699 | }) 700 | } 701 | 702 | func TestSet_EqualSliceSet(t *testing.T) { 703 | t.Run("empty empty", func(t *testing.T) { 704 | a := New[int](0) 705 | b := make([]int, 0) 706 | must.True(t, a.EqualSliceSet(b)) 707 | }) 708 | 709 | t.Run("empty some", func(t *testing.T) { 710 | a := New[int](0) 711 | b := []int{1, 2, 3} 712 | must.False(t, a.EqualSliceSet(b)) 713 | }) 714 | 715 | t.Run("some empty", func(t *testing.T) { 716 | a := From[int]([]int{1, 2, 3}) 717 | b := make([]int, 0) 718 | must.False(t, a.EqualSliceSet(b)) 719 | }) 720 | 721 | t.Run("equal", func(t *testing.T) { 722 | a := From[int]([]int{1, 2, 3}) 723 | b := []int{3, 2, 1} 724 | must.True(t, a.EqualSliceSet(b)) 725 | }) 726 | 727 | t.Run("not equal", func(t *testing.T) { 728 | a := From[int]([]int{1, 2, 3}) 729 | b := []int{2, 3, 4} 730 | must.False(t, a.EqualSliceSet(b)) 731 | }) 732 | 733 | t.Run("subset", func(t *testing.T) { 734 | a := From[int]([]int{1, 2, 3, 4, 5}) 735 | b := []int{2, 3, 4} 736 | must.False(t, a.EqualSliceSet(b)) 737 | }) 738 | 739 | t.Run("superset", func(t *testing.T) { 740 | a := From[int]([]int{2, 3, 4}) 741 | b := []int{1, 2, 3, 4, 5} 742 | must.False(t, a.EqualSliceSet(b)) 743 | }) 744 | 745 | t.Run("duplicates", func(t *testing.T) { 746 | a := From[int]([]int{1, 2, 3, 4, 5}) 747 | b := []int{1, 2, 2, 3, 3, 4, 5} 748 | must.False(t, a.EqualSliceSet(b)) 749 | }) 750 | } 751 | 752 | func TestSet_Items(t *testing.T) { 753 | s := From[int]([]int{1, 2, 3, 4, 5}) 754 | 755 | sum := 0 756 | for element := range s.Items() { 757 | sum += element 758 | } 759 | 760 | must.Eq(t, 15, sum) 761 | } 762 | -------------------------------------------------------------------------------- /stack.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | type stack[T any] struct { 7 | top *object[T] 8 | } 9 | 10 | type object[T any] struct { 11 | item T 12 | next *object[T] 13 | } 14 | 15 | func makeStack[T any]() *stack[T] { 16 | return new(stack[T]) 17 | } 18 | 19 | func (s *stack[T]) push(item T) { 20 | obj := &object[T]{ 21 | item: item, 22 | next: s.top, 23 | } 24 | s.top = obj 25 | } 26 | 27 | func (s *stack[T]) pop() T { 28 | obj := s.top 29 | s.top = obj.next 30 | obj.next = nil 31 | return obj.item 32 | } 33 | 34 | func (s *stack[T]) empty() bool { 35 | return s.top == nil 36 | } 37 | -------------------------------------------------------------------------------- /stack_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "testing" 8 | 9 | "github.com/shoenig/test/must" 10 | ) 11 | 12 | func TestStack_simple(t *testing.T) { 13 | s := makeStack[int]() 14 | must.True(t, s.empty()) 15 | 16 | s.push(1) 17 | must.False(t, s.empty()) 18 | 19 | value := s.pop() 20 | must.Eq(t, 1, value) 21 | must.True(t, s.empty()) 22 | } 23 | 24 | func TestStack_complex(t *testing.T) { 25 | s := makeStack[byte]() 26 | 27 | s.push('a') 28 | s.push('b') 29 | s.push('c') 30 | s.push('d') 31 | s.push('e') 32 | s.push('f') 33 | 34 | must.Eq(t, 'f', s.pop()) 35 | must.Eq(t, 'e', s.pop()) 36 | must.Eq(t, 'd', s.pop()) 37 | 38 | s.push('x') 39 | s.push('y') 40 | 41 | must.Eq(t, 'y', s.pop()) 42 | 43 | s.push('z') 44 | 45 | must.Eq(t, 'z', s.pop()) 46 | must.Eq(t, 'x', s.pop()) 47 | must.Eq(t, 'c', s.pop()) 48 | must.Eq(t, 'b', s.pop()) 49 | must.Eq(t, 'a', s.pop()) 50 | must.True(t, s.empty()) 51 | } 52 | -------------------------------------------------------------------------------- /treeset.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "fmt" 8 | "iter" 9 | ) 10 | 11 | // CompareFunc represents a function that compares two elements. 12 | // 13 | // Must return 14 | // < 0 if the first parameter is less than the second parameter 15 | // 0 if the two parameters are equal 16 | // > 0 if the first parameters is greater than the second parameter 17 | // 18 | // Often T will be a type that satisfies cmp.Ordered, and CompareFunc can 19 | // be implemented by using cmp.Compare. 20 | type CompareFunc[T any] func(T, T) int 21 | 22 | // TreeSet provides a generic sortable set implementation for Go. 23 | // Enables fast storage and retrieval of ordered information. Most effective 24 | // in cases where data is regularly being added and/or removed and fast 25 | // lookup properties must be maintained. 26 | // 27 | // The underlying data structure is a Red-Black Binary Search Tree. 28 | // https://en.wikipedia.org/wiki/Red–black_tree 29 | // 30 | // Not thread safe, and not safe for concurrent modification. 31 | type TreeSet[T any] struct { 32 | comparison CompareFunc[T] 33 | root *node[T] 34 | marker *node[T] 35 | size int 36 | } 37 | 38 | // NewTreeSet creates a TreeSet of type T, comparing elements via a given 39 | // CompareFunc[T]. 40 | // 41 | // T may be any type. 42 | // 43 | // For builtin types, CompareBuiltin provides a convenient CompareFunc implementation. 44 | func NewTreeSet[T any](compare CompareFunc[T]) *TreeSet[T] { 45 | return &TreeSet[T]{ 46 | comparison: compare, 47 | root: nil, 48 | marker: &node[T]{color: black}, 49 | size: 0, 50 | } 51 | } 52 | 53 | // TreeSetFrom creates a new TreeSet containing each item in items. 54 | // 55 | // T may be any type. 56 | // 57 | // C is an implementation of CompareFunc[T]. For builtin types, Cmp provides a 58 | // convenient Compare implementation. 59 | func TreeSetFrom[T any](items []T, compare CompareFunc[T]) *TreeSet[T] { 60 | s := NewTreeSet[T](compare) 61 | s.InsertSlice(items) 62 | return s 63 | } 64 | 65 | // Insert item into s. 66 | // 67 | // Returns true if s was modified (item was not already in s), false otherwise. 68 | func (s *TreeSet[T]) Insert(item T) bool { 69 | return s.insert(&node[T]{ 70 | element: item, 71 | color: red, 72 | }) 73 | } 74 | 75 | // InsertSlice will insert each item in items into s. 76 | // 77 | // Return true if s was modified (at least one item was not already in s), false otherwise. 78 | func (s *TreeSet[T]) InsertSlice(items []T) bool { 79 | modified := false 80 | for _, item := range items { 81 | if s.Insert(item) { 82 | modified = true 83 | } 84 | } 85 | return modified 86 | } 87 | 88 | // InsertSet will insert each element of col into s. 89 | // 90 | // Return true if s was modified (at least one item of o was not already in s), false otherwise. 91 | func (s *TreeSet[T]) InsertSet(col Collection[T]) bool { 92 | modified := false 93 | for item := range col.Items() { 94 | if s.Insert(item) { 95 | modified = true 96 | } 97 | } 98 | return modified 99 | } 100 | 101 | // Remove item from s. 102 | // 103 | // Returns true if s was modified (item was in s), false otherwise. 104 | func (s *TreeSet[T]) Remove(item T) bool { 105 | return s.delete(item) 106 | } 107 | 108 | // RemoveSlice will remove each item in items from s. 109 | // 110 | // Return true if s was modified (any item was in s), false otherwise. 111 | func (s *TreeSet[T]) RemoveSlice(items []T) bool { 112 | modified := false 113 | for _, item := range items { 114 | if s.Remove(item) { 115 | modified = true 116 | } 117 | } 118 | return modified 119 | } 120 | 121 | // RemoveSet will remove each element in col from s. 122 | // 123 | // Returns true if s was modified (at least one item in o was in s), false otherwise. 124 | func (s *TreeSet[T]) RemoveSet(col Collection[T]) bool { 125 | return removeSet(s, col) 126 | } 127 | 128 | // RemoveFunc will remove each element from s that satisifies condition f. 129 | // 130 | // Return true if s was modified, false otherwise. 131 | func (s *TreeSet[T]) RemoveFunc(f func(T) bool) bool { 132 | return removeFunc(s, f) 133 | } 134 | 135 | // Min returns the smallest item in the set. 136 | // 137 | // Must not be called on an empty set. 138 | func (s *TreeSet[T]) Min() T { 139 | if s.root == nil { 140 | panic("min: tree is empty") 141 | } 142 | n := s.min(s.root) 143 | return n.element 144 | } 145 | 146 | // Max returns the largest item in s. 147 | // 148 | // Must not be called on an empty set. 149 | func (s *TreeSet[T]) Max() T { 150 | if s.root == nil { 151 | panic("max: tree is empty") 152 | } 153 | n := s.max(s.root) 154 | return n.element 155 | } 156 | 157 | // TopK returns the top n (smallest) elements in s, in ascending order. 158 | func (s *TreeSet[T]) TopK(n int) []T { 159 | result := make([]T, 0, n) 160 | s.fillLeft(s.root, &result) 161 | return result 162 | } 163 | 164 | // BottomK returns the bottom n (largest) elements in s, in descending order. 165 | func (s *TreeSet[T]) BottomK(n int) []T { 166 | result := make([]T, 0, n) 167 | s.fillRight(s.root, &result) 168 | return result 169 | } 170 | 171 | // FirstBelow returns the first element strictly below item. 172 | // 173 | // A zero value and false are returned if no such element exists. 174 | func (s *TreeSet[T]) FirstBelow(item T) (T, bool) { 175 | var candidate *node[T] = nil 176 | var n = s.root 177 | for n != nil { 178 | c := s.comparison(item, n.element) 179 | switch { 180 | case c > 0: 181 | candidate = n 182 | n = n.right 183 | case c <= 0: 184 | n = n.left 185 | } 186 | } 187 | return candidate.get() 188 | } 189 | 190 | // FirstBelowEqual returns the first element below item (or item itself if present). 191 | // 192 | // A zero value and false are returned if no such element exists. 193 | func (s *TreeSet[T]) FirstBelowEqual(item T) (T, bool) { 194 | var candidate *node[T] = nil 195 | var n = s.root 196 | for n != nil { 197 | c := s.comparison(item, n.element) 198 | switch { 199 | case c == 0: 200 | return n.get() 201 | case c > 0: 202 | candidate = n 203 | n = n.right 204 | case c < 0: 205 | n = n.left 206 | } 207 | } 208 | return candidate.get() 209 | } 210 | 211 | // Below returns a TreeSet containing the elements of s that are < item. 212 | func (s *TreeSet[T]) Below(item T) *TreeSet[T] { 213 | result := NewTreeSet[T](s.comparison) 214 | s.filterLeft(s.root, func(element T) bool { 215 | return s.comparison(element, item) < 0 216 | }, result) 217 | return result 218 | } 219 | 220 | // BelowEqual returns a TreeSet containing the elements of s that are ≤ item. 221 | func (s *TreeSet[T]) BelowEqual(item T) *TreeSet[T] { 222 | result := NewTreeSet[T](s.comparison) 223 | s.filterLeft(s.root, func(element T) bool { 224 | return s.comparison(element, item) <= 0 225 | }, result) 226 | return result 227 | } 228 | 229 | // FirstAbove returns the first element strictly above item. 230 | // 231 | // A zero value and false are returned if no such element exists. 232 | func (s *TreeSet[T]) FirstAbove(item T) (T, bool) { 233 | var candidate *node[T] = nil 234 | var n = s.root 235 | for n != nil { 236 | c := s.comparison(item, n.element) 237 | switch { 238 | case c < 0: 239 | candidate = n 240 | n = n.left 241 | case c >= 0: 242 | n = n.right 243 | } 244 | } 245 | return candidate.get() 246 | } 247 | 248 | // FirstAboveEqual returns the first element above item (or item itself if present). 249 | // 250 | // A zero value and false are returned if no such element exists. 251 | func (s *TreeSet[T]) FirstAboveEqual(item T) (T, bool) { 252 | var candidate *node[T] 253 | var n = s.root 254 | for n != nil { 255 | c := s.comparison(item, n.element) 256 | switch { 257 | case c == 0: 258 | return n.get() 259 | case c < 0: 260 | candidate = n 261 | n = n.left 262 | case c > 0: 263 | n = n.right 264 | } 265 | } 266 | return candidate.get() 267 | } 268 | 269 | // After returns a TreeSet containing the elements of s that are > item. 270 | func (s *TreeSet[T]) Above(item T) *TreeSet[T] { 271 | result := NewTreeSet[T](s.comparison) 272 | s.filterRight(s.root, func(element T) bool { 273 | return s.comparison(element, item) > 0 274 | }, result) 275 | return result 276 | } 277 | 278 | // AfterEqual returns a TreeSet containing the elements of s that are ≥ item. 279 | func (s *TreeSet[T]) AboveEqual(item T) *TreeSet[T] { 280 | result := NewTreeSet[T](s.comparison) 281 | s.filterRight(s.root, func(element T) bool { 282 | return s.comparison(element, item) >= 0 283 | }, result) 284 | return result 285 | } 286 | 287 | // Contains returns whether item is present in s. 288 | func (s *TreeSet[T]) Contains(item T) bool { 289 | return s.locate(s.root, item) != nil 290 | } 291 | 292 | // ContainsSlice returns whether all elements in items are present in s. 293 | func (s *TreeSet[T]) ContainsSlice(items []T) bool { 294 | return containsSlice(s, items) 295 | } 296 | 297 | // Size returns the number of elements in s. 298 | func (s *TreeSet[T]) Size() int { 299 | return s.size 300 | } 301 | 302 | // Empty returns true if there are no elements in s. 303 | func (s *TreeSet[T]) Empty() bool { 304 | return s.Size() == 0 305 | } 306 | 307 | // Slice returns the elements of s as a slice, in order. 308 | func (s *TreeSet[T]) Slice() []T { 309 | result := make([]T, 0, s.Size()) 310 | for item := range s.Items() { 311 | result = append(result, item) 312 | } 313 | return result 314 | } 315 | 316 | // Subset returns whether col is a subset of s. 317 | func (s *TreeSet[T]) Subset(col Collection[T]) bool { 318 | // try the fast paths 319 | if col.Empty() { 320 | return true 321 | } 322 | if s.Empty() { 323 | return false 324 | } 325 | if s.Size() < col.Size() { 326 | return false 327 | } 328 | 329 | // iterate o, and increment s finding each element 330 | // i.e. merge algorithm but with channels 331 | iterO := col.(*TreeSet[T]).iterate() 332 | iterS := s.iterate() 333 | 334 | idxO := 0 335 | idxS := 0 336 | 337 | next: 338 | for ; idxO < col.Size(); idxO++ { 339 | nextO := iterO() 340 | for idxS < s.Size() { 341 | idxS++ 342 | nextS := iterS() 343 | cmp := s.compare(nextS, nextO) 344 | switch { 345 | case cmp > 0: 346 | return false 347 | case cmp < 0: 348 | continue 349 | default: 350 | continue next 351 | } 352 | } 353 | return false 354 | } 355 | return true 356 | } 357 | 358 | // ProperSubset returns whether col is a proper subset of s. 359 | func (s *TreeSet[T]) ProperSubset(col Collection[T]) bool { 360 | if s.Size() <= col.Size() { 361 | return false 362 | } 363 | return s.Subset(col) 364 | } 365 | 366 | // Union returns a set that contains all elements of s and col combined. 367 | func (s *TreeSet[T]) Union(col Collection[T]) Collection[T] { 368 | tree := NewTreeSet[T](s.comparison) 369 | f := func(n *node[T]) { tree.Insert(n.element) } 370 | s.prefix(f, s.root) 371 | oSet := col.(*TreeSet[T]) 372 | oSet.prefix(f, oSet.root) 373 | return tree 374 | } 375 | 376 | // Difference returns a set that contains elements of s that are not in col. 377 | func (s *TreeSet[T]) Difference(col Collection[T]) Collection[T] { 378 | tree := NewTreeSet[T](s.comparison) 379 | f := func(n *node[T]) { 380 | if !col.Contains(n.element) { 381 | tree.Insert(n.element) 382 | } 383 | } 384 | s.prefix(f, s.root) 385 | return tree 386 | } 387 | 388 | // Intersect returns a set that contains elements that are present in both s and col. 389 | func (s *TreeSet[T]) Intersect(col Collection[T]) Collection[T] { 390 | tree := NewTreeSet[T](s.comparison) 391 | f := func(n *node[T]) { 392 | if col.Contains(n.element) { 393 | tree.Insert(n.element) 394 | } 395 | } 396 | s.prefix(f, s.root) 397 | return tree 398 | } 399 | 400 | // Copy creates a copy of s. 401 | // 402 | // Individual elements are reference copies. 403 | func (s *TreeSet[T]) Copy() *TreeSet[T] { 404 | tree := NewTreeSet[T](s.comparison) 405 | f := func(n *node[T]) { 406 | tree.Insert(n.element) 407 | } 408 | s.prefix(f, s.root) 409 | return tree 410 | } 411 | 412 | // Equal return whether s and o contain the same elements. 413 | func (s *TreeSet[T]) Equal(o *TreeSet[T]) bool { 414 | // try the fast fail paths 415 | if s.Empty() || o.Empty() { 416 | return s.Size() == o.Size() 417 | } 418 | switch { 419 | case s.Size() != o.Size(): 420 | return false 421 | case s.comparison(s.Min(), o.Min()) != 0: 422 | return false 423 | case s.comparison(s.Max(), o.Max()) != 0: 424 | return false 425 | } 426 | 427 | iterS := s.iterate() 428 | iterO := o.iterate() 429 | for i := 0; i < s.Size(); i++ { 430 | nextS := iterS() 431 | nextO := iterO() 432 | if s.compare(nextS, nextO) != 0 { 433 | return false 434 | } 435 | } 436 | 437 | return true 438 | } 439 | 440 | // EqualSet returns s and col contain the same elements. 441 | func (s *TreeSet[T]) EqualSet(col Collection[T]) bool { 442 | return equalSet(s, col) 443 | } 444 | 445 | // EqualSlice returns whether s and items contain the same elements. 446 | // 447 | // The items slice may contain duplicates. 448 | // 449 | // If the items slice is known to contain no duplicates, EqualSliceSet may be 450 | // used instead as a faster implementation. 451 | // 452 | // To detect if a slice is a subset of s, use ContainsSlice. 453 | func (s *TreeSet[T]) EqualSlice(items []T) bool { 454 | other := TreeSetFrom[T](items, s.comparison) 455 | return s.Equal(other) 456 | } 457 | 458 | // EqualSliceSet returns whether s and items contain exactly the same elements. 459 | // 460 | // If items contains duplicates EqualSliceSet will return false. The elements of 461 | // items are assumed to be set-like. For comparing s to a slice that may contain 462 | // duplicate elements, use EqualSlice instead. 463 | // 464 | // To detect if a slice is a subset of s, use ContainsSlice. 465 | func (s *TreeSet[T]) EqualSliceSet(items []T) bool { 466 | // TODO optimize 467 | if s.Size() != len(items) { 468 | return false 469 | } 470 | return s.EqualSlice(items) 471 | } 472 | 473 | // String creates a string representation of s, using "%v" printf formatting 474 | // each element into a string. The result contains elements in order. 475 | func (s *TreeSet[T]) String() string { 476 | return s.StringFunc(func(element T) string { 477 | return fmt.Sprintf("%v", element) 478 | }) 479 | } 480 | 481 | // StringFunc creates a string representation of s, using f to transform each 482 | // element into a string. The result contains elements in order. 483 | func (s *TreeSet[T]) StringFunc(f func(T) string) string { 484 | l := make([]string, 0, s.Size()) 485 | for item := range s.Items() { 486 | l = append(l, f(item)) 487 | } 488 | return fmt.Sprintf("%s", l) 489 | } 490 | 491 | // Items returns a generator function for iterating each element in s by using 492 | // the range keyword. 493 | // 494 | // for i, element := range s.Items() { ... } 495 | func (s *TreeSet[T]) Items() iter.Seq[T] { 496 | return func(yield func(T) bool) { 497 | iter := s.iterate() 498 | n := iter() 499 | for i := 0; n != nil; i++ { 500 | if !yield(n.element) { 501 | return 502 | } 503 | n = iter() 504 | } 505 | } 506 | } 507 | 508 | // Red-Black Tree Invariants 509 | // 510 | // 1. each node is either red or black 511 | // 2. the root node is always black 512 | // 3. nil leaf nodes are always black 513 | // 4. a red node must not have red children 514 | // 5. all simple paths from a node to nil leaf contain the same number of 515 | // black nodes 516 | 517 | type color bool 518 | 519 | const ( 520 | red color = false 521 | black color = true 522 | ) 523 | 524 | type node[T any] struct { 525 | element T 526 | color color 527 | parent *node[T] 528 | left *node[T] 529 | right *node[T] 530 | } 531 | 532 | func (n *node[T]) black() bool { 533 | return n == nil || n.color == black 534 | } 535 | 536 | func (n *node[T]) red() bool { 537 | return n != nil && n.color == red 538 | } 539 | 540 | func (n *node[T]) get() (T, bool) { 541 | if n == nil { 542 | var zero T 543 | return zero, false 544 | } 545 | return n.element, true 546 | } 547 | 548 | func (s *TreeSet[T]) locate(start *node[T], target T) *node[T] { 549 | n := start 550 | for { 551 | if n == nil { 552 | return nil 553 | } 554 | cmp := s.compare(n, &node[T]{element: target}) 555 | switch { 556 | case cmp < 0: 557 | n = n.right 558 | case cmp > 0: 559 | n = n.left 560 | default: 561 | return n 562 | } 563 | } 564 | } 565 | 566 | func (s *TreeSet[T]) rotateRight(n *node[T]) { 567 | parent := n.parent 568 | leftChild := n.left 569 | 570 | n.left = leftChild.right 571 | if leftChild.right != nil { 572 | leftChild.right.parent = n 573 | } 574 | 575 | leftChild.right = n 576 | n.parent = leftChild 577 | 578 | s.replaceChild(parent, n, leftChild) 579 | } 580 | 581 | func (s *TreeSet[T]) rotateLeft(n *node[T]) { 582 | parent := n.parent 583 | rightChild := n.right 584 | 585 | n.right = rightChild.left 586 | if rightChild.left != nil { 587 | rightChild.left.parent = n 588 | } 589 | 590 | rightChild.left = n 591 | n.parent = rightChild 592 | 593 | s.replaceChild(parent, n, rightChild) 594 | } 595 | 596 | func (s *TreeSet[T]) replaceChild(parent, previous, next *node[T]) { 597 | switch { 598 | case parent == nil: 599 | s.root = next 600 | case parent.left == previous: 601 | parent.left = next 602 | case parent.right == previous: 603 | parent.right = next 604 | default: 605 | panic("node is not child of its parent") 606 | } 607 | 608 | if next != nil { 609 | next.parent = parent 610 | } 611 | } 612 | 613 | func (s *TreeSet[T]) insert(n *node[T]) bool { 614 | var ( 615 | parent *node[T] = nil 616 | tmp *node[T] = s.root 617 | ) 618 | 619 | for tmp != nil { 620 | parent = tmp 621 | 622 | cmp := s.compare(n, tmp) 623 | switch { 624 | case cmp < 0: 625 | tmp = tmp.left 626 | case cmp > 0: 627 | tmp = tmp.right 628 | default: 629 | // already exists in tree 630 | return false 631 | } 632 | } 633 | 634 | n.color = red 635 | switch { 636 | case parent == nil: 637 | s.root = n 638 | case s.compare(n, parent) < 0: 639 | parent.left = n 640 | default: 641 | parent.right = n 642 | } 643 | n.parent = parent 644 | 645 | s.rebalanceInsertion(n) 646 | s.size++ 647 | return true 648 | } 649 | 650 | func (s *TreeSet[T]) rebalanceInsertion(n *node[T]) { 651 | parent := n.parent 652 | 653 | // case 1: parent is nil 654 | // - means we are the root 655 | // - our color must be black 656 | if parent == nil { 657 | n.color = black 658 | return 659 | } 660 | 661 | // if parent is black there is nothing to do 662 | if parent.black() { 663 | return 664 | } 665 | 666 | // case 2: no grandparent 667 | // - implies the parent is root 668 | // - we must now be black 669 | grandparent := parent.parent 670 | if grandparent == nil { 671 | parent.color = black 672 | return 673 | } 674 | 675 | uncle := s.uncleOf(parent) 676 | 677 | switch { 678 | // case 3: uncle is red 679 | // - fix color of parent, grandparent, uncle 680 | // - recurse upwards as necessary 681 | case uncle != nil && uncle.red(): 682 | parent.color = black 683 | grandparent.color = red 684 | uncle.color = black 685 | s.rebalanceInsertion(grandparent) 686 | 687 | case parent == grandparent.left: 688 | // case 4a: uncle is black 689 | // + node is left->right child of its grandparent 690 | if n == parent.right { 691 | s.rotateLeft(parent) 692 | parent = n // recolor in case 5a 693 | } 694 | 695 | // case 5a: uncle is black 696 | // + node is left->left child of its grandparent 697 | s.rotateRight(grandparent) 698 | 699 | // fix color of original parent and grandparent 700 | parent.color = black 701 | grandparent.color = red 702 | 703 | // parent is right child of grandparent 704 | default: 705 | // case 4b: uncle is black 706 | // + node is right->left child of its grandparent 707 | if n == parent.left { 708 | s.rotateRight(parent) 709 | // points to root of rotated sub tree 710 | parent = n // recolor in case 5b 711 | } 712 | 713 | // case 5b: uncle is black 714 | // + node is right->right child of its grandparent 715 | s.rotateLeft(grandparent) 716 | 717 | // fix color of original parent and grandparent 718 | parent.color = black 719 | grandparent.color = red 720 | } 721 | } 722 | 723 | func (s *TreeSet[T]) delete(element T) bool { 724 | n := s.locate(s.root, element) 725 | if n == nil { 726 | return false 727 | } 728 | 729 | var ( 730 | moved *node[T] 731 | deleted color 732 | ) 733 | 734 | if n.left == nil || n.right == nil { 735 | // case where deleted node had zero or one child 736 | moved = s.delete01(n) 737 | deleted = n.color 738 | } else { 739 | // case where node has two children 740 | 741 | // find minimum of right subtree 742 | successor := s.min(n.right) 743 | 744 | // copy successor data into n 745 | n.element = successor.element 746 | 747 | // delete successor 748 | moved = s.delete01(successor) 749 | deleted = successor.color 750 | } 751 | 752 | // re-balance if the node was black 753 | if deleted == black { 754 | s.rebalanceDeletion(moved) 755 | 756 | // remove marker 757 | if moved == s.marker { 758 | s.replaceChild(moved.parent, moved, nil) 759 | } 760 | } 761 | 762 | // element was removed 763 | s.size-- 764 | s.marker.color = black 765 | s.marker.left = nil 766 | s.marker.right = nil 767 | s.marker.parent = nil 768 | return true 769 | } 770 | 771 | func (s *TreeSet[T]) delete01(n *node[T]) *node[T] { 772 | // node only has left child, replace by left child 773 | if n.left != nil { 774 | s.replaceChild(n.parent, n, n.left) 775 | return n.left 776 | } 777 | 778 | // node only has right child, replace by right child 779 | if n.right != nil { 780 | s.replaceChild(n.parent, n, n.right) 781 | return n.right 782 | } 783 | 784 | // node has both children 785 | // if node is black replace with marker 786 | // if node is red we just remove it 787 | if n.black() { 788 | s.replaceChild(n.parent, n, s.marker) 789 | return s.marker 790 | } else { 791 | s.replaceChild(n.parent, n, nil) 792 | return nil 793 | } 794 | } 795 | 796 | func (s *TreeSet[T]) rebalanceDeletion(n *node[T]) { 797 | // base case: node is root 798 | if n == s.root { 799 | n.color = black 800 | return 801 | } 802 | 803 | sibling := s.siblingOf(n) 804 | 805 | // case: sibling is red 806 | if sibling.red() { 807 | s.fixRedSibling(n, sibling) 808 | sibling = s.siblingOf(n) 809 | } 810 | 811 | // case: black sibling with two black children 812 | if sibling.left.black() && sibling.right.black() { 813 | sibling.color = red 814 | 815 | // case: black sibling with to black children and a red parent 816 | if n.parent.red() { 817 | n.parent.color = black 818 | } else { 819 | // case: black sibling with two black children and black parent 820 | s.rebalanceDeletion(n.parent) 821 | } 822 | } else { 823 | // case: black sibling with at least one red child 824 | s.fixBlackSibling(n, sibling) 825 | } 826 | } 827 | 828 | func (s *TreeSet[T]) fixRedSibling(n *node[T], sibling *node[T]) { 829 | sibling.color = black 830 | n.parent.color = red 831 | 832 | switch { 833 | case n == n.parent.left: 834 | s.rotateLeft(n.parent) 835 | default: 836 | s.rotateRight(n.parent) 837 | } 838 | } 839 | 840 | func (s *TreeSet[T]) fixBlackSibling(n, sibling *node[T]) { 841 | isLeftChild := n == n.parent.left 842 | 843 | if isLeftChild && sibling.right.black() { 844 | sibling.left.color = black 845 | sibling.color = red 846 | s.rotateRight(sibling) 847 | sibling = n.parent.right 848 | } else if !isLeftChild && sibling.left.black() { 849 | sibling.right.color = black 850 | sibling.color = red 851 | s.rotateLeft(sibling) 852 | sibling = n.parent.left 853 | } 854 | 855 | sibling.color = n.parent.color 856 | n.parent.color = black 857 | if isLeftChild { 858 | sibling.right.color = black 859 | s.rotateLeft(n.parent) 860 | } else { 861 | sibling.left.color = black 862 | s.rotateRight(n.parent) 863 | } 864 | } 865 | 866 | func (s *TreeSet[T]) siblingOf(n *node[T]) *node[T] { 867 | parent := n.parent 868 | switch { 869 | case n == parent.left: 870 | return parent.right 871 | case n == parent.right: 872 | return parent.left 873 | default: 874 | panic("bug: parent is not a child of its grandparent") 875 | } 876 | } 877 | 878 | func (*TreeSet[T]) uncleOf(n *node[T]) *node[T] { 879 | grandparent := n.parent 880 | switch { 881 | case grandparent.left == n: 882 | return grandparent.right 883 | case grandparent.right == n: 884 | return grandparent.left 885 | default: 886 | panic("bug: parent is not a child of our childs grandparent") 887 | } 888 | } 889 | 890 | func (s *TreeSet[T]) min(n *node[T]) *node[T] { 891 | for n.left != nil { 892 | n = n.left 893 | } 894 | return n 895 | } 896 | 897 | func (s *TreeSet[T]) max(n *node[T]) *node[T] { 898 | for n.right != nil { 899 | n = n.right 900 | } 901 | return n 902 | } 903 | 904 | func (s *TreeSet[T]) compare(a, b *node[T]) int { 905 | return s.comparison(a.element, b.element) 906 | } 907 | 908 | // TreeNodeVisit is a function that is called for each node in the tree. 909 | type TreeNodeVisit[T any] func(*node[T]) (next bool) 910 | 911 | func (s *TreeSet[T]) infix(visit TreeNodeVisit[T], n *node[T]) (next bool) { 912 | if n == nil { 913 | return true 914 | } 915 | if next = s.infix(visit, n.left); !next { 916 | return 917 | } 918 | if next = visit(n); !next { 919 | return 920 | } 921 | return s.infix(visit, n.right) 922 | } 923 | 924 | func (s *TreeSet[T]) fillLeft(n *node[T], k *[]T) { 925 | if n == nil { 926 | return 927 | } 928 | 929 | if len(*k) < cap(*k) { 930 | s.fillLeft(n.left, k) 931 | } 932 | 933 | if len(*k) < cap(*k) { 934 | *k = append(*k, n.element) 935 | } 936 | 937 | if len(*k) < cap(*k) { 938 | s.fillLeft(n.right, k) 939 | } 940 | } 941 | 942 | func (s *TreeSet[T]) fillRight(n *node[T], k *[]T) { 943 | if n == nil { 944 | return 945 | } 946 | 947 | if len(*k) < cap(*k) { 948 | s.fillRight(n.right, k) 949 | } 950 | 951 | if len(*k) < cap(*k) { 952 | *k = append(*k, n.element) 953 | } 954 | 955 | if len(*k) < cap(*k) { 956 | s.fillRight(n.left, k) 957 | } 958 | } 959 | 960 | func (s *TreeSet[T]) prefix(visit func(*node[T]), n *node[T]) { 961 | if n == nil { 962 | return 963 | } 964 | visit(n) 965 | s.prefix(visit, n.left) 966 | s.prefix(visit, n.right) 967 | } 968 | 969 | func (s *TreeSet[T]) iterate() func() *node[T] { 970 | stck := makeStack[*node[T]]() 971 | 972 | for n := s.root; n != nil; n = n.left { 973 | stck.push(n) 974 | } 975 | 976 | return func() *node[T] { 977 | if stck.empty() { 978 | return nil 979 | } 980 | n := stck.pop() 981 | for r := n.right; r != nil; r = r.left { 982 | stck.push(r) 983 | } 984 | return n 985 | } 986 | } 987 | 988 | // MarshalJSON implements the json.Marshaler interface. 989 | func (s *TreeSet[T]) MarshalJSON() ([]byte, error) { 990 | return marshalJSON[T](s) 991 | } 992 | 993 | // UnmarshalJSON implements the json.Unmarshaler interface. 994 | func (s *TreeSet[T]) UnmarshalJSON(data []byte) error { 995 | return unmarshalJSON[T](s, data) 996 | } 997 | 998 | func (s *TreeSet[T]) filterLeft(n *node[T], accept func(element T) bool, result *TreeSet[T]) { 999 | if n == nil { 1000 | return 1001 | } 1002 | 1003 | s.filterLeft(n.left, accept, result) 1004 | 1005 | if accept(n.element) { 1006 | result.Insert(n.element) 1007 | s.filterLeft(n.right, accept, result) 1008 | } 1009 | } 1010 | 1011 | func (s *TreeSet[T]) filterRight(n *node[T], accept func(element T) bool, result *TreeSet[T]) { 1012 | if n == nil { 1013 | return 1014 | } 1015 | 1016 | s.filterRight(n.right, accept, result) 1017 | 1018 | if accept(n.element) { 1019 | result.Insert(n.element) 1020 | s.filterRight(n.left, accept, result) 1021 | } 1022 | } 1023 | -------------------------------------------------------------------------------- /treeset_test.go: -------------------------------------------------------------------------------- 1 | // Copyright (c) HashiCorp, Inc. 2 | // SPDX-License-Identifier: MPL-2.0 3 | 4 | package set 5 | 6 | import ( 7 | "cmp" 8 | "fmt" 9 | "math/rand" 10 | "strings" 11 | "testing" 12 | 13 | "github.com/shoenig/test/must" 14 | ) 15 | 16 | const ( 17 | size = 1000 18 | ) 19 | 20 | type token struct { 21 | id string 22 | } 23 | 24 | func (t *token) String() string { 25 | return t.id 26 | } 27 | 28 | func compareTokens(a, b *token) int { 29 | return cmp.Compare(a.id, b.id) 30 | } 31 | 32 | var ( 33 | tokenA = &token{id: "A"} 34 | tokenB = &token{id: "B"} 35 | tokenC = &token{id: "C"} 36 | tokenD = &token{id: "D"} 37 | tokenE = &token{id: "E"} 38 | tokenF = &token{id: "F"} 39 | tokenG = &token{id: "G"} 40 | tokenH = &token{id: "H"} 41 | ) 42 | 43 | func TestNewTreeSet(t *testing.T) { 44 | ts := NewTreeSet[*token](compareTokens) 45 | must.NotNil(t, ts) 46 | ts.dump() 47 | } 48 | 49 | func TestTreeSetFrom(t *testing.T) { 50 | s := shuffle(ints(10)) 51 | ts := TreeSetFrom[int](s, cmp.Compare[int]) 52 | must.NotEmpty(t, ts) 53 | } 54 | 55 | func TestTreeSet_Empty(t *testing.T) { 56 | t.Run("empty", func(t *testing.T) { 57 | ts := NewTreeSet[int](cmp.Compare[int]) 58 | must.Empty(t, ts) 59 | }) 60 | 61 | t.Run("not empty", func(t *testing.T) { 62 | ts := NewTreeSet[int](cmp.Compare[int]) 63 | ts.Insert(1) 64 | must.NotEmpty(t, ts) 65 | }) 66 | } 67 | 68 | func TestTreeSet_Size(t *testing.T) { 69 | t.Run("empty", func(t *testing.T) { 70 | ts := NewTreeSet[int](cmp.Compare[int]) 71 | must.Size(t, 0, ts) 72 | }) 73 | t.Run("one", func(t *testing.T) { 74 | ts := NewTreeSet[int](cmp.Compare[int]) 75 | ts.Insert(42) 76 | must.Size(t, 1, ts) 77 | }) 78 | t.Run("ten", func(t *testing.T) { 79 | ts := NewTreeSet[int](cmp.Compare[int]) 80 | s := shuffle(ints(10)) 81 | for i := 0; i < len(s); i++ { 82 | ts.Insert(s[i]) 83 | must.Size(t, i+1, ts) 84 | } 85 | // insert again (all duplicates) 86 | s = shuffle(s) 87 | for i := 0; i < len(s); i++ { 88 | ts.Insert(s[i]) 89 | must.Size(t, 10, ts) 90 | } 91 | }) 92 | } 93 | 94 | func TestTreeSet_Insert_token(t *testing.T) { 95 | ts := NewTreeSet[*token](compareTokens) 96 | 97 | ts.Insert(tokenA) 98 | invariants(t, ts, compareTokens) 99 | 100 | ts.Insert(tokenB) 101 | invariants(t, ts, compareTokens) 102 | 103 | ts.Insert(tokenC) 104 | invariants(t, ts, compareTokens) 105 | 106 | ts.Insert(tokenD) 107 | invariants(t, ts, compareTokens) 108 | 109 | ts.Insert(tokenE) 110 | invariants(t, ts, compareTokens) 111 | 112 | ts.Insert(tokenF) 113 | invariants(t, ts, compareTokens) 114 | 115 | ts.Insert(tokenG) 116 | invariants(t, ts, compareTokens) 117 | 118 | ts.Insert(tokenH) 119 | invariants(t, ts, compareTokens) 120 | 121 | t.Log("dump: insert token") 122 | t.Log(ts.dump()) 123 | } 124 | 125 | func TestTreeSet_Insert_int(t *testing.T) { 126 | cmp := cmp.Compare[int] 127 | ts := NewTreeSet[int](cmp) 128 | 129 | numbers := ints(size) 130 | random := shuffle(numbers) 131 | 132 | for _, i := range random { 133 | ts.Insert(i) 134 | invariants(t, ts, cmp) 135 | } 136 | 137 | t.Log("dump: insert int") 138 | t.Log(ts.dump()) 139 | } 140 | 141 | func TestTreeSet_InsertSlice(t *testing.T) { 142 | cmp := cmp.Compare[int] 143 | 144 | numbers := ints(size) 145 | random := shuffle(numbers) 146 | 147 | ts := NewTreeSet[int](cmp) 148 | must.True(t, ts.InsertSlice(random)) 149 | must.Eq(t, numbers, ts.Slice()) 150 | must.False(t, ts.InsertSlice(numbers)) 151 | } 152 | 153 | func TestTreeSet_InsertSet(t *testing.T) { 154 | cmp := cmp.Compare[int] 155 | 156 | ts1 := TreeSetFrom[int]([]int{1, 3, 5, 7, 9}, cmp) 157 | ts2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp) 158 | 159 | must.True(t, ts1.InsertSet(ts2)) 160 | must.Eq(t, []int{1, 2, 3, 5, 7, 9}, ts1.Slice()) 161 | must.Eq(t, []int{1, 2, 3}, ts2.Slice()) 162 | } 163 | 164 | func TestTreeSet_Remove_int(t *testing.T) { 165 | cmp := cmp.Compare[int] 166 | ts := NewTreeSet[int](cmp) 167 | 168 | numbers := ints(size) 169 | rnd := shuffle(numbers) 170 | 171 | // insert in random order 172 | for _, i := range rnd { 173 | ts.Insert(i) 174 | } 175 | 176 | invariants(t, ts, cmp) 177 | 178 | // reshuffle 179 | rnd = shuffle(rnd) 180 | 181 | // remove every element in random order 182 | for _, i := range rnd { 183 | removed := ts.Remove(i) 184 | t.Log("dump: remove", i) 185 | t.Log(ts.dump()) 186 | must.True(t, removed) 187 | invariants(t, ts, cmp) 188 | } 189 | 190 | // all gone 191 | must.Empty(t, ts) 192 | } 193 | 194 | func TestTreeSet_RemoveSlice(t *testing.T) { 195 | cmp := cmp.Compare[int] 196 | ts := NewTreeSet[int](cmp) 197 | 198 | numbers := ints(size) 199 | random := shuffle(numbers) 200 | ts.InsertSlice(random) 201 | 202 | must.True(t, ts.RemoveSlice(numbers)) 203 | must.Empty(t, ts) 204 | } 205 | 206 | func TestTreeSet_RemoveSet(t *testing.T) { 207 | cmp := cmp.Compare[int] 208 | 209 | ts1 := NewTreeSet[int](cmp) 210 | ts2 := NewTreeSet[int](cmp) 211 | 212 | numbers := ints(size) 213 | random := shuffle(numbers) 214 | ts1.InsertSlice(random) 215 | 216 | random2 := shuffle(numbers[5:]) 217 | ts2.InsertSlice(random2) 218 | 219 | ts1.RemoveSet(ts2) 220 | result := ts1.Slice() 221 | must.Eq(t, []int{1, 2, 3, 4, 5}, result) 222 | } 223 | 224 | func TestTreeSet_RemoveFunc(t *testing.T) { 225 | cmp := cmp.Compare[byte] 226 | 227 | ts := TreeSetFrom[byte]([]byte{ 228 | 'a', 'b', '1', 'c', '2', 'd', 229 | }, cmp) 230 | 231 | notAlpha := func(c byte) bool { 232 | return c < 'a' || c > 'z' 233 | } 234 | 235 | ts.RemoveFunc(notAlpha) 236 | 237 | must.Eq(t, []byte{'a', 'b', 'c', 'd'}, ts.Slice()) 238 | } 239 | 240 | func TestTreeSet_Contains(t *testing.T) { 241 | t.Run("empty", func(t *testing.T) { 242 | ts := NewTreeSet[int](cmp.Compare[int]) 243 | must.False(t, ts.Contains(42)) 244 | }) 245 | 246 | t.Run("exists", func(t *testing.T) { 247 | ts := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 248 | must.Contains[int](t, 1, ts) 249 | must.Contains[int](t, 2, ts) 250 | must.Contains[int](t, 3, ts) 251 | must.Contains[int](t, 4, ts) 252 | must.Contains[int](t, 5, ts) 253 | }) 254 | 255 | t.Run("absent", func(t *testing.T) { 256 | ts := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 257 | must.NotContains[int](t, 0, ts) 258 | must.NotContains[int](t, 6, ts) 259 | }) 260 | } 261 | 262 | func TestTreeSet_ContainsSlice(t *testing.T) { 263 | t.Run("empty", func(t *testing.T) { 264 | ts := NewTreeSet[int](cmp.Compare[int]) 265 | must.False(t, ts.ContainsSlice([]int{42, 43, 44})) 266 | }) 267 | 268 | t.Run("exists", func(t *testing.T) { 269 | ts := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 270 | must.True(t, ts.ContainsSlice([]int{2, 1, 3})) 271 | must.True(t, ts.ContainsSlice([]int{5, 4, 3, 2, 1})) 272 | }) 273 | 274 | t.Run("absent", func(t *testing.T) { 275 | ts := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 276 | must.False(t, ts.ContainsSlice([]int{6, 7, 8})) 277 | must.False(t, ts.ContainsSlice([]int{4, 5, 6})) 278 | }) 279 | } 280 | 281 | func TestTreeSet_Subset(t *testing.T) { 282 | t.Run("empty empty", func(t *testing.T) { 283 | t1 := NewTreeSet[int](cmp.Compare[int]) 284 | t2 := NewTreeSet[int](cmp.Compare[int]) 285 | must.True(t, t1.Subset(t2)) 286 | }) 287 | 288 | t.Run("empty full", func(t *testing.T) { 289 | t1 := NewTreeSet[int](cmp.Compare[int]) 290 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 291 | must.False(t, t1.Subset(t2)) 292 | }) 293 | 294 | t.Run("full empty", func(t *testing.T) { 295 | t1 := NewTreeSet[int](cmp.Compare[int]) 296 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 297 | must.True(t, t2.Subset(t1)) 298 | }) 299 | 300 | t.Run("same", func(t *testing.T) { 301 | t1 := TreeSetFrom[int]([]int{2, 1, 3}, cmp.Compare[int]) 302 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 303 | must.True(t, t1.Subset(t2)) 304 | must.True(t, t2.Subset(t1)) 305 | }) 306 | 307 | t.Run("subset", func(t *testing.T) { 308 | t1 := TreeSetFrom[int]([]int{2, 1, 3}, cmp.Compare[int]) 309 | t2 := TreeSetFrom[int]([]int{5, 4, 1, 2, 3}, cmp.Compare[int]) 310 | must.False(t, t1.Subset(t2)) 311 | }) 312 | 313 | t.Run("superset", func(t *testing.T) { 314 | t1 := TreeSetFrom[int]([]int{9, 7, 8, 5, 4, 2, 1, 3}, cmp.Compare[int]) 315 | t2 := TreeSetFrom[int]([]int{5, 1, 2, 8, 3}, cmp.Compare[int]) 316 | must.True(t, t1.Subset(t2)) 317 | }) 318 | 319 | t.Run("diff set", func(t *testing.T) { 320 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 321 | t2 := TreeSetFrom[int]([]int{6, 7, 8, 9, 10}, cmp.Compare[int]) 322 | must.False(t, t1.Subset(t2)) 323 | }) 324 | 325 | t.Run("exhaust", func(t *testing.T) { 326 | s1 := TreeSetFrom[string]([]string{"a", "b", "c", "d", "e"}, cmp.Compare[string]) 327 | s2 := TreeSetFrom[string]([]string{"a", "z"}, cmp.Compare[string]) 328 | must.False(t, s1.Subset(s2)) 329 | }) 330 | } 331 | 332 | func TestTreeSet_ProperSubset(t *testing.T) { 333 | t.Run("empty empty", func(t *testing.T) { 334 | t1 := NewTreeSet[int](cmp.Compare[int]) 335 | t2 := NewTreeSet[int](cmp.Compare[int]) 336 | must.False(t, t1.ProperSubset(t2)) 337 | }) 338 | 339 | t.Run("empty full", func(t *testing.T) { 340 | t1 := NewTreeSet[int](cmp.Compare[int]) 341 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 342 | must.False(t, t1.ProperSubset(t2)) 343 | }) 344 | 345 | t.Run("full empty", func(t *testing.T) { 346 | t1 := NewTreeSet[int](cmp.Compare[int]) 347 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 348 | must.True(t, t2.ProperSubset(t1)) 349 | }) 350 | 351 | t.Run("same", func(t *testing.T) { 352 | t1 := TreeSetFrom[int]([]int{2, 1, 3}, cmp.Compare[int]) 353 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 354 | must.False(t, t1.ProperSubset(t2)) 355 | must.False(t, t2.ProperSubset(t1)) 356 | }) 357 | 358 | t.Run("subset", func(t *testing.T) { 359 | t1 := TreeSetFrom[int]([]int{2, 1, 3}, cmp.Compare[int]) 360 | t2 := TreeSetFrom[int]([]int{5, 4, 1, 2, 3}, cmp.Compare[int]) 361 | must.False(t, t1.ProperSubset(t2)) 362 | }) 363 | 364 | t.Run("superset", func(t *testing.T) { 365 | t1 := TreeSetFrom[int]([]int{9, 7, 8, 5, 4, 2, 1, 3}, cmp.Compare[int]) 366 | t2 := TreeSetFrom[int]([]int{5, 1, 2, 8, 3}, cmp.Compare[int]) 367 | must.True(t, t1.ProperSubset(t2)) 368 | }) 369 | 370 | t.Run("diff set", func(t *testing.T) { 371 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 372 | t2 := TreeSetFrom[int]([]int{6, 7, 8, 9, 10}, cmp.Compare[int]) 373 | must.False(t, t1.ProperSubset(t2)) 374 | }) 375 | 376 | t.Run("exhaust", func(t *testing.T) { 377 | s1 := TreeSetFrom[string]([]string{"a", "b", "c", "d", "e"}, cmp.Compare[string]) 378 | s2 := TreeSetFrom[string]([]string{"a", "z"}, cmp.Compare[string]) 379 | must.False(t, s1.ProperSubset(s2)) 380 | }) 381 | } 382 | 383 | func TestTreeSet_Union(t *testing.T) { 384 | t.Run("empty empty", func(t *testing.T) { 385 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 386 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 387 | result := t1.Union(t2) 388 | must.Empty(t, result) 389 | }) 390 | 391 | t.Run("empty full", func(t *testing.T) { 392 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 393 | t2 := TreeSetFrom[int]([]int{3, 1, 2}, cmp.Compare[int]) 394 | result := t1.Union(t2) 395 | must.NotEmpty(t, result) 396 | must.Eq(t, []int{1, 2, 3}, result.Slice()) 397 | }) 398 | 399 | t.Run("full empty", func(t *testing.T) { 400 | t1 := TreeSetFrom[int]([]int{2, 3, 1}, cmp.Compare[int]) 401 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 402 | result := t1.Union(t2) 403 | must.NotEmpty(t, result) 404 | must.Eq(t, []int{1, 2, 3}, result.Slice()) 405 | }) 406 | 407 | t.Run("subset", func(t *testing.T) { 408 | t1 := TreeSetFrom[int]([]int{2, 3, 1}, cmp.Compare[int]) 409 | t2 := TreeSetFrom[int]([]int{2}, cmp.Compare[int]) 410 | result := t1.Union(t2) 411 | must.NotEmpty(t, result) 412 | must.Eq(t, []int{1, 2, 3}, result.Slice()) 413 | }) 414 | 415 | t.Run("superset", func(t *testing.T) { 416 | t1 := TreeSetFrom[int]([]int{2, 3, 1}, cmp.Compare[int]) 417 | t2 := TreeSetFrom[int]([]int{2, 5, 1, 2, 4}, cmp.Compare[int]) 418 | result := t1.Union(t2) 419 | must.NotEmpty(t, result) 420 | must.Eq(t, []int{1, 2, 3, 4, 5}, result.Slice()) 421 | }) 422 | } 423 | 424 | func TestTreeSet_Difference(t *testing.T) { 425 | t.Run("empty empty", func(t *testing.T) { 426 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 427 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 428 | result := t1.Difference(t2) 429 | must.Empty(t, result) 430 | }) 431 | 432 | t.Run("empty full", func(t *testing.T) { 433 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 434 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 435 | result := t1.Difference(t2) 436 | must.Empty(t, result) 437 | }) 438 | 439 | t.Run("full empty", func(t *testing.T) { 440 | t1 := TreeSetFrom[int]([]int{2, 1, 3}, cmp.Compare[int]) 441 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 442 | result := t1.Difference(t2) 443 | must.NotEmpty(t, result) 444 | must.Eq(t, []int{1, 2, 3}, result.Slice()) 445 | }) 446 | 447 | t.Run("subset", func(t *testing.T) { 448 | t1 := TreeSetFrom[int]([]int{3, 2, 4}, cmp.Compare[int]) 449 | t2 := TreeSetFrom[int]([]int{1, 2, 3, 4, 5}, cmp.Compare[int]) 450 | result := t1.Difference(t2) 451 | must.Empty(t, result) 452 | }) 453 | 454 | t.Run("superset", func(t *testing.T) { 455 | t1 := TreeSetFrom[int]([]int{2, 1, 3, 4, 5}, cmp.Compare[int]) 456 | t2 := TreeSetFrom[int]([]int{1, 2, 5}, cmp.Compare[int]) 457 | result := t1.Difference(t2) 458 | must.NotEmpty(t, result) 459 | must.Eq(t, []int{3, 4}, result.Slice()) 460 | }) 461 | } 462 | 463 | func TestTreeSet_Intersect(t *testing.T) { 464 | t.Run("empty empty", func(t *testing.T) { 465 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 466 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 467 | result := t1.Intersect(t2) 468 | must.Empty(t, result) 469 | }) 470 | 471 | t.Run("empty full", func(t *testing.T) { 472 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 473 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 474 | result := t1.Intersect(t2) 475 | must.Empty(t, result) 476 | }) 477 | 478 | t.Run("full empty", func(t *testing.T) { 479 | t1 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 480 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 481 | result := t1.Intersect(t2) 482 | must.Empty(t, result) 483 | }) 484 | 485 | t.Run("overlap", func(t *testing.T) { 486 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4, 5, 6}, cmp.Compare[int]) 487 | t2 := TreeSetFrom[int]([]int{0, 4, 5, 7}, cmp.Compare[int]) 488 | result := t1.Intersect(t2) 489 | must.NotEmpty(t, result) 490 | must.Eq(t, []int{4, 5}, result.Slice()) 491 | }) 492 | } 493 | 494 | func TestTreeSet_Copy(t *testing.T) { 495 | t.Run("empty", func(t *testing.T) { 496 | t1 := NewTreeSet[int](cmp.Compare[int]) 497 | c := t1.Copy() 498 | must.Empty(t, c) 499 | }) 500 | 501 | t.Run("full", func(t *testing.T) { 502 | t1 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 503 | c := t1.Copy() 504 | must.NotEmpty(t, c) 505 | must.Eq(t, []int{1, 2, 3}, c.Slice()) 506 | }) 507 | 508 | t.Run("modify", func(t *testing.T) { 509 | t1 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 510 | c := t1.Copy() 511 | c.Insert(4) 512 | t1.Remove(2) 513 | must.Eq(t, []int{1, 3}, t1.Slice()) 514 | must.Eq(t, []int{1, 2, 3, 4}, c.Slice()) 515 | }) 516 | } 517 | 518 | func TestTreeSet_EqualSlice(t *testing.T) { 519 | t.Run("empty empty", func(t *testing.T) { 520 | ts := TreeSetFrom[int](nil, cmp.Compare[int]) 521 | must.True(t, ts.EqualSlice(nil)) 522 | }) 523 | 524 | t.Run("empty full", func(t *testing.T) { 525 | ts := TreeSetFrom[int](nil, cmp.Compare[int]) 526 | must.False(t, ts.EqualSlice([]int{1, 2, 3})) 527 | }) 528 | 529 | t.Run("matching", func(t *testing.T) { 530 | ts := TreeSetFrom[int]([]int{1, 2, 3, 4, 5, 6}, cmp.Compare[int]) 531 | must.True(t, ts.EqualSlice([]int{3, 2, 1, 6, 5, 4})) 532 | }) 533 | 534 | t.Run("different middle", func(t *testing.T) { 535 | ts := TreeSetFrom[int]([]int{1, 2, 3, 5, 6}, cmp.Compare[int]) 536 | must.False(t, ts.EqualSlice([]int{3, 2, 9, 6, 5, 4})) 537 | }) 538 | } 539 | 540 | func TestTreeSet_Equal(t *testing.T) { 541 | t.Run("empty empty", func(t *testing.T) { 542 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 543 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 544 | must.Equal(t, t1, t2) 545 | }) 546 | 547 | t.Run("empty full", func(t *testing.T) { 548 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 549 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 550 | must.NotEqual(t, t1, t2) 551 | }) 552 | 553 | t.Run("matching", func(t *testing.T) { 554 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4, 5, 6}, cmp.Compare[int]) 555 | t2 := TreeSetFrom[int]([]int{6, 5, 4, 3, 2, 1}, cmp.Compare[int]) 556 | must.Equal(t, t1, t2) 557 | }) 558 | 559 | t.Run("different min", func(t *testing.T) { 560 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4}, cmp.Compare[int]) 561 | t2 := TreeSetFrom[int]([]int{0, 2, 3, 4}, cmp.Compare[int]) 562 | must.NotEqual(t, t1, t2) 563 | }) 564 | 565 | t.Run("different max", func(t *testing.T) { 566 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 4}, cmp.Compare[int]) 567 | t2 := TreeSetFrom[int]([]int{5, 3, 2, 1}, cmp.Compare[int]) 568 | must.NotEqual(t, t1, t2) 569 | }) 570 | 571 | t.Run("different middle", func(t *testing.T) { 572 | t1 := TreeSetFrom[int]([]int{1, 2, 3, 5, 6}, cmp.Compare[int]) 573 | t2 := TreeSetFrom[int]([]int{1, 2, 4, 5, 6}, cmp.Compare[int]) 574 | must.NotEqual(t, t1, t2) 575 | }) 576 | } 577 | 578 | func TestTreeSet_EqualSet(t *testing.T) { 579 | t.Run("empty empty", func(t *testing.T) { 580 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 581 | t2 := TreeSetFrom[int](nil, cmp.Compare[int]) 582 | must.True(t, t1.EqualSet(t2)) 583 | }) 584 | 585 | t.Run("empty full", func(t *testing.T) { 586 | t1 := TreeSetFrom[int](nil, cmp.Compare[int]) 587 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 588 | must.False(t, t1.EqualSet(t2)) 589 | must.False(t, t2.EqualSet(t1)) 590 | }) 591 | 592 | t.Run("different", func(t *testing.T) { 593 | t1 := TreeSetFrom[int]([]int{1, 2, 4}, cmp.Compare[int]) 594 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 595 | must.False(t, t1.EqualSet(t2)) 596 | must.False(t, t2.EqualSet(t1)) 597 | }) 598 | 599 | t.Run("same", func(t *testing.T) { 600 | t1 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 601 | t2 := TreeSetFrom[int]([]int{1, 2, 3}, cmp.Compare[int]) 602 | must.True(t, t1.EqualSet(t2)) 603 | must.True(t, t2.EqualSet(t1)) 604 | }) 605 | } 606 | 607 | func TestTreeSet_TopK(t *testing.T) { 608 | t.Run("empty", func(t *testing.T) { 609 | ts := NewTreeSet[int](cmp.Compare[int]) 610 | result := ts.TopK(5) 611 | must.Eq(t, []int{}, result) 612 | }) 613 | 614 | t.Run("same size", func(t *testing.T) { 615 | ts := TreeSetFrom[int]([]int{3, 9, 1, 7, 5}, cmp.Compare[int]) 616 | result := ts.TopK(5) 617 | must.Eq(t, []int{1, 3, 5, 7, 9}, result) 618 | }) 619 | 620 | t.Run("smaller k", func(t *testing.T) { 621 | ts := TreeSetFrom[int]([]int{3, 9, 1, 7, 5}, cmp.Compare[int]) 622 | result := ts.TopK(3) 623 | must.Eq(t, []int{1, 3, 5}, result) 624 | }) 625 | } 626 | 627 | func TestTreeSet_BottomK(t *testing.T) { 628 | t.Run("empty", func(t *testing.T) { 629 | ts := NewTreeSet[int](cmp.Compare[int]) 630 | result := ts.BottomK(5) 631 | must.Eq(t, []int{}, result) 632 | }) 633 | 634 | t.Run("same size", func(t *testing.T) { 635 | ts := TreeSetFrom[int]([]int{3, 9, 1, 7, 5}, cmp.Compare[int]) 636 | result := ts.BottomK(5) 637 | must.Eq(t, []int{9, 7, 5, 3, 1}, result) 638 | }) 639 | 640 | t.Run("smaller k", func(t *testing.T) { 641 | ts := TreeSetFrom[int]([]int{3, 9, 1, 7, 5}, cmp.Compare[int]) 642 | result := ts.BottomK(3) 643 | must.Eq(t, []int{9, 7, 5}, result) 644 | }) 645 | } 646 | 647 | func TestTreeSet_FirstBelow(t *testing.T) { 648 | t.Run("empty", func(t *testing.T) { 649 | ts := NewTreeSet[int](cmp.Compare[int]) 650 | _, exists := ts.FirstBelow(5) 651 | must.False(t, exists) 652 | }) 653 | 654 | t.Run("basic", func(t *testing.T) { 655 | ts := TreeSetFrom[int]([]int{1, 3, 4, 5, 7, 8}, cmp.Compare[int]) 656 | v, exists := ts.FirstBelow(5) 657 | must.True(t, exists) 658 | must.Eq(t, 4, v) 659 | }) 660 | 661 | t.Run("many", func(t *testing.T) { 662 | ts := NewTreeSet[int](cmp.Compare[int]) 663 | nums := shuffle(ints(100)) 664 | ts.InsertSlice(nums) 665 | for i := 2; i < 100; i++ { 666 | v, exists := ts.FirstBelow(i) 667 | must.True(t, exists) 668 | must.Eq(t, i-1, v) 669 | } 670 | }) 671 | } 672 | 673 | func TestTreeSet_FirstBelowEqual(t *testing.T) { 674 | t.Run("empty", func(t *testing.T) { 675 | ts := NewTreeSet[int](cmp.Compare[int]) 676 | _, exists := ts.FirstBelowEqual(5) 677 | must.False(t, exists) 678 | }) 679 | 680 | t.Run("basic", func(t *testing.T) { 681 | ts := TreeSetFrom[int]([]int{1, 3, 4, 5, 7, 8}, cmp.Compare[int]) 682 | v, exists := ts.FirstBelowEqual(5) 683 | must.True(t, exists) 684 | must.Eq(t, 5, v) 685 | }) 686 | 687 | t.Run("many", func(t *testing.T) { 688 | ts := NewTreeSet[int](cmp.Compare[int]) 689 | nums := shuffle(ints(100)) 690 | ts.InsertSlice(nums) 691 | for i := 1; i < 100; i++ { 692 | v, exists := ts.FirstBelowEqual(i) 693 | must.True(t, exists) 694 | must.Eq(t, i, v) 695 | } 696 | }) 697 | } 698 | 699 | func TestTreeSet_Below(t *testing.T) { 700 | t.Run("empty", func(t *testing.T) { 701 | ts := TreeSetFrom[int]([]int{5, 6, 7, 8, 9}, cmp.Compare[int]) 702 | b := ts.Below(5) 703 | must.Empty(t, b) 704 | }) 705 | 706 | t.Run("basic", func(t *testing.T) { 707 | ts := TreeSetFrom[int]([]int{4, 7, 1, 5, 2, 8, 9, 3}, cmp.Compare[int]) 708 | b := ts.Below(5) 709 | result := b.Slice() 710 | must.Eq(t, []int{1, 2, 3, 4}, result) 711 | }) 712 | 713 | t.Run("many", func(t *testing.T) { 714 | ts := NewTreeSet[int](cmp.Compare[int]) 715 | nums := shuffle(ints(100)) 716 | ts.InsertSlice(nums) 717 | for i := 2; i < 100; i++ { 718 | below := ts.Below(i) 719 | must.Size(t, i-1, below) 720 | must.Min(t, 1, below) 721 | must.Max(t, i-1, below) 722 | } 723 | }) 724 | } 725 | 726 | func TestTreeSet_BelowEqual(t *testing.T) { 727 | t.Run("empty", func(t *testing.T) { 728 | ts := TreeSetFrom[int]([]int{5, 6, 7, 8, 9}, cmp.Compare[int]) 729 | b := ts.BelowEqual(4) 730 | must.Empty(t, b) 731 | }) 732 | 733 | t.Run("basic", func(t *testing.T) { 734 | ts := TreeSetFrom[int]([]int{4, 7, 1, 5, 2, 8, 9, 3}, cmp.Compare[int]) 735 | b := ts.BelowEqual(5) 736 | result := b.Slice() 737 | must.Eq(t, []int{1, 2, 3, 4, 5}, result) 738 | }) 739 | 740 | t.Run("many", func(t *testing.T) { 741 | ts := NewTreeSet[int](cmp.Compare[int]) 742 | nums := shuffle(ints(100)) 743 | ts.InsertSlice(nums) 744 | for i := 1; i < 100; i++ { 745 | below := ts.BelowEqual(i) 746 | must.Size(t, i, below) 747 | must.Min(t, 1, below) 748 | must.Max(t, i, below) 749 | } 750 | }) 751 | } 752 | 753 | func TestTreeSet_FirstAbove(t *testing.T) { 754 | t.Run("empty", func(t *testing.T) { 755 | ts := TreeSetFrom[int]([]int{2, 1, 3, 5, 4}, cmp.Compare[int]) 756 | _, exists := ts.FirstAbove(5) 757 | must.False(t, exists) 758 | }) 759 | 760 | t.Run("basic", func(t *testing.T) { 761 | ts := TreeSetFrom[int]([]int{2, 1, 4, 6, 5, 7, 8}, cmp.Compare[int]) 762 | v, exists := ts.FirstAbove(5) 763 | must.True(t, exists) 764 | must.Eq(t, 6, v) 765 | }) 766 | 767 | t.Run("many", func(t *testing.T) { 768 | ts := NewTreeSet[int](cmp.Compare[int]) 769 | nums := shuffle(ints(100)) 770 | ts.InsertSlice(nums) 771 | for i := 1; i < 100; i++ { 772 | v, exists := ts.FirstAbove(i) 773 | must.True(t, exists) 774 | must.Eq(t, i+1, v) 775 | } 776 | }) 777 | } 778 | 779 | func TestTreeSet_FirstAboveEqual(t *testing.T) { 780 | t.Run("empty", func(t *testing.T) { 781 | ts := TreeSetFrom[int]([]int{2, 1, 3, 4}, cmp.Compare[int]) 782 | _, exists := ts.FirstAboveEqual(5) 783 | must.False(t, exists) 784 | }) 785 | 786 | t.Run("basic", func(t *testing.T) { 787 | ts := TreeSetFrom[int]([]int{2, 1, 4, 6, 5, 7, 8}, cmp.Compare[int]) 788 | v, exists := ts.FirstAboveEqual(5) 789 | must.True(t, exists) 790 | must.Eq(t, 5, v) 791 | }) 792 | 793 | t.Run("many", func(t *testing.T) { 794 | ts := NewTreeSet[int](cmp.Compare[int]) 795 | nums := shuffle(ints(100)) 796 | ts.InsertSlice(nums) 797 | for i := 1; i < 100; i++ { 798 | v, exists := ts.FirstAboveEqual(i) 799 | must.True(t, exists) 800 | must.Eq(t, i, v) 801 | } 802 | }) 803 | } 804 | 805 | func TestTreeSet_Above(t *testing.T) { 806 | t.Run("empty", func(t *testing.T) { 807 | ts := TreeSetFrom[int]([]int{5, 6, 7, 8, 9}, cmp.Compare[int]) 808 | b := ts.Above(9) 809 | must.Empty(t, b) 810 | }) 811 | 812 | t.Run("basic", func(t *testing.T) { 813 | ts := TreeSetFrom[int]([]int{4, 7, 1, 5, 2, 8, 9, 3}, cmp.Compare[int]) 814 | b := ts.Above(5) 815 | result := b.Slice() 816 | must.Eq(t, []int{7, 8, 9}, result) 817 | }) 818 | 819 | t.Run("many", func(t *testing.T) { 820 | ts := NewTreeSet[int](cmp.Compare[int]) 821 | nums := shuffle(ints(100)) 822 | ts.InsertSlice(nums) 823 | for i := 1; i < 100; i++ { 824 | above := ts.Above(i) 825 | must.Size(t, 100-i, above) 826 | must.Min(t, i+1, above) 827 | must.Max(t, 100, above) 828 | } 829 | }) 830 | } 831 | 832 | func TestTreeSet_AboveEqual(t *testing.T) { 833 | t.Run("empty", func(t *testing.T) { 834 | ts := TreeSetFrom[int]([]int{5, 6, 7, 8, 9}, cmp.Compare[int]) 835 | b := ts.AboveEqual(10) 836 | must.Empty(t, b) 837 | }) 838 | 839 | t.Run("basic", func(t *testing.T) { 840 | ts := TreeSetFrom[int]([]int{4, 7, 1, 5, 2, 8, 9, 3}, cmp.Compare[int]) 841 | b := ts.AboveEqual(5) 842 | result := b.Slice() 843 | must.Eq(t, []int{5, 7, 8, 9}, result) 844 | }) 845 | 846 | t.Run("many", func(t *testing.T) { 847 | ts := NewTreeSet[int](cmp.Compare[int]) 848 | nums := shuffle(ints(100)) 849 | ts.InsertSlice(nums) 850 | for i := 1; i < 100; i++ { 851 | above := ts.AboveEqual(i) 852 | must.Size(t, 100-i+1, above) 853 | must.Min(t, i, above) 854 | must.Max(t, 100, above) 855 | } 856 | }) 857 | } 858 | 859 | func TestTreeSet_Slice(t *testing.T) { 860 | t.Run("empty", func(t *testing.T) { 861 | ts := NewTreeSet[int](cmp.Compare[int]) 862 | result := ts.Slice() 863 | must.Eq(t, []int{}, result) 864 | }) 865 | 866 | t.Run("full", func(t *testing.T) { 867 | ts := TreeSetFrom[int]([]int{4, 2, 6, 1}, cmp.Compare[int]) 868 | result := ts.Slice() 869 | must.Eq(t, []int{1, 2, 4, 6}, result) 870 | }) 871 | } 872 | 873 | func TestTreeSet_String(t *testing.T) { 874 | t.Run("empty", func(t *testing.T) { 875 | ts := NewTreeSet[int](cmp.Compare[int]) 876 | result := ts.String() 877 | must.Eq(t, "[]", result) 878 | }) 879 | 880 | t.Run("full", func(t *testing.T) { 881 | ts := TreeSetFrom[int]([]int{4, 2, 6, 1}, cmp.Compare[int]) 882 | result := ts.String() 883 | must.Eq(t, "[1 2 4 6]", result) 884 | }) 885 | } 886 | 887 | func TestTreeSet_StringFunc(t *testing.T) { 888 | f := func(i int) string { return fmt.Sprintf("%02d", i) } 889 | t.Run("empty", func(t *testing.T) { 890 | ts := NewTreeSet[int](cmp.Compare[int]) 891 | result := ts.StringFunc(f) 892 | must.Eq(t, "[]", result) 893 | }) 894 | 895 | t.Run("full", func(t *testing.T) { 896 | ts := TreeSetFrom[int]([]int{4, 2, 6, 1}, cmp.Compare[int]) 897 | result := ts.StringFunc(f) 898 | must.Eq(t, "[01 02 04 06]", result) 899 | }) 900 | } 901 | 902 | // create a colorful representation of the element in node 903 | func (n *node[T]) String() string { 904 | if n.red() { 905 | return fmt.Sprintf("\033[1;31m%v\033[0m", n.element) 906 | } 907 | return fmt.Sprintf("%v", n.element) 908 | } 909 | 910 | // output creates a colorful string representation of s 911 | func (s *TreeSet[T]) output(prefix, cprefix string, n *node[T], sb *strings.Builder) { 912 | if n == nil { 913 | return 914 | } 915 | 916 | sb.WriteString(prefix) 917 | sb.WriteString(n.String()) 918 | sb.WriteString("\n") 919 | 920 | if n.right != nil && n.left != nil { 921 | s.output(cprefix+"├── ", cprefix+"│ ", n.right, sb) 922 | } else if n.right != nil { 923 | s.output(cprefix+"└── ", cprefix+" ", n.right, sb) 924 | } 925 | if n.left != nil { 926 | s.output(cprefix+"└── ", cprefix+" ", n.left, sb) 927 | } 928 | if n.left == nil && n.right == nil { 929 | return 930 | } 931 | } 932 | 933 | // dump the output of s along with the slice string 934 | func (s *TreeSet[T]) dump() string { 935 | var sb strings.Builder 936 | sb.WriteString("\ntree:\n") 937 | s.output("", "", s.root, &sb) 938 | sb.WriteString("string:") 939 | sb.WriteString(s.String()) 940 | return sb.String() 941 | } 942 | 943 | // invariants makes basic assertions about tree 944 | func invariants[T any](t *testing.T, tree *TreeSet[T], cmp CompareFunc[T]) { 945 | // assert Slice elements are ascending 946 | slice := tree.Slice() 947 | must.AscendingCmp(t, slice, cmp) 948 | 949 | // assert size of tree 950 | size := tree.Size() 951 | must.Eq(t, size, len(slice), must.Sprint("tree is wrong size")) 952 | 953 | if size == 0 { 954 | return 955 | } 956 | 957 | // assert slice[0] is the minimum 958 | must.Min(t, slice[0], tree) 959 | 960 | // assert slice[len(slice)-1] is the maximum 961 | must.Max(t, slice[len(slice)-1], tree) 962 | } 963 | 964 | // ints will create a []int from 1 to n 965 | func ints(n int) []int { 966 | s := make([]int, n) 967 | for i := 0; i < n; i++ { 968 | s[i] = i + 1 969 | } 970 | return s 971 | } 972 | 973 | // create a copy of s and shuffle 974 | func shuffle(s []int) []int { 975 | c := make([]int, len(s)) 976 | copy(c, s) 977 | 978 | n := len(c) 979 | for i := 0; i < n; i++ { 980 | swp := rand.Int31n(int32(n)) 981 | c[i], c[swp] = c[swp], c[i] 982 | } 983 | return c 984 | } 985 | 986 | func TestTreeSet_infix(t *testing.T) { 987 | ts := TreeSetFrom[int]([]int{4, 7, 1, 5, 2, 8, 9, 3, 11, 13}, cmp.Compare[int]) 988 | isOdd := func(n *node[int]) bool { 989 | return n.element%2 == 1 990 | } 991 | odds := make([]int, 0, 5) 992 | ts.infix(func(n *node[int]) bool { 993 | if n.element == 8 { 994 | return false 995 | } 996 | if isOdd(n) { 997 | odds = append(odds, n.element) 998 | } 999 | 1000 | return true 1001 | }, ts.root) 1002 | must.Eq(t, []int{1, 3, 5, 7}, odds) 1003 | } 1004 | 1005 | func TestTreeSet_iterate2(t *testing.T) { 1006 | nums := shuffle(ints(11)) 1007 | s := TreeSetFrom[int](nums, cmp.Compare[int]) 1008 | 1009 | iter := s.iterate() 1010 | for i := 1; i <= 11; i++ { 1011 | must.Eq(t, i, iter().element) 1012 | } 1013 | must.Nil(t, iter()) 1014 | } 1015 | 1016 | func TestTreeSet_Items(t *testing.T) { 1017 | ts := TreeSetFrom[int]([]int{2, 1, 4, 3, 5}, cmp.Compare[int]) 1018 | 1019 | exp := []int{1, 2, 3, 4, 5} 1020 | result := []int{} 1021 | for element := range ts.Items() { 1022 | result = append(result, element) 1023 | } 1024 | 1025 | must.Eq(t, exp, result) 1026 | } 1027 | --------------------------------------------------------------------------------