From b1af6dda3856403c5249d0791facfd1c530b172b Mon Sep 17 00:00:00 2001
From: Akash Yadav <itsaky01@gmail.com>
Date: Fri, 21 Jan 2022 14:00:43 +0530
Subject: [PATCH] Improved switch-case completion

---
 .../java/providers/CompletionProvider.java    | 26 ++++++++++++++-----
 .../lsp/java/visitors/FindCompletionsAt.java  | 13 +++++++---
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/lsp/java/src/main/java/com/itsaky/lsp/java/providers/CompletionProvider.java b/lsp/java/src/main/java/com/itsaky/lsp/java/providers/CompletionProvider.java
index 6b25eec12..2f28b3cfa 100644
--- a/lsp/java/src/main/java/com/itsaky/lsp/java/providers/CompletionProvider.java
+++ b/lsp/java/src/main/java/com/itsaky/lsp/java/providers/CompletionProvider.java
@@ -158,7 +158,7 @@ public class CompletionProvider implements ICompletionProvider {
                     case MEMBER_REFERENCE:
                         return completeMemberReference(task, path, partial);
                     case SWITCH:
-                        return completeSwitchConstant(task, path, partial);
+                        return completeSwitchConstant(task, path, partial, endsWithParen);
                     case IMPORT:
                         return completeImport(qualifiedPartialIdentifier(contents, (int) cursor));
                     default:
@@ -245,7 +245,7 @@ public class CompletionProvider implements ICompletionProvider {
         addKeywords(path, partial, list);
         return list;
     }
-
+    
     private void addKeywords(TreePath path, String partial, CompletionResult list) {
         Tree level = findKeywordLevel(path);
         String[] keywords = {};
@@ -543,22 +543,36 @@ public class CompletionProvider implements ICompletionProvider {
         Objects.requireNonNull (methods.get (name)).add(method);
     }
 
-    private CompletionResult completeSwitchConstant(@NonNull CompileTask task, @NonNull TreePath path, String partial) {
+    private CompletionResult completeSwitchConstant(@NonNull CompileTask task, @NonNull TreePath path, String partial, boolean endsWithParen) {
         SwitchTree switchTree = (SwitchTree) path.getLeaf();
         path = new TreePath(path, switchTree.getExpression());
         TypeMirror type = Trees.instance(task.task).getTypeMirror(path);
-        LOG.info("...complete constants of type " + type);
-        if (!(type instanceof DeclaredType)) {
-            return NOT_SUPPORTED;
+        
+        if (type.getKind ().isPrimitive () || !(type instanceof DeclaredType)) {
+            // primitive types do not have any members
+            return completeIdentifier (task, path, partial, endsWithParen);
         }
+        
         DeclaredType declared = (DeclaredType) type;
         TypeElement element = (TypeElement) declared.asElement();
+        
+        if (element.getKind () != ElementKind.ENUM) {
+            // If the switch's expression is not an enum type
+            // we will not get any constants to complete
+            // In this case, we fall back to completing identifiers
+            // At this point, we are sure that the case expression will definitely be an identifier tree
+            // see visitCase (CaseTree, Long) in FindCompletionsAt.java
+            return completeIdentifier (task, path, partial, endsWithParen);
+        }
+        
+        LOG.info("...complete constants of type " + type);
         List<CompletionItem> list = new ArrayList<> ();
         for (Element member : task.task.getElements().getAllMembers(element)) {
             if (member.getKind() != ElementKind.ENUM_CONSTANT) continue;
             if (!StringSearch.matchesPartialName(member.getSimpleName(), partial)) continue;
             list.add(item(task, member));
         }
+        
         return new CompletionResult(false, list);
     }
 
diff --git a/lsp/java/src/main/java/com/itsaky/lsp/java/visitors/FindCompletionsAt.java b/lsp/java/src/main/java/com/itsaky/lsp/java/visitors/FindCompletionsAt.java
index f58e3fb09..b53389de6 100644
--- a/lsp/java/src/main/java/com/itsaky/lsp/java/visitors/FindCompletionsAt.java
+++ b/lsp/java/src/main/java/com/itsaky/lsp/java/visitors/FindCompletionsAt.java
@@ -84,16 +84,23 @@ public class FindCompletionsAt extends TreePathScanner<TreePath, Long> {
     public TreePath visitCase (CaseTree t, Long find) {
         SourcePositions pos = Trees.instance (task).getSourcePositions ();
         
+        // check if the cursor is in the case expression
         // default statements have null expression
-//        if (t.getExpression () != null) {
-//
-//        }
+        // In case of an identifier tree, we have to check for both, variables and switch constants in CompletionProvider
+        if (t.getExpression () != null && !(t.getExpression () instanceof IdentifierTree)) {
+            long start = pos.getStartPosition (root, t.getExpression ());
+            long end = pos.getEndPosition (root, t.getExpression ());
+            if (start <= find && find <= end) {
+                return new TreePath (getCurrentPath (), t.getExpression ());
+            }
+        }
         
         long start = pos.getStartPosition (root, t) + "case".length ();
         long end = pos.getEndPosition (root, t.getExpression ());
         if (start <= find && find <= end) {
             return getCurrentPath ().getParentPath ();
         }
+        
         return super.visitCase (t, find);
     }