Some refactoring for fixing the tests and also eliminating the use of

TreeMap but use a more light weight data structure.

Bug: 145493956
Test: atest FrameworksServicesTests:IntegrityFileManagerTest
Change-Id: I2e7dd5336c4aa3ee41c05d4af0388125751bc207

13 files changed, 124 insertions, 116 deletions

diff --git a/services/core/java/com/android/server/integrity/model/BitTrackedInputStream.java b/services/core/java/com/android/server/integrity/model/BitTrackedInputStream.java
index e555e3e5746e..4bf8fe8d93b6 100644
--- a/services/core/java/com/android/server/integrity/model/BitTrackedInputStream.java
+++ b/services/core/java/com/android/server/integrity/model/BitTrackedInputStream.java
@@ -27,30 +27,37 @@ import java.io.InputStream;
  */
 public class BitTrackedInputStream extends BitInputStream {
 
-    private static int sReadBitsCount;
+    private int mReadBitsCount;
 
     /** Constructor with byte array. */
     public BitTrackedInputStream(byte[] inputStream) {
         super(inputStream);
-        sReadBitsCount = 0;
+        mReadBitsCount = 0;
     }
 
     /** Constructor with input stream. */
     public BitTrackedInputStream(InputStream inputStream) {
         super(inputStream);
-        sReadBitsCount = 0;
+        mReadBitsCount = 0;
     }
 
     /** Obtains an integer value of the next {@code numOfBits}. */
     @Override
     public int getNext(int numOfBits) throws IOException {
-        sReadBitsCount += numOfBits;
+        mReadBitsCount += numOfBits;
         return super.getNext(numOfBits);
     }
 
     /** Returns the current cursor position showing the number of bits that are read. */
     public int getReadBitsCount() {
-        return sReadBitsCount;
+        return mReadBitsCount;
+    }
+
+    /**
+     * Returns true if we can read more rules by checking whether the end index is not reached yet.
+     */
+    public boolean canReadMoreRules(int endIndexBytes) {
+        return mReadBitsCount < endIndexBytes * 8;
     }
 
     /**
@@ -59,11 +66,11 @@ public class BitTrackedInputStream extends BitInputStream {
      * Note that the integer parameter specifies the location in bytes -- not bits.
      */
     public void setCursorToByteLocation(int byteLocation) throws IOException {
-        int bitCountToRead = byteLocation * 8 - sReadBitsCount;
+        int bitCountToRead = byteLocation * 8 - mReadBitsCount;
         if (bitCountToRead < 0) {
             throw new IllegalStateException("The byte position is already read.");
         }
         super.getNext(bitCountToRead);
-        sReadBitsCount = byteLocation * 8;
+        mReadBitsCount = byteLocation * 8;
     }
 }
diff --git a/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java b/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
index e744326c49db..2f285631b982 100644
--- a/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
+++ b/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
@@ -105,8 +105,7 @@ public class RuleBinaryParser implements RuleParser {
             bitTrackedInputStream.setCursorToByteLocation(range.getStartIndex());
 
             // Read the rules until we reach the end index.
-            while (bitTrackedInputStream.hasNext()
-                    && bitTrackedInputStream.getReadBitsCount() < range.getEndIndex()) {
+            while (bitTrackedInputStream.canReadMoreRules(range.getEndIndex())) {
                 if (bitTrackedInputStream.getNext(SIGNAL_BIT) == 1) {
                     parsedRules.add(parseRule(bitTrackedInputStream));
                 }
diff --git a/services/core/java/com/android/server/integrity/parser/RuleIndexRange.java b/services/core/java/com/android/server/integrity/parser/RuleIndexRange.java
index 8c8450e5fdc4..453fa5d84053 100644
--- a/services/core/java/com/android/server/integrity/parser/RuleIndexRange.java
+++ b/services/core/java/com/android/server/integrity/parser/RuleIndexRange.java
@@ -23,28 +23,28 @@ import android.annotation.Nullable;
  * RuleIndexingController}.
  */
 public class RuleIndexRange {
-    private static int sStartIndex;
-    private static int sEndIndex;
+    private int mStartIndex;
+    private int mEndIndex;
 
     /** Constructor with start and end indexes. */
     public RuleIndexRange(int startIndex, int endIndex) {
-        this.sStartIndex = startIndex;
-        this.sEndIndex = endIndex;
+        this.mStartIndex = startIndex;
+        this.mEndIndex = endIndex;
     }
 
     /** Returns the startIndex. */
     public int getStartIndex() {
-        return sStartIndex;
+        return mStartIndex;
     }
 
     /** Returns the end index. */
     public int getEndIndex() {
-        return sEndIndex;
+        return mEndIndex;
     }
 
     @Override
     public boolean equals(@Nullable Object object) {
-        return sStartIndex == ((RuleIndexRange) object).getStartIndex()
-                && sEndIndex == ((RuleIndexRange) object).getEndIndex();
+        return mStartIndex == ((RuleIndexRange) object).getStartIndex()
+                && mEndIndex == ((RuleIndexRange) object).getEndIndex();
     }
 }
diff --git a/services/core/java/com/android/server/integrity/parser/RuleIndexingController.java b/services/core/java/com/android/server/integrity/parser/RuleIndexingController.java
index c9713220d6e8..03392abf478f 100644
--- a/services/core/java/com/android/server/integrity/parser/RuleIndexingController.java
+++ b/services/core/java/com/android/server/integrity/parser/RuleIndexingController.java
@@ -56,7 +56,7 @@ public class RuleIndexingController {
      * read and evaluated.
      */
     public List<RuleIndexRange> identifyRulesToEvaluate(AppInstallMetadata appInstallMetadata) {
-        ArrayList<RuleIndexRange> indexRanges = new ArrayList();
+        List<RuleIndexRange> indexRanges = new ArrayList<>();
 
         // Add the range for package name indexes rules.
         indexRanges.add(
@@ -102,7 +102,7 @@ public class RuleIndexingController {
                         .collect(Collectors.toCollection(TreeSet::new));
 
         String minIndex = keyTreeSet.floor(searchedKey);
-        String maxIndex = keyTreeSet.ceiling(searchedKey);
+        String maxIndex = keyTreeSet.higher(searchedKey);
 
         return new RuleIndexRange(
                 indexMap.get(minIndex == null ? START_INDEXING_KEY : minIndex),
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java b/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
index d3588d38d976..8f53be7d87af 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
@@ -52,7 +52,7 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.TreeMap;
+import java.util.stream.Collectors;
 
 /** A helper class to serialize rules from the {@link Rule} model to Binary representation. */
 public class RuleBinarySerializer implements RuleSerializer {
@@ -80,21 +80,24 @@ public class RuleBinarySerializer implements RuleSerializer {
             throws RuleSerializeException {
         try {
             // Determine the indexing groups and the order of the rules within each indexed group.
-            Map<Integer, TreeMap<String, List<Rule>>> indexedRules =
+            Map<Integer, Map<String, List<Rule>>> indexedRules =
                     RuleIndexingDetailsIdentifier.splitRulesIntoIndexBuckets(rules);
 
             // Serialize the rules.
             ByteTrackedOutputStream ruleFileByteTrackedOutputStream =
                     new ByteTrackedOutputStream(rulesFileOutputStream);
             serializeRuleFileMetadata(formatVersion, ruleFileByteTrackedOutputStream);
-            Map<String, Integer> packageNameIndexes =
-                    serializeRuleList(indexedRules.get(PACKAGE_NAME_INDEXED),
+            LinkedHashMap<String, Integer> packageNameIndexes =
+                    serializeRuleList(
+                            indexedRules.get(PACKAGE_NAME_INDEXED),
                             ruleFileByteTrackedOutputStream);
-            Map<String, Integer> appCertificateIndexes =
-                    serializeRuleList(indexedRules.get(APP_CERTIFICATE_INDEXED),
+            LinkedHashMap<String, Integer> appCertificateIndexes =
+                    serializeRuleList(
+                            indexedRules.get(APP_CERTIFICATE_INDEXED),
                             ruleFileByteTrackedOutputStream);
-            Map<String, Integer> unindexedRulesIndexes =
-                    serializeRuleList(indexedRules.get(NOT_INDEXED),
+            LinkedHashMap<String, Integer> unindexedRulesIndexes =
+                    serializeRuleList(
+                            indexedRules.get(NOT_INDEXED),
                             ruleFileByteTrackedOutputStream);
 
             // Serialize their indexes.
@@ -104,6 +107,9 @@ public class RuleBinarySerializer implements RuleSerializer {
                     true);
             serializeIndexGroup(unindexedRulesIndexes, indexingBitOutputStream, /* isIndexed= */
                     false);
+            // TODO(b/147609625): This dummy bit is set for fixing the padding issue. Remove when
+            // the issue is fixed and correct the tests that does this padding too.
+            indexingBitOutputStream.setNext();
             indexingFileOutputStream.write(indexingBitOutputStream.toByteArray());
         } catch (Exception e) {
             throw new RuleSerializeException(e.getMessage(), e);
@@ -120,24 +126,25 @@ public class RuleBinarySerializer implements RuleSerializer {
         outputStream.write(bitOutputStream.toByteArray());
     }
 
-    private Map<String, Integer> serializeRuleList(TreeMap<String, List<Rule>> rulesMap,
-            ByteTrackedOutputStream outputStream)
+    private LinkedHashMap<String, Integer> serializeRuleList(
+            Map<String, List<Rule>> rulesMap, ByteTrackedOutputStream outputStream)
             throws IOException {
         Preconditions.checkArgument(rulesMap != null,
                 "serializeRuleList should never be called with null rule list.");
 
         BitOutputStream bitOutputStream = new BitOutputStream();
-        Map<String, Integer> indexMapping = new LinkedHashMap();
-        int indexTracker = 0;
-
+        LinkedHashMap<String, Integer> indexMapping = new LinkedHashMap();
         indexMapping.put(START_INDEXING_KEY, outputStream.getWrittenBytesCount());
-        for (Map.Entry<String, List<Rule>> entry : rulesMap.entrySet()) {
+
+        List<String> sortedKeys = rulesMap.keySet().stream().sorted().collect(Collectors.toList());
+        int indexTracker = 0;
+        for (String key : sortedKeys) {
             if (indexTracker >= INDEXING_BLOCK_SIZE) {
-                indexMapping.put(entry.getKey(), outputStream.getWrittenBytesCount());
+                indexMapping.put(key, outputStream.getWrittenBytesCount());
                 indexTracker = 0;
             }
 
-            for (Rule rule : entry.getValue()) {
+            for (Rule rule : rulesMap.get(key)) {
                 bitOutputStream.clear();
                 serializeRule(rule, bitOutputStream);
                 outputStream.write(bitOutputStream.toByteArray());
@@ -222,11 +229,13 @@ public class RuleBinarySerializer implements RuleSerializer {
     }
 
     private void serializeIndexGroup(
-            Map<String, Integer> indexes, BitOutputStream bitOutputStream, boolean isIndexed) {
+            LinkedHashMap<String, Integer> indexes,
+            BitOutputStream bitOutputStream,
+            boolean isIndexed) {
 
         // Output the starting location of this indexing group.
-        serializeStringValue(START_INDEXING_KEY, /* isHashedValue= */false,
-                bitOutputStream);
+        serializeStringValue(
+                START_INDEXING_KEY, /* isHashedValue= */false, bitOutputStream);
         serializeIntValue(indexes.get(START_INDEXING_KEY), bitOutputStream);
 
         // If the group is indexed, output the locations of the indexes.
@@ -244,9 +253,6 @@ public class RuleBinarySerializer implements RuleSerializer {
         // Output the end location of this indexing group.
         serializeStringValue(END_INDEXING_KEY, /*isHashedValue= */ false, bitOutputStream);
         serializeIntValue(indexes.get(END_INDEXING_KEY), bitOutputStream);
-
-        // This dummy bit is set for fixing the padding issue. songpan@ will fix it and remove it.
-        bitOutputStream.setNext();
     }
 
     private void serializeStringValue(
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetails.java b/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetails.java
index dd871e2bbe6c..2cbd4ede5214 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetails.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetails.java
@@ -28,6 +28,8 @@ class RuleIndexingDetails {
     static final int PACKAGE_NAME_INDEXED = 1;
     static final int APP_CERTIFICATE_INDEXED = 2;
 
+    static final String DEFAULT_RULE_KEY = "N/A";
+
     /** Represents which indexed file the rule should be located. */
     @IntDef(
             value = {
@@ -45,7 +47,7 @@ class RuleIndexingDetails {
     /** Constructor without a ruleKey for {@code NOT_INDEXED}. */
     RuleIndexingDetails(@IndexType int indexType) {
         this.mIndexType = indexType;
-        this.mRuleKey = null;
+        this.mRuleKey = DEFAULT_RULE_KEY;
     }
 
     /** Constructor with a ruleKey for indexed rules. */
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifier.java b/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifier.java
index cbc365e2c250..7d9a90188983 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifier.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifier.java
@@ -30,30 +30,27 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.TreeMap;
 
 /** A helper class for identifying the indexing type and key of a given rule. */
 class RuleIndexingDetailsIdentifier {
 
-    private static final String DEFAULT_RULE_KEY = "N/A";
-
     /**
      * Splits a given rule list into three indexing categories. Each rule category is returned as a
      * TreeMap that is sorted by their indexing keys -- where keys correspond to package name for
      * PACKAGE_NAME_INDEXED rules, app certificate for APP_CERTIFICATE_INDEXED rules and N/A for
      * NOT_INDEXED rules.
      */
-    public static Map<Integer, TreeMap<String, List<Rule>>> splitRulesIntoIndexBuckets(
+    public static Map<Integer, Map<String, List<Rule>>> splitRulesIntoIndexBuckets(
             List<Rule> rules) {
         if (rules == null) {
             throw new IllegalArgumentException(
                     "Index buckets cannot be created for null rule list.");
         }
 
-        Map<Integer, TreeMap<String, List<Rule>>> typeOrganizedRuleMap = new HashMap();
-        typeOrganizedRuleMap.put(NOT_INDEXED, new TreeMap());
-        typeOrganizedRuleMap.put(PACKAGE_NAME_INDEXED, new TreeMap());
-        typeOrganizedRuleMap.put(APP_CERTIFICATE_INDEXED, new TreeMap());
+        Map<Integer, Map<String, List<Rule>>> typeOrganizedRuleMap = new HashMap();
+        typeOrganizedRuleMap.put(NOT_INDEXED, new HashMap());
+        typeOrganizedRuleMap.put(PACKAGE_NAME_INDEXED, new HashMap<>());
+        typeOrganizedRuleMap.put(APP_CERTIFICATE_INDEXED, new HashMap<>());
 
         // Split the rules into the appropriate indexed pattern. The Tree Maps help us to keep the
         // entries sorted by their index key.
@@ -66,21 +63,14 @@ class RuleIndexingDetailsIdentifier {
                         String.format("Malformed rule identified. [%s]", rule.toString()));
             }
 
-            String ruleKey =
-                    indexingDetails.getIndexType() != NOT_INDEXED
-                            ? indexingDetails.getRuleKey()
-                            : DEFAULT_RULE_KEY;
+            int ruleIndexType = indexingDetails.getIndexType();
+            String ruleKey = indexingDetails.getRuleKey();
 
-            if (!typeOrganizedRuleMap.get(indexingDetails.getIndexType()).containsKey(ruleKey)) {
-                typeOrganizedRuleMap
-                        .get(indexingDetails.getIndexType())
-                        .put(ruleKey, new ArrayList());
+            if (!typeOrganizedRuleMap.get(ruleIndexType).containsKey(ruleKey)) {
+                typeOrganizedRuleMap.get(ruleIndexType).put(ruleKey, new ArrayList());
             }
 
-            typeOrganizedRuleMap
-                    .get(indexingDetails.getIndexType())
-                    .get(ruleKey)
-                    .add(rule);
+            typeOrganizedRuleMap.get(ruleIndexType).get(ruleKey).add(rule);
         }
 
         return typeOrganizedRuleMap;
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java b/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
index 4194432375b8..8f164e645434 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
@@ -35,7 +35,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.TreeMap;
+import java.util.stream.Collectors;
 
 /** A helper class to serialize rules from the {@link Rule} model to Xml representation. */
 public class RuleXmlSerializer implements RuleSerializer {
@@ -90,7 +90,7 @@ public class RuleXmlSerializer implements RuleSerializer {
             throws RuleSerializeException {
         try {
             // Determine the indexing groups and the order of the rules within each indexed group.
-            Map<Integer, TreeMap<String, List<Rule>>> indexedRules =
+            Map<Integer, Map<String, List<Rule>>> indexedRules =
                     RuleIndexingDetailsIdentifier.splitRulesIntoIndexBuckets(rules);
 
             // Write the XML formatted rules in order.
@@ -107,11 +107,12 @@ public class RuleXmlSerializer implements RuleSerializer {
         }
     }
 
-    private void serializeRuleList(TreeMap<String, List<Rule>> rulesMap,
-            XmlSerializer xmlSerializer)
+    private void serializeRuleList(Map<String, List<Rule>> rulesMap, XmlSerializer xmlSerializer)
             throws IOException {
-        for (Map.Entry<String, List<Rule>> entry : rulesMap.entrySet()) {
-            for (Rule rule : entry.getValue()) {
+        List<String> sortedKeyList =
+                rulesMap.keySet().stream().sorted().collect(Collectors.toList());
+        for (String key : sortedKeyList) {
+            for (Rule rule : rulesMap.get(key)) {
                 serializeRule(rule, xmlSerializer);
             }
         }
diff --git a/services/tests/servicestests/src/com/android/server/integrity/IntegrityFileManagerTest.java b/services/tests/servicestests/src/com/android/server/integrity/IntegrityFileManagerTest.java
index dd69c6613ab2..5aed194773f5 100644
--- a/services/tests/servicestests/src/com/android/server/integrity/IntegrityFileManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/integrity/IntegrityFileManagerTest.java
@@ -185,7 +185,7 @@ public class IntegrityFileManagerTest {
         List<Rule> rulesFetched = mIntegrityFileManager.readRules(appInstallMetadata);
 
         // Verify that we do not load all the rules and we have the necessary rules to evaluate.
-        assertThat(rulesFetched.size()).isEqualTo(170);
+        assertThat(rulesFetched.size()).isEqualTo(270);
         assertThat(rulesFetched)
                 .containsAllOf(
                         getPackageNameIndexedRule(installedPackageName),
diff --git a/services/tests/servicestests/src/com/android/server/integrity/model/BitTrackedInputStreamTest.java b/services/tests/servicestests/src/com/android/server/integrity/model/BitTrackedInputStreamTest.java
index 1eb5eb51504a..4fa73028ece1 100644
--- a/services/tests/servicestests/src/com/android/server/integrity/model/BitTrackedInputStreamTest.java
+++ b/services/tests/servicestests/src/com/android/server/integrity/model/BitTrackedInputStreamTest.java
@@ -121,6 +121,26 @@ public class BitTrackedInputStreamTest {
     }
 
     @Test
+    public void testBitTrackedInputStream_canReadMoreRules() throws IOException {
+        String packageName = "com.test.app";
+        byte[] testInput =
+                getBytes(
+                        IS_NOT_HASHED
+                                + getBits(packageName.length(), VALUE_SIZE_BITS)
+                                + getValueBits(packageName));
+
+        BitTrackedInputStream bitTrackedInputStream = new BitTrackedInputStream(testInput);
+        assertThat(bitTrackedInputStream.canReadMoreRules(2)).isTrue();
+
+        // Read until the string parameter.
+        String stringValue = BinaryFileOperations.getStringValue(bitTrackedInputStream);
+
+        // Verify that the read bytes are counted.
+        assertThat(stringValue).isEqualTo(packageName);
+        assertThat(bitTrackedInputStream.canReadMoreRules(2)).isFalse();
+    }
+
+    @Test
     public void testBitTrackedInputStream_moveCursorForwardFailsIfAlreadyRead() throws IOException {
         String packageName = "com.test.app";
         byte[] testInput =
diff --git a/services/tests/servicestests/src/com/android/server/integrity/parser/RuleIndexingControllerTest.java b/services/tests/servicestests/src/com/android/server/integrity/parser/RuleIndexingControllerTest.java
index 742952e056bc..291cfeea5bbc 100644
--- a/services/tests/servicestests/src/com/android/server/integrity/parser/RuleIndexingControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/integrity/parser/RuleIndexingControllerTest.java
@@ -115,7 +115,8 @@ public class RuleIndexingControllerTest {
                                 + getKeyValueString(START_INDEXING_KEY, 500)
                                 + getKeyValueString(END_INDEXING_KEY, 900)
                                 + getKeyValueString(START_INDEXING_KEY, 900)
-                                + getKeyValueString(END_INDEXING_KEY, 945));
+                                + getKeyValueString(END_INDEXING_KEY, 945)
+                                + getBits(1, 1));
         ByteBuffer rule = ByteBuffer.allocate(stringBytes.length);
         rule.put(stringBytes);
         InputStream inputStream = new ByteArrayInputStream(rule.array());
@@ -152,7 +153,8 @@ public class RuleIndexingControllerTest {
                                 + getKeyValueString("888", 800)
                                 + getKeyValueString(END_INDEXING_KEY, 900)
                                 + getKeyValueString(START_INDEXING_KEY, 900)
-                                + getKeyValueString(END_INDEXING_KEY, 945));
+                                + getKeyValueString(END_INDEXING_KEY, 945)
+                                + getBits(1, 1));
         ByteBuffer rule = ByteBuffer.allocate(stringBytes.length);
         rule.put(stringBytes);
         return new ByteArrayInputStream(rule.array());
diff --git a/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleBinarySerializerTest.java b/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleBinarySerializerTest.java
index eb6698b0d479..bc2aac0acf10 100644
--- a/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleBinarySerializerTest.java
+++ b/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleBinarySerializerTest.java
@@ -135,15 +135,14 @@ public class RuleBinarySerializerTest {
                 .isEqualTo(expectedRuleOutputStream.toByteArray());
 
         ByteArrayOutputStream expectedIndexingOutputStream = new ByteArrayOutputStream();
+        String serializedIndexingBytes =
+                SERIALIZED_START_INDEXING_KEY
+                        + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */ 32)
+                        + SERIALIZED_END_INDEXING_KEY
+                        + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */32);
         byte[] expectedIndexingBytes =
-                getBytes(
-                        SERIALIZED_START_INDEXING_KEY
-                                + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */ 32)
-                                + SERIALIZED_END_INDEXING_KEY
-                                + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */
-                                32));
-        expectedIndexingOutputStream.write(expectedIndexingBytes);
-        expectedIndexingOutputStream.write(expectedIndexingBytes);
+                getBytes(serializedIndexingBytes + serializedIndexingBytes
+                        + serializedIndexingBytes + getBits(1, 1));
         expectedIndexingOutputStream.write(expectedIndexingBytes);
         assertThat(indexingOutputStream.toByteArray())
                 .isEqualTo(expectedIndexingOutputStream.toByteArray());
@@ -197,15 +196,16 @@ public class RuleBinarySerializerTest {
                         + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */ 32)
                         + SERIALIZED_END_INDEXING_KEY
                         + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */ 32);
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBitsForIndexed));
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBitsForIndexed));
         String expectedIndexingBitsForUnindexed =
                 SERIALIZED_START_INDEXING_KEY
                         + getBits(DEFAULT_FORMAT_VERSION_BYTES.length, /* numOfBits= */ 32)
                         + SERIALIZED_END_INDEXING_KEY
-                        + getBits(DEFAULT_FORMAT_VERSION_BYTES.length + getBytes(
-                        expectedBits).length, /* numOfBits= */ 32);
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBitsForUnindexed));
+                        + getBits(DEFAULT_FORMAT_VERSION_BYTES.length
+                        + getBytes(expectedBits).length, /* numOfBits= */ 32);
+        expectedIndexingOutputStream.write(getBytes(
+                expectedIndexingBitsForIndexed + expectedIndexingBitsForIndexed
+                        + expectedIndexingBitsForUnindexed + getBits(1, 1)));
+
         assertThat(indexingOutputStream.toByteArray())
                 .isEqualTo(expectedIndexingOutputStream.toByteArray());
     }
@@ -564,7 +564,6 @@ public class RuleBinarySerializerTest {
         expectedIndexingBytesForPackageNameIndexed +=
                 SERIALIZED_END_INDEXING_KEY
                         + getBits(totalBytesWritten, /* numOfBits= */ 32);
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBytesForPackageNameIndexed));
 
         String expectedIndexingBytesForAppCertificateIndexed =
                 SERIALIZED_START_INDEXING_KEY
@@ -588,7 +587,6 @@ public class RuleBinarySerializerTest {
         expectedIndexingBytesForAppCertificateIndexed +=
                 SERIALIZED_END_INDEXING_KEY
                         + getBits(totalBytesWritten, /* numOfBits= */ 32);
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBytesForAppCertificateIndexed));
 
         String expectedIndexingBytesForUnindexed =
                 SERIALIZED_START_INDEXING_KEY
@@ -603,8 +601,11 @@ public class RuleBinarySerializerTest {
         expectedIndexingBytesForUnindexed +=
                 SERIALIZED_END_INDEXING_KEY
                         + getBits(totalBytesWritten, /* numOfBits= */ 32);
-        expectedIndexingOutputStream.write(getBytes(expectedIndexingBytesForUnindexed));
-
+        expectedIndexingOutputStream.write(
+                getBytes(expectedIndexingBytesForPackageNameIndexed
+                        + expectedIndexingBytesForAppCertificateIndexed
+                        + expectedIndexingBytesForUnindexed
+                        + getBits(1, 1)));
 
         assertThat(ruleOutputStream.toByteArray())
                 .isEqualTo(expectedOrderedRuleOutputStream.toByteArray());
diff --git a/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifierTest.java b/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifierTest.java
index 55fada44b99e..1674422f3af9 100644
--- a/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifierTest.java
+++ b/services/tests/servicestests/src/com/android/server/integrity/serializer/RuleIndexingDetailsIdentifierTest.java
@@ -38,10 +38,8 @@ import org.junit.runners.JUnit4;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.TreeMap;
 
 /** Unit tests for {@link RuleIndexingDetailsIdentifier}. */
 @RunWith(JUnit4.class)
@@ -140,7 +138,7 @@ public class RuleIndexingDetailsIdentifierTest {
         List<Rule> ruleList = new ArrayList();
         ruleList.add(RULE_WITH_PACKAGE_NAME);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         // Verify the resulting map content.
         assertThat(result.keySet())
@@ -157,7 +155,7 @@ public class RuleIndexingDetailsIdentifierTest {
         List<Rule> ruleList = new ArrayList();
         ruleList.add(RULE_WITH_APP_CERTIFICATE);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         assertThat(result.keySet())
                 .containsExactly(NOT_INDEXED, PACKAGE_NAME_INDEXED, APP_CERTIFICATE_INDEXED);
@@ -174,7 +172,7 @@ public class RuleIndexingDetailsIdentifierTest {
         List<Rule> ruleList = new ArrayList();
         ruleList.add(RULE_WITH_INSTALLER_RESTRICTIONS);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         assertThat(result.keySet())
                 .containsExactly(NOT_INDEXED, PACKAGE_NAME_INDEXED, APP_CERTIFICATE_INDEXED);
@@ -189,7 +187,7 @@ public class RuleIndexingDetailsIdentifierTest {
         List<Rule> ruleList = new ArrayList();
         ruleList.add(RULE_WITH_NONSTRING_RESTRICTIONS);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         assertThat(result.keySet())
                 .containsExactly(NOT_INDEXED, PACKAGE_NAME_INDEXED, APP_CERTIFICATE_INDEXED);
@@ -215,7 +213,7 @@ public class RuleIndexingDetailsIdentifierTest {
         List<Rule> ruleList = new ArrayList();
         ruleList.add(negatedRule);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         assertThat(result.keySet())
                 .containsExactly(NOT_INDEXED, PACKAGE_NAME_INDEXED, APP_CERTIFICATE_INDEXED);
@@ -225,7 +223,7 @@ public class RuleIndexingDetailsIdentifierTest {
     }
 
     @Test
-    public void getIndexType_allRulesTogetherInValidOrder() {
+    public void getIndexType_allRulesTogetherSplitCorrectly() {
         Rule packageNameRuleA = getRuleWithPackageName("aaa");
         Rule packageNameRuleB = getRuleWithPackageName("bbb");
         Rule packageNameRuleC = getRuleWithPackageName("ccc");
@@ -243,38 +241,20 @@ public class RuleIndexingDetailsIdentifierTest {
         ruleList.add(RULE_WITH_INSTALLER_RESTRICTIONS);
         ruleList.add(RULE_WITH_NONSTRING_RESTRICTIONS);
 
-        Map<Integer, TreeMap<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
+        Map<Integer, Map<String, List<Rule>>> result = splitRulesIntoIndexBuckets(ruleList);
 
         assertThat(result.keySet())
                 .containsExactly(NOT_INDEXED, PACKAGE_NAME_INDEXED, APP_CERTIFICATE_INDEXED);
 
         // We check asserts this way to ensure ordering based on package name.
         assertThat(result.get(PACKAGE_NAME_INDEXED).keySet()).containsExactly("aaa", "bbb", "ccc");
-        Iterator<String> keySetIterator = result.get(PACKAGE_NAME_INDEXED).keySet().iterator();
-        assertThat(keySetIterator.next()).isEqualTo("aaa");
-        assertThat(keySetIterator.next()).isEqualTo("bbb");
-        assertThat(keySetIterator.next()).isEqualTo("ccc");
-        assertThat(result.get(PACKAGE_NAME_INDEXED).get("aaa")).containsExactly(packageNameRuleA);
-        assertThat(result.get(PACKAGE_NAME_INDEXED).get("bbb")).containsExactly(packageNameRuleB);
-        assertThat(result.get(PACKAGE_NAME_INDEXED).get("ccc")).containsExactly(packageNameRuleC);
 
         // We check asserts this way to ensure ordering based on app certificate.
         assertThat(result.get(APP_CERTIFICATE_INDEXED).keySet()).containsExactly("cert1", "cert2",
                 "cert3");
-        keySetIterator = result.get(APP_CERTIFICATE_INDEXED).keySet().iterator();
-        assertThat(keySetIterator.next()).isEqualTo("cert1");
-        assertThat(keySetIterator.next()).isEqualTo("cert2");
-        assertThat(keySetIterator.next()).isEqualTo("cert3");
-        assertThat(result.get(APP_CERTIFICATE_INDEXED).get("cert1")).containsExactly(
-                certificateRule1);
-        assertThat(result.get(APP_CERTIFICATE_INDEXED).get("cert2")).containsExactly(
-                certificateRule2);
-        assertThat(result.get(APP_CERTIFICATE_INDEXED).get("cert3")).containsExactly(
-                certificateRule3);
 
         assertThat(result.get(NOT_INDEXED).get("N/A"))
-                .containsExactly(
-                        RULE_WITH_INSTALLER_RESTRICTIONS,
+                .containsExactly(RULE_WITH_INSTALLER_RESTRICTIONS,
                         RULE_WITH_NONSTRING_RESTRICTIONS);
     }