Don't Transform the Code, Code the Transforms: Towards Precise Code Rewriting using LLMs (2410.08806v1)

Abstract: Tools for rewriting, refactoring and optimizing code should be fast and correct. LLMs, by their nature, possess neither of these qualities. Yet, there remains tremendous opportunity in using LLMs to improve code. We explore the use of LLMs not to transform code, but to code transforms. We propose a chain-of-thought approach to synthesizing code transformations from a small number of input/output code examples that incorporates execution and feedback. Unlike the direct rewrite approach, LLM-generated transformations are easy to inspect, debug, and validate. The logic of the rewrite is explicitly coded and easy to adapt. The compute required to run code transformations is minute compared to that of LLM rewriting. We test our approach on 16 Python code transformations and find that LLM- generated transforms are perfectly precise for 7 of them and less imprecise than direct LLM rewriting on the others. We hope to encourage further research to improving the precision of LLM code rewriting.

• The paper introduces a novel LLM-based technique that synthesizes transformation functions from input/output examples, achieving a precision of 0.95.
• It leverages iterative chain-of-thought reasoning and feedback loops—up to 50 iterations—to refine transformation logic for accurate AST rewrites.
• The approach enhances software maintenance by producing transparent, inspectable transformation functions that simplify debugging and code refactoring.

Overview of "Don't Transform the Code, Code the Transforms"

The paper presents a novel methodology leveraging LLMs for generating precise code transformations. Moving away from direct code rewriting, the authors propose synthesizing transformation functions using LLMs through a chain-of-thought approach applied to input/output code examples.

Core Contributions

The authors introduce a method where LLMs are tasked with generating code transformations instead of performing the transformations themselves. This approach offers several benefits, including easy inspection, debugging, and validation of transformations, as compared to the opaque processes typical of direct LLM rewrites. The paper focuses on rewriting Python Abstract Syntax Trees (ASTs) with functions generated through a structured process involving execution and feedback.

Methodology

The approach is systematically structured and involves:

1. Providing LLMs with a small number of input/output examples.
2. Encouraging iterations to refine transformation logic using feedback, increasing up to 50 iterations if needed.
3. Generating code transformation implementations based on refined logic.
4. Executing these transformations in a controlled environment to validate their correctness.
5. Utilizing an introspection mechanism to diagnose and refine failed transformations.

This process aims to codify transformations that are both precise and adaptable. The authors emphasize that this method minimizes computation compared to direct LLM code rewriting.

Experimental Evaluation

The experimentation involves comparing the "Transform the Code" (TTC) approach with the proposed "Code the Transform" (CTT) across several Python code transformations. The dataset includes 480 input/output examples representing 16 transformation classes, ranging from simple arithmetic optimizations to complex API replacements.

• Precision and Recall: CTT significantly outperforms TTC in terms of precision (0.95 vs. 0.60), while recall remains high for both.
• Iterations for Transformation: Simple transformations require fewer iterations to achieve correctness compared to more complex ones.
• Impact of Chain-of-Thought: Removing introspection or stepwise problem description degrades performance, highlighting the efficacy of the proposed method.

Implications

This research suggests that having LLMs generate transformation functions rather than directly modifying code leads to higher precision with a more systematic and inspectable process. Such an approach could significantly enhance applications in software maintenance, optimization, and refactoring by reducing manual intervention in code review and validation processes.

Future Directions

There are promising avenues for refining this work. Integrating reinforcement learning or fine-tuning may enhance model performance in generating more sophisticated transformations. Additionally, exploring different domains beyond Python and scaling to larger codebases could further test the robustness and applicability of the approach.

This paper paves the way for exploiting LLM capabilities in creating reliable, efficient, and automated code transformation tools, potentially transforming practices in software engineering and development.

In conclusion, the paper provides a compelling alternative to direct code rewriting, presenting insights into how LLMs can be harnessed to generate precise and manageable code transformations. This research might influence continued explorations into the applications of machine learning in advancing software engineering methodologies.