C ODE T REK : F LEXIBLE M ODELING OF C ODE USING AN E XTENSIBLE R ELATIONAL R EPRESENTATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Designing a suitable representation for code-reasoning tasks is challenging in aspects such as the kinds of program information to model, how to combine them, and how much context to consider.
We propose CodeTrek, a deep learning approach that addresses these challenges by representing codebases as databases that conform to rich relational schemas.
The relational representation not only allows CodeTrek to uniformly represent diverse kinds of program information, but also to leverage program-analysis queries to derive new semantic relations, which can be readily incorporated without further architectural engineering.
CodeTrek embeds this relational representation using a set of walks that can traverse different relations in an unconstrained fashion, and incorporates all relevant attributes along the way.
We evaluate CodeTrek on four diverse and challenging Python tasks: variable misuse, exception prediction, unused deﬁnition, and variable shadowing.
CodeTrek achieves an accuracy of 91%, 63%, 98%, and 94% on these tasks respectively, and outperforms state-of-the-art neural models by 2-19% points.
I NTRODUCTION
Deep learning techniques are increasingly applied to code-reasoning tasks, including bug detection (Allamanis et al., 2018), type inference (Hellendoorn et al., 2018), code summarization (Alon et al., 2019b), program repair (Dinella et al., 2020), and code generation (Alon et al., 2020), among many others.
The successful application of these techniques to a given task depends heavily on the program representation that encompasses relevant program features and model architecture.
There are many crucial choices involved in designing a suitable representation for a new task.
Consider the example in Figure 1 which depicts an instance of an exception-prediction task (Kanade et al., 2020), whose goal is to predict the exception type in the placeholder “ [??] ” in the highlighted statement.
Predicting the correct type AssertionError from a range of 20 pre-deﬁned exception types in Python requires understanding the implementation of the check() function, which is deﬁned in a different class ( TestObject ).
Thus, the desirable context for the model goes beyond the immediate lexical neighborhood of the placeholder, possibly inside a chain of called functions.
The richness of information and extended scope inevitably imply that the relevant context may be very large.
Models tackling such tasks are challenged to either reduce the scope of a task—e.g., a single function, or a few contiguous lines of code text—or heuristically sample from larger scope to produce a small enough input to ﬁt inside the memory of a GPU.
For example, Transformers learn to reason about code from a sequence of tokens in the program; and GNNs with n layers—a hyper-parameter which, for message-passing architectures, determines how much of the graph is reachable from some immediately adjacent context to the task example—prune all nodes that are further than n graph hops away from the placeholder.
However, distance is not always the determining factor in collecting relevant information.
For instance, when trying to decide which exception type is applicable, it may be more important to follow control ﬂow edges until they meet the raising of an exception than, say, fetching all adjacent statements without exception handling.
We outline two key design goals motivated by these considerations.
First, we observe that programming languages have well-deﬁned semantics, which ensures that relevant information (module imports such as pickle , class inheritance, inter-procedural control ﬂow and data ﬂow, deeper analyses like def/use chains and object escape, etc.) can be extracted via a number of deterministic analyses.
So, instead of learning this kind of information indirectly from labeled data, we can make it directly
[Figure]
available to the model.
The model can thus focus on learning the information that can only be discovered from rich contexts.
Second, even when rich information is easily accessible, making well-informed predictions in code-reasoning tasks requires intelligent context collection to ﬁt the needs of the task.
So, instead of solely considering the model’s technical constraints, we must capture relevant context in a task-speciﬁc manner.
In this paper, we propose C ODE T REK , a deep learning approach that realizes these goals.
C ODE T REK leverages a declarative program analysis framework to produce a rich, easily extensible representation of context as a relational database, and a biased graph-walk mechanism for pruning that context in a task-speciﬁc way before presenting it to a model based on Transformers and DeepSets (Zaheer et al., 2017).
C ODE T REK builds upon Semmle (Avgustinov et al., 2016), which converts codebases in C, Java, Python, etc., into relational databases that capture the underlying structure and semantics of code, as well as a query language, CodeQL, for specifying program analyses to compute new semantic information.
C ODE T REK brings little modelling innovation—the architecture is reminiscent of neural techniques for knowledge-graph reasoning (Das et al., 2018; Perozzi et al., 2014).
However, it is useful both as a relation generator and a task generator.
As the former, it can harness existing and new analyses to add more inductive bias for hard tasks.
As the latter, it can help benchmark neural techniques on more challenging tasks, and generate auxiliary training objectives to pre-train unsupervised code-understanding models (Feng et al., 2020;
Guo et al., 2020; Kanade et al., 2020).
We evaluate C ODE T REK on four diverse tasks on real-world Python programs.
They include two existing tasks, variable misuse and exception prediction , as well as two newer ones, unused deﬁnition and variable shadowing .
The newer tasks are sophisticated CodeQL queries, written by program analysis experts, and enable testing the power of neural models: they both involve complex logical reasoning, and only 1.6% of the unused deﬁnition samples contain bugs, which is more in line with real-world settings.
C ODE T REK achieves an accuracy of 91%, 63%, 98%, and 94% on these tasks respectively, which is 2-19% points higher than state-of-the-art neural models CuBERT, GREAT,
GGNN, and Code2Seq.
We also demonstrate the robustness of C ODE T REK in two out-of-distribution scenarios: real-world variable misuse samples from GitHub and unused deﬁnition samples involving subtle code perturbations introduced using a systematic test -generation framework, Skeletal Program
Enumeration (Zhang et al., 2017).
C ODE T REK achieves an accuracy of 57% and ROC -AUC of 78%, respectively in these scenarios, which is 6–11% points and 14–36% points higher than the baselines.
In summary, this paper makes the following contributions:
We propose to represent programs as relational databases that make rich context readily available for code-reasoning tasks using deep learning.
We present a graph-walk mechanism that prunes the unrelated context in a task-speciﬁc manner.
We propose techniques to enable task designers to easily tailor and stress-test their models via program analysis queries, walk speciﬁcations, and systematic test-program generation.
We identify two new challenging tasks for neural code reasoning, unused deﬁnition and variable shadowing ; although sophisticated, non-neural static-analysis tools can solve them, these tasks pose a useful litmus test for neural code-reasoning frameworks and demonstrate the ability of C ODE T REK to generate hard tasks that follow real-world program distributions with modest effort.
We extensively evaluate our approach and demonstrate that deeper relational information about code helps neural models outperform the state-of-the-art in terms of accuracy and robustness.
T HE C ODE T REK F RAMEWORK 2. B ACKGROUND
Inspired by the idea of storing codebases as databases, C ODE T REK represents a program as a relational database.
Speciﬁcally, C ODE T REK leverages the per-language schema deﬁned by Semmle to uniformly store lexical, syntactic, and semantic program information as base relations in the database—we focus on Python in this paper, but the approach is language-agnostic, as long as
Semmle supports the language.
Each relation contains information—in the form of tuples—about a particular kind of program element, such as expressions, statements, and so on.
The columns of a relation specify its attributes.
For instance, in Figure 1, tuple (s1, except, c1) in the stmt relation speciﬁes that s1 is an except statement contained in a scope with identiﬁer c1 , and tuple (s4, o) in the def relation speciﬁes that variable o is deﬁned in some statement with identiﬁer s4 .
Some attributes enable tuples of different relations to be related to each other.
Because c1 is a unique identiﬁer for scopes, it can relate tuples that are in the same scope, and determine that o is deﬁned in an assign statement in that scope.
In Figure 1, consider foreign key CID in table stmt that references the primary key ID in table scope . A referential integrity constraint states that any non-null value in stmt. CID must occur in scope. ID .
Facilitated by C ODE T REK ’s uniform representation of programs, task developers can easily obtain new semantic information by writing program-analysis queries in CodeQL, an SQL-like language.
The newly derived information is also in the form of derived relations , which maintains the uniformity of the relational representation.
The derived information is stored in def , which, together with call , can bias the prediction of the best variable to replace a placeholder.
A task developer need not be a machine-learning expert to bring in more semantic information about programs: All they need do is write a CodeQL query, and the resulting derived information will be added to the existing richness of the program’s available features in C ODE T REK . [Figure]
C ODE T REK translates a relational database to a graph whose nodes correspond to tuples, and whose edges follow referential integrity constraints.
An example such graph is illustrated in Figure 2 where each node is depicted as a circle along with its type (e.g., func ) in white font.
Orange and green nodes correspond to tuples of base and derived relations, respectively.
The attributes (e.g., name , kind , etc.) of the node are shown in a box at the corner of the node.
An edge type (e.g., call func ) is the concatenation of the types of the two nodes (e.g., func and call ) that it connects.
This graph view of program semantics helps extract succinct context as input to a model.
Context extraction from the resulting C ODE T REK graph is done via biased random walks of the graph, in a fashion speciﬁed by the task deﬁnition.
The starting node—which we call an anchor node—may be example-speciﬁc (e.g., the node containing the placeholder) or task-speciﬁc (e.g., all nodes holding a variable declaration).
In
Figure 2, the node that represents tuple stmt(s1, except, c1) —which corresponds to the statement on line 137 in Figure 1 (left)—is the anchor because the goal of the task is to predict a suitable exception type in the except statement. The walk generator traverses the graph by biasing traversal of edges stmt scope func call ... except test.py call_chk ... scope_stmt func_scope call_func call_func ...
node types node values edge types subtokenize except test py call chk ... summation & embed using E v node vocabulary ( E R ) edge vocabulary ( E e ) mean pooling positional encoding (
[Figure]
[Figure] Transformer concat positional encoding positional encoding (n) (n) (n - 1) [Figure] according to each neighbor’s node type.
If no bias is speciﬁed, walks are simply fair random walks.
Different probability mixes for different node types encourage the model to sample walks more relevant to a task.
An example of such a walk is shown in Figure 2 using circles with thicker borders.
This walk reaches the ”assert” node which in fact determines the exception type that should be used in the except statement.
For instance, to spend more time traversing longer-range dependencies in other functions, the developer can assign a higher value to call nodes.
In our evaluation, we assign higher probabilities to nodes of types stmt and expr .
We could achieve improved accuracy for the tasks compared to baselines by only modifying the probabilities of up to 4 types of nodes.
Learning the walk speciﬁcation given the task without human input is exciting future work.
Finally, to convert random walks to a distributed representation, C ODE T REK embeds each walk (including the types and attributes of each node and edge in the walk) using a Transformer encoder, and then produces an order-invariant representation of the set of walks using the Deep Set architecture (Zaheer et al., 2017).
The resulting hidden representation can then be used by to make predictions for the particular code-reasoning task.
Random walks to embed a graph are perhaps a regressive choice, compared to more modern solutions such as GNNs or relational Transformers.
For instance, DeepWalk (Perozzi et al., 2014) uses random walks for distributed-representation learning in a transductive setting, and Code2Seq
(Alon et al., 2019a) uses shortest paths between pairs of AST leaf nodes.
We chose biased random walks for several reasons.
First, this enables C ODE T REK to choose task-speciﬁc strategies to heuristically fetch relevant context for a task, rather than choosing to embed all tokens of a function or class in lexical order.
Second, in contrast to GNNs, this enables C ODE T REK to potentially follow much longer chains through the semantic graph than what would be possible in a message-passing GNN of a tractable number of layers—this was, in fact, the motivation behind the model architectures by
Hellendoorn et al.
Finally, in contrast to other walk-based approaches such as Code2Seq and
AnyCodeGen, which share some of our motivation, our graph structure has much richer connectivity and is larger.
For instance, those approaches only consider paths ending at token-bearing leaf nodes, with only AST interior nodes in between, whereas C ODE T REK admits arbitrary paths through the graph.
In our example, paths with more than two token nodes, e.g., the walk illustrated in Figureusing circles with thicker borders is admissible for C ODE T REK but not for Code2Seq.
C ODE T REK builds upon the above techniques and extends to relational graphs, with different sampling strategies and neural architectures that would suit database representation better.
2. B UILDING B LOCKS OF C ODE T REK
We now describe brieﬂy the building blocks of C ODE T REK .
Pre-processing turns code and tasks into graphs.
Then a graph is processed into sets of graph walks to present to a model.
Finally, we train using a cross-entropy loss to implement a particular task.
Code as a Graph.
Pre-processing consists of two steps: (a) turning the codebase into a relational database ( Code2Rel ), using a pre-existing set of analyses and the developer’s own analyses, and (b) mapping the relational database to a graph ( Rel2Graph ).
Code2Rel applies to the codebase the system’s base program analysis queries, which make up the base relations, as well as those provided by the developer, which form the derived relations.
The result is a database comprising of a number of named tuples for each relation type.
For Python, we collect 95 base relations and 277 derived relations, although adding more derived relations is simply an exercise in writing a few lines of CodeQL.
Optionally, this step can limit the resulting database to only those relations that are reachable (through query-to-query dependencies) from some program analyses marked as required by the developer.
Appendix C.1 provides further details.
Rel2Graph interprets the relations produced by Code2Rel as a graph, as follows.
Each named tuple is represented by a node with the values of the tuple attributes as its features.
For any two relations R and S connected via referential integrity, an R S edge type is deﬁned, connecting the tuples of the two relations with the same value on the edge attributes.
Appendix C.2 provides further details.
Neural Representation.
Given a code database that is converted to graph G via the above
Rel2Graph, we propose to represent it by the embedding of a set of walks W , via the procedure
Graph2Walks .
Graph2Walks projects a code graph as produced by Rel2Graph to a set of walks, according to the task-speciﬁc walk speciﬁcation (anchor node predicate, traversal bias, and target walk length).
Graph2Walks samples from the distribution of such random walks, by repeatedly picking a node satisfying the anchor predicate, and traversing up to a maximum number of neighbors, following the transition probabilities speciﬁed.
The resulting walks are collected as token sequences of relation types of nodes, their attribute values, and the edge types traversed, in the order of traversal.
Figure 3 row (1) shows such a walk representation corresponding to the walk highlighted in Figure 2.
Appendix C.3 speciﬁes Graph2Walks precisely.
Given a walk w = [ v 0 , e 0 , v 1 , e 1 , . . .
, v n − 1 ] of length n steps consisting of n nodes and n − 1 edges, we produce an initial embedding X w ∈ R (3 n − 1) × d , where d is the embedding dimension.
It consists of three parts.
The ﬁrst n rows of the embedding tensor correspond to the embedding of node type (the relation name) from the embedding-lookup over E R ∈ R R × d , where R is the number of relations.
The next n rows are the bag-of-words embedding for the attribute sub-tokens obtained from token embedding E v ∈ R v × d , where v is our attribute-vocabulary size—we produce a WordPiece vocabulary for our attribute values.
The last n − 1 rows are the embedding lookup for the types of each edge in E e ∈ R R 2 × d —recall that an edge type is the pair of relation types it connects.
All three embedding matrices E R , E v , E e are learnable parameters.
Each individual part of the embedding tensor gets its own learnable positional encoding (denoted as PE w ∈ R (3 n − 1) × d ).
We use a Transformer encoder to represent the embedding of walk w as e w = POOLING (cid:16) Transformer ( X w +
PE w ) (cid:17) :
R (3 n − 1) ×
d (cid:55)→ R d , where after the last layer of Transformer we do mean-pooling over all 3 n − 1 elements of the walk, to obtain a d -dimensional e w .
The steps for embedding of a walk sampled from the graph in Figure 2 is illustrated in Figure 3.
Training and Inference.
An example consists of a set of walks and a ground-truth label, ( W, ˆ y ) .
Given an unordered set of walk embeddings { e w } w ∈ W , we build a classiﬁer by using the construction y =
MLP (cid:0) DeepSet ( { e w } w ∈ W ) (cid:1) , where y denotes the predicted label and we optimize for cross entropy loss.
However, for a binary classiﬁcation task, we can obtain a more interpretable model via y = (cid:80) w ∈ W α w σ ( MLP ( e w )) , where σ ( · ) is the sigmoid function, and α w =
[Equation]
This way, we can inspect the individual walks that contributed the most (i.e., the highest α w ) to the positive or negative predictions, and see if that aligns with human reasoning.
We refer to α w as the walk score .
We train using the Adam optimizer with 8 GPUs for distributed synchronized SGD training (see Appendix A for details).
E VALUATION
Tasks.
We consider two main criteria in selecting tasks.
The ﬁrst is locality , which is determined by whether reasoning within a function typically sufﬁces, or whether inter-procedural reasoning is required.
The second is declarativity —whether the task can be stated as a logic problem that can be solved using declarative queries.
For declarative tasks, we write queries in CodeQL (detailed in
Appendix B); for non-declarative tasks, we rely on available datasets.
We treat the following tasks:
V AR M ISUSE .
Given a function and a variable accessed in it, predict whether the variable is misused.
We also consider a variation of this task, V AR M ISUSE -F UN (Kanade et al., 2020), that takes only a function and predicts whether all variables are used correctly in the function.
Note that neither variation is declarative: given a well-formed program, no logic query can deterministically decide that a variable is misused, since that decision depends on the intended semantics. [Table]
E XCEPTION .
Given a module containing a masked exception type in an except clause, predict the most appropriate built-in exception type out of 20 choices.
We also consider a variation of this task, E XCEPTION -F
UN (Kanade et al., 2020), that is similar to E XCEPTION but takes a single function as scope.
Although E XCEPTION needs inter-procedural reasoning, neither variation is declarative, since the choice of the appropriate exception type is subjective in Python.
D EF
U SE .
Given a function and a variable deﬁnition in its scope, predict whether the deﬁnition is used.
We also consider a variation of this task, D EF U SE -F UN , that takes a function as its input and predicts whether any deﬁnitions are unused.
This task is especially interesting because the real-world distribution of programs that contain unused deﬁnitions is skewed.
Both variations are declarative (see Appendix B.1) and require only intra-procedural reasoning.
V AR S HADOW .
Given a module, predict whether any variable deﬁned within a certain scope has the same name as a variable deﬁned in an enclosing outer scope, thereby shadowing that latter variable.
Similar to E XCEPTION , this task requires inter-procedural analysis in order to reason over both local as well as global variables.
It has a declarative query (see Appendix B.2).
Benchmark.
We use the ETH Py150 Open corpus consisting of 125K Python modules 1 .
It is a de-duplicated and redistributable subset of ETH Py150 2 .
Speciﬁcally, for the non-declarative tasks, we use the datasets released by Kanade et al.
Since these are function-level samples but the
E XCEPTION task is module-level, we augment the function in each sample with the entire containing module for this task.
For the declarative tasks, we use analyses written in CodeQL to annotate all functions (or modules, as applicable) in ETH Py150 Open.
All datasets consist of real examples, except for V AR M ISUSE -F UN and V AR M ISUSE where variable misuses are synthetically introduced into real code.
We collected a number of apparent variable misuses from GitHub commits to test our models and baselines on a realistic dataset.
Dataset details are available in Appendix G.
Baselines.
To compare C ODE T REK ’s performance with state-of-the-art techniques, we select four baselines: we implement GGNN by Allamanis et al.
and Code2Seq by Alon et al.
(2019a)
, build classiﬁers on top of the GREAT encoder by Hellendoorn et al.
(2020), and ﬁne-tune the pre-trained Python model for CuBERT by Kanade et al.
(2020), which is essentially the Transformerbased classiﬁer implementation of BERT.
For Code2Seq we use ASTs as base program structures as described by Alon et al.
(2019a).
We sample leaf-to-leaf paths from these ASTs.
The number of paths we sample is the same as the number of walks we sample for training C ODE T REK models.
For GGNN and GREAT, we compute the data ﬂow, control ﬂow, and lexical information described by Allamanis et al.
and Hellendoorn et al.
(2020), respectively, using Semmle CodeQL and augment program ASTs with those edges.
We detail baseline hyperparameters in Appendix A, and ﬁne-grained information about the size of AST-based versus C ODE T REK graphs in Appendix F.1.
3. A CCURACY OF C ODE T REK
We evaluate the performance of C ODE T REK and the baseline techniques on all the tasks described in
Section 3, all presented as classiﬁcation tasks.
We perform 10-fold cross-validation and report the average of the metric that we use to measure the performance of each task.
We use ROC-AUC as the metric for D EF U SE and D EF U SE -F UN tasks due to their unbalanced datasets, and accuracy for the remaining tasks.
The results are reported in Table 1.
In 5 out of 7 tasks, C ODE T REK outperforms
GGNN, Code2Seq, GREAT, and CuBERT by 2–19% points.
[Table]
There are various reasons why C ODE T REK performs better than these approaches.
First, declarative tasks such as D EF
U SE -F UN (or D EF U SE ) require complex reasoning about the interactions between program variables.
For instance, one needs to reason about the uses and deﬁnitions of variables in a ﬂow-sensitive manner to determine whether any unused deﬁnitions exist in a program.
Consider the code snippet in Figure 4.
The deﬁnition of the variable month on line 2 is unused but that at line 3 is used in line 4.
C ODE T REK gives a majority of the walks sampled using the deﬁnition at line 3 a high score (around 0.99), indicating the existence of a use of that deﬁnition.
However, most of the walks sampled from the deﬁnition of month at line 2 were given a lower score, and so C ODE T REK determines that this deﬁnition is unused.
We observe that both CuBERT and GREAT fail to distinguish between the deﬁnitions on lines 2 and 3, and so they predict both to be used. [Figure]
Additionally, some tasks such as E XCEPTION require reasoning beyond the boundaries of a single function to make informed predictions.
Functions in the chain of function calls can be lexically far from each other, thus rendering the limited context size of CuBERT and GREAT insufﬁcient.
We observe that GGNN fails in the presence of such long call chains on par with ﬁndings of Alon & Yahav (2021).
C ODE T REK addresses this kind of mispredictions by readily using a call graph relation to connect the chain of function calls.
This enables C ODE T REK to traverse a long distance without the need to consider other statements in the program that have no effect in raising some exception.
However, C ODE T REK performs worse than CuBERT in E XCEPTION -F UN .
This could be attributed to the fact that CuBERT is pre-trained on around 7 million Python programs, and therefore is able to memorize tokens from several instances of try-except blocks.
An example of a heuristic that it learns is that in presence of tokens such as request or response in the context, it suggests catching
HTTPError , which is usually the correct choice.
However, its prediction is not robust against changes in the variable names.
For instance, changing the names of a few nearby variables to request or response forces CuBERT to predict HTTPError regardless of the semantics.
C ODE T REK on the other hand, does not rely on memorizing the tokens, but learns to assign high probabilities to walks that correctly traverse a chain of function calls starting from the try blocks to locations in programs (or their libraries) where the exception is originally raised.
C ODE T REK also performs worse than GREAT in V AR M ISUSE -F UN .
This is because every node that corresponds to a variable is selected to be an anchor for this task.
The total number of walks ( in this task) is divided among these variables.
However, there can be hundreds of variables in some programs, resulting in few walks generated for each variable in such cases, diminishing the ability of
C ODE T REK to learn sufﬁcient information about each variable.
3. R OBUSTNESS OF C ODE T REK
We evaluate the robustness of C ODE T REK on additional test data that does not follow the distribution of the training data.
This data includes two new datasets: one representing real-world bugs for the V AR M ISUSE -F UN task and the other consisting of programs mutated using a systematic testgeneration framework for the D EF U SE -F UN task.
The results are reported in Table 2.
Real-world bugs.
We manually collect 199 real -world instances containing a V AR M ISUSE -F UN bug and their corrected counterparts (a total of 398 samples) from commits on GitHub and use them as testing data for the V AR M ISUSE -F UN task.
We deﬁne a V AR M ISUSE -F UN bug as the occurrence of a misused variable that is changed to another in-scope variable in the commit.
We evaluate the baselines using this real-world set of bugs.
C ODE T REK outperforms baselines in detecting real-world variable misuse bugs (V AR M ISUSE -R EAL ) by achieving an accuracy of 57% which is 6% points better than the second best result obtained by GGNN.
Mutated programs.
There are several approaches to mutating existing datasets, including transforming existing data (Yang et al.
(1992)), generating synthetic programs, and fuzzing.
A representative approach that has been used to systematically evaluate the robustness of compilers is Skeletal
Program Enumeration (SPE), proposed by Zhang et al. (2017).
SPE parameterizes each program by a set of its variables, and replaces each variable name exhaustively with other in-scope variable names.
We generate variations of the D EF U SE -F UN testing data using this technique, and evaluate the baselines on this mutated dataset (D EF U SE -SPE).
C ODE T REK outperforms all baselines in classifying these perturbed programs by achieving the ROC-AUC score of 78% which is 15% points better than the second best result obtained by Code2Seq.
The poor performance of the baselines can be explained by the fact that the code generated by
SPE is out-of-distribution.
For example, the assignment a = a + a is unusual in real code, but occurs frequently in SPE-generated samples.
Despite this, the inductive bias borne by rich relational information during training remains applicable and prevails over the unusual-looking token sequences, thus explaining C ODE T REK ’s performance.
These results suggest that sampling walks can be a promising strategy for robustness.
Interestingly, the runner-up in this study is Code2Seq—another walk-based approach.
We inspected the paths in both approaches to understand the reason behind the difference in performance of Code2Seq and
C ODE T REK despite their similarities.
We identiﬁed two reasons: 1) the kinds of program information that can be captured from an AST are limited compared to the program graph we propose, and 2) several walks that C ODE T REK prioritizes for this task cannot be embedded by Code2Seq.
E FFECTIVENESS ON L ONGER -R ANGE T ASKS
We evaluate the effectiveness of C ODE T REK on tasks that require reasoning beyond function boundaries.
C ODE T REK achieves this ability by readily incorporating relations that capture inter-procedural or inter-modular dependencies such as call graphs.
To demonstrate this, we compare C ODE T REK ’s performance on E XCEPTION with vs. without incorporating the call graph information at training time.
C ODE T REK achieves an accuracy of 52% when call graph information is not provided, which increases to 63% after providing the call graph information between functions within a module.
[Figure] To illustrate the kinds of mistakes that the baselines (and also C ODE T REK in the absence of call graph information) make, consider the representative example in Figure 5, snipped and simpliﬁed from the zipﬁle package.
In this example, the model predicts the exception type that should be caught on line 11.
However, to make an informed prediction, the model must consider the exceptions that may be raised when calling the ZipFile constructor (line 10).
Hence, the deﬁnition of the constructor (line 2) must be taken into consideration.
This constructor calls another function check compression in which a NotImplementedError is raised on line 5.
This chain of function dependence can be easily represented using a call graph.
Therefore, C ODE T REK , once provided with call graph, will eventually traverse the path that reaches this raise statement from the exception statement through the call graph edges.
3. S ENSITIVITY TO N UMBER OF S AMPLED W ALKS
We evaluate the sensitivity of C ODE T REK at test time to the number of walks that are sampled from program graphs.
All the models for the considered tasks are trained on 100 sampled walks per program.
The results are reported in Figure 6.
We observe that the accuracies of the models increase with the number of sampled walks.
In some tasks, such as D EF U SE and V AR M ISUSE that involve local reasoning about one point in the program, reducing the number of walks from 100 to 50 reduces the accuracies of the models by a very small amount.
On the other hand, for tasks that require reasoning about numerous points in the program (e.g., D EF U SE -F UN ) or reasoning globally (e.g., E XCEPTION ) decreasing the number of sampled walks has a bigger impact on the accuracy.
3. I MPACT OF P ROGRAM R EPRESENTATION
To evaluate the impact of different code representations, we train two models for each task using
C ODE T REK ’s architecture: for one set of models, the walks are sampled from relational graphs whereas for the other set of models, the walks are sampled from ASTs.
The performance results are reported in Table 3.
Notably, the models trained on walks sampled from relational graphs are about 20–36% more accurate than models trained on walks sampled from ASTs. [Figure] [Figure] [Table]
We also evaluate the usefulness of the ability to bias random walks in C ODE T REK .
The accuracy results that are reported in Table 1 are all trained on biased random walks.
Speciﬁcally, in all of the tasks, nodes with types stmt , expr , and variable are biased such that they are 5 times more likely to be traversed compared to other kinds of neighboring nodes.
In addition, in the E XCEPTION task, we decrease the bias assigned to nodes of type module to 0 to avoid traveling from one function to another through the module node they have in common.
(See Appendix D for more details about speciﬁcation of the tasks.)
This forces the walks to only go to other functions by taking call graph edges between them.
We select one of the tasks, E XCEPTION -F UN , to measure the accuracy in the absence of said biases.
We re-train E XCEPTION -F UN using uniformly sampled walks and observe that the accuracy reduces from 65% to 58% as a result.
R ELATED W ORK
Learning to represent code.
There is a rich literature on using neural networks for code reasoning.
At the token sequence level, the Transformer and its variants (Hellendoorn et al., 2020; Dowdell &
Zhang, 2020) have been widely used (Berabi et al., 2021; Ahmad et al., 2020; Z¨ugner et al., 2021;
Kim et al., 2021; Wang et al., 2020).
Their performance can be further boosted via pretraining (Feng et al., 2020;
Guo et al., 2020; Kanade et al., 2020; Wang et al., 2021;
Peng et al., 2021; Liu et al., 2020).
Others have proposed to represent programs with ASTs and additional semantic edges (Allamanis et al., 2018;
Brockschmidt et al., 2018) or learned abstract relations (Johnson et al., 2020), using GNN or leaf-to-leaf sequence embeddings (Alon et al., 2019a;b).
Our work enables adding much richer semantic information while reducing the dependency on syntax structures.
Additionally, C ODE T REK can take advantage of program analysis queries on relational databases to eliminate the engineering burden of augmenting program graphs with additional semantic edges.
Graph representation learning.
Our work on learning program representations via relational databases is closely related to inductive representation learning on graphs (Hamilton et al., 2017) with graph neural networks (GNNs)
(Xu et al., 2018) or Transformers (Ying et al., 2021).
Although scalable GNNs via sampling (Chen et al., 2017; Zhou et al., 2020) have been proposed in the transductive setting, it is still challenging to represent large database graphs with 100k nodes in this inductive setting (Clement et al., 2021;
Yang & Kuang, 2021).
Techniques from transductive graph embedding based on skip-gram (Perozzi et al., 2014; Grover & Leskovec, 2016) or general knowledge graph embedding (Das et al., 2017; Hamilton et al., 2018; Zheng et al., 2020) are scalable but not directly applicable for inductive setting.
C ODE T REK achieves a good balance between modeling for large codebases and efﬁciency.
Feature selection techniques.
There are efforts in the data-mining literature to minimize human effort in feature augmentation and selection.
Chepurko et al.
(2020) discover joins that can improve the prediction accuracy for a single data table whereas C ODE T REK operates on multiple tables.
C ONCLUSIONS
We proposed C ODE T REK , a technique that represents programs as relational databases to make rich semantic information available to deep learning models for code-reasoning tasks.
We also introduced a ﬂexible walk-based mechanism to sample relevant contexts from large graphs which are constructed from relational databases.
We evaluated C ODE T REK on a variety of real-world tasks and datasets, and showed that it outperforms state-of-the-art neural models.
A T RAINING P ARAMETERS AND H YPERPARAMETERS
In this section we describe details on parameters and hyperparameters we used.
CodeTrek We train C ODE T REK models with a learning rate of 10 − 4 , 4 transformer layers, an embedding size of 256, 8 attention heads, and 512 hidden units.
We sample 100 walks with lengths of up to 24 in each graph for every task, except for the V AR M ISUSE -F UN task for which we sample 500 such walks per graph.
The reason is that the anchors we select for V AR M ISUSE -F UN task are all the variables in the given program which can be well over 100 variables.
So, we increase the total number of walks to include more random walks starting from each variable.
CuBERT We ﬁne-tune the CuBERT pre-trained model that is provided by Kanade et al.
(2020) with a learning rate of 10 − 4 , 4 transformer layers, and 512 hidden units.
We use the checkpoint that is pre-trained on examples of size 512 tokens.
GREAT We train GREAT models with a learning rate of 10 − 4 , 4 transformer layers, 8 attention heads and 512 hidden units.
The example size in 512 tokens.
Code2Seq
We train Code2Seq models with a learning rate of 10 − 3 , 4 layers, 512 hidden units, and embedding size of 256.
We sample 100 paths in each AST.
GGNN We train GGNN models with a learning rate of 10 − 4 , 5 layers, a latent dimension of size 128, and a message dimension of size 128.
B C ODE QL Q UERIES U SED FOR L ABELING
In this section, we present the CodeQL queries that we used to label the examples for newly added tasks.
Both queries are adapted from CodeQL’s query repository at https://github.com/ github/codeql .
[Table] [Table]
UN
[Table]
UERY
[Table]
B. V AR S HADOW Q UERY
[Table] [Table]
[Table]
C C ODEBASES AS D ATABASES C. T RANSLATING C ODE TO R ELATIONAL D ATABASE
C
ODE
T REK views program information as relations.
Table 4 shows the number of relations in three common programming languages.
There is a directed acyclic graph F that represents the dependencies between these relations.
For example, in F Python represents the dependencies between 95 base relations plus 277 derived relations.
Example C.1.
A certain Python task requires base relations { A, B } and derived relation { E } .
However, in F Python , E depends on derived relation D , which in turn depends on base relations { B, C } .
Therefore, Code2Rel computes all ﬁve relations: A , B , C , D , E in order.
Algorithm 1 ( Code2Rel )
Given a program P , a set of base relation names R B , and a set of derived relation names R Q , construct and return database D .
Note that the term node in this algorithm refers to nodes in the relation dependency graph, not to nodes in the program graph (e.g., in Algorithm 2) that the model will see.
Initialize D to the set of all base relations in R B by translating program P .
Let R S be
the set of relation names reachable from R Q in relation dependency graph F :
(a) R Q ⊆ R S (b) if r ∈ R S and r → r (cid:48) ∈ Edges ( F ) then r (cid:48) ∈ R S 3.
Let F be the sub-graph of F induced by set of nodes R S : (a) Nodes ( F ) =
R S (b) Edges ( F ) = { r → r (cid:
) ∈ Edges ( F ) | r, r (cid:
) ∈
Nodes ( F ) } 4.
Compute a topological ordering L =
[ r 1 , ..., r | Nodes ( F ) | ] of F . 5.
For each r in L in order:
Evaluate the query for computing relation r on database D and add the result to D .
Example C.2.
In Figure 7, the Python code snippet on the left side illustrates module test threadable.py .
Algorithm 1 translates this program to a set of base relations, e.g., (t2,ExceptHandler,s3) in table Statements that corresponds to the exception handling statement on line 141, and (f1,aMethod,s2) in table Functions that corresponds to the method deﬁned on line 5.
[Figure]
Then
[Figure]
it computes the required derived relation Calls to capture the fact that callMethod calls aMethod as a relation (f2,f1) .
C.
T RANSLATING R ELATIONS TO G RAPH Deﬁnition
C.1 (Referential Integrity) .
Consider relations R and S .
Referential integrity is a constraint deﬁned on a set of attributes A
= { a 1 , ..., a k } ∈ R to guarantee that the values assigned to these attributes match the values assigned to another set of attributes A (cid:48)
= { a (cid:48) 1 , ..., a (cid:48) k } ∈ S when R is referenced by S .
A tuple t 1 ∈ S references a tuple t 2 ∈ R if the values assigned to { a 1 , ..., a k } match the values assigned to { a (cid:48) 1 , ..., a (cid:48) k } .
Algorithm 2 ( Rel2Graph ) Given a database D , construct a program graph G .
Construct an undirected and labeled graph G as follows:
(a) Nodes ( G ) =
D (b) Add to Edges ( G ) each ( t 1 , t 2 , l ) that satisﬁes
the following conditions
: i. l : R. [ a 1 , ...a k ] → S. [ b 1 , ..., b k ] is a referential integrity constraint in the schema of database D ii.
t 1 is a tuple of relation named R in D iii.
t 2 is a tuple of relation named S in D iv.
for all i ∈ [1 ..k ] : t 1 .a
i =
t 2 .b
i
Example C.3.
Algorithm 2 takes the relational database—i.e., the set of named tuples—that correspond to the example program in Figure 7, and generates a program graph.
In this graph, nodes are named tuples such as n 1 =Functions(f1,aMethod,s2) , n 2 =Functions(f2,callMethod,s3) , and n 3 =Calls(f2,f1) .
Tuples are connected to each other if they have attributes in common, e.g., there is an edge between n 1 and n 3 since they have the attribute function ID in common.
This type of this edge is determined by a tuple of relation names and the common attributes, which is (cid:104) Functions.
ID,Calls.
Callee (cid:105) in this case.
C. T RANSLATING G RAPH TO A S ET OF W ALKS Deﬁnition
C.2 (Walk speciﬁcation) .
A walk speciﬁcation S =
(cid:104)C , B, min, max (cid:105) is a tuple in which
C is a conditional expression that ﬁlters walk anchors from the set of nodes Nodes ( G ) , B is a map of bias values that correspond to each relation name, and min, max ∈ R + specify the minimum and maximum length of the walks generated by the speciﬁcation.
Example C.4.
An example of a walk speciﬁcation is as follows.
[Equation]
Algorithm 3 ( Graph2Walks )
Given a program graph G , a walk speciﬁcation S =
(cid:104)C , B, min, max (cid:105) , and the number of walks w , sample a set of walks W .
Initialize
set
i.
ii.
Set t curr to a t ∈ Neighbors ( G, t prev ) with prob.
proportionate to S.B [ type ( t )] .
iii.
Let e curr =
( t prev , t curr , l ) ∈ Edges ( G ) iv.
If e curr / ∈ walk then extend walk by e curr .
Otherwise set t curr := t prev .
(c) If walk / ∈ W then add it to W .
Example C.5.
The sequence of edges that are numbered 1 through 5 in Figure 7, is an example of a walk sampled from program graphs.
The anchor of this walk is Statements ( t2 , ExceptHandler , s3 ) , and the length of this walk is 7—the number of nodes along the walk.
C.
B RINGING ALL THE PIECES TOGETHER
Deﬁnition C.3 (Task speciﬁcation) .
A task speciﬁcation T = (cid:104) R B , R Q , S, n (cid:105) is a tuple in which
R B is a set of base relation names, R Q is a set of derived relation names, S is a walk speciﬁcation as described in Deﬁnition C.2, and n is the number of walks to be generated.
Algorithm 4 ( Code2Walks )
Given a program P and a task speciﬁcation T = (cid:104) R B , R Q , S, n (cid:105) , generate a set of walks W .
D = Code2Rel
( P, T.R B , T.R Q ) 2.
G =
Rel2Graph ( D ) 3.
W = Graph2Walks ( G, T.S, T.n )
D T ASK S PECIFICATIONS
Using the notations deﬁned in Appendix C, we describe the speciﬁcations of each task that we used for evaluating C ODE T REK .
In the following speciﬁcations, ellipsis (...) indicate the rest of the universe of base relations as designed in Semmle framework.
T
[Equation]
R
[Equation]
= { stmt, expr,
[Equation]
, ssa-defn, ssa
[Equation]
use, successor,
[Equation]
R Q = {} ,
S = {
C =
{ t | t ∈ Nodes ( G ) ∧ t ∈ expr ∧ t.Kind = name } ,
B = { stmt : 5 , expr : 5 , variable : 5 } , min = 4 , max =} , n =}
T E XCEPTION -F UN
= {
R B = { stmt, expr, variable, ssa-defn, ssa-use, successor, ... } ,
R Q = {} ,
S = {
C =
{ t | t ∈ Nodes ( G )
∧ t ∈ stmt ∧ t.Kind = except ∧ t.Type = HoleException } ,
B = { stmt : 5 , expr : 5 , variable : 5 } , min = 4 , max =} , n =}
T E XCEPTION = {
R B = { stmt, expr, variable, ssa-defn, ssa-use, successor, ... } ,
R Q = { call-graph } ,
S = {
C =
{ t | t ∈ Nodes ( G )
∧ t ∈ stmt ∧ t.Kind = except ∧ t.Type = HoleException } ,
B = { stmt : 5 , expr : 5 , module : 0 } , min = 10 , max =} , n =}
T D EF U SE -F UN
= {
R B = { stmt, expr, variable, ssa-defn, ssa-use, successor, ... } ,
R Q = { variable-defs, local-variables } ,
S = {
C =
{ t | t ∈ Nodes ( G ) ∧ t ∈ expr ∧ t.Kind = name ∧ t.Context ∈ { write, param } } ,
B = { stmt : 5 , expr : 5 , variable : 5 } , min = 4 , max =} , n =}
T D EF U SE and T V AR M ISUSE are deﬁned similar to T D EF U SE -F UN and T V AR M ISUSE -F UN , respectively.
However, in both T
D EF U SE and T V AR M ISUSE their C has an additional restriction t.location = user selection.
Also, T V AR M ISUSE is evaluated with n = 100 .
T V AR S HADOW = {
R B = { stmt, expr, variable, ssa-defn, ssa-use, successor, ... } ,
R Q = { call-graph } ,
S = {
C =
{ t | t ∈ Nodes ( G ) ∧ t ∈ var ∧ t.Kind = name ∧ t.is global =
True } ,
B = { stmt : 5 , expr : 5 , variable : 5 } , min = 4 , max =} , n =
} E W ALK F ORMAT
We describe the format of a walk as the model views it before encoding using an example.
Each walk consists of three lists for node types, node values, and edges.
{ "anchor": "py_stmts(415098,6,415072,2)" "trajectory": { "node_types": ["py_stmts", "py_scopes", "py_Functions", "py_scopes", "py_stmts", "py_exprs", "py_exprs", "py_scopes", "py_exprs", "py_variables", "v_8"], "node_values": [["ExceptStmt"], "", ["pipe","line"], "", ["Assign"], ["Call"], ["Attribute"], ["Name"], ["metr","ic"]], "edges": ["(py_scopes.node,py_stmts.id)", "(py_Functions.id,py_scopes.scope)", "(py_Functions.id,py_scopes.scope)", "(py_scopes.node,py_stmts.id)", "(py_exprs.parent,py_stmts.id)", "(py_exprs.id,py_exprs.parent)", "(py_exprs.id,py_scopes.node)", "(py_exprs.id,py_scopes.node)", "(py_exprs.id,py_variables.parent)", "(py_variables.id,variable.id)"] }, }
The node values can consist of the value of any attribute of a node.
For instance, ExceptStmt above is the kind of the node with type py stmts , and pipeline is the name of the function that corresponds to the function which is illustrated by node with type py Functions .
The values that are identiﬁers, function names, etc. are subtokenized using a subword tokenizer (the tensor2tensor package).
For instance, pipeline will break into [pipe,line] .
Therefore, each value is encoded using the corresponding vector in the subword dictionary.
If the relation name of a tuple is variable , we assign an id to it to be able to distinguish between different variables.
So, the type of a variable node is determined by an id (e.g., v 8 ).
F A DDITIONAL I NFORMATION ABOUT G RAPHS
F. G RAPH S IZES
Relational
[Table]
[Table]
F. E DGES R EPRESENTED TO B ASELINE G RAPHS
[Table]
GGNN Graphs.
The edges that are represented to GGNN models are borrowed from Allamanis et al.
They include
AST edges 2.
NextToken edges 3.
LastRead/LastWrite/ComputedFrom/LastLexicalUse edges among variable accesses 4.
GuardedBy/GuardedByNegation edges between variables used in branches and their corresponding conditional expressions 5.
ReturnsTo edges from the return statement to the method declaration, and 6.
FormalArgName edges between method call arguments and their corresponding formal parameters.
GREAT Graphs.
For GREAT models, similar to Hellendoorn et al.
(2020), we borrow the edge types from Allamanis et al.
and augment them with function calls.
G D ATASET S IZES
The number of samples in each dataset is shown in Table 7.
The V AR S HADOW dataset consists of 49% samples with positive labels.
D EF U SE -F UN dataset is more skewed, with 15% positive labels, dataset has 1.6% labels.