Move to struct declaration: handle struct declaration via _var_ case

Replace single struct declaration with short var declaration, or add to existing struct literal in statement

Author: wbars

src/com/goide/intentions/GoMoveToStructInitializationIntention.java

@@ -61,11 +61,29 @@ private static Data getData(@NotNull PsiElement element) {
     if (!element.isValid() || !element.isWritable()) return null;
     GoAssignmentStatement assignment = getValidAssignmentParent(element);
     GoReferenceExpression selectedFieldReference = assignment != null ? getFieldReferenceExpression(element, assignment) : null;
-    GoCompositeLit compositeLit = selectedFieldReference != null ? getStructLiteralByReference(selectedFieldReference, assignment) : null;
-    if (compositeLit == null) return null;
+    if (selectedFieldReference == null) return null;
 
-    List<GoReferenceExpression> references = getUninitializedSingleFieldReferences(assignment, selectedFieldReference, compositeLit);
-    return !references.isEmpty() ? new Data(assignment, compositeLit, references) : null;
+    GoVarDefinition structDefinition = getDefinition(selectedFieldReference);
+    GoStatement previousStatement = PsiTreeUtil.getPrevSiblingOfType(assignment, GoStatement.class);
+    boolean needReplaceDeclarationWithShortVar = isUnassigned(getSingleVarSpecByDefinition(previousStatement, structDefinition));
+
+    GoCompositeLit compositeLit = getStructLiteralByReference(selectedFieldReference, assignment);
+    if (compositeLit == null && !needReplaceDeclarationWithShortVar || structDefinition == null) return null;
+
+    List<GoReferenceExpression> references = getUninitializedSingleFieldReferences(assignment, structDefinition, compositeLit);
+    return !references.isEmpty() ? new Data(assignment, compositeLit, references, previousStatement, structDefinition) : null;
+  }
+
+  @Nullable
+  private static GoCompositeLit createStructLiteral(@NotNull GoVarDefinition definition, @NotNull Project project) {
+    GoType type = definition.getGoType(null);
+    return type != null ? GoElementFactory.createCompositeLit(project, type) : null;
+  }
+
+  @Nullable
+  @Contract("null -> null")
+  private static GoVarDefinition getDefinition(@Nullable GoReferenceExpression referenceExpressions) {
+    return ObjectUtils.tryCast(resolveQualifier(referenceExpressions), GoVarDefinition.class);
   }
 
   @Nullable
@@ -104,6 +122,22 @@ private static GoReferenceExpression unwrapParensAndCast(@Nullable PsiElement e)
     return ObjectUtils.tryCast(e, GoReferenceExpression.class);
   }
 
+  @Nullable
+  @Contract("_, null -> null; null, _ -> null")
+  private static GoVarSpec getSingleVarSpecByDefinition(@Nullable GoStatement statement,
+                                                        @Nullable GoVarDefinition definition) {
+    GoVarDeclaration declaration = statement != null ? statement.getVarDeclaration() : null;
+    List<GoVarSpec> varSpecs = declaration != null ? declaration.getVarSpecList() : emptyList();
+    GoVarSpec singleVarSpec = varSpecs.size() == 1 ? getFirstItem(varSpecs) : null;
+    List<GoVarDefinition> varDefinitions = singleVarSpec != null ? singleVarSpec.getVarDefinitionList() : emptyList();
+    return varDefinitions.size() == 1 && definition == getFirstItem(varDefinitions) ? singleVarSpec : null;
+  }
+
+  @Contract("null -> false")
+  private static boolean isUnassigned(@Nullable GoVarSpec varSpec) {
+    return varSpec != null && varSpec.getExpressionList().isEmpty();
+  }
+
   @Contract("null -> false")
   private static boolean isFieldReferenceExpression(@Nullable PsiElement element) {
     return element instanceof GoReferenceExpression && isFieldDefinition(((GoReferenceExpression)element).resolve());
@@ -138,11 +172,10 @@ private static boolean isResolvedTo(@Nullable PsiElement e, @Nullable PsiElement
 
   @NotNull
   private static List<GoReferenceExpression> getUninitializedSingleFieldReferences(@NotNull GoAssignmentStatement assignment,
-                                                                                   @NotNull GoReferenceExpression fieldReferenceExpression,
-                                                                                   @NotNull GoCompositeLit compositeLit) {
-    PsiElement resolve = resolveQualifier(fieldReferenceExpression);
+                                                                                   @Nullable GoVarDefinition definition,
+                                                                                   @Nullable GoCompositeLit compositeLit) {
     List<GoReferenceExpression> uninitializedFieldReferencesByQualifier =
-      filter(getUninitializedFieldReferenceExpressions(assignment, compositeLit), e -> isResolvedTo(e.getQualifier(), resolve));
+      filter(getUninitializedFieldReferenceExpressions(assignment, compositeLit), e -> isResolvedTo(e.getQualifier(), definition));
     MultiMap<PsiElement, GoReferenceExpression> resolved = groupBy(uninitializedFieldReferencesByQualifier, GoReferenceExpression::resolve);
     return map(filter(resolved.entrySet(), set -> set.getValue().size() == 1), set -> getFirstItem(set.getValue()));
   }
@@ -151,13 +184,15 @@ private static List<GoReferenceExpression> getUninitializedSingleFieldReferences
   private static GoCompositeLit getStructLiteralByReference(@NotNull GoReferenceExpression fieldReferenceExpression,
                                                             @NotNull GoAssignmentStatement assignment) {
     GoStatement previousStatement = PsiTreeUtil.getPrevSiblingOfType(assignment, GoStatement.class);
+    if (previousStatement == null) return null;
+
     if (previousStatement instanceof GoSimpleStatement) {
       return getStructLiteral(fieldReferenceExpression, (GoSimpleStatement)previousStatement);
     }
     if (previousStatement instanceof GoAssignmentStatement) {
       return getStructLiteral(fieldReferenceExpression, (GoAssignmentStatement)previousStatement);
     }
-    return null;
+    return getStructLiteral(previousStatement, fieldReferenceExpression);
   }
 
   @Nullable
@@ -172,15 +207,16 @@ private static GoCompositeLit getStructLiteral(@NotNull GoReferenceExpression fi
   }
 
   @Nullable
-  private static PsiElement resolveQualifier(@NotNull GoReferenceExpression fieldReferenceExpression) {
-    GoReferenceExpression qualifier = fieldReferenceExpression.getQualifier();
+  @Contract("null -> null")
+  private static PsiElement resolveQualifier(@Nullable GoReferenceExpression fieldReferenceExpression) {
+    GoReferenceExpression qualifier = fieldReferenceExpression != null ? fieldReferenceExpression.getQualifier() : null;
     return qualifier != null ? qualifier.resolve() : null;
   }
 
   @Nullable
   private static GoCompositeLit getStructLiteral(@NotNull GoReferenceExpression fieldReferenceExpression,
                                                  @NotNull GoAssignmentStatement structAssignment) {
-    GoVarDefinition varDefinition = ObjectUtils.tryCast(resolveQualifier(fieldReferenceExpression), GoVarDefinition.class);
+    GoVarDefinition varDefinition = getDefinition(fieldReferenceExpression);
     PsiElement field = fieldReferenceExpression.resolve();
     if (varDefinition == null || !isFieldDefinition(field) || !hasStructTypeWithField(varDefinition, (GoNamedElement)field)) {
       return null;
@@ -193,6 +229,14 @@ private static GoCompositeLit getStructLiteral(@NotNull GoReferenceExpression fi
     return ObjectUtils.tryCast(compositeLit, GoCompositeLit.class);
   }
 
+  @Nullable
+  private static GoCompositeLit getStructLiteral(@NotNull GoStatement statement,
+                                                 @NotNull GoReferenceExpression fieldReferenceExpression) {
+    GoVarDefinition definition = getDefinition(fieldReferenceExpression);
+    GoVarSpec varSpec = definition != null ? getSingleVarSpecByDefinition(statement, definition) : null;
+    return varSpec != null ? ObjectUtils.tryCast(getFirstItem(varSpec.getRightExpressionsList()), GoCompositeLit.class) : null;
+  }
+
   private static boolean hasStructTypeWithField(@NotNull GoVarDefinition structVarDefinition, @NotNull GoNamedElement field) {
     GoType type = structVarDefinition.getGoType(null);
     GoStructType structType = type != null ? ObjectUtils.tryCast(type.getUnderlyingType(), GoStructType.class) : null;
@@ -207,15 +251,16 @@ private static boolean isFieldInitialization(@NotNull GoElement element, @NotNul
 
   @NotNull
   private static List<GoReferenceExpression> getUninitializedFieldReferenceExpressions(@NotNull GoAssignmentStatement assignment,
-                                                                                       @NotNull GoCompositeLit structLiteral) {
+                                                                                       @Nullable GoCompositeLit structLiteral) {
     return filter(getFieldReferenceExpressions(assignment), expression ->
       isUninitializedFieldReferenceExpression(expression, structLiteral) && !isAssignedInPreviousStatement(expression, assignment));
   }
 
-  @Contract("null, _-> false")
+  @Contract("null, _-> false; !null, null -> true")
   private static boolean isUninitializedFieldReferenceExpression(@Nullable GoReferenceExpression fieldReferenceExpression,
-                                                                 @NotNull GoCompositeLit structLiteral) {
+                                                                 @Nullable GoCompositeLit structLiteral) {
     if (fieldReferenceExpression == null) return false;
+    if (structLiteral == null) return true;
     GoLiteralValue literalValue = structLiteral.getLiteralValue();
     PsiElement resolve = fieldReferenceExpression.resolve();
     return literalValue != null && isFieldDefinition(resolve) &&
@@ -238,19 +283,31 @@ private static List<? extends PsiElement> getLeftHandElements(@NotNull GoStateme
   public void invoke(@NotNull Project project, Editor editor, @NotNull PsiElement element) throws IncorrectOperationException {
     Data data = getData(element);
     if (data == null) return;
-    moveFieldReferenceExpressions(data);
+
+    boolean needReplaceDeclarationWithShortVar = data.getCompositeLit() == null;
+    GoCompositeLit compositeLit =
+      needReplaceDeclarationWithShortVar ? createStructLiteral(data.getStructDefinition(), project) : data.getCompositeLit();
+    if (compositeLit == null || needReplaceDeclarationWithShortVar && data.getStructDeclaration() == null) return;
+
+    moveFieldReferenceExpressions(data.getReferenceExpressions(), compositeLit, data.getAssignment());
+    if (!needReplaceDeclarationWithShortVar) return;
+    GoStatement shortVarStatement =
+      GoElementFactory.createShortVarDeclarationStatement(project, data.getStructDefinition().getText(), compositeLit.getText());
+    data.getStructDeclaration().replace(shortVarStatement);
   }
 
-  private static void moveFieldReferenceExpressions(@NotNull Data data) {
-    GoLiteralValue literalValue = data.getCompositeLit().getLiteralValue();
+  private static void moveFieldReferenceExpressions(@NotNull List<GoReferenceExpression> referenceExpressions,
+                                                    @NotNull GoCompositeLit compositeLit,
+                                                    @NotNull GoAssignmentStatement parentAssignment) {
+    GoLiteralValue literalValue = compositeLit.getLiteralValue();
     if (literalValue == null) return;
 
-    for (GoReferenceExpression expression : data.getReferenceExpressions()) {
+    for (GoReferenceExpression expression : referenceExpressions) {
       GoExpression anchor = getTopmostExpression(expression);
-      GoExpression fieldValue = GoPsiImplUtil.getRightExpression(data.getAssignment(), anchor);
+      GoExpression fieldValue = GoPsiImplUtil.getRightExpression(parentAssignment, anchor);
       if (fieldValue == null) continue;
 
-      GoPsiImplUtil.deleteExpressionFromAssignment(data.getAssignment(), anchor);
+      GoPsiImplUtil.deleteExpressionFromAssignment(parentAssignment, anchor);
       addFieldDefinition(literalValue, expression.getIdentifier().getText(), fieldValue.getText());
     }
   }
@@ -272,25 +329,46 @@ private static class Data {
     private final GoCompositeLit myCompositeLit;
     private final GoAssignmentStatement myAssignment;
     private final List<GoReferenceExpression> myReferenceExpressions;
+    private final GoStatement myStructDeclaration;
+    private final GoVarDefinition myStructDefinition;
 
     public Data(@NotNull GoAssignmentStatement assignment,
-                @NotNull GoCompositeLit compositeLit,
-                @NotNull List<GoReferenceExpression> referenceExpressions) {
+                @Nullable GoCompositeLit compositeLit,
+                @NotNull List<GoReferenceExpression> referenceExpressions,
+                @Nullable GoStatement structDeclaration,
+                @NotNull GoVarDefinition structDefinition) {
       myCompositeLit = compositeLit;
       myAssignment = assignment;
       myReferenceExpressions = referenceExpressions;
+      myStructDeclaration = structDeclaration;
+      myStructDefinition = structDefinition;
     }
 
+    @Nullable
     public GoCompositeLit getCompositeLit() {
       return myCompositeLit;
     }
 
+    @NotNull
     public GoAssignmentStatement getAssignment() {
       return myAssignment;
     }
 
+    @NotNull
     public List<GoReferenceExpression> getReferenceExpressions() {
       return myReferenceExpressions;
     }
+
+    @Nullable
+    public GoStatement getStructDeclaration() {
+      return myStructDeclaration;
+    }
+
+    @NotNull
+    public GoVarDefinition getStructDefinition() {
+      return myStructDefinition;
+    }
   }
 }
+
+

src/com/goide/psi/impl/GoElementFactory.java

@@ -266,4 +266,10 @@ public static GoTypeDeclaration createTypeDeclaration(@NotNull Project project,
     GoFile file = createFileFromText(project, "package a; type " + name + " " + type.getText());
     return PsiTreeUtil.findChildOfType(file, GoTypeDeclaration.class);
   }
+
+  @NotNull
+  public static GoCompositeLit createCompositeLit(@NotNull Project project, @NotNull GoType type) {
+    GoFile file = createFileFromText(project, "package a; var _ = " + type.getText() + "{}");
+    return PsiTreeUtil.findChildOfType(file, GoCompositeLit.class);
+  }
 }

testData/intentions/move-to-struct-initialization/notStructDeclaration.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s string
+	s.foo <caret>= "bar"
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/parensStructDeclaration-after.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	s := S{foo: "bar"}
+
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/parensStructDeclaration.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var (s S)
+	s.foo <caret>= "bar"
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/structDeclaration-after.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	s := S{foo: "bar"}
+
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/structDeclaration.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s S
+	s.foo <caret>= "bar"
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/structDeclarationMultipleExpressions.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s, b S
+	s.foo <caret>= "bar"
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/structDeclarationWithEmptyInit-after.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s S = S{foo: "bar"}
+
+	print(s.foo)
+}

testData/intentions/move-to-struct-initialization/structDeclarationWithEmptyInit.go

@@ -0,0 +1,11 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s S = S{}
+	s.foo <caret>= "bar"
+	print(s.foo)
+}