Add a method to replace a value with a constant directly.

Author: raphw

byte-buddy-dep/src/main/java/net/bytebuddy/asm/MemberSubstitution.java

@@ -295,6 +295,49 @@ public MemberSubstitution stub() {
             return replaceWith(Substitution.Stubbing.INSTANCE);
         }
 
+        /**
+         * Replaces any interaction with a matched byte code element with the provided compile-time constant.
+         *
+         * @param value The compile-time constant to set.
+         * @return A member substitution that replaces any interaction with the supplied compile-time constant.
+         */
+        public MemberSubstitution replaceWithConstant(Object value) {
+            if (value instanceof JavaConstant) {
+                return replaceWith(new Substitution.ForValue(new JavaConstantValue((JavaConstant) value), ((JavaConstant) value).getTypeDescription().asGenericType()));
+            } else if (value instanceof TypeDescription) {
+                return replaceWith(new Substitution.ForValue(ClassConstant.of((TypeDescription) value), TypeDefinition.Sort.describe(Class.class)));
+            } else {
+                Class<?> type = value.getClass();
+                if (type == String.class) {
+                    return replaceWith(new Substitution.ForValue(new TextConstant((String) value), TypeDefinition.Sort.describe(String.class)));
+                } else if (type == Class.class) {
+                    return replaceWith(new Substitution.ForValue(ClassConstant.of(TypeDescription.ForLoadedType.of((Class<?>) value)), TypeDefinition.Sort.describe(Class.class)));
+                } else if (type == Boolean.class) {
+                    return replaceWith(new Substitution.ForValue(IntegerConstant.forValue((Boolean) value), TypeDefinition.Sort.describe(boolean.class)));
+                } else if (type == Byte.class) {
+                    return replaceWith(new Substitution.ForValue(IntegerConstant.forValue((Byte) value), TypeDefinition.Sort.describe(byte.class)));
+                } else if (type == Short.class) {
+                    return replaceWith(new Substitution.ForValue(IntegerConstant.forValue((Short) value), TypeDefinition.Sort.describe(short.class)));
+                } else if (type == Character.class) {
+                    return replaceWith(new Substitution.ForValue(IntegerConstant.forValue((Character) value), TypeDefinition.Sort.describe(char.class)));
+                } else if (type == Integer.class) {
+                    return replaceWith(new Substitution.ForValue(IntegerConstant.forValue((Integer) value), TypeDefinition.Sort.describe(int.class)));
+                } else if (type == Long.class) {
+                    return replaceWith(new Substitution.ForValue(LongConstant.forValue((Long) value), TypeDefinition.Sort.describe(long.class)));
+                } else if (type == Float.class) {
+                    return replaceWith(new Substitution.ForValue(FloatConstant.forValue((Float) value), TypeDefinition.Sort.describe(float.class)));
+                } else if (type == Double.class) {
+                    return replaceWith(new Substitution.ForValue(DoubleConstant.forValue((Double) value), TypeDefinition.Sort.describe(double.class)));
+                } else if (JavaType.METHOD_HANDLE.getTypeStub().isAssignableFrom(type)) {
+                    return replaceWith(new Substitution.ForValue(new JavaConstantValue(JavaConstant.MethodHandle.ofLoaded(value)), TypeDefinition.Sort.describe(type)));
+                } else if (JavaType.METHOD_TYPE.getTypeStub().represents(type)) {
+                    return replaceWith(new Substitution.ForValue(new JavaConstantValue(JavaConstant.MethodType.ofLoaded(value)), TypeDefinition.Sort.describe(type)));
+                } else {
+                    throw new IllegalArgumentException("Not a constant value: " + value);
+                }
+            }
+        }
+
         /**
          * <p>
          * Replaces any interaction with a matched byte code element by an interaction with the specified field. If a field
@@ -869,6 +912,43 @@ public StackManipulation resolve(TypeDescription targetType,
             }
         }
 
+        class ForValue implements Substitution, Factory {
+
+            private final StackManipulation stackManipulation;
+
+            private final TypeDescription.Generic typeDescription;
+
+            public ForValue(StackManipulation stackManipulation, TypeDescription.Generic typeDescription) {
+                this.stackManipulation = stackManipulation;
+                this.typeDescription = typeDescription;
+            }
+
+            /**
+             * {@inheritDoc}
+             */
+            public Substitution make(TypeDescription instrumentedType, MethodDescription instrumentedMethod, TypePool typePool) {
+                return this;
+            }
+
+            /**
+             * {@inheritDoc}
+             */
+            public StackManipulation resolve(TypeDescription targetType,
+                                             ByteCodeElement target,
+                                             TypeList.Generic parameters,
+                                             TypeDescription.Generic result,
+                                             int freeOffset) {
+                List<StackManipulation> stackManipulations = new ArrayList<StackManipulation>(parameters.size());
+                for (int index = parameters.size() - 1; index >= 0; index--) {
+                    stackManipulations.add(Removal.of(parameters.get(index)));
+                }
+                if (!typeDescription.asErasure().isAssignableTo(result.asErasure())) {
+                    throw new IllegalStateException("Cannot assign " + typeDescription + " to " + result);
+                }
+                return new StackManipulation.Compound(CompoundList.of(stackManipulations, stackManipulation));
+            }
+        }
+
         /**
          * A substitution with a field access.
          */
@@ -1624,7 +1704,7 @@ public Resolution resolve(TypeDescription targetType,
                                               int freeOffset) {
                         return targetType.represents(void.class)
                                 ? this
-                                : new Simple(new StackManipulation.Compound(Removal.of(targetType), stackManipulation), resultType);
+                                : new Simple(new StackManipulation.Compound(Removal.of(current), stackManipulation), resultType);
                     }
 
                     /**
@@ -1832,9 +1912,9 @@ class ForArgumentLoading implements Step {
                     /**
                      * Creates an argument loading step.
                      *
-                     * @param index             The index of the argument to load.
-                     * @param assigner          The assigner to use for assigning the argument.
-                     * @param typing            The typing to use for the argument assignment.
+                     * @param index    The index of the argument to load.
+                     * @param assigner The assigner to use for assigning the argument.
+                     * @param typing   The typing to use for the argument assignment.
                      */
                     protected ForArgumentLoading(int index,
                                                  Assigner assigner,

byte-buddy-dep/src/test/java/net/bytebuddy/asm/MemberSubstitutionTypeTest.java

@@ -45,19 +45,35 @@ public static Collection<Object[]> data() {
     }
 
     @Test
-    public void testSubstitutionChainSimple() throws Exception {
+    public void testSubstitutionReplacement() throws Exception {
         Class<?> type = new ByteBuddy()
                 .redefine(this.type)
-                .visit(MemberSubstitution.strict().field(named(FOO)).replaceWithChain(MemberSubstitution.Substitution.Chain.Step.Simple.of(value)).on(named(RUN)))
+                .visit(MemberSubstitution.strict().field(named(FOO)).replaceWithConstant(replacement).on(named(RUN)))
                 .make()
                 .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
                 .getLoaded();
         Object instance = type.getDeclaredConstructor().newInstance();
         assertThat(type.getDeclaredField(FOO).get(instance), is(value));
         assertThat(type.getDeclaredField(BAR).get(instance), is(value));
         assertThat(type.getDeclaredMethod(RUN).invoke(instance), nullValue(Object.class));
-        assertThat(type.getDeclaredField(FOO).get(instance), is(replacement));
+        assertThat(type.getDeclaredField(FOO).get(instance), is(value));
+        assertThat(type.getDeclaredField(BAR).get(instance), is(replacement));
+    }
+
+    @Test
+    public void testSubstitutionChainSimple() throws Exception {
+        Class<?> type = new ByteBuddy()
+                .redefine(this.type)
+                .visit(MemberSubstitution.strict().field(named(FOO)).replaceWithChain(MemberSubstitution.Substitution.Chain.Step.Simple.of(replacement)).on(named(RUN)))
+                .make()
+                .load(ClassLoadingStrategy.BOOTSTRAP_LOADER, ClassLoadingStrategy.Default.WRAPPER)
+                .getLoaded();
+        Object instance = type.getDeclaredConstructor().newInstance();
+        assertThat(type.getDeclaredField(FOO).get(instance), is(value));
         assertThat(type.getDeclaredField(BAR).get(instance), is(value));
+        assertThat(type.getDeclaredMethod(RUN).invoke(instance), nullValue(Object.class));
+        assertThat(type.getDeclaredField(FOO).get(instance), is(value));
+        assertThat(type.getDeclaredField(BAR).get(instance), is(replacement));
     }
 
     @Test