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

TL;DR

This paper enhances code search by introducing a simple code rewriting method called ReCo, which normalizes code style to improve retrieval accuracy in large language model augmented frameworks.

Contribution

The paper proposes ReCo, a straightforward code rewriting technique for style normalization, significantly improving code retrieval performance in LLM-augmented search frameworks.

Findings

ReCo boosts retrieval accuracy up to 35.7% in sparse settings.

ReCo improves zero-shot dense retrieval by up to 27.6%.

ReCo enhances fine-tuned dense retrieval by up to 23.6%.

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}.