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-ofthe-art approach.

heading: 1 Introduction
text: Pre-trained language models (LMs) of code (Chen et al., 2021; Nijkamp et al., 2023a,b; Allal et al., 2023; Li et al., 2023b) have shown remarkable performance in improving programming productivity (Kazemitabaar et al., 2023; Dakhel et al., 2023). Instead of using a single code file, well-designed programs emphasize separating complicated functionality into independent modules (Barnett and Constantine, 1968). While facilitating collaborative development and software maintenance, it introduces the real-world problem of repository-level code completion: given an unfinished code file in
∗ Corresponding author

heading: Dataflow-guided:
text: a private repository, complete the following pieces of code at the cursor position.
Despite pre-training on large-scale corpora, code LMs are still blind to unique naming conventions and programming styles in private repositories (Pei et al., 2023; Liu et al., 2024b; Ding et al., 2023). Previous works finetune LMs to leverage crossfile context (Ding et al., 2024; Shrivastava et al., 2023a,b), which requires additional training data and is difficult to work with larger LMs. Recently, retrieval-augmented generation (RAG) is widely used to aid pre-trained LMs with external knowledge and maintain their parameters intact (Lewis et al., 2020; Mallen et al., 2023; Trivedi et al., 2023). For repository-level code completion, the retrieval database is the current private repository. The state-of-the-art approach, RepoCoder (Zhang et al., 2023), iteratively incorporates a text similarity-based retriever and a code LM.
As shown in Figure 1, the CodeGen25 Python model (Nijkamp et al., 2023a) with 7 billion parameters assigns a value to the attribute channel of the object newSignal, which seems rational in
7957
the unfinished code but is actually outside the list of valid attributes. Due to the lack of similar code snippets in the repository, the text similarity-based approach (Zhang et al., 2023) also fails to complete the correct code line. From a programmer’s perspective, one would explore the data origin of the variable newSignal in the last line. It comes from the call signal.getSignalByName, where the variable type of signal is RecordSignal imported from the module RecordSignal. After providing relevant background knowledge in the private repository, the model would know that the variable type of newSignal is the class Signal and thus call the correct function.
Inspired by this programming behavior in private repositories, we propose DRACO, a novel dataflow-guided retrieval augmentation approach for repository-level code completion, which steers code LMs with relevant background knowledge rather than similar code snippets. Dataflow analysis is a static program analysis reacting to data dependency relations between variables in a program. In this work, we extend traditional dataflow analysis by setting type-sensitive dependency relations. We employ the standard RAG framework (Lewis et al., 2020): (i) Indexing, which parses a private repository into code entities and establishes their relations through dataflow analysis, forming a repospecific context graph for retrieval. (ii) Retrieval, which uses dataflow analysis to obtain fine-grained import information in the unfinished code and retrieves relevant code entities from the pre-built context graph. (iii) Generation, which organizes the relevant background knowledge as natural code and concatenates it with the unfinished code to generate well-formed prompts for querying code LMs.
In addition to the existing dataset CrossCodeEval (Ding et al., 2023) for repository-level code completion, we construct a new dataset, ReccEval, with diverse completion targets collected from Python Package Index (PyPI).1 We conduct experiments with popular LMs including adapted code LMs (Rozière et al., 2023), specialized code LMs (Nijkamp et al., 2023a,b; Allal et al., 2023; Li et al., 2023b), and GPT models (Ouyang et al., 2022; OpenAI, 2023). Our experiments demonstrate that DRACO achieves generally superior accuracy across all settings. Furthermore, DRACO is plug-and-play for various code LMs and efficient to real-time code completion.
1https://pypi.org/
Our main contributions are outlined as follows: • We design an extended dataflow analysis by
setting type-sensitive data dependency relations, which supports more precise retrieval. • We propose DRACO, a dataflow-guided retrieval augmentation approach for repositorylevel code completion. DRACO builds a repospecific context graph for retrieval and generates well-formed prompts with relevant background knowledge in real-time completion. • We construct a Python dataset ReccEval with diverse completion targets. The experimental results show that DRACO improves code exact match by 3.43%, identifier F1-score by 3.27%, and prompt generation time by 100× on average compared to the second-best approach RepoCoder (Zhang et al., 2023). Our source code and data are available at https: //github.com/nju-websoft/DraCo.

heading: 2 Related Work
text: Code completion. Early studies adopt statistical LMs (Raychev et al., 2014; Proksch et al., 2015; Raychev et al., 2016; He et al., 2021) and neural models (Li et al., 2018; Svyatkovskiy et al., 2019; Kim et al., 2021; Izadi et al., 2022; Tufano et al., 2023) for code completion. After pre-training on large-scale code corpora, code LMs are familiar with frequent code patterns and achieve superior performance (Chen et al., 2021; Lu et al., 2021; Wang et al., 2021; Le et al., 2022; Allal et al., 2023; Li et al., 2023b; Nijkamp et al., 2023a,b; Shen et al., 2023; Zheng et al., 2023). 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 out-

heading: Repository Dataflow analysis
text: dated in pre-trained LMs, RAG has achieved stateof-the-art performance in many NLP tasks (Cai et al., 2022; Feng et al., 2023; Mallen et al., 2023). Usually, RAG integrates the retrieved knowledge with frozen pre-trained LMs (Ram et al., 2023; Levine et al., 2022; Shi et al., 2023). There exist different types of retrievals including term-based sparse retriever (Robertson and Zaragoza, 2009; Trivedi et al., 2023), embedding-based dense retriever (Karpukhin et al., 2020; Lewis et al., 2020), commercial search engines (Nakano et al., 2021; Liu et al., 2023), and LMs themself (Yu et al., 2023; Sun et al., 2023). RAG is also broadly applied to code intelligence tasks such as code summarization (Liu et al., 2021; Zhang et al., 2020; Zhou et al., 2023) and code generation (Hashimoto et al., 2018; Parvez et al., 2021; Li et al., 2023a). In this work, we leverage dataflow analysis to guide retrieval, which mines more precise data dependency information for repository-level code completion.

heading: 3 Methodology
text: 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.

heading: 3.1 Dataflow Analysis
text: Dataflow analysis is a static program analysis that reacts to the data dependency relations between variables in a program, producing a dataflow graph (DFG), in which nodes represent the variables and edges indicate where the variables come from and where they go. It provides code semantic information that is not affected by personal naming conventions and programming styles.
We assume that the background knowledge relevant to variable types is crucial for code completion. Take the statement v = f(p) as an example, the parameter p has far less influence on the variable v than the call f does. Therefore, we extend traditional dataflow analysis by setting dependency relation types. We focus on five type-sensitive relations, which indicate what the variable type is or where it derives from:
• Assigns relation is a one-to-one correspondence in an assignment statement, which controls variable creation and mutation. • As relation is from with or except statements and similar with the assigns relation. • Refers relation represents a reference to an existing variable or its attribute. • Typeof relation is from the explicit type hints (van Rossum and Lehtosalo, 2022) written by programmers, indicating the data type of the (return) value of a variable or function. • Inherits relation is an implicit data dependency relation since a subclass inherits all the class members of its base classes.
Our DFG is a heterogeneous directed acyclic graph G = {(h, r, t) | h, t ∈ E, r ∈ R}, where E
denotes the entity set, R denotes the type-sensitive relation set, and a triplet (h, r, t) represents the head entity h pointing to the tail entity t with the relation r. The details of our DFG construction are shown in Appendix A.

heading: 3.2 Repo-specific Context Graph
text: Offline preprocessing is often used in RAG to index a retrieval database. Instead of treating source code as text (Lu et al., 2022; Zhang et al., 2023), we parse a private repository into code entities and establish their relations through our dataflow analysis, forming a repo-specific context graph.
For each code file in a repository, we traverse its abstract syntax tree (AST) to collect code entities including modules, classes, functions, and variables. A module entity stores its file path and docstring as properties. A class entity stores its name, signature, docstring, and starting line number. A function entity stores its name, signature, docstring, body, and starting line number. A variable entity stores its name, statement, and starting line number. There are natural contains relations between these entities, e.g., the class RecordSignal contains its member function getSignalByName. Based on the type-sensitive relations in DFG, we establish depends relations between the entity pairs in individual modules, e.g., Signal is the return type of the function getSignalByName. Eventually, we establish depends relations between the variables in local import statements and the pointing entities in other modules, e.g., the imported Signal points to the class Signal in another module.

heading: 3.3 Dataflow-Guided Retrieval
text: Given an unfinished code, we identify fine-grained import information by dataflow analysis and retrieve relevant entities from the repo-specific context graph. We do not intend to perform precise type inference (Peng et al., 2022) for a dynamically typed language like Python, but rather provide relevant background knowledge to code LMs, which provides the definitions of code entities such as class members and function arguments.
All cross-file context is indicated by local import statements in Python. However, only using such coarse-grained import information may overlook the knowledge of its specific usages (Ding et al., 2024). We denote import information by (module, name), where module indicates another code file in the repository and name indicates the specific code entity. Particularly, name can be ex-
panded by its refers relations in the extracted DFG. For example, we would obtain the fine-grained import information (module, name.attr) if there is a code statement containing name.attr.
For each local import statement, we collect a set of fine-grained import information. Import information points to code entities in the repo-specific context graph, which is achieved through directory structure and string matching. Let us see Figure 2 for example. Given the fine-grained import information (RecordSignal, RecordSignal.getSignalByName), we first identify the corresponding module entity pyPhasesRecordloader/RecordSignal.py through directory structure. Then, the code entity getSignalByName in class RecordSignal contained in the module is located by string matching. Finally, the relevant entities are retrieved along depends relations using a depth-first search. The retrieved entities provide comprehensive type-related background knowledge for both cross-file imports and usages in the unfinished code.

heading: 3.4 Prompt Generation
text: Before querying LMs, we restore the retrieved entities to the source code and concatenate it with the unfinished code to generate well-formed prompts.
As the maximum input lengths of LMs are finite and fixed, we employ the dynamic context allocation strategy in (Shrivastava et al., 2023b). It pre-allocates half of the total input lengths for the relevant background knowledge and the other half for the unfinished code. If either is shorter than the allocated length, the remaining tokens are allocated to the other. For the overlong unfinished code, we just truncate it to obtain the last tokens. Given the number n of tokens allocated to the relevant background knowledge, we next describe the generation process of background knowledge. Refer to Algorithms 1 and 2 in Appendix B for details.
Algorithm 1 describes the center control of the generation process. We consider the entities that have data relations with the line to be completed as the most relevant entities, which are denoted by Er and form the primary background knowledge. The entities from other local import statements are denoted by Eo and are incrementally added to the background knowledge until the length exceeds the allocated tokens n. This design helps LMs further understand the context while prevents the primary background knowledge from being truncated.
Our mission is to organize the prompts like the
original code to maintain the nature of programs (Hindle et al., 2012), as described in Algorithm 2. We group the relevant entities in modules and merge those with contains relations to avoid redundancy, e.g., class members would not be duplicated if the class already exists. The code entities in the same module are sorted by their starting line numbers. Therefore, we construct a module graph Gm, where an edge m1 → m2 indicates that there exists an entity in m1 that depends on an entity in m2. The modules in prompts are ranked by a pseudotopological sort with two priorities: (i) Dependent modules come first. The content of m2 should be placed in front of that of m1, which is consistent with programming conventions (Lines 10–15 in Algorithm 2). (ii) Once there are multiple options for topological sort in a directed cyclic graph, the relevant modules are placed ahead (Liu et al., 2024a). We prefer the modules that are reachable from the entities in Er than those in Eo (i.e., relevant or not). The entities within both Er and Eo are sorted in ascending order by their starting line numbers, resulting in an ordered list Ec of the import entities (Lines 3–9, 13). Note that a comment “# file path of the module” is put ahead of each module to indicate the relative directory structure. Finally, we place the relevant background knowledge inside a multi-line string, which precedes the unfinished code to generate the prompts for querying LMs.
Benefiting from the design of our repo-specific context graph, there are two prompt scopes, named definition and complete, to control the details of code entities. Compared with only definitions, prompts under the complete scope contain specific function bodies and variable statements.

heading: 4 Experiment Setup

heading: 4.1 Datasets
text: The widely-used datasets (Raychev et al., 2016; Lu et al., 2021; Peng et al., 2023) for code completion only provide a single unfinished code file as input. Several recent benchmarks (Zhang et al., 2023; Liu et al., 2024b) evaluate next-line prediction, which is different from our concern with the currently edited line. CrossCodeEval (Ding et al., 2023) is a multilingual benchmark for repository-level code completion, where the statement to be completed has at least one use of cross-file API. Since we focus on Python, we evaluate our DRACO on the Python subset of CrossCodeEval.
We further build a new Python dataset ReccE-
val with more diverse completion targets. See Appendix C.1 for details. The statistics of ReccEval and the Python subset of CrossCodeEval are shown in Table 1, where the number of tokens is calculated using the StarCoder tokenizer (Li et al., 2023b).

heading: 4.2 Implementation Details
text: We evaluate the retrieval-augmented methods that do not involve training, which excludes several works (Shrivastava et al., 2023a,b; Lu et al., 2022). See Appendix C.2 for more details:
• Zero-Shot directly feeds the unfinished code to code LMs, which evaluates their performance without any cross-file information. • CCFinder (Ding et al., 2024) is a crossfile context finder tool retrieving the relevant cross-file context from the pre-built project context graph by import statements. We conduct experiments for CCFinder-k (k = 1, 2), which indicates that CCFinder retrieves k-hop neighbors of cross-file code entities. • RG-1 and RepoCoder (Zhang et al., 2023) construct a retrieval database through a sliding window and retrieve similar code snippets using text similarity-based retrievers. RepoCoder is an iterative retrieval-generation framework, which retrieves the database with the results generated in the previous iteration. RG-1 represents the standard RAG and is the first iteration of RepoCoder.
As shown in Table 2, we conduct comprehensive experiments on seven popular LMs. For a method, we first preprocess all repositories in the datasets. Then, we generate prompts for the unfinished code and record the time used. Finally, we acquire the completion results by feeding the prompts to LMs. A prediction is the first line of a completion result.

heading: 4.3 Evaluation Metrics
text: We evaluate the accuracy of each method by code match and identifier match scores (Ding et al., 2023), as well as the efficiency by prompt generation time. We report the average of each metric. See Appendix C.4 for more details:
• Code match. We directly compare the prediction with the reference, which is measured using exact match (EM) and edit similarity (ES) (Lu et al., 2021; Zhang et al., 2023). • Identifier match. We evaluate the predicted APIs by identifier exact match (ID.EM) and F1-score (Ding et al., 2023). • Prompt generation time. We record the prompt generation time to evaluate the efficiency of each method, which is a new and significant metric for real-time completion.

heading: 5 Experimental Results and Analysis

heading: 5.1 Performance Comparison
text: The performance comparison on the CrossCodeEval and ReccEval datasets is listed in Tables 3 and 4, respectively. Additional results on other CodeGen models are supplemented in Appendix D.1. Overall, DRACO significantly improves the accuracy of various code LMs. Particularly, the CodeGen350M model integrated with DRACO even outperforms the zero-shot StarCoder-15.5B model.
In comparison to other retrieval-augmented methods, DRACO also shows generally superior accuracy across all settings. The average absolute improvement on EM, ES, ID.EM, and F1 versus RepoCoder is 3.43%, 1.00%, 3.62%, and 3.27%, respectively. RepoCoder retrieves similar code demonstrations that help increase the ES metric of completion results. But RepoCoder ignores the validity of its generated identifiers in private repositories, which decreases the metrics for code exact match and identifier match. Such almost correct completion results may introduce unconscious bugs
to the programmers who are unfamiliar with the repository. In contrast, DRACO presents the definitions of relevant code entities, providing better control over code LMs to generate valid identifiers. Moreover, the background knowledge can be used as a reference to help the programmers understand and review the completion results in IDEs. DRACO using the CodeGen-350M model is slightly worse than RepoCoder in terms of code match metrics on the ReccEval dataset, as the model may not be powerful enough to capture the data relations in our provided background knowledge.
CCFinder retrieves cross-file code entities by plain import relations. The entities retrieved by CCFinder were originally designed to be encoded for training code LMs. When used as a retrievalaugmented method, CCFinder retrieves too many code entities based on coarse-grained import information, resulting in truncation of truly relevant context. As a result, CCFinder-2 with more retrieved entities only exceeds CCFinder-1 slightly on the StarCoder model that supports longer inputs. Therefore, the subsequent analysis experiments are conducted with CCFinder-1. Guiding by our dataflow analysis, DRACO retrieves relevant code entities more precisely, leading to significantly superior accuracy.
The performance of code completion varies on the two datasets. First, according to the statistics in Table 1, the average reference length of ReccEval is significantly shorter than that of CrossCodeEval, leading to the higher EM metrics of both code and identifier on ReccEval. Moreover, all inputs of CrossCodeEval end with a dot where a correct API is required in the first place, which is more suitable for CCFinder and DRACO that retrieve code definitions. Many inputs of ReccEval end with partial names of the target APIs, which facilitates text similarity-based retrievals including RG-1 and RepoCoder. Therefore, the lead of DRACO on CrossCodeEval is more significant.

heading: 5.2 Efficiency Evaluation
text: The time spent on prompt generation is perceived by users whenever code completion is triggered. Table 5 shows the prompt generation time of each method using the CodeGen-350M model, and additional results are shown in Appendix D.2. CCFinder and DRACO require parsing the unfinished code into an AST or a DFG, which is slightly slower than RG-1 with text similarity-based retrieval but still comparable. RepoCoder relies on
RG-1 to generate sufficient content for the second retrieval, which results in more than 4 seconds even on the smallest CodeGen-350M model and may not be feasible for real-time code completion.
In summary, DRACO is efficient for real-time code completion in IDEs. Compared to the methods with comparable efficiency (i.e., excluding RepoCoder), DRACO is considerably ahead in the accuracy of repository-level code completion.

heading: 5.3 Ablation Study
text: To analyze the effectiveness of dataflow analysis in DRACO, we conduct an ablation study shown in Tables 6 and 17. “w/o cross_df” disables depends relations in the repo-specific context graph, making DRACO unable to handle the data dependency relations in other code files. “w/o intra_df” disables the dataflow analysis for the unfinished code, which only allows DRACO to retrieve coarsegrained import information in the order of their starting line numbers. “w/o dataflow” degenerates DRACO into a naive method that simply takes the imported cross-file entities in the unfinished code as the relevant background knowledge.
The ablation study demonstrates that the complete DRACO achieves the best performance, and all usages of dataflow analysis play a positive role in repository-level code completion. It can be observed that the enhancement of the “intra_df” component on the StarCoder model is less than that on other models. This component places the more relevant background knowledge in front of the prompt to prevent truncation, which is weakened to some extent on the StarCoder model with a maximum context length of 8K tokens.
The accuracy of DRACO without dataflow analysis is still comparable with CCFinder. CCFinder groups the relevant context in code entities, which is counter-intuitive for source code (see the example shown in Appendix E.1). The results reveal that the well-formed prompts generated by DRACO can better steer code LMs, even if the depth-first search for dependent code entities is absent.

heading: 5.4 Analysis of Prompt Scopes
text: The prompts generated by DRACO consist of the definitions of code entities, which provide options for the definition and complete scopes, as described in Section 3.4. We further conduct experiments to evaluate the influence of the two prompt scopes. The results on the CrossCodeEval and ReccEval datasets are shown in Figures 3 and 6, respectively.
DRACO with the complete scope achieves the best performance across all settings, which indicates that code implementation can further enhance LMs. Implementation details can provide a deeper understanding of code entities, along with the programming styles. Moreover, DRACO with the definition scope outperforms CCFinder and RG-1 in most settings (cf. Tables 3 and 4), suggesting that the definitions without specific implementations are also useful for code LMs. Since an implementation is usually much longer than its definitions, both prompt scopes are optional in practical applications, in a trade-off between accuracy and cost.

heading: 5.5 Analysis of Adapted LM
text: We analyze the effect of different maximum input lengths for different types of LMs. Specifically, we evaluate Code Llama-7B and StarCoder-15.5B with 2K, 4K, and 8K tokens. The performance changes are shown in Figure 4, and the complete results are presented in Tables 18–21.
With the increase of the maximum input length, the accuracy of DRACO applied to Code Llama decreases, which shows the opposite trend of the StarCoder model. Code Llama is created by further training Llama 2 on its code-specific datasets. It is
different from specialized code LMs such as StarCoder which are mainly pre-trained on code corpora. With similar background knowledge, Code Llama-7B does not have enough capability to capture data dependency relations in long Python code, leading to the degraded accuracy of DRACO. In contrast, specialized code LMs can understand longer code context and may be a better choice for repository-level code completion.

heading: 5.6 Analysis of GPT Models
text: We randomly sample 50 examples from CrossCodeEval and evaluate them with GPT-3.5, GPT4, and StarCoder-15.5B. The results shown in Table 7 reveal that: (i) DRACO can also significantly enhance GPT models and achieve superior accuracy. (ii) Code completion with GPT models suffers from the difficulties in output format and API
cost. Given the long context of repository-level code completion, there lacks sufficient length to place the demonstrations required for in-context learning (Brown et al., 2020). It is hard to control the output format of GPT models through instruction, which may introduce bias into the evaluation, especially for RG-1 with GPT-4. Moreover, the API cost for this evaluation (only 50 examples) is nearly 35 US dollars. (iii) Excluding the effect of wrong output format, we may assume “GPT-4 > StarCoder-15.5B > GPT-3.5” for this task.

heading: 6 Conclusions
text: 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.

heading: Acknowledgments
text: We thank the anonymous reviewers for their valuable comments. This work was supported by the National Natural Science Foundation of China (No. 62272219) and the Collaborative Innovation Center of Novel Software Technology & Industrialization.

heading: Ethical Considerations
text: The code generated by pre-trained LMs may contain non-existent APIs or even introduce potential bugs. The retrieval-augmented approaches including ours mitigate this issue only to some extent. We recommend presenting our retrieved background knowledge to programmers for review and taking appropriate care of these risks if deploying our approach in real-world applications.
All the datasets and code LMs used in this work are publicly available with permissive licenses. The CrossCodeEval dataset and CodeGen family are licensed under the Apache-2.0 License. The SantaCoder and StarCoder models are licensed under the BigCode OpenRAIL-M v1 license agreement.
Code Llama is governed by the Meta license.2 The repositories in our ReccEval dataset are all licensed under permissive licenses including MIT, Apache, and BSD licenses.

heading: Limitations
text: DRACO relies on a code LM to support long inputs and capture data dependency relations in the provided background knowledge. Thus, the performance of DRACO may be limited by the capability of the code LM. According to our experiments, DRACO still has a considerable improvement on the smallest CodeGen-350M model with 2K tokens, which mitigates this limitation.
The effectiveness of DRACO may degrade when the code intent is unclear. For new line or function body completion, the guidance of dataflow analysis is weakened since DRACO may not be able to set priorities for import information. We focus on the code completion of currently edited line and evaluate multi-line code completion in Appendix D.4. Future work may explore the role of dataflow analysis in different completion scenarios.
DRACO requires changes to migrate to other programming languages. Our idea of guiding retrieval with dataflow analysis is not limited to Python. However, due to the different characteristics of programming languages, DRACO needs to extend dataflow analysis for target languages. The variety of static analysis tools for common programming languages provides convenience for implementing multilingual DRACO.

heading: A Dataflow Graph Construction
text: We first parse the Python code into an AST by treesitter,3 which is feasible to parse incomplete code snippets. Based on the AST, we extract variables as the entity set and identify type-sensitive relations as triplets, forming our DFG. The implementation involves specific Python syntax features and is presented in our published code.
The type-sensitive relations are listed in Table 8. We also introduce an example to visualize our DFG construction, as shown in Figure 5. The extracted variables include identifiers (e.g., newSignal) and attributes (e.g., signal.getSignalByName). The type-sensitive relations are identified as described in Section 3.1. For example, the parameter signal: RecordSignal indicates a triplet (RecordSignal, Typeof, signal). The assignment statement forms a triplet (signal.getSignalByName, Assigns, newSignal), and the parameter newChannelName is a type-insensitive relation pruned in our DFG.

heading: B Generation of Background Knowledge
text: Algorithms 1 and 2 describe the generation process of the relevant background knowledge (Section 3.4), where the time complexities areO(|Eo|× (|M |+ |P |)) and O(|M |+ |P |), respectively.

heading: C Details of Experiment Setup

heading: C.1 Details of Dataset Construction
text: We collect the projects that are first released on PyPI between 2023-01-01 to 2023-04-28, which is after the releases of pre-training corpora (Husain et al., 2019; Chen et al., 2021; Kocetkov et al., 2023). We pick the projects with permissive licenses (i.e., MIT, Apache, and BSD) and filter out those that have fewer than 6 or more than 100 Python code files. We identify the usages of local imported resources and randomly select a subsequent token as the cursor position. The context before the cursor is the input, while the current line after the cursor is the reference. For the diversity of ReccEval, we limit the maximum number of examples to one per code file and 10 per repository. Moreover, we ensure that the reference is not in the unfinished code and feed the examples to StarCoderBase-1B model (Li et al., 2023b) to remove the exact matches (Ding et al., 2023), which excludes strong clues in the unfinished code to make ReccEval more challenging.
3https://github.com/tree-sitter/tree-sitter
T
(golemgpt.memory, BaseMemory)
self.memory 9
self.memory 6 memory 6
memory 5BaseMemory 5
BaseMemory 1
A
T CF
CF
CF
newSignal 4
newSignal 3
signal.getSignalByName 3
signal 2 RecordSignal 2
RecordSignal 1
CF
RecordSignal
RecordSignal.getSignalByName ...
A
CF
CF
Figure 5: An example of our DFG, which corresponds to the unfinished code in Figure 1. The numbers labeled in the DFG correspond to the line numbers of the variables. The labels on the edges are the initials of the relation names defined in Section 3.1.
C.2 Implementation Details of Baselines
We describe more implementation details of CCFinder, RG-1, and RepoCoder, which are in line with the experiment setup in their papers:
• CCFinder. Because CCFinder is not open source, we reproduce it according to its paper. We do not limit the number of retrieved code entities, as the cross-file context would be truncated if it exceeds the maximum length. We also re-order the retrieved entities, ensuring the entities from the same source file follow the original code order. • RG-1 and RepoCoder. In our experiments, we use a sparse bag-of-words model as their retriever, which calculates text similarity using the Jaccard index and achieves equivalent performance to the dense retriever. The line length of the sliding window and the sliding size are set to 20 and 10, respectively. According to the maximum input length of code LMs, the maximum number of the retrieved code snippets in prompts is set to 40 for the StarCoder model and 10 for other models. The number of iterations of RepoCoder is set to 2.

heading: C.3 Details of Used LMs
text: We categorize the used LMs into specialized code LMs, adapted code LMs, and GPT models. The
1 Ec ← Er; 2 bk ← Organize_BK(Ec); 3 foreach e ∈ Eo do 4 Append e to Ec; 5 temp← Organize_BK(Ec); 6 if Length(temp) ≤ n then
// as many entities as possible 7 bk ← temp; 8 else
// prevents the primary background knowledge from being truncated
9 break;
10 if Length(bk) > n then 11 bk ← the first n tokens of bk; 12 return bk
details are listed as follows: • CodeGen (Nijkamp et al., 2023a,b) is a family
of auto-regressive LMs for program synthesis. We use the CodeGen2.5 model with 7B parameters and the CodeGen models with 350M, 2.7B, 6.1B, and 16.1B parameters, which support a maximum context length of 2,048 (2K) tokens. We use their mono versions, which are further trained on additional Python tokens. • SantaCoder (Allal et al., 2023) is a code model with 1.1B parameters, which supports a maximum context length of 2K tokens. • StarCoder (Li et al., 2023b) is a 15.5B model trained on 80+ programming languages and further trained on Python, which supports a maximum context length of 8K tokens. • Code Llama (Rozière et al., 2023) is created by further training Llama 2 (Touvron et al., 2023) on its code-specific datasets, which supports a maximum context length of 16k tokens. Considering GPU member and efficiency, we chose the 7B Python specialist version named CodeLlama-7b-Python-hf and limit the maximum context length to 8K tokens. • GPT models (Ouyang et al., 2022; OpenAI, 2023) are commercial black box models released by OpenAI. The used API of GPT-3.5
Algorithm 2: Organize_BK(Ec) Input: the ordered import entities Ec Output: background knowledge bk // let Gm = (M,P ), where M is the
vertex set, and P is the edge set
1 Gm ← retrieve all dependent entities of Ec and group them in modules; 2 Me ← the corresponding modules of Ec; // set the priorities of modules M 3 Initialize an empty dictionary priorities; 4 rank ← 1; 5 foreach m ∈Me do // priorities ii 6 priorities[m]← rank; 7 rank ← rank + 1; 8 Md ←M −Me; // details omitted: pass priorities from
Me to the dependent modules Md
9 ∀m ∈Md, priorities[m]← the minimum priority of the direct predecessors of m;
// priority i: pseudo-topological sort 10 Initialize an empty list Ms; 11 while M is not empty do 12 Mc ← the candidate modules with the
least direct predecessors in Gm; // alphabetical order for equality
13 m← the module m ∈Mc with the maximum priorities[m]; 14 Pop m from M and append it to Ms;
15 bk ← combine all modules in the reverse order of Ms; 16 return bk
is gpt-3.5-turbo-0613 with a maximum context length of 4K tokens. The used API of GPT-4 is gpt-4-0613 with a maximum context length of 8K tokens. The analysis experiments with Code LLama and GPT models may suffer from data leakage issues. As mentioned in Appendix C.1, both CrossCodeEval and our ReccEval devote effort to collecting projects (e.g., released between 2023-01-01 to 2023-04-28 for ReccEval) that are not in pretraining corpora. However, Code Llama is trained on unknown datasets between January 2023 and July 2023, and the training data of GPT models is unknown and updated.
We set the temperature of these LMs as 0 to obtain deterministic results. The maximum generation length is set to 48 tokens, which is long enough
to accomplish line completions. An exception is RG-1, which asks LMs to generate 100 tokens since RepoCoder requires sufficient content for further retrieval. We run StarCoder-15.5B, CodeGen-16.1B, and Code Llama-7B on an NVIDIA A800 with 80GB memory and run other LMs on an NVIDIA GeForce RTX 4090 with 24GB memory.
The instruction of GPT models is “You are a Python expert. Please complete the last line of the following Python code:”. For RG-1, the instruction is “You are a Python expert. Please complete the following Python code:”. We set line feed as the stop token of LM generation, except for RG-1 (still 100 tokens).

heading: C.4 Details of Evaluation Metrics
text: • Code match. Given a prediction y and the reference y∗, we assess y using the exact match accuracy (EM) and the Levenshtein edit similarity (ES) (Lu et al., 2021; Zhang et al., 2023). EM is calculated by an indicator function whose value is 1 if y = y∗; otherwise, it is 0. ES = 1− Lev(y,y∗)max(||y||,||y∗||) , where || · || calculates the string length and Lev() calculates the Levenshtein distance. • Identifier match. Identifier exact match (ID.EM) and F1-score test the model’s ability to predict the correct APIs (Ding et al., 2023). We parse the code to extract the identifiers from y and y∗ and get two ordered identifier lists. ID.EM is calculated by an indicator function whose value is 1 if their elements are the same and in the same order; otherwise, it is 0. Then, we transform the lists into two sets (without repeated elements and unordered), which can be used to calculate the precision, recall, and F1-score, where F1-score is a combination of precision and recall. • Prompt generation time. To evaluate the efficiency of code completion, we record the prompt generation time, which contains the time to retrieve relevant context and the time to assemble final prompts. Note that we ignore the time spent by code LMs in generating predictions, which is determined by the used LMs rather than the methods.

heading: D Additional Evaluation

heading: D.1 More Performance Comparison Results
text: Beyond the experimental results of the main paper, we show additional evaluation results of other
CodeGen models in Tables 9 and 10. The additional results show consistent conclusions on performance comparisons in the main paper. Under the same architecture of the CodeGen-* models, the performance of all methods improves as the model parameters increase. Moreover, the improvement of DRACO for zero-shot code LMs increases as the model’s capability grows. It indicates that stronger LMs can better utilize the relevant background knowledge retrieved by DRACO.

heading: D.2 More Efficiency Evaluation Results
text: We also record the time spent on indexing the repositories of CrossCodeEval and ReccEval, as shown in Table 11. It is an offline preprocessing in RAG, which indicates the time required to activate a method. CCFinder and DRACO build retrieval databases by statically parsing code files, which are independent of the used code LMs. RG-1 and Re-
poCoder need to tokenize the code snippets within a sliding window, which requires the tokenizers of used LMs. Note that the tokenizers of CodeGen-* models are the same. DRACO is 3–6 times faster than RepoCoder in preprocessing time. As the size of the repository increases, the preprocessing time grows linearly. Therefore, RG-1 and RepoCoder may suffer from scalability challenges.
The prompt generation time of each method using other code LMs is shown in Tables 12 and 13, which show consistent conclusions with the main paper. For the methods with one retrieval, only the tokenizers have a subtle effect on efficiency when different models are employed. As a result, the prompt generation time using different CodeGen-* models is the same for CCFinder, RG-1, as well as DRACO. RepoCoder relies on RG-1 to generate sufficient content for the second retrieval, where the efficiency mainly depends on the generation time of code LMs. In general, the generation efficiency of RepoCoder decreases as the model parameters increase. Its average prompt generation time is more than 3 seconds on the most efficient SantaCoder model, which far exceeds the time spent by other retrieval-augmented methods. Note that the architectures of code LMs also matter in efficiency, e.g., SantaCoder-1.1B is faster than CodeGen-350M. The A800 GPU used to run the StarCoder-15.5B and CodeGen-16.1B models is superior to the RTX 4090 GPU used for the other models, so these are
not head-to-head comparisons for RepoCoder.

heading: D.3 Effect of Other Import Statements
text: As described in Section 3.4, DRACO includes as many other import statements that are not directly relevant as possible. To analyze the effect of other import statements, we implement a variant of DRACO for comparison, which adds the entities from other local import statements only when the primary prompt is empty. The experimental results are shown in Table 14. DRACO consistently outperforms this variant, especially on StarCoder15.5B, which has capability to understand longer code context.

heading: D.4 Evaluation of Multi-Line Completion
text: Although DRACO is designed for the code completion of currently edited line, it can help with
multi-line completion to some extent.
We begin by discussing the benchmarks with multi-line completion targets in code repositories. The input code of the function completion dataset in RepoEval (Zhang et al., 2023) is truncated for text similarity-based retrieval, which cannot be parsed to get import information. Other benchmarks (Yu et al., 2024; Li et al., 2024) are constructed for function-level code generation, which aims to generate code based on natural language descriptions, as distinct from code completion. Since there is no available dataset for multi-line completion, we create a "multi-line" variant of the CrossCodeEval dataset (Ding et al., 2023) by expanding the references to three non-empty lines. In addition to the baselines, we evaluate two variants of DRACO including “w/o dataflow” as described in Section 5.3 and “w/ steps”, which completes the unfinished code line by line (i.e., three times “retrieval-then-generation"). The maximum generation length of LMs is increased to 100 tokens.
The experimental results with SantaCoder and CodeGen25 models shown in Table 15 reveal that: (i) DRACO still delivers a significant improvement to the zero-shot setting, and dataflow analysis plays a positive role. As discussed in the Limitations section, unclear code intent in multi-line completion would hurt the accuracy and weaken the guidance of dataflow analysis. (ii) RepoCoder iteratively retrieves similar code snippets in the repository to narrow the gap with the intended completion target, which is slightly better than DRACO but still strug-
gles with multi-line code completion. (iii) Completing line by line is not as effective as completing it all at once. The intermediate generated line may not point to cross-file imports in the dataflow, which would weaken the relevance of the retrieved content. In addition, progressive completion may lead to error accumulation. Note that the abnormal results of “w/ step” with SantaCoder are attributed to the model itself, which usually generates a comment or another function definition when completing a new line.
Recent empirical studies have shown that the completion of currently edited statement is more commonly used and more acceptable (Barke et al., 2023; Wang et al., 2023). In contrast, long and multi-line suggestions are usually at best dismissed out of hand and at worst distract the programmer away from their flow. Based on this practical evidence, we believe that our focus on single-line code completion is meaningful and sufficient. Future work could explore taking into account the dynamic changes of dataflow during the decoding process, e.g., automatically determining when it is necessary to re-retrieve to update the prompt.

heading: E Case Study

heading: E.1 Prompt Examples
text: We show the prompts generated by each method for the example unfinished code (see Figure 1). The prompts are excerpted for viewing the individual format, as shown in Figure 7. It can be observed that the prompts generated by DRACO look like natural code, which is in line with the training corpora of code LMs. The prediction result of each method using the CodeGen25-7B model is shown in Table 16, and only our DRACO generates the correct code line.

heading: E.2 Study on Failed Cases
text: Despite the superior performance achieved by DRACO, there are still many failed cases. Based on our observations, they are caused by three major reasons: (i) The definitions of completion targets may be truncated. Even though the correct definition is retrieved through dataflow analysis, the completion target may be truncated due to the limited input length, e.g., target member function in a long class definition. (ii) The used code LMs may be not powerful enough, which has already been demonstrated in Section 5.5. It is crucial for code LMs to capture data dependency relations in the provided background knowledge, where the prompts are usually long Python code. (iii) Unclear code intent may lead to wrong generation, which is a common sore point of code completion. For example, completing the first line of a new function is uncertain even to a programmer.
In Listings 1, 2, and 3, we show an example of reasons (i) and (iii) from CrossCodeEval. The ground truth is gen_begin_reuse(input_ids), while the prediction of DRACO with CodeGen25 is in_beam_search = False. Although the member function gen_begin_reuse is in the retrieved background knowledge, it is truncated to be invisible to LMs. Moreover, the comment # Start generation is not clear enough for completion, where both gen_begin and gen_begin_reuse look like rational choices.
''' # pyPhasesRecordloader.RecordSignal.RecordSignal @classLogger class RecordSignal:
# pyPhasesRecordloader.RecordSignal.RecordSignal.__init__ def __init__(self, targetFrequency=200, recordId=None):
self.recordId = recordId self.signals: List[Signal] = [] self.labelSignals = [] self.signalNames = [] self.targetFrequency = targetFrequency
# pyPhasesRecordloader.RecordSignal.RecordSignal.getSignalByName def getSignalByName(self, name) -> Signal:
index = self.getSignalIndexByName(name) return self.signals[index]
... '''
(a) CCFinder-*.
# Here are some relevant code fragments from other files of the repo: # -------------------------------------------------- # the below code fragment can be found in: # pyPhasesRecordloader-0.3.12/pyPhasesRecordloader/RecordLoader.py # -------------------------------------------------- # signalTypeStr = self.signalTypeDict[signalName] # else: # self.logError("Signal '%s' had no type when initilizing the RecordLoader" % str(signalName)) # signalTypeStr = "unknown" # return signalTypeStr
... # -------------------------------------------------- # the below code fragment can be found in: # pyPhasesRecordloader-0.3.12/pyPhasesRecordloader/RecordSignal.py ...
(b) RepoCoder, same as RG-1.
''' # pyPhasesRecordloader/Signal.py class Signal:
def __init__( self, name="Unknown", signal: np.ndarray = None, frequency=100, type=SignalType.UNKNOWN, typeStr="unknown" ) -> None: self.name = name self.signal = signal self.frequency = frequency self.type = type self.typeStr = typeStr
def setSignalTypeFromTypeStr(self): if self.typeStr in signalTypeDict:
self.type = signalTypeDict[self.typeStr] else:
self.type = SignalType.UNKNOWN ...
'''
(c) Our DRACO.
''' # pyPhasesRecordloader/Signal.py class Signal:
def __init__( self, name="Unknown", signal: np.ndarray = None, frequency=100, type=SignalType.UNKNOWN, typeStr="unknown" ) -> None: self.name self.signal self.frequency self.type self.typeStr
def setSignalTypeFromTypeStr(self): def getFilterCoefficients(self, tansitionWidth=15.0, cutOffHz=30.0, rippleDB=40.0): def bandpass(self, low, high, order=10): def lowpass(self, value, order=10): ...
'''
(d) DRACO with the definition scope.
, top_k , top_p , min_p , typical , num = 1): ...
'''
Listing 2: The truncated background knowledge retrieved by DRACO.
''' # generator.py class ExLlamaGenerator:
class Settings: ... sequence: torch.Tensor or None sequence_actual: torch.Tensor or
None settings: Settings beams: int or None max_beam_length: int in_beam_search: True disallowed_tokens: list[int] or None lora: ExLlamaLora or None def __init__(self , model , tokenizer ,
cache): ...
def reset(self): ... def make_rep_mask(self , penalty_max , sustain , decay): ... def batched_sample(self , logits , temperature , top_k , top_p , min_p , typical , num = 1): ... def sample_current(self , logits , num = 1): ... def sample(self , logits , temperature , top_k , top_p , min_p , typical , num = 1): ... def disallow_tokens(self , tokens): ... def gen_begin(self , in_tokens , mask = None): ... def gen_begin_empty(self): ... def gen_begin_reuse(self , in_tokens , mask = None): self.end_beam_search () ... def gen_feed_tokens(self , in_tokens , mask = None): ... def gen_accept_token(self , token): ... def gen_rewind(self , num_tokens): ... def gen_prune_right(self , tokens , mask = None): ... def gen_prune_to(self , min_tokens_to_keep , token_id , mask = None): ... def gen_prune_left(self , num_tokens , mask = None): ... def gen_num_tokens(self):
... def generate_simple(self , prompt ,
max_new_tokens = 128): ...
def apply_rep_penalty(self , logits): ... def gen_single_token(self , constraints = None , mask = None) : ... class Beam: ... def begin_beam_search(self): ... def beam_search(self): ... def end_beam_search(self): if not self.in_beam_search:
return self.in_beam_search = False
def replace_last_token(self , token , seq = False): ... def sequence_ends_with(self , tokens) : ...
'''
Listing 3: The complete background knowledge retrieved by DRACO.

heading: Future Work
text: Future work could explore taking into account the dynamic changes of dataflow during the decoding process, e. g. , automatically determining when it is necessary to re-retrieve to update the prompt