Set with choice implementation in ponylang

main.pony

"""
A small test of our setch.
"""
use "collections"

actor Main
    new create(env: Env) =>
        let s: Setch[USize] = Setch[USize]

        for j in Range(0, 10) do
            s.set(j)
        end

        for _ in Range(0, 10) do
            env.out.print(
                try
                    s.choose()?.string()
                else
                    "error"
                end
            )
        end

        let r: Setch[USize] = Setch[USize]
        env.out.print(
            try
                r.choose()?.string()
            else
                "error"
            end
        )

setch.pony

use "collections"
use "random"
use "time"

type Setch[A: (Hashable #read & Equatable[A] #read)] is HashSetch[A, HashEq[A]]
    """
    Setch is a portmanteau of set and CHoice.
    A Setch is a set-like object that is optimized for choosing a random item from it.
    """

type SetchIs[A] is HashSetch[A, HashIs[A!]]

class HashSetch[A, H: HashFunction[A!] val] is Comparable[HashSetch[A, H] box]
    """
    A damn shame about that inheritance.  Most of this code is verbatim stdlib Set code with exceptions being:
        * create
        * clear
        * set
        * unset
        * extract
        * choose
    (Not including places that type HashSet has been replaced with HashSetch and SetValues with SetchValues.)
    """
    embed _map: HashMap[A!, USize, H]
    embed _items: Array[A]

    embed _rand: Rand

    new create(prealloc: USize = 8) =>
        _map = _map.create(prealloc)
        _items = Array[A](prealloc)

        _rand = Rand(Time.millis())

    fun size(): USize =>
        _map.size()

    fun space(): USize =>
        _map.space()

    fun apply(value: box->A!): this->A ? =>
        _items(_map(value)?)?

    fun contains(value: box->A!): Bool =>
        _map.contains(value)

    fun ref clear() =>
        _map.clear()
        _items.clear()

    fun ref set(value: A) =>
        if not _map.contains(value) then
            _map(value) = _items.size()
            _items.push(consume value)
        end

    fun ref unset(value: box->A!) =>
        try extract(value)? end

    fun ref extract(value: box->A!): A^ ? =>
        try
            let position: USize = _map.remove(consume value)?._2
            let last_item: A = _items.pop()?

            if position != _items.size() then
                _map.update(last_item, position)
                _items.update(consume position, consume last_item)?
            else
                consume last_item
            end
        else
            error
        end

    fun ref union(that: Iterator[A^]) =>
        for value in that do
            set(consume value)
        end

    fun ref intersect[K: HashFunction[box->A!] val = H](
        that: HashSetch[box->A!, K])
    =>
        let start_size = _map.size()
        var seen: USize = 0
        var i: USize = -1
        while seen < start_size do
            try
                i = next_index(i)?
                if not that.contains(index(i)?) then
                    unset(index(i)?)
                end
            end
            seen = seen + 1
        end

    fun ref difference(that: Iterator[A^]) =>
        for value in that do
            try
                extract(value)?
            else
                set(consume value)
            end
        end

    fun ref remove(that: Iterator[box->A!]) =>
        for value in that do
            unset(value)
        end

    fun add[K: HashFunction[this->A!] val = H](
        value: this->A!)
        : HashSetch[this->A!, K]^
    =>
        clone[K]() .> set(value)

    fun sub[K: HashFunction[this->A!] val = H](
        value: this->A!)
        : HashSetch[this->A!, K]^
    =>
        clone[K]() .> unset(value)

    fun op_or[K: HashFunction[this->A!] val = H](
        that: this->HashSetch[A, H])
        : HashSetch[this->A!, K]^
    =>
        let r = clone[K]()

        for value in that.values() do
            r.set(value)
        end
        r

    fun op_and[K: HashFunction[this->A!] val=H](
        that: this->HashSetch[A, H])
        : HashSetch[this->A!, K]^
    =>
        let r = HashSetch[this->A!, K](size().min(that.size()))

        for value in values() do
            try
                that(value)?
                r.set(value)
            end
        end
        r

    fun op_xor[K: HashFunction[this->A!] val = H](
        that: this->HashSetch[A, H])
        : HashSetch[this->A!, K]^
    =>
        let r = HashSetch[this->A!, K](size().max(that.size()))

        for value in values() do
            try
                that(value)?
            else
                r.set(value)
            end
        end

        for value in that.values() do
            try
                this(value)?
            else
                r.set(value)
            end
        end
        r

    fun without[K: HashFunction[this->A!] val = H](
        that: this->HashSetch[A, H])
        : HashSetch[this->A!, K]^
    =>
        let r = HashSetch[this->A!, K](size())

        for value in values() do
            try
                that(value)?
            else
                r.set(value)
            end
        end
        r

    fun clone[K: HashFunction[this->A!] val = H](): HashSetch[this->A!, K]^ =>
        let r = HashSetch[this->A!, K](size())

        for value in values() do
            r.set(value)
        end
        r

    fun eq(that: HashSetch[A, H] box): Bool =>
        (size() == that.size()) and (this <= that)

    fun ne(that: HashSetch[A, H] box): Bool =>
        not (this == that)

    fun lt(that: HashSetch[A, H] box): Bool =>
        (size() < that.size()) and (this <= that)

    fun le(that: HashSetch[A, H] box): Bool =>
        try
            for value in values() do
                that(value)?
            end
            true
        else
            false
        end

    fun gt(that: HashSetch[A, H] box): Bool =>
        (size() > that.size()) and (that <= this)

    fun ge(that: HashSetch[A, H] box): Bool =>
        that <= this

    fun next_index(prev: USize = -1): USize ? =>
        _map.next_index(prev)?

    fun index(i: USize): this->A ? =>
        _items(_map.index(i)?._2)?

    fun values(): SetchValues[A, H, this->HashSetch[A, H]]^ =>
        SetchValues[A, H, this->HashSetch[A, H]](this)

    fun ref choose(): A ? =>
        """
        Return a random item from this or an error if this is empty.
        """
        _items(_rand.usize() %% size())?


class SetchValues[A, H: HashFunction[A!] val, S: HashSetch[A, H] #read] is
    Iterator[S->A]
    let _set: S
    var _i: USize = -1
    var _count: USize = 0

    new create(set: S) =>
        _set = set

    fun has_next(): Bool =>
        _count < _set.size()

    fun ref next(): S->A ? =>
        _i = _set.next_index(_i)?
        _count = _count + 1
        _set.index(_i)?