Fix undefined reference (#465)

1. Add the header file of the function-under-test in the code generation
prompt.
2. Add the header file or source code file of the missing function in
the code fixing prompt (C projects mostly).
3. Related instructions.

Author: DonggeLiu

data_prep/project_context/context_introspector.py

@@ -4,7 +4,7 @@
 import logging
 import os
 from difflib import SequenceMatcher
-from typing import Any
+from typing import Any, Optional
 
 from data_prep import introspector
 from experiment import benchmark as benchmarklib
@@ -165,12 +165,14 @@ def get_context_info(self) -> dict:
     func_source = self._get_function_implementation()
     files = self._get_files_to_include()
     decl = self._get_embeddable_declaration()
+    header = self.get_prefixed_header_file()
 
     context_info = {
         'xrefs': xrefs,
         'func_source': func_source,
         'files': files,
-        'decl': decl
+        'decl': decl,
+        'header': header,
     }
 
     logging.debug('Context: %s', context_info)
@@ -229,7 +231,7 @@ def get_same_header_file_paths(self, wrong_file: str) -> list[str]:
     for header in header_list:
       correct_file_name = os.path.splitext(os.path.basename(header))
       if wrong_file_name == correct_file_name:
-        candidate_headers.append(header)
+        candidate_headers.append(os.path.normpath(header))
 
     return candidate_headers[:5]
 
@@ -245,11 +247,23 @@ def get_similar_header_file_paths(self, wrong_file: str) -> list[str]:
     candidate_headers = sorted(candidate_header_scores,
                                key=lambda x: candidate_header_scores[x],
                                reverse=True)
-    return candidate_headers[:5]
+    return [os.path.normpath(header) for header in candidate_headers[:5]]
+
+  def _get_header_files_to_include(self, func_sig: str) -> Optional[str]:
+    """Retrieves the header file of the function signature."""
+    header_file = introspector.query_introspector_header_files_to_include(
+        self._benchmark.project, func_sig)
+    return header_file[0] if header_file else None
 
-  def get_target_function_file_path(self) -> str:
+  def _get_target_function_file_path(self) -> str:
     """Retrieves the header/source file of the function under test."""
-    # Step 1: Find a header file that shares the same name as the source file.
+    # Step 1: Find a header file from the default API.
+    header_file = self._get_header_files_to_include(
+        self._benchmark.function_signature)
+    if header_file:
+      return header_file
+
+    # Step 2: Find a header file that shares the same name as the source file.
     # TODO: Make this more robust, e.g., when header file and base file do not
     # share the same basename.
     source_file = introspector.query_introspector_source_file_path(
@@ -264,5 +278,41 @@ def get_target_function_file_path(self) -> str:
     if candidate_headers:
       return candidate_headers[0]
 
-    # Step 2: Use the source file If it does not have a same-name-header.
+    # Step 3: Use the source file If it does not have a same-name-header.
     return source_file
+
+  def get_prefixed_header_file(self, func_sig: str = '') -> Optional[str]:
+    """Retrieves the header_file with `extern "C"` if needed."""
+    if func_sig:
+      header_file = self._get_header_files_to_include(func_sig)
+    else:
+      header_file = self._get_target_function_file_path()
+
+    if not header_file:
+      return None
+    include_statement = f'#include "{os.path.normpath(header_file)}"'
+    return (f'extern "C" {{\n{include_statement}\n}}'
+            if self._benchmark.needs_extern else include_statement)
+
+  def get_prefixed_header_file_by_name(self, func_name: str) -> Optional[str]:
+    """Retrieves the header file based on function name with `extern "C"` if
+    needed."""
+    func_sig = introspector.query_introspector_function_signature(
+        self._benchmark.project, func_name)
+    return self.get_prefixed_header_file(func_sig)
+
+  def get_prefixed_source_file(self,
+                               function_signature: str = '') -> Optional[str]:
+    """Retrieves the source file with `extern "C"` if needed."""
+    if function_signature:
+      source_file = introspector.query_introspector_source_file_path(
+          self._benchmark.project, function_signature)
+    else:
+      source_file = introspector.query_introspector_source_file_path(
+          self._benchmark.project, self._benchmark.function_signature)
+    if not source_file:
+      return None
+
+    include_statement = f'#include "{os.path.normpath(source_file)}"'
+    return (f'extern "C" {{\n{include_statement}\n}}'
+            if self._benchmark.needs_extern else include_statement)

experiment/benchmark.py

@@ -184,11 +184,30 @@ def file_type(self) -> FileType:
     """Returns the file type of the benchmark."""
     return get_file_type(self.target_path)
 
+  @property
+  def is_c_target(self) -> bool:
+    """Validates if the project is written in C."""
+    return self.file_type.value.lower() == 'c'
+
+  @property
+  def is_cpp_target(self) -> bool:
+    """Validates if the project is written in C++."""
+    return self.file_type.value.lower() == 'c++'
+
+  @property
+  def is_c_projcet(self) -> bool:
+    """Validates if the project is written in C."""
+    return self.language.lower() == 'c'
+
+  @property
+  def is_cpp_projcet(self) -> bool:
+    """Validates if the project is written in C++."""
+    return self.language.lower() == 'c++'
+
   @property
   def needs_extern(self) -> bool:
     """Checks if it is C++ fuzz target for a C project, which needs `extern`."""
-    return (self.file_type.value.lower() == 'c++' and
-            self.language.lower() == 'c')
+    return self.is_cpp_target and self.is_c_projcet
 
 
 def get_file_type(file_path: str) -> FileType:

llm_toolkit/code_fixer.py

@@ -30,6 +30,7 @@
 ERROR_LINES = 20
 NO_MEMBER_ERROR_REGEX = r"error: no member named '.*' in '([^':]*):?.*'"
 FILE_NOT_FOUND_ERROR_REGEX = r"fatal error: '([^']*)' file not found"
+UNDEFINED_REF_ERROR_REGEX = r"undefined reference to `([^']*)'"
 UNKNOWN_TYPE_ERROR = 'error: unknown type name'
 
 # The following strings identify errors when a C fuzz target attempts to use
@@ -461,6 +462,8 @@ def _collect_instructions(benchmark: benchmarklib.Benchmark, errors: list[str],
   for error in errors:
     instruction += _collect_instruction_file_not_found(benchmark, error,
                                                        fuzz_target_source_code)
+    instruction += _collect_instruction_undefined_reference(
+        benchmark, error, fuzz_target_source_code)
   instruction += _collect_instruction_fdp_in_c_target(benchmark, errors,
                                                       fuzz_target_source_code)
   instruction += _collect_instruction_no_goto(fuzz_target_source_code)
@@ -470,6 +473,62 @@ def _collect_instructions(benchmark: benchmarklib.Benchmark, errors: list[str],
   return instruction
 
 
+def _collect_instruction_undefined_reference(
+    benchmark: benchmarklib.Benchmark, error: str,
+    fuzz_target_source_code: str) -> str:
+  """Collects the instructions to fix the 'undefined reference' errors."""
+  matched_funcs = re.findall(UNDEFINED_REF_ERROR_REGEX, error)
+  if not matched_funcs:
+    return ''
+  instruction = ''
+  for undefined_func in matched_funcs:
+    if undefined_func == 'LLVMFuzzerTestOneInput':
+      continue
+    ci = context_introspector.ContextRetriever(benchmark)
+    header_file = ci.get_prefixed_header_file_by_name(undefined_func)
+    if header_file and header_file not in fuzz_target_source_code:
+      instruction += (
+          'You must add the following #include statement to fix the error of '
+          f'<error>undefined reference to {undefined_func}</error>:\n<code>\n'
+          f'{header_file}\n</code>.\n')
+    elif not header_file and benchmark.is_c_projcet:
+      instruction += (
+          f'You must remove the function <code>{undefined_func}</code> from the'
+          ' generated fuzz target, because the function does not exist.\n')
+    elif not header_file or header_file in fuzz_target_source_code:
+      # C project: NO header file found, or
+      # C++: Cannot map demangled C++ function name to signature
+      source_file = ci.get_prefixed_source_file(undefined_func)
+      if not source_file and benchmark.function_name in undefined_func:
+        source_file = ci.get_prefixed_source_file()
+      if source_file:
+        if header_file:
+          # To avoid redefinition.
+          instruction += ('You must remove the following statement\n<code>\n'
+                          f'{header_file}</code>\n')
+        instruction += (
+            'You must add the following #include statement to fix the error of '
+            f"<error>undefined reference to `{undefined_func}'</error>:\n"
+            f'<code>\n{source_file}\n</code>.\n')
+    else:
+      instruction += (
+          f"To fix <error>undefined reference to `{undefined_func}'</error>,"
+          'check the library documentation (e.g. README.md, comments) for '
+          'special instructions, such as required macros or specific inclusion '
+          'methods. Ensure any necessary definitions or inclusions are '
+          'correctly implemented in your generated fuzz target, following the '
+          "library's guidance.")
+    if not instruction:
+      instruction += (
+          f"To fix <error>undefined reference to `{undefined_func}'</error>,"
+          'check the library documentation (e.g. README.md, comments) for '
+          'special instructions, such as required macros or specific inclusion '
+          'methods. Ensure any necessary definitions or inclusions are '
+          'correctly implemented in your generated fuzz target, following the '
+          "library's guidance.")
+  return instruction
+
+
 def _collect_instruction_file_not_found(benchmark: benchmarklib.Benchmark,
                                         error: str,
                                         fuzz_target_source_code: str) -> str:
@@ -498,8 +557,8 @@ def _collect_instruction_file_not_found(benchmark: benchmarklib.Benchmark,
     return instruction
 
   # Step 3: Suggest the header/source file of the function under test.
-  function_file = ci.get_target_function_file_path()
-  if f'#include "{function_file}"' in fuzz_target_source_code:
+  function_file = ci.get_prefixed_header_file()
+  if function_file and f'#include "{function_file}"' in fuzz_target_source_code:
     function_file_base_name = os.path.basename(function_file)
 
     instruction += (

llm_toolkit/prompt_builder.py

@@ -198,6 +198,7 @@ def format_context(self, context_info: dict) -> str:
         must_insert=context_info['decl'],
         func_source=context_info['func_source'],
         xrefs='\n'.join(context_info['xrefs']),
+        include_statement=context_info['header'],
     )
 
   def _select_examples(self, examples: list[list],

prompts/template_xml/context.txt

@@ -1,14 +1,11 @@
-{% if headers %}
-
-The code block below contains the header files which you must include. The function requires these header files to successfully compile.
-If the headers are private, do not include them.
-Header files:
+{% if include_statement %}
+You must add the following #include statement in the fuzz target, its header file defines the function-under-test.
 <code>
-{{ headers }}
+{{ include_statement }}
 </code>
 {% endif %}
-{% if must_insert %}
 
+{% if must_insert %}
 You must insert code in the below code block before the function being tested:
 <code>
 {{ must_insert }}