Code Completion with Assistance from Token-Level Type Prediction

Code completion is a pivotal feature within Integrated Development Environments (IDEs), which enhances development efficiency by predicting the next code token.The majority of current code completion techniques forecast the value based on Abstract Syntax Trees (ASTs) and contextual source code. However, they often model relative positional information and type information independently, lacking an effective framework for their unified modeling, which consequently limits the ability to characterize syntactic constraints and structural dependencies of code. To address this problem, we propose TAVP, a token-based type prediction-assisted code completion method. This model jointly models leaf node values, node types, and root-path structural information within a unified framework, and introduces Rotary Positional Embedding (RoPE) into the attention mechanism to capture the relative positional relationships among leaf nodes. During the inference stage, the model modulates the injection intensity of type semantic constraints based on the correctness and confidence of type prediction, thereby effectively assisting in the prediction of leaf node values. Experimental results demonstrate that TAVP outperforms baseline models in terms of Top-k accuracy and Mean Reciprocal Rank (MRR) on Python and JavaScript datasets. Furthermore, ablation experiments confirm each essential component's efficacy. Keywords:Code Completion; Type Assisted; Rotary Position; Abstract Syntax Trees