LLVM一个简单的Pass

// Before :
//                            entry
//                              |
//                        ______v______
//                       |   switch    |
//                       |_____________|
//                              |
//              +---------------+
//              |         _____________
//              +------->|    case1    |-------+
//              |        |_____________|       |
//              +------->|    case2    |-------+
//              |        |_____________|       |
//              +------->|    case3    |-------+
//              |        |_____________|       |
//              +------->|    case4    |-------+
//              |        |_____________|       |
//              +------->|    case5    |-------+
//              |        |_____________|       |
//              +------->|   default   |-------+
//                       |_____________|       |
//                                             |
//                              +--------------+
//                              |
//                              v
//                           return
//
// After :
//                               entry
//                                 |
//                           ______v______
//          +-------(true)--|     if1     |
//          | +----(false)--|_____________|
//          | |              _____________  
//          | |             |     else    |
//          | | +---------->|_____________|------------------+
//          | | |            _____________                   |
//          | | +--(false)--|     if5     |                  |
//          | | +---------->|_____________|--(true)--------+ |
//          | | |            _____________                 | |
//          | | +--(false)--|     if4     |                | |
//          | | +---------->|_____________|--(true)------+ | |
//          | | |            _____________               | | |
//          | | +--(false)--|     if3     |              | | |
//          | | +---------->|_____________|--(true)----+ | | |
//          | | |            _____________             | | | |
//          | | +--(false)--|     if2     |            | | | |
//          | +------------>|_____________|--(true)--+ | | | |
//          |                _____________           | | | | |
//          +-------------->|     run1    |          | | | | |
//          +---------------|_____________|          | | | | |
//          |                _____________           | | | | |
//          +---------------|     run2    |<---------+ | | | |
//          |               |_____________|            | | | |
//          |                _____________             | | | |
//          +---------------|     run3    |<-----------+ | | |
//          |               |_____________|              | | |
//          |                _____________               | | |
//          +---------------|     run4    |<-------------+ | |
//          |               |_____________|                | |
//          |                _____________                 | |
//          +---------------|     run5    |<---------------+ |
//          |               |_____________|                  |
//          |                _____________                   |
//          +---------------|     run6    |<-----------------+
//          |               |_____________|
//          |              
//          +----------------------+           
//                                 |
//                                 v
//                               return
//                      
//        当前模块执行后的整体流程图如上
//    当前模块仅支持将LLVM识别的switch 语句修改成 if else 格式，
//    但是这种改动对IDA 没有太大影响

#include "stdafx.h"

#include "LowerSwitchInst.h"


static llvm::zoo::cl::opt_bool _lsi("LowerSwitchInst", "lsi", false, "改变Switch 语句成if else 格式");

static llvm::zoo::cl::opt_str _lsi_name("LowerSwitchInst", "lsi_name", "", "要替换的函数名字");


namespace zoollvm
{
    LowerSwitchInst::LowerSwitchInst()
    {
        _bWorked = false;
    }

    void LowerSwitchInst::ClearWork()
    {
        _bWorked = false;
    }

    bool LowerSwitchInst::IsWorked()
    {
        return _bWorked;
    }

    bool LowerSwitchInst::Run(llvm::Module &M, llvm::Function &F)
    {
        std::vector<llvm::BasicBlock*> _vecDelete;
        std::vector<llvm::SwitchInst*> _vecDeleteInst;
        //    删除列表

        ZooPrint("Function Name = [%s] \n", zoollvm::module::function::GetFunctionName(F).c_str());
        for (llvm::Function::iterator I = F.begin(), E = F.end(); I != E; I++)
        {
            llvm::BasicBlock &Cur = *I;
            auto i = Cur.getTerminator();
            if (i == NULL)
            {
                continue;
            }
            if (llvm::SwitchInst *SI = llvm::dyn_cast<llvm::SwitchInst>(i))
            {
                _bWorked = true;
                _vecDelete.push_back(&*I);
                _vecDeleteInst.push_back(SI);

                ZooPrint("SI = %p \n", SI);
                ZooPrint("    getOpcodeName : [%s] \n", SI->getOpcodeName());

                Lower(M, F, Cur, *SI);
                ZooPrint("\n");
            }
        }
        ZooPrint("Write Pass Success \n");
        for (auto it : _vecDeleteInst)
        {
            zoollvm::module::function::block::RemoveBlockSelf(it);
        }
        ZooPrint("Delete Inst Success \n");
        return true;
    }

    bool LowerSwitchInst::Lower(llvm::Module &M, llvm::Function &F, llvm::BasicBlock &B, llvm::SwitchInst &S)
    {
        llvm::Value *Val = S.getCondition();
        struct CaseRange
        {
            llvm::ConstantInt* cValue;                //
            llvm::BasicBlock* BB;                //    处理块
            CaseRange(llvm::ConstantInt *low, llvm::BasicBlock *bb) : cValue(low), BB(bb) {}
        };
        std::vector<CaseRange> Cases;
        ZooPrint("        Value = [%s] \n", Val->getName().str().c_str());
        ULONG NumSimpleCases = 0;
        //    找到所有case 块
        for (auto Case : S.cases())
        {
            if (Case.getCaseSuccessor() == S.getDefaultDest())
            {
                continue;
            }
            Cases.push_back(CaseRange(Case.getCaseValue(), Case.getCaseSuccessor()));
            ++NumSimpleCases;
        }
        ZooPrint("            NumSimpleCases = [%d] \n", NumSimpleCases);

        //    修改所有case 块
        llvm::BasicBlock *bb = &B;
        for (auto it : Cases)
        {
            int64_t nextValue = it.cValue->getSExtValue();

            //    创建一个基本块，目前这个基本块位置未知，只是先要有，用它占位
            auto ba = zoollvm::module::function::block::Create(B.getContext());

            //    创建一个逻辑指令
            //    v = (Val == nextValue )
            //    这个逻辑指令的位置，放在当前块的最后一条指令的位置，即   (llvm::BinaryOperator *)&*(bb->end() --)
            auto v = zoollvm::op_code::compare::CreateICmpEQ(Val, zoollvm::value::const_value::CreateInt64(M, nextValue), (llvm::BinaryOperator *)&*(bb->end() --));
            
            //    if v then exec else 空block end
            //    这个比较指令的位置，放在当前块的最后一条指令，结合上下文，就是放到上面的逻辑指令的下一条处，就是 bb 的 end 处
            zoollvm::op_code::branch::CreateCondBr(v, it.BB, ba, bb);

            //    插入刚才创建的基本块，位置是当前基本块的下一个位置，即顺序向下执行的话，就是当前新的基本块跟在上面的比较指令后面
            ba->insertInto(&F, B.getNextNode());

            //    修改要操作的基本块，下一个循环，下一条指令，都从这个新的基本块里面添加
            bb = ba;

            ZooPrint("                nextValue = [%I64d] \n", nextValue);
        }
        //    最后补一个跳到default 的指令，位置是在最初第一个基本块后面的第一个基本块中
        //    上面由于添加指令是从上往下加，而添加基本块是从下往上加，所以最后一个新的基本块在第一条指令的下面，这也是最后一个跳转的块
        zoollvm::op_code::branch::CreateBr(S.getDefaultDest(), B.getNextNode());


        return true;
    }


    class LowerSwitchInstPass : public llvm::ModulePass, public zoollvm::frame::Base
    {
    public:
        static char ID; // pass identification
        LowerSwitchInstPass() : ModulePass(ID), zoollvm::frame::Base("LowerSwitchInstPass")
        {
            zoollvm::opt::GetFunctionNameList(_lsi_name, _vec);
        }


        //    启动当前Pass
        virtual bool runOnModule(llvm::Module &M)
        {
            if (_lsi == false)
            {
                return true;
            }

            LowerSwitchInst lsi;

            for (auto it = M.begin(); it != M.end(); it++)
            {
                if (zoollvm::opt::IsFunctionInList(_vec, *it))
                {
                    lsi.Run(M, *it);
                }
            }

            return true;
        }

    private:
        std::vector<std::string> _vec;
    };
}

char zoollvm::LowerSwitchInstPass::ID = 0;
llvm::Pass * createLowerSwitchInst()
{
    return new zoollvm::LowerSwitchInstPass();
}

 

由于新版LLVM（9.0.0），无法直接调用Util 自带的 LowerSwitch 了，所以这里，我自己弄了一个，

主要功能是展开Switch，把它变成if else，理论上其实没啥区别，但是实际上在LLVM的角度上看，区别很大，

switch 实际上是一条 instruction ，后面的 case 是和switch 在一起的一条 instruction，

这就导致，如果不对它做处理，那么这个 instruction 是非常大的，由于这一条指令是在一个 block 里面的，所以结构也清晰，

if else 则不同，一片 if else 实际上是一片 block ，不管做什么操作也稍微更容易一些