Optimized modulo calculation with bitmask

Signed-off-by: Peter Alfonsi <[email protected]>

Author: Peter Alfonsi

server/src/main/java/org/opensearch/common/cache/tier/keystore/RBMIntKeyLookupStore.java

@@ -72,6 +72,8 @@ public int getValue() {
     }
 
     protected final int modulo;
+    protected final int modulo_bitmask;
+    // Since our modulo is always a power of two we can optimize it by ANDing with a particular bitmask
     KeyStoreStats stats;
     protected RoaringBitmap rbm;
     private HashMap<Integer, CounterMetric> collidedIntCounters;
@@ -92,15 +94,20 @@ public RBMIntKeyLookupStore(long memSizeCapInBytes) {
 
     public RBMIntKeyLookupStore(KeystoreModuloValue moduloValue, long memSizeCapInBytes) {
         this.modulo = moduloValue.getValue();
+        if (modulo > 0) {
+            this.modulo_bitmask = modulo - 1; // keep last log_2(modulo) bits
+        } else {
+            this.modulo_bitmask = -1; // -1 in twos complement is all ones -> includes all bits -> same as no modulo
+        }
         this.stats = new KeyStoreStats(memSizeCapInBytes);
         this.rbm = new RoaringBitmap();
         this.collidedIntCounters = new HashMap<>();
         this.removalSets = new HashMap<>();
         this.mostRecentByteEstimate = 0L;
     }
 
-    private final int transform(int value) {
-        return modulo == 0 ? value : value % modulo;
+    private int transform(int value) {
+        return value & modulo_bitmask;
     }
 
     private void handleCollisions(int transformedValue) {

server/src/test/java/org/opensearch/common/cache/tier/keystore/RBMIntKeyLookupStoreTests.java

@@ -61,14 +61,22 @@ public void testTransformationLogic() throws Exception {
         for (int i = 0; i < 4; i++) { // after this we run into max value, but thats not a flaw with the class design
             int posValue = i * modulo + offset;
             kls.add(posValue);
+            assertEquals(offset, (int) kls.getInternalRepresentation(posValue));
             int negValue = -(i * modulo + offset);
             kls.add(negValue);
+            assertEquals(modulo - offset, (int) kls.getInternalRepresentation(negValue));
         }
         assertEquals(2, kls.getSize());
         int[] testVals = new int[] { 0, 1, -1, -23495, 23058, modulo, -modulo, Integer.MAX_VALUE, Integer.MIN_VALUE };
         for (int value : testVals) {
             assertTrue(kls.getInternalRepresentation(value) < modulo);
-            assertTrue(kls.getInternalRepresentation(value) > -modulo);
+            assertTrue(kls.getInternalRepresentation(value) >= 0);
+        }
+        RBMIntKeyLookupStore no_modulo_kls = new RBMIntKeyLookupStore(RBMIntKeyLookupStore.KeystoreModuloValue.NONE, 0L);
+        Random rand = Randomness.get();
+        for (int i = 0; i < 100; i++) {
+            int val = rand.nextInt();
+            assertEquals(val, (int) no_modulo_kls.getInternalRepresentation(val));
         }
     }