update implementation TokenBundle

Author: Krishang Nadgauda

contracts/feature/TokenBundle.sol

@@ -4,51 +4,73 @@ pragma solidity ^0.8.0;
 import "./interface/ITokenBundle.sol";
 
 abstract contract TokenBundle is ITokenBundle {
-    /// @dev UID => asset count, bundle uri, and tokens contained in the bundle
+
+    /// @dev Mapping from bundle UID => bundle info.
     mapping(uint256 => BundleInfo) private bundle;
 
-    /// @dev Returns the count of assets in a bundle, given a bundle Id.
-    function getTokenCount(uint256 _bundleId) public view returns (uint256) {
+    /// @dev Returns the total number of assets in a particular bundle.
+    function getTokenCountOfBundle(uint256 _bundleId) public view returns (uint256) {
         return bundle[_bundleId].count;
     }
 
-    /// @dev Returns a token contained in a bundle, given a bundle Id and index of token.
-    function getToken(uint256 _bundleId, uint256 index) public view returns (Token memory) {
+    /// @dev Returns an asset contained in a particular bundle, at a particular index.
+    function getTokenOfBundle(uint256 _bundleId, uint256 index) public view returns (Token memory) {
         return bundle[_bundleId].tokens[index];
     }
 
-    /// @dev Returns the uri of bundle for a particular bundle Id.
-    function getUri(uint256 _bundleId) public view returns (string memory) {
+    /// @dev Returns the uri of a particular bundle.
+    function getUriOfBundle(uint256 _bundleId) public view returns (string memory) {
         return bundle[_bundleId].uri;
     }
 
     /// @dev Lets the calling contract create/update a bundle, by passing in a list of tokens and a unique id.
     function _setBundle(Token[] calldata _tokensToBind, uint256 _bundleId) internal {
-        // uint256 _bundleId = _getNextBundleId();
-        require(_tokensToBind.length > 0, "no tokens to bind");
-        for (uint256 i = 0; i < _tokensToBind.length; i += 1) {
-            bundle[_bundleId].tokens[i] = _tokensToBind[i];
+
+        require(_tokensToBind.length > 0, "TokenBundle: no tokens to bind.");
+
+        uint256 currentCount = bundle[_bundleId].count;
+        uint256 targetCount = _tokensToBind.length;
+        uint256 check = currentCount > targetCount ? currentCount : targetCount;
+        
+        for (uint256 i = 0; i < check; i += 1) {
+
+            if(i < targetCount) {
+                bundle[_bundleId].tokens[i] = _tokensToBind[i];
+            } else if (i < currentCount) {
+                delete bundle[_bundleId].tokens[i];
+            }
         }
-        bundle[_bundleId].count = _tokensToBind.length;
+
+        bundle[_bundleId].count = targetCount;
+    }
+
+    /// @dev Lets the calling contract add a token to a bundle for a unique bundle id and index.
+    function _addTokenInBundle(
+        Token memory _tokenToBind,
+        uint256 _bundleId
+    ) internal {
+        uint256 id = bundle[_bundleId].count;
+
+        bundle[_bundleId].tokens[id] = _tokenToBind;
+        bundle[_bundleId].count += 1;
     }
 
-    /// @dev Lets the calling contract set/update a token in a bundle for a unique bundle id and index.
-    function _setBundleToken(
+    /// @dev Lets the calling contract update a token in a bundle for a unique bundle id and index.
+    function _updateTokenInBundle(
         Token memory _tokenToBind,
         uint256 _bundleId,
-        uint256 index,
-        bool isUpdate
+        uint256 _index
     ) internal {
-        bundle[_bundleId].tokens[index] = _tokenToBind;
-        bundle[_bundleId].count += isUpdate ? 0 : 1;
+        require(_index < bundle[_bundleId].count, "TokenBundle: index DNE.");
+        bundle[_bundleId].tokens[_index] = _tokenToBind;
     }
 
-    /// @dev Lets the calling contract set/update the bundle uri for a particular bundle id.
-    function _setUri(string calldata _uri, uint256 _bundleId) internal {
+    /// @dev Lets the calling contract set/update the uri of a particular bundle.
+    function _setUriOfBundle(string calldata _uri, uint256 _bundleId) internal {
         bundle[_bundleId].uri = _uri;
     }
 
-    /// @dev Lets the calling contract delete a bundle with a given id.
+    /// @dev Lets the calling contract delete a particular bundle.
     function _deleteBundle(uint256 _bundleId) internal {
         delete bundle[_bundleId];
     }