Cleanup OUSD rebasing/nonrebasing accounting changes

contracts/contracts/token/OUSD.sol

@@ -457,7 +457,7 @@ contract OUSD is Initializable, InitializableERC20Detailed, Governable {
     function _isNonRebasingAccount(address _account) internal returns (bool) {
         bool isContract = Address.isContract(_account);
         if (isContract && rebaseState[_account] == RebaseOptions.NotSet) {
-            _ensureRebasingMigration(_account);
+            _ensureMigrationToNonRebasing(_account);
         }
         return nonRebasingCreditsPerToken[_account] > 0;
     }
@@ -466,25 +466,34 @@ contract OUSD is Initializable, InitializableERC20Detailed, Governable {
      * @dev Ensures internal account for rebasing and non-rebasing credits and
      *      supply is updated following deployment of frozen yield change.
      */
-    function _ensureRebasingMigration(address _account) internal {
-        if (nonRebasingCreditsPerToken[_account] == 0) {
-            if (_creditBalances[_account] == 0) {
-                // Since there is no existing balance, we can directly set to
-                // high resolution, and do not have to do any other bookkeeping
-                nonRebasingCreditsPerToken[_account] = 1e27;
-            } else {
-                // Migrate an existing account:
-
-                // Set fixed credits per token for this account
-                nonRebasingCreditsPerToken[_account] = _rebasingCreditsPerToken;
-                // Update non rebasing supply
-                nonRebasingSupply = nonRebasingSupply.add(balanceOf(_account));
-                // Update credit tallies
-                _rebasingCredits = _rebasingCredits.sub(
-                    _creditBalances[_account]
-                );
-            }
+    function _ensureMigrationToNonRebasing(address _account) internal {
+        if (nonRebasingCreditsPerToken[_account] != 0) {
+            return; // Account already is non-rebasing
+        }
+        if (_creditBalances[_account] == 0) {
+            // Since there is no existing balance, we can directly set to
+            // high resolution, and do not have to do any other bookkeeping
+            nonRebasingCreditsPerToken[_account] = 1e27;
+            return;
         }
+
+        // This does not change, but if it did, we would want to
+        // use the value before changes.
+        uint256 oldCredits = _creditBalances[_account];
+
+        // Atomicly update account information:
+        // It is important that balanceOf not be called inside updating
+        // account data, since it will give wrong answers if it does
+        // not have all an account's data in a consistent state.
+        //
+        // By setting a per account nonRebasingCreditsPerToken,
+        // this account will no longer follow with the global
+        // rebasing credits per token and will become non-rebasing.
+        nonRebasingCreditsPerToken[_account] = _rebasingCreditsPerToken;
+
+        // Update global totals
+        nonRebasingSupply = nonRebasingSupply.add(balanceOf(_account));
+        _rebasingCredits = _rebasingCredits.sub(oldCredits);
     }
 
     /**
@@ -495,24 +504,30 @@ contract OUSD is Initializable, InitializableERC20Detailed, Governable {
     function rebaseOptIn() public nonReentrant {
         require(_isNonRebasingAccount(msg.sender), "Account has not opted out");
 
+        uint256 oldBalance = balanceOf(msg.sender);
+
+        // Precalculate new credits, so that we avoid internal calls when
+        // atomicly updating account.
         // Convert balance into the same amount at the current exchange rate
         uint256 newCreditBalance = _creditBalances[msg.sender]
             .mul(_rebasingCreditsPerToken)
             .div(_creditsPerToken(msg.sender));
 
-        // Decreasing non rebasing supply
-        nonRebasingSupply = nonRebasingSupply.sub(balanceOf(msg.sender));
-
+        // Atomicly update this account:
+        // Important that no internal calls happen during this.
+        // Remove pinned fixed credits per token
+        delete nonRebasingCreditsPerToken[msg.sender];
+        // New credits
         _creditBalances[msg.sender] = newCreditBalance;
+        // Mark explicitly opted out of rebasing
+        rebaseState[msg.sender] = RebaseOptions.OptIn;
 
+        // Update global totals:
+        // Decrease non rebasing supply
+        nonRebasingSupply = nonRebasingSupply.sub(oldBalance);
         // Increase rebasing credits, totalSupply remains unchanged so no
         // adjustment necessary
         _rebasingCredits = _rebasingCredits.add(_creditBalances[msg.sender]);
-
-        rebaseState[msg.sender] = RebaseOptions.OptIn;
-
-        // Delete any fixed credits per token
-        delete nonRebasingCreditsPerToken[msg.sender];
     }
 
     /**
@@ -521,16 +536,8 @@ contract OUSD is Initializable, InitializableERC20Detailed, Governable {
     function rebaseOptOut() public nonReentrant {
         require(!_isNonRebasingAccount(msg.sender), "Account has not opted in");
 
-        // Increase non rebasing supply
-        nonRebasingSupply = nonRebasingSupply.add(balanceOf(msg.sender));
-        // Set fixed credits per token
-        nonRebasingCreditsPerToken[msg.sender] = _rebasingCreditsPerToken;
-
-        // Decrease rebasing credits, total supply remains unchanged so no
-        // adjustment necessary
-        _rebasingCredits = _rebasingCredits.sub(_creditBalances[msg.sender]);
+        _ensureMigrationToNonRebasing(msg.sender);
 
-        // Mark explicitly opted out of rebasing
         rebaseState[msg.sender] = RebaseOptions.OptOut;
     }