aboutsummaryrefslogtreecommitdiff
path: root/src/main/java/org/zwobble/couscous/util/asm/StackManipulationSwitch.java
blob: 2dbcd708a3c7b2cc074020e7bb3b84c432fa507d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
package org.zwobble.couscous.util.asm;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Ordering;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.jar.asm.Label;
import net.bytebuddy.jar.asm.MethodVisitor;

import java.util.List;

import static java.util.Arrays.asList;
import static org.zwobble.couscous.util.ExtraIterables.lazyMap;
import static org.zwobble.couscous.util.ExtraLists.list;

public class StackManipulationSwitch implements StackManipulation {
    public static Case switchCase(int key, StackManipulation stackManipulation) {
        return new Case(key, stackManipulation);
    }

    public static class Case {
        private final int key;
        private final StackManipulation stackManipulation;

        private Case(int key, StackManipulation stackManipulation) {
            this.key = key;
            this.stackManipulation = stackManipulation;
        }
    }

    private final StackManipulation defaultCase;
    private final List<Case> cases;

    public StackManipulationSwitch(StackManipulation defaultCase, List<Case> cases) {
        this.defaultCase = defaultCase;
        this.cases = Ordering.<Case>from((first, second) -> first.key - second.key).immutableSortedCopy(cases);
    }

    @Override
    public boolean isValid() {
        return true;
    }

    @Override
    public Size apply(MethodVisitor methodVisitor, Implementation.Context implementationContext) {
        Label defaultLabel = new Label();
        Label[] caseLabels = Iterables.toArray(lazyMap(cases, switchCase -> new Label()), Label.class);

        return new StackManipulation.Compound(ImmutableList.copyOf(Iterables.concat(
            list(
                new StackManipulationLookupTable(
                    defaultLabel,
                    cases.stream().mapToInt(switchCase -> switchCase.key).toArray(),
                    caseLabels
                )
            ),

            lazyMap(cases, asList(caseLabels), (switchCase, label) -> new StackManipulation.Compound(
                new StackManipulationLabel(label),
                StackManipulationFrame.SAME,
                switchCase.stackManipulation
            )),

            list(
                new StackManipulationLabel(defaultLabel),
                StackManipulationFrame.SAME,
                defaultCase
            )
        ))).apply(methodVisitor, implementationContext);
    }
}