harb/kraiken-lib/src/snatch.ts

import type { Positions } from './__generated__/graphql.js';
import { toBigIntId } from './ids.js';

export interface SnatchablePosition {
  id: bigint;
  stakeShares: bigint;
  taxRate: number;
  taxRateIndex?: number;
  owner?: string | null;
}

export interface SnatchSelectionOptions {
  shortfallShares: bigint;
  maxTaxRate: number;
  includeOwned?: boolean;
  recipientAddress?: string | null;
}

export interface SnatchSelectionResult {
  selected: SnatchablePosition[];
  coveredShares: bigint;
  remainingShortfall: bigint;
  maxSelectedTaxRate?: number;
  maxSelectedTaxRateIndex?: number;
}

const SHARE_SCALE = 1_000_000n;

function normaliseAddress(value?: string | null): string {
  return (value ?? '').toLowerCase();
}

export function minimumTaxRate<T extends { taxRate: number }>(positions: T[], fallback: number = 0): number {
  if (!positions.length) return fallback;
  return positions.reduce((min, position) => (position.taxRate < min ? position.taxRate : min), Number.POSITIVE_INFINITY);
}

export function selectSnatchPositions(candidates: SnatchablePosition[], options: SnatchSelectionOptions): SnatchSelectionResult {
  const { shortfallShares, maxTaxRate, includeOwned = false, recipientAddress = null } = options;

  if (shortfallShares <= 0n) {
    return {
      selected: [],
      coveredShares: 0n,
      remainingShortfall: 0n,
    };
  }

  const recipientNormalised = normaliseAddress(recipientAddress);

  const filtered = candidates.filter(candidate => {
    if (candidate.taxRate >= maxTaxRate) {
      return false;
    }

    if (!includeOwned && candidate.owner) {
      return normaliseAddress(candidate.owner) !== recipientNormalised;
    }

    return true;
  });

  const sorted = filtered.slice().sort((a, b) => {
    if (a.taxRate === b.taxRate) {
      return Number(a.stakeShares - b.stakeShares);
    }
    return a.taxRate - b.taxRate;
  });

  const selection: SnatchablePosition[] = [];
  let remaining = shortfallShares;
  let covered = 0n;
  let maxSelectedTaxRate: number | undefined;
  let maxSelectedTaxRateIndex: number | undefined;

  for (const candidate of sorted) {
    if (remaining <= 0n) break;
    if (candidate.stakeShares <= 0n) continue;

    selection.push(candidate);
    remaining -= candidate.stakeShares;
    if (remaining < 0n) {
      covered = shortfallShares;
      remaining = 0n;
    } else {
      covered = shortfallShares - remaining;
    }

    if (maxSelectedTaxRate === undefined || candidate.taxRate > maxSelectedTaxRate) {
      maxSelectedTaxRate = candidate.taxRate;
      if (typeof candidate.taxRateIndex === 'number') {
        maxSelectedTaxRateIndex = candidate.taxRateIndex;
      }
    } else if (
      candidate.taxRate === maxSelectedTaxRate &&
      typeof candidate.taxRateIndex === 'number' &&
      candidate.taxRateIndex > (maxSelectedTaxRateIndex ?? -1)
    ) {
      maxSelectedTaxRateIndex = candidate.taxRateIndex;
    }
  }

  return {
    selected: selection,
    coveredShares: covered,
    remainingShortfall: remaining > 0n ? remaining : 0n,
    maxSelectedTaxRate,
    maxSelectedTaxRateIndex,
  };
}

export function getSnatchList(
  positions: Positions[],
  needed: bigint,
  taxRate: number,
  stakeTotalSupply: bigint
): bigint[] {
  if (stakeTotalSupply <= 0n) {
    throw new Error('stakeTotalSupply must be greater than zero');
  }

  const candidates: SnatchablePosition[] = positions.map(position => {
    const shareRatio = Number(position.share);
    const scaledShares = BigInt(Math.round(shareRatio * Number(SHARE_SCALE)));

    return {
      id: toBigIntId(position.id),
      owner: position.owner,
      stakeShares: (stakeTotalSupply * scaledShares) / SHARE_SCALE,
      taxRate: Number(position.taxRate),
    };
  });

  const result = selectSnatchPositions(candidates, {
    shortfallShares: needed,
    maxTaxRate: taxRate,
    includeOwned: true,
  });

  if (result.remainingShortfall > 0n) {
    throw new Error('Not enough capacity');
  }

  return result.selected.map(candidate => candidate.id);
}