Refactor share redemption code into a few sensible functions

functions/src/redeem-shares.ts

@@ -4,93 +4,90 @@ import { partition, sumBy } from 'lodash'
 import { Bet } from '../../common/bet'
 import { getProbability } from '../../common/calculate'
 
-import { Contract } from '../../common/contract'
+import { Contract, CPMMContract } from '../../common/contract'
 import { noFees } from '../../common/fees'
 import { User } from '../../common/user'
 
+type CandidateBet<T extends Bet> = Omit<T, 'id' | 'userId'>
+type RedeemableBet = Pick<Bet, 'outcome' | 'shares' | 'loanAmount'>
+
+const getRedeemableAmount = (bets: RedeemableBet[]) => {
+  const [yesBets, noBets] = partition(bets, (b) => b.outcome === 'YES')
+  const yesShares = sumBy(yesBets, (b) => b.shares)
+  const noShares = sumBy(noBets, (b) => b.shares)
+  const shares = Math.max(Math.min(yesShares, noShares), 0)
+  const loanAmount = sumBy(bets, (bet) => bet.loanAmount ?? 0)
+  const loanPayment = Math.min(loanAmount, shares)
+  const netAmount = shares - loanPayment
+  return { shares, loanPayment, netAmount }
+}
+
+const getRedemptionBets = (
+  shares: number,
+  loanPayment: number,
+  contract: CPMMContract
+) => {
+  const p = getProbability(contract)
+  const createdTime = Date.now()
+  const yesBet: CandidateBet<Bet> = {
+    contractId: contract.id,
+    amount: p * -shares,
+    shares: -shares,
+    loanAmount: loanPayment ? -loanPayment / 2 : 0,
+    outcome: 'YES',
+    probBefore: p,
+    probAfter: p,
+    createdTime,
+    isRedemption: true,
+    fees: noFees,
+  }
+  const noBet: CandidateBet<Bet> = {
+    contractId: contract.id,
+    amount: (1 - p) * -shares,
+    shares: -shares,
+    loanAmount: loanPayment ? -loanPayment / 2 : 0,
+    outcome: 'NO',
+    probBefore: p,
+    probAfter: p,
+    createdTime,
+    isRedemption: true,
+    fees: noFees,
+  }
+  return [yesBet, noBet]
+}
+
 export const redeemShares = async (userId: string, contractId: string) => {
-  return await firestore.runTransaction(async (transaction) => {
+  return await firestore.runTransaction(async (trans) => {
     const contractDoc = firestore.doc(`contracts/${contractId}`)
-    const contractSnap = await transaction.get(contractDoc)
+    const contractSnap = await trans.get(contractDoc)
     if (!contractSnap.exists)
       return { status: 'error', message: 'Invalid contract' }
 
     const contract = contractSnap.data() as Contract
-    const { mechanism, outcomeType } = contract
+    const { mechanism } = contract
-    if (
+    if (mechanism !== 'cpmm-1') return { status: 'success' }
-      !(outcomeType === 'BINARY' || outcomeType === 'PSEUDO_NUMERIC') ||
-      mechanism !== 'cpmm-1'
-    )
-      return { status: 'success' }
 
-    const betsSnap = await transaction.get(
+    const betsColl = firestore.collection(`contracts/${contract.id}/bets`)
-      firestore
+    const betsSnap = await trans.get(betsColl.where('userId', '==', userId))
-        .collection(`contracts/${contract.id}/bets`)
-        .where('userId', '==', userId)
-    )
     const bets = betsSnap.docs.map((doc) => doc.data() as Bet)
-    const [yesBets, noBets] = partition(bets, (b) => b.outcome === 'YES')
+    const { shares, loanPayment, netAmount } = getRedeemableAmount(bets)
-    const yesShares = sumBy(yesBets, (b) => b.shares)
+    const [yesBet, noBet] = getRedemptionBets(shares, loanPayment, contract)
-    const noShares = sumBy(noBets, (b) => b.shares)
-
-    const amount = Math.min(yesShares, noShares)
-    if (amount <= 0) return
-
-    const prevLoanAmount = sumBy(bets, (bet) => bet.loanAmount ?? 0)
-    const loanPaid = Math.min(prevLoanAmount, amount)
-    const netAmount = amount - loanPaid
-
-    const p = getProbability(contract)
-    const createdTime = Date.now()
-
-    const yesDoc = firestore.collection(`contracts/${contract.id}/bets`).doc()
-    const yesBet: Bet = {
-      id: yesDoc.id,
-      userId: userId,
-      contractId: contract.id,
-      amount: p * -amount,
-      shares: -amount,
-      loanAmount: loanPaid ? -loanPaid / 2 : 0,
-      outcome: 'YES',
-      probBefore: p,
-      probAfter: p,
-      createdTime,
-      isRedemption: true,
-      fees: noFees,
-    }
-
-    const noDoc = firestore.collection(`contracts/${contract.id}/bets`).doc()
-    const noBet: Bet = {
-      id: noDoc.id,
-      userId: userId,
-      contractId: contract.id,
-      amount: (1 - p) * -amount,
-      shares: -amount,
-      loanAmount: loanPaid ? -loanPaid / 2 : 0,
-      outcome: 'NO',
-      probBefore: p,
-      probAfter: p,
-      createdTime,
-      isRedemption: true,
-      fees: noFees,
-    }
 
     const userDoc = firestore.doc(`users/${userId}`)
-    const userSnap = await transaction.get(userDoc)
+    const userSnap = await trans.get(userDoc)
     if (!userSnap.exists) return { status: 'error', message: 'User not found' }
-
     const user = userSnap.data() as User
-
     const newBalance = user.balance + netAmount
 
     if (!isFinite(newBalance)) {
       throw new Error('Invalid user balance for ' + user.username)
     }
 
-    transaction.update(userDoc, { balance: newBalance })
+    const yesDoc = betsColl.doc()
-
+    const noDoc = betsColl.doc()
-    transaction.create(yesDoc, yesBet)
+    trans.update(userDoc, { balance: newBalance })
-    transaction.create(noDoc, noBet)
+    trans.create(yesDoc, { id: yesDoc.id, userId, ...yesBet })
+    trans.create(noDoc, { id: noDoc.id, userId, ...noBet })
 
     return { status: 'success' }
   })