diff --git a/collect/src/main/kotlin/com/certora/collect/EmptyTreapSet.kt b/collect/src/main/kotlin/com/certora/collect/EmptyTreapSet.kt index 609fa57..73f2d26 100644 --- a/collect/src/main/kotlin/com/certora/collect/EmptyTreapSet.kt +++ b/collect/src/main/kotlin/com/certora/collect/EmptyTreapSet.kt @@ -16,6 +16,7 @@ internal class EmptyTreapSet<@Treapable E> private constructor() : TreapSet, override fun contains(element: E): Boolean = false override fun containsAll(elements: Collection): Boolean = elements.isEmpty() override fun containsAny(elements: Iterable): Boolean = false + override fun containsAny(predicate: (E) -> Boolean): Boolean = false override fun findEqual(element: E): E? = null override fun forEachElement(action: (element: E) -> Unit): Unit {} override fun remove(element: E): TreapSet = this diff --git a/collect/src/main/kotlin/com/certora/collect/HashTreapSet.kt b/collect/src/main/kotlin/com/certora/collect/HashTreapSet.kt index 0f6c8d5..987e503 100644 --- a/collect/src/main/kotlin/com/certora/collect/HashTreapSet.kt +++ b/collect/src/main/kotlin/com/certora/collect/HashTreapSet.kt @@ -240,6 +240,15 @@ internal class HashTreapSet<@Treapable E>( } return result!! } + + override fun containsAny(predicate: (E) -> Boolean): Boolean { + forEachNodeElement { + if (predicate(it)) { + return true + } + } + return left?.containsAny(predicate) == true || right?.containsAny(predicate) == true + } } internal interface ElementList { diff --git a/collect/src/main/kotlin/com/certora/collect/SortedTreapSet.kt b/collect/src/main/kotlin/com/certora/collect/SortedTreapSet.kt index 87e9f4b..5704f0a 100644 --- a/collect/src/main/kotlin/com/certora/collect/SortedTreapSet.kt +++ b/collect/src/main/kotlin/com/certora/collect/SortedTreapSet.kt @@ -53,4 +53,7 @@ internal class SortedTreapSet<@Treapable E>( override fun arbitraryOrNull(): E? = treapKey override fun shallowForEach(action: (element: E) -> Unit): Unit { action(treapKey) } override fun shallowMapReduce(map: (E) -> R, reduce: (R, R) -> R): R = map(treapKey) + + override fun containsAny(predicate: (E) -> Boolean): Boolean = + predicate(treapKey) || left?.containsAny(predicate) == true || right?.containsAny(predicate) == true } diff --git a/collect/src/main/kotlin/com/certora/collect/TreapSet.kt b/collect/src/main/kotlin/com/certora/collect/TreapSet.kt index 66de30b..7d8d643 100644 --- a/collect/src/main/kotlin/com/certora/collect/TreapSet.kt +++ b/collect/src/main/kotlin/com/certora/collect/TreapSet.kt @@ -32,6 +32,13 @@ public sealed interface TreapSet : PersistentSet { */ public fun containsAny(elements: Iterable<@UnsafeVariance T>): Boolean + /** + Checks if this set contains any element that satisfies the given [predicate]. + + This traverses the treap without allocating temporary storage, which may be more efficient than [any]. + */ + public fun containsAny(predicate: (T) -> Boolean): Boolean + /** If this set contains exactly one element, returns that element. Otherwise, throws [NoSuchElementException]. */