Enumerable Set

Create a set data structure on Vyper

Vyper does not have a set data structure. A set is like an array but only have unique members. We will create one in this article. Our set is an address data type set and it has a fixed size, which is 56 as you'll see later in the code. (Vyper does not have a dynamic array yet). You can change the data type and the size for your own project if you want to. As usual, initialize Mamba project directory.


(.venv) $ mkdir enumerable_set
(.venv) $ cd enumerable_set
(.venv) $ mamba init

Create a new file, contracts/EnumerableSet.vy. Add the following code to the file:


"""
@title An EnumerableSet design pattern
@license MIT
@author Arjuna Sky Kok
@notice Translated from OpenZeppelin's EnumerableSet code: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/utils/EnumerableSet.sol
"""

SET_LENGTH: constant(uint256) = 56

set_indexes: public(HashMap[address, int128])
struct Set:
    _values: address[SET_LENGTH]
    _index: int128

set: public(Set)

@internal
def contains_in_set(v: address) -> bool:
    return self.set_indexes[v] != 0

@internal
def add_to_set(v: address) -> bool:
    if not self.contains_in_set(v):
        self.set._values[self.set._index] = v
        self.set_indexes[v] = self.set._index + 1
        self.set._index = self.set._index + 1
        return True
    else:
        return False

@internal
def remove_from_set(v: address) -> bool:
    valueIndex : int128 = self.set_indexes[v]

    if valueIndex != 0:
        toDeleteIndex : int128 = valueIndex - 1
        lastIndex : int128 = self.set._index - 1
        lastValue : address = self.set._values[lastIndex]
        self.set._values[toDeleteIndex] = lastValue
        self.set_indexes[v] = 0
        self.set._values[lastIndex] = ZERO_ADDRESS
        self.set._index = self.set._index - 1
        return True
    else:
        return False

@internal
def length_of_set() -> int128:
    return self.set._index

@internal
def at_set(index: int128) -> address:
    return self.set._values[index]

@external
def __init__():
    pass

@external
def user_add(v: address) -> bool:
    return self.add_to_set(v)

@external
def user_remove(v: address) -> bool:
    return self.remove_from_set(v)

@external
def user_contains(v: address) -> bool:
    return self.contains_in_set(v)

@external
def user_length() -> int128:
    return self.length_of_set()

@external
def user_at(index: int128) -> address:
    return self.at_set(index)

We have a Set struct, which has a _values array and a _index pointer. We keep members in the _values array. The _index pointer points to the first empty space in the array. We also have set_indexes which is a map. The purpose of this map is to check whether there is a value or not in this set data structure. The API of the set data structure are contains_in_set, add_to_set, remove_from_set, length_of_set, at_set. The other methods are merely for interacting with this set data structure for demonstration purpose.

Let's compile it.


(.venv) $ mamba compile

Let's create a test file, test/test_enumerable_set.py.


from black_mamba.testlib import contract, eth_tester
import pytest


def test_init(eth_tester):
    accounts = eth_tester.get_accounts()
    enumset_contract = contract("EnumerableSet", [])
    assert enumset_contract.functions.user_length().call() == 0
    assert enumset_contract.functions.user_contains(accounts[0]).call() == False

def test_add(eth_tester):
    accounts = eth_tester.get_accounts()
    enumset_contract = contract("EnumerableSet", [])

    enumset_contract.functions.user_add(accounts[0]).transact()
    assert enumset_contract.functions.user_length().call() == 1
    assert enumset_contract.functions.user_contains(accounts[0]).call() == True
    assert enumset_contract.functions.user_at(0).call() == accounts[0]

    enumset_contract.functions.user_add(accounts[0]).transact()
    assert enumset_contract.functions.user_length().call() == 1
    assert enumset_contract.functions.user_contains(accounts[0]).call() == True
    assert enumset_contract.functions.user_contains(accounts[1]).call() == False
    assert enumset_contract.functions.user_at(0).call() == accounts[0]

    enumset_contract.functions.user_add(accounts[1]).transact()
    assert enumset_contract.functions.user_length().call() == 2
    assert enumset_contract.functions.user_contains(accounts[0]).call() == True
    assert enumset_contract.functions.user_at(0).call() == accounts[0]
    assert enumset_contract.functions.user_contains(accounts[1]).call() == True
    assert enumset_contract.functions.user_at(1).call() == accounts[1]

def test_remove(eth_tester):
    accounts = eth_tester.get_accounts()
    enumset_contract = contract("EnumerableSet", [])

    enumset_contract.functions.user_add(accounts[0]).transact()
    enumset_contract.functions.user_add(accounts[1]).transact()
    enumset_contract.functions.user_add(accounts[2]).transact()

    enumset_contract.functions.user_remove(accounts[1]).transact()
    assert enumset_contract.functions.user_length().call() == 2
    assert enumset_contract.functions.user_contains(accounts[0]).call() == True
    assert enumset_contract.functions.user_contains(accounts[1]).call() == False
    assert enumset_contract.functions.user_contains(accounts[2]).call() == True
    assert enumset_contract.functions.user_at(0).call() == accounts[0]
    assert enumset_contract.functions.user_at(1).call() == accounts[2]
    assert enumset_contract.functions.user_at(2).call() == "0x0000000000000000000000000000000000000000"

Now you can execute the test file. The test should run with a green result.


(.venv) $ py.test test/test_enumerable_set.py