一款方便、高效的基于tree-sitter的代码风格转换器，支持Python、C/C++和Java共100多种风格转换

​
源代码网址：https://github.com/rebibabo/SCTS/tree/main

如果有引用本文或者工具，请注明引用来源

如果觉得对您有帮助，还请各位帅哥美女在github点个免费的star🙏

工具介绍

实现了一款多语言代码等价语义转换器，基于tree-sitter开发工具包，对代码进行解析，并在具体语法树（concrete syntax tree, CST）树上进行节点的匹配和替换，支持Python、C/C++和Java共100多种风格转换，处理2000份代码仅需要13.6s。

代码风格转换器可以应用在以下的场景

• 后门/对抗攻击：实现隐蔽性高的攻击
• 数据集扩展：提高模型性能和鲁棒性，使模型更好的对抗扰动
• 模型水印：向数据集的代码注入水印，保护模型或数据集的产权
• 统一风格：增强代码的可读性
• 代码混淆：防止反编译
  ​

现有代码风格转换器方法可以大致分为两类：
语法分析器

• 优点：精确匹配语法规则，保证语义等价和语法正确，转换速度快，能够大批量处理数据集
• 缺点：需要设计转换规则，比较费时

深度学习模型

• 优点：不需要指定转换规则，只需要通过数据集训练模型即可，转换成功率高
• 缺点：token长度限制，生成的代码语法可能错误，并且生成速度非常慢。

本工具收集整理了四种语言C，C++，Python和Java四种语言的112条语义等价变换，设计了增删算法，能够高效处理大批量代码，提升转换速率。

环境搭建

运行命令pip install -r requirements.txt安装python环境，如果想可视化CST树，请安装graphviz，CST树样例如下

linux请运行sudo apt-get install git graphviz graphviz-doc安装，windows系统请从官网上下载，并配置环境变量
如果不想可视化CST树，可以注释掉所有和graphviz相关的代码

使用教程

构建scts类，设置代码的语言种类

from SCTS.change_program_style import SCTS
scts = SCTS('c')

然后调用change_file_style方法，第一个参数为风格列表，第二个参数为代码文本，返回值第一个元素为风格转换后的代码，第二个元素为转换是否成功，判断依据为代码是否发生变化且语法正确。

new_code, succ = scts.change_file_style(["8.11", "0.1"], code)

例如可以将C语言的for循环修改风格，将初始化语句提到for循环前，条件语句移到循环体内部，自增语句放在循环体末尾。

其他的风格转换样例

如果想要可视化CST树，执行以下代码

scts.see_tree(code)

如果想将代码进行分词，返回分词列表，执行下面的代码

scts.tokenize(code)

工具框架流程图

工具框架如下图所示，初始化模块输入代码的语言，根据指定语言构建对应的解析器和风格字典的键为风格名称，值为该风格的规则，根据每种编程语言的特点，构建相应的风格字典。风格字典中的键是风格名称，而值则是对应风格的转换规则。例如，对于C语言，风格名称可以是“0.1”，其中小数点左边的数字“0”表示该风格所属的类型，即“修改变量名”，而小数点右边的数字“1”表示该类型的子风格，即“驼峰命名法”。这样就构建了一个风格规则：“0.1: (‘val’, ‘camel’)”，表示将变量名转换为驼峰命名法。风格字典就是由该语言设计的所有风格组成。再输入源代码，经过格式化代码之后，生成CST树，最后输入风格列表，例如[0.1, 1.1, 3.2]，然后查询风格字典，得到所有风格的匹配规则和替换规则。

匹配模块：输入风格列表中所有风格的匹配规则，并初始化匹配节点列表为空，该列表存放所有匹配规则成功的节点。接着输入代码的CST树，递归遍历CST的每一个节点，如果该节点不符合匹配规则，如图中红色实线所示，则遍历该节点的所有子节点，如果符合匹配规则，则将其加入到匹配节点中，根据不同风格的匹配规则，可能继续嵌套的遍历子节点或者不继续往下遍历，直到遍历完所有CST节点，匹配阶段结束，最终输出匹配节点列表。

替换模块：输入风格列表中风格的替换规则以及匹配节点列表，对列表中的每一个节点，根据该风格对应的替换规则，得到修改操作（将在下一节中详细介绍），并将其加入到修改列表中，遍历完成后，根据修改列表，在源代码上直接进行修改，得到该风格转换后的代码。如果还有风格需要修改，则跳回到初始化模块，查询下一个风格的匹配规则和替换规则，重复上述的步骤，直到所有的风格都转换完毕，输出最终的修改后的代码，并检查语法，如果代码发生了改变并且语法正确，则转换成功，否则转换失败。

增删算法

tree-sitter编辑CST树的成本比较高，而且操作不方便，容易出错，因此考虑对源代码直接进行替换，而不是编辑CST树。在替换模块中，对源代码进行修改操作可以分为插入操作和删除操作两类，通过这两种操作的组合实现替换操作。

为了形式化定义这两种操作，假设原始代码为x，定义插入操作为 I(x, i, s)，表示在字节偏移量为i处的右边插入字符串s，则插入操作如公式3-1所示,而删除操作定义为 D(x, i, j)，表示在字节偏移量为i处，向左删除i个字节：

$I(x,i,s)=x[:i]+s+x[i+1:]$
$D\left(x,i,j\right)=x[:$