Foundry Snapshots

Vault.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

// From https://github.com/transmissions11/solmate/blob/main/src/tokens/ERC20.sol
import {ERC20} from './ERC20.sol';

contract Vault is ERC20 {
    // Underlying asset
    ERC20 public token;

    constructor(address _token) ERC20("VaultToken", "VLT", 18) {
        token = ERC20(_token);
    }

    function deposit(uint256 assetAmount) external {
        // Calculate share amount
        uint256 totalAssets = token.balanceOf(address(this));
        uint256 shareAmount = assetAmount;
        if (totalAssets != 0) {
            shareAmount = assetAmount * totalSupply / totalAssets;
        }
        
        // Mint shares to depositor 
        _mint(msg.sender, shareAmount);

        // Transfer underlying asset from depositor to vault
        token.transferFrom(msg.sender, address(this), assetAmount);
    }

    function withdraw(uint256 shareAmount) external {
        // Calculate asset amount
        uint256 totalAssets = token.balanceOf(address(this));
        uint256 assetAmount = shareAmount * totalAssets / totalSupply;

        // Burn shares from depositor
        _burn(msg.sender, shareAmount);

        // Transfer underlying asset from vault to depositor
        token.transfer(msg.sender, assetAmount);
    }
}

Vault.t.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import {Test, console2} from "forge-std/Test.sol";
import {Vault} from "../src/Counter.sol";
import {ERC20} from "../src/ERC20.sol";

contract VaultTest is Test {
    ERC20 token;
    Vault vault;

    function setUp() public {
        // Deploy token and vault contracts
        token = new Token();
        vault = new Vault(address(token));

        // Mint 100 tokens to this address
        deal(address(token), address(this), 100e18);
    }

    function testSnapshot() public {
        // Deposit 100 tokens
        token.approve(address(vault), 100e18);
        vault.deposit(100e18);

        // Take a snapshot of the current state
        uint256 snapshot = vm.snapshot();

        // Withdraw 100 tokens
        vault.withdraw(100e18);
        uint256 sharesAfterFirstWithdrawal = vault.balanceOf(address(this));
        uint256 tokensAfterFirstWithdrawal = token.balanceOf(address(this));

        // Restore the state until after the deposit
        vm.revertTo(snapshot);

        // Withdraw 50 tokens twice
        vault.withdraw(50e18);
        vault.withdraw(50e18);
        uint256 sharesAfterSecondWithdrawal = vault.balanceOf(address(this));
        uint256 tokensAfterSecondWithdrawal = token.balanceOf(address(this));

        // Ensure the end state is the same
        assertEq(sharesAfterFirstWithdrawal, sharesAfterSecondWithdrawal);
        assertEq(tokensAfterFirstWithdrawal, tokensAfterSecondWithdrawal);
    }
}

contract Token is ERC20("", "", 18) {}

Test passed
Knowledge gained

Thank you so much