add ability to write more options for a claim

src/OptionSettlement.sol

@@ -188,60 +188,72 @@ contract OptionSettlementEngine is ERC1155, IOptionSettlementEngine {
 
     /// @inheritdoc IOptionSettlementEngine
     function write(uint256 optionId, uint112 amount) external returns (uint256 claimId) {
-        if (amount == 0) {
-            revert AmountWrittenCannotBeZero();
-        }
+        /// supplying claimId as 0 to the overloaded write signifies that a new
+        /// claim NFT should be minted for the options lot, rather than being added
+        /// as an existing claim.
+        return write(optionId, amount, 0);
+    }
 
-        (uint160 _optionId, uint96 claimIndex) = getDecodedIdComponents(optionId);
-        if (claimIndex != 0) {
-            // claim index must be zero in lower 96b
+    /// @inheritdoc IOptionSettlementEngine
+    function write(uint256 optionId, uint112 amount, uint256 claimId) public returns (uint256) {
+        (uint160 _optionIdU160b, uint96 _optionIdL96b) = getDecodedIdComponents(optionId);
+
+        // optionId must be zero in lower 96b for provided option Id
+        if (_optionIdL96b != 0) {
             revert InvalidOption(optionId);
         }
 
-        Option storage optionRecord = _option[_optionId];
-
-        if (optionRecord.expiryTimestamp <= block.timestamp) {
-            revert ExpiredOption(optionId, optionRecord.expiryTimestamp);
+        // claim provided must match the option provided
+        if (claimId != 0 && ((claimId >> 96) != (optionId >> 96))) {
+            revert EncodedOptionIdInClaimIdDoesNotMatchProvidedOptionId(claimId, optionId);
         }
 
-        uint256 rxAmount = amount * optionRecord.underlyingAmount;
-        uint256 fee = ((rxAmount / 10000) * feeBps);
-        address underlyingAsset = optionRecord.underlyingAsset;
+        Option storage optionRecord = _writeOptions(_optionIdU160b, amount);
+        uint256 mintClaimNft = 0;
+
+        if (claimId == 0) {
+            // create new claim
+            // Increment the next token ID
+            uint96 claimIndex = optionRecord.nextClaimId++;
+            claimId = getTokenId(_optionIdU160b, claimIndex);
+            // Store info about the claim
+            _claim[claimId] = Claim({amountWritten: amount, amountExercised: 0, claimed: false});
+            unexercisedClaimsByOption[_optionIdU160b].push(claimId);
+            mintClaimNft = 1;
+        } else {
+            // check ownership of claim
+            uint256 balance = balanceOf[msg.sender][claimId];
+            if (balance != 1) {
+                revert CallerDoesNotOwnClaimId(claimId);
+            }
 
-        // Transfer the requisite underlying asset
-        SafeTransferLib.safeTransferFrom(ERC20(underlyingAsset), msg.sender, address(this), (rxAmount + fee));
+            // retrieve claim
+            Claim storage existingClaim = _claim[claimId];
 
-        claimIndex = optionRecord.nextClaimId;
-        claimId = getTokenId(_optionId, claimIndex);
+            if (existingClaim.claimed) {
+                revert AlreadyClaimed(claimId);
+            }
+
+            existingClaim.amountWritten += amount;
+        }
 
         // Mint the options contracts and claim token
         uint256[] memory tokens = new uint256[](2);
-        tokens[0] = uint256(_optionId) << 96;
+        tokens[0] = optionId;
         tokens[1] = claimId;
 
         uint256[] memory amounts = new uint256[](2);
         amounts[0] = amount;
-        amounts[1] = 1;
+        amounts[1] = mintClaimNft;
 
         bytes memory data = new bytes(0);
 
-        // Store info about the claim
-        _claim[claimId] = Claim({amountWritten: amount, amountExercised: 0, claimed: false});
-        unexercisedClaimsByOption[_optionId].push(claimId);
-
-        feeBalance[underlyingAsset] += fee;
-
-        // Increment the next token ID
-        optionRecord.nextClaimId++;
+        // Send tokens to writer
+        _batchMint(msg.sender, tokens, amounts, data);
 
-        emit FeeAccrued(underlyingAsset, msg.sender, fee);
-        // TODO: option ID in addition to claim ID is potentially redundant now
-        // either emit claim idx specifically, or just the entire option id
-        // with encoded claim idx
         emit OptionsWritten(optionId, msg.sender, claimId, amount);
 
-        // Send tokens to writer
-        _batchMint(msg.sender, tokens, amounts, data);
+        return claimId;
     }
 
     function assignExercise(uint160 optionId, uint96 claimsLen, uint112 amount, uint160 settlementSeed) internal {
@@ -364,7 +376,7 @@ contract OptionSettlementEngine is ERC1155, IOptionSettlementEngine {
         uint256 balance = this.balanceOf(msg.sender, claimId);
 
         if (balance != 1) {
-            revert RedeemerDoesNotOwnClaimId(claimId);
+            revert CallerDoesNotOwnClaimId(claimId);
         }
 
         Claim storage claimRecord = _claim[claimId];
@@ -441,6 +453,41 @@ contract OptionSettlementEngine is ERC1155, IOptionSettlementEngine {
         }
     }
 
+    // **********************************************************************
+    //                        INTERNAL HELPERS
+    // **********************************************************************
+    /**
+     * @dev Writes the specified number of options, transferring in the requisite
+     * underlying assets, and trasnsferring fungible ERC1155 tokens to caller.
+     * Reverts if insufficient underlying assets are not available from caller.
+     * @param _optionId The options to write.
+     * @param amount The amount of options to write.
+     */
+    function _writeOptions(uint160 _optionId, uint112 amount) internal returns (Option storage) {
+        if (amount == 0) {
+            revert AmountWrittenCannotBeZero();
+        }
+
+        Option storage optionRecord = _option[_optionId];
+
+        if (optionRecord.expiryTimestamp <= block.timestamp) {
+            revert ExpiredOption(uint256(_optionId) << 96, optionRecord.expiryTimestamp);
+        }
+
+        uint256 rxAmount = amount * optionRecord.underlyingAmount;
+        uint256 fee = ((rxAmount / 10000) * feeBps);
+        address underlyingAsset = optionRecord.underlyingAsset;
+
+        // Transfer the requisite underlying asset
+        SafeTransferLib.safeTransferFrom(ERC20(underlyingAsset), msg.sender, address(this), (rxAmount + fee));
+
+        feeBalance[underlyingAsset] += fee;
+
+        emit FeeAccrued(underlyingAsset, msg.sender, fee);
+
+        return optionRecord;
+    }
+
     // **********************************************************************
     //                    TOKEN ID ENCODING HELPERS
     // **********************************************************************

src/interfaces/IOptionSettlementEngine.sol

@@ -55,6 +55,14 @@ interface IOptionSettlementEngine {
      */
     error InvalidClaim(uint256 token);
 
+    /**
+     * @notice Provided claimId does not match provided option id in the upper 160b
+     * encoding the corresponding option ID for which the claim was written.
+     * @param claimId The provided claim ID.
+     * @param optionId The provided option ID.
+     */
+    error EncodedOptionIdInClaimIdDoesNotMatchProvidedOptionId(uint256 claimId, uint256 optionId);
+
     /**
      * @notice The optionId specified expired at expiry.
      * @param optionId The id of the expired option.
@@ -76,10 +84,10 @@ interface IOptionSettlementEngine {
     error NoClaims(uint256 optionId);
 
     /**
-     * @notice This account has no claims.
+     * @notice This claim is not owned by the caller.
      * @param claimId Supplied claim ID.
      */
-    error RedeemerDoesNotOwnClaimId(uint256 claimId);
+    error CallerDoesNotOwnClaimId(uint256 claimId);
 
     /**
      * @notice This claimId has already been claimed.
@@ -305,6 +313,16 @@ interface IOptionSettlementEngine {
      */
     function write(uint256 optionId, uint112 amount) external returns (uint256 claimId);
 
+    /**
+     * @notice This override allows additional options to be written against a particular
+     * claim id.
+     * @param optionId The desired option id to write.
+     * @param amount The desired number of options to write.
+     * @param claimId The claimId for the options lot to which the caller will add options
+     * @return claimId The claim NFT id for the option bundle.
+     */
+    function write(uint256 optionId, uint112 amount, uint256 claimId) external returns (uint256);
+
     /**
      * @notice Exercises specified amount of optionId, transferring in the exercise asset,
      * and transferring out the underlying asset if requirements are met.

src/test/OptionSettlement.t.sol

@@ -292,6 +292,37 @@ contract OptionSettlementTest is Test, NFTreceiver {
         _assertPosition(underlyingPositions.exercisePosition, 2 * testExerciseAmount);
     }
 
+    function testAddOptionsToExistingClaim() public {
+        // write some options, grab a claim
+        vm.startPrank(ALICE);
+        uint256 claimId = engine.write(testOptionId, 1);
+
+        IOptionSettlementEngine.Claim memory claimRecord = engine.claim(claimId);
+
+        assertEq(1, claimRecord.amountWritten);
+        assertEq(0, claimRecord.amountExercised);
+        assertEq(false, claimRecord.claimed);
+        assertEq(1, engine.balanceOf(ALICE, claimId));
+        assertEq(1, engine.balanceOf(ALICE, testOptionId));
+
+        // write some more options, get a new claim NFT
+        uint256 claimId2 = engine.write(testOptionId, 1);
+        assertFalse(claimId == claimId2);
+        assertEq(1, engine.balanceOf(ALICE, claimId2));
+        assertEq(2, engine.balanceOf(ALICE, testOptionId));
+
+        // write some more options, adding to existing claim
+        uint256 claimId3 = engine.write(testOptionId, 1, claimId);
+        assertEq(claimId, claimId3);
+        assertEq(1, engine.balanceOf(ALICE, claimId3));
+        assertEq(3, engine.balanceOf(ALICE, testOptionId));
+
+        claimRecord = engine.claim(claimId3);
+        assertEq(2, claimRecord.amountWritten);
+        assertEq(0, claimRecord.amountExercised);
+        assertEq(false, claimRecord.claimed);
+    }
+
     // **********************************************************************
     //                            FAIL TESTS
     // **********************************************************************
@@ -429,18 +460,18 @@ contract OptionSettlementTest is Test, NFTreceiver {
         vm.startPrank(ALICE);
         uint256 claimId = engine.write(testOptionId, 1);
         engine.safeTransferFrom(ALICE, BOB, testOptionId, 1, "");
-        vm.expectRevert(IOptionSettlementEngine.RedeemerDoesNotOwnClaimId.selector);
+        vm.expectRevert(IOptionSettlementEngine.CallerDoesNotOwnClaimId.selector);
         engine.redeem(claimId);
         vm.stopPrank();
         // Carol feels left out and tries to redeem what she can't
         vm.startPrank(CAROL);
-        vm.expectRevert(IOptionSettlementEngine.RedeemerDoesNotOwnClaimId.selector);
+        vm.expectRevert(IOptionSettlementEngine.CallerDoesNotOwnClaimId.selector);
         engine.redeem(claimId);
         vm.stopPrank();
         // Bob redeems, which should burn, and then be unable to redeem a second time
         vm.startPrank(BOB);
         engine.redeem(claimId);
-        vm.expectRevert(IOptionSettlementEngine.RedeemerDoesNotOwnClaimId.selector);
+        vm.expectRevert(IOptionSettlementEngine.CallerDoesNotOwnClaimId.selector);
         engine.redeem(claimId);
     }
 
@@ -477,6 +508,30 @@ contract OptionSettlementTest is Test, NFTreceiver {
         engine.uri(420);
     }
 
+    function testRevertIfClaimIdDoesNotEncodeOptionId() public {
+        uint256 option1Claim1 = engine.getTokenId(0xDEADBEEF1, 0xCAFECAFE1);
+        uint256 option2 = engine.getTokenId(0xDEADBEEF2, 0x0);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IOptionSettlementEngine.EncodedOptionIdInClaimIdDoesNotMatchProvidedOptionId.selector,
+                option1Claim1,
+                option2
+            )
+        );
+        engine.write(option2, 1, option1Claim1);
+    }
+
+    function testRevertIfWriterDoesNotOwnClaim() public {
+        vm.startPrank(ALICE);
+        uint256 claimId = engine.write(testOptionId, 1);
+        vm.stopPrank();
+
+        vm.startPrank(BOB);
+        vm.expectRevert(abi.encodeWithSelector(IOptionSettlementEngine.CallerDoesNotOwnClaimId.selector, claimId));
+        engine.write(testOptionId, 1, claimId);
+    }
+
     // **********************************************************************
     //                            FUZZ TESTS
     // **********************************************************************