From f5bb1106a88977a1a053042511d79906f80c56f0 Mon Sep 17 00:00:00 2001
From: Shiwani Gupta <jshiwanigupta@gmail.com>
Date: Fri, 25 Oct 2024 02:15:50 +0530
Subject: [PATCH] Add super class in reference context

---
 .../ai/completion/JeddictCompletionProvider.java    |  4 ++++
 .../jeddict/ai/scanner/ProjectClassListener.java    | 13 +++++++------
 .../jeddict/ai/scanner/ProjectClassScanner.java     |  9 ++++++---
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/main/java/io/github/jeddict/ai/completion/JeddictCompletionProvider.java b/src/main/java/io/github/jeddict/ai/completion/JeddictCompletionProvider.java
index 84826d3..1f751f4 100644
--- a/src/main/java/io/github/jeddict/ai/completion/JeddictCompletionProvider.java
+++ b/src/main/java/io/github/jeddict/ai/completion/JeddictCompletionProvider.java
@@ -244,6 +244,10 @@ private Set<String> findReferencedClasses(CompilationUnitTree compilationUnit) {
             for (Tree tree : compilationUnit.getTypeDecls()) {
                 if (tree instanceof ClassTree) {
                     ClassTree classTree = (ClassTree) tree;
+                    Tree superclass = classTree.getExtendsClause();
+                    if (superclass != null) {
+                        referencedClasses.add(superclass.toString());
+                    }
                     for (Tree member : classTree.getMembers()) {
                         if (member instanceof VariableTree) {
                             VariableTree variable = (VariableTree) member;
diff --git a/src/main/java/io/github/jeddict/ai/scanner/ProjectClassListener.java b/src/main/java/io/github/jeddict/ai/scanner/ProjectClassListener.java
index 648c69d..bec64af 100644
--- a/src/main/java/io/github/jeddict/ai/scanner/ProjectClassListener.java
+++ b/src/main/java/io/github/jeddict/ai/scanner/ProjectClassListener.java
@@ -21,18 +21,19 @@
  */
 public class ProjectClassListener {
 
-    private Project project;
+    private final Project project;
 
-    private Map<FileObject, ClassData> classDatas;
-    private Set<FileObject> pending = new HashSet<>();
+    private final Map<FileObject, ClassData> classDatas;
+    private final Set<DataObject> pendingDO = new HashSet<>();
 
     public ProjectClassListener(Project project, Map<FileObject, ClassData> classDatas) {
         this.project = project;
         this.classDatas = classDatas;
     }
 
-    public Set<FileObject> getPending() {
-        return pending;
+
+    public Set<DataObject> getPendingDataObject() {
+        return pendingDO;
     }
 
     public void register() {
@@ -62,7 +63,7 @@ public void register() {
             DataObject[] modifiedObjects = registry.getModified();
             for (DataObject dataObj : modifiedObjects) {
                 if (FileUtil.isParentOf(javaFolder2, dataObj.getPrimaryFile())) {
-                    pending.add(dataObj.getPrimaryFile());
+                    pendingDO.add(dataObj);
                     classDatas.remove(dataObj.getPrimaryFile());
                     System.out.println("Modified DataObject: " + dataObj.getName());
                 } else {
diff --git a/src/main/java/io/github/jeddict/ai/scanner/ProjectClassScanner.java b/src/main/java/io/github/jeddict/ai/scanner/ProjectClassScanner.java
index 53101f8..e468738 100644
--- a/src/main/java/io/github/jeddict/ai/scanner/ProjectClassScanner.java
+++ b/src/main/java/io/github/jeddict/ai/scanner/ProjectClassScanner.java
@@ -82,6 +82,9 @@ private static void scanFolder(FileObject folder, Map<FileObject, ClassData> cla
         }
     }
 
+    public static void scanJavaFile(DataObject javaFile, Map<FileObject, ClassData> classList) throws IOException {
+        scanJavaFile(javaFile.getPrimaryFile(), classList) ;
+    }
     public static void scanJavaFile(FileObject javaFile, Map<FileObject, ClassData> classList) throws IOException {
         JavaSource javaSource = JavaSource.forFileObject(javaFile);
 
@@ -257,10 +260,10 @@ public static List<ClassData> getClassData(FileObject fileObject, Set<String> fi
                     projectClassListeners.put(key, projectClassListener);
                 }
                 if (projectClassListeners.get(key) != null) {
-                    Iterator<FileObject> iterator = projectClassListeners.get(key).getPending().iterator();
+                    Iterator<DataObject> iterator = projectClassListeners.get(key).getPendingDataObject().iterator();
                     while (iterator.hasNext()) {
-                        FileObject javaFile = iterator.next();
-                        if (javaFile.equals(fileObject)) {
+                        DataObject javaFile = iterator.next();
+                        if (javaFile.getPrimaryFile().equals(fileObject)) {
                             // Ignore current editor
                             System.out.println("Ignoring " + fileObject.getName());
                         } else {