1515 */
1616package org .openrewrite .java .migrate .util ;
1717
18+ import lombok .EqualsAndHashCode ;
19+ import lombok .Value ;
1820import org .jspecify .annotations .Nullable ;
1921import org .openrewrite .*;
2022import org .openrewrite .java .JavaTemplate ;
3133import java .util .List ;
3234import java .util .StringJoiner ;
3335
36+ @ EqualsAndHashCode (callSuper = false )
37+ @ Value
3438public class UseEnumSetOf extends Recipe {
3539 private static final MethodMatcher SET_OF = new MethodMatcher ("java.util.Set of(..)" , true );
3640 private static final String METHOD_TYPE = "java.util.EnumSet" ;
3741
42+ @ Option (
43+ displayName = "Convert empty `Set.of()` to `EnumSet.noneOf()`" ,
44+ description = "When true, converts `Set.of()` with no arguments to `EnumSet.noneOf()`. Default true." ,
45+ example = "true" ,
46+ required = false
47+ )
48+ @ Nullable
49+ Boolean convertEmptySet ;
50+
3851 @ Override
3952 public String getDisplayName () {
4053 return "Prefer `EnumSet of(..)`" ;
@@ -52,73 +65,78 @@ public Duration getEstimatedEffortPerOccurrence() {
5265
5366 @ Override
5467 public TreeVisitor <?, ExecutionContext > getVisitor () {
55- return Preconditions .check (Preconditions .and (new UsesJavaVersion <>(9 ),
56- new UsesMethod <>(SET_OF )), new UseEnumSetOfVisitor ());
57- }
68+ return Preconditions .check (Preconditions .and (new UsesJavaVersion <>(9 ), new UsesMethod <>(SET_OF )), new JavaVisitor <ExecutionContext >() {
69+ @ Override
70+ public J .MethodInvocation visitMethodInvocation (J .MethodInvocation methodInvocation , ExecutionContext ctx ) {
71+ J .MethodInvocation mi = (J .MethodInvocation ) super .visitMethodInvocation (methodInvocation , ctx );
5872
59- private static class UseEnumSetOfVisitor extends JavaVisitor <ExecutionContext > {
60- @ Override
61- public J .MethodInvocation visitMethodInvocation (J .MethodInvocation methodInvocation , ExecutionContext ctx ) {
62- J .MethodInvocation mi = (J .MethodInvocation ) super .visitMethodInvocation (methodInvocation , ctx );
63-
64- if (SET_OF .matches (mi ) &&
65- mi .getType () instanceof JavaType .Parameterized &&
66- !TypeUtils .isOfClassType (mi .getType (), METHOD_TYPE )) {
67- Cursor parent = getCursor ().dropParentUntil (is -> is instanceof J .Assignment || is instanceof J .VariableDeclarations || is instanceof J .Block );
68- if (!(parent .getValue () instanceof J .Block )) {
69- JavaType type = parent .getValue () instanceof J .Assignment ?
70- ((J .Assignment ) parent .getValue ()).getType () : ((J .VariableDeclarations ) parent .getValue ()).getVariables ().get (0 ).getType ();
71- if (isAssignmentSetOfEnum (type )) {
72- maybeAddImport (METHOD_TYPE );
73-
74- List <Expression > args = mi .getArguments ();
75- if (isArrayParameter (args )) {
76- return mi ;
77- }
73+ if (SET_OF .matches (mi ) &&
74+ mi .getType () instanceof JavaType .Parameterized &&
75+ !TypeUtils .isOfClassType (mi .getType (), METHOD_TYPE ) &&
76+ convertEmptySet (mi )) {
77+ Cursor parent = getCursor ().dropParentUntil (is -> is instanceof J .Assignment || is instanceof J .VariableDeclarations || is instanceof J .Block );
78+ if (!(parent .getValue () instanceof J .Block )) {
79+ JavaType type = parent .getValue () instanceof J .Assignment ?
80+ ((J .Assignment ) parent .getValue ()).getType () : ((J .VariableDeclarations ) parent .getValue ()).getVariables ().get (0 ).getType ();
81+ if (isAssignmentSetOfEnum (type )) {
82+ maybeAddImport (METHOD_TYPE );
7883
79- if (args .get (0 ) instanceof J .Empty ) {
80- JavaType firstTypeParameter = ((JavaType .Parameterized ) type ).getTypeParameters ().get (0 );
81- JavaType .ShallowClass shallowClass = JavaType .ShallowClass .build (firstTypeParameter .toString ());
82- return JavaTemplate .builder ("EnumSet.noneOf(" + shallowClass .getClassName () + ".class)" )
84+ List <Expression > args = mi .getArguments ();
85+ if (isArrayParameter (args )) {
86+ return mi ;
87+ }
88+
89+ if (args .get (0 ) instanceof J .Empty ) {
90+ JavaType firstTypeParameter = ((JavaType .Parameterized ) type ).getTypeParameters ().get (0 );
91+ JavaType .ShallowClass shallowClass = JavaType .ShallowClass .build (firstTypeParameter .toString ());
92+ return JavaTemplate .builder ("EnumSet.noneOf(" + shallowClass .getClassName () + ".class)" )
93+ .contextSensitive ()
94+ .imports (METHOD_TYPE )
95+ .build ()
96+ .apply (updateCursor (mi ), mi .getCoordinates ().replace ());
97+ }
98+
99+ StringJoiner setOf = new StringJoiner (", " , "EnumSet.of(" , ")" );
100+ args .forEach (o -> setOf .add ("#{any()}" ));
101+ return JavaTemplate .builder (setOf .toString ())
83102 .contextSensitive ()
84103 .imports (METHOD_TYPE )
85104 .build ()
86- .apply (updateCursor (mi ), mi .getCoordinates ().replace ());
105+ .apply (updateCursor (mi ), mi .getCoordinates ().replace (), args . toArray () );
87106 }
88-
89- StringJoiner setOf = new StringJoiner (", " , "EnumSet.of(" , ")" );
90- args .forEach (o -> setOf .add ("#{any()}" ));
91- return JavaTemplate .builder (setOf .toString ())
92- .contextSensitive ()
93- .imports (METHOD_TYPE )
94- .build ()
95- .apply (updateCursor (mi ), mi .getCoordinates ().replace (), args .toArray ());
96107 }
97108 }
109+ return mi ;
98110 }
99- return mi ;
100- }
101-
102- private boolean isAssignmentSetOfEnum (@ Nullable JavaType type ) {
103- if (type instanceof JavaType .Parameterized ) {
104- JavaType .Parameterized parameterized = (JavaType .Parameterized ) type ;
105- if (TypeUtils .isOfClassType (parameterized .getType (), "java.util.Set" )) {
106- return ((JavaType .Parameterized ) type ).getTypeParameters ().stream ()
107- .filter (org .openrewrite .java .tree .JavaType .Class .class ::isInstance )
108- .map (org .openrewrite .java .tree .JavaType .Class .class ::cast )
109- .anyMatch (o -> o .getKind () == JavaType .FullyQualified .Kind .Enum );
111+
112+ private boolean convertEmptySet (J .MethodInvocation mi ) {
113+ if (convertEmptySet == null || convertEmptySet ) {
114+ return true ;
110115 }
116+ return !mi .getArguments ().isEmpty () && !(mi .getArguments ().get (0 ) instanceof J .Empty );
111117 }
112- return false ;
113- }
114118
115- private boolean isArrayParameter (final List <Expression > args ) {
116- if (args .size () != 1 ) {
119+ private boolean isAssignmentSetOfEnum (@ Nullable JavaType type ) {
120+ if (type instanceof JavaType .Parameterized ) {
121+ JavaType .Parameterized parameterized = (JavaType .Parameterized ) type ;
122+ if (TypeUtils .isOfClassType (parameterized .getType (), "java.util.Set" )) {
123+ return ((JavaType .Parameterized ) type ).getTypeParameters ().stream ()
124+ .filter (JavaType .Class .class ::isInstance )
125+ .map (JavaType .Class .class ::cast )
126+ .anyMatch (o -> o .getKind () == JavaType .FullyQualified .Kind .Enum );
127+ }
128+ }
117129 return false ;
118130 }
119- JavaType type = args .get (0 ).getType ();
120- return TypeUtils .asArray (type ) != null ;
121- }
131+
132+ private boolean isArrayParameter (final List <Expression > args ) {
133+ if (args .size () != 1 ) {
134+ return false ;
135+ }
136+ JavaType type = args .get (0 ).getType ();
137+ return TypeUtils .asArray (type ) != null ;
138+ }
139+ });
122140 }
123141
124142}
0 commit comments