Python玩转C语言AST：用Tree-sitter写一个简易的“代码相似度检测”工具

Python玩转C语言AST用Tree-sitter构建代码相似度检测器在代码审查、作业查重或开源项目分析场景中文本层面的简单比对往往会产生大量误判——变量重命名、注释调整或格式修改都会导致传统diff工具失效。而抽象语法树AST作为代码的结构化表示能穿透表面差异直击逻辑本质。本文将演示如何用PythonTree-sitter构建一个能识别C代码结构相似性的智能工具其核心思路是将代码解析为AST并提取关键节点特征设计抗干扰的相似度计算算法输出可量化的相似度评分1. 环境配置与AST解析1.1 安装Tree-sitter全家桶pip install tree-sitter tree-sitter-c需要预编译C语言语法规则以Linux/macOS为例git clone https://github.com/tree-sitter/tree-sitter-c python -m tree_sitter build --outputbuild/my-languages.so tree-sitter-c验证安装是否成功from tree_sitter import Language, Parser LANG_C Language(build/my-languages.so, c) parser Parser() parser.set_language(LANG_C) # 若无报错则配置成功1.2 代码到AST的转换准备测试文件example.c// 示例A计算斐波那契数列 int fib(int n) { if (n 1) return n; return fib(n-1) fib(n-2); } // 示例B阶乘计算与A有相似结构 int fact(int x) { if (x 0) return 1; return x * fact(x-1); }解析代码生成ASTdef parse_code(file_path): with open(file_path, rb) as f: return parser.parse(f.read()) tree parse_code(example.c) print(tree.root_node.sexp()) # 查看AST的S表达式表示关键节点类型说明节点类型对应代码结构相似度分析价值function_definition函数定义★★★★★if_statement条件判断★★★★☆binary_expression二元运算(/-/*等)★★★☆☆call_expression函数调用★★★★☆return_statement返回语句★★☆☆☆2. AST特征提取策略2.1 节点类型序列指纹提取所有节点的类型名称构成特征序列def extract_node_types(node, types[]): if node.type not in [comment, , ;]: # 过滤噪音节点 types.append(node.type) for child in node.children: extract_node_types(child, types) return types types_a extract_node_types(parse_code(fib.c).root_node) types_b extract_node_types(parse_code(fact.c).root_node)得到的特征序列示例[function_definition, primitive_type, identifier, parameter_list, compound_statement, if_statement, ...]2.2 关键子树哈希对特定节点生成结构哈希以if语句为例import hashlib def hash_if_statement(node): if node.type ! if_statement: return None components [ node.child_by_field_name(condition).text, node.child_by_field_name(consequence).text ] return hashlib.md5(|.join(components).encode()).hexdigest() # 收集所有if语句哈希 if_hashes [hash_if_statement(n) for n in tree.root_node.descendants_of_type(if_statement)]3. 相似度算法实现3.1 基于编辑距离的序列比对from Levenshtein import distance def sequence_similarity(seq_a, seq_b): max_len max(len(seq_a), len(seq_b)) return 1 - distance(seq_a, seq_b) / max_len sim_score sequence_similarity(types_a, types_b)3.2 基于集合的重叠度计算def jaccard_similarity(set_a, set_b): intersection set_a set_b union set_a | set_b return len(intersection) / len(union) # 比较函数调用关系 calls_a {n.text for n in tree_a.descendants_of_type(call_expression)} calls_b {n.text for n in tree_b.descendants_of_type(call_expression)} call_sim jaccard_similarity(calls_a, calls_b)3.3 综合评分算法def overall_similarity(tree1, tree2): weights { node_sequence: 0.4, control_flow: 0.3, data_flow: 0.2, api_usage: 0.1 } # 各维度子评分实际实现需补充细节 scores { node_sequence: sequence_similarity(...), control_flow: compare_control_flow(...), data_flow: analyze_data_dependencies(...), api_usage: jaccard_similarity(...) } return sum(w*scores[k] for k,w in weights.items())4. 实战优化技巧4.1 抗干扰处理标准化策略移除所有标识符名称差异统一空白字符和格式化忽略注释和预处理指令def normalize_code(code): # 用占位符替换所有标识符 normalized re.sub(r[a-zA-Z_]\w*, VAR, code) # 移除注释 normalized re.sub(r//.*?$|/\*.*?\*/, , normalized, flagsre.DOTALL) return normalized4.2 性能优化增量分析对大型代码库先进行文件级粗筛对候选对再作精细AST比对并行处理from concurrent.futures import ThreadPoolExecutor def batch_compare(file_pairs): with ThreadPoolExecutor() as executor: results list(executor.map( lambda p: compare_files(p[0], p[1]), file_pairs)) return results4.3 结果可视化生成差异报告示例def generate_report(file_a, file_b, score): print(f## 代码相似度分析报告) print(f- 文件A: {file_a}) print(f- 文件B: {file_b}) print(f\n**综合相似度**: {score:.1%}) print(\n### 关键相似点) print(1. 相同的控制流结构if-return模式) print(2. 近似的递归调用方式) print(3. 匹配的参数传递模式)5. 进阶方向5.1 机器学习增强使用AST节点序列训练Siamese网络基于图神经网络处理AST拓扑结构# 伪代码示例PyTorch模型架构 class ASTSimModel(nn.Module): def __init__(self): super().__init__() self.embed nn.Embedding(num_node_types, 128) self.lstm nn.LSTM(128, 256) self.head nn.Sequential( nn.Linear(512, 128), nn.ReLU(), nn.Linear(128, 1), nn.Sigmoid()) def forward(self, ast1, ast2): h1 self.lstm(self.embed(ast1))[1][0] h2 self.lstm(self.embed(ast2))[1][0] return self.head(torch.cat([h1, h2], dim-1))5.2 多语言支持通过Tree-sitter的多语言能力扩展工具languages { c: Language(build/my-languages.so, c), python: Language(build/my-languages.so, python), java: Language(build/my-languages.so, java) } def parse_with_lang(file_path, lang): parser Parser() parser.set_language(languages[lang]) return parser.parse(open(file_path, rb).read())在实际项目中这种AST级相似度检测器比传统文本比对更能识别出以下典型场景变量重命名后的逻辑等价代码不同格式化风格的同源代码经过简单语法转换的衍生代码保留核心逻辑的非实质修改