ATokenWithDelegation.diff

diff --git a/docs/3.4/ATokenWithDelegation.sol b/src/contracts/protocol/tokenization/ATokenWithDelegation.sol
index b3967209..0d3d4b7b 100644
--- a/docs/3.4/ATokenWithDelegation.sol
+++ b/src/contracts/protocol/tokenization/ATokenWithDelegation.sol
@@ -1,17 +1,24 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.8.0;
+// SPDX-License-Identifier: BUSL-1.1
+pragma solidity ^0.8.10;
+
+import {WadRayMath} from '../libraries/math/WadRayMath.sol';
+import {IPool} from '../../interfaces/IPool.sol';

-import {IPool} from 'aave-v3-core/contracts/interfaces/IPool.sol';
-import {BaseDelegation} from 'aave-token-v3/BaseDelegation.sol';
 import {AToken} from './AToken.sol';

+import {DelegationMode} from './base/DelegationMode.sol';
+import {BaseDelegation} from './delegation/BaseDelegation.sol';
+
 /**
+ * @title Aave ERC20 ATokenWithDelegation
  * @author BGD Labs
  * @notice contract that gives a tokens the delegation functionality. For now should only be used for AAVE aToken
  * @dev uint sizes are used taken into account that is tailored for AAVE token. In this AToken child we only update
         delegation balances. Balances amount is taken care of by AToken contract
  */
-contract ATokenWithDelegation is AToken, BaseDelegation {
+abstract contract ATokenWithDelegation is AToken, BaseDelegation {
+  using WadRayMath for uint256;
+
   struct ATokenDelegationState {
     uint72 delegatedPropositionBalance;
     uint72 delegatedVotingBalance;
@@ -19,7 +26,17 @@ contract ATokenWithDelegation is AToken, BaseDelegation {

   mapping(address => ATokenDelegationState) internal _delegatedState;

-  constructor(IPool pool) AToken(pool) {}
+  /**
+   * @dev Constructor.
+   * @param pool The address of the Pool contract
+   * @param rewardsController The address of the rewards controller contract
+   * @param treasury The address of the treasury. This is where accrued interest is sent.
+   */
+  constructor(
+    IPool pool,
+    address rewardsController,
+    address treasury
+  ) AToken(pool, rewardsController, treasury) {}

   function _getDomainSeparator() internal view override returns (bytes32) {
     return DOMAIN_SEPARATOR();
@@ -57,20 +74,36 @@ contract ATokenWithDelegation is AToken, BaseDelegation {
   }

   /**
-   * @notice Overrides the parent _transfer to force validated transfer() and delegation balance transfers
-   * @param from The source address
-   * @param to The destination address
-   * @param amount The amount getting transferred
+   * @notice Transfers tokens and updates delegation balances.  This function overrides the parent `_transfer`
+   *         to include delegation logic. It first updates the delegation balances based on the transfer
+   *         and then calls the parent's `_transfer` function to perform the actual token transfer.
+   * @dev The amount is divided by the index inside this function to perform the scaling.
+   * @param from The sender's address.
+   * @param to The recipient's address.
+   * @param amount The amount of tokens to transfer (non-scaled).
+   * @param index The current liquidity index of the reserve.
    */
-  function _transfer(address from, address to, uint256 amount, bool validate) internal override {
-    _delegationChangeOnTransfer(from, to, _getBalance(from), _getBalance(to), amount);
-    super._transfer(from, to, amount, validate);
+  function _transfer(
+    address from,
+    address to,
+    uint256 amount,
+    uint256 index
+  ) internal virtual override {
+    _delegationChangeOnTransfer({
+      from: from,
+      to: to,
+      fromBalanceBefore: _userState[from].balance,
+      toBalanceBefore: _userState[to].balance,
+      amount: uint256(amount).rayDiv(index)
+    });
+
+    super._transfer(from, to, amount, index);
   }

   /**
    * @notice Overrides the parent _mint to force delegation balance transfers
    * @param account The address receiving tokens
-   * @param amount The amount of tokens to mint
+   * @param amount The amount of tokens to mint (scaled)
    */
   function _mint(address account, uint120 amount) internal override {
     _delegationChangeOnTransfer(address(0), account, 0, _getBalance(account), amount);
@@ -80,7 +113,7 @@ contract ATokenWithDelegation is AToken, BaseDelegation {
   /**
    * @notice Overrides the parent _burn to force delegation balance transfers
    * @param account The account whose tokens are burnt
-   * @param amount The amount of tokens to burn
+   * @param amount The amount of tokens to burn (scaled)
    */
   function _burn(address account, uint120 amount) internal override {
     _delegationChangeOnTransfer(account, address(0), _getBalance(account), 0, amount);