precompile/stableswap/contract.go

// Copyright (C) 2026, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

// Package stableswap implements Curve-style StableSwap AMM as a precompile.
// Address: 0x0400000000000000000000000000000000000060 (DEX extended range)
//
// Implements the StableSwap invariant:
//
//	An^n * sum(x_i) + D = A * D^(n+1) / (n^n * prod(x_i))
//
// Operations:
//   - 0x01: GetDy(i, j, dx, balances, amp) -> output amount for swap i->j
//   - 0x02: AddLiquidity(amounts, balances, amp, totalSupply) -> LP tokens minted
//   - 0x03: RemoveLiquidity(lpAmount, balances, totalSupply) -> underlying amounts
//   - 0x04: GetD(balances, amp) -> invariant D
//
// Gas: 5,000 base (vs ~50,000 in Solidity -- 10x savings)
// Used for: LUSD/USDC, LETH/stETH, any pegged-asset pools
package stableswap

import (
	"encoding/binary"
	"errors"
	"math/big"

	"github.com/luxfi/geth/common"
	"github.com/luxfi/precompile/contract"
)

var (
	ContractAddress = common.HexToAddress("0x0400000000000000000000000000000000000060")

	Precompile = &stableSwapPrecompile{}

	_ contract.StatefulPrecompiledContract = Precompile

	ErrInvalidInput  = contract.ErrInvalidInput
	ErrInvalidIndex  = errors.New("invalid token index")
	ErrZeroLiquidity = errors.New("zero liquidity")
	ErrConvergence   = errors.New("newton method did not converge")
	ErrUnknownOp     = errors.New("unknown stableswap operation")
	ErrTooManyTokens = errors.New("token count exceeds maxTokens")
	ErrInvalidAmp    = errors.New("amplification coefficient must be positive")
)

const (
	OpGetDy           = 0x01
	OpAddLiquidity    = 0x02
	OpRemoveLiquidity = 0x03
	OpGetD            = 0x04

	GasBase = 5000

	maxIterations = 256

	// maxTokens caps the pool's token count. Curve's MAX_COINS is 8; real
	// pegged-asset pools never exceed it. The cap is a DETERMINISM/DoS bound, not a
	// feature limit: the Newton solvers below run up to maxIterations passes, each a
	// nested loop of big.Int mul/div over n tokens whose operands grow ~O(n*256) bits.
	// Without an upper bound on n an attacker supplies a huge token count (paying only
	// flat calldata gas) and forces super-linear in-consensus compute on every
	// validator — a one-tx halt. Capping n makes the worst-case work a fixed constant.
	maxTokens = 8

	// gasPerTokenIter prices ONE token's contribution to ONE Newton iteration. Gas is
	// charged as solves*n*maxIterations*gasPerTokenIter so it always reflects the
	// BOUNDED WORST CASE (every iteration runs), never the lucky early-convergence
	// cost — an attacker who crafts balances that maximize iterations cannot underpay.
	gasPerTokenIter = 10
)

var (
	zero = big.NewInt(0)
	one  = big.NewInt(1)
)

type stableSwapPrecompile struct{}

func (p *stableSwapPrecompile) Run(
	accessibleState contract.AccessibleState,
	caller common.Address,
	addr common.Address,
	input []byte,
	suppliedGas uint64,
	readOnly bool,
) ([]byte, uint64, error) {
	gas, err := contract.DeductGas(suppliedGas, GasBase)
	if err != nil {
		return nil, 0, err
	}
	if len(input) < 1 {
		return nil, gas, ErrInvalidInput
	}

	switch input[0] {
	case OpGetDy:
		return getDy(input[1:], gas)
	case OpAddLiquidity:
		return addLiquidity(input[1:], gas)
	case OpRemoveLiquidity:
		return removeLiquidity(input[1:], gas)
	case OpGetD:
		return getD(input[1:], gas)
	default:
		return nil, gas, ErrUnknownOp
	}
}

// readAmp parses the amplification coefficient A and refuses a non-positive one.
//
// A is the ONE caller-supplied scalar the Newton solvers divide by, and at A == 0 all
// three of these divisions are by zero — reachable from calldata, so a panic here is a
// chain halt, not a failed call:
//
//	computeY: c.Div(c, ann*n)      // ann == A*n == 0
//	computeY: b = s + D/ann        // ann == 0
//	computeD: d.Div(num, denom)    // denom == (A*n - 1)*D + (n+1)*dP == (n+1)*dP - D,
//	                               // which crosses zero exactly for reachable balances
//
// Refusing is also the mathematically honest answer: the StableSwap invariant
// A*n^n*S + D = A*D*n^n + D^(n+1)/(n^n*prod(x)) is not defined at A = 0 — the pool
// degenerates to constant product, a different curve with a different solver.
func readAmp(b []byte) (*big.Int, error) {
	amp := new(big.Int).SetBytes(b)
	if amp.Sign() <= 0 {
		return nil, ErrInvalidAmp
	}
	return amp, nil
}

// chargeNewtonGas deducts gas for `solves` Newton solves over n tokens at the bounded
// worst case (maxIterations passes each). It MUST be called AFTER n is capped at
// maxTokens (so the product cannot overflow or be unbounded) and BEFORE any solve runs,
// so an attacker pays for the worst case up front rather than after the work lands.
func chargeNewtonGas(gas uint64, n, solves int) (uint64, error) {
	work := uint64(solves) * uint64(n) * uint64(maxIterations) * gasPerTokenIter
	return contract.DeductGas(gas, work)
}

// Input layout for GetDy:
//
//	[0:4]   i (uint32)
//	[4:8]   j (uint32)
//	[8:40]  dx (uint256)
//	[40:44] n (uint32, number of tokens)
//	[44:76] amp (uint256)
//	[76:]   balances (n * 32 bytes)
func getDy(data []byte, gas uint64) ([]byte, uint64, error) {
	if len(data) < 76 {
		return nil, gas, ErrInvalidInput
	}
	i := binary.BigEndian.Uint32(data[0:4])
	j := binary.BigEndian.Uint32(data[4:8])
	dx := new(big.Int).SetBytes(data[8:40])
	n := int(binary.BigEndian.Uint32(data[40:44]))
	amp, err := readAmp(data[44:76])
	if err != nil {
		return nil, gas, err
	}

	if n < 2 || len(data) < 76+n*32 {
		return nil, gas, ErrInvalidInput
	}
	if n > maxTokens {
		return nil, gas, ErrTooManyTokens
	}
	if int(i) >= n || int(j) >= n || i == j {
		return nil, gas, ErrInvalidIndex
	}
	// getDy runs TWO Newton solves: computeD then computeY. Charge both up front.
	gas, err = chargeNewtonGas(gas, n, 2)
	if err != nil {
		return nil, 0, err
	}

	balances := make([]*big.Int, n)
	for k := range n {
		off := 76 + k*32
		balances[k] = new(big.Int).SetBytes(data[off : off+32])
	}

	d, err := computeD(balances, amp, n)
	if err != nil {
		return nil, gas, err
	}

	// Compute new balance for token j after adding dx to token i
	newBalances := make([]*big.Int, n)
	for k := range n {
		newBalances[k] = new(big.Int).Set(balances[k])
	}
	newBalances[int(i)].Add(newBalances[int(i)], dx)

	y, err := computeY(newBalances, amp, n, int(j), d)
	if err != nil {
		return nil, gas, err
	}

	// dy = balances[j] - y - 1 (subtract 1 for rounding)
	dy := new(big.Int).Sub(balances[int(j)], y)
	dy.Sub(dy, one)
	if dy.Sign() < 0 {
		dy.SetInt64(0)
	}
	return padTo32(dy.Bytes()), gas, nil
}

// Input layout for AddLiquidity:
//
//	[0:4]   n (uint32)
//	[4:36]  amp (uint256)
//	[36:68] totalSupply (uint256)
//	[68:]   amounts (n * 32) + balances (n * 32)
func addLiquidity(data []byte, gas uint64) ([]byte, uint64, error) {
	if len(data) < 68 {
		return nil, gas, ErrInvalidInput
	}
	n := int(binary.BigEndian.Uint32(data[0:4]))
	amp, err := readAmp(data[4:36])
	if err != nil {
		return nil, gas, err
	}
	totalSupply := new(big.Int).SetBytes(data[36:68])

	if n < 2 || len(data) < 68+2*n*32 {
		return nil, gas, ErrInvalidInput
	}
	if n > maxTokens {
		return nil, gas, ErrTooManyTokens
	}
	// addLiquidity runs TWO Newton solves: computeD(balances) then computeD(newBalances).
	gas, err = chargeNewtonGas(gas, n, 2)
	if err != nil {
		return nil, 0, err
	}

	amounts := make([]*big.Int, n)
	balances := make([]*big.Int, n)
	for k := range n {
		off := 68 + k*32
		amounts[k] = new(big.Int).SetBytes(data[off : off+32])
	}
	for k := range n {
		off := 68 + n*32 + k*32
		balances[k] = new(big.Int).SetBytes(data[off : off+32])
	}

	d0, err := computeD(balances, amp, n)
	if err != nil {
		return nil, gas, err
	}

	newBalances := make([]*big.Int, n)
	for k := range n {
		newBalances[k] = new(big.Int).Add(balances[k], amounts[k])
	}
	d1, err := computeD(newBalances, amp, n)
	if err != nil {
		return nil, gas, err
	}

	if totalSupply.Sign() == 0 {
		// First deposit: mint D1 tokens
		return padTo32(d1.Bytes()), gas, nil
	}

	// D0 is the divisor below, and computeD returns 0 for an all-zero pool. A caller
	// supplies balances and totalSupply independently, so "every balance zero yet LP
	// tokens outstanding" is reachable from calldata — and dividing by it is a panic,
	// i.e. a chain halt. It is also a contradictory pool: shares backed by nothing, with
	// no ratio to price the deposit against. Refuse rather than invent one.
	if d0.Sign() == 0 {
		return nil, gas, ErrZeroLiquidity
	}

	// mint = totalSupply * (D1 - D0) / D0, rounded DOWN so the depositor is never
	// credited a share the pool did not receive (rounding favours the pool).
	diff := new(big.Int).Sub(d1, d0)
	mint := new(big.Int).Mul(totalSupply, diff)
	mint.Div(mint, d0)
	return padTo32(mint.Bytes()), gas, nil
}

// Input layout for RemoveLiquidity:
//
//	[0:4]   n (uint32)
//	[4:36]  lpAmount (uint256)
//	[36:68] totalSupply (uint256)
//	[68:]   balances (n * 32)
func removeLiquidity(data []byte, gas uint64) ([]byte, uint64, error) {
	if len(data) < 68 {
		return nil, gas, ErrInvalidInput
	}
	n := int(binary.BigEndian.Uint32(data[0:4]))
	lpAmount := new(big.Int).SetBytes(data[4:36])
	totalSupply := new(big.Int).SetBytes(data[36:68])

	if n < 2 || len(data) < 68+n*32 {
		return nil, gas, ErrInvalidInput
	}
	if n > maxTokens {
		return nil, gas, ErrTooManyTokens
	}
	// removeLiquidity runs NO Newton solve — just a bounded O(n) proportional split, so
	// GasBase covers it once n is capped. No extra per-token charge needed.
	if totalSupply.Sign() == 0 {
		return nil, gas, ErrZeroLiquidity
	}

	balances := make([]*big.Int, n)
	for k := range n {
		off := 68 + k*32
		balances[k] = new(big.Int).SetBytes(data[off : off+32])
	}

	// Return proportional share: amounts[i] = balances[i] * lpAmount / totalSupply
	result := make([]byte, n*32)
	for k := range n {
		amt := new(big.Int).Mul(balances[k], lpAmount)
		amt.Div(amt, totalSupply)
		b := padTo32(amt.Bytes())
		copy(result[k*32:(k+1)*32], b)
	}
	return result, gas, nil
}

// Input layout for GetD:
//
//	[0:4]  n (uint32)
//	[4:36] amp (uint256)
//	[36:]  balances (n * 32)
func getD(data []byte, gas uint64) ([]byte, uint64, error) {
	if len(data) < 36 {
		return nil, gas, ErrInvalidInput
	}
	n := int(binary.BigEndian.Uint32(data[0:4]))
	amp, err := readAmp(data[4:36])
	if err != nil {
		return nil, gas, err
	}
	if n < 2 || len(data) < 36+n*32 {
		return nil, gas, ErrInvalidInput
	}
	if n > maxTokens {
		return nil, gas, ErrTooManyTokens
	}
	// getD runs ONE Newton solve: computeD.
	gas, err = chargeNewtonGas(gas, n, 1)
	if err != nil {
		return nil, 0, err
	}

	balances := make([]*big.Int, n)
	for k := range n {
		off := 36 + k*32
		balances[k] = new(big.Int).SetBytes(data[off : off+32])
	}

	d, err := computeD(balances, amp, n)
	if err != nil {
		return nil, gas, err
	}
	return padTo32(d.Bytes()), gas, nil
}

// computeD solves for D using Newton's method on the StableSwap invariant:
//
//	A * n^n * S + D = A * D * n^n + D^(n+1) / (n^n * prod(x_i))
//
// where S = sum(x_i).
func computeD(balances []*big.Int, amp *big.Int, n int) (*big.Int, error) {
	s := new(big.Int)
	for _, b := range balances {
		s.Add(s, b)
	}
	if s.Sign() == 0 {
		return big.NewInt(0), nil
	}

	nn := big.NewInt(int64(n))
	ann := new(big.Int).Mul(amp, nn) // A * n

	d := new(big.Int).Set(s)
	for range maxIterations {
		// dP = D^(n+1) / (n^n * prod(x_i))
		// Computed iteratively: dP = D, then dP = dP * D / (n * x_i) for each i
		dP := new(big.Int).Set(d)
		for _, x := range balances {
			if x.Sign() == 0 {
				return nil, ErrZeroLiquidity
			}
			// dP = dP * D / (x * n)
			dP.Mul(dP, d)
			dP.Div(dP, new(big.Int).Mul(x, nn))
		}

		prevD := new(big.Int).Set(d)

		// numerator   = (A*n*S + n*dP) * D
		// denominator = (A*n - 1) * D + (n+1) * dP
		num := new(big.Int).Mul(ann, s)
		num.Add(num, new(big.Int).Mul(nn, dP))
		num.Mul(num, d)

		denom := new(big.Int).Sub(ann, one)
		denom.Mul(denom, d)
		nPlus1 := new(big.Int).Add(nn, one)
		denom.Add(denom, new(big.Int).Mul(nPlus1, dP))

		d.Div(num, denom)

		// Check convergence: |d - prevD| <= 1
		diff := new(big.Int).Sub(d, prevD)
		diff.Abs(diff)
		if diff.Cmp(one) <= 0 {
			return d, nil
		}
	}
	return nil, ErrConvergence
}

// computeY solves for the new balance of token j given:
// - updated balances (with token j excluded)
// - A, n, D
// Uses Newton's method on the StableSwap invariant, solving for x_j.
func computeY(balances []*big.Int, amp *big.Int, n int, j int, d *big.Int) (*big.Int, error) {
	nn := big.NewInt(int64(n))
	ann := new(big.Int).Mul(amp, nn)

	// c = D^(n+1) / (n^n * A * n * prod(x_i, i != j))
	c := new(big.Int).Set(d)
	s := new(big.Int)
	for k := range n {
		if k == j {
			continue
		}
		x := balances[k]
		if x.Sign() == 0 {
			return nil, ErrZeroLiquidity
		}
		s.Add(s, x)
		c.Mul(c, d)
		c.Div(c, new(big.Int).Mul(x, nn))
	}
	c.Mul(c, d)
	c.Div(c, new(big.Int).Mul(ann, nn))

	b := new(big.Int).Add(s, new(big.Int).Div(d, ann))

	y := new(big.Int).Set(d)
	for range maxIterations {
		prevY := new(big.Int).Set(y)

		// y_next = ceil((y^2 + c) / (2*y + b - D))
		yy := new(big.Int).Mul(y, y)
		yy.Add(yy, c)

		denom := new(big.Int).Add(new(big.Int).Mul(big.NewInt(2), y), b)
		denom.Sub(denom, d)
		// A non-positive denominator cannot yield a meaningful balance: zero divides,
		// and negative flips the sign of y — whose magnitude then flows into
		// dy = balance - y - 1 as if it were positive, quoting MORE than the pool holds
		// (big.Int.Bytes() drops the sign, so the wrong answer looks plausible).
		if denom.Sign() <= 0 {
			return nil, ErrConvergence
		}

		divCeil(y, yy, denom)

		diff := new(big.Int).Sub(y, prevY)
		diff.Abs(diff)
		if diff.Cmp(one) <= 0 {
			return y, nil
		}
	}
	return nil, ErrConvergence
}

// divCeil sets z = ceil(num/denom) for denom > 0.
//
// The rounding DIRECTION here is a solvency property, not a style choice. y is the
// balance token j must retain for the invariant to hold, and the caller is paid
// balance[j] - y - 1: rounding y DOWN pays the caller MORE.
//
// Newton on f(y) = y^2 + (b-D)y - c converges monotonically downward to the root
// from y0 = D, because f is convex — so every exact iterate stays ABOVE the root.
// Floor division broke that: it dragged each integer iterate below its exact
// counterpart until the sequence crossed under the root, measured at up to 2 wei
// low in 55% of solves, which made the single `-1` margin in getDy insufficient and
// let a swap DECREASE the pool invariant D. Ceiling division restores the
// one-sidedness the proof depends on: each integer iterate is >= the exact iterate
// >= the root, so the quote is always conservative.
func divCeil(z, num, denom *big.Int) {
	var rem big.Int
	z.QuoRem(num, denom, &rem)
	if rem.Sign() > 0 {
		z.Add(z, one)
	}
}

func padTo32(b []byte) []byte {
	if len(b) >= 32 {
		return b[len(b)-32:]
	}
	padded := make([]byte, 32)
	copy(padded[32-len(b):], b)
	return padded
}