feat: merge createDataset and addPieces

Author: Chaitu-Tatipamula

src/PDPVerifier.sol

@@ -52,6 +52,7 @@ contract PDPVerifier is Initializable, UUPSUpgradeable, OwnableUpgradeable {
     bytes32 public constant FIL_USD_PRICE_FEED_ID = 0x150ac9b959aee0051e4091f0ef5216d941f590e1c5e7f91cf7635b5c11628c0e;
     uint256 public constant NO_CHALLENGE_SCHEDULED = 0;
     uint256 public constant NO_PROVEN_EPOCH = 0;
+    uint256 public constant NEW_DATA_SET_SENTINEL = type(uint256).max;
 
     // Events
     event DataSetCreated(uint256 indexed setId, address indexed storageProvider);
@@ -382,35 +383,6 @@ contract PDPVerifier is Initializable, UUPSUpgradeable, OwnableUpgradeable {
         }
     }
 
-    // A data set is created empty, with no pieces. Creation yields a data set ID
-    // for referring to the data set later.
-    // Sender of create message is storage provider.
-    function createDataSet(address listenerAddr, bytes calldata extraData) public payable returns (uint256) {
-        require(extraData.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
-        uint256 sybilFee = PDPFees.sybilFee();
-        require(msg.value >= sybilFee, "sybil fee not met");
-        burnFee(sybilFee);
-
-        uint256 setId = nextDataSetId++;
-        dataSetLeafCount[setId] = 0;
-        nextChallengeEpoch[setId] = NO_CHALLENGE_SCHEDULED; // Initialized on first call to NextProvingPeriod
-        storageProvider[setId] = msg.sender;
-        dataSetListener[setId] = listenerAddr;
-        dataSetLastProvenEpoch[setId] = NO_PROVEN_EPOCH;
-
-        if (listenerAddr != address(0)) {
-            PDPListener(listenerAddr).dataSetCreated(setId, msg.sender, extraData);
-        }
-        emit DataSetCreated(setId, msg.sender);
-
-        // Return the at the end to avoid any possible re-entrency issues.
-        if (msg.value > sybilFee) {
-            (bool success,) = msg.sender.call{value: msg.value - sybilFee}("");
-            require(success, "Transfer failed.");
-        }
-        return setId;
-    }
-
     // Removes a data set. Must be called by the storage provider.
     function deleteDataSet(uint256 setId, bytes calldata extraData) public {
         require(extraData.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
@@ -432,35 +404,90 @@ contract PDPVerifier is Initializable, UUPSUpgradeable, OwnableUpgradeable {
         emit DataSetDeleted(setId, deletedLeafCount);
     }
 
-    // Appends new pieces to the collection managed by a data set.
-    // These pieces won't be challenged until the next proving period is
-    // started by calling nextProvingPeriod.
-    function addPieces(uint256 setId, Cids.Cid[] calldata pieceData, bytes calldata extraData)
+    // Create Dataset and Add Pieces, When setId == NEW_DATA_SET_SENTINEL, this will create a new dataset with piece data provided
+    // with the provided listenerAddr and expect extraData to be abi.encode(bytes createPayload, bytes addPayload).
+    // When adding to an existing set, pass listenerAddr == address(0) and setId to the live dataset.
+    function addPieces(uint256 setId, address listenerAddr, Cids.Cid[] calldata pieceData, bytes calldata extraData)
         public
+        payable
         returns (uint256)
     {
-        uint256 nPieces = pieceData.length;
-        require(extraData.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
-        require(dataSetLive(setId), "Data set not live");
-        require(nPieces > 0, "Must add at least one piece");
-        require(storageProvider[setId] == msg.sender, "Only the storage provider can add pieces");
-        uint256 firstAdded = nextPieceId[setId];
-        uint256[] memory pieceIds = new uint256[](pieceData.length);
-        Cids.Cid[] memory pieceCidsAdded = new Cids.Cid[](pieceData.length);
-
-        for (uint256 i = 0; i < nPieces; i++) {
-            addOnePiece(setId, i, pieceData[i]);
-            pieceIds[i] = firstAdded + i;
-            pieceCidsAdded[i] = pieceData[i];
-        }
-        emit PiecesAdded(setId, pieceIds, pieceCidsAdded);
+        if (setId == NEW_DATA_SET_SENTINEL) {
+            (bytes memory createPayload, bytes memory addPayload) = abi.decode(extraData, (bytes, bytes));
+
+            require(createPayload.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
+            uint256 sybilFee = PDPFees.sybilFee();
+            require(msg.value >= sybilFee, "sybil fee not met");
+            burnFee(sybilFee);
+
+            require(listenerAddr != address(0), "listener required for new dataset");
+            uint256 newSetId = nextDataSetId++;
+            dataSetLeafCount[newSetId] = 0;
+            nextChallengeEpoch[newSetId] = NO_CHALLENGE_SCHEDULED; // Initialized on first call to NextProvingPeriod
+            storageProvider[newSetId] = msg.sender;
+            dataSetListener[newSetId] = listenerAddr;
+            dataSetLastProvenEpoch[newSetId] = NO_PROVEN_EPOCH;
 
-        address listenerAddr = dataSetListener[setId];
-        if (listenerAddr != address(0)) {
-            PDPListener(listenerAddr).piecesAdded(setId, firstAdded, pieceData, extraData);
-        }
+            if (listenerAddr != address(0)) {
+                PDPListener(listenerAddr).dataSetCreated(newSetId, msg.sender, createPayload);
+            }
+            emit DataSetCreated(newSetId, msg.sender);
+
+            // Add pieces to the newly created data set (if any)
+            uint256 nPieces = pieceData.length;
+            require(addPayload.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
+            uint256 firstAddedNew = nextPieceId[newSetId];
 
-        return firstAdded;
+            if (nPieces > 0) {
+                uint256[] memory pieceIdsNew = new uint256[](pieceData.length);
+                Cids.Cid[] memory pieceCidsAddedNew = new Cids.Cid[](pieceData.length);
+
+                for (uint256 i = 0; i < nPieces; i++) {
+                    addOnePiece(newSetId, i, pieceData[i]);
+                    pieceIdsNew[i] = firstAddedNew + i;
+                    pieceCidsAddedNew[i] = pieceData[i];
+                }
+                emit PiecesAdded(newSetId, pieceIdsNew, pieceCidsAddedNew);
+
+                address listenerAddrNew = dataSetListener[newSetId];
+                if (listenerAddrNew != address(0)) {
+                    PDPListener(listenerAddrNew).piecesAdded(newSetId, firstAddedNew, pieceData, addPayload);
+                }
+            }
+
+            // Return the at the end to avoid any possible re-entrency issues.
+            if (msg.value > sybilFee) {
+                (bool success,) = msg.sender.call{value: msg.value - sybilFee}("");
+                require(success, "Transfer failed.");
+            }
+
+            return newSetId;
+        } else {
+            // Adding to an existing set; no fee should be sent and listenerAddr must be zero
+            require(listenerAddr == address(0), "listener must be zero for existing dataset");
+            require(msg.value == 0, "no fee on add to existing dataset");
+
+            uint256 nPieces = pieceData.length;
+            require(extraData.length <= EXTRA_DATA_MAX_SIZE, "Extra data too large");
+            require(dataSetLive(setId), "Data set not live");
+            require(nPieces > 0, "Must add at least one piece");
+            require(storageProvider[setId] == msg.sender, "Only the storage provider can add pieces");
+            uint256 firstAdded = nextPieceId[setId];
+            uint256[] memory pieceIds = new uint256[](pieceData.length);
+            Cids.Cid[] memory pieceCidsAdded = new Cids.Cid[](pieceData.length);
+
+            for (uint256 i = 0; i < nPieces; i++) {
+                addOnePiece(setId, i, pieceData[i]);
+                pieceIds[i] = firstAdded + i;
+                pieceCidsAdded[i] = pieceData[i];
+            }
+            emit PiecesAdded(setId, pieceIds, pieceCidsAdded);
+            address listenerAddrExisting = dataSetListener[setId];
+            if (listenerAddrExisting != address(0)) {
+                PDPListener(listenerAddrExisting).piecesAdded(setId, firstAdded, pieceData, extraData);
+            }
+            return firstAdded;
+        }
     }
 
     error IndexedError(uint256 idx, string msg);