Merge pull request #42647 from nipunayf/fix-42591-master

[Master] Add the union of the types in the `Create variable with type` code action
ballerina-platform · Apr 26, 2024 · 1cd80a4 · 1cd80a4
2 parents 939d0cb + 0a934bf
commit 1cd80a4
Show file tree

Hide file tree

Showing 18 changed files with 768 additions and 15 deletions.
diff --git a/...erinalang/langserver/codeaction/providers/createvar/CreateVariableWithTypeCodeAction.java b/...erinalang/langserver/codeaction/providers/createvar/CreateVariableWithTypeCodeAction.java
@@ -17,7 +17,6 @@
 
 import io.ballerina.compiler.api.ModuleID;
 import io.ballerina.compiler.api.SemanticModel;
-import io.ballerina.compiler.api.TypeBuilder;
 import io.ballerina.compiler.api.Types;
 import io.ballerina.compiler.api.symbols.MethodSymbol;
 import io.ballerina.compiler.api.symbols.ModuleSymbol;
@@ -59,10 +58,12 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * Create variable code action when type infer diagnostic is presented.
@@ -178,6 +179,8 @@ private List<String> getPossibleTypes(TypeSymbol typeSymbol, CodeActionContext c
                                           ImportsAcceptor importsAcceptor) {
         typeSymbol = getRawType(typeSymbol, context);
         Set<String> possibleTypes = new HashSet<>();
+        Set<String> unionTypes = new HashSet<>();
+        String unionType = null;
         List<TypeSymbol> errorTypes = new ArrayList<>();
         if (typeSymbol.typeKind() == TypeDescKind.UNION) {
             ((UnionTypeSymbol) typeSymbol)
@@ -186,39 +189,47 @@ private List<String> getPossibleTypes(TypeSymbol typeSymbol, CodeActionContext c
                     .map(memberTypeSymbol -> getRawType(memberTypeSymbol, context))
                     .forEach(memberTypeSymbol -> {
                         if (memberTypeSymbol.typeKind() == TypeDescKind.UNION) {
+                            unionTypes.add(getTypeName(memberTypeSymbol, context, importsAcceptor));
                             possibleTypes.addAll(
                                     ((UnionTypeSymbol) memberTypeSymbol)
                                             .memberTypeDescriptors()
                                             .stream()
                                             .map(memberTSymbol -> getRawType(memberTSymbol, context))
                                             .map(symbol -> getTypeName(symbol, context, importsAcceptor))
-                                            .collect(Collectors.toList()));
+                                            .toList());
                         } else if (memberTypeSymbol.typeKind() == TypeDescKind.ERROR ||
                                 CommonUtil.getRawType(memberTypeSymbol).typeKind() == TypeDescKind.ERROR) {
                             errorTypes.add(memberTypeSymbol);
                         } else {
                             possibleTypes.add(getTypeName(memberTypeSymbol, context, importsAcceptor));
                         }
                     });
+            if (unionTypes.isEmpty()) {
+                unionType = getTypeName(typeSymbol, context, importsAcceptor);
+            }
         } else {
             String type = getTypeName(typeSymbol, context, importsAcceptor);
             if (!"any".equals(type)) {
                 return Collections.singletonList(type);
             }
         }
 
+        Set<String> typesSet = new LinkedHashSet<>(unionTypes);
+        typesSet.addAll(possibleTypes);
+        Stream<String> typeStream = typesSet.stream().filter(type -> !"any".equals(type));
+
         if (!errorTypes.isEmpty()) {
             String errorTypeStr = errorTypes.stream()
                     .map(type -> getTypeName(type, context, importsAcceptor))
                     .collect(Collectors.joining("|"));
-            return possibleTypes.stream()
-                    .filter(type -> !"any".equals(type))
-                    .map(type -> type + "|" + errorTypeStr)
-                    .collect(Collectors.toList());
+            typeStream = typeStream.map(type -> type + "|" + errorTypeStr);
         }
-        return possibleTypes.stream()
-                .filter(type -> !"any".equals(type))
-                .collect(Collectors.toList());
+
+        List<String> typesList = typeStream.collect(Collectors.toList());
+        if (unionType != null) {
+            typesList.add(0, unionType);
+        }
+        return typesList;
     }
 
     private boolean isInRemoteMethodCallOrResourceAccess(CodeActionContext context) {
@@ -272,14 +283,17 @@ private boolean isLangAnnotationModule(ModuleID moduleID) {
     private TypeSymbol getRawType(TypeSymbol typeSymbol, CodeActionContext context) {
         TypeSymbol rawType = CommonUtil.getRawType(typeSymbol);
         Types types = context.currentSemanticModel().get().types();
-        TypeBuilder builder = types.builder();
-        RecordTypeSymbol recordTypeSymbol = builder.RECORD_TYPE.withRestField(types.ANY).build();
-        if (rawType.subtypeOf(types.ERROR) || rawType.subtypeOf(recordTypeSymbol)) {
+        if (rawType.subtypeOf(types.ERROR) || subTypeOfRecord(types, rawType)) {
             return typeSymbol;
         }
         return rawType;
     }
 
+    private boolean subTypeOfRecord(Types types, TypeSymbol rawType) {
+        RecordTypeSymbol recordTypeSymbol = types.builder().RECORD_TYPE.withRestField(types.ANY).build();
+        return rawType.typeKind() != TypeDescKind.UNION && rawType.subtypeOf(recordTypeSymbol);
+    }
+
     /**
      * Finds and returns the action node in focus.
      *

diff --git a/...ore/src/test/java/org/ballerinalang/langserver/codeaction/CreateVariableWithTypeTest.java b/...ore/src/test/java/org/ballerinalang/langserver/codeaction/CreateVariableWithTypeTest.java
@@ -54,7 +54,15 @@ public Object[][] dataProvider() {
                 {"type_infer_for_resource_action_config.json"},
                 {"type_infer_for_resource_action_config2.json"},
                 {"type_infer_for_resource_action_with_check1.json"},
-                {"type_infer_for_resource_action_with_check2.json"}
+                {"type_infer_for_resource_action_with_check2.json"},
+                {"type_infer_for_action_with_union_config1.json"},
+                {"type_infer_for_action_with_union_config2.json"},
+                {"type_infer_for_action_with_union_config3.json"},
+                {"type_infer_for_action_with_union_config4.json"},
+                {"type_infer_for_action_with_union_config5.json"},
+                {"type_infer_for_action_with_union_config6.json"},
+                {"type_infer_for_action_with_union_config7.json"},
+                {"type_infer_for_action_with_union_config8.json"}
         };
     }
 

diff --git a/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config1.json b/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config1.json
@@ -0,0 +1,69 @@
+{
+  "position": {
+    "line": 50,
+    "character": 18
+  },
+  "source": "source2.bal",
+  "expected": [
+    {
+      "title": "Create variable with 'Rec1|Rec2'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 50,
+              "character": 4
+            },
+            "end": {
+              "line": 50,
+              "character": 4
+            }
+          },
+          "newText": "Rec1|Rec2 targetType = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'Rec1'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 50,
+              "character": 4
+            },
+            "end": {
+              "line": 50,
+              "character": 4
+            }
+          },
+          "newText": "Rec1 targetType = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'Rec2'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 50,
+              "character": 4
+            },
+            "end": {
+              "line": 50,
+              "character": 4
+            }
+          },
+          "newText": "Rec2 targetType = "
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}
diff --git a/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config2.json b/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config2.json
@@ -0,0 +1,69 @@
+{
+  "position": {
+    "line": 51,
+    "character": 18
+  },
+  "source": "source2.bal",
+  "expected": [
+    {
+      "title": "Create variable with 'Rec1|string|Error2'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 51,
+              "character": 4
+            },
+            "end": {
+              "line": 51,
+              "character": 4
+            }
+          },
+          "newText": "Rec1|string|Error2 unionResult = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'string|Error2'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 51,
+              "character": 4
+            },
+            "end": {
+              "line": 51,
+              "character": 4
+            }
+          },
+          "newText": "string|Error2 unionResult = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'Rec1|Error2'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 51,
+              "character": 4
+            },
+            "end": {
+              "line": 51,
+              "character": 4
+            }
+          },
+          "newText": "Rec1|Error2 unionResult = "
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}
diff --git a/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config3.json b/...rces/codeaction/create-var-with-type/config/type_infer_for_action_with_union_config3.json
@@ -0,0 +1,69 @@
+{
+  "position": {
+    "line": 52,
+    "character": 18
+  },
+  "source": "source2.bal",
+  "expected": [
+    {
+      "title": "Create variable with 'Rec2|Rec3|error'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 52,
+              "character": 4
+            },
+            "end": {
+              "line": 52,
+              "character": 4
+            }
+          },
+          "newText": "Rec2|Rec3|error unionResult = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'Rec2|error'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 52,
+              "character": 4
+            },
+            "end": {
+              "line": 52,
+              "character": 4
+            }
+          },
+          "newText": "Rec2|error unionResult = "
+        }
+      ],
+      "resolvable": false
+    },
+    {
+      "title": "Create variable with 'Rec3|error'",
+      "kind": "quickfix",
+      "edits": [
+        {
+          "range": {
+            "start": {
+              "line": 52,
+              "character": 4
+            },
+            "end": {
+              "line": 52,
+              "character": 4
+            }
+          },
+          "newText": "Rec3|error unionResult = "
+        }
+      ],
+      "resolvable": false
+    }
+  ]
+}