From 7bf96cd87e783e7a7514ad562658bec28cedaf67 Mon Sep 17 00:00:00 2001
From: harpreet singh <happybiostockcode07040@gmail.com>
Date: Mon, 7 Oct 2024 17:41:50 +0530
Subject: [PATCH 1/7] Add swap hooks contract

---
 packages/foundry/contracts/hooks/SafeSwap.sol | 160 ++++++++++++++++++
 packages/foundry/script/SafeSwap.s.sol        |  26 +++
 2 files changed, 186 insertions(+)
 create mode 100644 packages/foundry/contracts/hooks/SafeSwap.sol
 create mode 100644 packages/foundry/script/SafeSwap.s.sol

diff --git a/packages/foundry/contracts/hooks/SafeSwap.sol b/packages/foundry/contracts/hooks/SafeSwap.sol
new file mode 100644
index 00000000..7a4a63d4
--- /dev/null
+++ b/packages/foundry/contracts/hooks/SafeSwap.sol
@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.24;
+// import {BalancerPoolToken} from "../lib/balancer-v3-monorepo/pkg/vault/contracts/BalancerPoolToken.sol";
+import { BalancerPoolToken } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BalancerPoolToken.sol";
+import { IPoolLiquidity } from "@balancer-labs/balancer-v3-monorepo/pkg/interfaces/contracts/vault/IPoolLiquidity.sol";
+import { FixedPoint } from "@balancer-labs/balancer-v3-monorepo/pkg/solidity-utils/contracts/math/FixedPoint.sol";
+import { Math } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BasePoolMath.sol";
+
+contract SafeSwap is IBasePool, IPoolLiquidity, BalancerPoolToken {
+    using FixedPoint for uint256;
+
+    uint256 private constant _MIN_SWAP_FEE_PERCENTAGE = 0;
+    uint256 private constant _MAX_SWAP_FEE_PERCENTAGE = 0.1e18; // 10%
+
+    constructor(IVault vault, string memory name, string memory symbol) BalancerPoolToken(vault, name, symbol) {}
+
+    /**
+     * @notice Execute a swap in the pool.
+     * @param params Swap parameters
+     * @return amountCalculatedScaled18 Calculated amount for the swap
+     */
+    function onSwap(PoolSwapParams calldata params) external pure returns (uint256 amountCalculatedScaled18) {
+        amountCalculatedScaled18 =
+            (params.balancesScaled18[params.indexOut] * params.amountGivenScaled18) /
+            (params.balancesScaled18[params.indexIn] + params.amountGivenScaled18);
+    }
+
+    /**
+     * @notice Computes and returns the pool's invariant.
+     * @dev This function computes the invariant based on current balances
+     * @param balancesLiveScaled18 Array of current pool balances for each token in the pool, scaled to 18 decimals
+     * @return invariant The calculated invariant of the pool, represented as a uint256
+     */
+    function computeInvariant(uint256[] memory balancesLiveScaled18) public pure returns (uint256 invariant) {
+        // expected to work with 2 tokens only
+        invariant = FixedPoint.ONE;
+        for (uint256 i = 0; i < balancesLiveScaled18.length; ++i) {
+            invariant = invariant.mulDown(balancesLiveScaled18[i]);
+        }
+        // scale the invariant to 1e18
+        invariant = Math.sqrt(invariant) * 1e9;
+    }
+
+    /**
+     * @dev Computes the new balance of a token after an operation, given the invariant growth ratio and all other balances.
+     * @param balancesLiveScaled18 Current live balances (adjusted for decimals, rates, etc.)
+     * @param tokenInIndex The index of the token we're computing the balance for, in token registration order
+     * @param invariantRatio The ratio of the new invariant (after an operation) to the old
+     * @return newBalance The new balance of the selected token, after the operation
+     */
+    function computeBalance(
+        uint256[] memory balancesLiveScaled18,
+        uint256 tokenInIndex,
+        uint256 invariantRatio
+    ) external pure returns (uint256 newBalance) {
+        uint256 otherTokenIndex = tokenInIndex == 0 ? 1 : 0;
+
+        uint256 newInvariant = computeInvariant(balancesLiveScaled18).mulDown(invariantRatio);
+
+        newBalance = (newInvariant * newInvariant) / balancesLiveScaled18[otherTokenIndex];
+    }
+
+    /// @return minimumSwapFeePercentage The minimum swap fee percentage for a pool
+    function getMinimumSwapFeePercentage() external pure returns (uint256) {
+        return _MIN_SWAP_FEE_PERCENTAGE;
+    }
+
+    /// @return maximumSwapFeePercentage The maximum swap fee percentage for a pool
+    function getMaximumSwapFeePercentage() external pure returns (uint256) {
+        return _MAX_SWAP_FEE_PERCENTAGE;
+    }
+
+    /**
+     * @notice Custom add liquidity hook.
+     * @param router The address that initiated the operation
+     * @param maxAmountsInScaled18 Maximum input amounts, sorted in token registration order
+     * @param minBptAmountOut Minimum amount of output pool tokens
+     * @param balancesScaled18 Current pool balances, sorted in token registration order
+     * @param userData Arbitrary data sent with the request
+     * @return amountsInScaled18 Input token amounts, sorted in token registration order
+     * @return bptAmountOut Calculated pool token amount to receive
+     * @return swapFeeAmountsScaled18 The amount of swap fees charged for each token
+     * @return returnData Arbitrary data with an encoded response from the pool
+     */
+    function onAddLiquidityCustom(
+        address router,
+        uint256[] memory maxAmountsInScaled18,
+        uint256 minBptAmountOut,
+        uint256[] memory balancesScaled18,
+        bytes memory userData
+    )
+        external
+        override
+        returns (
+            uint256[] memory amountsInScaled18,
+            uint256 bptAmountOut,
+            uint256[] memory swapFeeAmountsScaled18,
+            bytes memory returnData
+        )
+    {
+        // Custom logic for adding liquidity
+        uint256 invariantBefore = computeInvariant(balancesScaled18);
+        amountsInScaled18 = maxAmountsInScaled18; // You can modify this based on custom liquidity logic
+        swapFeeAmountsScaled18 = new uint256[](balancesScaled18.length); // Placeholder for swap fees
+
+        // Update balances after adding liquidity
+        for (uint256 i = 0; i < balancesScaled18.length; ++i) {
+            balancesScaled18[i] += amountsInScaled18[i];
+        }
+
+        uint256 invariantAfter = computeInvariant(balancesScaled18);
+        bptAmountOut = invariantAfter - invariantBefore; // Example calculation
+
+        returnData = userData; // Custom return data
+    }
+
+    /**
+     * @notice Custom remove liquidity hook.
+     * @param router The address that initiated the operation
+     * @param maxBptAmountIn Maximum amount of input pool tokens
+     * @param minAmountsOutScaled18 Minimum output amounts, sorted in token registration order
+     * @param balancesScaled18 Current pool balances, sorted in token registration order
+     * @param userData Arbitrary data sent with the request
+     * @return bptAmountIn Calculated pool token amount to burn
+     * @return amountsOutScaled18 Amount of tokens to receive, sorted in token registration order
+     * @return swapFeeAmountsScaled18 The amount of swap fees charged for each token
+     * @return returnData Arbitrary data with an encoded response from the pool
+     */
+    function onRemoveLiquidityCustom(
+        address router,
+        uint256 maxBptAmountIn,
+        uint256[] memory minAmountsOutScaled18,
+        uint256[] memory balancesScaled18,
+        bytes memory userData
+    )
+        external
+        override
+        returns (
+            uint256 bptAmountIn,
+            uint256[] memory amountsOutScaled18,
+            uint256[] memory swapFeeAmountsScaled18,
+            bytes memory returnData
+        )
+    {
+        // Custom logic for removing liquidity
+        uint256 invariantBefore = computeInvariant(balancesScaled18);
+        amountsOutScaled18 = minAmountsOutScaled18; // Modify this based on custom logic
+        swapFeeAmountsScaled18 = new uint256[](balancesScaled18.length); // Placeholder for swap fees
+
+        // Update balances after removing liquidity
+        for (uint256 i = 0; i < balancesScaled18.length; ++i) {
+            balancesScaled18[i] -= amountsOutScaled18[i];
+        }
+
+        uint256 invariantAfter = computeInvariant(balancesScaled18);
+        bptAmountIn = invariantBefore - invariantAfter; // Example calculation
+
+        returnData = userData; // Custom return data
+    }
+}
diff --git a/packages/foundry/script/SafeSwap.s.sol b/packages/foundry/script/SafeSwap.s.sol
new file mode 100644
index 00000000..76e6fd15
--- /dev/null
+++ b/packages/foundry/script/SafeSwap.s.sol
@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import { Script, console } from "../lib/forge-std/src/Script.sol";
+import { ConstantProductPool } from "../src/ConstantProductPool.sol";
+import { IVault } from "../balancer-v3-monorepo/pkg/interfaces/contracts/vault/IVault.sol";
+
+contract SafeSwap is Script {
+    ConstantProductPool public constantProductPool;
+
+    address public vaultAddress = "0x7966FE92C59295EcE7FB5D9EfDB271967BFe2fbA";
+    string public poolName = "MyConstantProductPool";
+    string public poolSymbol = "MCP";
+
+    function setUp() public {}
+
+    function run() public {
+        vm.startBroadcast();
+
+        constantProductPool = new ConstantProductPool(IVault(vaultAddress), poolName, poolSymbol);
+
+        vm.stopBroadcast();
+    }
+}
+
+//forge script script/ConstantProductPool.s.sol --rpc-url https://sepolia.infura.io/v3/2de477c3b1b74816ae5475da6d289208 --private-key f46e7f0936b479bba879c9f764259d1e5838aa015232f0018a1c07214e491812

From 83c32290a1f4bf47d0a4eb2296f95bee682222b1 Mon Sep 17 00:00:00 2001
From: harpreet singh <happybiostockcode07040@gmail.com>
Date: Mon, 7 Oct 2024 17:59:31 +0530
Subject: [PATCH 2/7] update README

---
 README.md | 402 ++++++++++++++----------------------------------------
 1 file changed, 105 insertions(+), 297 deletions(-)

diff --git a/README.md b/README.md
index 1397b2da..de3174fd 100644
--- a/README.md
+++ b/README.md
@@ -1,301 +1,109 @@
-# 🏗︎ Scaffold Balancer v3
 
-A starter kit for building on top of Balancer v3. Accelerate the process of creating custom pools and hooks contracts. Concentrate on mastering the core concepts within a swift and responsive environment augmented by a local fork and a frontend pool operations playground.
+# SafeSwap Contract
+
+The **SafeSwap** contract is an implementation based on the **Balancer V3 protocol**, designed to manage token swaps, liquidity additions, and liquidity removals in a decentralized way. This contract introduces custom hooks for liquidity operations and operates within the Balancer V3 framework, leveraging mathematical invariants and swap fees.
+
+## Overview
+
+The contract extends the core components of Balancer V3:
+- **BalancerPoolToken**: This provides the functionality of an ERC-20 compliant pool token, which represents the user's share in the liquidity pool.
+- **IPoolLiquidity**: This is an interface that manages the core pool liquidity operations.
+- **FixedPoint**: A math library that allows for accurate fixed-point arithmetic.
+
+### Key Constants
+
+- `_MIN_SWAP_FEE_PERCENTAGE`: Minimum swap fee percentage, set to `0`.
+- `_MAX_SWAP_FEE_PERCENTAGE`: Maximum swap fee percentage, capped at `10%` (0.1e18).
+
+## Constructor
+
+The constructor initializes the `BalancerPoolToken` with the provided vault, pool name, and symbol.
+
+### Parameters:
+- **vault**: The Balancer vault where the pool is registered.
+- **name**: Name of the pool token.
+- **symbol**: Symbol of the pool token.
+
+---
+
+## Core Functions
+
+### 1. **`onSwap`** 
+   - **Purpose**: Executes a swap in the pool by calculating the output token amount.
+   - **Parameters**: 
+     - `PoolSwapParams calldata params`: Contains balances and amount details for the swap.
+   - **Returns**: `amountCalculatedScaled18`: The calculated amount of the output token.
+   - **Logic**: 
+     - It uses a formula to compute the new balance of the output token based on the input token and respective balances.
+     - Formula: 
+       ```
+       amountCalculatedScaled18 = (params.balancesScaled18[params.indexOut] * params.amountGivenScaled18) 
+                                  / (params.balancesScaled18[params.indexIn] + params.amountGivenScaled18)
+       ```
+
+### 2. **`computeInvariant`** 
+   - **Purpose**: Computes the pool’s invariant, which is a mathematical property representing the pool’s balance state.
+   - **Parameters**: 
+     - `balancesLiveScaled18`: An array of current pool balances (scaled to 18 decimals).
+   - **Returns**: `invariant`: The calculated invariant as a `uint256`.
+   - **Logic**: 
+     - The invariant is computed by multiplying all token balances and taking the square root to get a balance ratio between tokens.
+
+### 3. **`computeBalance`** 
+   - **Purpose**: Computes the new balance of a token after a liquidity operation based on the invariant growth ratio.
+   - **Parameters**: 
+     - `balancesLiveScaled18`: Current balances for tokens.
+     - `tokenInIndex`: The index of the token for which the balance is computed.
+     - `invariantRatio`: The ratio of the new invariant to the old one.
+   - **Returns**: `newBalance`: The new balance of the token.
+   - **Logic**: 
+     - Uses the invariant formula to calculate how the balance of a token would change after a liquidity change.
+
+### 4. **`getMinimumSwapFeePercentage`** 
+   - **Purpose**: Returns the minimum swap fee percentage for a pool (which is 0).
+   - **Returns**: `_MIN_SWAP_FEE_PERCENTAGE`.
+
+### 5. **`getMaximumSwapFeePercentage`** 
+   - **Purpose**: Returns the maximum swap fee percentage (10%).
+   - **Returns**: `_MAX_SWAP_FEE_PERCENTAGE`.
+
+---
+
+## Custom Liquidity Hooks
+
+### 6. **`onAddLiquidityCustom`** 
+   - **Purpose**: Custom logic for adding liquidity to the pool.
+   - **Parameters**:
+     - `router`: The address that initiated the add liquidity operation.
+     - `maxAmountsInScaled18`: Maximum input amounts per token (scaled to 18 decimals).
+     - `minBptAmountOut`: Minimum output of pool tokens (BPT).
+     - `balancesScaled18`: Current pool balances (scaled to 18 decimals).
+     - `userData`: Arbitrary data sent with the request.
+   - **Returns**: 
+     - `amountsInScaled18`: Actual input amounts of tokens.
+     - `bptAmountOut`: Amount of pool tokens to be minted.
+     - `swapFeeAmountsScaled18`: Swap fees charged for each token.
+     - `returnData`: Custom return data.
+   - **Logic**: 
+     - It calculates the invariant before and after the liquidity addition.
+     - Updates the balances of tokens and returns the amount of pool tokens to mint.
+
+### 7. **`onRemoveLiquidityCustom`** 
+   - **Purpose**: Custom logic for removing liquidity from the pool.
+   - **Parameters**:
+     - `router`: The address that initiated the remove liquidity operation.
+     - `maxBptAmountIn`: Maximum amount of pool tokens to burn.
+     - `minAmountsOutScaled18`: Minimum output amounts of each token.
+     - `balancesScaled18`: Current pool balances (scaled to 18 decimals).
+     - `userData`: Arbitrary data sent with the request.
+   - **Returns**:
+     - `bptAmountIn`: Amount of pool tokens burned.
+     - `amountsOutScaled18`: Amount of tokens withdrawn.
+     - `swapFeeAmountsScaled18`: Swap fees charged for each token.
+     - `returnData`: Custom return data.
+   - **Logic**: 
+     - It calculates the invariant before and after the liquidity removal.
+     - Reduces the pool token balances and returns the number of pool tokens burned.
 
-[![intro-to-scaffold-balancer](https://github.com/user-attachments/assets/f862091d-2fe9-4b4b-8d70-cb2fdc667384)](https://www.youtube.com/watch?v=m6q5M34ZdXw)
 
-### 🔁 Development Life Cycle
 
-1. Learn the core concepts for building on top of Balancer v3
-2. Configure and deploy factories, pools, and hooks contracts to a local anvil fork of Sepolia
-3. Interact with pools via a frontend that runs at [localhost:3000](http://localhost:3000/)
-
-### 🪧 Table Of Contents
-
-- [🧑‍💻 Environment Setup](#-environment-setup)
-- [👩‍🏫 Learn Core Concepts](#-learn-core-concepts)
-- [🕵️ Explore the Examples](#-explore-the-examples)
-- [🌊 Create a Custom Pool](#-create-a-custom-pool)
-- [🏭 Create a Pool Factory](#-create-a-pool-factory)
-- [🪝 Create a Pool Hook](#-create-a-pool-hook)
-- [🚢 Deploy the Contracts](#-deploy-the-contracts)
-- [🧪 Test the Contracts](#-test-the-contracts)
-
-## 🧑‍💻 Environment Setup
-
-### 1. Requirements 📜
-
-- [Node (>= v18.17)](https://nodejs.org/en/download/)
-- Yarn ([v1](https://classic.yarnpkg.com/en/docs/install/) or [v2+](https://yarnpkg.com/getting-started/install))
-- [Git](https://git-scm.com/downloads)
-- [Foundry](https://book.getfoundry.sh/getting-started/installation) (>= v0.2.0)
-
-### 2. Quickstart 🏃
-
-1. Ensure you have the latest version of foundry installed
-
-```
-foundryup
-```
-
-2. Clone this repo & install dependencies
-
-```bash
-git clone https://github.com/balancer/scaffold-balancer-v3.git
-cd scaffold-balancer-v3
-yarn install
-```
-
-3. Set a `SEPOLIA_RPC_URL` in the `packages/foundry/.env` file
-
-```
-SEPOLIA_RPC_URL=...
-```
-
-4. Start a local anvil fork of the Sepolia testnet
-
-```bash
-yarn fork
-```
-
-5. Deploy the mock tokens, pool factories, pool hooks, and custom pools contracts
-   > By default, the anvil account #0 will be the deployer and recieve the mock tokens and BPT from pool initialization
-
-```bash
-yarn deploy
-```
-
-6. Start the nextjs frontend
-
-```bash
-yarn start
-```
-
-7. Explore the frontend
-
-- Navigate to http://localhost:3000 to see the home page
-- Visit the [Pools Page](http://localhost:3000/pools) to search by address or select using the pool buttons
-- Vist the [Debug Page](http://localhost:3000/debug) to see the mock tokens, factory, and hooks contracts
-
-8. Run the Foundry tests
-
-```
-yarn test
-```
-
-### 3. Scaffold ETH 2 Tips 🏗️
-
-SE-2 offers a variety of configuration options for connecting an account, choosing networks, and deploying contracts
-
-<details><summary><strong>🔥 Burner Wallet</strong></summary>
-
-If you do not have an active wallet extension connected to your web browser, then scaffold eth will automatically connect to a "burner wallet" that is randomly generated on the frontend and saved to the browser's local storage. When using the burner wallet, transactions will be instantly signed, which is convenient for quick iterative development.
-
-To force the use of burner wallet, disable your browsers wallet extensions and refresh the page. Note that the burner wallet comes with 0 ETH to pay for gas so you will need to click the faucet button in top right corner. Also the mock tokens for the pool are minted to your deployer account set in `.env` so you will want to navigate to the "Debug Contracts" page to mint your burner wallet some mock tokens to use with the pool.
-
-![Burner Wallet](https://github.com/Dev-Rel-as-a-Service/scaffold-balancer-v3/assets/73561520/0a1f3456-f22a-46b5-9e05-0ef5cd17cce7)
-
-![Debug Tab Mint](https://github.com/Dev-Rel-as-a-Service/scaffold-balancer-v3/assets/73561520/fbb53772-8f6d-454d-a153-0e7a2925ef9f)
-
-</details>
-
-<details><summary><strong>👛 Browser Extension Wallet</strong></summary>
-    
-- To use your preferred browser extension wallet, ensure that the account you are using matches the PK you previously provided in the `foundry/.env` file
-- You may need to add a local development network with rpc url `http://127.0.0.1:8545/` and chain id `31337`. Also, you may need to reset the nonce data for your wallet exension if it gets out of sync.
-
-</details>
-
-<details><summary><strong>🐛 Debug Contracts Page </strong></summary>
-
-The [Debug Contracts Page](http://localhost:3000/debug) can be useful for viewing and interacting with all of the externally avaiable read and write functions of a contract. The page will automatically hot reload with contracts that are deployed via the `01_DeployConstantSumFactory.s.sol` script. We use this handy setup to mint `mockERC20` tokens to any connected wallet
-
-</details>
-
-<details><summary><strong>🌐 Changing The Frontend Network Connection</strong></summary>
-
-- The network the frontend points at is set via `targetNetworks` in the `scaffold.config.ts` file using `chains` from viem.
-- By default, the frontend runs on a local node at `http://127.0.0.1:8545`
-
-```typescript
-const scaffoldConfig = {
-  targetNetworks: [chains.foundry],
-```
-
-</details>
-
-<details><summary><strong>🍴 Changing The Forked Network</strong></summary>
-
-- By default, the `yarn fork` command points at sepolia, but any of the network aliases from the `[rpc_endpoints]` of `foundry.toml` can be used to modify the `"fork"` alias in the `packages/foundry/package.json` file
-
-```json
-	"fork": "anvil --fork-url ${0:-sepolia} --chain-id 31337 --config-out localhost.json",
-```
-
-- To point the frontend at a different forked network, change the `targetFork` in `scaffold.config.ts`
-
-```typescript
-const scaffoldConfig = {
-  // The networks the frontend can connect to
-  targetNetworks: [chains.foundry],
-
-  // If using chains.foundry as your targetNetwork, you must specify a network to fork
-  targetFork: chains.sepolia,
-```
-
-</details>
-
-## 👩‍🏫 Learn Core Concepts
-
-- [Contract Architecture](https://docs-v3.balancer.fi/concepts/core-concepts/architecture.html)
-- [Balancer Pool Tokens](https://docs-v3.balancer.fi/concepts/core-concepts/balancer-pool-tokens.html)
-- [Balancer Pool Types](https://docs-v3.balancer.fi/concepts/explore-available-balancer-pools/)
-- [Building Custom AMMs](https://docs-v3.balancer.fi/build-a-custom-amm/)
-- [Exploring Hooks and Custom Routers](https://pitchandrolls.com/2024/08/30/unlocking-the-power-of-balancer-v3-exploring-hooks-and-custom-routers/)
-- [Hook Development Tips](https://medium.com/@johngrant/unlocking-the-power-of-balancer-v3-hook-development-made-simple-831391a68296)
-
-![v3-components](https://github.com/user-attachments/assets/ccda9323-790f-4276-b092-c867fd80bf9e)
-
-## 🕵️ Explore the Examples
-
-Each of the following examples have turn key deploy scripts that can be found in the [foundry/script/](https://github.com/balancer/scaffold-balancer-v3/tree/main/packages/foundry/script) directory
-
-### 1. Constant Sum Pool with Dynamic Swap Fee Hook
-
-The swap fee percentage is altered by the hook contract before the pool calculates the amount for the swap
-
-![dynamic-fee-hook](https://github.com/user-attachments/assets/5ba69ea3-6894-4eeb-befa-ed87cfeb6b13)
-
-### 2. Constant Product Pool with Lottery Hook
-
-An after swap hook makes a request to an oracle contract for a random number
-
-![after-swap-hook](https://github.com/user-attachments/assets/594ce1ac-2edc-4d16-9631-14feb2d085f8)
-
-### 3. Weighted Pool with Exit Fee Hook
-
-An after remove liquidity hook adjusts the amounts before the vault transfers tokens to the user
-
-![after-remove-liquidity-hook](https://github.com/user-attachments/assets/2e8f4a5c-f168-4021-b316-28a79472c8d1)
-
-## 🌊 Create a Custom Pool
-
-[![custom-amm-video](https://github.com/user-attachments/assets/e6069a51-f1b5-4f98-a2a9-3a2098696f96)](https://www.youtube.com/watch?v=kXynS3jAu0M)
-
-### 1. Review the Docs 📖
-
-- [Create a custom AMM with a novel invariant](https://docs-v3.balancer.fi/build-a-custom-amm/build-an-amm/create-custom-amm-with-novel-invariant.html)
-
-### 2. Recall the Key Requirements 🔑
-
-- Must inherit from `IBasePool` and `BalancerPoolToken`
-- Must implement `onSwap`, `computeInvariant`, and `computeBalance`
-- Must implement `getMaximumSwapFeePercentage` and `getMinimumSwapFeePercentage`
-
-### 3. Write a Custom Pool Contract 📝
-
-- To get started, edit the`ConstantSumPool.sol` contract directly or make a copy
-
-## 🏭 Create a Pool Factory
-
-After designing a pool contract, the next step is to prepare a factory contract because Balancer's off-chain infrastructure uses the factory address as a means to identify the type of pool, which is important for integration into the UI, SDK, and external aggregators
-
-### 1. Review the Docs 📖
-
-- [Deploy a Custom AMM Using a Factory](https://docs-v3.balancer.fi/build-a-custom-amm/build-an-amm/deploy-custom-amm-using-factory.html)
-
-### 2. Recall the Key Requirements 🔑
-
-- A pool factory contract must inherit from [BasePoolFactory](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/vault/contracts/factories/BasePoolFactory.sol)
-- Use the internal `_create` function to deploy a new pool
-- Use the internal `_registerPoolWithVault` fuction to register a pool immediately after creation
-
-### 3. Write a Factory Contract 📝
-
-- To get started, edit the`ConstantSumFactory.sol` contract directly or make a copy
-
-## 🪝 Create a Pool Hook
-
-[![hook-video](https://github.com/user-attachments/assets/96e12c29-53c2-4a52-9437-e477f6d992d1)](https://www.youtube.com/watch?v=kaz6duliRPA)
-
-### 1. Review the Docs 📖
-
-- [Extend an Existing Pool Type Using Hooks](https://docs-v3.balancer.fi/build-a-custom-amm/build-an-amm/extend-existing-pool-type-using-hooks.html)
-
-### 2. Recall the Key Requirements 🔑
-
-- A hooks contract must inherit from [BasePoolHooks.sol](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/vault/contracts/BaseHooks.sol)
-- A hooks contract should also inherit from [VaultGuard.sol](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/vault/contracts/VaultGuard.sol)
-- Must implement `onRegister` to determine if a pool is allowed to use the hook contract
-- Must implement `getHookFlags` to define which hooks are supported
-- The `onlyVault` modifier should be applied to all hooks functions (i.e. `onRegister`, `onBeforeSwap`, `onAfterSwap` ect.)
-
-### 3. Write a Hook Contract 📝
-
-- To get started, edit the `VeBALFeeDiscountHook.sol` contract directly or make a copy
-
-## 🚢 Deploy the Contracts
-
-The deploy scripts are located in the [foundry/script/](https://github.com/balancer/scaffold-balancer-v3/tree/main/packages/foundry/script) directory. To better understand the lifecycle of deploying a pool that uses a hooks contract, see the diagram below
-
-![pool-deploy-scripts](https://github.com/user-attachments/assets/bb906080-8f42-46c0-af90-ba01ba1754fc)
-
-### 1. Modifying the Deploy Scripts 🛠️
-
-For all the scaffold integrations to work properly, each deploy script must be imported into `Deploy.s.sol` and inherited by the `DeployScript` contract in `Deploy.s.sol`
-
-### 2. Broadcast the Transactions 📡
-
-#### Deploy to local fork
-
-1. Run the following command
-
-```bash
-yarn deploy
-```
-
-#### Deploy to a live network
-
-1. Add a `DEPLOYER_PRIVATE_KEY` to the `packages/foundry/.env` file
-
-```
-DEPLOYER_PRIVATE_KEY=0x...
-SEPOLIA_RPC_URL=...
-```
-
-> The `DEPLOYER_PRIVATE_KEY` must start with `0x` and must hold enough Sepolia ETH to deploy the contracts. This account will receive the BPT from pool initialization
-
-2. Run the following command
-
-```
-yarn deploy --network sepolia
-```
-
-## 🧪 Test the Contracts
-
-The [balancer-v3-monorepo](https://github.com/balancer/balancer-v3-monorepo) provides testing utility contracts like [BasePoolTest](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/vault/test/foundry/utils/BasePoolTest.sol) and [BaseVaultTest](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/vault/test/foundry/utils/BaseVaultTest.sol). Therefore, the best way to begin writing tests for custom factories, pools, and hooks contracts is to leverage the examples established by the source code.
-
-### 1. Testing Factories 👨‍🔬
-
-The `ConstantSumFactoryTest` roughly mirrors the [WeightedPool8020FactoryTest
-](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/pool-weighted/test/foundry/WeightedPool8020Factory.t.sol)
-
-```
-yarn test --match-contract ConstantSumFactoryTest
-```
-
-### 2. Testing Pools 🏊
-
-The `ConstantSumPoolTest` roughly mirrors the [WeightedPoolTest](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/pool-weighted/test/foundry/WeightedPool.t.sol)
-
-```
-yarn test --match-contract ConstantSumPoolTest
-```
-
-### 3. Testing Hooks 🎣
-
-The `VeBALFeeDiscountHookExampleTest` mirrors the [VeBALFeeDiscountHookExampleTest](https://github.com/balancer/balancer-v3-monorepo/blob/main/pkg/pool-hooks/test/foundry/VeBALFeeDiscountHookExample.t.sol)
-
-```
-yarn test --match-contract VeBALFeeDiscountHookExampleTest
-```

From d3eba99bcc61f0820ae6a862adf6d18a50ead5ed Mon Sep 17 00:00:00 2001
From: harpreet singh <happybiostockcode07040@gmail.com>
Date: Mon, 7 Oct 2024 18:40:00 +0530
Subject: [PATCH 3/7] update readme and contract

---
 README.md                                     | 259 +++++++++++-------
 packages/foundry/contracts/hooks/SafeSwap.sol |  28 +-
 2 files changed, 179 insertions(+), 108 deletions(-)

diff --git a/README.md b/README.md
index de3174fd..57ac0d6c 100644
--- a/README.md
+++ b/README.md
@@ -1,109 +1,158 @@
 
-# SafeSwap Contract
-
-The **SafeSwap** contract is an implementation based on the **Balancer V3 protocol**, designed to manage token swaps, liquidity additions, and liquidity removals in a decentralized way. This contract introduces custom hooks for liquidity operations and operates within the Balancer V3 framework, leveraging mathematical invariants and swap fees.
-
-## Overview
-
-The contract extends the core components of Balancer V3:
-- **BalancerPoolToken**: This provides the functionality of an ERC-20 compliant pool token, which represents the user's share in the liquidity pool.
-- **IPoolLiquidity**: This is an interface that manages the core pool liquidity operations.
-- **FixedPoint**: A math library that allows for accurate fixed-point arithmetic.
-
-### Key Constants
-
-- `_MIN_SWAP_FEE_PERCENTAGE`: Minimum swap fee percentage, set to `0`.
-- `_MAX_SWAP_FEE_PERCENTAGE`: Maximum swap fee percentage, capped at `10%` (0.1e18).
-
-## Constructor
-
-The constructor initializes the `BalancerPoolToken` with the provided vault, pool name, and symbol.
-
-### Parameters:
-- **vault**: The Balancer vault where the pool is registered.
-- **name**: Name of the pool token.
-- **symbol**: Symbol of the pool token.
-
----
-
-## Core Functions
-
-### 1. **`onSwap`** 
-   - **Purpose**: Executes a swap in the pool by calculating the output token amount.
-   - **Parameters**: 
-     - `PoolSwapParams calldata params`: Contains balances and amount details for the swap.
-   - **Returns**: `amountCalculatedScaled18`: The calculated amount of the output token.
-   - **Logic**: 
-     - It uses a formula to compute the new balance of the output token based on the input token and respective balances.
-     - Formula: 
-       ```
-       amountCalculatedScaled18 = (params.balancesScaled18[params.indexOut] * params.amountGivenScaled18) 
-                                  / (params.balancesScaled18[params.indexIn] + params.amountGivenScaled18)
-       ```
-
-### 2. **`computeInvariant`** 
-   - **Purpose**: Computes the pool’s invariant, which is a mathematical property representing the pool’s balance state.
-   - **Parameters**: 
-     - `balancesLiveScaled18`: An array of current pool balances (scaled to 18 decimals).
-   - **Returns**: `invariant`: The calculated invariant as a `uint256`.
-   - **Logic**: 
-     - The invariant is computed by multiplying all token balances and taking the square root to get a balance ratio between tokens.
-
-### 3. **`computeBalance`** 
-   - **Purpose**: Computes the new balance of a token after a liquidity operation based on the invariant growth ratio.
-   - **Parameters**: 
-     - `balancesLiveScaled18`: Current balances for tokens.
-     - `tokenInIndex`: The index of the token for which the balance is computed.
-     - `invariantRatio`: The ratio of the new invariant to the old one.
-   - **Returns**: `newBalance`: The new balance of the token.
-   - **Logic**: 
-     - Uses the invariant formula to calculate how the balance of a token would change after a liquidity change.
-
-### 4. **`getMinimumSwapFeePercentage`** 
-   - **Purpose**: Returns the minimum swap fee percentage for a pool (which is 0).
-   - **Returns**: `_MIN_SWAP_FEE_PERCENTAGE`.
-
-### 5. **`getMaximumSwapFeePercentage`** 
-   - **Purpose**: Returns the maximum swap fee percentage (10%).
-   - **Returns**: `_MAX_SWAP_FEE_PERCENTAGE`.
-
----
-
-## Custom Liquidity Hooks
-
-### 6. **`onAddLiquidityCustom`** 
-   - **Purpose**: Custom logic for adding liquidity to the pool.
-   - **Parameters**:
-     - `router`: The address that initiated the add liquidity operation.
-     - `maxAmountsInScaled18`: Maximum input amounts per token (scaled to 18 decimals).
-     - `minBptAmountOut`: Minimum output of pool tokens (BPT).
-     - `balancesScaled18`: Current pool balances (scaled to 18 decimals).
-     - `userData`: Arbitrary data sent with the request.
-   - **Returns**: 
-     - `amountsInScaled18`: Actual input amounts of tokens.
-     - `bptAmountOut`: Amount of pool tokens to be minted.
-     - `swapFeeAmountsScaled18`: Swap fees charged for each token.
-     - `returnData`: Custom return data.
-   - **Logic**: 
-     - It calculates the invariant before and after the liquidity addition.
-     - Updates the balances of tokens and returns the amount of pool tokens to mint.
-
-### 7. **`onRemoveLiquidityCustom`** 
-   - **Purpose**: Custom logic for removing liquidity from the pool.
-   - **Parameters**:
-     - `router`: The address that initiated the remove liquidity operation.
-     - `maxBptAmountIn`: Maximum amount of pool tokens to burn.
-     - `minAmountsOutScaled18`: Minimum output amounts of each token.
-     - `balancesScaled18`: Current pool balances (scaled to 18 decimals).
-     - `userData`: Arbitrary data sent with the request.
-   - **Returns**:
-     - `bptAmountIn`: Amount of pool tokens burned.
-     - `amountsOutScaled18`: Amount of tokens withdrawn.
-     - `swapFeeAmountsScaled18`: Swap fees charged for each token.
-     - `returnData`: Custom return data.
-   - **Logic**: 
-     - It calculates the invariant before and after the liquidity removal.
-     - Reduces the pool token balances and returns the number of pool tokens burned.
+# SafeSwap
 
+SafeSwap is a decentralized automated market maker (AMM) built on the Balancer V3 protocol. It enables liquidity providers to add and remove liquidity in a pool while allowing users to swap between tokens. Additionally, SafeSwap introduces a unique discount mechanism for users holding a specific token, providing them with a 10% discount on swap fees.
 
+## Key Features
 
+- **Balancer V3 Integration**: Built on top of the Balancer V3 protocol with BalancerPoolToken, allowing flexible pool operations.
+- **Custom Liquidity Operations**: Custom hooks for adding and removing liquidity.
+- **Discount on Swap Fees**: Users holding a specific ERC-20 token are eligible for a 10% discount on the swap fee.
+- **Optimized Fee Structure**: The contract enforces a swap fee between 0% and 10%, with a default fee of 10%, reduced if the user qualifies for a discount.
+
+## Contract Structure
+
+### SafeSwap
+
+The `SafeSwap` contract extends the Balancer V3 pool functionality with added customization for liquidity and swaps:
+
+- **Discount Token**: The contract allows setting a discount token, an ERC-20 token. Users holding this token in their wallets receive a reduced swap fee.
+- **Swap Fee Logic**: The default swap fee is 10%, which can be reduced by 10% (i.e., down to 9%) if the user holds the discount token.
+- **Custom Liquidity Hooks**: Implementations for custom liquidity operations (`onAddLiquidityCustom` and `onRemoveLiquidityCustom`).
+
+### Key Functions
+
+1. **`onSwap`**: Executes a token swap within the pool. If the user holds the discount token, they receive a 10% discount on the swap fee.
+   - Parameters:
+     - `params`: Pool swap parameters, including token balances and amounts.
+   - Returns:
+     - `amountCalculatedScaled18`: The calculated amount for the swap, adjusted for the discounted swap fee.
+
+2. **`computeInvariant`**: Calculates the pool's invariant based on the current token balances.
+
+3. **`onAddLiquidityCustom`**: Custom implementation for adding liquidity to the pool.
+   - Returns:
+     - `amountsInScaled18`: Amount of tokens being added to the pool.
+     - `bptAmountOut`: Calculated pool token (BPT) amount the user receives.
+     - `swapFeeAmountsScaled18`: The swap fee charged for each token.
+
+4. **`onRemoveLiquidityCustom`**: Custom implementation for removing liquidity from the pool.
+   - Returns:
+     - `bptAmountIn`: Amount of pool tokens (BPT) burned.
+     - `amountsOutScaled18`: Amount of tokens the user receives.
+     - `swapFeeAmountsScaled18`: The swap fee charged for each token.
+
+### Discount Logic
+
+The discount logic checks whether the user holds a specific ERC-20 token (the "discount token"). If the user holds this token, they are eligible for a 10% discount on the swap fee.
+
+### Constructor Parameters
+
+The `SafeSwap` contract constructor takes the following parameters:
+
+1. **`IVault vault`**: The Balancer Vault that manages the pool.
+2. **`string name`**: Name of the pool token (BPT).
+3. **`string symbol`**: Symbol of the pool token (BPT).
+4. **`address discountToken`**: The address of the ERC-20 token that provides users with the discount on swap fees.
+
+### Deployment
+
+To deploy the contract, provide the following arguments to the constructor:
+
+```solidity
+constructor(
+    IVault vault,
+    string memory name,
+    string memory symbol,
+    address _discountToken
+)
+```
+
+- **`vault`**: Address of the Balancer Vault.
+- **`name`**: The name of the Balancer Pool Token (BPT).
+- **`symbol`**: The symbol of the Balancer Pool Token (BPT).
+- **`_discountToken`**: Address of the ERC-20 token that gives users a swap fee discount.
+
+### Example
+
+Deploying the `SafeSwap` contract:
+
+```solidity
+SafeSwap safeSwap = new SafeSwap(
+    vaultAddress,
+    "SafeSwap Pool Token",
+    "SSPT",
+    discountTokenAddress
+);
+```
+
+### Swap Fee Discount Example
+
+If a user holds the specified `discountToken`, they will receive a 10% discount on the swap fee:
+
+- Default swap fee: 10%
+- Discounted swap fee: 9% (if the user holds the discount token)
+
+## Installation
+
+To install and run the SafeSwap contract, follow these steps:
+
+1. Clone the repository:
+
+```bash
+git clone https://github.com/your-repo/safeswap.git
+```
+
+2. Install dependencies:
+
+```bash
+npm install
+```
+
+3. Compile the contracts:
+
+```bash
+npx hardhat compile
+```
+
+4. Run tests:
+
+```bash
+npx hardhat test
+```
+
+## Usage
+
+Once deployed, users can interact with the `SafeSwap` contract to:
+
+1. **Swap tokens** using the `onSwap` function.
+2. **Add liquidity** with the `onAddLiquidityCustom` function.
+3. **Remove liquidity** with the `onRemoveLiquidityCustom` function.
+
+### Interacting with the Contract
+
+#### Swap Tokens
+
+To execute a swap, call the `onSwap` function:
+
+```solidity
+safeSwap.onSwap(swapParams);
+```
+
+If the user holds the discount token, the swap fee will automatically be reduced.
+
+#### Add Liquidity
+
+To add liquidity, call the `onAddLiquidityCustom` function:
+
+```solidity
+safeSwap.onAddLiquidityCustom(router, maxAmountsInScaled18, minBptAmountOut, balancesScaled18, userData);
+```
+
+#### Remove Liquidity
+
+To remove liquidity, call the `onRemoveLiquidityCustom` function:
+
+```solidity
+safeSwap.onRemoveLiquidityCustom(router, maxBptAmountIn, minAmountsOutScaled18, balancesScaled18, userData);
+```
diff --git a/packages/foundry/contracts/hooks/SafeSwap.sol b/packages/foundry/contracts/hooks/SafeSwap.sol
index 7a4a63d4..9dddd663 100644
--- a/packages/foundry/contracts/hooks/SafeSwap.sol
+++ b/packages/foundry/contracts/hooks/SafeSwap.sol
@@ -1,28 +1,50 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 pragma solidity ^0.8.24;
-// import {BalancerPoolToken} from "../lib/balancer-v3-monorepo/pkg/vault/contracts/BalancerPoolToken.sol";
+
 import { BalancerPoolToken } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BalancerPoolToken.sol";
 import { IPoolLiquidity } from "@balancer-labs/balancer-v3-monorepo/pkg/interfaces/contracts/vault/IPoolLiquidity.sol";
 import { FixedPoint } from "@balancer-labs/balancer-v3-monorepo/pkg/solidity-utils/contracts/math/FixedPoint.sol";
 import { Math } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BasePoolMath.sol";
+import { IERC20 } from "@balancer-labs/balancer-v3-monorepo/pkg/interfaces/contracts/tokens/IERC20.sol";
 
 contract SafeSwap is IBasePool, IPoolLiquidity, BalancerPoolToken {
     using FixedPoint for uint256;
 
     uint256 private constant _MIN_SWAP_FEE_PERCENTAGE = 0;
     uint256 private constant _MAX_SWAP_FEE_PERCENTAGE = 0.1e18; // 10%
+    uint256 private constant _DISCOUNT_PERCENTAGE = 10; // 10% discount
+
+    address public discountToken; // Address of the token that provides a discount
 
-    constructor(IVault vault, string memory name, string memory symbol) BalancerPoolToken(vault, name, symbol) {}
+    constructor(
+        IVault vault,
+        string memory name,
+        string memory symbol,
+        address _discountToken // Address of the discount token to be passed in constructor
+    ) BalancerPoolToken(vault, name, symbol) {
+        discountToken = _discountToken;
+    }
 
     /**
      * @notice Execute a swap in the pool.
      * @param params Swap parameters
      * @return amountCalculatedScaled18 Calculated amount for the swap
      */
-    function onSwap(PoolSwapParams calldata params) external pure returns (uint256 amountCalculatedScaled18) {
+    function onSwap(PoolSwapParams calldata params) external view returns (uint256 amountCalculatedScaled18) {
+        uint256 swapFee = _MAX_SWAP_FEE_PERCENTAGE; // Default fee
+
+        // Check if the user holds the discount token and apply the discount if applicable
+        if (IERC20(discountToken).balanceOf(msg.sender) > 0) {
+            swapFee = swapFee - ((swapFee * _DISCOUNT_PERCENTAGE) / 100); // Apply 10% discount
+        }
+
+        // Swap logic with discounted swap fee
         amountCalculatedScaled18 =
             (params.balancesScaled18[params.indexOut] * params.amountGivenScaled18) /
             (params.balancesScaled18[params.indexIn] + params.amountGivenScaled18);
+
+        // Apply the swap fee (using the discounted or full fee)
+        amountCalculatedScaled18 = amountCalculatedScaled18.mulDown(FixedPoint.ONE - swapFee);
     }
 
     /**

From fdb2b6b179519f70f3e1deea06e6047837b52db4 Mon Sep 17 00:00:00 2001
From: harpreet singh <happybiostockcode07040@gmail.com>
Date: Tue, 8 Oct 2024 15:52:13 +0530
Subject: [PATCH 4/7] deployed and test swap

---
 packages/foundry/.env.example                 |   2 +-
 packages/foundry/contracts/hooks/SafeSwap.sol | 182 ------------------
 .../contracts/hooks/SwapDiscountHook.sol      |  95 +++++++++
 packages/foundry/script/SafeSwap.s.sol        |  26 ---
 .../foundry/script/SafeSwapDiscount.s.sol     |  50 +++++
 packages/foundry/test/SafeSwapDiscount.t.sol  | 142 ++++++++++++++
 6 files changed, 288 insertions(+), 209 deletions(-)
 delete mode 100644 packages/foundry/contracts/hooks/SafeSwap.sol
 create mode 100644 packages/foundry/contracts/hooks/SwapDiscountHook.sol
 delete mode 100644 packages/foundry/script/SafeSwap.s.sol
 create mode 100644 packages/foundry/script/SafeSwapDiscount.s.sol
 create mode 100644 packages/foundry/test/SafeSwapDiscount.t.sol

diff --git a/packages/foundry/.env.example b/packages/foundry/.env.example
index 05cf662f..e544d074 100644
--- a/packages/foundry/.env.example
+++ b/packages/foundry/.env.example
@@ -1,2 +1,2 @@
 DEPLOYER_PRIVATE_KEY=
-SEPOLIA_RPC_URL=
\ No newline at end of file
+SEPOLIA_RPC_URL=https://sepolia.infura.io/v3/
\ No newline at end of file
diff --git a/packages/foundry/contracts/hooks/SafeSwap.sol b/packages/foundry/contracts/hooks/SafeSwap.sol
deleted file mode 100644
index 9dddd663..00000000
--- a/packages/foundry/contracts/hooks/SafeSwap.sol
+++ /dev/null
@@ -1,182 +0,0 @@
-// SPDX-License-Identifier: GPL-3.0-or-later
-pragma solidity ^0.8.24;
-
-import { BalancerPoolToken } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BalancerPoolToken.sol";
-import { IPoolLiquidity } from "@balancer-labs/balancer-v3-monorepo/pkg/interfaces/contracts/vault/IPoolLiquidity.sol";
-import { FixedPoint } from "@balancer-labs/balancer-v3-monorepo/pkg/solidity-utils/contracts/math/FixedPoint.sol";
-import { Math } from "@balancer-labs/balancer-v3-monorepo/pkg/vault/contracts/BasePoolMath.sol";
-import { IERC20 } from "@balancer-labs/balancer-v3-monorepo/pkg/interfaces/contracts/tokens/IERC20.sol";
-
-contract SafeSwap is IBasePool, IPoolLiquidity, BalancerPoolToken {
-    using FixedPoint for uint256;
-
-    uint256 private constant _MIN_SWAP_FEE_PERCENTAGE = 0;
-    uint256 private constant _MAX_SWAP_FEE_PERCENTAGE = 0.1e18; // 10%
-    uint256 private constant _DISCOUNT_PERCENTAGE = 10; // 10% discount
-
-    address public discountToken; // Address of the token that provides a discount
-
-    constructor(
-        IVault vault,
-        string memory name,
-        string memory symbol,
-        address _discountToken // Address of the discount token to be passed in constructor
-    ) BalancerPoolToken(vault, name, symbol) {
-        discountToken = _discountToken;
-    }
-
-    /**
-     * @notice Execute a swap in the pool.
-     * @param params Swap parameters
-     * @return amountCalculatedScaled18 Calculated amount for the swap
-     */
-    function onSwap(PoolSwapParams calldata params) external view returns (uint256 amountCalculatedScaled18) {
-        uint256 swapFee = _MAX_SWAP_FEE_PERCENTAGE; // Default fee
-
-        // Check if the user holds the discount token and apply the discount if applicable
-        if (IERC20(discountToken).balanceOf(msg.sender) > 0) {
-            swapFee = swapFee - ((swapFee * _DISCOUNT_PERCENTAGE) / 100); // Apply 10% discount
-        }
-
-        // Swap logic with discounted swap fee
-        amountCalculatedScaled18 =
-            (params.balancesScaled18[params.indexOut] * params.amountGivenScaled18) /
-            (params.balancesScaled18[params.indexIn] + params.amountGivenScaled18);
-
-        // Apply the swap fee (using the discounted or full fee)
-        amountCalculatedScaled18 = amountCalculatedScaled18.mulDown(FixedPoint.ONE - swapFee);
-    }
-
-    /**
-     * @notice Computes and returns the pool's invariant.
-     * @dev This function computes the invariant based on current balances
-     * @param balancesLiveScaled18 Array of current pool balances for each token in the pool, scaled to 18 decimals
-     * @return invariant The calculated invariant of the pool, represented as a uint256
-     */
-    function computeInvariant(uint256[] memory balancesLiveScaled18) public pure returns (uint256 invariant) {
-        // expected to work with 2 tokens only
-        invariant = FixedPoint.ONE;
-        for (uint256 i = 0; i < balancesLiveScaled18.length; ++i) {
-            invariant = invariant.mulDown(balancesLiveScaled18[i]);
-        }
-        // scale the invariant to 1e18
-        invariant = Math.sqrt(invariant) * 1e9;
-    }
-
-    /**
-     * @dev Computes the new balance of a token after an operation, given the invariant growth ratio and all other balances.
-     * @param balancesLiveScaled18 Current live balances (adjusted for decimals, rates, etc.)
-     * @param tokenInIndex The index of the token we're computing the balance for, in token registration order
-     * @param invariantRatio The ratio of the new invariant (after an operation) to the old
-     * @return newBalance The new balance of the selected token, after the operation
-     */
-    function computeBalance(
-        uint256[] memory balancesLiveScaled18,
-        uint256 tokenInIndex,
-        uint256 invariantRatio
-    ) external pure returns (uint256 newBalance) {
-        uint256 otherTokenIndex = tokenInIndex == 0 ? 1 : 0;
-
-        uint256 newInvariant = computeInvariant(balancesLiveScaled18).mulDown(invariantRatio);
-
-        newBalance = (newInvariant * newInvariant) / balancesLiveScaled18[otherTokenIndex];
-    }
-
-    /// @return minimumSwapFeePercentage The minimum swap fee percentage for a pool
-    function getMinimumSwapFeePercentage() external pure returns (uint256) {
-        return _MIN_SWAP_FEE_PERCENTAGE;
-    }
-
-    /// @return maximumSwapFeePercentage The maximum swap fee percentage for a pool
-    function getMaximumSwapFeePercentage() external pure returns (uint256) {
-        return _MAX_SWAP_FEE_PERCENTAGE;
-    }
-
-    /**
-     * @notice Custom add liquidity hook.
-     * @param router The address that initiated the operation
-     * @param maxAmountsInScaled18 Maximum input amounts, sorted in token registration order
-     * @param minBptAmountOut Minimum amount of output pool tokens
-     * @param balancesScaled18 Current pool balances, sorted in token registration order
-     * @param userData Arbitrary data sent with the request
-     * @return amountsInScaled18 Input token amounts, sorted in token registration order
-     * @return bptAmountOut Calculated pool token amount to receive
-     * @return swapFeeAmountsScaled18 The amount of swap fees charged for each token
-     * @return returnData Arbitrary data with an encoded response from the pool
-     */
-    function onAddLiquidityCustom(
-        address router,
-        uint256[] memory maxAmountsInScaled18,
-        uint256 minBptAmountOut,
-        uint256[] memory balancesScaled18,
-        bytes memory userData
-    )
-        external
-        override
-        returns (
-            uint256[] memory amountsInScaled18,
-            uint256 bptAmountOut,
-            uint256[] memory swapFeeAmountsScaled18,
-            bytes memory returnData
-        )
-    {
-        // Custom logic for adding liquidity
-        uint256 invariantBefore = computeInvariant(balancesScaled18);
-        amountsInScaled18 = maxAmountsInScaled18; // You can modify this based on custom liquidity logic
-        swapFeeAmountsScaled18 = new uint256[](balancesScaled18.length); // Placeholder for swap fees
-
-        // Update balances after adding liquidity
-        for (uint256 i = 0; i < balancesScaled18.length; ++i) {
-            balancesScaled18[i] += amountsInScaled18[i];
-        }
-
-        uint256 invariantAfter = computeInvariant(balancesScaled18);
-        bptAmountOut = invariantAfter - invariantBefore; // Example calculation
-
-        returnData = userData; // Custom return data
-    }
-
-    /**
-     * @notice Custom remove liquidity hook.
-     * @param router The address that initiated the operation
-     * @param maxBptAmountIn Maximum amount of input pool tokens
-     * @param minAmountsOutScaled18 Minimum output amounts, sorted in token registration order
-     * @param balancesScaled18 Current pool balances, sorted in token registration order
-     * @param userData Arbitrary data sent with the request
-     * @return bptAmountIn Calculated pool token amount to burn
-     * @return amountsOutScaled18 Amount of tokens to receive, sorted in token registration order
-     * @return swapFeeAmountsScaled18 The amount of swap fees charged for each token
-     * @return returnData Arbitrary data with an encoded response from the pool
-     */
-    function onRemoveLiquidityCustom(
-        address router,
-        uint256 maxBptAmountIn,
-        uint256[] memory minAmountsOutScaled18,
-        uint256[] memory balancesScaled18,
-        bytes memory userData
-    )
-        external
-        override
-        returns (
-            uint256 bptAmountIn,
-            uint256[] memory amountsOutScaled18,
-            uint256[] memory swapFeeAmountsScaled18,
-            bytes memory returnData
-        )
-    {
-        // Custom logic for removing liquidity
-        uint256 invariantBefore = computeInvariant(balancesScaled18);
-        amountsOutScaled18 = minAmountsOutScaled18; // Modify this based on custom logic
-        swapFeeAmountsScaled18 = new uint256[](balancesScaled18.length); // Placeholder for swap fees
-
-        // Update balances after removing liquidity
-        for (uint256 i = 0; i < balancesScaled18.length; ++i) {
-            balancesScaled18[i] -= amountsOutScaled18[i];
-        }
-
-        uint256 invariantAfter = computeInvariant(balancesScaled18);
-        bptAmountIn = invariantBefore - invariantAfter; // Example calculation
-
-        returnData = userData; // Custom return data
-    }
-}
diff --git a/packages/foundry/contracts/hooks/SwapDiscountHook.sol b/packages/foundry/contracts/hooks/SwapDiscountHook.sol
new file mode 100644
index 00000000..3494e103
--- /dev/null
+++ b/packages/foundry/contracts/hooks/SwapDiscountHook.sol
@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
+import { IHooks } from "@balancer-labs/v3-interfaces/contracts/vault/IHooks.sol";
+import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
+import { IRouterCommon } from "@balancer-labs/v3-interfaces/contracts/vault/IRouterCommon.sol"; // Importing IRouterCommon
+import {
+    AddLiquidityKind,
+    LiquidityManagement,
+    RemoveLiquidityKind,
+    AfterSwapParams,
+    SwapKind,
+    TokenConfig,
+    HookFlags
+} from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
+import { IBasePoolFactory } from "@balancer-labs/v3-interfaces/contracts/vault/IBasePoolFactory.sol";
+
+import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
+import { BaseHooks } from "@balancer-labs/v3-vault/contracts/BaseHooks.sol";
+
+contract SwapDiscountHook is BaseHooks {
+    using FixedPoint for uint256;
+    using SafeERC20 for IERC20;
+
+    address private immutable _allowedFactory;
+    address private immutable _trustedRouter;
+    address public discountToken; // The token to check for discount eligibility (BAL)
+    uint256 public requiredBalance; // The balance of BAL required for discount eligibility
+    uint64 public hookSwapDiscountPercentage; // The discount percentage to apply
+
+    // modifier onlyVault() {
+    //     require(msg.sender == address(IVault), "Caller is not the vault");
+    //     _;
+    // }
+
+    constructor(
+        IVault vault,
+        address allowedFactory,
+        address trustedRouter,
+        address _discountToken,
+        uint64 _hookSwapDiscountPercentage,
+        uint256 _requiredBalance // Setting the required balance for discount eligibility
+    ) BaseHooks() {
+        _allowedFactory = allowedFactory;
+        _trustedRouter = trustedRouter;
+        discountToken = _discountToken;
+        hookSwapDiscountPercentage = _hookSwapDiscountPercentage;
+        requiredBalance = _requiredBalance; // Store the required balance
+    }
+
+    /// @inheritdoc IHooks
+    function onRegister(
+        address factory,
+        address pool,
+        TokenConfig[] memory,
+        LiquidityManagement calldata
+    ) public view override returns (bool) {
+        return factory == _allowedFactory && IBasePoolFactory(factory).isPoolFromFactory(pool);
+    }
+
+    /// @inheritdoc IHooks
+    function getHookFlags() public pure override returns (HookFlags memory hookFlags) {
+        hookFlags.shouldCallAfterSwap = true;
+        return hookFlags;
+    }
+
+    /// @inheritdoc IHooks
+    function onAfterSwap(
+        AfterSwapParams calldata params
+    ) public view override returns (bool success, uint256 discountedAmount) {
+        discountedAmount = params.amountCalculatedRaw;
+
+        // Check if the user holds enough BAL tokens for a discount
+        if (
+            hookSwapDiscountPercentage > 0 &&
+            address(params.tokenOut) == discountToken &&
+            params.kind == SwapKind.EXACT_IN
+        ) {
+            // Get the sender's balance of the discount token
+            uint256 userBalance = IERC20(discountToken).balanceOf(IRouterCommon(params.router).getSender());
+
+            // Apply discount if user holds enough BAL tokens
+            if (userBalance >= requiredBalance) {
+                discountedAmount = discountedAmount.mulDown(hookSwapDiscountPercentage);
+                return (true, discountedAmount);
+            }
+        }
+        return (true, discountedAmount); // No discount applied, return original amount
+    }
+}
diff --git a/packages/foundry/script/SafeSwap.s.sol b/packages/foundry/script/SafeSwap.s.sol
deleted file mode 100644
index 76e6fd15..00000000
--- a/packages/foundry/script/SafeSwap.s.sol
+++ /dev/null
@@ -1,26 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.13;
-
-import { Script, console } from "../lib/forge-std/src/Script.sol";
-import { ConstantProductPool } from "../src/ConstantProductPool.sol";
-import { IVault } from "../balancer-v3-monorepo/pkg/interfaces/contracts/vault/IVault.sol";
-
-contract SafeSwap is Script {
-    ConstantProductPool public constantProductPool;
-
-    address public vaultAddress = "0x7966FE92C59295EcE7FB5D9EfDB271967BFe2fbA";
-    string public poolName = "MyConstantProductPool";
-    string public poolSymbol = "MCP";
-
-    function setUp() public {}
-
-    function run() public {
-        vm.startBroadcast();
-
-        constantProductPool = new ConstantProductPool(IVault(vaultAddress), poolName, poolSymbol);
-
-        vm.stopBroadcast();
-    }
-}
-
-//forge script script/ConstantProductPool.s.sol --rpc-url https://sepolia.infura.io/v3/2de477c3b1b74816ae5475da6d289208 --private-key f46e7f0936b479bba879c9f764259d1e5838aa015232f0018a1c07214e491812
diff --git a/packages/foundry/script/SafeSwapDiscount.s.sol b/packages/foundry/script/SafeSwapDiscount.s.sol
new file mode 100644
index 00000000..078e5286
--- /dev/null
+++ b/packages/foundry/script/SafeSwapDiscount.s.sol
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import "forge-std/Script.sol";
+
+import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
+import { PoolFactoryMock } from "@balancer-labs/v3-vault/contracts/test/PoolFactoryMock.sol";
+import { SwapDiscountHook } from "../contracts/hooks/SwapDiscountHook.sol";
+
+contract DeploySwapDiscountHook is Script {
+    IVault internal vault;
+    address internal allowedFactory;
+    address internal trustedRouter;
+    address internal discountToken;
+    uint64 internal hookSwapDiscountPercentage; // Discount percentage
+    uint256 internal requiredBalance; // Required balance for discount eligibility
+
+    function run() external {
+        // Set the addresses and parameters
+        vault = IVault(0xBA12222222228d8Ba445958a75a0704d566BF2C8); // Replace with actual vault address
+        allowedFactory = 0xed52D8E202401645eDAD1c0AA21e872498ce47D0; // Replace with actual factory address
+        trustedRouter = 0x886A3Ec7bcC508B8795990B60Fa21f85F9dB7948; // Replace with actual router address
+        discountToken = 0xba100000625a3754423978a60c9317c58a424e3D; // Replace with actual discount token address
+        hookSwapDiscountPercentage = 50e16; // 50% discount
+        requiredBalance = 100e18; // 100 of the discount token required for eligibility
+
+        // Start the deployment process
+        vm.startBroadcast();
+
+        // Deploy the SwapDiscountHook contract
+        SwapDiscountHook swapDiscountHook = new SwapDiscountHook(
+            vault,
+            allowedFactory,
+            trustedRouter,
+            discountToken,
+            hookSwapDiscountPercentage,
+            requiredBalance
+        );
+
+        // Label the deployed contract for easier identification
+        vm.label(address(swapDiscountHook), "Swap Discount Hook");
+
+        // End the deployment process
+        vm.stopBroadcast();
+
+        // Output the deployed contract address
+        console.log("SwapDiscountHook deployed at:", address(swapDiscountHook));
+    }
+}
diff --git a/packages/foundry/test/SafeSwapDiscount.t.sol b/packages/foundry/test/SafeSwapDiscount.t.sol
new file mode 100644
index 00000000..0b09a60a
--- /dev/null
+++ b/packages/foundry/test/SafeSwapDiscount.t.sol
@@ -0,0 +1,142 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import "forge-std/Test.sol";
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
+import { IVaultAdmin } from "@balancer-labs/v3-interfaces/contracts/vault/IVaultAdmin.sol";
+import { IVaultErrors } from "@balancer-labs/v3-interfaces/contracts/vault/IVaultErrors.sol";
+import {
+    HooksConfig,
+    LiquidityManagement,
+    TokenConfig
+} from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
+
+import { CastingHelpers } from "@balancer-labs/v3-solidity-utils/contracts/helpers/CastingHelpers.sol";
+import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
+import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
+import { PoolMock } from "@balancer-labs/v3-vault/contracts/test/PoolMock.sol";
+import { PoolFactoryMock } from "@balancer-labs/v3-vault/contracts/test/PoolFactoryMock.sol";
+import { RouterMock } from "@balancer-labs/v3-vault/contracts/test/RouterMock.sol";
+
+import { SwapDiscountHook } from "packages/foundry/contracts/hooks/SwapDiscountHook.sol";
+import { console } from "forge-std/console.sol";
+
+contract SwapDiscountHookTest is BaseVaultTest {
+    using CastingHelpers for address[];
+    using FixedPoint for uint256;
+
+    uint256 internal daiIdx;
+    uint256 internal usdcIdx;
+
+    uint64 public constant MAX_SWAP_DISCOUNT_PERCENTAGE = 50e16; // 50%
+    uint256 public constant REQUIRED_BALANCE = 100e18; // Minimum required balance for discount (100 BAL)
+
+    address payable internal trustedRouter;
+
+    function setUp() public override {
+        super.setUp();
+
+        (daiIdx, usdcIdx) = getSortedIndexes(address(dai), address(usdc));
+
+        // Mint 100 BAL tokens for the test account (Bob)
+        deal(address(dai), bob, REQUIRED_BALANCE);
+    }
+
+    function createHook() internal override returns (address) {
+        trustedRouter = payable(router);
+
+        // LP will be the owner of the hook. Only LP is able to set hook fee percentages.
+        vm.prank(lp);
+        address swapDiscountHook = address(
+            new SwapDiscountHook(
+                IVault(address(vault)),
+                address(factoryMock),
+                trustedRouter,
+                address(dai),
+                MAX_SWAP_DISCOUNT_PERCENTAGE,
+                REQUIRED_BALANCE
+            )
+        );
+        vm.label(swapDiscountHook, "Swap Discount Hook");
+        return swapDiscountHook;
+    }
+
+    function testSuccessfulRegistrySwap() public {
+        // Registering with allowed factory
+        address swapDHookPool = factoryMock.createPool("Test Pool", "TEST");
+        TokenConfig[] memory tokenConfig = vault.buildTokenConfig(
+            [address(dai), address(usdc)].toMemoryArray().asIERC20()
+        );
+
+        _registerPoolWithHook(swapDHookPool, tokenConfig, address(factoryMock));
+
+        HooksConfig memory hooksConfig = vault.getHooksConfig(swapDHookPool);
+
+        assertEq(hooksConfig.hooksContract, poolHooksContract, "Wrong poolHooksContract");
+        assertEq(hooksConfig.shouldCallAfterSwap, true, "shouldCallAfterSwap is false");
+    }
+
+    function testSwapDiscountWithSufficientBalance() public {
+        // Bob has enough DAI for the discount
+        vm.prank(bob);
+        RouterMock(trustedRouter).swapSingleTokenExactIn(
+            pool,
+            usdc,
+            dai,
+            REQUIRED_BALANCE,
+            REQUIRED_BALANCE,
+            MAX_UINT256,
+            false,
+            bytes("")
+        );
+
+        // Check Bob's DAI balance after the swap
+        uint256 expectedDAIAmount = REQUIRED_BALANCE.mulDown(1 - MAX_SWAP_DISCOUNT_PERCENTAGE);
+        assertEq(dai.balanceOf(bob), expectedDAIAmount, "Bob's DAI balance is incorrect after discount");
+    }
+
+    function testSwapDiscountWithInsufficientBalance() public {
+        // Adjust Bob's DAI balance to be less than the required balance for discount
+        vm.prank(bob);
+        deal(address(dai), bob, REQUIRED_BALANCE / 2); // Only 50 DAI
+
+        vm.prank(bob);
+        RouterMock(trustedRouter).swapSingleTokenExactIn(
+            pool,
+            usdc,
+            dai,
+            REQUIRED_BALANCE,
+            REQUIRED_BALANCE,
+            MAX_UINT256,
+            false,
+            bytes("")
+        );
+
+        // Check Bob's DAI balance after the swap
+        assertEq(
+            dai.balanceOf(bob),
+            REQUIRED_BALANCE,
+            "Bob's DAI balance should be unchanged due to insufficient balance"
+        );
+    }
+
+    // Registry tests require a new pool, because an existing pool may already be registered
+    function _createPoolToRegister() private returns (address newPool) {
+        newPool = address(new PoolMock(IVault(address(vault)), "SwapD Hook Pool", "swapDHookPool"));
+        vm.label(newPool, "SwapD Hook Pool");
+    }
+
+    function _registerPoolWithHook(address swapDhookPool, TokenConfig[] memory tokenConfig, address factory) private {
+        LiquidityManagement memory liquidityManagement;
+        PoolFactoryMock(factory).registerPool(
+            swapDhookPool,
+            tokenConfig,
+            PoolRoleAccounts({ lp: lp, guardian: address(0), rewards: address(0) }),
+            poolHooksContract,
+            liquidityManagement
+        );
+    }
+}

From e66e010f3aa429aa22c2a535da796ac2adca05e4 Mon Sep 17 00:00:00 2001
From: harpreet singh <happybiostockcode07040@gmail.com>
Date: Wed, 9 Oct 2024 17:55:22 +0530
Subject: [PATCH 5/7] add contract address

---
 .gitmodules                                             | 3 +++
 packages/foundry/.env.example                           | 3 ++-
 packages/foundry/lib/openzeppelin-contracts-upgradeable | 1 +
 3 files changed, 6 insertions(+), 1 deletion(-)
 create mode 160000 packages/foundry/lib/openzeppelin-contracts-upgradeable

diff --git a/.gitmodules b/.gitmodules
index bfd357b9..dff455ab 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -16,3 +16,6 @@
 [submodule "packages/foundry/lib/balancer-v3-monorepo"]
 	path = packages/foundry/lib/balancer-v3-monorepo
 	url = https://github.com/balancer/balancer-v3-monorepo
+[submodule "packages/foundry/lib/openzeppelin-contracts-upgradeable"]
+	path = packages/foundry/lib/openzeppelin-contracts-upgradeable
+	url = https://github.com/OpenZeppelin/openzeppelin-contracts-upgradeable
diff --git a/packages/foundry/.env.example b/packages/foundry/.env.example
index e544d074..4cd43cd0 100644
--- a/packages/foundry/.env.example
+++ b/packages/foundry/.env.example
@@ -1,2 +1,3 @@
 DEPLOYER_PRIVATE_KEY=
-SEPOLIA_RPC_URL=https://sepolia.infura.io/v3/
\ No newline at end of file
+SEPOLIA_RPC_URL=https://sepolia.infura.io/v3/
+COTRACT_ADDRESS=0x90193C961A926261B756D1E5bb255e67ff9498A1
\ No newline at end of file
diff --git a/packages/foundry/lib/openzeppelin-contracts-upgradeable b/packages/foundry/lib/openzeppelin-contracts-upgradeable
new file mode 160000
index 00000000..df85a5ac
--- /dev/null
+++ b/packages/foundry/lib/openzeppelin-contracts-upgradeable
@@ -0,0 +1 @@
+Subproject commit df85a5acb5daae4dbc576b7e4ce2ca8600c646ff

From fee069488f95a48f44289df2f6361dfe384fbcef Mon Sep 17 00:00:00 2001
From: Happy <gillharpreetsingh211@gmail.com>
Date: Sat, 19 Oct 2024 15:19:56 +0530
Subject: [PATCH 6/7] Update SwapDiscountHook.sol

---
 .../contracts/hooks/SwapDiscountHook.sol      | 86 +++++++------------
 1 file changed, 30 insertions(+), 56 deletions(-)

diff --git a/packages/foundry/contracts/hooks/SwapDiscountHook.sol b/packages/foundry/contracts/hooks/SwapDiscountHook.sol
index 3494e103..1ea6f4c7 100644
--- a/packages/foundry/contracts/hooks/SwapDiscountHook.sol
+++ b/packages/foundry/contracts/hooks/SwapDiscountHook.sol
@@ -3,54 +3,37 @@
 pragma solidity ^0.8.24;
 
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
-import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
-import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
+import { IBasePoolFactory } from "@balancer-labs/v3-interfaces/contracts/vault/IBasePoolFactory.sol";
 import { IHooks } from "@balancer-labs/v3-interfaces/contracts/vault/IHooks.sol";
+import { IRouterCommon } from "@balancer-labs/v3-interfaces/contracts/vault/IRouterCommon.sol";
 import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
-import { IRouterCommon } from "@balancer-labs/v3-interfaces/contracts/vault/IRouterCommon.sol"; // Importing IRouterCommon
 import {
-    AddLiquidityKind,
     LiquidityManagement,
-    RemoveLiquidityKind,
-    AfterSwapParams,
-    SwapKind,
     TokenConfig,
+    PoolSwapParams,
     HookFlags
 } from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
-import { IBasePoolFactory } from "@balancer-labs/v3-interfaces/contracts/vault/IBasePoolFactory.sol";
 
-import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
+import { VaultGuard } from "@balancer-labs/v3-vault/contracts/VaultGuard.sol";
 import { BaseHooks } from "@balancer-labs/v3-vault/contracts/BaseHooks.sol";
 
-contract SwapDiscountHook is BaseHooks {
-    using FixedPoint for uint256;
-    using SafeERC20 for IERC20;
-
+contract SafeSwapDiscount is BaseHooks, VaultGuard {
     address private immutable _allowedFactory;
     address private immutable _trustedRouter;
-    address public discountToken; // The token to check for discount eligibility (BAL)
-    uint256 public requiredBalance; // The balance of BAL required for discount eligibility
-    uint64 public hookSwapDiscountPercentage; // The discount percentage to apply
+    IERC20 private immutable _discountToken;
 
-    // modifier onlyVault() {
-    //     require(msg.sender == address(IVault), "Caller is not the vault");
-    //     _;
-    // }
+    event SwapDiscountHookRegistered(address indexed hooksContract, address indexed factory, address indexed pool);
 
-    constructor(
-        IVault vault,
-        address allowedFactory,
-        address trustedRouter,
-        address _discountToken,
-        uint64 _hookSwapDiscountPercentage,
-        uint256 _requiredBalance // Setting the required balance for discount eligibility
-    ) BaseHooks() {
+    constructor(IVault vault, address allowedFactory, address discountToken, address trustedRouter) VaultGuard(vault) {
         _allowedFactory = allowedFactory;
         _trustedRouter = trustedRouter;
-        discountToken = _discountToken;
-        hookSwapDiscountPercentage = _hookSwapDiscountPercentage;
-        requiredBalance = _requiredBalance; // Store the required balance
+        _discountToken = IERC20(discountToken);
+    }
+
+    /// @inheritdoc IHooks
+    function getHookFlags() public pure override returns (HookFlags memory hookFlags) {
+        hookFlags.shouldCallComputeDynamicSwapFee = true;
     }
 
     /// @inheritdoc IHooks
@@ -59,37 +42,28 @@ contract SwapDiscountHook is BaseHooks {
         address pool,
         TokenConfig[] memory,
         LiquidityManagement calldata
-    ) public view override returns (bool) {
-        return factory == _allowedFactory && IBasePoolFactory(factory).isPoolFromFactory(pool);
-    }
+    ) public override onlyVault returns (bool) {
+        emit SwapDiscountHookRegistered(address(this), factory, pool);
 
-    /// @inheritdoc IHooks
-    function getHookFlags() public pure override returns (HookFlags memory hookFlags) {
-        hookFlags.shouldCallAfterSwap = true;
-        return hookFlags;
+        return factory == _allowedFactory && IBasePoolFactory(factory).isPoolFromFactory(pool);
     }
 
     /// @inheritdoc IHooks
-    function onAfterSwap(
-        AfterSwapParams calldata params
-    ) public view override returns (bool success, uint256 discountedAmount) {
-        discountedAmount = params.amountCalculatedRaw;
+    function onComputeDynamicSwapFeePercentage(
+        PoolSwapParams calldata params,
+        address,
+        uint256 staticSwapFeePercentage
+    ) public view override onlyVault returns (bool, uint256) {
+        if (params.router != _trustedRouter) {
+            return (true, staticSwapFeePercentage);
+        }
 
-        // Check if the user holds enough BAL tokens for a discount
-        if (
-            hookSwapDiscountPercentage > 0 &&
-            address(params.tokenOut) == discountToken &&
-            params.kind == SwapKind.EXACT_IN
-        ) {
-            // Get the sender's balance of the discount token
-            uint256 userBalance = IERC20(discountToken).balanceOf(IRouterCommon(params.router).getSender());
+        address user = IRouterCommon(params.router).getSender();
 
-            // Apply discount if user holds enough BAL tokens
-            if (userBalance >= requiredBalance) {
-                discountedAmount = discountedAmount.mulDown(hookSwapDiscountPercentage);
-                return (true, discountedAmount);
-            }
+        if (_discountToken.balanceOf(user) > 0) {
+            return (true, staticSwapFeePercentage / 2);
         }
-        return (true, discountedAmount); // No discount applied, return original amount
+
+        return (true, staticSwapFeePercentage);
     }
 }

From 8ce7ba69afb16cb93ab5e466947a9529a17e7aca Mon Sep 17 00:00:00 2001
From: Happy <gillharpreetsingh211@gmail.com>
Date: Sat, 19 Oct 2024 15:20:27 +0530
Subject: [PATCH 7/7] Update SafeSwapDiscount.t.sol

---
 packages/foundry/test/SafeSwapDiscount.t.sol | 215 ++++++++++++++-----
 1 file changed, 158 insertions(+), 57 deletions(-)

diff --git a/packages/foundry/test/SafeSwapDiscount.t.sol b/packages/foundry/test/SafeSwapDiscount.t.sol
index 0b09a60a..4e1eea0c 100644
--- a/packages/foundry/test/SafeSwapDiscount.t.sol
+++ b/packages/foundry/test/SafeSwapDiscount.t.sol
@@ -4,35 +4,37 @@ pragma solidity ^0.8.24;
 
 import "forge-std/Test.sol";
 
-import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
 import { IVaultAdmin } from "@balancer-labs/v3-interfaces/contracts/vault/IVaultAdmin.sol";
 import { IVaultErrors } from "@balancer-labs/v3-interfaces/contracts/vault/IVaultErrors.sol";
 import {
     HooksConfig,
     LiquidityManagement,
+    PoolRoleAccounts,
     TokenConfig
 } from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
 
 import { CastingHelpers } from "@balancer-labs/v3-solidity-utils/contracts/helpers/CastingHelpers.sol";
+import { ArrayHelpers } from "@balancer-labs/v3-solidity-utils/contracts/test/ArrayHelpers.sol";
 import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
+
 import { BaseVaultTest } from "@balancer-labs/v3-vault/test/foundry/utils/BaseVaultTest.sol";
 import { PoolMock } from "@balancer-labs/v3-vault/contracts/test/PoolMock.sol";
 import { PoolFactoryMock } from "@balancer-labs/v3-vault/contracts/test/PoolFactoryMock.sol";
 import { RouterMock } from "@balancer-labs/v3-vault/contracts/test/RouterMock.sol";
 
-import { SwapDiscountHook } from "packages/foundry/contracts/hooks/SwapDiscountHook.sol";
-import { console } from "forge-std/console.sol";
+import { SafeSwapDiscount } from "../contracts/hooks/SafeSwapDiscount.sol";
 
 contract SwapDiscountHookTest is BaseVaultTest {
     using CastingHelpers for address[];
     using FixedPoint for uint256;
+    using ArrayHelpers for *;
 
     uint256 internal daiIdx;
     uint256 internal usdcIdx;
 
-    uint64 public constant MAX_SWAP_DISCOUNT_PERCENTAGE = 50e16; // 50%
-    uint256 public constant REQUIRED_BALANCE = 100e18; // Minimum required balance for discount (100 BAL)
+    // Maximum swap fee of 10%
+    uint64 public constant MAX_SWAP_FEE_PERCENTAGE = 10e16;
 
     address payable internal trustedRouter;
 
@@ -41,100 +43,199 @@ contract SwapDiscountHookTest is BaseVaultTest {
 
         (daiIdx, usdcIdx) = getSortedIndexes(address(dai), address(usdc));
 
-        // Mint 100 BAL tokens for the test account (Bob)
-        deal(address(dai), bob, REQUIRED_BALANCE);
+        // Grants LP the ability to change the static swap fee percentage.
+        authorizer.grantRole(vault.getActionId(IVaultAdmin.setStaticSwapFeePercentage.selector), lp);
     }
 
     function createHook() internal override returns (address) {
         trustedRouter = payable(router);
 
-        // LP will be the owner of the hook. Only LP is able to set hook fee percentages.
+        // lp will be the owner of the hook. Only LP is able to set hook fee percentages.
         vm.prank(lp);
-        address swapDiscountHook = address(
-            new SwapDiscountHook(
-                IVault(address(vault)),
-                address(factoryMock),
-                trustedRouter,
-                address(dai),
-                MAX_SWAP_DISCOUNT_PERCENTAGE,
-                REQUIRED_BALANCE
+        address veBALFeeHook = address(
+            new SafeSwapDiscount(IVault(address(vault)), address(factoryMock), address(veBAL), trustedRouter)
+        );
+        vm.label(veBALFeeHook, "VeBAL Fee Hook");
+        return veBALFeeHook;
+    }
+
+    function testRegistryWithWrongFactory() public {
+        address veBALFeePool = _createPoolToRegister();
+        TokenConfig[] memory tokenConfig = vault.buildTokenConfig(
+            [address(dai), address(usdc)].toMemoryArray().asIERC20()
+        );
+
+        uint32 pauseWindowEndTime = IVaultAdmin(address(vault)).getPauseWindowEndTime();
+        uint32 bufferPeriodDuration = IVaultAdmin(address(vault)).getBufferPeriodDuration();
+        uint32 pauseWindowDuration = pauseWindowEndTime - bufferPeriodDuration;
+        address unauthorizedFactory = address(new PoolFactoryMock(IVault(address(vault)), pauseWindowDuration));
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IVaultErrors.HookRegistrationFailed.selector,
+                poolHooksContract,
+                veBALFeePool,
+                unauthorizedFactory
+            )
+        );
+        _registerPoolWithHook(veBALFeePool, tokenConfig, unauthorizedFactory);
+    }
+
+    function testCreationWithWrongFactory() public {
+        address veBALFeePool = _createPoolToRegister();
+        TokenConfig[] memory tokenConfig = vault.buildTokenConfig(
+            [address(dai), address(usdc)].toMemoryArray().asIERC20()
+        );
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IVaultErrors.HookRegistrationFailed.selector,
+                poolHooksContract,
+                veBALFeePool,
+                address(factoryMock)
             )
         );
-        vm.label(swapDiscountHook, "Swap Discount Hook");
-        return swapDiscountHook;
+        _registerPoolWithHook(veBALFeePool, tokenConfig, address(factoryMock));
     }
 
-    function testSuccessfulRegistrySwap() public {
-        // Registering with allowed factory
-        address swapDHookPool = factoryMock.createPool("Test Pool", "TEST");
+    function testSuccessfulRegistry() public {
+        // Register with the allowed factory.
+        address veBALFeePool = factoryMock.createPool("Test Pool", "TEST");
         TokenConfig[] memory tokenConfig = vault.buildTokenConfig(
             [address(dai), address(usdc)].toMemoryArray().asIERC20()
         );
 
-        _registerPoolWithHook(swapDHookPool, tokenConfig, address(factoryMock));
+        vm.expectEmit();
+        emit SafeSwapDiscount.SwapDiscountHookRegistered(
+            poolHooksContract,
+            address(factoryMock),
+            veBALFeePool
+        );
 
-        HooksConfig memory hooksConfig = vault.getHooksConfig(swapDHookPool);
+        _registerPoolWithHook(veBALFeePool, tokenConfig, address(factoryMock));
+
+        HooksConfig memory hooksConfig = vault.getHooksConfig(veBALFeePool);
 
         assertEq(hooksConfig.hooksContract, poolHooksContract, "Wrong poolHooksContract");
-        assertEq(hooksConfig.shouldCallAfterSwap, true, "shouldCallAfterSwap is false");
+        assertEq(hooksConfig.shouldCallComputeDynamicSwapFee, true, "shouldCallComputeDynamicSwapFee is false");
     }
 
-    function testSwapDiscountWithSufficientBalance() public {
-        // Bob has enough DAI for the discount
-        vm.prank(bob);
-        RouterMock(trustedRouter).swapSingleTokenExactIn(
-            pool,
-            usdc,
-            dai,
-            REQUIRED_BALANCE,
-            REQUIRED_BALANCE,
-            MAX_UINT256,
-            false,
-            bytes("")
-        );
+    function testSwapWithoutVeBal() public {
+        assertEq(veBAL.balanceOf(bob), 0, "Bob still has veBAL");
 
-        // Check Bob's DAI balance after the swap
-        uint256 expectedDAIAmount = REQUIRED_BALANCE.mulDown(1 - MAX_SWAP_DISCOUNT_PERCENTAGE);
-        assertEq(dai.balanceOf(bob), expectedDAIAmount, "Bob's DAI balance is incorrect after discount");
+        _doSwapAndCheckBalances(trustedRouter);
     }
 
-    function testSwapDiscountWithInsufficientBalance() public {
-        // Adjust Bob's DAI balance to be less than the required balance for discount
+    function testSwapWithVeBal() public {
+        // Mint 1 veBAL to Bob, so he's able to receive the fee discount.
+        veBAL.mint(bob, 1);
+        assertGt(veBAL.balanceOf(bob), 0, "Bob does not have veBAL");
+
+        _doSwapAndCheckBalances(trustedRouter);
+    }
+
+    function testSwapWithVeBalAndUntrustedRouter() public {
+        // Mint 1 veBAL to Bob, so he's able to receive the fee discount.
+        veBAL.mint(bob, 1);
+        assertGt(veBAL.balanceOf(bob), 0, "Bob does not have veBAL");
+
+        // Create an untrusted router
+        address payable untrustedRouter = payable(new RouterMock(IVault(address(vault)), weth, permit2));
+        vm.label(untrustedRouter, "untrusted router");
+
+        // Allows permit2 to move DAI tokens from Bob to untrustedRouter.
         vm.prank(bob);
-        deal(address(dai), bob, REQUIRED_BALANCE / 2); // Only 50 DAI
+        permit2.approve(address(dai), untrustedRouter, type(uint160).max, type(uint48).max);
+
+        // Even if Bob has veBAL, since he is using an untrusted router, he will get no discount.
+        _doSwapAndCheckBalances(untrustedRouter);
+    }
+
+    function _doSwapAndCheckBalances(address payable routerToUse) private {
+        // Since the Vault has no swap fee, the fee will stay in the pool.
+        uint256 swapFeePercentage = MAX_SWAP_FEE_PERCENTAGE;
+
+        vm.prank(lp);
+        vault.setStaticSwapFeePercentage(pool, swapFeePercentage);
+
+        uint256 exactAmountIn = poolInitAmount / 100;
+        // PoolMock uses linear math with a rate of 1, so amountIn == amountOut when no fees are applied.
+        uint256 expectedAmountOut = exactAmountIn;
+        // If Bob has veBAL and the router is trusted, Bob gets a 50% discount.
+        bool shouldGetDiscount = routerToUse == trustedRouter && veBAL.balanceOf(bob) > 0;
+        uint256 expectedHookFee = exactAmountIn.mulDown(swapFeePercentage) / (shouldGetDiscount ? 2 : 1);
+        // The hook fee will remain in the pool, so the expected amountOut discounts the fees.
+        expectedAmountOut -= expectedHookFee;
+
+        BaseVaultTest.Balances memory balancesBefore = getBalances(bob);
 
         vm.prank(bob);
-        RouterMock(trustedRouter).swapSingleTokenExactIn(
+        RouterMock(routerToUse).swapSingleTokenExactIn(
             pool,
-            usdc,
             dai,
-            REQUIRED_BALANCE,
-            REQUIRED_BALANCE,
+            usdc,
+            exactAmountIn,
+            expectedAmountOut,
             MAX_UINT256,
             false,
             bytes("")
         );
 
-        // Check Bob's DAI balance after the swap
+        BaseVaultTest.Balances memory balancesAfter = getBalances(bob);
+
+        // Bob's balance of DAI is supposed to decrease, since DAI is the token in
         assertEq(
-            dai.balanceOf(bob),
-            REQUIRED_BALANCE,
-            "Bob's DAI balance should be unchanged due to insufficient balance"
+            balancesBefore.userTokens[daiIdx] - balancesAfter.userTokens[daiIdx],
+            exactAmountIn,
+            "Bob's DAI balance is wrong"
+        );
+        // Bob's balance of USDC is supposed to increase, since USDC is the token out
+        assertEq(
+            balancesAfter.userTokens[usdcIdx] - balancesBefore.userTokens[usdcIdx],
+            expectedAmountOut,
+            "Bob's USDC balance is wrong"
+        );
+
+        // Vault's balance of DAI is supposed to increase, since DAI was added by Bob
+        assertEq(
+            balancesAfter.vaultTokens[daiIdx] - balancesBefore.vaultTokens[daiIdx],
+            exactAmountIn,
+            "Vault's DAI balance is wrong"
+        );
+        // Vault's balance of USDC is supposed to decrease, since USDC was given to Bob
+        assertEq(
+            balancesBefore.vaultTokens[usdcIdx] - balancesAfter.vaultTokens[usdcIdx],
+            expectedAmountOut,
+            "Vault's USDC balance is wrong"
+        );
+
+        // Pool deltas should equal vault's deltas
+        assertEq(
+            balancesAfter.poolTokens[daiIdx] - balancesBefore.poolTokens[daiIdx],
+            exactAmountIn,
+            "Pool's DAI balance is wrong"
+        );
+        assertEq(
+            balancesBefore.poolTokens[usdcIdx] - balancesAfter.poolTokens[usdcIdx],
+            expectedAmountOut,
+            "Pool's USDC balance is wrong"
         );
     }
 
-    // Registry tests require a new pool, because an existing pool may already be registered
+    // Registry tests require a new pool, because an existing pool may be already registered
     function _createPoolToRegister() private returns (address newPool) {
-        newPool = address(new PoolMock(IVault(address(vault)), "SwapD Hook Pool", "swapDHookPool"));
-        vm.label(newPool, "SwapD Hook Pool");
+        newPool = address(new PoolMock(IVault(address(vault)), "VeBAL Fee Pool", "veBALFeePool"));
+        vm.label(newPool, "VeBAL Fee Pool");
     }
 
-    function _registerPoolWithHook(address swapDhookPool, TokenConfig[] memory tokenConfig, address factory) private {
+    function _registerPoolWithHook(address exitFeePool, TokenConfig[] memory tokenConfig, address factory) private {
+        PoolRoleAccounts memory roleAccounts;
         LiquidityManagement memory liquidityManagement;
+
         PoolFactoryMock(factory).registerPool(
-            swapDhookPool,
+            exitFeePool,
             tokenConfig,
-            PoolRoleAccounts({ lp: lp, guardian: address(0), rewards: address(0) }),
+            roleAccounts,
             poolHooksContract,
             liquidityManagement
         );