NonDex instrumentation of IdentityHashMap can break with `removeAll` calls

[lycoris106]

Problem

When running under NonDex, using IdentityHashMap.keySet().removeAll(...) (or entrySet().removeAll(...))
can throw ArrayIndexOutOfBoundsException inside java.util.IdentityHashMap$IdentityHashMapIterator.

This does not fail under the normal JDK and can be reproduced in a minimal test case.

Root Cause

The NonDex instrumentation inIdentityHashMapShufflingAdder.java uses a shuffled list of keys keys and idx to iterate, but it does not modify the original remove method for iterator accordingly. The problematic trace:

1. removeAll() iterates using the map's iterator while calling remove()
2. Sometimes in remove, the traversed traversalTable is cut-off to avoid revisiting elements due to replacement in the closeDeletion process.
3. This works completely fine in the original jdk's version because the iterator is running sequentially from the start of table, so that after the cut-off, all non-visited elements are correctly placed in the sub-table (the later half of the original table).
4. However, in nondex's version, it is using idx instead of index to run through the shuffled keys. Now when the remove cut-offs the table, the sub-table does not contain the correct non-visited elements since it's visited in shuffled order.
5. The idx keeps increasing as if the table is of the original size, leading to out-of-index error

Reproduce

import org.junit.Test;
import java.util.*;

public class TestIdentityHashMap {

    @Test
    public void nondexIdentityKeySetTest() {
        for (int i = 0; i < 10_000; i++) {
            IdentityHashMap<Object, Boolean> map = new IdentityHashMap<>();

            Object a = new Object();
            Object b = new Object();
            Object c = new Object();

            map.put(a, Boolean.TRUE);
            map.put(b, Boolean.TRUE);
            map.put(c, Boolean.TRUE);

            Set<Object> s = map.keySet();

            s.removeAll(Collections.singleton(a));
        }
    }
}

Possible Fix

A possible fix is to delete the block in the remove method through bytecode injection:

if (i < deletedSlot && d >= deletedSlot &&
                        traversalTable == IdentityHashMap.this.table) {
                        int remaining = len - deletedSlot;
                        Object[] newTable = new Object[remaining];
                        System.arraycopy(tab, deletedSlot,
                                         newTable, 0, remaining);
                        traversalTable = newTable;
                        index = 0;
                    }

The purpose of this block is to avoid revisiting keys after some keys are moved due to the closeDeletion process. However, this closeDeletion process would not affect nondex's version of iteration nextIndex() since it always loops through the whole table to find the key.