基于結(jié)構(gòu)和語義的代碼分類以及聚類方法

金巖磊 秦冠軍 姜 凱 甘 迪 史志成 周 宇*

1(南京南瑞繼保電氣有限公司 江蘇 南京 211106) 2(南京航空航天大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)學(xué)院 江蘇 南京 211106)

0 引 言

如何利用深度學(xué)習(xí)對(duì)代碼進(jìn)行高質(zhì)量的特征提取是近幾年研究的一個(gè)熱點(diǎn)問題。代碼語言與自然語言有很多相似的特征,都是由單詞組成的,都可以生成AST的樹形結(jié)構(gòu),二者最大的不同就是代碼元素之間的依賴關(guān)系更強(qiáng),且代碼段的長(zhǎng)度遠(yuǎn)遠(yuǎn)超過自然語句的長(zhǎng)度,因此如何提取代碼中的結(jié)構(gòu)信息一直是軟件工程領(lǐng)域研究的熱點(diǎn)和難點(diǎn)。Tai等[1]提出了TreeLSTM模型,該模型的核心思想是通過將LSTM[2]的鏈狀結(jié)構(gòu)拓展為樹狀結(jié)構(gòu)來提取自然語言的結(jié)構(gòu)信息,軟件工程領(lǐng)域內(nèi)的很多模型將Tree-LSTM作為它們模型的一個(gè)子模塊,例如Wei等[3]就將Tree-LSTM作為克隆檢測(cè)模型CDLH的子模塊來協(xié)助提取代碼的結(jié)構(gòu)信息,Wan等[4]則借助Tree-LSTM。Alon等[5-6]通過提取AST的路徑,利用路徑的中間節(jié)點(diǎn)提取代碼的結(jié)構(gòu)信息,路徑的兩端節(jié)點(diǎn)提取代碼的語義信息,以此進(jìn)行方法名生成、變量名預(yù)測(cè)工作。Ou等[7]、Allamanis等[8]則使用圖網(wǎng)絡(luò)對(duì)代碼進(jìn)行特征提取。

本文將基于卷積和循環(huán)神經(jīng)網(wǎng)絡(luò)生成高質(zhì)量的代碼向量完成代碼分類以及聚類任務(wù)。該任務(wù)的主要目的是根據(jù)代碼的結(jié)構(gòu)和語義信息給出對(duì)應(yīng)的代碼標(biāo)記。自動(dòng)代碼分類能夠推動(dòng)軟件工程領(lǐng)域內(nèi)眾多任務(wù)的發(fā)展,如程序理解(program comprehension)[9]、概念定位(concept location)[10]、代碼剽竊檢測(cè)(code plagiarism detection)[11]、惡意軟件檢測(cè)(malware detection)等。代碼分類的標(biāo)簽可以看作代碼內(nèi)部含義的一個(gè)總結(jié),該標(biāo)簽?zāi)軌騾f(xié)助后續(xù)代碼模塊化、分析、重用等工作的進(jìn)行。當(dāng)下代碼分類的模型主要以深度學(xué)習(xí)模型為主,例如:Zhang等[12]設(shè)計(jì)的基于遞歸以及循環(huán)神經(jīng)網(wǎng)絡(luò)的ASTNN模型;Mou等[13]構(gòu)建的基于自定義樹卷積的TBCNN模型。除了借助AST提取代碼信息之外,Ben-Num等[14]則借助LLVM獲得代碼的中間表示并結(jié)合循環(huán)神經(jīng)網(wǎng)絡(luò)對(duì)代碼進(jìn)行特征提取。通過中間表示可以生成代碼的程序依賴圖。然而中間表示一個(gè)明顯的弊端就是無法獲得變量最原始定義的標(biāo)識(shí)符,這類模型在強(qiáng)化結(jié)構(gòu)信息提取的同時(shí)卻丟失了代碼的語義信息。該領(lǐng)域使用最廣泛的數(shù)據(jù)集就是Mou等[13]整理的104代碼分類數(shù)據(jù)集,該數(shù)據(jù)集共有104個(gè)類,每個(gè)類別下有500條代碼。當(dāng)下在該數(shù)據(jù)集取得最高分類精度的就是ASTNN模型,分類的準(zhǔn)確率達(dá)到98.2%。本文在后面的實(shí)驗(yàn)也將與ASTNN模型進(jìn)行對(duì)比。

與很多學(xué)者的研究方法類似,本文也借助AST來提取代碼的特征,通過對(duì)Zhang等提出的ASTNN模型以及Mou等提出的TBCNN模型進(jìn)行結(jié)合,將ASTNN中的遞歸神經(jīng)網(wǎng)絡(luò)替換成TBCNN中自定義的速度更快的卷積網(wǎng)絡(luò),構(gòu)建了一個(gè)基于卷積和循環(huán)神經(jīng)網(wǎng)絡(luò)的深度學(xué)習(xí)模型CVRNN,并基于該模型對(duì)代碼分類展開研究。該模型在與ASTNN模型分類準(zhǔn)確率十分接近的情況下,速度卻是該模型的1.55倍。在我們之前相似代碼搜索任務(wù)上[15]該模型也得到了成功的應(yīng)用,此模型是原模型改進(jìn)后的版本。

在使用分類訓(xùn)練好的CVRNN模型對(duì)代碼進(jìn)行編碼,并將得到的代碼向量進(jìn)行降維打印后發(fā)現(xiàn),代碼功能越相似,對(duì)應(yīng)代碼向量之間的幾何距離就越短。由于K-means的損失函數(shù)是樣本與其所屬類的中心之間的距離總和,具體距離的定義是歐氏距離的平方。因此,本文還基于CVRNN生成的代碼向量利用K-means進(jìn)行了聚類實(shí)驗(yàn),聚類效果對(duì)比TBCNN有明顯的優(yōu)勢(shì),除了FMI指標(biāo)與ASTNN持平,其他指標(biāo)均高于ASTNN。

電力監(jiān)控系統(tǒng)是整個(gè)電力系統(tǒng)運(yùn)轉(zhuǎn)的核心部分,為輸電、變電、配電提供基礎(chǔ)性的服務(wù),若其因受到惡意代碼攻擊進(jìn)而導(dǎo)致系統(tǒng)癱瘓將會(huì)造成巨大的經(jīng)濟(jì)損失。隨著計(jì)算機(jī)技術(shù)的飛速進(jìn)步,電力監(jiān)控系統(tǒng)面臨惡意代碼攻擊的問題愈發(fā)嚴(yán)峻。在該領(lǐng)域主要使用可信軟件白名單對(duì)惡意代碼進(jìn)行防御[16],其中白名單維護(hù)所有可信軟件的版本號(hào)、Hash指紋等相關(guān)的參數(shù),通過比對(duì)這些參數(shù)是否匹配判斷該軟件是否是可信軟件。很明顯該方式是一個(gè)完全靜態(tài)的方法,一方面,這個(gè)白名單很難進(jìn)行實(shí)時(shí)維護(hù),另一方面,將導(dǎo)致所有未記錄在名單上的軟件都給攔截掉。本文的工作是解決該問題的一個(gè)探索性工作,未來我們將試圖將代碼壓縮成代碼向量,首先使用白名單進(jìn)行攔截,若在名單內(nèi),則可直接運(yùn)行,否則通過計(jì)算代碼向量之間的距離并設(shè)置一個(gè)閾值來判斷該代碼是否是惡意代碼,使系統(tǒng)更加智能化,使得未在白名單中出現(xiàn)的非惡意代碼也能夠正常運(yùn)行。矩陣的運(yùn)算速度極快,經(jīng)我們測(cè)試,即使一個(gè)千萬級(jí)別的代碼數(shù)據(jù)庫,也能在毫秒級(jí)別完成計(jì)算,在保證系統(tǒng)正常運(yùn)行的同時(shí)也保證系統(tǒng)的運(yùn)行效率。

1 方法介紹

代碼向量是本文進(jìn)行分類和聚類任務(wù)的基石,代碼向量具體的生成過程如圖1所示。首先使用srcml工具生成代碼的AST,然后通過先序遍歷獲得AST節(jié)點(diǎn)的展開序列,根據(jù)該序列應(yīng)用Word2vec生成AST節(jié)點(diǎn)的詞向量并存入到一個(gè)矩陣當(dāng)中,訓(xùn)練得到的詞向量能夠?qū)⒐?jié)點(diǎn)之間的某些潛藏關(guān)系編碼進(jìn)詞向量中,如“if”“while”“for”表示程序的控制流關(guān)系,編碼后得到的這三個(gè)詞向量在向量空間上也是彼此鄰近的。為了解決因AST規(guī)模過于龐大而引發(fā)的梯度消失問題,本文并沒有簡(jiǎn)單地將完整的AST作為模型的輸入,而是將AST切割成由循環(huán)以及條件子樹構(gòu)成的序列。圖2給出了具體的切割樣例,每個(gè)代碼塊都與一棵AST相對(duì)應(yīng),①號(hào)代碼塊將被切分成②③④⑤四個(gè)部分,其中②號(hào)代碼塊主要包含不屬于條件判斷以及循環(huán)執(zhí)行的剩余代碼,例如異常拋出語句沒有包含在任何循環(huán)和條件判斷的代碼塊中,因此將其放入到第一個(gè)代碼塊中,剩余代碼塊均屬于條件判斷或者循環(huán)執(zhí)行,它們AST對(duì)應(yīng)的根節(jié)點(diǎn)一定在{“if”,”while”,“for”}集合當(dāng)中。由此可以得到一個(gè)子樹序列,這個(gè)序列可以看作是源代碼的控制流。通過之前得到的詞向量矩陣,可以將這些子樹節(jié)點(diǎn)由原先的字符串形式映射成詞向量的形式,然后將該子樹序列輸入到我們?cè)O(shè)計(jì)的CVRNN(Convolutional and Recurrent Neural Network)模型中得到代碼的向量表示。

圖1 生成代碼向量流程

圖2 切割樣例

圖3展示了CVRNN模型的執(zhí)行過程,該模型使用TBCNN對(duì)輸入到模型中的AST進(jìn)行編碼,得到每棵AST的向量表示。圖4展示了TBCNN樹卷積的過程。該模型設(shè)置了一個(gè)固定深度的滑動(dòng)窗口,自上而下地遍歷AST中的所有節(jié)點(diǎn),圖4中滑動(dòng)窗口的深度是2,滑動(dòng)窗口內(nèi)的節(jié)點(diǎn)將被編碼成一個(gè)向量,如圖4左邊①號(hào)三角形區(qū)域中的節(jié)點(diǎn)將被編碼進(jìn)右邊所標(biāo)注的①號(hào)向量,②號(hào)三角形區(qū)域同理。具體的編碼公式如式(1)-式(5)所示。

圖3 CVRNN模型

圖4 TBCNN樹卷積

經(jīng)過樹卷積之后,原先AST中的特征將被提取到新生成的AST中,新AST中每個(gè)節(jié)點(diǎn)向量包含了一個(gè)區(qū)域內(nèi)的節(jié)點(diǎn)信息,該AST與原AST有相同的結(jié)構(gòu),之后再經(jīng)過動(dòng)態(tài)池化[16]以及一層全連接神經(jīng)網(wǎng)絡(luò)將AST轉(zhuǎn)化為一個(gè)固定長(zhǎng)度的代碼向量,因此子樹序列將轉(zhuǎn)化成一個(gè)向量序列[v1,v2,…,vn],在TBCNN之上再疊加一層雙向循環(huán)神經(jīng)網(wǎng)絡(luò)(Bidirectional Recurrent Neural Network,BiRNN)以提取代碼塊之間的序列信息。與單向RNN不同的是,BiRNN將從過去和未來兩個(gè)時(shí)間方向?qū)Ξ?dāng)前時(shí)間步進(jìn)行編碼,因此會(huì)綜合考慮代碼上文和下文的信息。BiRNN內(nèi)部神經(jīng)元選取的是LSTM,相對(duì)于GRU[17],LSTM多了一個(gè)記憶單元,能夠更好地保存過去的時(shí)間信息以緩解長(zhǎng)期依賴導(dǎo)致的梯度消失問題,最終每個(gè)時(shí)間步的輸出向量是兩個(gè)時(shí)間方向隱層單元拼接得到的向量,例如時(shí)間步t最終的輸出向量可以使用以下公式得出:

將每個(gè)時(shí)間步所生成的代碼向量ht都存放到矩陣中,然后使用最大池化提取每個(gè)特征維度的最大值作為該維度上的值,將矩陣壓縮成一個(gè)代碼向量。最大池化可結(jié)合圖5進(jìn)行理解,其中有陰影的圓圈表示這個(gè)維度上的最大值,空白圓圈則表示非最大值,依次選取當(dāng)前維度的最大值作為代碼向量在該維度的值。

圖5 最大池化

2 實(shí)驗(yàn)驗(yàn)證

2.1 數(shù)據(jù)集

代碼分類以及聚類均使用的是Mou等[13]公開的104代碼分類數(shù)據(jù)集。表1列出了使用葉子節(jié)點(diǎn)和不使用葉子節(jié)點(diǎn)的情況下,AST節(jié)點(diǎn)數(shù)目以及深度的最大值和平均值。從統(tǒng)計(jì)信息也可以看出原始的AST的結(jié)構(gòu)十分龐大,若不對(duì)其進(jìn)行切割而直接對(duì)整棵樹進(jìn)行提取必然會(huì)影響代碼特征提取的效果。

表1 數(shù)據(jù)集統(tǒng)計(jì)信息

2.2 評(píng)估指標(biāo)

2.2.1代碼分類指標(biāo)

Accuracy準(zhǔn)確率,簡(jiǎn)稱ACC:

式中:k表示代碼類別的數(shù)目,在該任務(wù)中,值為104;nTP,i表示第i個(gè)類別被正確分類的數(shù)目;n表示樣本總量。

2.2.2代碼聚類度量指標(biāo)

(1) Jaccard系數(shù),簡(jiǎn)稱JC:

(2) FM指數(shù),簡(jiǎn)稱FMI:

(3) ACC。計(jì)算公式與代碼分類準(zhǔn)確率公式相同。在計(jì)算該指標(biāo)的時(shí)候,聚類標(biāo)簽要與真實(shí)標(biāo)簽相統(tǒng)一,因此實(shí)驗(yàn)過程中還增加了一個(gè)標(biāo)簽映射的過程,這里映射方法使用的是Kuhn-Munkres算法。

2.3 CVRNN參數(shù)設(shè)置以及實(shí)驗(yàn)平臺(tái)

該模型核心配置參數(shù)如表2所示,使用葉子節(jié)點(diǎn)和不使用葉子節(jié)點(diǎn)在參數(shù)配置上唯一的不同就是詞向量的維度。因?yàn)椴皇褂萌~子節(jié)點(diǎn)所需要訓(xùn)練的詞匯表只有60,而使用葉子節(jié)點(diǎn)詞匯表擴(kuò)充至6 156。為了降低頻詞對(duì)模型的影響,具體的實(shí)驗(yàn)過程中只選取了詞頻大于等于5的單詞納入詞匯表。在分類模型的訓(xùn)練過程中,以交叉熵作為訓(xùn)練的損失函數(shù)并選取Adam作為優(yōu)化器。仿照Mou等[13]的工作,該模型同樣沒有設(shè)置正則化項(xiàng)來降低模型過擬合的風(fēng)險(xiǎn)。深度學(xué)習(xí)框架使用的是TensorFlow 1.14,操作系統(tǒng)是Ubuntu 16.04,GPU型號(hào)是RTX2060,CPU型號(hào)是AMD Ryzen 7 2700。

表2 模型參數(shù)

2.4 任 務(wù)

2.4.1代碼分類

該實(shí)驗(yàn)將104數(shù)據(jù)集按3∶1∶1劃分訓(xùn)練集、驗(yàn)證集和測(cè)試集,對(duì)比模型選取的是Ben-Num等[14]的inst2vec模型、Zhang等[12]的ASTNN模型,具體的實(shí)驗(yàn)結(jié)果如表3所示。因?yàn)閕nst2vec模型并沒有使用AST,所以表3直接展示了原論文中的實(shí)驗(yàn)結(jié)果。ASTNN原文中使用的AST生成工具是javalang,為了對(duì)比的公平性,表3中展示的結(jié)果是使用srcml作為AST生成工具重新訓(xùn)練得到的實(shí)驗(yàn)結(jié)果。兩個(gè)模型均訓(xùn)練了30輪,表3和表5分別展示了實(shí)驗(yàn)結(jié)果和訓(xùn)練時(shí)間?？梢钥闯?ASTNN與CVRNN分類的準(zhǔn)確率很接近,而CVRNN的運(yùn)行速度大約是其的1.55倍,此外還進(jìn)行了另一個(gè)與ASTNN的對(duì)比實(shí)驗(yàn),不使用AST的葉子節(jié)點(diǎn),僅使用srcml符號(hào)表中定義的節(jié)點(diǎn)對(duì)代碼進(jìn)行分類,表4和表5分別展示了實(shí)驗(yàn)結(jié)果和訓(xùn)練時(shí)間。CVRNN在速度和準(zhǔn)確率上都要優(yōu)于ASTNN模型。

表3 代碼分類測(cè)試集準(zhǔn)確率

表4 不使用葉子節(jié)點(diǎn)測(cè)試集代碼分類的準(zhǔn)確率

表5 30輪的訓(xùn)練時(shí)間

使用已訓(xùn)練的CVRNN對(duì)測(cè)試集中的代碼進(jìn)行編碼生成代碼向量(對(duì)應(yīng)圖3中的r向量)并使用TSNE進(jìn)行降維打印,打印結(jié)果如圖6所示,相同灰度的表示屬于同一個(gè)類。可以看出,屬于同種類別的代碼形成一個(gè)簇,功能相近的代碼經(jīng)過CVRNN編碼后得到的代碼向量之間的幾何距離也是彼此臨近的。

圖6 降維后測(cè)試集代碼向量

2.4.2代碼聚類

上面的代碼分類任務(wù)中,盡管測(cè)試集中的代碼與訓(xùn)練集樣本不重疊,然而兩個(gè)數(shù)據(jù)集代碼的類別相同,在實(shí)際的代碼分類場(chǎng)景下,測(cè)試集代碼的類別大部分與訓(xùn)練集不同,為了探究CVRNN模型是否能夠?qū)τ?xùn)練過程中未出現(xiàn)過類別的代碼進(jìn)行正確分類,我們借助K-means設(shè)計(jì)了代碼聚類實(shí)驗(yàn),并根據(jù)代碼的真實(shí)標(biāo)簽對(duì)聚類效果進(jìn)行度量。

聚類模型整體架構(gòu)如圖7所示。首先將104個(gè)類別的代碼數(shù)據(jù)庫分為兩部分,其中84個(gè)類別作為訓(xùn)練集,剩下的20個(gè)類別作為測(cè)試集。首先使用84個(gè)類別的代碼數(shù)據(jù)庫對(duì)CVRNN模型進(jìn)行訓(xùn)練,然后應(yīng)用訓(xùn)練好的CVRNN對(duì)20個(gè)類別構(gòu)成的代碼庫進(jìn)行編碼得到的代碼向量再使用TSNE進(jìn)行降維,將降維后的代碼向量輸入到K-means中進(jìn)行聚類,由此可以得到每個(gè)代碼的預(yù)測(cè)標(biāo)記。計(jì)算JC和FMI度量指標(biāo)不需要預(yù)測(cè)標(biāo)記和真實(shí)標(biāo)記之間實(shí)現(xiàn)統(tǒng)一,因此可以直接計(jì)算出這兩個(gè)度量指標(biāo)的值。由于準(zhǔn)確率ACC要統(tǒng)一真實(shí)標(biāo)簽和預(yù)測(cè)標(biāo)簽,例如得到的預(yù)測(cè)進(jìn)行聚類,聚類的數(shù)目設(shè)為20。標(biāo)簽是(0,0,1,2,2,2,2,2),而真實(shí)標(biāo)簽是(1,1,2,0,0,0,2,0),就要將預(yù)測(cè)標(biāo)簽0映射為1,標(biāo)簽1映射為2,標(biāo)簽2映射為0,再計(jì)算準(zhǔn)確率。這里使用的是Kuhn-Munkres算法完成標(biāo)簽的映射。具體的做法是將預(yù)測(cè)標(biāo)簽與真實(shí)標(biāo)簽進(jìn)行組合得到一個(gè)矩陣,這里測(cè)試集共有20類,因此可以得到一個(gè)20×20的矩陣,矩陣中的每個(gè)值表示該標(biāo)簽組合情況下數(shù)據(jù)的重合數(shù)量,例如下面得到的標(biāo)簽混合矩陣:

圖7 聚類整體架構(gòu)

第1行第2列數(shù)字為380,表示真實(shí)標(biāo)簽為1與預(yù)測(cè)標(biāo)簽為2的樣本重合數(shù)為380,然后基于該矩陣計(jì)算最低代價(jià)的映射。將預(yù)測(cè)標(biāo)簽與真實(shí)標(biāo)簽實(shí)現(xiàn)統(tǒng)一之后,便可以計(jì)算準(zhǔn)確率的值。對(duì)比模型選取的依然是TBCNN和ASTNN,因?yàn)镵-means聚類度量指標(biāo)可能會(huì)出現(xiàn)4到5百分點(diǎn)的波動(dòng),因此最終取10次實(shí)驗(yàn)結(jié)果的平均值來避免偶然性,結(jié)果如表6所示?？梢钥闯鯟VRNN在3個(gè)模型的對(duì)比過程中各項(xiàng)度量指標(biāo)均占有優(yōu)勢(shì)。圖8給出了CVRNN模型聚類后的區(qū)域劃分,同一灰度的點(diǎn)表示對(duì)應(yīng)的代碼向量屬于同一個(gè)類?？梢钥闯鐾粋€(gè)區(qū)域下的點(diǎn)的灰度大致相同。

表6 聚類結(jié)果(%)

圖8 CVRNN 聚類效果

3 結(jié) 語

本文通過卷積以及循環(huán)神經(jīng)網(wǎng)絡(luò)對(duì)代碼的結(jié)構(gòu)以及語義信息進(jìn)行提取來獲得高質(zhì)量的代碼向量并基于此對(duì)代碼分類以及聚類展開研究,中間還應(yīng)用Word2vec預(yù)訓(xùn)練AST節(jié)點(diǎn)的詞向量來增強(qiáng)其初始語義。在代碼分類以及聚類任務(wù)上對(duì)比當(dāng)下最前沿的幾個(gè)模型占有優(yōu)勢(shì)。實(shí)驗(yàn)選取的數(shù)據(jù)集是該領(lǐng)域內(nèi)公開的標(biāo)準(zhǔn)數(shù)據(jù)集,代碼也在GitHub上開源,方便后續(xù)工作的復(fù)現(xiàn)參考。因?yàn)樵撗芯可婕暗膶?shí)驗(yàn)所采用的數(shù)據(jù)均來源于Mou等[13]提供的104代碼分類數(shù)據(jù)集,然而該數(shù)據(jù)集經(jīng)過精細(xì)化的人工預(yù)處理,數(shù)據(jù)集的噪聲很少,而現(xiàn)實(shí)環(huán)境中很難有如此良好的數(shù)據(jù)集。因此在后面的工作中,我們還將對(duì)包含眾多噪音的數(shù)據(jù)集展開研究,提高模型的魯棒性。此外,該模型只在C/C++語言上進(jìn)行了測(cè)試,未來還將在其他語言上進(jìn)行嘗試。當(dāng)下不少的研究還制定了一系列啟發(fā)式規(guī)則將AST由原來的樹形結(jié)構(gòu)轉(zhuǎn)化成具有更強(qiáng)表達(dá)能力的圖的結(jié)構(gòu),未來我們也將在這一方向進(jìn)行深入的探索,根據(jù)AST的樹形結(jié)構(gòu),進(jìn)一步挖掘代碼元素之間的依賴關(guān)系,將原來的樹形結(jié)構(gòu)拓展為圖形結(jié)構(gòu),并設(shè)計(jì)構(gòu)建圖神經(jīng)網(wǎng)絡(luò)對(duì)代碼的特征進(jìn)行更深入的提取,將更多的代碼特征編碼進(jìn)生成的代碼向量中。

未來將對(duì)電力監(jiān)控領(lǐng)域內(nèi)的源代碼以及常見的惡意代碼進(jìn)行分析,使得模型更適合該領(lǐng)域內(nèi)代碼的特征,生成能夠提取該領(lǐng)域內(nèi)代碼特征的代碼向量,在保證電力系統(tǒng)穩(wěn)定運(yùn)行的前提下,實(shí)現(xiàn)對(duì)損害電力系統(tǒng)的惡意代碼進(jìn)行有效攔截。