Fold2Vec: Towards a Statement-Based Representation of Code for Code Comprehension

Premise. Code comprehension is one of the most challenging and time-consuming tasks in software development. Often code is poorly documented, and its functionalities are tangled and scattered throughout the whole codebase. This makes tasks like debugging and maintenance difficult and time consuming. Therefore, the need for tools to support the developer in code comprehension is becoming critical. This need is also reflected by recent research trends where the focus is on automatic code generation [8, 46, 48, 55], automatic code completion [52, 66], code summarization [1, 5, 6, 7, 65], and automatic test generation [16, 23].

Tangled intents. One aspect that renders code comprehension difficult is that the code of a single functionality is often polluted by the code of other functionalities. We refer to this phenomenon as tangled intents. This is where most of the models based on code comprehension have limits. Comprehension occurs at coarse grain usually set at the method/function level [1, 5, 6, 7, 65], which limits their applicability. Let us demonstrate the problem on an example involving intent identification. Figure 1 shows the method whose behavior we have to guess by giving it a proper description. Even if the code is quite short, it suffers from the preceding entangled intents—the code belonging to different intents has a different colored background. To tell apart and name the code of every single intent is difficult and often ambiguous. Moreover, a wrong classification presents obstacles to proper reuse, refactoring, and maintenance, among others.

Motivation. Models like code2seq [5] and code2vec [7] have demonstrated that training a neural network (NN) with specifically designed representations for code helps to achieve better results with respect to standard natural language processing techniques. However, both code2seq and code2vec are bounded by their representation, leaf-to-leaf abstract syntax tree (AST) path. In our view, a fine-grained prediction that considers potentially non-consecutive statements should be designed to deal with disconnected sub-trees of the AST. In these cases, leaves from different sub-trees cannot be connected by leaf-to-leaf paths. This limitation leads to shorter and narrower paths, for which it has been shown to hurt the performance [6]. These motivations lead us to design a novel representation that aggregates all of the information available from the different sub-trees. This information can be used to gather statements into coherent groups. These groups, called intents, can be used to detect bad code practices.

Research questions. To address the mentioned issues and limitations, we try to answer the research questions:

Statement-based representation. Kiros et al. [45] and Dwivedi and Shrivastava [24] build a hidden representation aligned with its semantics for each natural language phrase of the document. All of the hidden representations contribute together to perform a specific task on the whole document. In our view, a similar approach could be exploited to improve code comprehension and provide an answer to RQ₁ where the analogous of a natural language phrase is a code statement. Our NN, called fold2vec, builds a hidden representation for each code statement separately from the others. Then, the combined hidden representations contribute to the classification of the method. In this direction, our main contribution is a novel code representation. Given a code snippet, its AST is considered. We split the AST into sub-trees, each representing a single code statement. Each sub-tree is folded into two sequences of tokens—one for the terminals and one for the non-terminals—via a pre-order visit. All token sequences are then fed to the NN, and they contribute to the classification.

Code summarization and code search. To answer RQ₁, we show that our code representation improves the state of the art for the code summarization and code search tasks. It is worth noting that for code summarization, the dataset used to perform this task exploits the fact that method names often summarize the method behavior. This assumption can be seen as a rough approximation, but it permits to access a vast amount of available data on which the training can be performed. Usually, more data lead to better and more robust results [29].

Statement separation. An intent can be any subset of statements, not necessarily adjacent in the code, pursuing the same goal. In our view, intent identification could be possible by applying a NN specifically trained for the task of statement separation. This NN can be used to induce a metric over statements regarding the statement topics, meaning that statements with the same topic would be close to each other, and statements with different topics would end up far apart from each other. With an hierarchical clustering algorithm, the induced metric can be used to cluster statements that share the same topic. To answer RQ₂, we train fold2vec for the statement separation task. We call this model stmt-fold2vec. However, the novelty of the task limits the availability of baselines. To address this issue, we adapt the previously mentioned work (code2vec) for the statement separation task. We demonstrate that the novel code representation leads to improvements also for this task.

Overview. The rest of the article is organized as follows. Section 2 introduces the terminology. Section 3 illustrates the design of our source code representation and the NN architecture used on top of it. Section 4 discusses and compares fold2vec results to the state-of-the-art models for the code summarization and code search. It also brings some evidence that fold2vec can classify single or groups of statements to enable intent identification. We present our ablation study in Section 5 and discussion in Section 6. Finally, in Sections 7 and 8, we examine some related work and draw our conclusions, respectively.

The aim of this section is to set a common terminology used in this article by recalling the definition of some core concepts as NN layers and AST.

Neural network. A NN is “a computing system made up of a number of simple, highly interconnected processing elements, which process information by their dynamic state response to external inputs [21]”. Usually, NNs are organized in layers. In a supervised learning environment, the NN is trained through gradient descent techniques to minimize a loss from a set of examples called the training set. Additionally, a validation set and test set are used to evaluate the NN in a developing phase and at the end of the process, respectively. These phases are meant to limit the effect of overfitting.

Autoregressive architecture. A neural architecture is said to be autoregressive when each prediction is part of several single classifications, each based on the previous ones. For example, while predicting how a text continues, an autoregressive architecture predicts the next token by considering also the previous predictions.

Pre-trained and end-to-end models. A model is said to be pre-trained when it was already trained on a specific task. These models can exploit a task for which there is a high availability of data to learn intrinsic features of the media (text, code, images, videos, etc.) that can be useful in a wide range of tasks. The learned features can then be reused and fine-tuned for tasks of which there is less availability of data. Pre-trained models are usually trained on extremely large datasets that require expensive hardware during training. In contrast, end-to-end models are trained from scratch directly on a task of interest. These models are usually smaller, customizable, and easily reproducible. Using end-to-end models is convenient when the task of interest already has a large dataset available.

Multi-layer perceptron. The multi-layer perceptron (MLP) is the composition of several dense layers. A dense layer applies to a real vector \(\vec{x}\in \mathbb {R}^{1\times n}\) the transformation:where \(\mathbf {W} \in \mathbb {R}^{n\times m}\) and \(b \in \mathbb {R}^{m}\) are trainable parameters randomly initialized. Usually, a non-linear activation function is applied to the output \(\vec{y}\) (e.g., \(\vec{z} = \tanh (\vec{y})\) ).

Embeddings layer. The embeddings layer maps tokens to real vectors. This kind of layer is a matrix \(\mathbf {E}\in \mathbb {R}^{t\times d}\) , where \(t\) is the number of tokens to be considered and \(d\) is the size of each embedding. Pennington et al. [60] showed that values for \(d\) between 50 and 500 are a good fit in most cases. Each token, indexed by \(i\) such that \(0\le i\lt t\) is mapped to the \(i\) -th row of \(\mathbf {E}\) . \(\mathbf {E}\) is made up of trainable parameters randomly initialized. \(\mathbf {E}\) can be either pre-trained as in the work of Mikolov et al. [51] and Pennington et al. [60] or learned during network training like in the work of Alon et al. [5, 7].

Layer normalization. Layer normalization [9] is a regularization technique used to improve the network training time and test error. Apart from the input layer, each layer learns its input distribution from the output of the preceding layer. During training, the output distribution from a layer fed to another is constantly changing as the NN is updated; this is known as the covariate shift problem. Given a set of features, layer normalization computes the mean and standard deviation, then re-centers and re-scales the features using an extra normalization parameter. By normalizing the distribution, the covariate shift is reduced and the network learns faster.

Recurrent neural network. A recurrent neural network (RNN) is the neural architecture commonly used when processing input sequences. It processes each element of the input sequence considering the output of the preceding computation. Given the input sequence \(\vec{x}_1, \vec{x}_2, \ldots , \vec{x}_n\) , it computes the \(t\) -th stage aswhere, \(\vec{h}_{t-1}\) is a hidden state and \(\theta\) represents the network parameters. The hidden state is a sort of network memory updated as an input element is processed. RNN training can suffer from both the vanishing and the exploding gradient problem as explained in [12, 59]—that is, parameter updates can become either too small or too big, respectively.

Long short-term memory. Long short-term memory (LSTM) [26] is a type of RNN that deals with the gradient problem [59] so that the network can make long-term associations. Each layer can also use two LSTMs at the same time, where one consumes the input sequence and the other the reversed input sequence; this configuration is called bidirectional long short-term memory (BiLSTM). As well, an additional LSTM layer can consume the output of a previous LSTM layer; this configuration is called stacked LSTM.

Attention mechanism. An attention mechanism helps the network focus on the input portion that should be more relevant for its task. For example, articles are less informative than nouns when you are trying to understand a sentence. The attention mechanisms used in this article are self-attention (SA) [72] and global attention (GA) [47]:

Abstract syntax tree. An AST is a unique tree-shaped representation of a program. Formally, it is a quintuple \((N,T,X,s,\delta ,\psi)\) , where the following apply:

In this section, we present the proposed model, named fold2vec, and the steps needed to train and test it. In Section 3.1, we explain how to extract the features to feed the NN. In Section 3.2, we describe how the neural architecture is composed and how its components interact. In Section 3.3, we discuss the tuning process of fold2vec.

Introduction. To evaluate the proposed representation with the proposed architecture, we consider three different classification tasks: code summarization, statements separation, and code search. The first two tasks are evaluated using datasets derived from a collection of GitHub projects. As we will show, models trained for these tasks can be combined to build tools that could improve code readability. The third one uses a dataset collected by Husain et al. [36]. We evaluate our proposal against several baselines. Among these, several are completely replicated in our study to achieve a fair comparison. The only difference between the replicated baselines and the original ones is the used deep learning framework (originally TensorFlow,¹ and here, PyTorch²).

Hardware setup. All deep learning models are trained, evaluated, and tested on a single T4 NVIDIA GPU. This ensures complete and free reproducibility using services such as Google Colab³ or Kaggle Kernels.⁴ We provide a Google Colab notebook with the necessary steps to reproduce our study.⁵

Baselines. To achieve a fair comparison, we focus on baselines that fall in these categories:

Introduction. In this section, to answer RQ₃, we are going to evaluate how the model fold2vec behaves with and without some of its parts. Moreover, we are going to study the dependence of fold2vec from code identifiers and how identifiers relate to the method name.

Identifier dependence. First, let us study the contribution of the two types of information (terminals and non-terminals) fed to fold2vec. Let us call term2vec the fold2vec variant stripped of the architecture path from ⑦ to ⑩ of Figure 4 and non-term2vec the fold2vec variant stripped of the architecture path from ③ to ⑥ of Figure 4. Table 6 summarizes the result. term2vec achieves a similar result to the complete architecture, whereas non-term2vec falls behind by several percentage points (–32.18%, –23.52%, and –27.22% with respect to precision, recall, and F1-score). This fact should not come as a surprise. Developers put a lot of effort into using relevant identifiers, and often one can guess the method behaviors just by looking at them. Moreover, without identifiers, all method calls look alike. Therefore, the information on what a sub-call is doing is lost. If we test non-term2vec on methods that do not perform other method calls (which are 120,551 samples), we can even raise the scores to 37.49%, 27.93%, and 32.01% with respect to precision, recall, and F1-score. Alongside method calls, information about types and reused variables are also lost, which explains the ongoing low result.

Method name token dependence. The previous paragraph showed and explained the dependence of fold2vec from code identifiers. However, not all identifiers are equal. In particular, fold2vec is dependent on a specific kind of identifier, those that appear in method names. Let us consider only methods that have among their code identifier tokens that appear in the method name itself. By doing so, we can build a subset of the test split composed of 155,030 samples. fold2vec evaluated on the set scores 67.04%, 48.32%, and 56.16% with respect to precision, recall, and F1-score, which represent an improvement of around 5 percentage points.

Positional embeddings. Let us consider the effect of positional embeddings (PEs) within statements (see Table 6). Notably, most of the time, PEs do not yield improvements and when they do, the improvements are marginal. This fact suggests that positional information is not a necessary component for the CUIRs. For non-term2vec, this should not come as a surprise since CUIRs are generated with a BiLSTM; however, the same appears to be true for term2vec and fold2vec.

LSTM variants. Next, consider Table 7. We used increasingly bigger LSTM layers processing non-terminals. Namely, we trained fold2vec using an LSTM, a BiLSTM, and a 2LayerBiLSTM for non-terminals. The best scores are achieved using a single layer of BiLSTM (51.34% with respect to F1-score). The 2LayerBiLSTM leads to almost identical results (51.31% with respect to F1-score). This study suggests that adding more BiLSTM layers does not lead to better results.

Terminals and non-terminal processing variants. To validate the different choices of using different architectures for terminals and non-terminals, consider Table 8. We evaluated three additional variants of fold2vec: fold2vec+GA|GA uses GA instead of the BiLSTM, fold2vec+BiLSTM|GA has the architecture for terminals swapped with the architecture for non-terminals, and fold2vec+BiLSTM|BiLSTM uses a BiLSTM instead of GA. The most noticeable result is that swapping architecture has a detrimental effect. When exclusively using BiLSTMs and GAs, a mild version of the same effect seems to appear. These facts seem to suggest that terminals and non-terminals can benefit from different architectural paths of execution, each leveraging specific properties of the CUs.

Not all methods are equal. Not all methods are equally easy to classify. In fact, some methods are much easier than others. Among these methods, we have getters and setters. Removing these methods from the test set has negative effect on the results: –.11.35%, –9.65%, and, –10.57% with respect to precision, recall, and F1-score. This is mainly caused by two factors: the intrinsically hardness of getters and setters (which can be considered low), and the abundance of these methods in the dataset (which is high for java-large).

In this section, we discuss strengths and shortcomings of the considered approaches and the considered tasks.

In this section, we examine approaches that share the same goals, techniques, or application domain as fold2vec.

Code summarization. In the non-autoregressive family of models, we can enumerate code2vec [7] and Paths+CRFs [6], which we have already discussed extensively. Autoregressive models produce a sequence of predictions where each prediction is based on the previous one. Even if these are promising approaches, they better fit when the classification can be split into multiple steps. code2seq [5] and ConvAttention [3] fall into this category. The former is, once again, based on the leaf-to-leaf path representation, whereas the latter is based on source code tokenization. Both use NNs to achieve their purpose: attention based and convolution based, respectively. Although we focused on deep learning methods, there are other techniques using more traditional approaches. For example, Hellendoorn and Devanbu [30] use an n-gram-based language model. Raychev et al. [64] use support vector machines constrained to predict only few method names. Wang et al. [74] train an autoregressive model using reinforcement learning using two different datasets for Python and Java code.

Intent identification. Since intents and concerns can be overlapping concepts, in this section we present some work on concern identification. Kästner et al. [42] proposed a semi-automatic tool for extracting features from code. It relies on a configuration given from a domain expert to correctly recognize the features. In the same research area, Valente et al. [71] proposed an approach to semi-automatically annotate optional features. UML model variants have been used to automatically identify model-based product lines in the work of Martinez et al. [50]. Qu and Liu [62], Breu and Krinke [13], and Tonella and Ceccato [70] proposed approaches based on program tracing to automatically mine concerns. Moldovan and Şerban [53] proposed a clustering-based approach. All of these are semi-automatic or manually driven approaches to feature identification, whereas fold2vec is automatic; however, apart from this aspect, they share the same goal.

Machine learning on code. Allamanis and Sutton [4] presented an early analysis of source code based on n-gram language models. Their results have been used to build a Java-source code dataset that has been used as the basis to build the java-large dataset we use. From the same research group, graph NNs were applied for code generation and representation learning [2, 14]. The task of code generation is tackled in the work of Hu et al. [33] using transformer-generated representation from AST paths. For the same task, a massive transformer language model developed in the work of Chen et al. [17] achieved outstanding results. CodeBERT uses a large transformer language model for the task of mask language modeling [25]. Pre-trained models of the size of CodeBERT usually have high training costs. However, these models can be generalized to several tasks through fine-tuning with lower cost (both in terms of data and hardware). Usually, these large models are resource heavy and are accessible through online APIs. The main benefit of this approach is that it leads to better results. Again, for representation learning, Wang et al. [75, 76] used graph NNs and program traces to achieve the same scope. Instead, Ben-Nun et al. [11] used recurrent architectures trained from features extracted from the LLVM representation of source code. Raychev et al. [65] exploited CRF on a dependency network built from JavaScript code to automatically predict program properties such as type annotation and variable identifiers. Their approach led to JSNice, which predicts correct name identifiers in \(63\%\) of the cases and correct type annotations in \(81\%\) of the cases. Hu et al. [34], Chen and Zhou [18], and Movshovitz-Attias and Cohen [54] developed autoregressive models to automatically generate comments from source code. Iyer et al. [37] used attention networks to generate natural language descriptions from code, and Haiduc et al. [28] shared the same goal. Jiang et al. [40] used NN to automatically generate commit messages from diffs. Piech et al. [61] translated programs into real valued embeddings. Chen et al. [19] used tree decoders and encoders to translate programs from one language to another. Oda et al. [57] tried to translate formal code into pseudo code. Another, interesting research field applies machine learning techniques to spot bugs, as in the work of Dam et al. [22] and Shi et al. [68]. Related to bug detection, much work focuses on program repair. For example, Jiang et al. [39] used a transformer-based architecture (GPT) to pre-train a large language model for the task. For the same task, a recurrent neural architecture was used by Wang et al. [73]. All of these approaches exploit code representation and machine learning, and represent potential application domains for fold2vec. A different type of code representation based on code updates was developed by Hoang et al. [32]. Last, Kang et al. [41], Keim et al. [43], and Rabin et al. [63] assessed generalizability of several deep learning baselines as code2vec, code2seq, and CodeBERT.

This work introduced a source code representation to be used with NNs that differs from traditional approaches for its granularity. On such a representation, we developed the fold2vec model and compared it to the state-of-the-art techniques for the task of code search, code summarization, and statements separation. Such a representation could also bring an improvement in other code tasks, such as automatic defect detection and automatic comment generation. We addressed and answered RQ₁ by showing improvements on the task of code summarization, code search, and statements separation. Moreover, by addressing RQ₂, we showed that a NN trained on the statements separation task can be used to measure the code quality with respect to the SRP. To answer RQ₃, we evaluated several variants of fold2vec . It appears that the different properties of terminal and non-terminal can potentially benefit from tailored architectures. Both the model and data needed to reproduce our experiment are available at the following URL:

https://cazzola.di.unimi.it/fold2vec.html .

Additionally, we provide a Google Colab notebook to replicate our experiments:

https://colab.research.google.com/drive/1y383wyfNemYO7QYlmp7Nh7L_IHSMPvo4?usp=sharing .

In the future, we will move fold2vec to the autoregressive model family and compare it to code2seq . We must note that other valid directions are positional-like embedding that encode the parent-child relations between tokens [58]. In addition, Tree-LSTMs [69] could be applied for both inter-statements and intra-statements. Another available approach is to introduce an auto-encoder for statements to be either pre-trained or trained altogether with fold2vec .

The authors wish to thank Uri Alon for his precious suggestions and cooperation and Sudipto Ghosh for his thorough proofreading and comments.