router.mpc

exec(open("../cfmms.mpc").read())
exec(open("../optimizers.mpc").read())

"""
sfix.round_nearest = True
cfix.set_precision(f=16, k=40)
sfix.set_precision(f=16, k=40)

program.use_edabit(True)
program.use_split(4)
"""

class ArbitrageRouter:

    """
    cfmms: List[CFMM]
    n_tokens: sfix
    tendered_assets: List[List[sfix]]
    received_assets: List[List[sfix]]
    price_vector: List[sfix]
    """
    def __init__(self, cfmms, n_tokens, ln_objective):
        self.cfmms = cfmms
        self.big_delta = MultiArray([len(cfmms), 2], sfix)
        self.big_lambda = MultiArray([len(cfmms), 2], sfix)
        self.linearnn_objective = ln_objective

        @for_range_opt(len(cfmms))
        def _(i):
            self.big_delta.assign_all(sfix(0))
            self.big_lambda.assign_all(sfix(0))

        self.v = sfix.Array(n_tokens)
        self.v.assign_all(sfix(0))

    def find_arb(self, price_vector):
        for i, cfmm in zip(list(range(len(self.cfmms))), self.cfmms):
            cfmm.find_arb(self.big_delta[i], self.big_lambda[i], price_vector)
            

    def route(self, optimizer, price_vector=None):
        if price_vector is None:
            self.v.assign_all(sfix(1/2)) # OG code says to use initial marginal price here
        else:
            self.v.assign(price_vector)

        def fn(v):
            tmp = Matrix(len(self.v), 1, sfix)
            @for_range_opt(len(self.v))
            def _(i):
                tmp[i].assign(((v[i][0] <= self.v[i][0]).if_else(0, 1)))

            @if_((tmp[:] == 0).reveal())
            def _():
                self.find_arb(v)
                self.v.assign(v)

            accumulator = sfix(0)

            for (big_delta, big_lambda, cfmm) in zip(self.big_delta, self.big_lambda, self.cfmms):
                tmp_big_delta = Matrix(len(big_delta), 1, sfix)
                tmp_big_lambda = Matrix(len(big_lambda), 1, sfix)
                accumulator += (tmp_big_lambda.dot(v.get_vector_by_indices(*cfmm.Ai)) - tmp_big_delta.dot(v.get_vector_by_indices(*cfmm.Ai)))

            return self.linearnn_objective.conjugate(v)[:] + accumulator

        def g(v):
            G = Matrix(len(self.v), 1, sfix) 

            tmp = Matrix(len(self.v), 1, sfix)
            @for_range_opt(len(self.v))
            def _(i):
                tmp[i].assign(((v[i][0] <= self.v[i][0]).if_else(0, 1)))

            @if_((tmp[:] == 0).reveal())
            def _():
                #tmp_v = Array.create_from(v)
                self.find_arb(v)
                self.v.assign(v)

            G.assign(self.linearnn_objective.grad(v))

            for big_delta, big_lambda, c in zip(self.big_delta, self.big_lambda, self.cfmms):
                tmp2 = G.get_slice_vector(c.Ai) 
                tmp2 += (big_lambda - big_delta)
                G.assign(tmp2)

            return G
        

        self.find_arb(self.v)
        new_price_vector = optimizer.optimize(self.v, fn, g, n_iter=2)[0]
        self.v.assign(new_price_vector)
        self.find_arb(self.v)


    def netflows(self):
        psi = Matrix(len(self.v), 1, sfix)
        for (big_delta, big_lambda, cfmm) in zip(self.big_delta, self.big_lambda, self.cfmms):
            tmp = big_lambda - big_delta
            psi.assign_vector_by_indices(tmp, *cfmm.Ai)
        return psi