1616
1717package com .goide .inspections .unresolved ;
1818
19- import com .goide .psi .GoFieldDeclaration ;
20- import com .goide .psi .GoStructType ;
21- import com .goide .psi .impl .GoElementFactory ;
22- import com .intellij .codeInspection .LocalQuickFixOnPsiElement ;
19+ import com .goide .GoConstants ;
20+ import com .goide .psi .*;
21+ import com .goide .psi .impl .GoPsiImplUtil ;
22+ import com .goide .util .GoUtil ;
23+ import com .intellij .codeInsight .template .Template ;
24+ import com .intellij .codeInsight .template .TemplateManager ;
25+ import com .intellij .codeInsight .template .impl .ConstantNode ;
26+ import com .intellij .codeInspection .LocalQuickFixAndIntentionActionOnPsiElement ;
27+ import com .intellij .openapi .editor .Editor ;
2328import com .intellij .openapi .project .Project ;
2429import com .intellij .psi .PsiElement ;
2530import com .intellij .psi .PsiFile ;
31+ import com .intellij .psi .util .PsiTreeUtil ;
2632import com .intellij .util .ObjectUtils ;
2733import com .intellij .util .containers .ContainerUtil ;
2834import org .jetbrains .annotations .Nls ;
2935import org .jetbrains .annotations .NotNull ;
36+ import org .jetbrains .annotations .Nullable ;
3037
3138import java .util .List ;
3239
33- public class GoAddStructFieldFix extends LocalQuickFixOnPsiElement {
40+ public class GoAddStructFieldFix extends LocalQuickFixAndIntentionActionOnPsiElement {
3441 public static final String QUICK_FIX_NAME = "Add missing field" ;
35- private final String myFieldText ;
36- private final String myTypeText ;
3742
38- protected GoAddStructFieldFix (String fieldText , String typeText , @ NotNull GoStructType element ) {
43+ protected GoAddStructFieldFix (@ NotNull PsiElement element ) {
3944 super (element );
40- myFieldText = fieldText ;
41- myTypeText = typeText ;
4245 }
4346
4447 @ NotNull
@@ -48,12 +51,59 @@ public String getText() {
4851 }
4952
5053 @ Override
51- public void invoke (@ NotNull Project project , @ NotNull PsiFile file , @ NotNull PsiElement startElement , @ NotNull PsiElement endElement ) {
52- GoStructType structType = ObjectUtils .tryCast (startElement , GoStructType .class );
54+ public void invoke (@ NotNull Project project ,
55+ @ NotNull PsiFile file ,
56+ @ Nullable Editor editor ,
57+ @ NotNull PsiElement startElement ,
58+ @ NotNull PsiElement endElement ) {
59+ if (editor == null ) return ;
60+ GoReferenceExpression referenceExpression = ObjectUtils .tryCast (startElement , GoReferenceExpression .class );
61+ GoStructType structType = referenceExpression != null ? resolveStructType (referenceExpression ) : null ;
5362 if (structType == null ) return ;
63+
5464 List <GoFieldDeclaration > declarations = structType .getFieldDeclarationList ();
5565 PsiElement anchor = !declarations .isEmpty () ? ContainerUtil .getLastItem (declarations ) : structType .getLbrace ();
56- if (anchor != null ) structType .addAfter (GoElementFactory .createFieldDeclaration (project , myFieldText , myTypeText ), anchor );
66+ if (anchor == null ) return ;
67+
68+ startTemplate (project , editor , file , referenceExpression , anchor );
69+ }
70+
71+ private static void startTemplate (@ NotNull Project project ,
72+ @ NotNull Editor editor ,
73+ @ NotNull PsiFile file ,
74+ GoReferenceExpression referenceExpression ,
75+ PsiElement anchor ) {
76+ Template template = TemplateManager .getInstance (project ).createTemplate ("" , "" );
77+ template .addTextSegment (referenceExpression .getReference ().getCanonicalText () + " " );
78+ template .addVariable (new ConstantNode (getTypeName (referenceExpression , file )), true );
79+ template .addTextSegment ("\n " );
80+ editor .getCaretModel ().moveToOffset (anchor .getTextRange ().getEndOffset () + 1 );
81+ template .setToReformat (true );
82+ TemplateManager .getInstance (project ).startTemplate (editor , template );
83+ }
84+
85+
86+ private static String getTypeName (GoReferenceExpression referenceExpression , PsiFile file ) {
87+ GoAssignmentStatement assignment = PsiTreeUtil .getParentOfType (referenceExpression , GoAssignmentStatement .class );
88+ if (assignment == null ) return GoConstants .INTERFACE_TYPE ;
89+ GoExpression expression = GoPsiImplUtil .getRightExpression (assignment , referenceExpression );
90+ GoType type = expression != null ? expression .getGoType (null ) : null ;
91+
92+ if (type instanceof GoSpecType ) {
93+ GoSpecType spec = (GoSpecType )type ;
94+ GoFile typeFile = ObjectUtils .tryCast (spec .getContainingFile (), GoFile .class );
95+ if (typeFile != null && (file .isEquivalentTo (typeFile ) || GoUtil .inSamePackage (typeFile , file ))) {
96+ return spec .getIdentifier ().getText ();
97+ }
98+ }
99+ return type != null ? GoPsiImplUtil .getText (type ) : GoConstants .INTERFACE_TYPE ;
100+ }
101+
102+ @ Nullable
103+ private static GoStructType resolveStructType (@ NotNull GoReferenceExpression referenceExpression ) {
104+ GoReferenceExpression qualifier = referenceExpression .getQualifier ();
105+ GoSpecType type = qualifier != null ? ObjectUtils .tryCast (qualifier .getGoType (null ), GoSpecType .class ) : null ;
106+ return type != null ? ObjectUtils .tryCast (type .getType (), GoStructType .class ) : null ;
57107 }
58108
59109 @ Nls
0 commit comments