feat: Add referral code, events, NatSpec, README (SC-448) (#4)

* feat: first test working

* feat: use larger numbers:

* feat: test with initial burn amount passing

* feat: update tests to work with updated burn logic, move conversion functions around and use previews

* feat: remove todos

* fix: update to remove console and update comment

* feat: get swap tests working

* feat: get all swap tests working

* fix: update for three assets in logic

* feat: all tests passing

* fix: rm commented out test

* feat: add preview swap tests

* feat: move logic out of single use internal and use conversion rate everywhere

* feat: move divRoundUp out of single use internal

* feat: add full coverage for conversion tests

* feat: add more preview cases

* feat: refactor PSM to use three assets

* fix: rm comment

* feat: add interface, natspec, events, referral code, tests passing

* fix: update to rm consolegp

* feat: add events testing

* feat: make precisions internal and add state var natspec

* feat: finish natspec

* feat: add readme

* feat: add referral code note

* fix: update constructor test

* fix: update links

* fix: reformatting

* fix: update testing section

* fix: improve overview

* feat: add emojis

* feat: remove all share burn logic, get all non inflation attack tests to pass

* fix: cleanup diff

* fix: update to use initial deposit instead of burn

* feat: add readme section explaining attack

* fix: minimize diff

* fix: address bartek comments

* feat: update to address comments outside sharesToBurn

* feat: update inflation attack test and readme

* fix: update readme

* feat: update test to constrain deposit/withdraw

* feat: update to add both cases

* feat: update per review

* feat: update to use underscore bound, fix test

* feat: add overrides, remove referrals, update referral type

* fix: update expect emit

* fix: overrides

README.md

@@ -1,31 +1,90 @@
-# Spark PSM
+# ⚡ Spark PSM ⚡
 
-![Foundry CI](https://github.com/mars-foundation/spark-psm/actions/workflows/ci.yml/badge.svg)
+![Foundry CI](https://github.com/marsfoundation/spark-psm/actions/workflows/ci.yml/badge.svg)
 [![Foundry][foundry-badge]][foundry]
-[![License: AGPL v3](https://img.shields.io/badge/License-AGPL%20v3-blue.svg)](https://github.com/mars-foundation/spark-psm/blob/master/LICENSE)
+[![License: AGPL v3](https://img.shields.io/badge/License-AGPL%20v3-blue.svg)](https://github.com/marsfoundation/spark-psm/blob/master/LICENSE)
 
 [foundry]: https://getfoundry.sh/
 [foundry-badge]: https://img.shields.io/badge/Built%20with-Foundry-FFDB1C.svg
 
-PSM contracts to either:
-- Convert between a tokenization of an asset (ex. USDC) and a yield-bearing version of the asset (ex. sDAI).
-- Convert one to one between directly correlated assets (ex. USDC-DAI).
+## Overview
+
+This repository contains the implementation of a Peg Stability Module (PSM) contract, which facilitates the swapping, depositing, and withdrawing of three given assets to maintain stability and ensure the peg of involved assets. The PSM supports both yield-bearing and non-yield-bearing assets.
+
+`asset0` and `asset1` are two ERC20 tokens that are directly correlated and are non-yield-bearing, referred to as "base assets". `asset2` is a yield-bearing version of both `asset0` and `asset1`. The PSM contract allows users to swap between these assets, deposit any of the assets to mint shares, and withdraw any of the assets by burning shares.
+
+The conversion between a base asset and `asset2` is provided by a rate provider contract. The rate provider returns the conversion rate between `asset2` and the base asset in 1e27 precision. The conversion between the base assets is one to one.
+
+The conversion rate between assets and shares is based on the total value of assets held by the PSM. The total value is calculated by converting the assets to their equivalent value in the base asset with 18 decimal precision. The shares represent the ownership of the underlying assets in the PSM. Since three assets are used, each with different precisions and values, they are converted to a common base asset-denominated value for share conversions.
+
+For detailed implementation, refer to the contract code and `IPSM` interface documentation.
+
+## Contracts
+
+- **`src/PSM.sol`**: The core contract implementing the `IPSM` interface, providing functionality for swapping, depositing, and withdrawing assets.
+- **`src/interfaces/IPSM.sol`**: Defines the essential functions and events that the PSM contract implements.
 
 ## [CRITICAL]: First Depositor Attack Prevention on Deployment
 
 On the deployment of the PSM, the deployer **MUST make an initial deposit to get AT LEAST 1e18 shares in order to protect the first depositor from getting attacked with a share inflation attack**. This is outlined further [here](https://github.com/marsfoundation/spark-automations/assets/44272939/9472a6d2-0361-48b0-b534-96a0614330d3). Technical details related to this can be found in `test/InflationAttack.t.sol`. The deployment script [TODO] in this repo contains logic for the deployer to perform this initial deposit, so it is **HIGHLY RECOMMENDED** to use this deployment script when deploying the PSM. Reasoning for the technical implementation approach taken is outlined in more detail [here](https://github.com/marsfoundation/spark-psm/pull/2).
 
-## Usage
+## PSM Contract Details
 
-```bash
-forge build
-```
+### State Variables and Immutables
+
+- **`asset0`**: Non-yield-bearing base asset (e.g., USDC).
+- **`asset1`**: Another non-yield-bearing base asset that is directly correlated to `asset0` (e.g., DAI).
+- **`asset2`**: Yield-bearing version of both `asset0` and `asset1` (e.g., sDAI).
+- **`rateProvider`**: Contract that returns a conversion rate between and `asset2` and the base asset (e.g., sDAI to USD) in 1e27 precision.
+- **`totalShares`**: Total shares in the PSM. Shares represent the ownership of the underlying assets in the PSM.
+- **`shares`**: Mapping of user addresses to their shares.
+
+### Functions
+
+#### Swap Functions
+
+- **`swap`**: Allows swapping of assets based on current conversion rates. Ensures the output amount meets the minimum required before executing the transfer and emitting the swap event. Includes a referral code.
 
-## Test
+#### Liquidity Provision Functions
+
+- **`deposit`**: Deposits assets into the PSM, minting new shares. Includes a referral code.
+- **`withdraw`**: Withdraws assets from the PSM by burning shares. Ensures the user has sufficient shares for the withdrawal and adjusts the total shares accordingly. Includes a referral code.
+
+#### Preview Functions
+
+- **`previewDeposit`**: Estimates the number of shares minted for a given deposit amount.
+- **`previewWithdraw`**: Estimates the number of shares burned and the amount of assets withdrawn for a specified amount.
+- **`previewSwap`**: Estimates the amount of one asset received for a given amount of another asset in a swap.
+
+#### Conversion Functions
+
+NOTE: These functions do not round in the same way as preview functions, so they are meant to be used for general quoting purposes.
+
+- **`convertToAssets`**: Converts shares to the equivalent amount of a specified asset.
+- **`convertToAssetValue`**: Converts shares to their equivalent value in base asset terms with 18 decimal precision (e.g., USD).
+- **`convertToShares`**: Converts asset values to shares based on the current exchange rate.
+
+#### Asset Value Functions
+
+- **`getPsmTotalValue`**: Returns the total value of all assets held by the PSM denominated in the base asset with 18 decimal precision. (e.g., USD).
+
+### Events
+
+- **`Swap`**: Emitted on asset swaps.
+- **`Deposit`**: Emitted on asset deposits.
+- **`Withdraw`**: Emitted on asset withdrawals.
+
+## Running Tests
+
+To run tests in this repo, run:
 
 ```bash
 forge test
 ```
 
 ***
-*The IP in this repository was assigned to Mars SPC Limited in respect of the MarsOne SP*
+*The IP in this repository was assigned to Mars SPC Limited in respect of the MarsOne SP.*
+
+<p align="center">
+  <img src="https://1827921443-files.gitbook.io/~/files/v0/b/gitbook-x-prod.appspot.com/o/spaces%2FjvdfbhgN5UCpMtP1l8r5%2Fuploads%2Fgit-blob-c029bb6c918f8c042400dbcef7102c4e5c1caf38%2Flogomark%20colour.svg?alt=media" height="150" />
+</p>

src/PSM.sol

@@ -5,35 +5,35 @@ import { IERC20 } from "erc20-helpers/interfaces/IERC20.sol";
 
 import { SafeERC20 } from "erc20-helpers/SafeERC20.sol";
 
+import { IPSM } from "src/interfaces/IPSM.sol";
+
 interface IRateProviderLike {
     function getConversionRate() external view returns (uint256);
 }
 
-// TODO: Add events and corresponding tests
 // TODO: Determine what admin functionality we want (fees?)
 // TODO: Refactor into inheritance structure
-// TODO: Add interface with natspec and inherit
 // TODO: Prove that we're always rounding against user
 // TODO: Add receiver to deposit/withdraw
-contract PSM {
+contract PSM is IPSM {
 
     using SafeERC20 for IERC20;
 
+    uint256 internal immutable _asset0Precision;
+    uint256 internal immutable _asset1Precision;
+    uint256 internal immutable _asset2Precision;
+
     // NOTE: Assumption is made that asset2 is the yield-bearing counterpart of asset0 and asset1.
     //       Examples: asset0 = USDC, asset1 = DAI, asset2 = sDAI
-    IERC20 public immutable asset0;
-    IERC20 public immutable asset1;
-    IERC20 public immutable asset2;
-
-    address public immutable rateProvider;
+    IERC20 public override immutable asset0;
+    IERC20 public override immutable asset1;
+    IERC20 public override immutable asset2;
 
-    uint256 public immutable asset0Precision;
-    uint256 public immutable asset1Precision;
-    uint256 public immutable asset2Precision;
+    address public override immutable rateProvider;
 
-    uint256 public totalShares;
+    uint256 public override totalShares;
 
-    mapping(address user => uint256 shares) public shares;
+    mapping(address user => uint256 shares) public override shares;
 
     constructor(address asset0_, address asset1_, address asset2_, address rateProvider_) {
         require(asset0_       != address(0), "PSM/invalid-asset0");
@@ -51,9 +51,9 @@ contract PSM {
 
         rateProvider = rateProvider_;
 
-        asset0Precision = 10 ** IERC20(asset0_).decimals();
-        asset1Precision = 10 ** IERC20(asset1_).decimals();
-        asset2Precision = 10 ** IERC20(asset2_).decimals();
+        _asset0Precision = 10 ** IERC20(asset0_).decimals();
+        _asset1Precision = 10 ** IERC20(asset1_).decimals();
+        _asset2Precision = 10 ** IERC20(asset2_).decimals();
     }
 
     /**********************************************************************************************/
@@ -65,9 +65,10 @@ contract PSM {
         address assetOut,
         uint256 amountIn,
         uint256 minAmountOut,
-        address receiver
+        address receiver,
+        uint256 referralCode
     )
-        external
+        external override
     {
         require(amountIn != 0,          "PSM/invalid-amountIn");
         require(receiver != address(0), "PSM/invalid-receiver");
@@ -78,25 +79,29 @@ contract PSM {
 
         IERC20(assetIn).safeTransferFrom(msg.sender, address(this), amountIn);
         IERC20(assetOut).safeTransfer(receiver, amountOut);
+
+        emit Swap(assetIn, assetOut, msg.sender, receiver, amountIn, amountOut, referralCode);
     }
 
     /**********************************************************************************************/
     /*** Liquidity provision functions                                                          ***/
     /**********************************************************************************************/
 
     function deposit(address asset, uint256 assetsToDeposit)
-        external returns (uint256 newShares)
+        external override returns (uint256 newShares)
     {
         newShares = previewDeposit(asset, assetsToDeposit);
 
         shares[msg.sender] += newShares;
         totalShares        += newShares;
 
         IERC20(asset).safeTransferFrom(msg.sender, address(this), assetsToDeposit);
+
+        emit Deposit(asset, msg.sender, assetsToDeposit, newShares);
     }
 
     function withdraw(address asset, uint256 maxAssetsToWithdraw)
-        external returns (uint256 assetsWithdrawn)
+        external override returns (uint256 assetsWithdrawn)
     {
         uint256 sharesToBurn;
 
@@ -108,21 +113,25 @@ contract PSM {
         }
 
         IERC20(asset).safeTransfer(msg.sender, assetsWithdrawn);
+
+        emit Withdraw(asset, msg.sender, assetsWithdrawn, sharesToBurn);
     }
 
     /**********************************************************************************************/
     /*** Deposit/withdraw preview functions                                                     ***/
     /**********************************************************************************************/
 
-    function previewDeposit(address asset, uint256 assetsToDeposit) public view returns (uint256) {
+    function previewDeposit(address asset, uint256 assetsToDeposit)
+        public view override returns (uint256)
+    {
         require(_isValidAsset(asset), "PSM/invalid-asset");
 
         // Convert amount to 1e18 precision denominated in value of asset0 then convert to shares.
         return convertToShares(_getAssetValue(asset, assetsToDeposit));
     }
 
     function previewWithdraw(address asset, uint256 maxAssetsToWithdraw)
-        public view returns (uint256 sharesToBurn, uint256 assetsWithdrawn)
+        public view override returns (uint256 sharesToBurn, uint256 assetsWithdrawn)
     {
         require(_isValidAsset(asset), "PSM/invalid-asset");
 
@@ -147,46 +156,48 @@ contract PSM {
     /**********************************************************************************************/
 
     function previewSwap(address assetIn, address assetOut, uint256 amountIn)
-        public view returns (uint256 amountOut)
+        public view override returns (uint256 amountOut)
     {
         if (assetIn == address(asset0)) {
-            if      (assetOut == address(asset1)) return _previewOneToOneSwap(amountIn, asset0Precision, asset1Precision);
-            else if (assetOut == address(asset2)) return _previewSwapToAsset2(amountIn, asset0Precision);
+            if      (assetOut == address(asset1)) return _previewOneToOneSwap(amountIn, _asset0Precision, _asset1Precision);
+            else if (assetOut == address(asset2)) return _previewSwapToAsset2(amountIn, _asset0Precision);
         }
 
         else if (assetIn == address(asset1)) {
-            if      (assetOut == address(asset0)) return _previewOneToOneSwap(amountIn, asset1Precision, asset0Precision);
-            else if (assetOut == address(asset2)) return _previewSwapToAsset2(amountIn, asset1Precision);
+            if      (assetOut == address(asset0)) return _previewOneToOneSwap(amountIn, _asset1Precision, _asset0Precision);
+            else if (assetOut == address(asset2)) return _previewSwapToAsset2(amountIn, _asset1Precision);
         }
 
         else if (assetIn == address(asset2)) {
-            if      (assetOut == address(asset0)) return _previewSwapFromAsset2(amountIn, asset0Precision);
-            else if (assetOut == address(asset1)) return _previewSwapFromAsset2(amountIn, asset1Precision);
+            if      (assetOut == address(asset0)) return _previewSwapFromAsset2(amountIn, _asset0Precision);
+            else if (assetOut == address(asset1)) return _previewSwapFromAsset2(amountIn, _asset1Precision);
         }
 
         revert("PSM/invalid-asset");
     }
 
     /**********************************************************************************************/
-    /*** Swap preview functions                                                                 ***/
+    /*** Conversion functions                                                                   ***/
     /**********************************************************************************************/
 
-    function convertToAssets(address asset, uint256 numShares) public view returns (uint256) {
+    function convertToAssets(address asset, uint256 numShares)
+        public view override returns (uint256)
+    {
         require(_isValidAsset(asset), "PSM/invalid-asset");
 
         uint256 assetValue = convertToAssetValue(numShares);
 
-        if      (asset == address(asset0)) return assetValue * asset0Precision / 1e18;
-        else if (asset == address(asset1)) return assetValue * asset1Precision / 1e18;
+        if      (asset == address(asset0)) return assetValue * _asset0Precision / 1e18;
+        else if (asset == address(asset1)) return assetValue * _asset1Precision / 1e18;
 
         // NOTE: Multiplying by 1e27 and dividing by 1e18 cancels to 1e9 in numerator
         return assetValue
             * 1e9
-            * asset2Precision
+            * _asset2Precision
             / IRateProviderLike(rateProvider).getConversionRate();
     }
 
-    function convertToAssetValue(uint256 numShares) public view returns (uint256) {
+    function convertToAssetValue(uint256 numShares) public view override returns (uint256) {
         uint256 totalShares_ = totalShares;
 
         if (totalShares_ != 0) {
@@ -195,15 +206,15 @@ contract PSM {
         return numShares;
     }
 
-    function convertToShares(uint256 assetValue) public view returns (uint256) {
+    function convertToShares(uint256 assetValue) public view override returns (uint256) {
         uint256 totalValue = getPsmTotalValue();
         if (totalValue != 0) {
             return assetValue * totalShares / totalValue;
         }
         return assetValue;
     }
 
-    function convertToShares(address asset, uint256 assets) public view returns (uint256) {
+    function convertToShares(address asset, uint256 assets) public view override returns (uint256) {
         require(_isValidAsset(asset), "PSM/invalid-asset");
         return convertToShares(_getAssetValue(asset, assets));
     }
@@ -212,7 +223,7 @@ contract PSM {
     /*** Asset value functions                                                                  ***/
     /**********************************************************************************************/
 
-    function getPsmTotalValue() public view returns (uint256) {
+    function getPsmTotalValue() public view override returns (uint256) {
         return _getAsset0Value(asset0.balanceOf(address(this)))
             +  _getAsset1Value(asset1.balanceOf(address(this)))
             +  _getAsset2Value(asset2.balanceOf(address(this)));
@@ -244,16 +255,19 @@ contract PSM {
     }
 
     function _getAsset0Value(uint256 amount) internal view returns (uint256) {
-        return amount * 1e18 / asset0Precision;
+        return amount * 1e18 / _asset0Precision;
     }
 
     function _getAsset1Value(uint256 amount) internal view returns (uint256) {
-        return amount * 1e18 / asset1Precision;
+        return amount * 1e18 / _asset1Precision;
     }
 
     function _getAsset2Value(uint256 amount) internal view returns (uint256) {
-        // NOTE: Multiplying by 1e18 and dividing by 1e9 cancels to 1e9 in denominator
-        return amount * IRateProviderLike(rateProvider).getConversionRate() / 1e9 / asset2Precision;
+        // NOTE: Multiplying by 1e18 and dividing by 1e27 cancels to 1e9 in denominator
+        return amount
+            * IRateProviderLike(rateProvider).getConversionRate()
+            / 1e9
+            / _asset2Precision;
     }
 
     function _isValidAsset(address asset) internal view returns (bool) {
@@ -266,7 +280,7 @@ contract PSM {
         return amountIn
             * 1e27
             / IRateProviderLike(rateProvider).getConversionRate()
-            * asset2Precision
+            * _asset2Precision
             / assetInPrecision;
     }
 
@@ -277,7 +291,7 @@ contract PSM {
             * IRateProviderLike(rateProvider).getConversionRate()
             / 1e27
             * assetInPrecision
-            / asset2Precision;
+            / _asset2Precision;
     }
 
     function _previewOneToOneSwap(