Move to struct declaration: handle struct declaration via _var_ case

wbars · wbars · commit 1be6770e371a · 2016-11-02T17:41:08.000+10:00
Replace single struct declaration with short var declaration, or add to existing struct literal in statement
diff --git a/src/com/goide/intentions/GoMoveToStructInitializationIntention.java b/src/com/goide/intentions/GoMoveToStructInitializationIntention.java
@@ -61,11 +61,32 @@ 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 = needReplaceDeclarationWithShortVar
+                                  ? createStructLiteral(structDefinition, previousStatement.getProject())
+                                  : getStructLiteralByReference(selectedFieldReference, assignment);
+    if (compositeLit == null || structDefinition == null) return null;
+
+    List<GoReferenceExpression> references = getUninitializedSingleFieldReferences(assignment, structDefinition, compositeLit);
+    return !references.isEmpty() ? new Data(assignment, compositeLit, references, previousStatement, structDefinition.getText(),
+                                            needReplaceDeclarationWithShortVar) : 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 +125,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 +175,10 @@ private static boolean isResolvedTo(@Nullable PsiElement e, @Nullable PsiElement
 
   @NotNull
   private static List<GoReferenceExpression> getUninitializedSingleFieldReferences(@NotNull GoAssignmentStatement assignment,
-                                                                                   @NotNull GoReferenceExpression fieldReferenceExpression,
+                                                                                   @Nullable GoVarDefinition definition,
                                                                                    @NotNull GoCompositeLit compositeLit) {
-    PsiElement resolve = resolveQualifier(fieldReferenceExpression);
     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 +187,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 +210,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 +232,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;
@@ -239,6 +286,10 @@ public void invoke(@NotNull Project project, Editor editor, @NotNull PsiElement
     Data data = getData(element);
     if (data == null) return;
     moveFieldReferenceExpressions(data);
+
+    if (!data.needReplaceDeclarationWithShortVar()) return;
+    data.getStructDeclaration()
+      .replace(GoElementFactory.createShortVarDeclarationStatement(project, data.getStructVarName(), data.getCompositeLit().getText()));
   }
 
   private static void moveFieldReferenceExpressions(@NotNull Data data) {
@@ -272,13 +323,22 @@ private static class Data {
     private final GoCompositeLit myCompositeLit;
     private final GoAssignmentStatement myAssignment;
     private final List<GoReferenceExpression> myReferenceExpressions;
+    private final GoStatement myStructDeclaration;
+    private final String myStructVarName;
+    private final boolean myNeedReplaceDeclarationWithShortVar;
 
     public Data(@NotNull GoAssignmentStatement assignment,
-                @NotNull GoCompositeLit compositeLit,
-                @NotNull List<GoReferenceExpression> referenceExpressions) {
+                @Nullable GoCompositeLit compositeLit,
+                @NotNull List<GoReferenceExpression> referenceExpressions,
+                @Nullable GoStatement structDeclaration,
+                @Nullable String structVarName,
+                boolean needReplaceDeclarationWithShortVar) {
       myCompositeLit = compositeLit;
       myAssignment = assignment;
       myReferenceExpressions = referenceExpressions;
+      myStructDeclaration = structDeclaration;
+      myStructVarName = structVarName;
+      myNeedReplaceDeclarationWithShortVar = needReplaceDeclarationWithShortVar;
     }
 
     public GoCompositeLit getCompositeLit() {
@@ -292,5 +352,19 @@ public GoAssignmentStatement getAssignment() {
     public List<GoReferenceExpression> getReferenceExpressions() {
       return myReferenceExpressions;
     }
+
+    public GoStatement getStructDeclaration() {
+      return myStructDeclaration;
+    }
+
+    public String getStructVarName() {
+      return myStructVarName;
+    }
+
+    public boolean needReplaceDeclarationWithShortVar() {
+      return myNeedReplaceDeclarationWithShortVar;
+    }
   }
 }
+
+
diff --git a/src/com/goide/psi/impl/GoElementFactory.java b/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);
+  }
 }
diff --git a/testData/intentions/move-to-struct-initialization/notStructDeclaration.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/parensStructDeclaration-after.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/parensStructDeclaration.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclaration-after.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclaration.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationMultipleExpressions.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationWithEmptyInit-after.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationWithEmptyInit.go b/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)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationWithNonEmptyInit-after.go b/testData/intentions/move-to-struct-initialization/structDeclarationWithNonEmptyInit-after.go
@@ -0,0 +1,12 @@
+package main
+
+type S struct {
+	foo string
+	bar string
+}
+
+func main() {
+	var s S = S{bar: "a", foo: "bar"}
+
+	print(s.foo)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationWithNonEmptyInit.go b/testData/intentions/move-to-struct-initialization/structDeclarationWithNonEmptyInit.go
@@ -0,0 +1,12 @@
+package main
+
+type S struct {
+	foo string
+	bar string
+}
+
+func main() {
+	var s S = S{bar: "a"}
+	s.foo <caret>= "bar"
+	print(s.foo)
+}
diff --git a/testData/intentions/move-to-struct-initialization/structDeclarationWrongQualifier.go b/testData/intentions/move-to-struct-initialization/structDeclarationWrongQualifier.go
@@ -0,0 +1,12 @@
+package main
+
+type S struct {
+	foo string
+}
+
+func main() {
+	var s S
+	var b S
+	s.foo <caret>= "bar"
+	print(b.foo)
+}
diff --git a/tests/com/goide/quickfix/GoMoveToStructInitializationIntentionTest.java b/tests/com/goide/quickfix/GoMoveToStructInitializationIntentionTest.java
@@ -50,16 +50,23 @@ protected String getBasePath() {
   public void testMultipleFieldsPartlyAssigned()            { doTest(); }
   public void testWithParens()                              { doTest(); }
   public void testFieldExtractedFromParens()                { doTest(); }
+  public void testStructDeclaration()                       { doTest(); }
+  public void testParensStructDeclaration()                 { doTest(); }
+  public void testStructDeclarationWithEmptyInit()          { doTest(); }
+  public void testStructDeclarationWithNonEmptyInit()       { doTest(); }
 
   public void testDuplicateFields()                         { doTest(); }
   public void testMultiReturnFunction()                     { doTestNoFix(); }
   public void testWrongStruct()                             { doTestNoFix(); }
+  public void testNotStructDeclaration()                    { doTestNoFix(); }
   public void testExistingDeclaration()                     { doTestNoFix(); }
   public void testNotExistingField()                        { doTestNoFix(); }
   public void testJustAssignedVarWrongCaret()               { doTestNoFix(); }
   public void testJustAssignedVarWrongCaretWithParens()     { doTestNoFix(); }
   public void testJustInitializedVarWrongCaret()            { doTestNoFix(); }
   public void testJustAssignedVarBothParens()               { doTestNoFix(); }
   public void testJustAssignedFieldParens()                 { doTestNoFix(); }
+  public void testStructDeclarationMultipleExpressions()    { doTestNoFix(); }
+  public void testStructDeclarationWrongQualifier()         { doTestNoFix(); }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -266,4 +266,10 @@ public static GoTypeDeclaration createTypeDeclaration(@NotNull Project project,`
`266`	`266`	`GoFile file = createFileFromText(project, "package a; type " + name + " " + type.getText());`
`267`	`267`	`return PsiTreeUtil.findChildOfType(file, GoTypeDeclaration.class);`
`268`	`268`	`}`
	`269`	`+`
	`270`	`+ @NotNull`
	`271`	`+ public static GoCompositeLit createCompositeLit(@NotNull Project project, @NotNull GoType type) {`
	`272`	`+ GoFile file = createFileFromText(project, "package a; var _ = " + type.getText() + "{}");`
	`273`	`+ return PsiTreeUtil.findChildOfType(file, GoCompositeLit.class);`
	`274`	`+ }`
`269`	`275`	`}`