Dataflow-Guided Retrieval Augmentation for Repository-Level Code Completion (2024)

Abstract: Recent years have witnessed the deployment of code language models (LMs) in various code intelligence tasks such as code completion. Yet, it is challenging for pre-trained LMs to generate correct completions in private repositories. Previous studies retrieve cross-file context based on import relations or text similarity, which is insufficiently relevant to completion targets. In this paper, we propose a dataflow-guided retrieval augmentation approach, called DraCo, for repository-level code completion. DraCo parses a private repository into code entities and establishes their relations through an extended dataflow analysis, forming a repo-specific context graph. Whenever triggering code completion, DraCo precisely retrieves relevant background knowledge from the repo-specific context graph and generates well-formed prompts to query code LMs. Furthermore, we construct a large Python dataset, ReccEval, with more diverse completion targets. Our experiments demonstrate the superior accuracy and applicable efficiency of DraCo, improving code exact match by 3.43% and identifier F1-score by 3.27% on average compared to the state-of-the-art approach.

## introduction pre-trained language models (lms) of code @xcite @xcite have shown remarkable performance in improving programming productivity @xcite . instead of using a single code file, well-designed programs emphasize separating complicated functionality into independent modules @xcite . while facilitating collaborative development and software maintenance, it introduces the real-world problem of repository-level code completion: given an unfinished code file in ## related work code completion. early studies adopt statistical lms @xcite @xcite and neural models @xcite @xcite @xcite for code completion. after pre-training on large-scale code corpora, code lms are familiar with frequent code patterns and achieve superior performance @xcite @xcite @xcite @xcite @xcite . unlike singlefile code completion, repository-level code completion draws much attention to practical development. ding et al. (2024) learn in-file and cross-file context jointly on top of pre-trained lms. lu et al. (2022) present reacc to train a code-to-code search retriever and a code completion generator with an external source code database. shrivastava et al. (2023b) generate example-specific prompts using a prompt proposal classifier and further propose repofusion (shrivastava et al., 2023a) to incorporate relevant repository context by training code lms. repocoder (zhang et al., 2023) is an iterative retrieval-generation framework to approximate the intended completion target. despite their good performance, these methods are limited by the high overhead of extra training or iterative generation. retrieval-augmented generation. for the scenarios where required knowledge is missing or outdated in pre-trained lms, rag has achieved stateof-the-art performance in many nlp tasks @xcite @xcite . usually, rag integrates the retrieved knowledge with frozen pre-trained lms @xcite @xcite . there exist different types of retrievals including term-based sparse retriever @xcite , embedding-based dense retriever @xcite , commercial search engines @xcite , and lms themself @xcite . rag is also broadly applied to code intelligence tasks such as code summarization @xcite @xcite and code generation @xcite @xcite . in this work, we leverage dataflow analysis to guide retrieval, which mines more precise data dependency information for repository-level code completion. ## methodology as shown in figure 2 , draco employs the standard rag framework (lewis et al., 2020) including indexing ( §3.2), retrieval ( §3.3), and generation ( §3.4). because both indexing and retrieval in draco rely on the extended dataflow analysis, we first introduce it in §3.1. in this work, we focus on python and the code completion of currently edited line, which simulates real-world scenarios where users are programming in integrated development environments (ides) and only the context before the cursor is visible. ## conclusions this paper proposes draco, a dataflow-guided retrieval augmentation approach for repositorylevel code completion. to guide more precise retrieval, an extended dataflow analysis is designed by setting type-sensitive data dependency relations. draco indexes private repositories to form repospecific context graphs and retrieves relevant background knowledge from them, which is assembled with the unfinished code to generate well-formed prompts for querying code lms. the experiments on the crosscodeeval dataset and our recceval dataset demonstrate the superior accuracy and applicable efficiency of draco. in future work, we will explore other structured code representations.