Refactor share redemption code into a few sensible functions

functions/src/redeem-shares.ts

@@ -4,53 +4,36 @@ 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'
 
-export const redeemShares = async (userId: string, contractId: string) => {
+type CandidateBet<T extends Bet> = Omit<T, 'id' | 'userId'>
-  return await firestore.runTransaction(async (transaction) => {
+type RedeemableBet = Pick<Bet, 'outcome' | 'shares' | 'loanAmount'>
-    const contractDoc = firestore.doc(`contracts/${contractId}`)
-    const contractSnap = await transaction.get(contractDoc)
-    if (!contractSnap.exists)
-      return { status: 'error', message: 'Invalid contract' }
 
-    const contract = contractSnap.data() as Contract
+const getRedeemableAmount = (bets: RedeemableBet[]) => {
-    const { mechanism, outcomeType } = contract
-    if (
-      !(outcomeType === 'BINARY' || outcomeType === 'PSEUDO_NUMERIC') ||
-      mechanism !== 'cpmm-1'
-    )
-      return { status: 'success' }
-
-    const betsSnap = await transaction.get(
-      firestore
-        .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 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 amount = Math.min(yesShares, noShares)
+const getRedemptionBets = (
-    if (amount <= 0) return
+  shares: number,
-
+  loanPayment: number,
-    const prevLoanAmount = sumBy(bets, (bet) => bet.loanAmount ?? 0)
+  contract: CPMMContract
-    const loanPaid = Math.min(prevLoanAmount, amount)
+) => {
-    const netAmount = amount - loanPaid
-
   const p = getProbability(contract)
   const createdTime = Date.now()
-
+  const yesBet: CandidateBet<Bet> = {
-    const yesDoc = firestore.collection(`contracts/${contract.id}/bets`).doc()
-    const yesBet: Bet = {
-      id: yesDoc.id,
-      userId: userId,
     contractId: contract.id,
-      amount: p * -amount,
+    amount: p * -shares,
-      shares: -amount,
+    shares: -shares,
-      loanAmount: loanPaid ? -loanPaid / 2 : 0,
+    loanAmount: loanPayment ? -loanPayment / 2 : 0,
     outcome: 'YES',
     probBefore: p,
     probAfter: p,
@@ -58,15 +41,11 @@ export const redeemShares = async (userId: string, contractId: string) => {
     isRedemption: true,
     fees: noFees,
   }
-
+  const noBet: CandidateBet<Bet> = {
-    const noDoc = firestore.collection(`contracts/${contract.id}/bets`).doc()
-    const noBet: Bet = {
-      id: noDoc.id,
-      userId: userId,
     contractId: contract.id,
-      amount: (1 - p) * -amount,
+    amount: (1 - p) * -shares,
-      shares: -amount,
+    shares: -shares,
-      loanAmount: loanPaid ? -loanPaid / 2 : 0,
+    loanAmount: loanPayment ? -loanPayment / 2 : 0,
     outcome: 'NO',
     probBefore: p,
     probAfter: p,
@@ -74,23 +53,41 @@ export const redeemShares = async (userId: string, contractId: string) => {
     isRedemption: true,
     fees: noFees,
   }
+  return [yesBet, noBet]
+}
+
+export const redeemShares = async (userId: string, contractId: string) => {
+  return await firestore.runTransaction(async (trans) => {
+    const contractDoc = firestore.doc(`contracts/${contractId}`)
+    const contractSnap = await trans.get(contractDoc)
+    if (!contractSnap.exists)
+      return { status: 'error', message: 'Invalid contract' }
+
+    const contract = contractSnap.data() as Contract
+    const { mechanism } = contract
+    if (mechanism !== 'cpmm-1') return { status: 'success' }
+
+    const betsColl = firestore.collection(`contracts/${contract.id}/bets`)
+    const betsSnap = await trans.get(betsColl.where('userId', '==', userId))
+    const bets = betsSnap.docs.map((doc) => doc.data() as Bet)
+    const { shares, loanPayment, netAmount } = getRedeemableAmount(bets)
+    const [yesBet, noBet] = getRedemptionBets(shares, loanPayment, contract)
 
     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' }
   })