PaperlayerView exact revision on arXiv

Source-linked AI summary

Dataflow-Guided Retrieval Augmentation for Repository-Level Code Completion

Wei Cheng, Yuhan Wu, Wei Hu

arXiv:2405.19782v1cs.SEcs.CL

TL;DR

Pre-trained code LMs can lack the private-repository knowledge needed for correct completion, and existing retrieval based on imports or text similarity may not match the target. DraCo builds a repository context graph using extended dataflow analysis and retrieves relevant knowledge for code-LM prompts. It reports improved code exact match over RepoCoder, with limits for unclear intent and some completion settings.

  • Problem

    Pre-trained code LMs may miss private repositories’ naming conventions and styles, while prior retrieval based on imports or text similarity can be insufficiently relevant to completion targets.

    Source: §1 · p. 1

  • Method

    DraCo uses type-sensitive dataflow to build a repository context graph, retrieve relevant code entities, and assemble them with unfinished code into prompts for code LMs.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    Across settings, DraCo reports generally superior accuracy; against RepoCoder, it improves code exact match by 3.43% on average.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DraCo is described as plug-and-play across code LMs and efficient for real-time code completion.

    Source: §1 · p. 2

    Effectiveness may degrade when code intent is unclear, and dataflow guidance is weaker for new-line or function-body completion; the study focuses on currently edited-line completion.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Pre-trained code LMs struggle to supply private-repository conventions and relevant cross-file context, while similarity-based retrieval can miss the knowledge needed for a completion. DraCo uses type-sensitive dataflow to retrieve repository knowledge for prompts, and the paper reports improved accuracy and efficiency.

1 Introduction

  • Code LMs remain blind to private repositories’ naming conventions and styles, while similarity-based retrieval can fail when relevant code snippets are absent.The newSignal example shows a plausible but invalid attribute choice; following its data origin reveals its type and correct function.Source: §1 · pp. 1–2
  • DraCo indexes repository code entities and their dataflow relations in a context graph, retrieves relevant entities for unfinished code, and assembles them into prompts for code LMs.Its extended dataflow analysis sets type-sensitive dependency relations.Source: §1 · p. 2
  • Experiments report generally superior accuracy across settings and describe DraCo as plug-and-play across code LMs and efficient for real-time completion.The paper also introduces ReccEval, a Python dataset with diverse completion targets collected from PyPI.Source: §1 · p. 2
  • 3.43% code exact match, 3.27% identifier F1-score, and 100× prompt generation time improvements are reported over RepoCoder on average.The contributions also include the ReccEval Python dataset with diverse completion targets.Source: §1 · p. 2

2 Related Work

Related work spans code-completion methods using cross-file context and retrieval augmentation, while noting overhead in prior repository-level approaches. This paper applies dataflow-guided retrieval to repository completion.

2 Related Work

  • Repository-level completion methods include jointly learned in-file and cross-file context, code search, trained prompt construction, and iterative retrieval-generation.These approaches contrast with single-file completion and address repository context in different ways.Source: §2 · p. 2
  • Prior repository-level methods are limited by the high overhead of extra training or iterative generation.Source: §2 · p. 2
  • RAG integrates retrieved knowledge with frozen pretrained LMs, and this work uses dataflow analysis to mine dependency information for repository-level completion.Source: §2 · p. 3

3 Methodology

DraCo focuses on Python completion of the currently edited line, using extended dataflow analysis to build a repository context graph and retrieve entities for structured prompts. Its prompt construction allocates context, prioritizes relevant entities, and offers two entity-detail scopes.

3 Methodology

  • DraCo focuses on Python completion of the currently edited line, with only code before the cursor visible as in an IDE.Source: §3 · p. 3

3.1 Dataflow Analysis

  • Dataflow analysis builds a graph of variable dependencies; DraCo adds five typed relations—assigns, as, refers, typeof, and inherits—to indicate variable type or origin.The design assumes background knowledge relevant to variable types is crucial for completion.Source: §3.1 · p. 3

3.2 Repo-specific Context Graph

  • AST traversal indexes modules, classes, functions, and variables, linking them through contains relations and dataflow-derived depends relations, including across imported modules.Source: §3.2 · p. 4

3.3 Dataflow-Guided Retrieval

  • For an unfinished line, DraCo expands import information through refers relations, maps it to graph entities, and follows depends links to retrieve type-related definitions rather than perform precise Python type inference.Source: §3.3 · p. 4

3.4 Prompt Generation

  • Prompt construction allocates half the input length to background knowledge and half to unfinished code, reallocating unused tokens and truncating overlong unfinished code to its last tokens.Source: §3.4 · p. 4
  • Entities linked to the completion line are prioritized, other imports are added until the token limit, and merged, ordered module context is placed before unfinished code.Definition prompts contain entity definitions; complete prompts also include function bodies and variable statements.Source: §3.4 · pp. 4–5

4 Experiment Setup

The experiments use repository-level Python completion datasets, compare retrieval-augmented methods across popular LMs, and assess accuracy and prompt-generation efficiency.

4.1 Datasets

  • The Python subset of CrossCodeEval requires cross-file API use, while ReccEval adds more diverse completion targets.Table 1 reports dataset statistics, including token counts calculated with the StarCoder tokenizer.Source: §4.1 · p. 5, Table 1 · p. 5

4.2 Implementation Details

  • Evaluation covers non-training retrieval-augmented methods across seven popular LMs, timing prompt generation and taking the first completion line as the prediction.Compared methods include Zero-Shot, CCFinder-k, RG-1, and iterative RepoCoder.Source: §4.2 · p. 5, Table 2 · p. 6

4.3 Evaluation Metrics

  • Metrics are code exact match (EM) and edit similarity (ES), identifier exact match (ID.EM) and F1-score, plus prompt generation time.The study reports the average of each metric.Source: §4.3 · p. 6

5 Experimental Results and Analysis

The experiments compare DRACO with other repository-level completion methods across two datasets, then examine efficiency, retrieval ablations, prompt scopes, adapted LMs, and GPT models. DRACO generally improves accuracy, with results varying by dataset and model context capacity.

5.1 Performance Comparison

  • Against RepoCoder, DRACO improves average EM by 3.43%, ES by 1.00%, ID.EM by 3.62%, and F1 by 3.27%.CodeGen-350M with DRACO outperforms zero-shot StarCoder-15.5B, though it is slightly worse than RepoCoder on ReccEval code-match metrics; DRACO’s lead is more significant on CrossCodeEval, whose dot-ended targets suit definition retrieval, whereas ReccEval’s partial API names favor similarity retrieval.Source: §5.1 · p. 6, Table 3 · p. 7, Table 4 · p. 7

5.2 Efficiency Evaluation

  • DRACO is efficient for real-time IDE completion and is more accurate than methods with comparable efficiency, while RepoCoder takes more than 4 seconds even with CodeGen-350M.Parsing unfinished code into an AST or dataflow graph is slightly slower than RG-1’s text-similarity retrieval but remains comparable.Source: §5.2 · pp. 6–7, Table 5 · p. 7

5.3 Ablation Study

  • The complete DRACO performs best, and each dataflow component contributes positively; without dataflow, its accuracy remains comparable to CCFinder.The intra-file dataflow component contributes less on StarCoder, whose 8K-token context weakens its truncation-prevention benefit when relevant knowledge is placed first.Source: §5.3 · p. 7, Table 6 · p. 8

5.4 Analysis of Prompt Scopes

  • Complete-scope prompts perform best, while definition-only prompts outperform CCFinder and RG-1 in most settings; the two scopes offer an accuracy-cost choice.Complete scope includes function bodies and variable statements, unlike definition scope.Source: Figure 3 · p. 7, §5.4 · p. 8

5.5 Analysis of Adapted LM

  • As maximum input length increases, DRACO’s accuracy decreases with Code Llama but increases with StarCoder.The authors attribute Code Llama’s decline to difficulty capturing dependencies in long Python code and suggest specialized code LMs may better handle longer contexts.Source: §5.5 · p. 8, Figure 4 · p. 8

5.6 Analysis of GPT Models

  • On 50 CrossCodeEval examples, DRACO also significantly enhances GPT-model accuracy, but evaluation faces output-format problems, limited room for demonstrations, and nearly $35 in API costs.The authors caution that output-format errors may bias evaluation, especially for RG-1 with GPT-4.Source: §5.6 · pp. 8–9, Table 7 · p. 8

6 Conclusions

The conclusion presents DRACO as a dataflow-guided retrieval approach for repository-level code completion. Its experiments on CrossCodeEval and ReccEval show superior accuracy and applicable efficiency, while future work will explore other structured code representations.

6 Conclusions

  • DRACO extends dataflow analysis with type-sensitive data-dependency relations to guide more precise retrieval.It indexes private repositories into context graphs and retrieves relevant background knowledge to form prompts for code LMs.Source: §6 · p. 9
  • Experiments on CrossCodeEval and ReccEval demonstrate DRACO’s superior accuracy and applicable efficiency.Source: §6 · p. 9
  • Future work will explore other structured code representations.Source: §6 · p. 9

Ethical Considerations

The paper cautions that retrieval augmentation mitigates, but does not eliminate, risks in generated code. It also reports that the datasets and models used are publicly available under permissive licenses.

Ethical Considerations

  • Generated code may contain non-existent APIs or potential bugs, and retrieval-augmented approaches mitigate these risks only to some extent.The authors recommend presenting retrieved background knowledge to programmers for review when deploying the approach.Source: Ethical Considerations · p. 9
  • The datasets and code LMs used in the work are publicly available with permissive licenses.The named licenses include Apache-2.0 for CrossCodeEval and CodeGen, and MIT, Apache, or BSD licenses for ReccEval repositories.Source: Ethical Considerations · p. 9
  • SantaCoder and StarCoder use the BigCode OpenRAIL-M v1 license, while Code Llama is governed by the Meta license.Source: Ethical Considerations · p. 9

Limitations

DRACO’s effectiveness has three stated boundaries: dependence on code-LM capabilities, unclear code intent, and migration work for other languages.

Limitations

  • DRACO depends on code LMs that support long inputs and capture data dependencies in background knowledge.Experiments still found considerable improvement with CodeGen-350M at 2K tokens, mitigating this limitation.Source: Limitations · p. 9
  • Effectiveness may degrade when intent is unclear, as in new-line or function-body completion where dataflow provides weaker import priorities.The paper focuses on currently edited lines and suggests exploring dataflow across completion scenarios.Source: Limitations · p. 9
  • Migration beyond Python requires extending dataflow analysis to suit each target language’s characteristics.Static-analysis tools for common languages may facilitate multilingual implementations.Source: Limitations · p. 9

A Dataflow Graph Construction

The appendix describes how Python code is parsed into a dataflow graph and gives the time complexities of the background-knowledge generation algorithms.

A Dataflow Graph Construction

  • Python AST parsing extracts variables and type-sensitive relations as triplets to form the DFG.The example identifies identifiers and attributes, represents parameter types and assignments as relations, and prunes a type-insensitive parameter relation; figure labels map variables to line numbers and edge initials to relation names.Source: A Dataflow Graph Construction · p. 13

B Generation of Background Knowledge

C.1 Details of Dataset Construction

The appendices detail ReccEval construction and evaluation procedures, then report additional performance, efficiency, and import-context comparisons.

C.1 Details of Dataset Construction

  • ReccEval samples Python projects released on PyPI from 2023-01-01 to 2023-04-28, filters project size and licenses, and removes examples with exact matches from StarCoderBase-1B.It selects local-import usages and subsequent-token cursor positions, limits examples per file and repository, and ensures references are absent from unfinished code.Source: §C.1 · p. 13

C.2 Implementation Details of Baselines

  • The baseline setup reproduces CCFinder and uses sparse bag-of-words Jaccard retrieval for RG-1 and RepoCoder.Sliding-window line length and step are 20 and 10; the retrieval snippet cap is 40 for StarCoder and 10 for other models, and RepoCoder uses two iterations.Source: §C.2 · p. 14

C.3 Details of Used LMs

  • The authors warn that Code Llama and GPT analyses may face data leakage because their training data overlap or are unknown.Code Llama training data spans January–July 2023, while GPT training data is unknown and updated.Source: §C.3 · p. 15

C.4 Details of Evaluation Metrics

  • Evaluation measures code exact match and edit similarity, identifier exact match and F1, and prompt-generation time excluding LM prediction generation.Prompt time includes retrieval and prompt assembly; identifier F1 is computed from extracted identifier sets.Source: §C.4 · p. 15

D.1 More Performance Comparison Results

  • Across additional CodeGen models, performance improves with model size, and DRACO’s gains for zero-shot LMs increase with model capability.The authors infer that stronger LMs better use DRACO’s retrieved background knowledge.Source: §D.1 · pp. 15–16, Table 9 · p. 16, Table 10 · p. 16

D.2 More Efficiency Evaluation Results

  • DRACO’s repository preprocessing is 3–6 times faster than RepoCoder’s, while preprocessing time grows linearly with repository size.The authors note that this growth may create scalability challenges for RG-1 and RepoCoder.Source: Table 11 · p. 16, §D.2 · p. 16
  • RepoCoder’s average prompt-generation time exceeds 3 seconds even with SantaCoder, and its efficiency mainly depends on LM generation time for its second retrieval.The comparison across A800 and RTX 4090 runs is not head-to-head for RepoCoder.Source: §D.2 · pp. 16–17, Table 12 · p. 17, Table 13 · p. 17

D.3 Effect of Other Import Statements

  • DRACO consistently outperforms the variant that adds other local-import entities only when the primary prompt is empty.The difference is especially pronounced on StarCoder-15.5B, which can handle longer code context.Source: §D.3 · p. 17

D.4 Evaluation of Multi-Line Completion

The appendix tests DRACO on multi-line completion and illustrates the prompts and prediction for one example. DRACO helps in the multi-line setting, while the example comparison shows only DRACO producing the correct line.

D.4 Evaluation of Multi-Line Completion

  • The authors create a three-line CrossCodeEval variant because they report no available dataset for multi-line completion.They compare DRACO with and without dataflow, line-by-line completion, and retrieval-iterative RepoCoder.Source: §D.4 · p. 17
  • On the multi-line variant, DRACO significantly improves the zero-shot setting and dataflow analysis plays a positive role, while RepoCoder is slightly better.RepoCoder still struggles with multi-line completion.Source: §D.4 · pp. 17–18, Table 15 · p. 18

E.1 Prompt Examples

  • For the illustrated CodeGen25-7B example, DRACO’s prompt resembles natural code and it alone produces the correct completion line.The prompt formats are shown in Figure 7 and predictions in Table 16.Source: Table 16 · p. 18, §E.1 · p. 18

D.4 Evaluation of Multi-Line Completion

  • Completing multi-line code line by line is less effective because intermediate lines can weaken retrieval relevance and progressive completion can accumulate errors.The authors also report abnormal SantaCoder step-variant results because the model may generate a comment or another function definition when completing a new line.Source: §D.4 · p. 18

E.2 Study on Failed Cases

The authors attribute remaining failures to target truncation under input limits, code LMs’ difficulty capturing dependencies in long prompts, and unclear completion intent.

E.2 Study on Failed Cases

  • Retrieved definitions can be truncated by input-length limits, hiding even correctly retrieved completion targets.In the example, the retrieved gen_begin_reuse function is truncated and invisible to the model.Source: §E.2 · p. 18
  • The code LMs used may be too weak to capture data dependencies in the long Python prompts.Source: §E.2 · p. 18
  • Unclear code intent can lead to wrong generation, including when a new function’s first line is uncertain even to a programmer.In the example, both gen_begin and gen_begin_reuse are rational choices given the comment.Source: §E.2 · p. 18

Related papers

Loading 2405.19782v1…