3 files changed, 46 insertions, 8 deletions

diff --git a/src/main/java/org/javacs/ParseFile.java b/src/main/java/org/javacs/ParseFile.java
index 55c7111..0f7780a 100644
--- a/src/main/java/org/javacs/ParseFile.java
+++ b/src/main/java/org/javacs/ParseFile.java
@@ -69,14 +69,38 @@ public class ParseFile {
     public List<TreePath> declarations() {
         var found = new ArrayList<TreePath>();
         class FindDeclarations extends TreePathScanner<Void, Void> {
+            boolean isClass(Tree t) {
+                if (!(t instanceof ClassTree)) return false;
+                var cls = (ClassTree) t;
+                return cls.getKind() == Tree.Kind.CLASS;
+            }
+
             @Override 
-            public Void visitClass​(ClassTree node, Void __) {
+            public Void visitClass​(ClassTree t, Void __) {
+                found.add(getCurrentPath());
+                return super.visitClass(t, null);
+            }
+
+            @Override
+            public Void visitMethod(MethodTree t, Void __) {
+                var path = getCurrentPath();
+                var parent = path.getParentPath().getLeaf();
+                if (isClass(parent)) {
+                    found.add(path);
+                }
+                // Skip code lenses for local classes
+                return null;
+            }
+
+            @Override
+            public Void visitVariable(VariableTree t, Void __) {
                 var path = getCurrentPath();
-                found.add(path);
-                for (var m : node.getMembers()) {
-                    found.add(new TreePath(path, m));
+                var parent = path.getParentPath().getLeaf();
+                if (isClass(parent)) {
+                    found.add(path);
                 }
-                return super.visitClass(node, null);
+                // Skip code lenses for local classes
+                return null;
             }
         }
         new FindDeclarations().scan(root, null);
@@ -88,7 +112,6 @@ public class ParseFile {
         return range(task, contents, path);
     }
 
-    // TODO maybe this should return TreePath?
     public Optional<CompletionContext> completionContext(int line, int character) {
         LOG.info(String.format("Finding completion position near %s(%d,%d)...", file, line, character));
         
diff --git a/src/test/java/org/javacs/CodeLensTest.java b/src/test/java/org/javacs/CodeLensTest.java
index 45244c7..9383de8 100644
--- a/src/test/java/org/javacs/CodeLensTest.java
+++ b/src/test/java/org/javacs/CodeLensTest.java
@@ -55,7 +55,9 @@ public class CodeLensTest {
     private List<String> titles(List<? extends CodeLens> lenses) {
         var titles = new ArrayList<String>();
         for (var lens : lenses) {
-            titles.add(lens.getCommand().getTitle());
+            var line = lens.getRange().getStart().getLine() + 1;
+            var title = lens.getCommand().getTitle();
+            titles.add(line + ":" + title);
         }
         return titles;
     }
@@ -77,6 +79,14 @@ public class CodeLensTest {
         assertThat(lenses, not(empty()));
 
         var titles = titles(lenses);
-        assertThat(titles, hasItem("2 references"));
+        assertThat(titles, hasItem("4:2 references"));
+        assertThat(titles, hasItem("6:2 references"));
+    }
+
+    @Test
+    public void enumConstants() {
+        var lenses = lenses("/org/javacs/example/DontShowEnumConstRefs.java");
+        var titles = titles(lenses);
+        assertThat(titles, not(hasItem("4:0 references")));
     }
 }
diff --git a/src/test/test-project/workspace/src/org/javacs/example/DontShowEnumConstRefs.java b/src/test/test-project/workspace/src/org/javacs/example/DontShowEnumConstRefs.java
new file mode 100644
index 0000000..75d3c8f
--- /dev/null
+++ b/src/test/test-project/workspace/src/org/javacs/example/DontShowEnumConstRefs.java
@@ -0,0 +1,5 @@
+package org.javacs.example;
+
+enum DontShowEnumConstRefs {
+    EnumConstWithNoReferences;
+}
\ No newline at end of file