RAWL: Retrieval-Augmented Smart Contract Repair using Weisfeiler-Lehman-Based Retrieval

Abstract

Smart contracts are at risk of significant financial losses due to vulnerabilities in their program logic. To mitigate such damages, various automated program repair tools have been proposed. However, existing tools, especially those based on Large Language Models (LLMs), may introduce unintended changes to the original functionality even when they successfully remove vulnerabilities. To improve the repair performance of LLM-based approaches while preserving functionality, we propose RAWL, a Retrieval-Augmented Generation (RAG)-based approach for smart contract vulnerability repair. RAWL incorporates a Weisfeiler-Lehman (WL) algorithm-based Abstract Syntax Tree (AST) retrieval mechanism as its RAG component. By applying the WL algorithm to ASTs, RAWL encodes the structural information of the vulnerability-related portion of the AST into a compact representation, enabling the retrieval of structurally relevant contracts as repair references. We evaluate RAWL and a non-RAG-based approach on 43 vulnerable contracts, executing each approach five times per contract, resulting in 215 generated outputs for both methods. RAWL produced 203 valid patches (93%) and at least one valid patch for all 43 contracts (100%), whereas the baseline approach produced 170 valid patches (79%) and at least one valid patch for 40 contracts (93%). The results indicate that RAWL effectively leveraged retrieved contracts when necessary, achieving higher repair performance than the non-RAG-based approach.