diff --git a/src/main/java/com/alibaba/testable/generator/StaticNewClassGenerator.java b/src/main/java/com/alibaba/testable/generator/StaticNewClassGenerator.java index e6499b4..bfb62a1 100644 --- a/src/main/java/com/alibaba/testable/generator/StaticNewClassGenerator.java +++ b/src/main/java/com/alibaba/testable/generator/StaticNewClassGenerator.java @@ -5,7 +5,9 @@ import com.squareup.javapoet.*; import javax.lang.model.element.Modifier; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; /** * Generate global n.e class code @@ -18,11 +20,19 @@ public class StaticNewClassGenerator { return JavaFile.builder(ConstPool.SN_PKG, TypeSpec.classBuilder(ConstPool.SN_CLS) .addModifiers(Modifier.PUBLIC, Modifier.FINAL) + .addField(buildStaticPoolField()) .addMethod(buildStaticNewMethod()) .build()) .build().toString(); } + private FieldSpec buildStaticPoolField() { + return FieldSpec.builder(ParameterizedTypeName.get(Map.class, Class.class, Object.class), "pool", Modifier.PUBLIC) + .addModifiers(Modifier.STATIC) + .initializer("new $T<>()", HashMap.class) + .build(); + } + private MethodSpec buildStaticNewMethod() { TypeVariableName typeVariable = TypeVariableName.get("T"); MethodSpec.Builder builder = MethodSpec.methodBuilder(ConstPool.SN_METHOD) @@ -41,6 +51,15 @@ public class StaticNewClassGenerator { .beginControlFlow("for (Object o : args)") .addStatement("pts.add(o.getClass())") .endControlFlow() + .beginControlFlow("if (!pool.isEmpty())") + .beginControlFlow("try") + .addStatement("T obj = (T)pool.get(type)") + .beginControlFlow("if (obj != null)") + .addStatement("return obj") + .endControlFlow() + .nextControlFlow("catch (Exception e)") + .endControlFlow() + .endControlFlow() .beginControlFlow("try") .addStatement("return type.getConstructor(pts.toArray(new Class[0])).newInstance(args)") .nextControlFlow("catch (Exception e)")