Apply needed fixes and improvements

contracts/APRCalculator/APRCalculator.sol

@@ -99,58 +99,60 @@ contract APRCalculator is IAPRCalculator, MacroFactor, RSIndex {
      * @notice Initializes vesting bonus for each week.
      */
     function _initializeVestingBonus() private {
-        vestingBonus[0] = 6;
-        vestingBonus[1] = 16;
-        vestingBonus[2] = 30;
-        vestingBonus[3] = 46;
-        vestingBonus[4] = 65;
-        vestingBonus[5] = 85;
-        vestingBonus[6] = 108;
-        vestingBonus[7] = 131;
-        vestingBonus[8] = 157;
-        vestingBonus[9] = 184;
-        vestingBonus[10] = 212;
-        vestingBonus[11] = 241;
-        vestingBonus[12] = 272;
-        vestingBonus[13] = 304;
-        vestingBonus[14] = 338;
-        vestingBonus[15] = 372;
-        vestingBonus[16] = 407;
-        vestingBonus[17] = 444;
-        vestingBonus[18] = 481;
-        vestingBonus[19] = 520;
-        vestingBonus[20] = 559;
-        vestingBonus[21] = 599;
-        vestingBonus[22] = 641;
-        vestingBonus[23] = 683;
-        vestingBonus[24] = 726;
-        vestingBonus[25] = 770;
-        vestingBonus[26] = 815;
-        vestingBonus[27] = 861;
-        vestingBonus[28] = 907;
-        vestingBonus[29] = 955;
-        vestingBonus[30] = 1003;
-        vestingBonus[31] = 1052;
-        vestingBonus[32] = 1101;
-        vestingBonus[33] = 1152;
-        vestingBonus[34] = 1203;
-        vestingBonus[35] = 1255;
-        vestingBonus[36] = 1307;
-        vestingBonus[37] = 1361;
-        vestingBonus[38] = 1415;
-        vestingBonus[39] = 1470;
-        vestingBonus[40] = 1525;
-        vestingBonus[41] = 1581;
-        vestingBonus[42] = 1638;
-        vestingBonus[43] = 1696;
-        vestingBonus[44] = 1754;
-        vestingBonus[45] = 1812;
-        vestingBonus[46] = 1872;
-        vestingBonus[47] = 1932;
-        vestingBonus[48] = 1993;
-        vestingBonus[49] = 2054;
-        vestingBonus[50] = 2116;
-        vestingBonus[51] = 2178;
+        vestingBonus = [
+            6,
+            16,
+            30,
+            46,
+            65,
+            85,
+            108,
+            131,
+            157,
+            184,
+            212,
+            241,
+            272,
+            304,
+            338,
+            372,
+            407,
+            444,
+            481,
+            520,
+            559,
+            599,
+            641,
+            683,
+            726,
+            770,
+            815,
+            861,
+            907,
+            955,
+            1003,
+            1052,
+            1101,
+            1152,
+            1203,
+            1255,
+            1307,
+            1361,
+            1415,
+            1470,
+            1525,
+            1581,
+            1638,
+            1696,
+            1754,
+            1812,
+            1872,
+            1932,
+            1993,
+            2054,
+            2116,
+            2178
+        ];
     }
 
     // slither-disable-next-line unused-state,naming-convention

contracts/HydraChain/HydraChain.sol

@@ -13,14 +13,17 @@ import {IHydraChain} from "./IHydraChain.sol";
 import {Epoch} from "./IHydraChain.sol";
 
 contract HydraChain is IHydraChain, Inspector, DaoIncentive, ValidatorsData {
+    // @note unsafe arrays are used in the code
     using ArraysUpgradeable for uint256[];
 
     uint256 public currentEpochId;
+    // @note this is unsafe array
     /// @notice Array with epoch ending blocks
     uint256[] public epochEndBlocks;
     /// @notice Epoch data linked with the epoch id
     mapping(uint256 => Epoch) public epochs;
 
+    // @audit isn't this the same as the epochEndBlocks?
     mapping(uint256 => uint256) internal _commitBlockNumbers;
 
     // _______________ Initializer _______________
@@ -68,6 +71,7 @@ contract HydraChain is IHydraChain, Inspector, DaoIncentive, ValidatorsData {
         if (epoch.startBlock >= epoch.endBlock) {
             revert CommitEpochFailed("NO_BLOCKS_COMMITTED");
         }
+        // @audit this is going to be a problem when slashing is added
         if ((epoch.endBlock - epoch.startBlock + 1) % epochSize != 0) {
             revert CommitEpochFailed("EPOCH_MUST_BE_DIVISIBLE_BY_EPOCH_SIZE");
         }
@@ -145,8 +149,10 @@ contract HydraChain is IHydraChain, Inspector, DaoIncentive, ValidatorsData {
         uint256 validatorParticipation = validatorsParticipation[validator];
         // check if the validator is active and the last participation is less than the threshold
         if (
+            // @audit here we check if the validator is active but this doesn't mean it is actually part of the active set (check if this is breakable)
             validators[validator].status == ValidatorStatus.Active &&
             lastCommittedEndBlock > validatorParticipation &&
+            // @audit at block 1 what is the value of validatorParticipation? lastCommittedEndBlock is 0 in the first 500 blocks
             lastCommittedEndBlock - validatorParticipation >= initiateBanThreshold
         ) {
             return true;

contracts/HydraChain/modules/DaoIncentive/DaoIncentive.sol

@@ -39,6 +39,7 @@ abstract contract DaoIncentive is IDaoIncentive, System, Initializable, RewardWa
      * @inheritdoc IDaoIncentive
      */
     function distributeDAOIncentive() external onlySystemCall {
+        // @audit put all multiplications before division to avoid rounding errors
         uint256 reward = (((hydraStakingContract.totalBalance() * 200) / 10000) *
             (block.timestamp - lastDistribution)) / 365 days;
         lastDistribution = block.timestamp;
@@ -51,6 +52,7 @@ abstract contract DaoIncentive is IDaoIncentive, System, Initializable, RewardWa
      * @inheritdoc IDaoIncentive
      */
     function claimVaultFunds() external {
+        // @audit reward can be renamed to incentive or something
         uint256 reward = vaultDistribution;
         if (reward == 0) {
             revert NoVaultFundsToClaim();

contracts/HydraChain/modules/Inspector/Inspector.sol

@@ -12,9 +12,13 @@ abstract contract Inspector is IInspector, ValidatorManager {
     uint256 public validatorPenalty;
     /// @notice The reward for the person who reports a validator that have to be banned
     uint256 public reporterReward;
-    /// @notice Validator inactiveness (in blocks) threshold that needs to be passed to initiate ban for a validator
+    /// @notice Validator inactiveness (in blocks) threshold
+    /// that needs to be reached or passed to initiate ban for a validator
+    /// @dev must be always bigger than the epoch length (better bigger than at least 4 epochs),
+    /// otherwise all validators can be banned
     uint256 public initiateBanThreshold;
-    /// @notice Validator inactiveness (in seconds) threshold that needs to be passed to ban a validator
+    /// @notice Validator inactiveness (in seconds) threshold.
+    /// A validator can be banned if in "ban initiated" state for a duration equal to or exceeding this threshold.
     uint256 public banThreshold;
     /// @notice Mapping of the validators that bans has been initiated for (validator => timestamp)
     mapping(address => uint256) public bansInitiated;

contracts/HydraChain/modules/ValidatorManager/ValidatorManager.sol

@@ -36,6 +36,7 @@ abstract contract ValidatorManager is
     mapping(address => Validator) public validators;
     /**
      * @notice Mapping that keeps the last time when a validator has participated in the consensus
+     * @dev Updated on epoch-ending blocks only
      * @dev Keep in mind that the validator will initially be set active when stake,
      * but it will be able to participate in the next epoch. So, the validator will have
      * less blocks to participate before getting eligible for ban.
@@ -72,22 +73,6 @@ abstract contract ValidatorManager is
         }
     }
 
-    // _______________ Modifiers _______________
-
-    modifier onlyActiveValidator(address validator) {
-        if (validators[validator].status != ValidatorStatus.Active) revert Unauthorized("INACTIVE_VALIDATOR");
-        _;
-    }
-
-    /// @notice Modifier to check if the validator is registered or active
-    modifier onlyValidator(address validator) {
-        if (
-            validators[validator].status != ValidatorStatus.Registered &&
-            validators[validator].status != ValidatorStatus.Active
-        ) revert Unauthorized("INVALID_VALIDATOR");
-        _;
-    }
-
     // _______________ External functions _______________
 
     /**

contracts/HydraDelegation/Delegation.sol

@@ -216,12 +216,12 @@ contract Delegation is
     }
 
     /**
-     * @notice Base delegation funds to a staker
+     * @notice Core logic for delegating funds to a staker
      * @param staker Address of the validator
      * @param delegator Address of the delegator
      * @param amount Amount to delegate
      */
-    function _baseDelegate(address staker, address delegator, uint256 amount) internal virtual {
+    function _baseDelegate(address staker, address delegator, uint256 amount) internal {
         if (amount == 0) revert DelegateRequirement({src: "delegate", msg: "DELEGATING_AMOUNT_ZERO"});
         DelegationPool storage delegation = delegationPools[staker];
         uint256 delegatedAmount = delegation.balanceOf(delegator);
@@ -253,6 +253,7 @@ contract Delegation is
 
     /**
      * @notice Undelegates funds from a staker
+     * @dev overriden in child contracts to extend core undelegate behaviour
      * @param staker Address of the validator
      * @param delegator Address of the delegator
      * @param amount Amount to delegate
@@ -262,12 +263,12 @@ contract Delegation is
     }
 
     /**
-     * @notice Base undelegating funds from a staker
+     * @notice Core logic for undelegating funds from a staker
      * @param staker Address of the validator
      * @param delegator Address of the delegator
      * @param amount Amount to delegate
      */
-    function _baseUndelegate(address staker, address delegator, uint256 amount) internal virtual {
+    function _baseUndelegate(address staker, address delegator, uint256 amount) internal {
         DelegationPool storage delegation = delegationPools[staker];
         uint256 delegatedAmount = delegation.balanceOf(delegator);
         if (amount > delegatedAmount) revert DelegateRequirement({src: "undelegate", msg: "INSUFFICIENT_BALANCE"});
@@ -290,6 +291,7 @@ contract Delegation is
 
     /**
      * @notice Withdraws funds from stakers pool
+     * @dev overriden in child contracts to extend core withdraw behaviour
      * @param delegation Delegation pool
      * @param delegator Address of the delegator
      * @param amount Amount to withdraw

contracts/HydraDelegation/HydraDelegation.sol

@@ -39,7 +39,6 @@ contract HydraDelegation is IHydraDelegation, System, Delegation, LiquidDelegati
         __Liquid_init(liquidToken);
         __Vesting_init_unchained();
         __VestingManagerFactoryConnector_init(vestingManagerFactoryAddr);
-        __VestedDelegation_init_unchained();
     }
 
     // _______________ External functions _______________
@@ -60,7 +59,7 @@ contract HydraDelegation is IHydraDelegation, System, Delegation, LiquidDelegati
         address staker,
         address delegator,
         uint256 amount
-    ) internal virtual override(Delegation, LiquidDelegation, VestedDelegation) {
+    ) internal virtual override(Delegation, LiquidDelegation) {
         super._delegate(staker, delegator, amount);
     }
 
@@ -71,7 +70,7 @@ contract HydraDelegation is IHydraDelegation, System, Delegation, LiquidDelegati
         address staker,
         address delegator,
         uint256 amount
-    ) internal virtual override(Delegation, LiquidDelegation, VestedDelegation) {
+    ) internal virtual override(Delegation, LiquidDelegation) {
         super._undelegate(staker, delegator, amount);
     }
 
@@ -83,7 +82,18 @@ contract HydraDelegation is IHydraDelegation, System, Delegation, LiquidDelegati
      */
     function _distributeTokens(address staker, address account, uint256 amount) internal virtual override {
         VestingPosition memory position = vestedDelegationPositions[staker][msg.sender];
+        // This check works because if position has already been opened, the restrictions on delegateWithVesting()
+        // will prevent entering this check again
         if (_isOpeningPosition(position)) {
+            uint256 previousDelegation = delegationOf(staker, account) - amount;
+            if (previousDelegation != 0) {
+                // We want all previously distributed tokens to be collected,
+                // because the new position vesting period can be different from the previous one
+                // meaning the tokens to distribute amount will be different
+                _collectTokens(staker, previousDelegation);
+                amount += previousDelegation;
+            }
+
             uint256 debt = _calculatePositionDebt(amount, position.duration);
             liquidityDebts[account] -= debt.toInt256Safe(); // Add negative debt
             amount -= debt;

contracts/HydraDelegation/modules/DelegationPoolLib/DelegationPoolLib.sol

@@ -245,7 +245,6 @@ library DelegationPoolLib {
      * @param currentEpochNum the current epoch number
      * @return bool whether the balance change is made
      */
-    // TODO: Check if the commitEpoch is the last transaction in the epoch, otherwise bug may occur
     function isBalanceChangeMade(
         DelegationPool storage pool,
         address delegator,
@@ -277,7 +276,8 @@ library DelegationPoolLib {
     // _______________ Private functions _______________
 
     /**
-     * @notice Saves the RPS for the given staker for the epoch
+     * @notice Saves the RPS for the given staker's delegation pool
+     * @dev must be called when new reward is distributed (at the end of every epoch in our case)
      * @param pool the DelegationPool to save the RPS for
      * @param rewardPerShare Amount of tokens to be withdrawn
      * @param epochNumber Epoch number