Methods visibility corrected

Require reason message removed
Inline documentation in NatSpec added
Snapshot added to _mint and _burn method

contracts/mocks/ERC20SnapshotMock.sol

@@ -9,4 +9,11 @@ contract ERC20SnapshotMock is ERC20Snapshot {
     _mint(initialAccount, initialBalance);
   }
 
+  function mint(address account, uint256 amount) public {
+    _mint(account, amount);
+  }
+
+  function burn(address account, uint256 amount) public {
+    _burn(account, amount);
+  }
 }

contracts/token/ERC20/ERC20Snapshot.sol

@@ -6,42 +6,46 @@ import "../../utils/Arrays.sol";
 
 /**
  * @title ERC20Snapshot token
- * @dev An ERC20 token which enables taking snapshots of accounts' balances. This can be useful to
- * safely implement voting weighed by balance.
+ * @dev An ERC20 token which enables taking snapshots of account balances.
+ * This can be useful to safely implement voting weighed by balance.
  */
 contract ERC20Snapshot is ERC20 {
 
   using Arrays for uint256[];
 
   // The 0 id represents no snapshot was taken yet.
-  uint256 public currentSnapshotId;
+  uint256 private currentSnapshotId;
 
   mapping (address => uint256[]) private snapshotIds;
   mapping (address => uint256[]) private snapshotBalances;
 
   event Snapshot(uint256 id);
 
+  /**
+   * @dev Increments current snapshot. Emites Snapshot event.
+   * @return An uint256 representing current snapshot id.
+   */
   function snapshot() external returns (uint256) {
     currentSnapshotId += 1;
     emit Snapshot(currentSnapshotId);
     return currentSnapshotId;
   }
 
-  function snapshotsLength(address account) external view returns (uint256) {
-    return snapshotIds[account].length;
-  }
-
+  /**
+   * @dev Returns account balance for specific snapshot.
+   * @param account address The address to query the balance of.
+   * @param snapshotId uint256 The snapshot id for which to query the balance of.
+   * @return An uint256 representing the amount owned by the passed address for specific snapshot.
+   */
   function balanceOfAt(
     address account,
     uint256 snapshotId
   )
-    external 
-    view 
-    returns (uint256) 
+    public
+    view
+    returns (uint256)
   {
-    require(
-      snapshotId > 0 && snapshotId <= currentSnapshotId, 
-      "Parameter snapshotId has to be greater than 0 and lower/equal currentSnapshot");
+    require(snapshotId > 0 && snapshotId <= currentSnapshotId);
 
     uint256 idx = snapshotIds[account].findUpperBound(snapshotId);
 
@@ -52,39 +56,69 @@ contract ERC20Snapshot is ERC20 {
     }
   }
 
-  function transfer(
-    address to, 
-    uint256 value
-  ) 
-    public 
-    returns (bool) 
-  {
+  /**
+  * @dev Transfer token for a specified address. It takes balance snapshot for the sender and recipient account
+  * before transfer is done.
+  * @param to The address to transfer to.
+  * @param value The amount to be transferred.
+  */
+  function transfer(address to, uint256 value) public returns (bool) {
     updateSnapshot(msg.sender);
     updateSnapshot(to);
     return super.transfer(to, value);
   }
 
+  /**
+   * @dev Transfer tokens from one address to another. It takes balance snapshot of both accounts before
+   * the transfer is done.
+   * @param from address The address which you want to send tokens from
+   * @param to address The address which you want to transfer to
+   * @param value uint256 the amount of tokens to be transferred
+   */
   function transferFrom(
-    address from, 
-    address to, 
+    address from,
+    address to,
     uint256 value
   ) 
-    public 
-    returns (bool) 
+    public
+    returns (bool)
   {
     updateSnapshot(from);
     updateSnapshot(to);
     return super.transferFrom(from, to, value);
   }
 
-  function updateSnapshot(address account) internal {
+  /**
+   * @dev Internal function that mints an amount of the token and assigns it to
+   * an account. This encapsulates the modification of balances such that the
+   * proper events are emitted. Takes snapshot before tokens are minted.
+   * @param account The account that will receive the created tokens.
+   * @param amount The amount that will be created.
+   */
+  function _mint(address account, uint256 amount) internal {
+    updateSnapshot(account);
+    super._mint(account, amount);
+  }
+
+  /**
+   * @dev Internal function that burns an amount of the token of a given
+   * account. Takes snapshot before tokens are burned.
+   * @param account The account whose tokens will be burnt.
+   * @param amount The amount that will be burnt.
+   */
+  function _burn(address account, uint256 amount) internal {
+    updateSnapshot(account);
+    super._burn(account, amount);
+  }
+
+  function updateSnapshot(address account) private {
     if (lastSnapshotId(account) < currentSnapshotId) {
       snapshotIds[account].push(currentSnapshotId);
       snapshotBalances[account].push(balanceOf(account));
     }
   }
 
-  function lastSnapshotId(address account) internal view returns (uint256) {
+  function lastSnapshotId(address account) private view returns (uint256) {
     uint256[] storage snapshots = snapshotIds[account];
     if (snapshots.length == 0) {
       return 0;

contracts/utils/Arrays.sol

@@ -3,14 +3,22 @@ pragma solidity ^0.4.23;
 import "../math/Math.sol";
 
 
+/**
+ * @title Arrays
+ * @dev Utility library of inline array functions
+ */
 library Arrays {
+
+  /**
+   * @dev Find upper bound for searching element. Utilize binary search function with O(log(n)) cost.
+   */
   function findUpperBound(
-    uint256[] storage array, 
+    uint256[] storage array,
     uint256 element
   ) 
-    internal 
-    view 
-    returns (uint256) 
+    internal
+    view
+    returns (uint256)
   {
     uint256 low = 0;
     uint256 high = array.length;

test/token/ERC20/ERC20Snapshot.test.js

@@ -38,11 +38,6 @@ contract('ERC20Snapshot', function ([_, owner, recipient, spender]) {
         (await this.token.balanceOf(owner)).should.be.bignumber.equal(90);
         (await this.token.balanceOf(recipient)).should.be.bignumber.equal(10);
       });
-
-      it('snapshots do not exist', async function () {
-        (await this.token.snapshotsLength(owner)).should.be.bignumber.equal(0);
-        (await this.token.snapshotsLength(recipient)).should.be.bignumber.equal(0);
-      });
     });
   });
 
@@ -109,5 +104,33 @@ contract('ERC20Snapshot', function ([_, owner, recipient, spender]) {
         (await this.token.balanceOfAt(recipient, 1)).should.be.bignumber.equal(0);
       });
     });
+
+    describe('new tokens are minted', function () {
+      beforeEach(async function () {
+        await this.token.mint(owner, 50);
+      });
+
+      it('snapshot keeps balance before mint', async function () {
+        (await this.token.balanceOfAt(owner, 1)).should.be.bignumber.equal(100);
+      });
+
+      it('current balance is greater after mint', async function () {
+        (await this.token.balanceOf(owner)).should.be.bignumber.equal(150);
+      });
+    });
+
+    describe('chunk tokens are burned', function () {
+      beforeEach(async function () {
+        await this.token.burn(owner, 30);
+      });
+
+      it('snapshot keeps balance before burn', async function () {
+        (await this.token.balanceOfAt(owner, 1)).should.be.bignumber.equal(100);
+      });
+
+      it('crrent balance is lower after burn', async function () {
+        (await this.token.balanceOf(owner)).should.be.bignumber.equal(70);
+      });
+    });
   });
 });