Track account balances

schema.graphql

@@ -2,8 +2,21 @@
 # Main entities
 #
 
+type Account @entity {
+  id: ID!
+  address: Bytes!
+}
+
+type AccountBalance @entity {
+  id: ID!
+  account: Account!
+  token: Token!
+  amount: BigDecimal!
+  # TODO: timestamp: BigInt
+}
+
 type Token @entity {
-  id: ID! # Same as token address
+  id: ID!
 
   " Token address "
   address: Bytes!

src/mappings/account.ts

@@ -0,0 +1,49 @@
+import { BigDecimal, BigInt, Bytes } from '@graphprotocol/graph-ts'
+
+import { Account, AccountBalance, Token } from '../../generated/schema'
+
+export function getOrCreateAccount(accountAddress: Bytes): Account {
+  let accountId = accountAddress.toHex()
+  let existingEntity = Account.load(accountId)
+
+  if (existingEntity != null) {
+    return existingEntity as Account
+  }
+
+  let newEntity = new Account(accountId)
+  newEntity.address = accountAddress
+
+  return newEntity
+}
+
+export function getOrCreateAccountBalance(account: Account, token: Token): AccountBalance {
+  let balanceId = account.id + '-' + token.id
+  let previousBalance = AccountBalance.load(balanceId)
+
+  if (previousBalance != null) {
+    return previousBalance as AccountBalance
+  }
+
+  let newBalance = new AccountBalance(balanceId)
+  newBalance.account = account.id
+  newBalance.token = token.id
+  newBalance.amount = BigInt.fromI32(0).toBigDecimal()
+
+  return newBalance
+}
+
+export function increaseAccountBalance(account: Account, token: Token, amount: BigDecimal): AccountBalance {
+  let balance = getOrCreateAccountBalance(account, token)
+  balance.amount = balance.amount.plus(amount)
+  // TODO: balance.updated = timestamp
+
+  return balance
+}
+
+export function decreaseAccountBalance(account: Account, token: Token, amount: BigDecimal): AccountBalance {
+  let balance = getOrCreateAccountBalance(account, token)
+  balance.amount = balance.amount.minus(amount)
+  // TODO: balance.updated = timestamp
+
+  return balance
+}

src/mappings/registry.ts

@@ -1,6 +1,6 @@
 import { Address, JSONValue, Value, log, ipfs } from '@graphprotocol/graph-ts'
 
-import * as schema from '../../generated/schema'
+import { Token } from '../../generated/schema'
 import { Unknown } from '../../generated/TokenRegistry/TokenRegistry'
 import { BurnableToken, MintableToken, StandardToken } from '../../generated/TokenRegistry/templates'
 
@@ -27,10 +27,10 @@ export function createToken(value: JSONValue, userData: Value): void {
   if (address) {
     let contractAddress = Address.fromString(address)
 
-    let token = schema.Token.load(contractAddress.toHex())
+    let token = Token.load(contractAddress.toHex())
 
     if (!token) {
-      token = new schema.Token(contractAddress.toHex())
+      token = new Token(contractAddress.toHex())
       token.address = contractAddress
       token.name = name
       token.symbol = symbol

src/mappings/token.ts

@@ -1,101 +1,145 @@
 import { BigInt, BigDecimal, Bytes, EthereumEvent } from '@graphprotocol/graph-ts'
 
-import * as schema from '../../generated/schema'
-
 import { Transfer } from '../../generated/TokenRegistry/templates/StandardToken/ERC20'
 import { Burn } from '../../generated/TokenRegistry/templates/BurnableToken/Burnable'
 import { Mint } from '../../generated/TokenRegistry/templates/MintableToken/Mintable'
 
+import { Account, BurnEvent, MintEvent, Token, TransferEvent } from '../../generated/schema'
+
+import { getOrCreateAccount, decreaseAccountBalance, increaseAccountBalance } from './account'
+
 const GENESIS_ADDRESS = '0x0000000000000000000000000000000000000000'
 
 export function handleTransfer(event: Transfer): void {
-  let token = schema.Token.load(event.address.toHex())
+  let token = Token.load(event.address.toHex())
 
-  if (token) {
+  if (token != null) {
     let amount = toDecimal(event.params.value, token.decimals)
 
-    if (event.params.from.toHex() === GENESIS_ADDRESS) {
-      let entity = mintEvent(event, amount, event.params.to)
-      entity.save()
-    } else if (event.params.to.toHex() === GENESIS_ADDRESS) {
-      let entity = burnEvent(event, amount, event.params.from)
-      entity.save()
-    } else {
-      let entity = transferEvent(event, amount, event.params.from, event.params.to)
-      entity.save()
+    let isBurn = event.params.to.toHex() == GENESIS_ADDRESS
+    let isMint = event.params.from.toHex() == GENESIS_ADDRESS
+    let isTransfer = !(isBurn || isMint)
+
+    // Update token event logs
+    if (isBurn) {
+      let eventEntity = createBurnEvent(event, amount, event.params.from)
+      eventEntity.save()
+    }
+
+    if (isMint) {
+      let eventEntity = createMintEvent(event, amount, event.params.to)
+      eventEntity.save()
+    }
+
+    if (isTransfer) {
+      let eventEntity = createTransferEvent(event, amount, event.params.from, event.params.to)
+      eventEntity.save()
+    }
+
+    // Updates balances of accounts
+    if (isTransfer || isBurn) {
+      let sourceAccount = getOrCreateAccount(event.params.from)
+      let accountBalance = decreaseAccountBalance(sourceAccount, token as Token, amount)
+
+      sourceAccount.save()
+      accountBalance.save()
+    }
+
+    if (isTransfer || isMint) {
+      let destinationAccount = getOrCreateAccount(event.params.to)
+      let accountBalance = increaseAccountBalance(destinationAccount, token as Token, amount)
+
+      destinationAccount.save()
+      accountBalance.save()
     }
   }
 }
 
 export function handleBurn(event: Burn): void {
-  let token = schema.Token.load(event.address.toHex())
+  let token = Token.load(event.address.toHex())
 
-  if (token) {
+  if (token != null) {
     let amount = toDecimal(event.params.value, token.decimals)
-    let entity = burnEvent(event, amount, event.params.burner)
 
-    entity.save()
+    // Persist burn event log
+    let eventEntity = createBurnEvent(event, amount, event.params.burner)
+    eventEntity.save()
+
+    // Update source account balance
+    let account = getOrCreateAccount(event.params.burner)
+    account.save()
+
+    let accountBalance = decreaseAccountBalance(account, token as Token, amount)
+    accountBalance.save()
   }
 }
 
 export function handleMint(event: Mint): void {
-  let token = schema.Token.load(event.address.toHex())
+  let token = Token.load(event.address.toHex())
 
-  if (token) {
+  if (token != null) {
     let amount = toDecimal(event.params.amount, token.decimals)
-    let entity = mintEvent(event, amount, event.params.to)
 
-    entity.save()
+    // Persist mint event log
+    let eventEntity = createMintEvent(event, amount, event.params.to)
+    eventEntity.save()
+
+    // Update destination account balance
+    let account: Account = getOrCreateAccount(event.params.to)
+    account.save()
+
+    let accountBalance = increaseAccountBalance(account, token as Token, amount)
+    accountBalance.save()
   }
 }
 
-function burnEvent(event: EthereumEvent, amount: BigDecimal, burner: Bytes): schema.BurnEvent {
-  let entity = new schema.BurnEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
-  entity.token = event.address.toHex()
-  entity.amount = amount
-  entity.sender = event.transaction.from
-  entity.burner = burner
+function createBurnEvent(event: EthereumEvent, amount: BigDecimal, burner: Bytes): BurnEvent {
+  let eventEntity = new BurnEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
+  eventEntity.token = event.address.toHex()
+  eventEntity.amount = amount
+  eventEntity.sender = event.transaction.from
+  eventEntity.burner = burner
 
-  entity.block = event.block.number
-  entity.transaction = event.transaction.hash
-  entity.timestamp = event.block.timestamp
+  eventEntity.block = event.block.number
+  eventEntity.transaction = event.transaction.hash
+  eventEntity.timestamp = event.block.timestamp
 
-  return entity
+  return eventEntity
 }
 
-function mintEvent(event: EthereumEvent, amount: BigDecimal, destination: Bytes): schema.MintEvent {
-  let entity = new schema.MintEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
-  entity.token = event.address.toHex()
-  entity.amount = amount
-  entity.sender = event.transaction.from
-  entity.destination = destination
-  entity.minter = event.transaction.from
+function createMintEvent(event: EthereumEvent, amount: BigDecimal, destination: Bytes): MintEvent {
+  let eventEntity = new MintEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
+  eventEntity.token = event.address.toHex()
+  eventEntity.amount = amount
+  eventEntity.sender = event.transaction.from
+  eventEntity.destination = destination
+  eventEntity.minter = event.transaction.from
 
-  entity.block = event.block.number
-  entity.transaction = event.transaction.hash
-  entity.timestamp = event.block.timestamp
+  eventEntity.block = event.block.number
+  eventEntity.transaction = event.transaction.hash
+  eventEntity.timestamp = event.block.timestamp
 
-  return entity
+  return eventEntity
 }
 
-function transferEvent(
+function createTransferEvent(
   event: EthereumEvent,
   amount: BigDecimal,
   source: Bytes,
   destination: Bytes
-): schema.TransferEvent {
-  let entity = new schema.TransferEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
-  entity.token = event.address.toHex()
-  entity.amount = amount
-  entity.sender = source
-  entity.source = source
-  entity.destination = destination
+): TransferEvent {
+  let eventEntity = new TransferEvent(event.transaction.hash.toHex() + '-' + event.logIndex.toString())
+  eventEntity.token = event.address.toHex()
+  eventEntity.amount = amount
+  eventEntity.sender = source
+  eventEntity.source = source
+  eventEntity.destination = destination
 
-  entity.block = event.block.number
-  entity.transaction = event.transaction.hash
-  entity.timestamp = event.block.timestamp
+  eventEntity.block = event.block.number
+  eventEntity.transaction = event.transaction.hash
+  eventEntity.timestamp = event.block.timestamp
 
-  return entity
+  return eventEntity
 }
 
 function toDecimal(value: BigInt, decimals: u32): BigDecimal {