Fix and extend hasOverflowed

dexlib2/src/main/java/org/jf/dexlib2/writer/DexWriter.java

@@ -107,6 +107,8 @@ public abstract class DexWriter<
     public static final int NO_INDEX = -1;
     public static final int NO_OFFSET = 0;
 
+    public static final int MAX_POOL_SIZE = (1 << 16);
+
     protected final Opcodes opcodes;
 
     protected int stringIndexSectionOffset = NO_OFFSET;
@@ -150,6 +152,8 @@ public abstract class DexWriter<
     public final AnnotationSetSectionType annotationSetSection;
     public final EncodedArraySectionType encodedArraySection;
 
+    private final IndexSection<?>[] overflowableSections;
+
     protected DexWriter(Opcodes opcodes) {
         this.opcodes = opcodes;
 
@@ -166,6 +170,17 @@ protected DexWriter(Opcodes opcodes) {
         this.annotationSection = sectionProvider.getAnnotationSection();
         this.annotationSetSection = sectionProvider.getAnnotationSetSection();
         this.encodedArraySection = sectionProvider.getEncodedArraySection();
+
+        overflowableSections = new IndexSection<?>[] {
+                //stringSection,            // supports jumbo indexes
+                typeSection,
+                protoSection,
+                fieldSection,
+                methodSection,
+                //classSection,             // redundant check: cannot be larger than typeSection
+                callSiteSection,
+                methodHandleSection,
+        };
     }
 
     @Nonnull protected abstract SectionProvider getSectionProvider();
@@ -250,19 +265,28 @@ public List<String> getTypeReferences() {
     }
 
     /**
-     * Checks whether any of the size-sensitive constant pools have overflowed.
-     *
-     * This checks whether the type, method, field pools are larger than 64k entries.
+     * Checks whether any of the size-sensitive constant pools have overflowed and have more than 64Ki entries.
      *
      * Note that even if this returns true, it may still be possible to successfully write the dex file, if the
-     * overflowed items are not referenced anywhere that uses a 16-bit index
+     * overflowed items are not referenced anywhere that uses a 16-bit index.
      *
      * @return true if any of the size-sensitive constant pools have overflowed
      */
     public boolean hasOverflowed() {
-        return methodSection.getItemCount() > (1 << 16) ||
-                typeSection.getItemCount() > (1 << 16) ||
-                fieldSection.getItemCount() > (1 << 16);
+        return hasOverflowed(MAX_POOL_SIZE);
+    }
+
+    /**
+     * Checks whether any of the size-sensitive constant pools have more than the supplied maximum number of entries.
+     *
+     * @param maxPoolSize the maximum number of entries allowed in any of the size-sensitive constant pools
+     * @return true if any of the size-sensitive constant pools have overflowed the supplied size limit
+     */
+    public boolean hasOverflowed(int maxPoolSize) {
+        for (IndexSection section: overflowableSections) {
+            if (section.getItemCount() > maxPoolSize) return true;
+        }
+        return false;
     }
 
     public void writeTo(@Nonnull DexDataStore dest) throws IOException {