Javassit StackMapTable.Walker that shifts the first stack map frame.

Shifter.java

/**
 * A {@link StackMapTable.Walker} that shifts the offset_delta of the first frame in the table.
 *
 * The JVM Spec Section 4.7.4 says:
 *
 * "Each stack map frame described in the entries table relies on the previous frame for some of its
 * semantics. The first stack map frame of a method is implicit, and computed from the method
 * descriptor by the type checker. The stack_map_frame structure at entries[0] therefore describes
 * the second stack map frame of the method."
 *
 * It also says:
 *
 * "The bytecode offset at which a stack map frame applies is calculated by taking the value
 * offset_delta specified in the frame (either explicitly or implicitly), and adding
 * offset_delta + 1 to the bytecode offset of the previous frame, unless the previous frame is the
 * initial frame of the method. In that case, the bytecode offset at which the stack map frame
 * applies is the value offset_delta specified in the frame."
 *
 * Given the statements above we can conclude that it is is enough to only offset the frame at
 * entries[0] since all subsequent frames will consequentially be shifted due to the formula
 * 'offset_delta + 1' when inserting the simple `invokespecial` instruction at the start of
 * a method.
 *
 * @see <a href="https://docs.oracle.com/javase/specs/jvms/se8/html/jvms-4.html#jvms-4.7.4">JVM Spec §4.7.4</a>
 */
public class Shifter extends StackMapTable.Walker {

    private final AttributeInfo attributeInfo;
    private final byte[] info;
    private final int shiftAmount;
    private boolean isFrameUpdated = false;
    private byte[] updatedInfo = null;

    public Shifter(StackMapTable smt, int shiftAmount) {
        super(smt);
        this.attributeInfo = smt;
        this.info = smt.get();
        this.shiftAmount = shiftAmount;
    }

    public void performShift() {
        try {
            parse();
            if (updatedInfo != null) {
                attributeInfo.set(updatedInfo);
            }
        } catch (BadBytecode badBytecode) {
            throw new RuntimeException(badBytecode);
        }
    }

    @Override
    public void sameFrame(int pos, int offsetDelta) throws BadBytecode {
        update(pos, offsetDelta, 0,251 /* same_frame_extended */);
    }

    @Override
    public void sameLocals(int pos, int offsetDelta, int stackTag, int stackData) throws BadBytecode {
        update(pos, offsetDelta, 64, 247 /* same_locals_1_stack_item_frame_extended */);
    }

    @Override
    public void chopFrame(int pos, int offsetDelta, int k) throws BadBytecode {
        update(pos, offsetDelta);
    }

    @Override
    public void appendFrame(int pos, int offsetDelta, int[] tags, int[] data) throws BadBytecode {
        update(pos, offsetDelta);
    }

    @Override
    public void fullFrame(int pos, int offsetDelta, int[] localTags, int[] localData, int[] stackTags, int[] stackData) throws BadBytecode {
        update(pos, offsetDelta);
    }

    private void update(int pos, int offsetDelta, int deltaBase, int extraEntryType) {
        if (isFrameUpdated) {
            // We are only interested in the first non-implicit frame,
            // i.e. the second stack map frame of the method.
            return;
        }
        isFrameUpdated = true;

        int newDelta = offsetDelta + shiftAmount;
        if (newDelta < 64) {
            // Updates frame_type for both same_frame and same_locals_1_stack_item_frame.
            info[pos] = (byte) (newDelta + deltaBase);
        } else if (offsetDelta < 64) {
            // Creates a new frame, either same_frame_extended or same_locals_1_stack_item_frame_extended.
            byte[] newInfo = insertGap(info, pos, 2);
            newInfo[pos] = (byte) extraEntryType;
            ByteArray.write16bit(newDelta, newInfo, pos + 1);
            updatedInfo = newInfo;
        } else {
            // Updates offset_delta for same_frame_extended or same_locals_1_stack_item_frame_extended.
            ByteArray.write16bit(newDelta, info, pos + 1);
        }
    }

    private void update(int pos, int offsetDelta) {
        if (isFrameUpdated) {
            // We are only interested in the first non-implicit frame,
            // i.e. the second stack map frame of the method.
            return;
        }
        isFrameUpdated = true;
        
        // Updates offset_delta for chop_frame, append_frame or full_frame.
        int newDelta = offsetDelta + shiftAmount;
        ByteArray.write16bit(newDelta, info, pos + 1);
    }

    private static byte[] insertGap(byte[] info, int where, int gap) {
        if (where < -1) {
            throw new IllegalArgumentException("gap location can't be negative");
        }
        if (where > info.length) {
            throw new IllegalArgumentException("gap location outside array");
        }
        if (gap <= 0) {
            throw new IllegalArgumentException("gap size has to be greater than zero");
        }
        byte[] newInfo = new byte[info.length + gap];
        // Copy first half up to gap location
        System.arraycopy(info, 0, newInfo, 0, where);
        // Copy second half, leaving the gap in between
        System.arraycopy(info, where, newInfo, where + gap, info.length - where);
        return newInfo;
    }
}