add voting and total power at height response queries

src/formulas/formulas/contract/daoCore/base.ts

@@ -1,22 +1,36 @@
 import { ContractFormula } from '@/types'
 
-import { ContractInfo, Expiration, ProposalModule } from '../../types'
+import {
+  ContractInfo,
+  Expiration,
+  ProposalModule,
+  TotalPowerAtHeight,
+  VotingPowerAtHeight,
+} from '../../types'
 import { isExpirationExpired } from '../../utils'
 import { info } from '../common'
 import { balance } from '../external/cw20'
 import { dao as daoPreProposeBaseDao } from '../prePropose/daoPreProposeBase'
 import {
-  totalPower as daoVotingCw20StakedTotalPower,
-  votingPower as daoVotingCw20StakedVotingPower,
+  totalPowerAtHeight as daoVotingCw20StakedTotalPowerAtHeight,
+  votingPowerAtHeight as daoVotingCw20StakedVotingPowerAtHeight,
 } from '../voting/daoVotingCw20Staked'
 import {
-  totalPower as daoVotingCw4TotalPower,
-  votingPower as daoVotingCw4VotingPower,
+  totalPowerAtHeight as daoVotingCw4TotalPowerAtHeight,
+  votingPowerAtHeight as daoVotingCw4VotingPowerAtHeight,
 } from '../voting/daoVotingCw4'
 import {
-  totalPower as daoVotingCw721StakedTotalPower,
-  votingPower as daoVotingCw721StakedVotingPower,
+  totalPowerAtHeight as daoVotingCw721StakedTotalPowerAtHeight,
+  votingPowerAtHeight as daoVotingCw721StakedVotingPowerAtHeight,
 } from '../voting/daoVotingCw721Staked'
+import {
+  totalPowerAtHeight as daoVotingNativeStakedTotalPowerAtHeight,
+  votingPowerAtHeight as daoVotingNativeStakedVotingPowerAtHeight,
+} from '../voting/daoVotingNativeStaked'
+import {
+  totalPowerAtHeight as daoVotingTokenStakedTotalPowerAtHeight,
+  votingPowerAtHeight as daoVotingTokenStakedVotingPowerAtHeight,
+} from '../voting/daoVotingTokenStaked'
 
 export type Config = {
   automatically_add_cw20s: boolean
@@ -288,8 +302,19 @@ export const daoUri: ContractFormula<string> = {
   compute: async (env) => (await config.compute(env))?.dao_uri ?? '',
 }
 
-export const votingPower: ContractFormula<
-  string | undefined,
+const VOTING_POWER_AT_HEIGHT_FORMULAS: ContractFormula<
+  VotingPowerAtHeight | undefined,
+  { address: string }
+>[] = [
+  daoVotingCw4VotingPowerAtHeight,
+  daoVotingCw20StakedVotingPowerAtHeight,
+  daoVotingCw721StakedVotingPowerAtHeight,
+  daoVotingNativeStakedVotingPowerAtHeight,
+  daoVotingTokenStakedVotingPowerAtHeight,
+]
+
+export const votingPowerAtHeight: ContractFormula<
+  VotingPowerAtHeight | undefined,
   { address: string }
 > = {
   compute: async (env) => {
@@ -305,17 +330,34 @@ export const votingPower: ContractFormula<
     }
 
     // Find formula matching code ID key.
-    const votingPowerFormula = VOTING_POWER_FORMULAS.find((formula) =>
-      formula.filter?.codeIdsKeys?.includes(codeIdKey)
+    const votingPowerAtHeightFormula = VOTING_POWER_AT_HEIGHT_FORMULAS.find(
+      (formula) => formula.filter?.codeIdsKeys?.includes(codeIdKey)
     )
-    return await votingPowerFormula?.compute({
+    return await votingPowerAtHeightFormula?.compute({
       ...env,
       contractAddress: votingModuleAddress,
     })
   },
 }
 
-export const totalPower: ContractFormula<string | undefined> = {
+export const votingPower: ContractFormula<
+  string | undefined,
+  { address: string }
+> = {
+  compute: async (env) => (await votingPowerAtHeight.compute(env))?.power,
+}
+
+const TOTAL_POWER_AT_HEIGHT_FORMULAS: ContractFormula<TotalPowerAtHeight>[] = [
+  daoVotingCw4TotalPowerAtHeight,
+  daoVotingCw20StakedTotalPowerAtHeight,
+  daoVotingCw721StakedTotalPowerAtHeight,
+  daoVotingNativeStakedTotalPowerAtHeight,
+  daoVotingTokenStakedTotalPowerAtHeight,
+]
+
+export const totalPowerAtHeight: ContractFormula<
+  TotalPowerAtHeight | undefined
+> = {
   compute: async (env) => {
     const votingModuleAddress = (await votingModule.compute(env)) ?? ''
     if (!votingModuleAddress) {
@@ -329,30 +371,22 @@ export const totalPower: ContractFormula<string | undefined> = {
     }
 
     // Find formula matching code ID key.
-    const totalPowerFormula = TOTAL_POWER_FORMULAS.find((formula) =>
-      formula.filter?.codeIdsKeys?.includes(codeIdKey)
+    const totalPowerAtHeightFormula = TOTAL_POWER_AT_HEIGHT_FORMULAS.find(
+      (formula) => formula.filter?.codeIdsKeys?.includes(codeIdKey)
     )
-    return await totalPowerFormula?.compute({
+    return await totalPowerAtHeightFormula?.compute({
       ...env,
       contractAddress: votingModuleAddress,
     })
   },
 }
 
-// These formulas must have code ID filters set so we can match them.
-const VOTING_POWER_FORMULAS: ContractFormula<
+export const totalPower: ContractFormula<
   string | undefined,
   { address: string }
->[] = [
-  daoVotingCw4VotingPower,
-  daoVotingCw20StakedVotingPower,
-  daoVotingCw721StakedVotingPower,
-]
-const TOTAL_POWER_FORMULAS: ContractFormula<string>[] = [
-  daoVotingCw4TotalPower,
-  daoVotingCw20StakedTotalPower,
-  daoVotingCw721StakedTotalPower,
-]
+> = {
+  compute: async (env) => (await totalPowerAtHeight.compute(env))?.power,
+}
 
 // Returns contracts with an admin state key set to this DAO. Hopefully these
 // are mostly DAO contracts.

src/formulas/formulas/contract/staking/cw20Stake.ts

@@ -7,13 +7,23 @@ export type StakerBalance = {
   balance: string
 }
 
+export type StakedBalanceAtHeight = {
+  balance: string
+  height: number
+}
+
+export type TotalStakedAtHeight = {
+  total: string
+  height: number
+}
+
 export const config: ContractFormula<any | undefined> = {
   compute: async ({ contractAddress, get }) =>
     await get(contractAddress, 'config'),
 }
 
-export const stakedBalance: ContractFormula<
-  string,
+export const stakedBalanceAtHeight: ContractFormula<
+  StakedBalanceAtHeight,
   {
     address: string
     // Required when querying oraichain-cw20-staking contract directly.
@@ -31,6 +41,7 @@ export const stakedBalance: ContractFormula<
     const {
       getTransformationMatch,
       args: { address },
+      block,
     } = env
     if (!address) {
       throw new Error('missing `address`')
@@ -64,24 +75,47 @@ export const stakedBalance: ContractFormula<
 
     keys.push(address)
 
-    return (
+    const balance =
       (
         await getTransformationMatch<string | undefined>(
           contractAddress,
           keys.join(':')
         )
       )?.value ?? '0'
-    )
+
+    return {
+      balance,
+      height: Number(block.height),
+    }
   },
 }
 
-export const totalStaked: ContractFormula<
+export const stakedBalance: ContractFormula<
   string,
   {
+    address: string
     // Required when querying oraichain-cw20-staking contract directly.
     oraichainStakingToken?: string
   }
 > = {
+  filter: stakedBalanceAtHeight.filter,
+  compute: async (env) => (await stakedBalanceAtHeight.compute(env)).balance,
+}
+
+export const totalStakedAtHeight: ContractFormula<
+  TotalStakedAtHeight,
+  {
+    // Required when querying oraichain-cw20-staking contract directly.
+    oraichainStakingToken?: string
+  }
+> = {
+  filter: {
+    codeIdsKeys: [
+      'cw20-stake',
+      'oraichain-cw20-staking',
+      'oraichain-cw20-staking-proxy-snapshot',
+    ],
+  },
   compute: async (env) => {
     let contractAddress = env.contractAddress
     const keys: (string | Uint8Array)[] = ['total_staked']
@@ -109,10 +143,27 @@ export const totalStaked: ContractFormula<
       keys.push(fromBech32(env.args.oraichainStakingToken).data)
     }
 
-    return (await env.get<string | undefined>(contractAddress, ...keys)) || '0'
+    const total =
+      (await env.get<string | undefined>(contractAddress, ...keys)) || '0'
+
+    return {
+      total,
+      height: Number(env.block.height),
+    }
   },
 }
 
+export const totalStaked: ContractFormula<
+  string,
+  {
+    // Required when querying oraichain-cw20-staking contract directly.
+    oraichainStakingToken?: string
+  }
+> = {
+  filter: totalStakedAtHeight.filter,
+  compute: async (env) => (await totalStakedAtHeight.compute(env)).total,
+}
+
 export const stakedValue: ContractFormula<string, { address: string }> = {
   filter: {
     codeIdsKeys: ['cw20-stake'],

src/formulas/formulas/contract/voting/daoVotingCw20Staked.ts

@@ -1,12 +1,15 @@
 import { ContractFormula } from '@/types'
 
+import { TotalPowerAtHeight, VotingPowerAtHeight } from '../../types'
 import {
   StakerBalance,
   topStakers as cw20StakeTopStakers,
   stakedBalance,
   totalStaked,
 } from '../staking/cw20Stake'
 
+const CODE_IDS_KEYS = ['dao-voting-cw20-staked']
+
 export { activeThreshold } from './common'
 
 export const tokenContract: ContractFormula<string | undefined> = {
@@ -24,10 +27,15 @@ export const stakingContract: ContractFormula<string | undefined> = {
     (await get<string>(contractAddress, 'staking_contract')),
 }
 
-export const votingPower: ContractFormula<
-  string | undefined,
+export const votingPowerAtHeight: ContractFormula<
+  VotingPowerAtHeight | undefined,
   { address: string }
 > = {
+  // Filter by code ID since someone may modify the contract. This is also used
+  // in DAO core to match the voting module and pass the query through.
+  filter: {
+    codeIdsKeys: CODE_IDS_KEYS,
+  },
   compute: async (env) => {
     if (!env.args.address) {
       throw new Error('missing `address`')
@@ -48,27 +56,53 @@ export const votingPower: ContractFormula<
       return
     }
 
-    const power = await stakedBalance.compute({
-      ...env,
-      contractAddress: stakingContractAddress,
-    })
+    const power =
+      (await stakedBalance.compute({
+        ...env,
+        contractAddress: stakingContractAddress,
+      })) || '0'
 
-    return power || '0'
+    return {
+      power,
+      height: Number(env.block.height),
+    }
   },
 }
 
-export const totalPower: ContractFormula<string> = {
+export const votingPower: ContractFormula<
+  string | undefined,
+  { address: string }
+> = {
+  filter: votingPowerAtHeight.filter,
+  compute: async (env) => (await votingPowerAtHeight.compute(env))?.power,
+}
+
+export const totalPowerAtHeight: ContractFormula<TotalPowerAtHeight> = {
+  // Filter by code ID since someone may modify the contract. This is also used
+  // in DAO core to match the voting module and pass the query through.
+  filter: {
+    codeIdsKeys: CODE_IDS_KEYS,
+  },
   compute: async (env) => {
     const stakingContractAddress = (await stakingContract.compute(env)) ?? ''
-    const power = await totalStaked.compute({
-      ...env,
-      contractAddress: stakingContractAddress,
-    })
-
-    return power || '0'
+    const power =
+      (await totalStaked.compute({
+        ...env,
+        contractAddress: stakingContractAddress,
+      })) || '0'
+
+    return {
+      power,
+      height: Number(env.block.height),
+    }
   },
 }
 
+export const totalPower: ContractFormula<string> = {
+  filter: totalPowerAtHeight.filter,
+  compute: async (env) => (await totalPowerAtHeight.compute(env)).power,
+}
+
 export const dao: ContractFormula<string | undefined> = {
   compute: async ({ contractAddress, get, getTransformationMatch }) =>
     (await getTransformationMatch<string>(contractAddress, 'dao'))?.value ??