Merge pull request #21 from tenetxyz/support-composite-keys-keyswithv…

…alue chore: update handleSet func signature; add changeset
tenetxyz · Jul 27, 2023 · c98d96b · c98d96b
2 parents bd4edd0 + 49adfbe
commit c98d96b
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 23 deletions.
diff --git a/.changeset/tall-tools-admire.md b/.changeset/tall-tools-admire.md
@@ -0,0 +1,5 @@
+---
+"@latticexyz/world": minor
+---
+
+Add support for composite keys in KeysWithValue module
diff --git a/packages/world/src/modules/keyswithvalue/KeysWithValueHook.sol b/packages/world/src/modules/keyswithvalue/KeysWithValueHook.sol
@@ -29,24 +29,6 @@ contract KeysWithValueHook is IStoreHook {
     return IBaseWorld(StoreSwitch.getStoreAddress());
   }
 
-  function handleSet(bytes32 tableId, bytes32 valueHash, bytes32[] memory key) internal {
-    if (key.length > 0) {
-      KeysWithValue.pushKeys0(tableId, valueHash, key[0]);
-      if (key.length > 1) {
-        KeysWithValue.pushKeys1(tableId, valueHash, key[1]);
-        if (key.length > 2) {
-          KeysWithValue.pushKeys2(tableId, valueHash, key[2]);
-          if (key.length > 3) {
-            KeysWithValue.pushKeys3(tableId, valueHash, key[3]);
-            if (key.length > 4) {
-              KeysWithValue.pushKeys4(tableId, valueHash, key[4]);
-            }
-          }
-        }
-      }
-    }
-  }
-
   function onSetRecord(bytes32 sourceTableId, bytes32[] memory key, bytes memory data) public {
     bytes32 targetTableId = getTargetTableSelector(MODULE_NAMESPACE, sourceTableId);
 
@@ -57,7 +39,7 @@ contract KeysWithValueHook is IStoreHook {
     _removeKeyFromList(targetTableId, key, previousValue);
 
     // Push the key to the list of keys with the new value
-    handleSet(targetTableId, keccak256(data), key);
+    _addKeyToList(targetTableId, key, keccak256(data));
   }
 
   function onBeforeSetField(bytes32 sourceTableId, bytes32[] memory key, uint8, bytes memory) public {
@@ -71,7 +53,7 @@ contract KeysWithValueHook is IStoreHook {
     // Add the key to the list of keys with the new value
     bytes32 newValue = keccak256(_world().getRecord(sourceTableId, key));
     bytes32 targetTableId = getTargetTableSelector(MODULE_NAMESPACE, sourceTableId);
-    handleSet(targetTableId, newValue, key);
+    _addKeyToList(targetTableId, key, newValue);
   }
 
   function onDeleteRecord(bytes32 sourceTableId, bytes32[] memory key) public {
@@ -81,6 +63,24 @@ contract KeysWithValueHook is IStoreHook {
     _removeKeyFromList(targetTableId, key, previousValue);
   }
 
+  function _addKeyToList(bytes32 targetTableId, bytes32[] memory key, bytes32 valueHash) internal {
+    if (key.length > 0) {
+      KeysWithValue.pushKeys0(targetTableId, valueHash, key[0]);
+      if (key.length > 1) {
+        KeysWithValue.pushKeys1(targetTableId, valueHash, key[1]);
+        if (key.length > 2) {
+          KeysWithValue.pushKeys2(targetTableId, valueHash, key[2]);
+          if (key.length > 3) {
+            KeysWithValue.pushKeys3(targetTableId, valueHash, key[3]);
+            if (key.length > 4) {
+              KeysWithValue.pushKeys4(targetTableId, valueHash, key[4]);
+            }
+          }
+        }
+      }
+    }
+  }
+
   function _removeKeyFromList(bytes32 targetTableId, bytes32[] memory key, bytes32 valueHash) internal {
     // Get the keys with the previous value
     KeysWithValueData memory keysWithPreviousValue = KeysWithValue.get(targetTableId, valueHash);

diff --git a/packages/world/test/KeysWithValueModule.t.sol b/packages/world/test/KeysWithValueModule.t.sol
@@ -329,13 +329,13 @@ contract KeysWithValueModuleTest is Test, GasReporter {
     assertEq(keysWithValue[0].length, 1);
     assertEq(keysWithValue[0][0], key1);
 
-    // // Set a another key with the same value
+    // Set a another key with the same value
     world.setRecord(namespace, sourceName, keyTuple2, abi.encodePacked(value));
 
-    // // Get the list of keys with value from the target table
+    // Get the list of keys with value from the target table
     keysWithValue = getKeysWithValue(world, sourceTableId, abi.encode(value));
 
-    // // Assert that the list is correct
+    // Assert that the list is correct
     assertEq(keysWithValue.length, 2);
     assertEq(keysWithValue[0].length, 1);
     assertEq(keysWithValue[1].length, 1);