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);
}