# COMEX: A Tool for Generating Customized Source Code Representations

arXiv: https://arxiv.org/abs/2307.04693 · HTML: https://mlsanity.com/paper/2307.04693/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 24 references that are in the map.
- First posted on arXiv 10 Jul 2023.

## Builds on (most cited first, 13 of 13)

- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [CodeSearchNet Challenge: Evaluating the State of Semantic Code Search](https://mlsanity.com/paper/1909.09436/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks](https://mlsanity.com/paper/1909.03496/) (2019)
- [A Survey of Machine Learning for Big Code and Naturalness](https://mlsanity.com/paper/1709.06182/) (2017)
- [Learning to Represent Programs with Graphs](https://mlsanity.com/paper/1711.00740/) (2017)
- [Neural Program Generation Modulo Static Analysis](https://mlsanity.com/paper/2111.01633/) (2021)
- [Static Prediction of Runtime Errors by Learning to Execute Programs with External Resource Descriptions](https://mlsanity.com/paper/2203.03771/) (2022)
- [Learning Graph Structure With A Finite-State Automaton Layer](https://mlsanity.com/paper/2007.04929/) (2020)
- [Git Blame Who?: Stylistic Authorship Attribution of Small, Incomplete Source Code Fragments](https://mlsanity.com/paper/1701.05681/) (2017)
- [Multi-View Graph Representation for Programming Language Processing: An Investigation into Algorithm Detection](https://mlsanity.com/paper/2202.12481/) (2022)
- [A Library for Representing Python Programs as Graphs for Machine Learning](https://mlsanity.com/paper/2208.07461/) (2022)

## Built on by (most cited first, 3 of 3)

- [LLM Test Generation via Iterative Hybrid Program Analysis](https://mlsanity.com/paper/2503.13580/) (2025)
- [Refactoring with LLMs: Bridging Human Expertise and Machine Understanding](https://mlsanity.com/paper/2510.03914/) (2025)
- [Bridging Code Graphs and Large Language Models for Better Code Understanding](https://mlsanity.com/paper/2512.07666/) (2025)

## Abstract (arXiv)

Learning effective representations of source code is critical for any Machine Learning for Software Engineering (ML4SE) system. Inspired by natural language processing, large language models (LLMs) like Codex and CodeGen treat code as generic sequences of text and are trained on huge corpora of code data, achieving state of the art performance on several software engineering (SE) tasks. However, valid source code, unlike natural language, follows a strict structure and pattern governed by the underlying grammar of the programming language. Current LLMs do not exploit this property of the source code as they treat code like a sequence of tokens and overlook key structural and semantic properties of code that can be extracted from code-views like the Control Flow Graph (CFG), Data Flow Graph (DFG), Abstract Syntax Tree (AST), etc. Unfortunately, the process of generating and integrating code-views for every programming language is cumbersome and time consuming. To overcome this barrier, we propose our tool COMEX - a framework that allows researchers and developers to create and combine multiple code-views which can be used by machine learning (ML) models for various SE tasks. Some salient features of our tool are: (i) it works directly on source code (which need not be compilable), (ii) it currently supports Java and C#, (iii) it can analyze both method-level snippets and program-level snippets by using both intra-procedural and inter-procedural analysis, and (iv) it is easily extendable to other languages as it is built on tree-sitter - a widely used incremental parser that supports over 40 languages. We believe this easy-to-use code-view generation and customization tool will give impetus to research in source code representation learning methods and ML4SE. Tool: https://pypi.org/project/comex - GitHub: https://github.com/IBM/tree-sitter-codeviews - Demo: https://youtu.be/GER6U87FVbU