From a78d0ee72b1b5f247356ac2b3b9e9b4d7f011ce7 Mon Sep 17 00:00:00 2001
From: Yash Goyal <yashgo0018@gmail.com>
Date: Sat, 16 Dec 2023 14:25:26 +0800
Subject: [PATCH] fixed a bug

---
 contracts/Vault.sol | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/contracts/Vault.sol b/contracts/Vault.sol
index b4d02c0..004fdf7 100644
--- a/contracts/Vault.sol
+++ b/contracts/Vault.sol
@@ -98,8 +98,14 @@ contract Vault {
     ) public view returns (uint256 share) {
         uint256 totalValue = calculateTotalValue();
         uint256 currentTotalSupply = shareToken.totalSupply();
-        if (totalValue == 0 || currentTotalSupply == 0) share = amount;
-        else share = (currentTotalSupply * amount) / totalValue;
+
+        if (totalValue == 0 || currentTotalSupply == 0) {
+            uint256 stableTokenUnit = 10 ** stableToken.decimals();
+            uint256 shareTokenUnit = 10 ** shareToken.decimals();
+            share = (amount * shareTokenUnit) / stableTokenUnit;
+        } else {
+            share = (currentTotalSupply * amount) / totalValue;
+        }
     }
 
     function calculateTotalValue() public view returns (uint256 total) {