Very simple model of Uniswap V3 Fees

uniswap-v3-fees.py

from attrs import define


@define
class Tick:
    fee_growth_outside: int

    def cross(self, fee_growth_global: int):
        self.fee_growth_outside = fee_growth_global - self.fee_growth_outside


@define
class Pool:
    current_tick: int
    fee_growth_global: int
    ticks: dict[int, Tick]

    def init_tick(self, tick: int) -> Tick:
        if self.current_tick >= tick:
            fee_growth_outside = self.fee_growth_global
        else:
            fee_growth_outside = 0
        new_tick = Tick(fee_growth_outside)
        self.ticks[tick] = new_tick
        return new_tick

    def up(self):
        self.current_tick += 1
        # self.get(self.current_tick).cross(self.fee_growth_global)

    def down(self):
        self.get(self.current_tick).cross(self.fee_growth_global)
        self.current_tick -= 1

    def collect(self, fees: int):
        self.fee_growth_global += fees

    def get(self, tick: int) -> Tick:
        if (tobj := self.ticks.get(tick)) is not None:
            return tobj
        return self.init_tick(tick)

    def get_ticks(self, start_tick: int, end_tick: int) -> list[Tick]:
        return [
            self.get(tick)
            for tick in range(start_tick, end_tick + 1)
        ]

    def fee_above(self, tick: int) -> int:
        tobj = self.get(tick)
        if self.current_tick >= tick:
            return self.fee_growth_global - tobj.fee_growth_outside
        else:
            return tobj.fee_growth_outside

    def fee_below(self, tick: int) -> int:
        tobj = self.get(tick)
        if self.current_tick >= tick:
            return tobj.fee_growth_outside
        else:
            return self.fee_growth_global - tobj.fee_growth_outside

    def tick_fees(self, tick: int) -> int:
        return self.fee_rate(tick, tick + 1)

    def fee_rate(self, lower: int, upper: int) -> int:
        return self.fee_growth_global - self.fee_above(upper) - self.fee_below(lower)

    def total_fees(self, start: int, end: int) -> int:
        return sum(
            self.tick_fees(tick)
            for tick in range(start, end + 1)
        )


def main():
    pool = Pool(3, 0, {})

    pool.collect(2)

    for tick in range(1, 5 + 1):
        pool.init_tick(tick)

    print(f'total: {pool.total_fees(1, 5)}')
    for tick in range(1, 5 + 1):
        print(f'[{tick}] fees: {pool.tick_fees(tick)}')

    pool.collect(1)

    print(f'\ntotal: {pool.total_fees(1, 5)}')
    for tick in range(1, 5 + 1):
        print(f'[{tick}] fees: {pool.tick_fees(tick)}')

    pool.up()

    print(f'\ntotal: {pool.total_fees(1, 5)}')
    for tick in range(1, 5 + 1):
        print(f'[{tick}] fees: {pool.tick_fees(tick)}')

    pool.up()

    print(f'\ntotal: {pool.total_fees(1, 5)}')
    for tick in range(1, 5 + 1):
        print(f'[{tick}] fees: {pool.tick_fees(tick)}')

    pool.down()

    print(f'\ntotal: {pool.total_fees(1, 5)}')
    for tick in range(1, 5 + 1):
        print(f'[{tick}] fees: {pool.tick_fees(tick)}')


if __name__ == '__main__':
    main()