From 2a40b6911b09cfc27f626fb44f6d4d255d91ad73 Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Sat, 24 Aug 2019 14:23:27 +0300
Subject: [PATCH 1/8] add RewardsDistributor

---
 .../dividend/RewardsDistributor.sol           | 176 ++++++++++++++++++
 1 file changed, 176 insertions(+)
 create mode 100644 contracts/BondingCurve/dividend/RewardsDistributor.sol

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
new file mode 100644
index 0000000..1a5c98b
--- /dev/null
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -0,0 +1,176 @@
+
+pragma solidity ^0.5.6;
+
+import "@openzeppelin/contracts-ethereum-package/contracts/math/SafeMath.sol";
+
+
+/// @title RewardsDistributor - Distribute pro rata rewards (dividends)
+/// @author Bogdan Batog (https://batog.info)
+/// @dev Distribute pro rata rewards (dividends) to token holders in O(1) time.
+///      Based on http://batog.info/papers/scalable-reward-distribution.pdf
+///      And on https://solmaz.io/2019/02/24/scalable-reward-changing/
+contract RewardsDistributor {
+    using SafeMath for uint256;
+
+    /// @notice ELIGIBLE_UNIT is the smallest eligible unit for reward. Minimum
+    ///  possible distribution is 1 (wei for Ether) PER ELIGIBLE_UNIT.
+    ///
+    ///  Only multiple of ELIGIBLE_UNIT will be subject to reward
+    ///  distribution. Any fractional part of deposit, smaller than
+    ///  ELIGIBLE_UNIT, won't receive any reward.
+    ///
+    ///  Recommended value 10**(decimals / 2), that is 10**9 for most ERC20.
+    uint256 public ELIGIBLE_UNIT = 10**9;
+
+    /// @notice Stake per address.
+    mapping(address => uint256) stake;
+
+    /// @notice Stake reminder per address.
+    mapping(address => uint256) stakeReminder;
+
+    /// @notice Total staked tokens. In ELIGIBLE_UNIT units.
+    uint256 public stakeTotal;
+
+    /// @notice Total reward since the beginning of time, in units per
+    ///  ELIGIBLE_UNIT.
+    uint256 rewardTotal;
+
+    /// @notice Reminder from the last reward distribution.
+    uint256 rewardRemainder;
+
+    /// @notice Proportional rewards awarded *before* this stake was created.
+    mapping(address => int256) rewardOffset;
+
+
+    event DepositMade(address _from, uint value);
+    event DistributionMade(address _from, uint value);
+    event RewardWithdrawalMade(address _to, uint value);
+    event StakeWithdrawalMade(address _to, uint value);
+
+
+    /// Initialize the contract.
+    constructor() public {
+        stakeTotal = 0;
+        rewardTotal = 0;
+        rewardRemainder = 0;
+    }
+
+
+    /// @notice Deposit funds into contract.
+    function _deposit(address staker, uint256 tokens) internal returns (bool success) {
+
+        uint256 _tokensToAdd = tokens.add(stakeReminder[staker]);
+
+        uint256 _eligibleUnitsToAdd = _tokensToAdd.div(ELIGIBLE_UNIT);
+
+        // update the new reminder for this address
+        stakeReminder[staker] = _tokensToAdd.mod(ELIGIBLE_UNIT);
+
+        // set the current stake for this address
+        stake[staker] = stake[staker].add(_eligibleUnitsToAdd);
+
+        // update total eligible stake units
+        stakeTotal = stakeTotal.add(_eligibleUnitsToAdd);
+
+        // update reward offset
+        rewardOffset[staker] += (int256)(rewardTotal * _eligibleUnitsToAdd);
+
+        emit DepositMade(staker, tokens);
+        return true;
+    }
+
+
+    /// @notice Distribute tokens pro rata to all stakers.
+    function _distribute(uint tokens) internal returns (bool success) {
+        require(tokens > 0);
+        require(stakeTotal > 0);
+
+        // add past distribution reminder
+        uint256 _amountToDistribute = tokens.add(rewardRemainder);
+
+        // determine rewards per eligible stake
+        uint256 _ratio = _amountToDistribute.div(stakeTotal);
+
+        // carry on reminder
+        rewardRemainder = _amountToDistribute.mod(stakeTotal);
+
+        // increase total rewards per stake unit
+        rewardTotal = rewardTotal.add(_ratio);
+
+        emit DistributionMade(msg.sender, tokens);
+        return true;
+    }
+
+
+    /// @notice Withdraw accumulated reward for the staker address.
+    function _withdrawReward(address staker) internal returns (uint256 tokens) {
+
+        uint256 _reward = getReward(staker);
+
+        // refresh reward offset (so a new call to getReward returns 0)
+        rewardOffset[staker] = (int256) (rewardTotal.mul(stake[staker]));
+
+        emit RewardWithdrawalMade(staker, _reward);
+        return _reward;
+    }
+
+
+    /// @notice Withdraw stake for the staker address
+    function _withdrawStake(address staker, uint256 tokens) internal returns (bool) {
+
+        uint256 _currentStake = getStake(staker);
+
+        require(tokens <= _currentStake);
+
+        // update stake and reminder for this address
+        uint256 _newStake = _currentStake.sub(tokens);
+
+        stakeReminder[staker] = _newStake.mod(ELIGIBLE_UNIT);
+
+        uint256 _eligibleUnitsDelta = stake[staker].sub(
+            _newStake.div(ELIGIBLE_UNIT)
+        );
+
+        stake[staker] = stake[staker].sub(_eligibleUnitsDelta);
+
+        // update total stake
+        stakeTotal = stakeTotal.sub(_eligibleUnitsDelta);
+
+        // update reward offset
+        rewardOffset[staker] -= (int256) (rewardTotal.mul(_eligibleUnitsDelta));
+
+        emit StakeWithdrawalMade(staker, tokens);
+        return true;
+    }
+
+    ///
+    /// READ ONLY
+    ///
+
+    /// @notice Read current stake for address.
+    function getStake(address staker) public view returns (uint256 tokens) {
+        tokens = (
+            stake[staker].mul(ELIGIBLE_UNIT)
+        ).add(
+            stakeReminder[staker]
+        );
+
+        return tokens;
+    }
+
+
+    /// @notice Read current accumulated reward for address.
+    function getReward(address staker) public view returns (uint256 tokens) {
+        int256 _tokens = (
+            (int256)(
+                stake[staker].mul(rewardTotal)
+            ) - rewardOffset[staker]
+        );
+
+        tokens = (uint256) (_tokens);
+
+        return tokens;
+    }
+
+}
+

From f12c27b4e0673049f29211b9c6b1564a10b07991 Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Mon, 26 Aug 2019 14:31:38 +0300
Subject: [PATCH 2/8] wrapper, first unit tests

---
 .../dividend/RewardsDistributor.sol           | 75 ++++++++++---------
 .../dividend/RewardsDistributorWrapper.sol    | 39 ++++++++++
 index.js                                      | 18 ++++-
 test/unit/rewardDistributor.spec.ts           | 52 +++++++++++++
 4 files changed, 148 insertions(+), 36 deletions(-)
 create mode 100644 contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
 create mode 100644 test/unit/rewardDistributor.spec.ts

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
index 1a5c98b..ff2d5e1 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributor.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -20,60 +20,60 @@ contract RewardsDistributor {
     ///  ELIGIBLE_UNIT, won't receive any reward.
     ///
     ///  Recommended value 10**(decimals / 2), that is 10**9 for most ERC20.
-    uint256 public ELIGIBLE_UNIT = 10**9;
+    uint256 public constant ELIGIBLE_UNIT = 10**9;
 
     /// @notice Stake per address.
-    mapping(address => uint256) stake;
+    mapping(address => uint256) internal _stake;
 
     /// @notice Stake reminder per address.
-    mapping(address => uint256) stakeReminder;
+    mapping(address => uint256) internal _stakeReminder;
 
     /// @notice Total staked tokens. In ELIGIBLE_UNIT units.
-    uint256 public stakeTotal;
+    uint256 internal _stakeTotal;
 
     /// @notice Total reward since the beginning of time, in units per
     ///  ELIGIBLE_UNIT.
-    uint256 rewardTotal;
+    uint256 internal _rewardTotal;
 
     /// @notice Reminder from the last reward distribution.
-    uint256 rewardRemainder;
+    uint256 internal _rewardRemainder;
 
     /// @notice Proportional rewards awarded *before* this stake was created.
-    mapping(address => int256) rewardOffset;
+    mapping(address => int256) _rewardOffset;
 
 
-    event DepositMade(address _from, uint value);
-    event DistributionMade(address _from, uint value);
-    event RewardWithdrawalMade(address _to, uint value);
-    event StakeWithdrawalMade(address _to, uint value);
+    event DepositMade(address indexed _from, uint256 value);
+    event DistributionMade(address indexed _from, uint256 value);
+    event RewardWithdrawalMade(address indexed _to, uint256 value);
+    event StakeWithdrawalMade(address indexed _to, uint256 value);
 
 
     /// Initialize the contract.
     constructor() public {
-        stakeTotal = 0;
-        rewardTotal = 0;
-        rewardRemainder = 0;
+        _stakeTotal = 0;
+        _rewardTotal = 0;
+        _rewardRemainder = 0;
     }
 
 
     /// @notice Deposit funds into contract.
     function _deposit(address staker, uint256 tokens) internal returns (bool success) {
 
-        uint256 _tokensToAdd = tokens.add(stakeReminder[staker]);
+        uint256 _tokensToAdd = tokens.add(_stakeReminder[staker]);
 
         uint256 _eligibleUnitsToAdd = _tokensToAdd.div(ELIGIBLE_UNIT);
 
         // update the new reminder for this address
-        stakeReminder[staker] = _tokensToAdd.mod(ELIGIBLE_UNIT);
+        _stakeReminder[staker] = _tokensToAdd.mod(ELIGIBLE_UNIT);
 
         // set the current stake for this address
-        stake[staker] = stake[staker].add(_eligibleUnitsToAdd);
+        _stake[staker] = _stake[staker].add(_eligibleUnitsToAdd);
 
         // update total eligible stake units
-        stakeTotal = stakeTotal.add(_eligibleUnitsToAdd);
+        _stakeTotal = _stakeTotal.add(_eligibleUnitsToAdd);
 
         // update reward offset
-        rewardOffset[staker] += (int256)(rewardTotal * _eligibleUnitsToAdd);
+        _rewardOffset[staker] += (int256)(_rewardTotal * _eligibleUnitsToAdd);
 
         emit DepositMade(staker, tokens);
         return true;
@@ -83,19 +83,19 @@ contract RewardsDistributor {
     /// @notice Distribute tokens pro rata to all stakers.
     function _distribute(uint tokens) internal returns (bool success) {
         require(tokens > 0);
-        require(stakeTotal > 0);
+        require(_stakeTotal > 0);
 
         // add past distribution reminder
-        uint256 _amountToDistribute = tokens.add(rewardRemainder);
+        uint256 _amountToDistribute = tokens.add(_rewardRemainder);
 
         // determine rewards per eligible stake
-        uint256 _ratio = _amountToDistribute.div(stakeTotal);
+        uint256 _ratio = _amountToDistribute.div(_stakeTotal);
 
         // carry on reminder
-        rewardRemainder = _amountToDistribute.mod(stakeTotal);
+        _rewardRemainder = _amountToDistribute.mod(_stakeTotal);
 
         // increase total rewards per stake unit
-        rewardTotal = rewardTotal.add(_ratio);
+        _rewardTotal = _rewardTotal.add(_ratio);
 
         emit DistributionMade(msg.sender, tokens);
         return true;
@@ -108,7 +108,7 @@ contract RewardsDistributor {
         uint256 _reward = getReward(staker);
 
         // refresh reward offset (so a new call to getReward returns 0)
-        rewardOffset[staker] = (int256) (rewardTotal.mul(stake[staker]));
+        _rewardOffset[staker] = (int256) (_rewardTotal.mul(_stake[staker]));
 
         emit RewardWithdrawalMade(staker, _reward);
         return _reward;
@@ -125,19 +125,19 @@ contract RewardsDistributor {
         // update stake and reminder for this address
         uint256 _newStake = _currentStake.sub(tokens);
 
-        stakeReminder[staker] = _newStake.mod(ELIGIBLE_UNIT);
+        _stakeReminder[staker] = _newStake.mod(ELIGIBLE_UNIT);
 
-        uint256 _eligibleUnitsDelta = stake[staker].sub(
+        uint256 _eligibleUnitsDelta = _stake[staker].sub(
             _newStake.div(ELIGIBLE_UNIT)
         );
 
-        stake[staker] = stake[staker].sub(_eligibleUnitsDelta);
+        _stake[staker] = _stake[staker].sub(_eligibleUnitsDelta);
 
         // update total stake
-        stakeTotal = stakeTotal.sub(_eligibleUnitsDelta);
+        _stakeTotal = _stakeTotal.sub(_eligibleUnitsDelta);
 
         // update reward offset
-        rewardOffset[staker] -= (int256) (rewardTotal.mul(_eligibleUnitsDelta));
+        _rewardOffset[staker] -= (int256) (_rewardTotal.mul(_eligibleUnitsDelta));
 
         emit StakeWithdrawalMade(staker, tokens);
         return true;
@@ -147,12 +147,19 @@ contract RewardsDistributor {
     /// READ ONLY
     ///
 
+
+    /// @notice Read total stake.
+    function getStakeTotal() public returns (uint256) {
+        return _stakeTotal;
+    }
+
+
     /// @notice Read current stake for address.
     function getStake(address staker) public view returns (uint256 tokens) {
         tokens = (
-            stake[staker].mul(ELIGIBLE_UNIT)
+            _stake[staker].mul(ELIGIBLE_UNIT)
         ).add(
-            stakeReminder[staker]
+            _stakeReminder[staker]
         );
 
         return tokens;
@@ -163,8 +170,8 @@ contract RewardsDistributor {
     function getReward(address staker) public view returns (uint256 tokens) {
         int256 _tokens = (
             (int256)(
-                stake[staker].mul(rewardTotal)
-            ) - rewardOffset[staker]
+                _stake[staker].mul(_rewardTotal)
+            ) - _rewardOffset[staker]
         );
 
         tokens = (uint256) (_tokens);
diff --git a/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol b/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
new file mode 100644
index 0000000..9efdb85
--- /dev/null
+++ b/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
@@ -0,0 +1,39 @@
+pragma solidity ^0.5.6;
+
+import "./RewardsDistributor.sol";
+
+
+/// @title RewardsDistributorWrapper
+/// @author Bogdan Batog (https://batog.info)
+/// @dev ONLY FOR TESTING. DO NOT DEPLOY THIS!!!!!
+contract RewardsDistributorWrapper is RewardsDistributor {
+
+    /// @notice Deposit funds into contract.
+    function deposit(address staker, uint tokens) public returns (bool success) {
+        return _deposit(staker, tokens);
+    }
+
+    /// @notice Distribute tokens pro rata to all stakers.
+    function distribute(uint tokens) public returns (bool success) {
+        return _distribute(tokens);
+    }
+
+    /// @notice Withdraw accumulated reward for the staker address.
+    function withdrawReward(address staker) public returns (uint tokens) {
+        return _withdrawReward(staker);
+    }
+
+    /// @notice Withdraw all stake for the staker address.
+    function withdrawAllStake(address staker) public returns (uint tokens) {
+        tokens = getStake(staker);
+        _withdrawStake(staker, tokens);
+
+        return tokens;
+    }
+
+    /// @notice Withdraw stake for the staker address.
+    function withdrawStake(address staker, uint tokens) public returns (bool success) {
+        return _withdrawStake(staker, tokens);
+    }
+}
+
diff --git a/index.js b/index.js
index c47f620..107bcb7 100644
--- a/index.js
+++ b/index.js
@@ -12,6 +12,7 @@ const DividendPool = Contracts.getFromLocal('DividendPool');
 const BondingCurve = Contracts.getFromLocal('BondingCurve');
 const BondingCurveFactory = Contracts.getFromLocal('BondingCurveFactory');
 const BondedToken = Contracts.getFromLocal('BondedToken');
+const RewardsDistributorWrapper = Contracts.getFromLocal('RewardsDistributorWrapper');
 
 const CONTRACT_ABIS = {
   BondingCurve,
@@ -20,7 +21,8 @@ const CONTRACT_ABIS = {
   StaticCurveLogic,
   BondedToken,
   DividendPool,
-  BancorCurveService
+  BancorCurveService,
+  RewardsDistributorWrapper
 };
 
 const CONTRACT_NAMES = {
@@ -30,7 +32,8 @@ const CONTRACT_NAMES = {
   StaticCurveLogic: 'StaticCurveLogic',
   BondedToken: 'BondedToken',
   DividendPool: 'DividendPool',
-  BancorCurveService: 'BancorCurveService'
+  BancorCurveService: 'BancorCurveService',
+  RewardsDistributorWrapper: 'RewardsDistributorWrapper'
 };
 
 const PACKAGE_NAMES = {
@@ -55,6 +58,7 @@ async function setupApp(txParams) {
   await appProject.setImplementation(BondedToken, CONTRACT_NAMES.BondedToken);
   await appProject.setImplementation(DividendPool, CONTRACT_NAMES.DividendPool);
   await appProject.setImplementation(BancorCurveService, CONTRACT_NAMES.BancorCurveService);
+  await appProject.setImplementation(RewardsDistributor, CONTRACT_NAMES.RewardsDistributor);
 
   return appProject;
 }
@@ -166,6 +170,15 @@ async function deployStandaloneERC20(myProject, initArgs) {
   //   });
 }
 
+async function deployRewardsDistributorWrapper(myProject, initArgs) {
+  ZWeb3.initialize(web3.currentProvider);
+
+  const instance = await myProject.createProxy(RewardsDistributorWrapper, {
+    initArgs
+  });
+  return instance;
+}
+
 async function getImplementation(project, contractName) {
   const directory = await project.getCurrentDirectory();
   const implementation = await directory.getImplementation(contractName);
@@ -203,6 +216,7 @@ module.exports = {
   deployBondingCurveFactory,
   deployBondedToken,
   deployStandaloneERC20,
+  deployRewardsDistributorWrapper,
   CONTRACT_NAMES,
   CONTRACT_ABIS,
   getImplementation,
diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
new file mode 100644
index 0000000..35080bd
--- /dev/null
+++ b/test/unit/rewardDistributor.spec.ts
@@ -0,0 +1,52 @@
+// Import all required modules from openzeppelin-test-helpers
+const {BN, constants, expectEvent, expectRevert} = require('openzeppelin-test-helpers');
+
+// Import preferred chai flavor: both expect and should are supported
+const expect = require('chai').expect;
+const should = require('chai').should();
+
+require('../setup');
+const {deployProject, deployRewardsDistributorWrapper} = require('../../index.js');
+
+var TEN18 = new BN(String(10**18))
+
+var PPB = new BN(String(10**9))
+
+
+contract('RewardsDistributorWrapper', accounts => {
+  let tx;
+  let project;
+
+  const creator = accounts[0];
+  const initializer = accounts[1];
+
+  beforeEach(async function() {
+    project = await deployProject();
+    rd = await deployRewardsDistributorWrapper(project);
+  });
+
+  it("deploys and initializes", async function() {
+    let stakeTotal = await rd.methods.getStakeTotal().call({from: initializer});
+    assert.equal(stakeTotal, 0, "stakeTotal is NOT zero!");
+  });
+
+  it("deposit A, getStake, withdrawAllStake, getStake", async function() {
+    var acct_a = accounts[1];
+
+    var amount = (new BN("100")).mul(TEN18);
+
+    let r1 = await rd.methods.deposit(acct_a, amount.toString()).call({from: acct_a});
+    expect(r1).to.be.equal(true);
+
+    let stakeTotal = await rd.methods.getStakeTotal().call({from: initializer});
+    expect(new BN(stakeTotal)).to.be.bignumber.equal(amount);
+
+    let stake = await rd.methods.getStake(acct_a).call({from: acct_a});
+    expect(new BN(stake)).to.be.bignumber.equal(amount);
+
+    await rd.methods.withdrawAllStake().call({from: acct_a});
+    let _stakeFinal = await rd.methods.getStake(acct_a).call({from: acct_a});
+    expect(new BN(stakeTotal)).to.be.bignumber.equal(0);
+  });
+
+})
\ No newline at end of file

From 2b55825380ff898b8d1f9d925d171f446c3f6a4a Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Wed, 28 Aug 2019 13:12:54 +0300
Subject: [PATCH 3/8] fix tests

---
 test/unit/rewardDistributor.spec.ts | 70 +++++++++++++++++++++--------
 1 file changed, 52 insertions(+), 18 deletions(-)

diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
index 35080bd..5e10bde 100644
--- a/test/unit/rewardDistributor.spec.ts
+++ b/test/unit/rewardDistributor.spec.ts
@@ -1,21 +1,24 @@
 // Import all required modules from openzeppelin-test-helpers
-const {BN, constants, expectEvent, expectRevert} = require('openzeppelin-test-helpers');
+const {BN, constants, expectRevert} = require('openzeppelin-test-helpers');
 
 // Import preferred chai flavor: both expect and should are supported
 const expect = require('chai').expect;
 const should = require('chai').should();
+const expectEvent = require('../expectEvent');
 
 require('../setup');
-const {deployProject, deployRewardsDistributorWrapper} = require('../../index.js');
 
-var TEN18 = new BN(String(10**18))
+const {deployProject, deployRewardsDistributorWrapper} = require('../../index.js');
 
-var PPB = new BN(String(10**9))
+var TEN18 = new BN(String(10 ** 18));
 
+var PPB = new BN(String(10 ** 9));
 
 contract('RewardsDistributorWrapper', accounts => {
-  let tx;
   let project;
+  let rd;
+  let tx;
+  let ELIGIBLE_UNIT;
 
   const creator = accounts[0];
   const initializer = accounts[1];
@@ -23,30 +26,61 @@ contract('RewardsDistributorWrapper', accounts => {
   beforeEach(async function() {
     project = await deployProject();
     rd = await deployRewardsDistributorWrapper(project);
+    ELIGIBLE_UNIT = rd.ELIGIBLE_UNIT;
   });
 
-  it("deploys and initializes", async function() {
+  it('deploys and initializes', async function() {
     let stakeTotal = await rd.methods.getStakeTotal().call({from: initializer});
-    assert.equal(stakeTotal, 0, "stakeTotal is NOT zero!");
+    expect(new BN(stakeTotal)).to.be.bignumber.equal(new BN(0));
   });
 
-  it("deposit A, getStake, withdrawAllStake, getStake", async function() {
+  it('accepts deposit A, gets stake, withdraws all stake, gets stake again', async function() {
     var acct_a = accounts[1];
+    var amount = new BN('100').mul(TEN18);
 
-    var amount = (new BN("100")).mul(TEN18);
+    tx = await rd.methods
+      .deposit(acct_a, amount.toString())
+      .send({from: acct_a});
 
-    let r1 = await rd.methods.deposit(acct_a, amount.toString()).call({from: acct_a});
-    expect(r1).to.be.equal(true);
-
-    let stakeTotal = await rd.methods.getStakeTotal().call({from: initializer});
-    expect(new BN(stakeTotal)).to.be.bignumber.equal(amount);
+    expectEvent.inLogs(tx.events, 'DepositMade', {
+      _from: acct_a,
+      value: amount
+    });
 
     let stake = await rd.methods.getStake(acct_a).call({from: acct_a});
     expect(new BN(stake)).to.be.bignumber.equal(amount);
 
-    await rd.methods.withdrawAllStake().call({from: acct_a});
-    let _stakeFinal = await rd.methods.getStake(acct_a).call({from: acct_a});
-    expect(new BN(stakeTotal)).to.be.bignumber.equal(0);
+    await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
+
+    let stakeFinal = await rd.methods.getStake(acct_a).call({from: acct_a});
+    expect(new BN(stakeFinal)).to.be.bignumber.equal(new BN(0));
   });
 
-})
\ No newline at end of file
+  it('deposit A, distribute, getReward, withdrawReward and getReward', async function() {
+    var acct_a = accounts[1];
+    var amountDeposit = new BN('100').mul(TEN18);
+    var amountDistribute = new BN('200').mul(TEN18);
+
+    await rd.methods.deposit(acct_a, amountDeposit.toString()).send({from: acct_a});
+
+    tx = await rd.methods.distribute(amountDistribute.toString()).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'DistributionMade', {
+      _from: acct_a,
+      value: amountDistribute
+    });
+
+    // all reward is allocated to the single staker
+    var currentReward = await rd.methods.getReward(acct_a).call({from: acct_a});
+    expect(new BN(currentReward)).to.be.bignumber.equal(amountDistribute);
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_a,
+      value: amountDistribute
+    });
+
+    currentReward = await rd.methods.getReward(acct_a).call({from: acct_a});
+    expect(new BN(currentReward)).to.be.bignumber.equal(new BN(0));
+  });
+    
+});

From a1af84a9df35dcd50ac6ecda19741180363cad41 Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Wed, 28 Aug 2019 23:39:48 +0300
Subject: [PATCH 4/8] more tests

---
 test/unit/rewardDistributor.spec.ts | 150 +++++++++++++++++++++++++++-
 1 file changed, 147 insertions(+), 3 deletions(-)

diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
index 5e10bde..8594c97 100644
--- a/test/unit/rewardDistributor.spec.ts
+++ b/test/unit/rewardDistributor.spec.ts
@@ -34,7 +34,24 @@ contract('RewardsDistributorWrapper', accounts => {
     expect(new BN(stakeTotal)).to.be.bignumber.equal(new BN(0));
   });
 
-  it('accepts deposit A, gets stake, withdraws all stake, gets stake again', async function() {
+  it("withdraws ZERO reward", async function() {
+    var acct_a = accounts[1];
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_a,
+      value: new BN('0')
+    });
+  });
+
+  it("reads ZERO stake", async function() {
+    var acct_a = accounts[1];
+
+    let stake = await rd.methods.getStake(acct_a).call({from: acct_a});
+    expect(new BN(stake)).to.be.bignumber.equal(new BN('0'));
+  });
+
+  it('deposits A, gets stake, withdraws all stake, gets stake again', async function() {
     var acct_a = accounts[1];
     var amount = new BN('100').mul(TEN18);
 
@@ -56,7 +73,7 @@ contract('RewardsDistributorWrapper', accounts => {
     expect(new BN(stakeFinal)).to.be.bignumber.equal(new BN(0));
   });
 
-  it('deposit A, distribute, getReward, withdrawReward and getReward', async function() {
+  it('deposits A, distributes, gets Reward, withdraws Reward and gets Reward', async function() {
     var acct_a = accounts[1];
     var amountDeposit = new BN('100').mul(TEN18);
     var amountDistribute = new BN('200').mul(TEN18);
@@ -82,5 +99,132 @@ contract('RewardsDistributorWrapper', accounts => {
     currentReward = await rd.methods.getReward(acct_a).call({from: acct_a});
     expect(new BN(currentReward)).to.be.bignumber.equal(new BN(0));
   });
-    
+
+ it("deposits A, deposits B, distributes, withdraws stake, distributes, gets reward", async function() {
+    var acct_a = accounts[1];
+    var acct_b = accounts[2];
+
+    var depositA = new BN('100').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+
+    var depositB = new BN('300').mul(TEN18);
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+
+    var distribute1 = new BN('400').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+
+    var stakeA = await rd.methods.getStake(acct_a).call({from: acct_a});
+    expect(new BN(stakeA)).to.be.bignumber.equal(depositA);
+    await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
+
+    // a second distribution after A has withdrawn entirely
+    var distribute2 = new BN('900').mul(TEN18);
+    await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
+
+    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
+    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('1200').mul(TEN18));
+  });
+
+  it("deposits A, deposits B, distributes and withdraws reward", async function() {
+    var acct_a = accounts[1];
+    var acct_b = accounts[2];
+
+    var depositA = new BN('100').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+
+    var depositB = new BN('300').mul(TEN18);
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+
+    var distribute1 = new BN('40').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+
+    tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_b,
+      value: new BN('30').mul(TEN18)
+    });
+  });
+
+  it("deposits A, deposits B, distributes X 2 and withdraws reward", async function() {
+    var acct_a = accounts[1];
+    var acct_b = accounts[2];
+
+    var depositA = new BN('100').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+
+    var depositB = new BN('300').mul(TEN18);
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+
+    var distribute1 = new BN('400').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+
+    var distribute2 = new BN('4000').mul(TEN18);
+    await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_a,
+      value: new BN('1100').mul(TEN18)
+    });
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_a,
+      value: new BN('0')
+    });
+
+    tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      _from: acct_b,
+      value: new BN('3300').mul(TEN18)
+    });
+  });
+
+  it("deposits A, distributes, deposits B, distributes and reads reward", async function() {
+    var acct_a = accounts[1];
+    var acct_b = accounts[2];
+
+    var depositA = new BN('100').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+
+    var distribute1 = new BN('100').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+
+    var depositB = new BN('300').mul(TEN18);
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+
+    var stakeTotal = await rd.methods.getStakeTotal().call({from: acct_a});
+    expect(new BN(stakeTotal)).to.be.bignumber.equal(new BN('400').mul(TEN18));
+
+    var distribute2 = new BN('100').mul(TEN18);
+    await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
+
+    var rewardA = await rd.methods.getReward(acct_a).call({from: acct_a});
+    expect(new BN(rewardA)).to.be.bignumber.equal(new BN('125').mul(TEN18));
+
+    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
+    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('75').mul(TEN18));
+  });
+
+  it("handles magnitudes: A deposits 9999, B deposits 1, distribute, withdraw stake, distribute, withdraw reward", async function() {
+    var acct_a = accounts[1];
+    var acct_b = accounts[2];
+
+    var depositA = new BN('9999').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+
+    var depositB = new BN('1').mul(TEN18);
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+
+    var distribute1 = new BN('1').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+
+    await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
+
+    var distribute2 = new BN('2').mul(TEN18);
+    await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
+
+    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
+    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('2000100000000000000'));
+  });
 });

From 59bfdfe463e6142364841d29b7beaed10dc4490a Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Thu, 29 Aug 2019 12:51:12 +0300
Subject: [PATCH 5/8] better tests

---
 .../dividend/RewardsDistributor.sol           |  12 +-
 test/unit/rewardDistributor.spec.ts           | 337 +++++++++++++-----
 2 files changed, 257 insertions(+), 92 deletions(-)

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
index ff2d5e1..8eb3dea 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributor.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -42,10 +42,10 @@ contract RewardsDistributor {
     mapping(address => int256) _rewardOffset;
 
 
-    event DepositMade(address indexed _from, uint256 value);
-    event DistributionMade(address indexed _from, uint256 value);
-    event RewardWithdrawalMade(address indexed _to, uint256 value);
-    event StakeWithdrawalMade(address indexed _to, uint256 value);
+    event DepositMade(address indexed from, uint256 value);
+    event DistributionMade(uint256 value);
+    event RewardWithdrawalMade(address indexed to, uint256 value);
+    event StakeWithdrawalMade(address indexed to, uint256 value);
 
 
     /// Initialize the contract.
@@ -97,7 +97,7 @@ contract RewardsDistributor {
         // increase total rewards per stake unit
         _rewardTotal = _rewardTotal.add(_ratio);
 
-        emit DistributionMade(msg.sender, tokens);
+        emit DistributionMade(tokens);
         return true;
     }
 
@@ -150,7 +150,7 @@ contract RewardsDistributor {
 
     /// @notice Read total stake.
     function getStakeTotal() public returns (uint256) {
-        return _stakeTotal;
+        return _stakeTotal.mul(ELIGIBLE_UNIT);
     }
 
 
diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
index 8594c97..3f60047 100644
--- a/test/unit/rewardDistributor.spec.ts
+++ b/test/unit/rewardDistributor.spec.ts
@@ -20,8 +20,10 @@ contract('RewardsDistributorWrapper', accounts => {
   let tx;
   let ELIGIBLE_UNIT;
 
-  const creator = accounts[0];
-  const initializer = accounts[1];
+  let acct_a = accounts[1];
+  let acct_b = accounts[2];
+  let acct_c = accounts[3];
+  let acct_d = accounts[4];
 
   beforeEach(async function() {
     project = await deployProject();
@@ -30,29 +32,32 @@ contract('RewardsDistributorWrapper', accounts => {
   });
 
   it('deploys and initializes', async function() {
-    let stakeTotal = await rd.methods.getStakeTotal().call({from: initializer});
-    expect(new BN(stakeTotal)).to.be.bignumber.equal(new BN(0));
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal('0');
   });
 
   it("withdraws ZERO reward", async function() {
-    var acct_a = accounts[1];
-
     tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_a,
+      to: acct_a,
       value: new BN('0')
     });
   });
 
   it("reads ZERO stake", async function() {
-    var acct_a = accounts[1];
+    expect(
+      await rd.methods.getStake(acct_a).call({from: acct_a})
+    ).to.be.equal('0');
+  });
 
-    let stake = await rd.methods.getStake(acct_a).call({from: acct_a});
-    expect(new BN(stake)).to.be.bignumber.equal(new BN('0'));
+  it("reverts if trying to withdraw amount > stake", async function() {
+    await expectRevert.unspecified(
+      rd.methods.withdrawStake(acct_a, '1').send({from: acct_a})
+    );
   });
 
-  it('deposits A, gets stake, withdraws all stake, gets stake again', async function() {
-    var acct_a = accounts[1];
+  it('updates total stake after deposit > ELIGIBLE_UNIT', async function() {
     var amount = new BN('100').mul(TEN18);
 
     tx = await rd.methods
@@ -60,21 +65,83 @@ contract('RewardsDistributorWrapper', accounts => {
       .send({from: acct_a});
 
     expectEvent.inLogs(tx.events, 'DepositMade', {
-      _from: acct_a,
+      from: acct_a,
       value: amount
     });
 
-    let stake = await rd.methods.getStake(acct_a).call({from: acct_a});
-    expect(new BN(stake)).to.be.bignumber.equal(amount);
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal(amount.toString());
+  });
+
+  it('doesn\'t update total stake after deposit < ELIGIBLE_UNIT', async function() {
+    var amount = new BN('100');
+
+    tx = await rd.methods
+      .deposit(acct_a, amount.toString())
+      .send({from: acct_a});
+
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal('0');
+  });
+
+  it('deposits A, gets stake, withdraws all stake, gets stake again', async function() {
+    var amount = new BN('100').mul(TEN18);
+
+    tx = await rd.methods
+      .deposit(acct_a, amount.toString())
+      .send({from: acct_a});
+
+    expect(
+      await rd.methods.getStake(acct_a).call({from: acct_a})
+    ).to.be.equal(amount.toString());
 
     await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
 
-    let stakeFinal = await rd.methods.getStake(acct_a).call({from: acct_a});
-    expect(new BN(stakeFinal)).to.be.bignumber.equal(new BN(0));
+    expect(
+      await rd.methods.getStake(acct_a).call({from: acct_a})
+    ).to.be.equal('0');
+  });
+
+  it("does no distribution if no stake >= ELIGIBLE_UNIT", async function() {
+    await rd.methods.deposit(acct_a, '1234').send({from: acct_a});
+
+    await expectRevert.unspecified(
+      rd.methods.distribute('1000').send({from: acct_a}),
+      'no deposit greater than 1 ELIGIBLE_UNIT'
+    );
   });
 
-  it('deposits A, distributes, gets Reward, withdraws Reward and gets Reward', async function() {
-    var acct_a = accounts[1];
+  it("does no distribution if stake becomes ineligible after withdrawl", async function() {
+    var amountDeposit = new BN('100').mul(TEN18);
+    var amountWithdraw = new BN('100').mul(TEN18).sub(new BN('10000'));
+
+    await rd.methods.deposit(acct_a, amountDeposit.toString()).send({from: acct_a});
+
+    tx = await rd.methods.withdrawStake(acct_a, amountWithdraw.toString()).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'StakeWithdrawalMade', {
+      to: acct_b,
+      value: amountWithdraw
+    });
+
+    await expectRevert.unspecified(
+      rd.methods.distribute('1000').send({from: acct_a}),
+      'no deposit greater than 1 ELIGIBLE_UNIT'
+    );
+
+    // stake of A is 10000 but because it is < ELIGIBLE_UNIT total stake
+    // should be zero
+    expect(
+      await rd.methods.getStake(acct_a).call({from: acct_a})
+    ).to.be.equal('10000');
+
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal('0');
+  });
+
+  it('allocates all reward to a single staker and allow its withdrawl', async function() {
     var amountDeposit = new BN('100').mul(TEN18);
     var amountDistribute = new BN('200').mul(TEN18);
 
@@ -82,80 +149,87 @@ contract('RewardsDistributorWrapper', accounts => {
 
     tx = await rd.methods.distribute(amountDistribute.toString()).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'DistributionMade', {
-      _from: acct_a,
       value: amountDistribute
     });
 
-    // all reward is allocated to the single staker
-    var currentReward = await rd.methods.getReward(acct_a).call({from: acct_a});
-    expect(new BN(currentReward)).to.be.bignumber.equal(amountDistribute);
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal(amountDistribute.toString())
 
     tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_a,
+      to: acct_a,
       value: amountDistribute
     });
 
-    currentReward = await rd.methods.getReward(acct_a).call({from: acct_a});
-    expect(new BN(currentReward)).to.be.bignumber.equal(new BN(0));
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal('0')
   });
 
- it("deposits A, deposits B, distributes, withdraws stake, distributes, gets reward", async function() {
-    var acct_a = accounts[1];
-    var acct_b = accounts[2];
+  it("allocates no reward to stake <= ELIGIBLE_UNIT", async function() {
+    await rd.methods.deposit(acct_a, String(10 ** 9)).send({from: acct_a});
+    await rd.methods.deposit(acct_b, String(100)).send({from: acct_a});
 
-    var depositA = new BN('100').mul(TEN18);
-    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+    var distribute1 = new BN('100').mul(TEN18);
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
-    var depositB = new BN('300').mul(TEN18);
-    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+    // all reward goes to A
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal(distribute1.toString())
 
-    var distribute1 = new BN('400').mul(TEN18);
+    expect (
+      await rd.methods.getReward(acct_b).call({from: acct_a})
+    ).to.be.equal('0')
+  });
+
+ it("allocates 1st reward proportionally to 2 stakers and 2nd reward to remaining staker after the other withdrew", async function() {
+    var depositA = new BN(String(10 * 10 ** 9));  // 10 ELIGIBLE_UNITS
+    var depositB = new BN(String(30 * 10 ** 9));
+    var distribute1 = new BN('400');  // notice this is in wei
+    var distribute2 = new BN('900');
+
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
     await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
-    var stakeA = await rd.methods.getStake(acct_a).call({from: acct_a});
-    expect(new BN(stakeA)).to.be.bignumber.equal(depositA);
+    // second distribution after A has withdrawn entirely
     await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
-
-    // a second distribution after A has withdrawn entirely
-    var distribute2 = new BN('900').mul(TEN18);
     await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
 
-    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
-    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('1200').mul(TEN18));
-  });
+    expect (
+      await rd.methods.getReward(acct_b).call({from: acct_a})
+    ).to.be.equal('1200')
 
-  it("deposits A, deposits B, distributes and withdraws reward", async function() {
-    var acct_a = accounts[1];
-    var acct_b = accounts[2];
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal('100')
+  });
 
+  it("withdraws reward after proportional reward distribution", async function() {
     var depositA = new BN('100').mul(TEN18);
-    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
-
     var depositB = new BN('300').mul(TEN18);
-    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
-
     var distribute1 = new BN('40').mul(TEN18);
+
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
     await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
     tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_b,
+      to: acct_b,
       value: new BN('30').mul(TEN18)
     });
   });
 
-  it("deposits A, deposits B, distributes X 2 and withdraws reward", async function() {
-    var acct_a = accounts[1];
-    var acct_b = accounts[2];
-
+  it("withdraws reward after two consecutive reward distributions", async function() {
     var depositA = new BN('100').mul(TEN18);
-    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
-
     var depositB = new BN('300').mul(TEN18);
-    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
-
     var distribute1 = new BN('400').mul(TEN18);
+
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+    await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
     await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
     var distribute2 = new BN('4000').mul(TEN18);
@@ -163,68 +237,159 @@ contract('RewardsDistributorWrapper', accounts => {
 
     tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_a,
+      to: acct_a,
       value: new BN('1100').mul(TEN18)
     });
 
     tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_a,
+      to: acct_a,
       value: new BN('0')
     });
 
     tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
-      _from: acct_b,
+      to: acct_b,
       value: new BN('3300').mul(TEN18)
     });
   });
 
-  it("deposits A, distributes, deposits B, distributes and reads reward", async function() {
-    var acct_a = accounts[1];
-    var acct_b = accounts[2];
-
+  it("distributes after partial stake withdrawal and reads reward", async function() {
     var depositA = new BN('100').mul(TEN18);
-    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
-
+    var depositB = new BN('300').mul(TEN18);
     var distribute1 = new BN('100').mul(TEN18);
-    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
+    var withdrawB = new BN('200').mul(TEN18);
+    var distribute2 = new BN('100').mul(TEN18);
 
-    var depositB = new BN('300').mul(TEN18);
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
     await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
-    var stakeTotal = await rd.methods.getStakeTotal().call({from: acct_a});
-    expect(new BN(stakeTotal)).to.be.bignumber.equal(new BN('400').mul(TEN18));
+    tx = await rd.methods.withdrawStake(acct_b, withdrawB.toString()).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'StakeWithdrawalMade', {
+      to: acct_b,
+      value: withdrawB
+    });
 
-    var distribute2 = new BN('100').mul(TEN18);
     await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
 
-    var rewardA = await rd.methods.getReward(acct_a).call({from: acct_a});
-    expect(new BN(rewardA)).to.be.bignumber.equal(new BN('125').mul(TEN18));
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal(String(200 * 10 ** 18));
 
-    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
-    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('75').mul(TEN18));
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal(String(75 * 10 ** 18));
+
+    expect (
+      await rd.methods.getReward(acct_b).call({from: acct_a})
+    ).to.be.equal(String(125 * 10 ** 18));
   });
 
-  it("handles magnitudes: A deposits 9999, B deposits 1, distribute, withdraw stake, distribute, withdraw reward", async function() {
-    var acct_a = accounts[1];
-    var acct_b = accounts[2];
+  it("withdraws reward after stake has been withdrawn", async function() {
+    var depositA = new BN('100').mul(TEN18);
+    var distribute1 = new BN('10').mul(TEN18);
 
-    var depositA = new BN('9999').mul(TEN18);
     await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
+    await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
 
+    tx = await rd.methods.withdrawStake(acct_a, depositA.toString()).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'StakeWithdrawalMade', {
+      to: acct_b,
+      value: depositA
+    });
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      to: acct_a,
+      value: distribute1
+    });
+
+    expect(
+      await rd.methods.getStakeTotal().call({from: acct_a})
+    ).to.be.equal('0');
+  });
+
+  it("handles magnitude: A deposits 9999, B deposits 1, distribute, withdraw stake, distribute, withdraw reward", async function() {
+    var depositA = new BN('9999').mul(TEN18);
     var depositB = new BN('1').mul(TEN18);
+    var distribute1 = new BN('1').mul(TEN18);
+    var distribute2 = new BN('2').mul(TEN18);
+
+    await rd.methods.deposit(acct_a, depositA.toString()).send({from: acct_a});
     await rd.methods.deposit(acct_b, depositB.toString()).send({from: acct_a});
 
-    var distribute1 = new BN('1').mul(TEN18);
     await rd.methods.distribute(distribute1.toString()).send({from: acct_a});
-
     await rd.methods.withdrawAllStake(acct_a).send({from: acct_a});
-
-    var distribute2 = new BN('2').mul(TEN18);
     await rd.methods.distribute(distribute2.toString()).send({from: acct_a});
 
-    var rewardB = await rd.methods.getReward(acct_b).call({from: acct_a});
-    expect(new BN(rewardB)).to.be.bignumber.equal(new BN('2000100000000000000'));
+    expect (
+      await rd.methods.getReward(acct_b).call({from: acct_a})
+    ).to.be.equal('2000100000000000000');
+  });
+
+  it("handles magnitude: deposit 10**6 10**9 10**12 10**15, distribute, withdraw reward", async function() {
+    await rd.methods.deposit(acct_a, String(10 ** 6)).send({from: acct_a});
+    await rd.methods.deposit(acct_b, String(10 ** 9)).send({from: acct_a});
+    await rd.methods.deposit(acct_c, String(10 ** 12)).send({from: acct_a});
+    // 10**6 is NOT substracted so B + C + D stakes sum up to 10 ** 15
+    await rd.methods.deposit(acct_d, String(10 ** 15 - 10 ** 12 - 10 ** 9)).send({from: acct_a});
+
+    await rd.methods.distribute(String(10 ** 9)).send({from: acct_a});
+
+    // A gets no reward because its stake is below ELIGIBLE_UNIT
+    expect (
+      await rd.methods.getReward(acct_a).call({from: acct_a})
+    ).to.be.equal('0');
+
+    expect (
+      await rd.methods.getReward(acct_b).call({from: acct_a})
+    ).to.be.equal('1000');
+
+    expect (
+      await rd.methods.getReward(acct_c).call({from: acct_a})
+    ).to.be.equal('1000000');
+
+    expect (
+      await rd.methods.getReward(acct_d).call({from: acct_a})
+    ).to.be.equal(String(10 ** 9 - 10 ** 6 - 10 ** 3));
+  });
+
+  it("carries reminder to second distribution and withdraws reward", async function() {
+    await rd.methods.deposit(acct_a, String(10 ** 9)).send({from: acct_a});
+    await rd.methods.deposit(acct_b, String(9 * 10 ** 9)).send({from: acct_a});
+
+    // 19 wei can not be divided to 10 ELIGIBLE_UNITS; So only 10 wei
+    // will be distributed and 9 will be stored as remainder and
+    // added to the next distribution
+    await rd.methods.distribute(String(19)).send({from: acct_a});
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      to: acct_a,
+      value: 1
+    });
+
+    tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      to: acct_b,
+      value: 9
+    });
+
+    // 9 wei reminder + 1 new wei can now be divided to 10 EILIGIBLE_UNITS
+    await rd.methods.distribute(String(1)).send({from: acct_a});
+
+    tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      to: acct_a,
+      value: 1
+    });
+
+    tx = await rd.methods.withdrawReward(acct_b).send({from: acct_a});
+    expectEvent.inLogs(tx.events, 'RewardWithdrawalMade', {
+      to: acct_b,
+      value: 9
+    });
   });
+
 });

From c51be6058769bf6579b8045cac4ae389e3beda46 Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Sun, 1 Sep 2019 15:48:55 +0300
Subject: [PATCH 6/8] bring back sender address for distribute events

---
 contracts/BondingCurve/dividend/RewardsDistributor.sol      | 6 +++---
 .../BondingCurve/dividend/RewardsDistributorWrapper.sol     | 2 +-
 test/unit/rewardDistributor.spec.ts                         | 1 +
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
index 8eb3dea..fa9a0b0 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributor.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -43,7 +43,7 @@ contract RewardsDistributor {
 
 
     event DepositMade(address indexed from, uint256 value);
-    event DistributionMade(uint256 value);
+    event DistributionMade(address indexed from, uint256 value);
     event RewardWithdrawalMade(address indexed to, uint256 value);
     event StakeWithdrawalMade(address indexed to, uint256 value);
 
@@ -81,7 +81,7 @@ contract RewardsDistributor {
 
 
     /// @notice Distribute tokens pro rata to all stakers.
-    function _distribute(uint tokens) internal returns (bool success) {
+    function _distribute(address from, uint tokens) internal returns (bool success) {
         require(tokens > 0);
         require(_stakeTotal > 0);
 
@@ -97,7 +97,7 @@ contract RewardsDistributor {
         // increase total rewards per stake unit
         _rewardTotal = _rewardTotal.add(_ratio);
 
-        emit DistributionMade(tokens);
+        emit DistributionMade(from, tokens);
         return true;
     }
 
diff --git a/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol b/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
index 9efdb85..60b1b41 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributorWrapper.sol
@@ -15,7 +15,7 @@ contract RewardsDistributorWrapper is RewardsDistributor {
 
     /// @notice Distribute tokens pro rata to all stakers.
     function distribute(uint tokens) public returns (bool success) {
-        return _distribute(tokens);
+        return _distribute(address(0), tokens);
     }
 
     /// @notice Withdraw accumulated reward for the staker address.
diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
index 3f60047..bb39d6e 100644
--- a/test/unit/rewardDistributor.spec.ts
+++ b/test/unit/rewardDistributor.spec.ts
@@ -149,6 +149,7 @@ contract('RewardsDistributorWrapper', accounts => {
 
     tx = await rd.methods.distribute(amountDistribute.toString()).send({from: acct_a});
     expectEvent.inLogs(tx.events, 'DistributionMade', {
+      from: 0,
       value: amountDistribute
     });
 

From 99bbe27d496b596c0077998a603a442468fff4d0 Mon Sep 17 00:00:00 2001
From: Bogdan Batog <bogdan@batog.info>
Date: Sun, 1 Sep 2019 16:04:17 +0300
Subject: [PATCH 7/8] comments

---
 .../dividend/RewardsDistributor.sol           | 27 ++++++++++++-------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
index fa9a0b0..648e1d0 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributor.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -4,11 +4,13 @@ pragma solidity ^0.5.6;
 import "@openzeppelin/contracts-ethereum-package/contracts/math/SafeMath.sol";
 
 
-/// @title RewardsDistributor - Distribute pro rata rewards (dividends)
-/// @author Bogdan Batog (https://batog.info)
-/// @dev Distribute pro rata rewards (dividends) to token holders in O(1) time.
-///      Based on http://batog.info/papers/scalable-reward-distribution.pdf
-///      And on https://solmaz.io/2019/02/24/scalable-reward-changing/
+/**
+ * @title RewardsDistributor - Distribute pro rata rewards (dividends)
+ * @author Bogdan Batog (https://batog.info)
+ * @dev Distribute pro rata rewards (dividends) to token holders in O(1) time.
+ *      Based on [1] http://batog.info/papers/scalable-reward-distribution.pdf
+ *      And on [2] https://solmaz.io/2019/02/24/scalable-reward-changing/
+ */
 contract RewardsDistributor {
     using SafeMath for uint256;
 
@@ -17,7 +19,7 @@ contract RewardsDistributor {
     ///
     ///  Only multiple of ELIGIBLE_UNIT will be subject to reward
     ///  distribution. Any fractional part of deposit, smaller than
-    ///  ELIGIBLE_UNIT, won't receive any reward.
+    ///  ELIGIBLE_UNIT, won't receive any reward, but it will be tracked.
     ///
     ///  Recommended value 10**(decimals / 2), that is 10**9 for most ERC20.
     uint256 public constant ELIGIBLE_UNIT = 10**9;
@@ -25,20 +27,25 @@ contract RewardsDistributor {
     /// @notice Stake per address.
     mapping(address => uint256) internal _stake;
 
-    /// @notice Stake reminder per address.
+    /// @notice Stake reminder per address, smaller than ELIGIBLE_UNIT.
     mapping(address => uint256) internal _stakeReminder;
 
     /// @notice Total staked tokens. In ELIGIBLE_UNIT units.
     uint256 internal _stakeTotal;
 
-    /// @notice Total reward since the beginning of time, in units per
-    ///  ELIGIBLE_UNIT.
+    /// @notice Total accumulated reward since the beginning of time, in units
+    /// per ELIGIBLE_UNIT.
     uint256 internal _rewardTotal;
 
-    /// @notice Reminder from the last reward distribution.
+    /// @notice Reminder from the last _distribute() call, this amount was not
+    /// enough to award at least 1 wei to every staked ELIGIBLE_UNIT. At the
+    /// time of last _distribute() call _rewardRemainder < _stakeTotal.
+    /// Note that later, _stakeTotal can decrease, but _rewardRemainder will
+    /// stay unchanged until the next call to _distribute().
     uint256 internal _rewardRemainder;
 
     /// @notice Proportional rewards awarded *before* this stake was created.
+    /// See [2] for more details.
     mapping(address => int256) _rewardOffset;
 
 

From 9d1e0b2fb1e7f35fdd0af0f75a6360b1c4237b46 Mon Sep 17 00:00:00 2001
From: Ori Shimony <Ori@oris-mbp.intern>
Date: Mon, 2 Sep 2019 13:02:47 +0200
Subject: [PATCH 8/8] (spelling) reminder -> remainder

---
 .../dividend/RewardsDistributor.sol           | 22 +++++++++----------
 test/unit/rewardDistributor.spec.ts           |  4 ++--
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/contracts/BondingCurve/dividend/RewardsDistributor.sol b/contracts/BondingCurve/dividend/RewardsDistributor.sol
index 648e1d0..ece0455 100644
--- a/contracts/BondingCurve/dividend/RewardsDistributor.sol
+++ b/contracts/BondingCurve/dividend/RewardsDistributor.sol
@@ -27,8 +27,8 @@ contract RewardsDistributor {
     /// @notice Stake per address.
     mapping(address => uint256) internal _stake;
 
-    /// @notice Stake reminder per address, smaller than ELIGIBLE_UNIT.
-    mapping(address => uint256) internal _stakeReminder;
+    /// @notice Stake remainder per address, smaller than ELIGIBLE_UNIT.
+    mapping(address => uint256) internal _stakeRemainder;
 
     /// @notice Total staked tokens. In ELIGIBLE_UNIT units.
     uint256 internal _stakeTotal;
@@ -37,7 +37,7 @@ contract RewardsDistributor {
     /// per ELIGIBLE_UNIT.
     uint256 internal _rewardTotal;
 
-    /// @notice Reminder from the last _distribute() call, this amount was not
+    /// @notice Remainder from the last _distribute() call, this amount was not
     /// enough to award at least 1 wei to every staked ELIGIBLE_UNIT. At the
     /// time of last _distribute() call _rewardRemainder < _stakeTotal.
     /// Note that later, _stakeTotal can decrease, but _rewardRemainder will
@@ -66,12 +66,12 @@ contract RewardsDistributor {
     /// @notice Deposit funds into contract.
     function _deposit(address staker, uint256 tokens) internal returns (bool success) {
 
-        uint256 _tokensToAdd = tokens.add(_stakeReminder[staker]);
+        uint256 _tokensToAdd = tokens.add(_stakeRemainder[staker]);
 
         uint256 _eligibleUnitsToAdd = _tokensToAdd.div(ELIGIBLE_UNIT);
 
-        // update the new reminder for this address
-        _stakeReminder[staker] = _tokensToAdd.mod(ELIGIBLE_UNIT);
+        // update the new remainder for this address
+        _stakeRemainder[staker] = _tokensToAdd.mod(ELIGIBLE_UNIT);
 
         // set the current stake for this address
         _stake[staker] = _stake[staker].add(_eligibleUnitsToAdd);
@@ -92,13 +92,13 @@ contract RewardsDistributor {
         require(tokens > 0);
         require(_stakeTotal > 0);
 
-        // add past distribution reminder
+        // add past distribution remainder
         uint256 _amountToDistribute = tokens.add(_rewardRemainder);
 
         // determine rewards per eligible stake
         uint256 _ratio = _amountToDistribute.div(_stakeTotal);
 
-        // carry on reminder
+        // carry on remainder
         _rewardRemainder = _amountToDistribute.mod(_stakeTotal);
 
         // increase total rewards per stake unit
@@ -129,10 +129,10 @@ contract RewardsDistributor {
 
         require(tokens <= _currentStake);
 
-        // update stake and reminder for this address
+        // update stake and remainder for this address
         uint256 _newStake = _currentStake.sub(tokens);
 
-        _stakeReminder[staker] = _newStake.mod(ELIGIBLE_UNIT);
+        _stakeRemainder[staker] = _newStake.mod(ELIGIBLE_UNIT);
 
         uint256 _eligibleUnitsDelta = _stake[staker].sub(
             _newStake.div(ELIGIBLE_UNIT)
@@ -166,7 +166,7 @@ contract RewardsDistributor {
         tokens = (
             _stake[staker].mul(ELIGIBLE_UNIT)
         ).add(
-            _stakeReminder[staker]
+            _stakeRemainder[staker]
         );
 
         return tokens;
diff --git a/test/unit/rewardDistributor.spec.ts b/test/unit/rewardDistributor.spec.ts
index bb39d6e..5a01b88 100644
--- a/test/unit/rewardDistributor.spec.ts
+++ b/test/unit/rewardDistributor.spec.ts
@@ -356,7 +356,7 @@ contract('RewardsDistributorWrapper', accounts => {
     ).to.be.equal(String(10 ** 9 - 10 ** 6 - 10 ** 3));
   });
 
-  it("carries reminder to second distribution and withdraws reward", async function() {
+  it("carries remainder to second distribution and withdraws reward", async function() {
     await rd.methods.deposit(acct_a, String(10 ** 9)).send({from: acct_a});
     await rd.methods.deposit(acct_b, String(9 * 10 ** 9)).send({from: acct_a});
 
@@ -377,7 +377,7 @@ contract('RewardsDistributorWrapper', accounts => {
       value: 9
     });
 
-    // 9 wei reminder + 1 new wei can now be divided to 10 EILIGIBLE_UNITS
+    // 9 wei remainder + 1 new wei can now be divided to 10 EILIGIBLE_UNITS
     await rd.methods.distribute(String(1)).send({from: acct_a});
 
     tx = await rd.methods.withdrawReward(acct_a).send({from: acct_a});