feat: add ema module

curve_std/constants.vy

@@ -0,0 +1,4 @@
+# pragma version 0.4.3
+
+WAD: constant(uint256) = 10**18
+SWAD: constant(int256) = 10**18

curve_std/ema.vy

@@ -0,0 +1,88 @@
+from curve_std import constants as c
+from snekmate.utils import math
+
+WAD: constant(uint256) = c.WAD
+MAX_EMAS: constant(uint256) = 10
+
+
+struct EMA:
+    ema_time: uint256
+    prev_value: uint256
+    prev_timestamp: uint256
+
+# List of allowed EMAs ids, useful to expose in the contract importing
+# this module if it allows to pass arbitrary ids.
+ALLOWED_EMAS: public(immutable(DynArray[String[4], MAX_EMAS]))
+
+# Vyper doesn't support passing value by reference yet, so we use
+# a mapping and a 4 character string as a pointer to the corresponding
+# storage slot.
+emas: public(HashMap[String[4], EMA])
+
+
+@deploy
+def __init__(_allowed_emas: DynArray[String[4], MAX_EMAS]):
+    ALLOWED_EMAS = _allowed_emas
+
+
+@internal
+@view
+def _is_allowed(_ema_id: String[4]) -> bool:
+    for ema_id: String[4] in ALLOWED_EMAS:
+        if ema_id == _ema_id:
+            return True
+    return False
+
+
+@internal
+def setup(_ema_id: String[4], _initial_value: uint256, _ema_time: uint256):
+    # Setting an ema_time of 0 is not allowed, as it would break the math
+    # Setting an ema_time of 1 is equivalent to no smoothing at all
+    assert self._is_allowed(_ema_id) # dev: id not allowed
+    assert _ema_time > 0 # dev: invalid ema_time
+    ema: EMA = self.emas[_ema_id]
+    ema.ema_time = _ema_time
+    ema.prev_value = _initial_value
+    ema.prev_timestamp = block.timestamp
+    self.emas[_ema_id] = ema
+
+
+@internal
+def set_ema_time(_ema_id: String[4], _ema_time: uint256):
+    # Setting an ema_time of 0 is not allowed, as it would break the math
+    # Setting an ema_time of 1 is equivalent to no smoothing at all
+    assert self._is_allowed(_ema_id) # dev: id not allowed
+    assert _ema_time > 0 # dev: invalid ema_time
+    ema: EMA = self.emas[_ema_id]
+    ema.ema_time = _ema_time
+    self.emas[_ema_id] = ema
+
+
+@internal
+@view
+def compute(_ema_id: String[4], _new_value: uint256) -> uint256:
+    assert self._is_allowed(_ema_id) # dev: id not allowed
+    ema: EMA = self.emas[_ema_id]
+    assert ema.ema_time > 0 # dev: ema not initialized
+    dt: uint256 = block.timestamp - ema.prev_timestamp
+
+    if dt == 0:
+        # The math below would return prev_value in this case anyway,
+        # but let's save some gas by skipping it
+        return ema.prev_value
+
+    mul: uint256 = convert(
+        math._wad_exp(-convert(dt * WAD // ema.ema_time, int256)), uint256
+    )
+    return (ema.prev_value * mul + _new_value * (WAD - mul)) // WAD
+
+
+@internal
+def update(_ema_id: String[4], _new_value: uint256) -> uint256:
+    smoothed: uint256 = self.compute(_ema_id, _new_value)
+    ema: EMA = self.emas[_ema_id]
+    assert ema.ema_time > 0 # dev: ema not initialized
+    ema.prev_value = smoothed
+    ema.prev_timestamp = block.timestamp
+    self.emas[_ema_id] = ema
+    return smoothed

tests/fuzzing/ema/test_fuzz.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import math
+
+import boa
+import pytest
+from hypothesis import given, settings, strategies as st
+
+from tests.utils.deployers import EMA_DEPLOYER
+
+TEST_EMA_ID = "fz00"
+
+
+def _float_reference(prev_value: int, new_value: int, ema_time: int, dt: int) -> float:
+    if dt == 0:
+        return float(prev_value)
+    weight = math.exp(-dt / ema_time)
+    return prev_value * weight + new_value * (1 - weight)
+
+
+def _deploy_ema():
+    return EMA_DEPLOYER.deploy([TEST_EMA_ID])
+
+
+@given(
+    prev_value=st.integers(min_value=0, max_value=10**20),
+    new_value=st.integers(min_value=0, max_value=10**20),
+    ema_time=st.integers(min_value=1, max_value=365 * 24 * 60 * 60),
+    dt=st.integers(min_value=0, max_value=365 * 24 * 60 * 60),
+)
+@settings(max_examples=200, deadline=None)
+def test_compute_matches_float_reference(prev_value, new_value, ema_time, dt):
+    with boa.env.anchor():
+        ema = _deploy_ema()
+        ema.internal.setup(TEST_EMA_ID, prev_value, ema_time)
+        if dt:
+            boa.env.time_travel(dt)
+        observed = ema.internal.compute(TEST_EMA_ID, new_value)
+
+    expected = _float_reference(prev_value, new_value, ema_time, dt)
+    tolerance = max(1.0, abs(expected) * 1e-9)
+    assert observed == pytest.approx(expected, rel=1e-6, abs=tolerance)

tests/unitary/ema/conftest.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+import pytest
+
+from tests.utils.deployers import EMA_DEPLOYER
+
+DUMMY_ALLOWED_EMAS = ["ema0", "ema1", "ema2", "ema3", "ema4"]
+
+
+@pytest.fixture
+def ema():
+    """Deploy an EMA contract seeded with predictable ids."""
+    return EMA_DEPLOYER.deploy(DUMMY_ALLOWED_EMAS)

tests/unitary/ema/helpers.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from decimal import ROUND_FLOOR, Decimal, localcontext
+
+WAD = 10**18
+
+
+def reference_compute(prev_value: int, new_value: int, ema_time: int, dt: int) -> int:
+    """Pure Python replica of the EMA compute logic."""
+    if dt == 0:
+        return prev_value
+
+    exponent_wad = -((dt * WAD) // ema_time)
+    with localcontext() as ctx:
+        ctx.prec = 80
+        exponent = Decimal(exponent_wad) / Decimal(WAD)
+        scaled = exponent.exp() * Decimal(WAD)
+        mul = int(scaled.to_integral_value(rounding=ROUND_FLOOR))
+    return (prev_value * mul + new_value * (WAD - mul)) // WAD

tests/unitary/ema/test_allowed_emas.py

@@ -0,0 +1,7 @@
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+
+
+def test_default_behavior(ema):
+    """Fixture-provided EMA exposes the configured id list."""
+    observed = [ema.ALLOWED_EMAS(i) for i in range(len(DUMMY_ALLOWED_EMAS))]
+    assert observed == DUMMY_ALLOWED_EMAS

tests/unitary/ema/test_ctor_ema.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from tests.utils.constants import MAX_EMAS
+from tests.utils.deployers import EMA_DEPLOYER
+
+
+def test_default_behavior():
+    """Constructor should persist the allowed EMA ids as provided."""
+    for size in range(MAX_EMAS + 1):
+        allowed = [f"e{index:03d}" for index in range(size)]
+        ema = EMA_DEPLOYER.deploy(allowed)
+        observed = [ema.ALLOWED_EMAS(i) for i in range(size)]
+        assert observed == allowed

tests/unitary/ema/test_internal_compute.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import boa
+
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+from tests.unitary.ema.helpers import reference_compute
+
+
+def test_default_behavior_no_time_elapsed(ema):
+    """When no time has passed the previous value is returned."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    ema.internal.setup(ema_id, 1234, 10)
+
+    assert ema.internal.compute(ema_id, 9999) == 1234
+
+
+def test_default_behavior_time_elapsed(ema):
+    """EMA output should match the Python reference implementation."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    prev_value = 1_000_000
+    new_value = 2_000_000
+    ema_time = 30
+    elapsed = 9
+
+    ema.internal.setup(ema_id, prev_value, ema_time)
+    boa.env.time_travel(elapsed)
+
+    observed = ema.internal.compute(ema_id, new_value)
+    expected = reference_compute(prev_value, new_value, ema_time, elapsed)
+    assert observed == expected
+
+
+def test_ema_not_initialized(ema):
+    """Computing without a setup should revert with the proper dev string."""
+    with boa.reverts(dev="ema not initialized"):
+        ema.internal.compute(DUMMY_ALLOWED_EMAS[0], 1)

tests/unitary/ema/test_internal_is_allowed.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+
+
+def test_default_behavior_known_id(ema):
+    """Internal check should return True for whitelisted ids."""
+    assert ema.internal._is_allowed(DUMMY_ALLOWED_EMAS[0]) is True
+
+
+def test_default_behavior_unknown_id(ema):
+    """Unknown ids should be rejected."""
+    assert ema.internal._is_allowed("nope") is False

tests/unitary/ema/test_internal_set_ema_time.py

@@ -0,0 +1,24 @@
+import boa
+
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+
+
+def test_default_behavior(ema):
+    """Internal setter updates the stored ema_time value."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    ema.internal.set_ema_time(ema_id, 42)
+
+    slot = ema.emas(ema_id)
+    assert slot.ema_time == 42
+
+
+def test_allowed_ema(ema):
+    """Setting a non-whitelisted id should revert."""
+    with boa.reverts(dev="id not allowed"):
+        ema.internal.set_ema_time("nope", 1)
+
+
+def test_non_zero_ema_time(ema):
+    """ema_time must be strictly positive."""
+    with boa.reverts(dev="invalid ema_time"):
+        ema.internal.set_ema_time(DUMMY_ALLOWED_EMAS[0], 0)

tests/unitary/ema/test_internal_setup.py

@@ -0,0 +1,29 @@
+import boa
+
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+
+
+def test_default_behavior(ema):
+    """setup should populate storage with the provided values."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    initial_value = 123456789
+    ema_time = 15
+
+    ema.internal.setup(ema_id, initial_value, ema_time)
+
+    slot = ema.emas(ema_id)
+    assert slot.ema_time == ema_time
+    assert slot.prev_value == initial_value
+    assert slot.prev_timestamp > 0
+
+
+def test_allowed_ema(ema):
+    """setup rejects ids that are not whitelisted."""
+    with boa.reverts(dev="id not allowed"):
+        ema.internal.setup("nope", 1, 1)
+
+
+def test_non_zero_ema_time(ema):
+    """setup enforces ema_time > 0."""
+    with boa.reverts(dev="invalid ema_time"):
+        ema.internal.setup(DUMMY_ALLOWED_EMAS[0], 1, 0)

tests/unitary/ema/test_internal_update.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import boa
+
+from tests.unitary.ema.conftest import DUMMY_ALLOWED_EMAS
+from tests.unitary.ema.helpers import reference_compute
+
+
+def test_default_behavior_no_time_elapsed(ema):
+    """Update should be a no-op if block.timestamp did not move."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    prev_value = 10**9
+    ema.internal.setup(ema_id, prev_value, 42)
+
+    before_slot = ema.emas(ema_id)
+    result = ema.internal.update(ema_id, 5 * 10**9)
+
+    after_slot = ema.emas(ema_id)
+    assert result == prev_value
+    assert after_slot.prev_value == prev_value
+    assert after_slot.prev_timestamp == before_slot.prev_timestamp
+
+
+def test_default_behavior_time_elapsed(ema):
+    """Update should smooth the new value and persist it in storage."""
+    ema_id = DUMMY_ALLOWED_EMAS[0]
+    prev_value = 500_000
+    new_value = 1_500_000
+    ema_time = 20
+    elapsed = 7
+
+    ema.internal.setup(ema_id, prev_value, ema_time)
+    boa.env.time_travel(elapsed)
+
+    result = ema.internal.update(ema_id, new_value)
+    expected = reference_compute(prev_value, new_value, ema_time, elapsed)
+
+    slot = ema.emas(ema_id)
+    assert result == expected
+    assert slot.prev_value == expected
+    assert slot.prev_timestamp == boa.env.timestamp
+
+
+def test_ema_not_initialized(ema):
+    """Calling update without setup should revert."""
+    with boa.reverts(dev="ema not initialized"):
+        ema.internal.update(DUMMY_ALLOWED_EMAS[0], 1)

tests/utils/constants.py

@@ -0,0 +1,5 @@
+import boa
+
+from tests.utils.deployers import EMA_DEPLOYER
+
+MAX_EMAS = EMA_DEPLOYER._constants.MAX_EMAS

tests/utils/deployers.py

@@ -0,0 +1,5 @@
+import boa
+
+CONTRACTS_PATH = "curve_std/"
+
+EMA_DEPLOYER = boa.load_partial(CONTRACTS_PATH + "ema.vy")