1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package transformer;
18 
19 import annotations.ConstantMethodHandle;
20 import annotations.ConstantMethodType;
21 import java.io.InputStream;
22 import java.io.OutputStream;
23 import java.lang.invoke.MethodHandle;
24 import java.lang.invoke.MethodType;
25 import java.lang.reflect.Method;
26 import java.lang.reflect.Modifier;
27 import java.net.URL;
28 import java.net.URLClassLoader;
29 import java.nio.file.Files;
30 import java.nio.file.Path;
31 import java.nio.file.Paths;
32 import java.util.HashMap;
33 import java.util.Map;
34 import org.objectweb.asm.ClassReader;
35 import org.objectweb.asm.ClassVisitor;
36 import org.objectweb.asm.ClassWriter;
37 import org.objectweb.asm.Handle;
38 import org.objectweb.asm.MethodVisitor;
39 import org.objectweb.asm.Opcodes;
40 import org.objectweb.asm.Type;
41 
42 /**
43  * Class for transforming invoke static bytecodes into constant method handle loads and and constant
44  * method type loads.
45  *
46  * <p>When a parameterless private static method returning a MethodHandle is defined and annotated
47  * with {@code ConstantMethodHandle}, this transformer will replace static invocations of the method
48  * with a load constant bytecode with a method handle in the constant pool.
49  *
50  * <p>Suppose a method is annotated as: <code>
51  *  @ConstantMethodHandle(
52  *      kind = ConstantMethodHandle.STATIC_GET,
53  *      owner = "java/lang/Math",
54  *      fieldOrMethodName = "E",
55  *      descriptor = "D"
56  *  )
57  *  private static MethodHandle getMathE() {
58  *      unreachable();
59  *      return null;
60  *  }
61  * </code> Then invocations of {@code getMathE} will be replaced by a load from the constant pool
62  * with the constant method handle described in the {@code ConstantMethodHandle} annotation.
63  *
64  * <p>Similarly, a parameterless private static method returning a {@code MethodType} and annotated
65  * with {@code ConstantMethodType}, will have invocations replaced by a load constant bytecode with
66  * a method type in the constant pool.
67  */
68 class ConstantTransformer {
69     static class ConstantBuilder extends ClassVisitor {
70         private final Map<String, ConstantMethodHandle> constantMethodHandles;
71         private final Map<String, ConstantMethodType> constantMethodTypes;
72 
ConstantBuilder( int api, ClassVisitor cv, Map<String, ConstantMethodHandle> constantMethodHandles, Map<String, ConstantMethodType> constantMethodTypes)73         ConstantBuilder(
74                 int api,
75                 ClassVisitor cv,
76                 Map<String, ConstantMethodHandle> constantMethodHandles,
77                 Map<String, ConstantMethodType> constantMethodTypes) {
78             super(api, cv);
79             this.constantMethodHandles = constantMethodHandles;
80             this.constantMethodTypes = constantMethodTypes;
81         }
82 
83         @Override
visitMethod( int access, String name, String desc, String signature, String[] exceptions)84         public MethodVisitor visitMethod(
85                 int access, String name, String desc, String signature, String[] exceptions) {
86             MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);
87             return new MethodVisitor(this.api, mv) {
88                 @Override
89                 public void visitMethodInsn(
90                         int opcode, String owner, String name, String desc, boolean itf) {
91                     if (opcode == org.objectweb.asm.Opcodes.INVOKESTATIC) {
92                         ConstantMethodHandle constantMethodHandle = constantMethodHandles.get(name);
93                         if (constantMethodHandle != null) {
94                             insertConstantMethodHandle(constantMethodHandle);
95                             return;
96                         }
97                         ConstantMethodType constantMethodType = constantMethodTypes.get(name);
98                         if (constantMethodType != null) {
99                             insertConstantMethodType(constantMethodType);
100                             return;
101                         }
102                     }
103                     mv.visitMethodInsn(opcode, owner, name, desc, itf);
104                 }
105 
106                 private Type buildMethodType(Class<?> returnType, Class<?>[] parameterTypes) {
107                     Type rType = Type.getType(returnType);
108                     Type[] pTypes = new Type[parameterTypes.length];
109                     for (int i = 0; i < pTypes.length; ++i) {
110                         pTypes[i] = Type.getType(parameterTypes[i]);
111                     }
112                     return Type.getMethodType(rType, pTypes);
113                 }
114 
115                 private int getHandleTag(int kind) {
116                     switch (kind) {
117                         case ConstantMethodHandle.STATIC_PUT:
118                             return Opcodes.H_PUTSTATIC;
119                         case ConstantMethodHandle.STATIC_GET:
120                             return Opcodes.H_GETSTATIC;
121                         case ConstantMethodHandle.INSTANCE_PUT:
122                             return Opcodes.H_PUTFIELD;
123                         case ConstantMethodHandle.INSTANCE_GET:
124                             return Opcodes.H_GETFIELD;
125                         case ConstantMethodHandle.INVOKE_STATIC:
126                             return Opcodes.H_INVOKESTATIC;
127                         case ConstantMethodHandle.INVOKE_VIRTUAL:
128                             return Opcodes.H_INVOKEVIRTUAL;
129                         case ConstantMethodHandle.INVOKE_SPECIAL:
130                             return Opcodes.H_INVOKESPECIAL;
131                         case ConstantMethodHandle.NEW_INVOKE_SPECIAL:
132                             return Opcodes.H_NEWINVOKESPECIAL;
133                         case ConstantMethodHandle.INVOKE_INTERFACE:
134                             return Opcodes.H_INVOKEINTERFACE;
135                     }
136                     throw new Error("Unhandled kind " + kind);
137                 }
138 
139                 private void insertConstantMethodHandle(ConstantMethodHandle constantMethodHandle) {
140                     Handle handle =
141                             new Handle(
142                                     getHandleTag(constantMethodHandle.kind()),
143                                     constantMethodHandle.owner(),
144                                     constantMethodHandle.fieldOrMethodName(),
145                                     constantMethodHandle.descriptor(),
146                                     constantMethodHandle.ownerIsInterface());
147                     mv.visitLdcInsn(handle);
148                 }
149 
150                 private void insertConstantMethodType(ConstantMethodType constantMethodType) {
151                     Type methodType =
152                             buildMethodType(
153                                     constantMethodType.returnType(),
154                                     constantMethodType.parameterTypes());
155                     mv.visitLdcInsn(methodType);
156                 }
157             };
158         }
159     }
160 
161     private static void throwAnnotationError(
162             Method method, Class<?> annotationClass, String reason) {
163         StringBuilder sb = new StringBuilder();
164         sb.append("Error in annotation ")
165                 .append(annotationClass)
166                 .append(" on method ")
167                 .append(method)
168                 .append(": ")
169                 .append(reason);
170         throw new Error(sb.toString());
171     }
172 
173     private static void checkMethodToBeReplaced(
174             Method method, Class<?> annotationClass, Class<?> returnType) {
175         final int PRIVATE_STATIC = Modifier.STATIC | Modifier.PRIVATE;
176         if ((method.getModifiers() & PRIVATE_STATIC) != PRIVATE_STATIC) {
177             throwAnnotationError(method, annotationClass, " method is not private and static");
178         }
179         if (method.getTypeParameters().length != 0) {
180             throwAnnotationError(method, annotationClass, " method expects parameters");
181         }
182         if (!method.getReturnType().equals(returnType)) {
183             throwAnnotationError(method, annotationClass, " wrong return type");
184         }
185     }
186 
187     private static void transform(Path inputClassPath, Path outputClassPath) throws Throwable {
188         Path classLoadPath = inputClassPath.toAbsolutePath().getParent();
189         URLClassLoader classLoader =
190                 new URLClassLoader(new URL[] {classLoadPath.toUri().toURL()},
191                                    ClassLoader.getSystemClassLoader());
192         String inputClassName = inputClassPath.getFileName().toString().replace(".class", "");
193         Class<?> inputClass = classLoader.loadClass(inputClassName);
194 
195         final Map<String, ConstantMethodHandle> constantMethodHandles = new HashMap<>();
196         final Map<String, ConstantMethodType> constantMethodTypes = new HashMap<>();
197 
198         for (Method m : inputClass.getDeclaredMethods()) {
199             ConstantMethodHandle constantMethodHandle = m.getAnnotation(ConstantMethodHandle.class);
200             if (constantMethodHandle != null) {
201                 checkMethodToBeReplaced(m, ConstantMethodHandle.class, MethodHandle.class);
202                 constantMethodHandles.put(m.getName(), constantMethodHandle);
203                 continue;
204             }
205 
206             ConstantMethodType constantMethodType = m.getAnnotation(ConstantMethodType.class);
207             if (constantMethodType != null) {
208                 checkMethodToBeReplaced(m, ConstantMethodType.class, MethodType.class);
209                 constantMethodTypes.put(m.getName(), constantMethodType);
210                 continue;
211             }
212         }
213         ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
214         try (InputStream is = Files.newInputStream(inputClassPath)) {
215             ClassReader cr = new ClassReader(is);
216             ConstantBuilder cb =
217                     new ConstantBuilder(
218                             Opcodes.ASM6, cw, constantMethodHandles, constantMethodTypes);
219             cr.accept(cb, 0);
220         }
221         try (OutputStream os = Files.newOutputStream(outputClassPath)) {
222             os.write(cw.toByteArray());
223         }
224     }
225 
226     public static void main(String[] args) throws Throwable {
227         transform(Paths.get(args[0]), Paths.get(args[1]));
228     }
229 }
230