Skip to content

Commit

Permalink
Add a method to replace a value with a constant directly.
Browse files Browse the repository at this point in the history
  • Loading branch information
raphw committed Feb 4, 2023
1 parent cb8a9be commit db80456
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
*/
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit db80456

Please sign in to comment.