From 511275032e207bfcb2bca56ba1e9d3f0544a38c5 Mon Sep 17 00:00:00 2001
From: Arthur Chan <arthur.chan@adalogics.com>
Date: Wed, 10 Jul 2024 17:57:57 +0100
Subject: [PATCH] JVM: Add constructor informations in the prompt (#429)

This PR adds the constructors information of the target method with the
use of the new introspector query introduced in #426.
**This PR depends on #416 #425 #426 and needed rebase after these three
PRs are merged.**

---------

Signed-off-by: Arthur Chan <arthur.chan@adalogics.com>
---
 llm_toolkit/prompt_builder.py        | 48 ++++++++++++++++++++++++++--
 prompts/template_xml/jvm_problem.txt |  1 +
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/llm_toolkit/prompt_builder.py b/llm_toolkit/prompt_builder.py
index 0b17a923f..b0197f0a5 100644
--- a/llm_toolkit/prompt_builder.py
+++ b/llm_toolkit/prompt_builder.py
@@ -681,8 +681,8 @@ def _format_argument(self, count: int, arg_type: str) -> str:
     # java.lang.Object argument
     if 'java.lang.Object' in arg_type:
       base = self._get_template(self.object_arg_description_template_file)
-      prefix = 'Argument \#{count} requires an Object instance\n'
-      argument = '<argument>' + prefix + base + '</argument>'
+      prefix = f'Argument #{count} requires an Object instance\n'
+      argument = f'<argument>{prefix}{base}</argument>'
       return argument
 
     # Simple arguments
@@ -784,6 +784,49 @@ def _format_arguments(self) -> str:
 
     return '<arguments>' + '\n'.join(argument_descriptions) + '</arguments>'
 
+  def _format_constructors(self) -> str:
+    """Formats a list of functions / constructors to create the object for
+    invoking the target method."""
+    if self.benchmark.is_jvm_static:
+      return ''
+
+    constructors = []
+    ctrs = introspector.query_introspector_matching_function_constructor_type(
+        self.benchmark.project, self.benchmark.return_type, False)
+    for ctr in ctrs:
+      constructor_sig = ctr.get('function_signature')
+      if constructor_sig:
+        constructors.append(f'<signature>{constructor_sig}</signature>')
+
+    if constructors:
+      ctr_str = '\n'.join(constructors)
+      return f'<constructors>{ctr_str}</constructors>'
+
+    functions = []
+    funcs = introspector.query_introspector_matching_function_constructor_type(
+        self.benchmark.project, self.benchmark.return_type, True)
+    for func in funcs:
+      is_static = func.get('is_static', False)
+      function_sig = func.get('function_signature')
+      if not function_sig:
+        continue
+      if is_static:
+        functions.append(f'<item><signature>{function_sig}</signature></item>')
+      else:
+        function_class = function_sig[1:].split(']')[0]
+        function_str = f'<signature>{function_sig}</signature>'
+        function_str = function_str + (
+            '<prerequisite>You MUST create an '
+            f'{function_class} object before calling this constructing method.'
+            '</prerequisite>')
+        function_str = f'<item>{function_str}</item>'
+        functions.append(function_str)
+    if functions:
+      func_str = '\n'.join(functions)
+      return f'<constructors>{func_str}</constructors>'
+
+    return ''
+
   def _format_source_reference(self, signature: str) -> Tuple[str, str]:
     """Formats the source code reference for this target."""
     # Query for source code of the target method
@@ -807,6 +850,7 @@ def _format_problem(self, signature: str) -> str:
     problem = problem.replace('{REQUIREMENTS}',
                               self._format_requirement(signature))
     problem = problem.replace('{ARGUMENTS}', self._format_arguments())
+    problem = problem.replace('{CONSTRUCTORS}', self._format_constructors())
     problem = problem.replace('{EXCEPTIONS}', self._format_exceptions())
 
     self_source, cross_source = self._format_source_reference(signature)
diff --git a/prompts/template_xml/jvm_problem.txt b/prompts/template_xml/jvm_problem.txt
index 1d4248dab..9389f9825 100644
--- a/prompts/template_xml/jvm_problem.txt
+++ b/prompts/template_xml/jvm_problem.txt
@@ -2,6 +2,7 @@
 Your goal is to write a fuzzing harness for the provided method signature using Jazzer framework from Code Intellengence. It is important that the provided solution compiles and actually calls the function specified by the method signature:
 {TARGET}
 {ARGUMENTS}
+{CONSTRUCTORS}
 {EXCEPTIONS}
 Here is the source code of the target constructor for reference.
 <code>