Rewriting the Code: A Simple Method for Large Language Model Augmented Code Search

TL;DR

The paper addresses the challenge of modality misalignment in code search by extending Generation-Augmented Retrieval (GAR) with ReCo, a code-style normalization step that rewrites ground-truth code to match exemplar styles. ReCo boosts retrieval performance across sparse and dense settings and introduces Code Style Similarity (CSSim) as a dedicated metric for stylistic alignment, showing stronger explanatory power than existing metrics. The approach leverages prompt-based exemplar generation and a summarize-then-generate rewrite to create augmented queries and codes, optimizing both sparse and dense objectives via an InfoNCE loss. The work demonstrates practical gains across multiple datasets and languages, highlights the importance of code style in retrieval, and provides a foundation for future research on efficient, style-aware code search and related tasks.

Abstract

In code search, the Generation-Augmented Retrieval (GAR) framework, which generates exemplar code snippets to augment queries, has emerged as a promising strategy to address the principal challenge of modality misalignment between code snippets and natural language queries, particularly with the demonstrated code generation capabilities of Large Language Models (LLMs). Nevertheless, our preliminary investigations indicate that the improvements conferred by such an LLM-augmented framework are somewhat constrained. This limitation could potentially be ascribed to the fact that the generated codes, albeit functionally accurate, frequently display a pronounced stylistic deviation from the ground truth code in the codebase. In this paper, we extend the foundational GAR framework and propose a simple yet effective method that additionally Rewrites the Code (ReCo) within the codebase for style normalization. Experimental results demonstrate that ReCo significantly boosts retrieval accuracy across sparse (up to 35.7%), zero-shot dense (up to 27.6%), and fine-tuned dense (up to 23.6%) retrieval settings in diverse search scenarios. To further elucidate the advantages of ReCo and stimulate research in code style normalization, we introduce Code Style Similarity, the first metric tailored to quantify stylistic similarities in code. Notably, our empirical findings reveal the inadequacy of existing metrics in capturing stylistic nuances. The source code and data are available at \url{https://github.com/Alex-HaochenLi/ReCo}.

Figures (5)

• Figure 1: Comparison of GAR between passage retrieval and code search. In passage retrieval, the truth (yellow) is included in the generated content. In code search, despite the generated exemplar code satisfies the description of the query, it exhibits noticeable dissimilarity to the true code.
• Figure 2: An illustration of the ReCo method. It initially prompts LLMs to generate exemplar codes based on the search query. Subsequently, the original query and these exemplar codes are synthesized to formulate an augmented query. Analogously, the rewritten codes, produced by the LLMs, are merged with the original code, thereby creating a candidate for retrieval. The example delineated in this figure aligns with the one depicted in Fig. \ref{['fig1']}.
• Figure 3: Regression plots between $\Delta \mathrm{MRR}(=\mathrm{MRR}_{\mathrm{ReCo}}-\mathrm{MRR}_{\mathrm{GAR}})$ and $\Delta \mathrm{Metric Score}(=\mathrm{Metric}(c_q, c_c)-\mathrm{Metric}(c_q, c))$ under different evaluation metrics. The data points are from BM25 results on four datasets with four LLMs.
• Figure 4: Performance of BM25 + ReCo with different numbers of generated codes.
• Figure 5: Case study between Code Style Similarity and the existing metrics. The first exemplar code is preferred by existing metrics while the second one is preferred by Code Style Similarity. The second exemplar code is more similar to the true code from the perspective of style.