code embedding研究系列一-基于token的embedding

NLP中embedding无处不在，embedding将大型稀疏向量转换为保留语义关系的低维空间。在源代码研究领域(源代码分类，源代码克隆检测，变量，方法等重命名等)也会需要用到embedding，即code embedding。

code embedding可以分成很多种：
1.把源代码文本当成纯文本进行embedding (text embedding)
2.把源代码经过词法分析过后生成的token序列进行embedding (token embedding)
3.利用源代码解析成的抽象语法树**(AST), 从AST上提取信息进行embedding
4.利用源代码的控制流图(CFG)**和数据依赖图(PDG)进行embedding (基于graph的embedding)

这里Zimin Chen等人做了一个关于code embedding的统计研究

这里列举一些应用了token embedding的论文：

1.Automated software vulnerability detection with machine learning

这篇文章结局的问题是c/c++代码的bug检测，任务则是二分类，代码分为Good和Buggy，检测的粒度是function (即分析代码中的某个function是Good还是Buggy)

采用的方法结合了build-based（先编译C代码）和source-based（直接分析源代码文本），build-based分析了编译后的代码的CFG（控制流图）,source-based则是基于token embedding的分析，这里重点讨论source-based (数据集中有一部分不能编译成功)。

数据集来源

数据集来自2部分：
1.Debian包
2.Github

数据集的统计信息如下：

项目	Debian	Github	Source
Train Good	429807	250436	754150
Train Buggy	14616	9556	31790
Test Good	66100	38500	94268
Test Buggy	1878	1220	3724
Val Good	55102	32109	94268
val Buggy	1857	1203	3724

其中Source为Debian和Github的总和

数据集预处理

1.数据下下来后首先要去除重复的function**(长得一模一样)**
2.用Clangstatic analyzer (SA) 分析，将分析的输出去除和漏洞，bug无关的警告。去除过后没有警告信息的被标记为Good，有信息的被标记为Buggy。

分类方法

获取token

将源代码文本用自定义文法分析器解析成token序列，比如：

for (int i = 0; i <= 10; ++i){
	arr[i] += i;
}

生成的token序列如下：

for int i = 0 ; i <= 10 ; ++ i { arr [ i ] += i ; }

生成的token分类为：字符串值， 单字符值， 字符数组， 数值， 操作符， names等， names又细分为key（if for等等），系统调用（malloc），类型（bool, int, char），变量名。变量名又会被都映射为同一个标识符（具体参考论文），string值，float值都会被映射。

向量化表示

上述映射过程后会生成一个token序列，之后便是要对每个token进行向量化，用到的方法有：
1.词袋模型
2.word2vec

这样便可以生成非监督下的embedding表示，一个word2vec的token表示二维化后如下：

分类模型

论文设计的如下：

word2vec

TextCNN Conv layer

TextCNN fully connected layer

Extra Tree Classifier

Bag of Words

在提取到token序列后，作者采用了3种策略来分类：
1.word2vec向量化 + 卷积提取特征 + 全连接层分类
2.word2vec向量化 + 卷积提取特征 + 树分类器分类
3.Bag of Words向量化 + 树分类器分类

实验结果

评测标准采用ROC AUC和PR AUC。结果如下：

Model	ROC AUC	P-R AUC
BoW + ET	0.85	0.44
word2vec + CNN	0.87	0.45
word2vec + CNN feat. + ET	0.87	0.49

build-based部分

这篇论文还有一部分关于build-based就不提了，大致过程就是根据汇编代码生成控制流图（CFG），一个简单的代码片段CFG如下：

CFG由一系列的基本块(block)构成，每个基本块相当于图的一个node，每个block包括一些汇编代码（opcode），这些opcode不包括分支语句。整个CFG的特征由2部分组成:

1.use-def矩阵
2.op-vec

use-def矩阵表示变量定义和使用，假设一个use-def矩阵matrix。如果一个局部变量在基本块i中定义，基本块j中调用，则matrix[j][i]设为i

op-vec，每个基本块用1个op-vec表示。这个op-vec是一个9维二进制向量。分别对应 conditional, arrogate,binary, bit binary, conversion, memory address, termination, vector operation, and other 9个特征，如果基本块中存在对应该特征的指令，则该特征对应的值为1，否则为0。

之后结合CFG的邻接矩阵，每个代码块的op-vec以及CFG的use-def矩阵构建1个116维向量（怎么构建的就看不太懂了）来作为一个function的向量表示。通过随机森林对116维向量进行分类。

2.CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching

腾讯安全科恩实验室首次提出了基于AI的二进制代码/源代码端到端匹配算法。大概就是给你一个二进制代码，找出尽可能与之匹配的源代码，对搞逆向的人来说是个好东西，最新论文研究成果也将应用于腾讯安全科恩实验室研发的代码检索工具BinaryAI。

研究背景如下：在给定二进制代码的情况下，逆向分析研究人员希望找到它对应的源代码，从而提升逆向分析的效率和准确率（以后IDA逆向出来的伪代码可读性会更强）。

当然，这里我们更加关注token embedding部分。

总体架构

论文的总体架构如下：
大致流程就是对一个c语言函数，把它编译形成的二进制代码与c语言代码本身形成一个pair，对这样一个（c语言，二进制）pair,分别把源代码和二进制代码向量化后，计算Triplet Loss，具体可以参考论文。

这里source code（源代码）被解析成一个很长的token序列。token序列被embedding后通过DPCNN后形成一个source语义向量，如下图：

而源代码的处理不止于此，还需要提取String特征和Integer特征，即源代码文本中所有的字符串的值和数值，对于下面的代码。
if条件中的(void *)0中的0和exit(-1)中的-1会组成0 -1这样的数值序列输入Integer LSTM中生成一个Integer向量。

而源代码中的字符串有fatal error in SolveMap_nproc（%p）以及bad input，对于string特征采用一个hierarchical-LSTM 模型处理string变量。来得到一个string向量。

把上述的token序列向量,integer向量以及string向量拼接在一起就形成了源代码向量，对于二进制代码向量的处理也类似，不同的是采用HBMP + GGNN + Set2Set 处理二进制代码本身（关于二进制代码就不展开了）。

至于为什么不采用AST处理源代码，论文中提到：
1.需要耗费更多时间
2.解析采用的算法不一定能成功处理整段代码，比如用gcc来解析论文中使用的数据集，只有10%的代码被成功解析。
3.需要更多的专家知识。

数据集

数据集地址
数据集分成了2部分，一部分由gcc-x64-O0编译后的二进制代码和源代码组成pairs，另一部分由clang-arm-O3编译后的二进制代码和源代码组成的pairs组成。每部分包含50000个source-binary pairs, 30000用来training, 10000用来validation，10000用来testing。

实验结果这里就不提了（都是大佬写的）。

总结

上述引用2篇论文，讲述了一下code embedding之中的token embedding技术。大致流程就是将源代码文法处理后的token序列作为源代码的表示，应用到SE（软件工程）任务。接下来我会分析下基于AST，CFG等的code embedding技术。内容稍有欠缺，欢迎大佬们前来指点。

参考文献

Harer J A , Kim L Y , Russell R L , et al. Automated software vulnerability detection with machine learning[J]. 2018.

Zeping Yu, Wenxin Zheng, Jiaqi Wang, Qiyi Tang, Sen Nie, Shi Wu. CodeCMR: Cross-Modal Retrieval For Function-Level Binary Source Code Matching

Chen Z , Monperrus M . A Literature Study of Embeddings on Source Code[J]. 2019.