Skip to content

Add variance annotations #28

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ericeil merged 2 commits into from
Feb 27, 2026
Merged

Add variance annotations #28

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d7cc2c0
Add variance annotations
ericeil Feb 26, 2026
44a261e
Add comment about @UnsafeVariance annotations
ericeil Feb 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions collect/src/main/kotlin/com/certora/collect/AbstractTreapMap.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,19 +151,19 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
override operator fun iterator() = entrySequence().map { it.value }.iterator()
}

override fun union(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> =
override fun union(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> =
m.useAsTreap(
{ otherTreap -> self.unionWith(otherTreap, getShallowUnionMerger(merger)) ?: clear() },
{ fallbackUnion(m, merger) }
)

override fun parallelUnion(m: Map<K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> =
override fun parallelUnion(m: Map<out K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> =
m.useAsTreap(
{ otherTreap -> self.parallelUnionWith(otherTreap, parallelThresholdLog2, getShallowUnionMerger(merger)) ?: clear() },
{ fallbackUnion(m, merger) }
)

private fun fallbackUnion(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> {
private fun fallbackUnion(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> {
var newThis = this as TreapMap<K, V>
for ((k, v) in m.entries) {
if (k in this) {
Expand All @@ -175,19 +175,19 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
return newThis
}

override fun intersect(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> =
override fun intersect(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> =
m.useAsTreap(
{ otherTreap -> self.intersectWith(otherTreap, getShallowIntersectMerger(merger)) ?: clear() },
{ fallbackIntersect(m, merger) }
)

override fun parallelIntersect(m: Map<K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> =
override fun parallelIntersect(m: Map<out K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> =
m.useAsTreap(
{ otherTreap -> self.parallelIntersectWith(otherTreap, parallelThresholdLog2, getShallowIntersectMerger(merger)) ?: clear() },
{ fallbackIntersect(m, merger) }
)

private fun fallbackIntersect(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> {
private fun fallbackIntersect(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> {
var newThis = clear()
for ((k, v) in m.entries) {
if (k in this) {
Expand All @@ -201,7 +201,7 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
Merges the entries in `m` with the entries in this AbstractTreapMap, applying the "merger" function to get the
new values for each key.
*/
override fun merge(m: Map<K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> =
override fun merge(m: Map<out K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> =
m.useAsTreap(
{ otherTreap ->
self.mergeWith(otherTreap, mode, getShallowMerger(mode, merger))
Expand All @@ -220,7 +220,7 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
@param[merger] The merge function to apply to each pair of entries. Must be pure and thread-safe.
*/
override fun parallelMerge(
m: Map<K, V>,
m: Map<out K, V>,
mode: MergeMode,
parallelThresholdLog2: Int,
merger: (K, V?, V?) -> V?
Expand All @@ -233,7 +233,7 @@ internal sealed class AbstractTreapMap<@Treapable K, V, @Treapable S : AbstractT
{ fallbackMerge(m, mode, merger) }
)

private fun fallbackMerge(m: Map<K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> {
private fun fallbackMerge(m: Map<out K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> {
var newThis = clear()
val keys = when(mode) {
MergeMode.UNION -> this.keys union m.keys
Expand Down
12 changes: 6 additions & 6 deletions collect/src/main/kotlin/com/certora/collect/EmptyTreapMap.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,13 +44,13 @@ internal class EmptyTreapMap<@Treapable K, V> private constructor() : TreapMap<K
else -> put(key, v)
}

override fun union(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> = putAll(m)
override fun parallelUnion(m: Map<K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> = putAll(m)
override fun union(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> = putAll(m)
override fun parallelUnion(m: Map<out K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> = putAll(m)

override fun intersect(m: Map<K, V>, merger: (K, V, V) -> V): TreapMap<K, V> = this
override fun parallelIntersect(m: Map<K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> = this
override fun intersect(m: Map<out K, V>, merger: (K, V, V) -> V): TreapMap<K, V> = this
override fun parallelIntersect(m: Map<out K, V>, parallelThresholdLog2: Int, merger: (K, V, V) -> V): TreapMap<K, V> = this

override fun merge(m: Map<K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> {
override fun merge(m: Map<out K, V>, mode: MergeMode, merger: (K, V?, V?) -> V?): TreapMap<K, V> {
var map: TreapMap<K, V> = this
if (mode == MergeMode.UNION) {
for ((key, value) in m) {
Expand All @@ -65,7 +65,7 @@ internal class EmptyTreapMap<@Treapable K, V> private constructor() : TreapMap<K


override fun parallelMerge(
m: Map<K, V>,
m: Map<out K, V>,
mode: MergeMode,
parallelThresholdLog2: Int,
merger: (K, V?, V?) -> V?
Expand Down
28 changes: 14 additions & 14 deletions collect/src/main/kotlin/com/certora/collect/TreapList.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,17 +6,17 @@ import kotlinx.collections.immutable.PersistentList
A PersistentList implemented as a [Treap](https://en.wikipedia.org/wiki/Treap).
*/
@Treapable
public sealed interface TreapList<E> : PersistentList<E> {
override fun add(element: E): TreapList<E> = addLast(element)
override fun addAll(elements: Collection<E>): TreapList<E>
override fun remove(element: E): TreapList<E>
override fun removeAll(elements: Collection<E>): TreapList<E>
public sealed interface TreapList<out E> : PersistentList<E> {
override fun add(element: @UnsafeVariance E): TreapList<E> = addLast(element)
override fun addAll(elements: Collection<@UnsafeVariance E>): TreapList<E>
override fun remove(element: @UnsafeVariance E): TreapList<E>
override fun removeAll(elements: Collection<@UnsafeVariance E>): TreapList<E>
override fun removeAll(predicate: (E) -> Boolean): TreapList<E>
override fun retainAll(elements: Collection<E>): TreapList<E>
override fun retainAll(elements: Collection<@UnsafeVariance E>): TreapList<E>
override fun clear(): TreapList<E>
override fun addAll(index: Int, c: Collection<E>): TreapList<E>
override fun set(index: Int, element: E): TreapList<E>
override fun add(index: Int, element: E): TreapList<E>
override fun addAll(index: Int, c: Collection<@UnsafeVariance E>): TreapList<E>
override fun set(index: Int, element: @UnsafeVariance E): TreapList<E>
override fun add(index: Int, element: @UnsafeVariance E): TreapList<E>
override fun removeAt(index: Int): TreapList<E>
override fun subList(fromIndex: Int, toIndex: Int): TreapList<E>

Expand All @@ -28,14 +28,14 @@ public sealed interface TreapList<E> : PersistentList<E> {
public fun last(): E
public fun lastOrNull(): E?

public fun addFirst(element: E): TreapList<E>
public fun addLast(element: E): TreapList<E>
public fun addFirst(element: @UnsafeVariance E): TreapList<E>
public fun addLast(element: @UnsafeVariance E): TreapList<E>

public fun removeFirst(): TreapList<E>
public fun removeLast(): TreapList<E>

public fun updateElements(transform: (E) -> E?): TreapList<E>
public fun updateElementsIndexed(transform: (Int, E) -> E?): TreapList<E>
public fun updateElements(transform: (E) -> @UnsafeVariance E?): TreapList<E>
public fun updateElementsIndexed(transform: (Int, E) -> @UnsafeVariance E?): TreapList<E>

public fun <R : Any> mapReduce(map: (E) -> R, reduce: (R, R) -> R): R?
public fun <R : Any> parallelMapReduce(map: (E) -> R, reduce: (R, R) -> R, parallelThresholdLog2: Int = 5): R?
Expand All @@ -57,7 +57,7 @@ public sealed interface TreapList<E> : PersistentList<E> {
public fun removeLast(): E
}

override fun builder(): Builder<E>
override fun builder(): Builder<@UnsafeVariance E>
}

public fun <T> treapListOf(): TreapList<T> = EmptyTreapList<T>()
Expand Down
42 changes: 23 additions & 19 deletions collect/src/main/kotlin/com/certora/collect/TreapMap.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@ import kotlinx.collections.immutable.PersistentMap
A [PersistentMap] implemented as a [Treap](https://en.wikipedia.org/wiki/Treap) - a kind of balanced binary tree.
*/
@Treapable
public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
public sealed interface TreapMap<K, out V> : PersistentMap<K, V> {

// Note: all the @UnsafeVariance annotations are safe because we never mutate anything. Any value appearing in an
// "in" position can only be a) compared to other values, or b) returned in an "out" position.

override fun put(key: K, value: @UnsafeVariance V): TreapMap<K, V>
override fun remove(key: K): TreapMap<K, V>
override fun remove(key: K, value: @UnsafeVariance V): TreapMap<K, V>
Expand Down Expand Up @@ -54,8 +58,8 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
the value from [m], in that order, and the returned value will appear in the resulting map.
*/
public fun union(
m: Map<K, V>,
merger: (K, V, V) -> V
m: Map<out K, @UnsafeVariance V>,
merger: (K, V, V) -> @UnsafeVariance V
): TreapMap<K, V>

/**
Expand All @@ -68,9 +72,9 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
Merge operations are performed in parallel for maps larger than (approximately) 2^parallelThresholdLog2.
*/
public fun parallelUnion(
m: Map<K, V>,
m: Map<out K, @UnsafeVariance V>,
parallelThresholdLog2: Int = 4,
merger: (K, V, V) -> V
merger: (K, V, V) -> @UnsafeVariance V
): TreapMap<K, V>

/**
Expand All @@ -80,8 +84,8 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
the key, the value from this map, and the value from [m], in that order.
*/
public fun intersect(
m: Map<K, V>,
merger: (K, V, V) -> V
m: Map<out K, @UnsafeVariance V>,
merger: (K, V, V) -> @UnsafeVariance V
): TreapMap<K, V>

/**
Expand All @@ -93,9 +97,9 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
Merge operations are performed in parallel for maps larger than (approximately) 2^parallelThresholdLog2.
*/
public fun parallelIntersect(
m: Map<K, V>,
m: Map<out K, @UnsafeVariance V>,
parallelThresholdLog2: Int = 4,
merger: (K, V, V) -> V
merger: (K, V, V) -> @UnsafeVariance V
): TreapMap<K, V>

/**
Expand All @@ -115,9 +119,9 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
[union] or [intersect] instead.
*/
public fun merge(
m: Map<K, V>,
m: Map<out K, @UnsafeVariance V>,
mode: MergeMode,
merger: (K, V?, V?) -> V?
merger: (K, V?, V?) -> @UnsafeVariance V?
): TreapMap<K, V>

/**
Expand All @@ -134,8 +138,8 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
[union] or [intersect] instead.
*/
public fun merge(
m: Map<K, V>,
merger: (K, V?, V?) -> V?
m: Map<out K, @UnsafeVariance V>,
merger: (K, V?, V?) -> @UnsafeVariance V?
): TreapMap<K, V> = merge(m, MergeMode.UNION, merger)

/**
Expand All @@ -157,10 +161,10 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
[parallelUnion] or [parallelIntersect] instead.
*/
public fun parallelMerge(
m: Map<K, V>,
m: Map<out K, @UnsafeVariance V>,
mode: MergeMode,
parallelThresholdLog2: Int = 4,
merger: (K, V?, V?) -> V?
merger: (K, V?, V?) -> @UnsafeVariance V?
): TreapMap<K, V>

/**
Expand All @@ -179,9 +183,9 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
[parallelUnion] or [parallelIntersect] instead.
*/
public fun parallelMerge(
m: Map<K, V>,
m: Map<out K, @UnsafeVariance V>,
parallelThresholdLog2: Int = 4,
merger: (K, V?, V?) -> V?
merger: (K, V?, V?) -> @UnsafeVariance V?
): TreapMap<K, V> = parallelMerge(m, MergeMode.UNION, parallelThresholdLog2, merger)

/**
Expand Down Expand Up @@ -231,15 +235,15 @@ public sealed interface TreapMap<K, V> : PersistentMap<K, V> {
public fun <U> updateEntry(
key: K,
value: U,
merger: (V?, U) -> V?
merger: (V?, U) -> @UnsafeVariance V?
): TreapMap<K, V>

/**
Produces a sequence from the entries of this map and another map. For each key, the result is an entry mapping
the key to a pair of values. Each value may be null, if the key is not present in the corresponding map.
*/
public fun zip(
m: Map<out K, V>
m: Map<out K, @UnsafeVariance V>
): Sequence<Map.Entry<K, Pair<V?, V?>>>

/**
Expand Down
Loading