Add macro wrapped info in function context

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

Author: arthurscchan

data_prep/project_context/context_introspector.py

@@ -162,6 +162,28 @@ def _get_function_implementation(self) -> str:
 
     return function_source
 
+  def _get_macro_block(self) -> list[str]:
+    """Queries FI to check macro block around this function."""
+    project = self._benchmark.project
+    source = self._benchmark.function_dict.get('function_filename', '')
+    start = self._benchmark.function_dict.get('source_line_begin', 0)
+    end = self._benchmark.function_dict.get('source_end_begin', 99999)
+    macro_block = introspector.query_introspector_macro_block(
+        project, source, start, end)
+
+    results = set()
+    for macro in macro_block:
+      conditions = macro.get('conditions', [])
+      for condition in conditions:
+        cond_type = condition.get('type')
+        cond_detail = condition.get('condition')
+        if not cond_type or not cond_detail:
+          continue
+
+        results.add(f'#{cond_type} {cond_detail}')
+
+    return list(results)
+
   def _get_xrefs_to_function(self) -> list[str]:
     """Queries FI for function being fuzzed."""
     project = self._benchmark.project
@@ -177,13 +199,15 @@ def _get_xrefs_to_function(self) -> list[str]:
   def get_context_info(self) -> dict:
     """Retrieves contextual information and stores them in a dictionary."""
     xrefs = self._get_xrefs_to_function()
+    macro = self._get_macro_block()
     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,
+        'macro_block': macro,
         'func_source': func_source,
         'files': files,
         'decl': decl,

llm_toolkit/prompt_builder.py

@@ -203,6 +203,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']),
+        macro='\n'.join(context_info['macro']),
         include_statement=context_info['header'],
     )
 

prompts/template_xml/context.txt

@@ -20,6 +20,11 @@ Here is the source code of the function being tested:
 {% endif %}
 {% if xrefs %}
 
+Here are all the macro conditions wrapped around or present in the target functions:
+<code>
+{{ macro }}
+</code>
+
 Here is the source code for functions which reference the function being tested:
 <code>
 {{ xrefs }}