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