From 43b404e7ed356515863b9acc51c6a08d2543d972 Mon Sep 17 00:00:00 2001
From: James Grugett <jahooma@gmail.com>
Date: Mon, 28 Mar 2022 21:39:40 -0500
Subject: [PATCH] Compute sale value of shares with binary search to keep k
 constant.

---
 common/calculate-cpmm.ts     | 58 ++++++++++++++++++++++++++++--------
 common/sell-bet.ts           |  3 +-
 functions/src/sell-shares.ts | 11 ++++++-
 3 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/common/calculate-cpmm.ts b/common/calculate-cpmm.ts
index 1d698872..44e195a2 100644
--- a/common/calculate-cpmm.ts
+++ b/common/calculate-cpmm.ts
@@ -114,19 +114,47 @@ export function calculateCpmmPurchase(
   return { shares, newPool, newP, fees }
 }
 
-export function calculateCpmmShareValue(
+function computeK(y: number, n: number, p: number) {
+  return y ** p * n ** (1 - p)
+}
+
+function sellSharesK(
+  y: number,
+  n: number,
+  p: number,
+  s: number,
+  outcome: 'YES' | 'NO',
+  b: number
+) {
+  return outcome === 'YES'
+    ? computeK(y - b + s, n - b, p)
+    : computeK(y - b, n - b + s, p)
+}
+
+function calculateCpmmShareValue(
   contract: FullContract<CPMM, Binary>,
   shares: number,
-  outcome: string
+  outcome: 'YES' | 'NO'
 ) {
-  const { pool } = contract
-  const { YES: y, NO: n } = pool
+  const { pool, p } = contract
 
-  // TODO: calculate using new function
-  const poolChange = outcome === 'YES' ? shares + y - n : shares + n - y
-  const k = y * n
-  const shareValue = 0.5 * (shares + y + n - Math.sqrt(4 * k + poolChange ** 2))
-  return shareValue
+  const k = computeK(pool.YES, pool.NO, p)
+
+  // Find bet amount that preserves k after selling shares.
+  let lowAmount = 0
+  let highAmount = shares
+  let mid = 0
+  let kGuess = 0
+  while (Math.abs(k - kGuess) > 0.00000000001) {
+    mid = lowAmount + (highAmount - lowAmount) / 2
+    kGuess = sellSharesK(pool.YES, pool.NO, p, shares, outcome, mid)
+    if (kGuess < k) {
+      highAmount = mid
+    } else {
+      lowAmount = mid
+    }
+  }
+  return mid
 }
 
 export function calculateCpmmSale(
@@ -135,7 +163,11 @@ export function calculateCpmmSale(
 ) {
   const { shares, outcome } = bet
 
-  const rawSaleValue = calculateCpmmShareValue(contract, shares, outcome)
+  const rawSaleValue = calculateCpmmShareValue(
+    contract,
+    Math.abs(shares),
+    outcome as 'YES' | 'NO'
+  )
 
   const { fees, remainingBet: saleValue } = getCpmmLiquidityFee(
     contract,
@@ -153,9 +185,11 @@ export function calculateCpmmSale(
       ? [y + shares - saleValue + fee, n - saleValue + fee]
       : [y - saleValue + fee, n + shares - saleValue + fee]
 
-  const newPool = { YES: newY, NO: newN }
+  const postBetPool = { YES: newY, NO: newN }
 
-  return { saleValue, newPool, fees }
+  const { newPool, newP } = addCpmmLiquidity(postBetPool, contract.p, fee)
+
+  return { saleValue, newPool, newP, fees }
 }
 
 export function getCpmmProbabilityAfterSale(
diff --git a/common/sell-bet.ts b/common/sell-bet.ts
index 31c18d7f..750cfb39 100644
--- a/common/sell-bet.ts
+++ b/common/sell-bet.ts
@@ -92,7 +92,7 @@ export const getCpmmSellBetInfo = (
 ) => {
   const { pool, p } = contract
 
-  const { saleValue, newPool, fees } = calculateCpmmSale(contract, {
+  const { saleValue, newPool, newP, fees } = calculateCpmmSale(contract, {
     shares,
     outcome,
   })
@@ -128,6 +128,7 @@ export const getCpmmSellBetInfo = (
   return {
     newBet,
     newPool,
+    newP,
     newBalance,
     fees,
   }
diff --git a/functions/src/sell-shares.ts b/functions/src/sell-shares.ts
index 9ad4eed7..91d88fef 100644
--- a/functions/src/sell-shares.ts
+++ b/functions/src/sell-shares.ts
@@ -48,7 +48,7 @@ export const sellShares = functions.runWith({ minInstances: 1 }).https.onCall(
         .collection(`contracts/${contractId}/bets`)
         .doc()
 
-      const { newBet, newPool, newBalance, fees } = getCpmmSellBetInfo(
+      const { newBet, newPool, newP, newBalance, fees } = getCpmmSellBetInfo(
         user,
         shares,
         outcome,
@@ -56,15 +56,24 @@ export const sellShares = functions.runWith({ minInstances: 1 }).https.onCall(
         newBetDoc.id
       )
 
+      if (!isFinite(newP)) {
+        return {
+          status: 'error',
+          message: 'Trade rejected due to overflow error.',
+        }
+      }
+
       if (!isFinite(newBalance)) {
         throw new Error('Invalid user balance for ' + user.username)
       }
+
       transaction.update(userDoc, { balance: newBalance })
       transaction.create(newBetDoc, newBet)
       transaction.update(
         contractDoc,
         removeUndefinedProps({
           pool: newPool,
+          p: newP,
           collectedFees: addObjects(fees ?? {}, collectedFees ?? {}),
           volume: volume + Math.abs(newBet.amount),
         })