[jvm] Include other import statements for test to harness approach (#735

)

This PR fixes the prompts for jvm test-to-harness approach to add in
other import statement in the prompts for reference. This help LLM to
include more necessary import statements to the generated harness.

---------

Signed-off-by: Arthur Chan <[email protected]>

llm_toolkit/prompt_builder.py

@@ -1334,7 +1334,8 @@ def _get_template(self, template_file: str) -> str:
     with open(template_file) as file:
       return file.read()
 
-  def _extract_jvm_imports(self, src: str, cls: list[str]) -> list[str]:
+  def _extract_jvm_imports(self, src: str,
+                           cls: list[str]) -> tuple[list[str], list[str]]:
     """Extract and interpret import statements from java source."""
 
     # Extract import statements
@@ -1358,19 +1359,27 @@ def _extract_jvm_imports(self, src: str, cls: list[str]) -> list[str]:
         cls_map[package].append(name)
 
     results = set()
-    for full_import, is_static, cls_name in imports:
+    others = set()
+    for full_import, _, cls_name in imports:
+      found = False
       if '*' in cls_name:
         # Resolve generic import
         package = cls_name.rstrip('.*')
         if package in cls_map:
           for name in cls_map[package]:
             if name in used_cls:
               results.add(f"import {package}.{name};")
+              found = True
+
+        if not found:
+          others.add(full_import)
       else:
-        # Keep all other import
-        results.add(full_import)
+        if cls_name in cls:
+          results.add(full_import)
+        else:
+          others.add(full_import)
 
-    return list(sorted(results))
+    return list(sorted(results)), list(sorted(others))
 
   def extract_header_files(self, text):
     # Include any weird macros defined that does not have any values. This
@@ -1426,9 +1435,12 @@ def build(self,
                                         harness_sample_text)
 
       # Extract import list
-      import_list = self._extract_jvm_imports(test_source_code, classes)
+      import_list, other_import_list = self._extract_jvm_imports(
+          test_source_code, classes)
       prompt_text = prompt_text.replace('{IMPORT_STATEMENTS}',
                                         '\n'.join(import_list))
+      prompt_text = prompt_text.replace('{OTHER_IMPORT_STATEMENTS}',
+                                        '\n'.join(other_import_list))
     else:
       included_header_files = self.extract_header_files(test_source_code)
       if included_header_files:

prompts/template_xml/jvm_requirement_test_to_harness.txt

@@ -48,10 +48,14 @@ Here is a comma-separated list of all publicly accessible classes in this projec
 {PUBLIC_CLASSES}
 </item>
 <item>
-Here is a list of import statements that you many needed for create the harness. If use use any classes that is not from java.lang package, you must import it by referencing these import statements below.
+Here is a list of import statements that you many needed for create the harness. If use use any classes that is not from java.lang package, you must import it by referencing these import statements below. You may need to add other import stataments from other classes that does not covered by these import statements here.
 {IMPORT_STATEMENTS}
 </item>
 <item>
+Here is a list of other import statements that you may need to make the generated harness compiles. Please only add them if necessary.
+{OTHER_IMPORT_STATEMENTS}
+</item>
+<item>
 You MUST ONLY use any of the following methods from the FuzzedDataProvider of the Jazzer framework for generating random data for fuzzing.
 If the needed return value is not found in the table, try use constructors or methods to create the needed random object. But you MUST try your best to randomise the random object with the methods in the table.