Add setCapacity and removeRange to WatchedSparseBooleanMatrix

- This CL adds #setCapacity and #removeRange APIs to
  WatchedSparseBooleanMatrix for AppsFilter to update filter
  cache more efficiently when a user is created or deleted.

- Fix a crash in the #pack function when the caller invokes #compat
  to shrink the matrix.

Test: atest WatcherTest
Bug: 187853334
Change-Id: I6a3c87a7492793767f888487381182955d564d62

services/core/java/com/android/server/utils/WatchedSparseBooleanMatrix.java

@@ -262,6 +262,33 @@ public class WatchedSparseBooleanMatrix extends WatchableImpl implements Snappab
         onChanged();
     }
 
+    /**
+     * Removes all of the mappings whose index is between {@code fromIndex}, inclusive, and
+     * {@code toIndex}, exclusive. The matrix does not shrink.
+     */
+    public void removeRange(int fromIndex, int toIndex) {
+        if (toIndex < fromIndex) {
+            throw new ArrayIndexOutOfBoundsException("toIndex < fromIndex");
+        }
+        final int num = toIndex - fromIndex;
+        if (num == 0) {
+            return;
+        }
+        validateIndex(fromIndex);
+        validateIndex(toIndex - 1);
+        for (int i = fromIndex; i < toIndex; i++) {
+            mInUse[mMap[i]] = false;
+        }
+        System.arraycopy(mKeys, toIndex, mKeys, fromIndex, mSize - toIndex);
+        System.arraycopy(mMap, toIndex, mMap, fromIndex, mSize - toIndex);
+        for (int i = mSize - num; i < mSize; i++) {
+            mKeys[i] = 0;
+            mMap[i] = 0;
+        }
+        mSize -= num;
+        onChanged();
+    }
+
     /**
      * Returns the number of key-value mappings that this WatchedSparseBooleanMatrix
      * currently stores.
@@ -371,7 +398,7 @@ public class WatchedSparseBooleanMatrix extends WatchableImpl implements Snappab
                 // Preemptively grow the matrix, which also grows the free list.
                 growMatrix();
             }
-            int newIndex = nextFree();
+            int newIndex = nextFree(true /* acquire */);
             mKeys = GrowingArrayUtils.insert(mKeys, mSize, i, key);
             mMap = GrowingArrayUtils.insert(mMap, mSize, i, newIndex);
             mSize++;
@@ -447,12 +474,12 @@ public class WatchedSparseBooleanMatrix extends WatchableImpl implements Snappab
     }
 
     /**
-     * Find an unused storage index, mark it in-use, and return it.
+     * Find an unused storage index, and return it. Mark it in-use if the {@code acquire} is true.
      */
-    private int nextFree() {
+    private int nextFree(boolean acquire) {
         for (int i = 0; i < mInUse.length; i++) {
             if (!mInUse[i]) {
-                mInUse[i] = true;
+                mInUse[i] = acquire;
                 return i;
             }
         }
@@ -488,7 +515,8 @@ public class WatchedSparseBooleanMatrix extends WatchableImpl implements Snappab
         }
         // dst and src are identify raw (row, col) in mValues.  srcIndex is the index (as
         // in the result of keyAt()) of the key being relocated.
-        for (int dst = nextFree(); dst < mSize; dst = nextFree()) {
+        for (int dst = nextFree(false); dst < mSize; dst = nextFree(false)) {
+            mInUse[dst] = true;
             int srcIndex = lastInuse();
             int src = mMap[srcIndex];
             mInUse[src] = false;
@@ -538,6 +566,20 @@ public class WatchedSparseBooleanMatrix extends WatchableImpl implements Snappab
         return mOrder;
     }
 
+    /**
+     * Set capacity to enlarge the size of the 2D matrix. Capacity less than the {@link #capacity()}
+     * is not supported.
+     */
+    public void setCapacity(int capacity) {
+        if (capacity <= mOrder) {
+            return;
+        }
+        if (capacity % STEP != 0) {
+            capacity = ((capacity / STEP) + 1) * STEP;
+        }
+        resizeMatrix(capacity);
+    }
+
     /**
      * {@inheritDoc}
      */

services/tests/servicestests/src/com/android/server/utils/WatcherTest.java

@@ -915,6 +915,23 @@ public class WatcherTest {
         }
     }
 
+    // Fill new cells in the matrix which has enlarged capacity.
+    private void fillNew(WatchedSparseBooleanMatrix matrix, int initialCapacity,
+            int newCapacity, int[] indexes) {
+        final int size = newCapacity;
+        for (int i = 0; i < size; i++) {
+            for (int j = 0; j < size; j++) {
+                if (i < initialCapacity && j < initialCapacity) {
+                    // Do not touch old cells
+                    continue;
+                }
+                final int row = indexes[i];
+                final int col = indexes[j];
+                matrix.put(row, col, cellValue(i, j));
+            }
+        }
+    }
+
     // Verify the content of a matrix.  This asserts on mismatch.  Selected indices may
     // have been deleted.
     private void verify(WatchedSparseBooleanMatrix matrix, int[] indexes, boolean[] absent) {
@@ -989,6 +1006,24 @@ public class WatcherTest {
         assertTrue("Matrix shrink", finalCapacity - matrix.size() < matrix.STEP);
     }
 
+    private void matrixSetCapacity(WatchedSparseBooleanMatrix matrix, int newCapacity,
+            IndexGenerator indexer) {
+        final int initialCapacity = matrix.capacity();
+        final int[] indexes = indexer.indexes(Math.max(initialCapacity, newCapacity));
+        fill(matrix, initialCapacity, indexes);
+
+        matrix.setCapacity(newCapacity);
+        fillNew(matrix, initialCapacity, newCapacity, indexes);
+
+        assertEquals(matrix.size(), indexes.length);
+        verify(matrix, indexes, null);
+        // Test the keyAt/indexOfKey methods
+        for (int i = 0; i < matrix.size(); i++) {
+            int key = indexes[i];
+            assertEquals(matrix.keyAt(matrix.indexOfKey(key)), key);
+        }
+    }
+
     @Test
     public void testWatchedSparseBooleanMatrix() {
         final String name = "WatchedSparseBooleanMatrix";
@@ -1051,6 +1086,46 @@ public class WatcherTest {
         assertEquals(a.equals(s), false);
     }
 
+    @Test
+    public void testWatchedSparseBooleanMatrix_setCapacity() {
+        final IndexGenerator indexer = new IndexGenerator(3);
+        matrixSetCapacity(new WatchedSparseBooleanMatrix(500), 1000, indexer);
+        matrixSetCapacity(new WatchedSparseBooleanMatrix(1000), 500, indexer);
+    }
+
+    @Test
+    public void testWatchedSparseBooleanMatrix_removeRangeAndShrink() {
+        final IndexGenerator indexer = new IndexGenerator(3);
+        final int initialCapacity = 500;
+        final int removeCounts = 33;
+        final WatchedSparseBooleanMatrix matrix = new WatchedSparseBooleanMatrix(initialCapacity);
+        final int[] indexes = indexer.indexes(initialCapacity);
+        final boolean[] absents = new boolean[initialCapacity];
+        fill(matrix, initialCapacity, indexes);
+        assertEquals(matrix.size(), initialCapacity);
+
+        for (int i = 0; i < initialCapacity / removeCounts; i++) {
+            final int size = matrix.size();
+            final int fromIndex = (size / 2 < removeCounts ? 0 : size / 2 - removeCounts);
+            final int toIndex = (fromIndex + removeCounts > size ? size : fromIndex + removeCounts);
+            for (int index = fromIndex; index < toIndex; index++) {
+                final int key = matrix.keyAt(index);
+                for (int j = 0; j < indexes.length; j++) {
+                    if (key == indexes[j]) {
+                        absents[j] = true;
+                        break;
+                    }
+                }
+            }
+            matrix.removeRange(fromIndex, toIndex);
+            assertEquals(matrix.size(), size - (toIndex - fromIndex));
+            verify(matrix, indexes, absents);
+
+            matrix.compact();
+            verify(matrix, indexes, absents);
+        }
+    }
+
     @Test
     public void testNestedArrays() {
         final String name = "NestedArrays";