Peephole #2: Folding Arithmetic, and the Fold You Must Refuse

This is the second post in a series where we build upon a simple Java bytecode peephole optimiser: each post takes a single peephole optimisation for Java bytecode, shows how to implement it, and looks at the pitfalls to watch out for. If you’re new to the Class-File API, start with Java Class-File API Hello World.

We previously implemented a peephole that removed additions of zero like x + 0 and ended that post with “try implementing more arithmetic optimisations where constants are involved” as an exercise.

Now we’ll take on that exercise: peepholes that simplify arithmetic on constants, calculating the result while optimising where it is safe to do so. We’ll also take a look at one arithmetic fold that has to be refused, and at where in a peephole a refusal belongs.

The pattern

The shape we’re after is a pair of constants followed by an operation:

iconst_2
iconst_3
imul

If you compile something like return 2 * 3; with javac, you’ll notice that javac already does constant folding during compilation:

static int mul();
  Code:
       0: bipush        6
       2: ireturn

Folding like this is normally done in a compiler front end, rather than on the resulting bytecode. jlox, a Lox compiler I built for the JVM, folds on the syntax tree: its optimiser replaces a binary expression whose operands are both literals with a literal, and simplifies x + 0, x - 0 and x * 0 alongside it. That is this post’s two rules, one level up.

My C++ implementation of Lox, cpplox, implements no folding at all and gets it anyway, because LLVM’s IRBuilder folds constant operands as the IR is built. Ask it for fadd 2.0, 3.0 and you get the constant 5.0 back; you have to ask for an IRBuilder<NoFolder> to see the instruction you thought you were making.

So bytecode (or LLVM IR, or machine code) as produced by a compiler may contain no constant folding opportunities at all. The pattern turns up in code the compiler didn’t write: more interestingly, bytecode that another optimisation has just rewritten; that’s where peephole optimisations work together, each one small in itself, but potentially opening up opportunities for other peepholes.

The rewrite

The peephole matches three elements, two ConstantInstructions followed by an Instruction, and replaces them with a single constant:

class ConstantFold implements Peephole {
    public String name() { return "constantFold"; }

    public Rewrite apply(CodeElement[] window, ConstantPoolBuilder pool) {
        if (window[0] instanceof ConstantInstruction a &&
                window[1] instanceof ConstantInstruction b &&
                window[2] instanceof Instruction op) {
            if (a.constantValue() instanceof Integer x && b.constantValue() instanceof Integer y) {
                var result = foldInt(op.opcode(), x, y);
                if (result != null) {
                    return new Rewrite(3, List.of(intConstant(result, pool)));
                }
            }
        }
        return null;
    }
}

constantValue() gives us the constant that an instruction pushes, as a ConstantDesc: an Integer for the int forms (iconst_*, bipush, sipush and an ldc of an int), and a Long, Float, Double or String for the rest.

The foldInt method computes the result for the instructions it supports, and returns null for anything else.

static Integer foldInt(Opcode opcode, int x, int y) {
    return switch (opcode) {
        case IADD -> x + y;
        case ISUB -> x - y;
        case IMUL -> x * y;
        case IAND -> x & y;
        case IOR  -> x | y;
        case IXOR -> x ^ y;
        case ISHL -> x << y;
        case ISHR -> x >> y;
        case IUSHR -> x >>> y;
        default -> null; // Unsupported opcode for folding.
    };
}

The intConstant method is a helper that takes a Java integer and returns a Class-File API CodeElement instruction. Depending on the magnitude of the integer, a different instruction can be used. There are compact instructions for small numbers, such as iconst_0 for pushing a 0 onto the stack, while larger numbers require a combination of instruction and operand (e.g. bipush 42).

Even larger values go in the class file’s constant pool rather than inline in the bytecode. The instructions that carry a value inline have a small operand: bipush takes one byte and sipush two, so anything outside the range of a short is pushed with ldc, whose operand is an index into the pool.

static CodeElement intConstant(int value, ConstantPoolBuilder pool) {
    return switch (value) {
        // Compact instructions for small integers
        case -1 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_M1);
        case 0 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_0);
        case 1 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_1);
        case 2 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_2);
        case 3 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_3);
        case 4 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_4);
        case 5 -> ConstantInstruction.ofIntrinsic(Opcode.ICONST_5);
        default -> {
            if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) {
                // A byte sized integer
                yield ConstantInstruction.ofArgument(BIPUSH, value);
            }
            if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) {
                // A short sized integer
                yield ConstantInstruction.ofArgument(SIPUSH, value);
            }
            // Larger than short
            var entry = pool.intEntry(value);
            // ldc takes a one byte constant pool index; use ldc_w otherwise.
            yield ConstantInstruction.ofLoad(entry.index() <= 0xFF ? LDC : LDC_W, entry);
        }
    };
}

Arithmetic that does nothing

ConstantFold computes a result, so it needs both operands to be constants. A related family needs only one: the operation gives back the value it was handed, so it can be deleted instead of folded. Adding or subtracting zero, multiplying or dividing by one, and shifting by zero all leave the value unchanged, and so does negating twice.

The add-zero peephole from the original post was a member of that family, so we can extend it with the rest. IdentityFold matches two elements instead of three, and replaces them with nothing.

Show IdentityFold

class IdentityFold implements Peephole {
    public String name() { return "identityFold"; }

    public Rewrite apply(CodeElement[] window, ConstantPoolBuilder pool) {
        // double negation
        if (window[0] instanceof Instruction i0 && i0.opcode() == Opcode.INEG &&
                window[1] instanceof Instruction i1 && i1.opcode() == Opcode.INEG) {
            return new Rewrite(2, List.of());
        }

        if (window[0] instanceof ConstantInstruction c &&
                c.constantValue() instanceof Integer value &&
                window[1] instanceof Instruction op) {
            var identity = switch (op.opcode()) {
                // Add, subtract, or shift by zero
                case IADD, ISUB, ISHL, ISHR, IUSHR -> value == 0;
                // Multiply or divide by one
                case IMUL, IDIV -> value == 1;
                default -> false;
            };
            if (identity) {
                return new Rewrite(2, List.of());
            }
        }

        return null;
    }
}

Applying the folds

Both peepholes go in the optimiser’s list, and the optimiser will try them at every window position. The TestJarGenerator builds a Test class with the Class-File API rather than with javac, since javac folds constants itself, so this makes sure our peepholes have something to fold.

Generate the jar, optimise it, and run both:

$ java TestJarGenerator.java test.jar
$ java Optimizer.java test.jar test-opt.jar
appendMerge: 2
constantFold: 1
identityFold: 1
storeLoad: 1
$ java -jar test-opt.jar arg1 arg2
The length of the arguments array is 2 and 2 * 3 is 6

Disassembling main with javap -c shows what constantFold did to the 2 * 3. Here are the three instructions before:

iconst_2
iconst_3
imul

which became a single instruction after constant folding:

      24: bipush        6

identityFold matched once too, on the args.length + 0, removing both of these instructions:

      26: iconst_0
      27: iadd

The appendMerge and storeLoad counts come from peepholes elsewhere in the series: appendMerge joined the three constant append calls into one, and storeLoad replaced the reload with a dup.

The catch: division by zero

You might have noticed that we didn’t add the division and remainder instructions to the ConstantFold peephole: that’s because they need special treatment, depending on what the constants are.

If you were to fold the following, what would happen?

static int divZero(int);
  Code:
       0: bipush        10
       2: iconst_0
       3: idiv
       4: ireturn

Dividing an integer by zero throws an ArithmeticException, so there is no constant that 10 / 0 can be folded to, the way 10 / 1 folds to 10. You could rewrite the three instructions into code that constructs and throws the exception instead, but that leaves us with more instructions than we started with.

Therefore, if the denominator is a constant zero, we refuse to fold:

static Integer foldInt(Opcode opcode, int x, int y) {
    return switch (opcode) {
        case IADD -> x + y;
        case ISUB -> x - y;
        case IMUL -> x * y;
        case IAND -> x & y;
        case IOR  -> x | y;
        case IXOR -> x ^ y;
        case ISHL -> x << y;
        case ISHR -> x >> y;
        case IUSHR -> x >>> y;
        case IDIV -> y == 0 ? null : x / y; // Unsupported if y == 0
        case IREM -> y == 0 ? null : x % y; // Unsupported if y == 0
        default -> null; // Unsupported opcode for folding.
    };
}

It might be tempting to put a try/catch around the division instead, on the grounds that if the fold throws then the value can’t have mattered, but that deletes an exception the program is required to throw, and the method returns a value where it should have failed.

Folding floats and doubles

We’ve only considered integer constant folding so far, but folding float and double constants is very similar, except in the case of division: dividing by zero with ddiv or fdiv does not throw an ArithmeticException. A non-zero dividend gives an infinity, and 0.0 / 0.0 gives NaN. There is no exception to preserve, so the same fold needs no guard.

Try adding foldDouble and foldFloat alongside foldInt. The operands arrive from the same constantValue() call as Double and Float, and the replacement comes from pool.doubleEntry(value) or pool.floatEntry(value), with the same question about the smallest instruction that will do: dconst_0 and dconst_1 for doubles, fconst_0, fconst_1 and fconst_2 for floats. Anything else goes in the constant pool: a float is loaded with ldc or ldc_w as before, and a double always with ldc2_w.

Next steps

We’ve added two more peepholes to the optimiser: one that folds a pair of constants into their result, and one that deletes arithmetic that changes nothing. Both are small, and the code comes out of the optimiser with fewer instructions than it went in with. The division case is the one to remember: a fold has to preserve the exception, not just the value.

As an exercise, extend the folder to the comparison instructions, lcmp and fcmpl/fcmpg (and dcmpl/dcmpg), where the ordering of NaN is its own small trap: the two float comparisons exist because they disagree about it, fcmpg pushing 1 and fcmpl pushing -1 when either operand is NaN.

If you haven’t already read them, check out the related posts:

Peephole #1: Store, Reload, Rewrite - the first in this series, and why a match must stay inside one basic block
Build a Peephole Optimiser using the Java 24 Class-File API - the optimiser this series builds on, with the full code on GitHub