Implementation of an ERC20 token with snapshots

contracts/math/Math.sol

@@ -21,4 +21,9 @@ library Math {
   function min256(uint256 a, uint256 b) internal pure returns (uint256) {
     return a < b ? a : b;
   }
+
+  function average(uint256 a, uint256 b) internal pure returns (uint256) {
+    // (a + b) / 2 can overflow, so we distribute
+    return (a / 2) + (b / 2) + ((a % 2 + b % 2) / 2);
+  }
 }

contracts/mocks/SnapshotTokenMock.sol

@@ -0,0 +1,13 @@
+pragma solidity ^0.4.23;
+
+import { SnapshotToken } from "../token/ERC20/SnapshotToken.sol";
+
+
+contract SnapshotTokenMock is SnapshotToken {
+  constructor() public {
+    uint256 initial_balance = 10000;
+    balances[msg.sender] = initial_balance;
+    totalSupply_ = initial_balance;
+    emit Transfer(address(0), msg.sender, initial_balance);
+  }
+}

contracts/token/ERC20/SnapshotToken.sol

@@ -0,0 +1,70 @@
+pragma solidity ^0.4.23;
+
+import { StandardToken } from "./StandardToken.sol";
+import { ArrayUtils } from "../../utils/ArrayUtils.sol";
+
+
+/**
+ * @title SnapshotToken
+ *
+ * @dev An ERC20 token which enables taking snapshots of accounts' balances.
+ * @dev This can be useful to safely implement voting weighed by balance.
+ */
+contract SnapshotToken is StandardToken {
+  using ArrayUtils for uint256[];
+
+  // The 0 id represents no snapshot was taken yet.
+  uint256 private currSnapshotId;
+
+  mapping (address => uint256[]) private snapshotIds;
+  mapping (address => uint256[]) private snapshotBalances;
+
+  event Snapshot(uint256 id);
+
+  function transfer(address _to, uint256 _value) public returns (bool) {
+    _updateSnapshot(msg.sender);
+    _updateSnapshot(_to);
+    return super.transfer(_to, _value);
+  }
+
+  function transferFrom(address _from, address _to, uint256 _value) public returns (bool) {
+    _updateSnapshot(_from);
+    _updateSnapshot(_to);
+    return super.transferFrom(_from, _to, _value);
+  }
+
+  function snapshot() public returns (uint256) {
+    currSnapshotId += 1;
+    emit Snapshot(currSnapshotId);
+    return currSnapshotId;
+  }
+
+  function balanceOfAt(address _account, uint256 _snapshotId) public view returns (uint256) {
+    require(_snapshotId > 0 && _snapshotId <= currSnapshotId);
+
+    uint256 idx = snapshotIds[_account].findUpperBound(_snapshotId);
+
+    if (idx == snapshotIds[_account].length) {
+      return balanceOf(_account);
+    } else {
+      return snapshotBalances[_account][idx];
+    }
+  }
+
+  function _updateSnapshot(address _account) internal {
+    if (_lastSnapshotId(_account) < currSnapshotId) {
+      snapshotIds[_account].push(currSnapshotId);
+      snapshotBalances[_account].push(balanceOf(_account));
+    }
+  }
+
+  function _lastSnapshotId(address _account) internal returns (uint256) {
+    uint256[] storage snapshots = snapshotIds[_account];
+
+    if (snapshots.length == 0) {
+      return 0;
+    } else {
+      return snapshots[snapshots.length - 1];
+    }
+  }
+}

contracts/utils/ArrayUtils.sol

@@ -0,0 +1,29 @@
+pragma solidity ^0.4.23;
+
+import { Math } from "../math/Math.sol";
+
+
+library ArrayUtils {
+  function findUpperBound(uint256[] storage _array, uint256 _element) internal view returns (uint256) {
+    uint256 low = 0;
+    uint256 high = _array.length;
+
+    while (low < high) {
+      uint256 mid = Math.average(low, high);
+
+      if (_array[mid] > _element) {
+        high = mid;
+      } else {
+        low = mid + 1;
+      }
+    }
+
+    // At this point at `low` is the exclusive upper bound. We will return the inclusive upper bound.
+
+    if (low > 0 && _array[low - 1] == _element) {
+      return low - 1;
+    } else {
+      return low;
+    }
+  }
+}

test/token/ERC20/SnapshotToken.test.js

@@ -0,0 +1,38 @@
+const SnapshotToken = artifacts.require('SnapshotTokenMock');
+
+contract('SnapshotToken', function ([_, account1, account2]) {
+  beforeEach(async function () {
+    this.token = await SnapshotToken.new({ from: account1 });
+  });
+
+  it('can snapshot!', async function () {
+    assert.equal(await this.token.balanceOf(account1), 10000);
+    assert.equal(await this.token.balanceOf(account2), 0);
+
+    const { logs: logs1 } = await this.token.snapshot();
+    const snapshotId1 = logs1[0].args.id;
+
+    assert.equal(await this.token.balanceOf(account1), 10000);
+    assert.equal(await this.token.balanceOf(account2), 0);
+
+    assert.equal(await this.token.balanceOfAt(account1, snapshotId1), 10000);
+    assert.equal(await this.token.balanceOfAt(account2, snapshotId1), 0);
+
+    await this.token.transfer(account2, 500, { from: account1 });
+
+    assert.equal(await this.token.balanceOf(account1), 9500);
+    assert.equal(await this.token.balanceOf(account2), 500);
+
+    assert.equal(await this.token.balanceOfAt(account1, snapshotId1), 10000);
+    assert.equal(await this.token.balanceOfAt(account2, snapshotId1), 0);
+
+    const { logs: logs2 } = await this.token.snapshot();
+    const snapshotId2 = logs2[0].args.id;
+
+    assert.equal(await this.token.balanceOf(account1), 9500);
+    assert.equal(await this.token.balanceOf(account2), 500);
+
+    assert.equal(await this.token.balanceOfAt(account1, snapshotId2), 9500);
+    assert.equal(await this.token.balanceOfAt(account2, snapshotId2), 500);
+  });
+});