Improving AST-Level Code Completion with Graph Retrieval and Multi-Field Attention

Code completion, which provides code suggestions by generating code snippets or structures, has become an essential feature of integrated development environments (IDEs). Recently, some studies have begun to use graph neural networks to complete AST-level code, and shown that it is promising to introduce GNNs into AST-level completion. However, these methods do not fully exploit the potential of reference codes with similar structures nor solve out-of-vocabulary (OOV). We propose Retrieval-Assisted Graph Code Completion (ReGCC) to enhance AST-level code completion further. ReGCC integrates a retrieval model that searches for similar code graphs to generate graph nodes and a completion model that leverages information from multiple domains. The key component of both the retrieval and completion models is the Multi-field Graph Attention Block, which consists of three layers of stacked attention: (1) Neighborhood Attention: preserves the heterogeneity and local dependency of the graph, enabling nodes to exchange information within their neighborhood. (2) Global & Memory Attention: addresses the long-distance dependency problem by providing nodes with a global view and the ability to extract information from the memory domain. (3) Reference Attention: lets nodes obtain valuable information from structurally similar reference code graphs. Furthermore, we tackle the OOV issue by employing feature matching and copying values from existing nodes. Specifically, we predict edges between nodes beyond the vocabulary, enabling effective information transfer. Experimental results demonstrate the superiority of our approach over state-of-the-art AST-level completion methods and generative language models.