@ocap/statedb-memory 1.20.15 → 1.21.0

package/lib/db.js

@@ -3,6 +3,7 @@ const Table = require('./table/base');
 const Account = require('./table/account');
 const Token = require('./table/token');
 const Rollup = require('./table/rollup');
+const Balance = require('./table/balance');
 
 const { name, version } = require('../package.json');
 
@@ -13,17 +14,23 @@ class MemoryStateDB extends StateDB {
     this.name = name;
     this.version = version;
 
-    this.account = new Account('account', 'address');
-    this.asset = new Table('asset', 'address');
-    this.factory = new Table('factory', 'address');
-    this.delegation = new Table('delegation', 'address');
-    this.tx = new Table('tx', 'hash');
-    this.token = new Token('token', 'address');
-    this.chain = new Table('chain', 'address');
-    this.stake = new Table('stake', 'address');
-    this.rollup = new Rollup('rollup', 'address');
-    this.rollupBlock = new Table('rollupBlock', 'hash');
-    this.evidence = new Table('evidence', 'hash');
+    this.balance = new Balance({ name: 'balance', uniqIndex: ['address', 'tokenAddress'] });
+    this.account = new Account({
+      name: 'account',
+      uniqIndex: 'address',
+      balanceTable: this.balance,
+      syncBalance: true,
+    });
+    this.factory = new Table({ name: 'factory', uniqIndex: 'address', syncBalance: true, balanceTable: this.balance });
+    this.stake = new Table({ name: 'stake', uniqIndex: 'address', syncBalance: true, balanceTable: this.balance });
+    this.asset = new Table({ name: 'asset', uniqIndex: 'address' });
+    this.delegation = new Table({ name: 'delegation', uniqIndex: 'address' });
+    this.tx = new Table({ name: 'tx', uniqIndex: 'hash' });
+    this.token = new Token({ name: 'token', uniqIndex: 'address' });
+    this.chain = new Table({ name: 'chain', uniqIndex: 'address' });
+    this.rollup = new Rollup({ name: 'rollup', uniqIndex: 'address' });
+    this.rollupBlock = new Table({ name: 'rollupBlock', uniqIndex: 'hash' });
+    this.evidence = new Table({ name: 'evidence', uniqIndex: 'hash' });
 
     this.attachReadyListeners();
   }

package/lib/table/account.js

@@ -1,27 +1,19 @@
 const { ensureChecksumAddress } = require('@ocap/state/lib/states/account');
-const debug = require('debug')(require('../../package.json').name);
 const MemoryStateDB = require('./base');
 
 class AccountStateDB extends MemoryStateDB {
-  _get(address, { traceMigration = true } = {}) {
-    const current = this.collection.by(this.uniqIndex, ensureChecksumAddress(address)) || null;
+  async _get(address, { traceMigration = true } = {}) {
+    const current = await super._get(ensureChecksumAddress(address));
     if (current && traceMigration && Array.isArray(current.migratedTo) && current.migratedTo.length) {
       return this._get(current.migratedTo[0]);
     }
 
-    return current || null;
+    return current;
   }
 
-  _create(id, attrs = {}) {
-    if (this._get(id)) {
-      throw new Error(`${this.name} already exists: ${id}`);
-    }
-
-    const newDoc = { [this.uniqIndex]: id, ...attrs };
-    newDoc[this.uniqIndex] = ensureChecksumAddress(newDoc[this.uniqIndex]);
-    this.collection.insert(newDoc);
-    debug(`${this.name} create`, { id, ...attrs });
-    return this._get(id);
+  _create(key, attrs = {}) {
+    const address = ensureChecksumAddress(key);
+    return super._create(address, { ...attrs, [this.uniqIndex]: address });
   }
 }
 

package/lib/table/balance.js

@@ -0,0 +1,42 @@
+const { ensureChecksumAddress } = require('@ocap/state/lib/states/account');
+const MemoryStateDB = require('./base');
+
+class BalanceTable extends MemoryStateDB {
+  async getBalance(address) {
+    const tokens = await this.collection.find({ address: ensureChecksumAddress(address) });
+    return (tokens || []).reduce((acc, token) => {
+      acc[token.tokenAddress] = token.balance;
+      return acc;
+    }, {});
+  }
+
+  async updateBalance({ address, tokens, context = {} }, ctx) {
+    if (!Object.keys(tokens).length) return {};
+
+    const tokenBalances = await this.getBalance(address);
+    const updatedTokens = Object.keys(tokens)
+      .filter(
+        (token) => (tokenBalances[token] && tokenBalances[token] !== '0') || (tokens[token] && tokens[token] !== '0')
+      )
+      .filter((token) => tokenBalances[token] !== tokens[token]);
+
+    // eslint-disable-next-line array-callback-return
+    await Promise.all(
+      updatedTokens.map(async (token) => {
+        const key = { address, tokenAddress: token };
+
+        if (tokenBalances[token]) {
+          await this.update(key, { ...key, balance: tokens[token], context }, ctx);
+        } else {
+          await this.create(key, { ...key, balance: tokens[token], context }, ctx);
+        }
+
+        tokenBalances[token] = tokens[token];
+      })
+    );
+
+    return tokens;
+  }
+}
+
+module.exports = BalanceTable;

package/lib/table/base.js

@@ -8,55 +8,122 @@ const db = new Lokijs('ocap-memory-statedb.db');
 const debug = require('debug')(require('../../package.json').name);
 
 class MemoryTable extends StateDBTable {
-  constructor(name, uniqIndex) {
-    super();
+  /**
+   * @param {Object} param
+   * @param {string} param.name db name
+   * @param {string|string[]} param.uniqIndex primary keys
+   * @param {boolean} param.syncBalance sync balance table when create / update / get
+   * @param {BalanceTable} param.balanceTable balance table
+   */
+  constructor({ name, uniqIndex, balanceTable, syncBalance = false }) {
+    super(uniqIndex);
 
     this.name = name;
     this.uniqIndex = uniqIndex;
+    this.balanceTable = balanceTable;
+    this.syncBalance = syncBalance;
 
-    this.collection = db.addCollection(name, { unique: [uniqIndex], clone: true });
+    if (this.syncBalance && !this.balanceTable) {
+      throw new Error('balanceTable is required when syncBalance is true');
+    }
+
+    this.collection = db.addCollection(name, { unique: [this.primaryKey], clone: true });
     this.markReady();
   }
 
-  _create(id, attrs = {}) {
-    if (this._get(id)) {
-      throw new Error(`${this.name} already exists: ${id}`);
+  async _create(key, attrs) {
+    const id = this.generatePrimaryKey(key);
+    // Don't call this._get here cause _get may be overwritten
+    const doc = await MemoryTable.prototype._get.call(this, id);
+    if (doc) {
+      throw new Error(`${this.name} already exists: ${key}`);
+    }
+
+    debug(`insert ${this.name}`, attrs);
+
+    const insertAttrs = { ...attrs };
+    // insert without tokens cause we've migrate token to balance table
+    if (this.syncBalance) {
+      delete insertAttrs.tokens;
     }
 
-    this.collection.insert({ [this.uniqIndex]: id, ...attrs });
-    debug(`${this.name} create`, { id, ...attrs });
-    return this._get(id);
+    const result = await this.collection.insert({ [this.primaryKey]: id, ...insertAttrs });
+
+    if (this.syncBalance && attrs.tokens) {
+      debug(`update balance for ${this.name}`, attrs);
+
+      result.tokens = await this.balanceTable.updateBalance({
+        address: attrs.address,
+        tokens: attrs.tokens,
+        context: attrs.context,
+      });
+    }
+
+    return result;
   }
 
-  _get(id) {
-    return id ? this.collection.by(this.uniqIndex, id) : null;
+  async _get(key) {
+    if (!key) return null;
+
+    const id = this.generatePrimaryKey(key);
+    const result = await this.collection.by(this.primaryKey, id);
+
+    if (result && this.syncBalance) {
+      const balance = await this.balanceTable.getBalance(result.address);
+      result.tokens = { ...(result.tokens || {}), ...balance };
+    }
+
+    return result;
   }
 
   _history() {
     return [];
   }
 
-  _update(id, updates) {
-    const doc = this.collection.by(this.uniqIndex, id);
+  async _update(key, updates) {
+    const id = this.generatePrimaryKey(key);
+    // Don't call this._get here cause _get may be overwritten
+    const doc = await MemoryTable.prototype._get.call(this, id);
     if (!doc) {
-      throw new Error(`${this.name} does not exists: ${id}`);
+      throw new Error(`${this.name} does not exists: ${key}`);
     }
 
     Object.assign(doc, updates);
-    this.collection.update(doc);
+
+    const updateAttrs = { ...doc };
+    // insert without tokens cause we've migrate token to balance table
+    if (this.syncBalance) {
+      delete updateAttrs.tokens;
+    }
+
+    await this.collection.update(updateAttrs);
+
+    if (this.syncBalance && doc.tokens) {
+      debug(`update balance for ${this.name}`, doc);
+
+      doc.tokens = await this.balanceTable.updateBalance({
+        address: doc.address,
+        tokens: doc.tokens,
+        context: doc.context,
+      });
+    }
+
     return doc;
   }
 
   updateOrCreate(exist, state, ctx) {
-    if (!state[this.uniqIndex]) {
+    const id = this.generatePrimaryKey(state);
+    const attrs = omit(state, this.uniqIndex);
+
+    if (!id) {
       throw new Error('Cannot update or create without uniq index');
     }
 
     if (exist) {
-      return this.update(state[this.uniqIndex], omit(state, [this.uniqIndex]), ctx);
+      return this.update(id, attrs, ctx);
     }
 
-    return this.create(state[this.uniqIndex], omit(state, [this.uniqIndex]), ctx);
+    return this.create(id, attrs, ctx);
   }
 
   _reset() {

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@ocap/statedb-memory",
   "description": "OCAP statedb adapter that uses memory as backend statedb, just for test purpose",
-  "version": "1.20.15",
+  "version": "1.21.0",
   "author": "wangshijun <shijun@arcblock.io> (https://www.arcblock.io)",
   "bugs": {
     "url": "https://github.com/ArcBlock/blockchain/issues",
@@ -36,9 +36,9 @@
     "debug": "^4.3.6",
     "lodash": "^4.17.21",
     "lokijs": "^1.5.12",
-    "@ocap/statedb": "1.20.15",
-    "@ocap/util": "1.20.15",
-    "@ocap/state": "1.20.15"
+    "@ocap/statedb": "1.21.0",
+    "@ocap/util": "1.21.0",
+    "@ocap/state": "1.21.0"
   },
   "scripts": {
     "lint": "eslint tests lib",