CodeTrek: Flexible Modeling of Code using an Extensible Relational Representation

1 INTRODUCTION . 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-defined exception types in Python requires understanding the implementation of the check ( ) function , which is defined 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 fit 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 flow 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-defined semantics , which ensures that relevant information ( module imports such as pickle , class inheritance , inter-procedural control flow and data flow , 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 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 fit the needs of the task . So , instead of solely considering the model ’ s technical constraints , we must capture relevant context in a task-specific manner . In this paper , we propose CODETREK , a deep learning approach that realizes these goals . CODETREK 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-specific way before presenting it to a model based on Transformers and DeepSets ( Zaheer et al. , 2017 ) . CODETREK 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 . CODETREK 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 CODETREK 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 definition 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 definition samples contain bugs , which is more in line with real-world settings . CODETREK 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 CODETREK in two out-of-distribution scenarios : real-world variable misuse samples from GitHub and unused definition samples involving subtle code perturbations introduced using a systematic test-generation framework , Skeletal Program Enumeration ( Zhang et al. , 2017 ) . CODETREK 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 : 1 . We propose to represent programs as relational databases that make rich context readily available for code-reasoning tasks using deep learning . 2 . We present a graph-walk mechanism that prunes the unrelated context in a task-specific manner . 3 . We propose techniques to enable task designers to easily tailor and stress-test their models via program analysis queries , walk specifications , and systematic test-program generation . 4 . We identify two new challenging tasks for neural code reasoning , unused definition 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 CODETREK to generate hard tasks that follow real-world program distributions with modest effort . 5 . 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 . 2 THE CODETREK FRAMEWORK . 2.1 BACKGROUND . Inspired by the idea of storing codebases as databases , CODETREK represents a program as a relational database . Specifically , CODETREK leverages the per-language schema defined 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 specifies that s1 is an except statement contained in a scope with identifier c1 , and tuple ( s4 , o ) in the def relation specifies that variable o is defined in some statement with identifier s4 . The schema also defines referential integrity constraints of the form R.A → S.B where A is called a foreign key of referencing relation R , and B is a unique attribute ( e.g . a primary key ) of referenced relation S. For example , in Figure 1 , we have stmt.CID→ scope.ID . 2.2 A BIRD ’ S-EYE VIEW OF CODETREK . Facilitated by CODETREK ’ 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 CODETREK . CODETREK translates a relational database to a graph whose nodes correspond to tuples , and whose edges follow referential integrity constraints . An example of such a 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 . For each referential integrity constraint R.A→ S.B , an edge type R.A S.B is defined , connecting the tuples of the two relations with the same value on the edge attributes R.A and S.B . For brevity of presentation in Figure 2 , when there is a single such constraint between a pair of relations R and S , we omit the attributes from the edge type . But we do not omit them when there are multiple such constraints , such as in the case between relations func and call , namely , call.caller → func.id and call.callee → func.id . This graph view of program semantics helps extract succinct context as input to a model . Context extraction from the resulting CODETREK graph is done via biased random walks of the graph , in a fashion specified by the task definition . The starting node—which we call an anchor node—may be example-specific ( e.g. , the node containing the placeholder ) or task-specific ( 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 according to each neighbor ’ s node type . If no bias is specified , 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 specification given the task without human input is exciting future work . Finally , to convert random walks to a distributed representation , CODETREK 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 CODETREK to choose task-specific 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 CODETREK 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 . ( 2020 ) . 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 CODETREK admits arbitrary paths through the graph . In our example , paths with more than two token nodes , e.g. , the walk illustrated in Figure 2 using circles with thicker borders is admissible for CODETREK but not for Code2Seq . CODETREK builds upon the above techniques and extends to relational graphs , with different sampling strategies and neural architectures that would suit database representation better .

Summary: The paper presents a framework for creating task-specific program representations for deep learning based on program analysis and random graph walks. The basic idea is to use program analysis to create a graph-based representation of the program, capturing both syntactic and semantic relationships between elements. This graph-based representation of the program is then encoded by embedding a set of (random) walks in the graph. The paper evaluates the framework using four tasks: three tasks that only require intra-procedural program analysis (var-misuse, infer exception type, predict where a definition is used) and one that required inter-procedural analysis (variable shadowing). CodeTrek is shown to outperform existing techniques in most benchmarks.