[lifting] Allow defining llvm function name

JonathanSalwan · Feb 14, 2022 · 4e42a22 · 4e42a22
1 parent 05b05cf
commit 4e42a22
Show file tree

Hide file tree

Showing 8 changed files with 53 additions and 33 deletions.
diff --git a/.build_number b/.build_number
@@ -1 +1 @@
-1518
+1519
diff --git a/src/libtriton/api/api.cpp b/src/libtriton/api/api.cpp
@@ -1373,17 +1373,17 @@ namespace triton {
 
   /* Lifters engine API ================================================================================= */
 
-  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node) {
+  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname) {
     this->checkLifting();
     #ifdef TRITON_LLVM_INTERFACE
-    return this->lifting->liftToLLVM(stream, node);
+    return this->lifting->liftToLLVM(stream, node, fname);
     #endif
     throw triton::exceptions::API("API::liftToLLVM(): Triton not built with LLVM");
   }
 
 
-  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr) {
-    return this->liftToLLVM(stream, expr->getAst());
+  std::ostream& API::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname) {
+    return this->liftToLLVM(stream, expr->getAst(), fname);
   }
 
 

diff --git a/src/libtriton/bindings/python/objects/pyTritonContext.cpp b/src/libtriton/bindings/python/objects/pyTritonContext.cpp
@@ -286,8 +286,8 @@ Returns true if the taint engine is enabled.
 - <b>bool isThumb(void)</b><br>
 Returns true if execution mode is Thumb (only valid for ARM32).
 
-- <b>string liftToLLVM(\ref py_AstNode_page node)</b><br>
-Lifts an AST node and all its references to LLVM IR.
+- <b>string liftToLLVM(\ref py_AstNode_page node, string fname="__triton")</b><br>
+Lifts an AST node and all its references to LLVM IR. `fname` is the name of the LLVM function, by default it's `__triton`.
 
 - <b>string liftToLLVM(\ref py_SymbolicExpression_page expr)</b><br>
 Lifts a symbolic expression and all its references to LLVM IR.
@@ -2231,17 +2231,37 @@ namespace triton {
       }
 
 
-      static PyObject* TritonContext_liftToLLVM(PyObject* self, PyObject* arg) {
-        if (!PySymbolicExpression_Check(arg) && !PyAstNode_Check(arg))
+      static PyObject* TritonContext_liftToLLVM(PyObject* self, PyObject* args, PyObject* kwargs) {
+        PyObject* node  = nullptr;
+        PyObject* fname = nullptr;
+
+        static char* keywords[] = {
+          (char*)"node",
+          (char*)"fname",
+          nullptr
+        };
+
+        /* Extract keywords */
+        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OO", keywords, &node, &fname) == false) {
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Invalid number of arguments");
+        }
+
+        if (node == nullptr || (!PySymbolicExpression_Check(node) && !PyAstNode_Check(node)))
           return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a SymbolicExpression or a AstNode as first argument.");
 
+        if (fname != nullptr && !PyStr_Check(fname))
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a string as fname parameter.");
+
+        if (fname == nullptr)
+          fname = PyStr_FromString("__triton");
+
         try {
           std::ostringstream stream;
-          if (PySymbolicExpression_Check(arg)) {
-            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PySymbolicExpression_AsSymbolicExpression(arg));
+          if (PySymbolicExpression_Check(node)) {
+            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PySymbolicExpression_AsSymbolicExpression(node), PyStr_AsString(fname));
           }
           else {
-            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PyAstNode_AsAstNode(arg));
+            PyTritonContext_AsTritonContext(self)->liftToLLVM(stream, PyAstNode_AsAstNode(node), PyStr_AsString(fname));
           }
           return xPyString_FromString(stream.str().c_str());
         }
@@ -3405,7 +3425,7 @@ namespace triton {
         {"isSymbolicExpressionExists",          (PyCFunction)TritonContext_isSymbolicExpressionExists,                  METH_O,                        ""},
         {"isTaintEngineEnabled",                (PyCFunction)TritonContext_isTaintEngineEnabled,                        METH_NOARGS,                   ""},
         {"isThumb",                             (PyCFunction)TritonContext_isThumb,                                     METH_NOARGS,                   ""},
-        {"liftToLLVM",                          (PyCFunction)TritonContext_liftToLLVM,                                  METH_O,                        ""},
+        {"liftToLLVM",                          (PyCFunction)(void*)(PyCFunctionWithKeywords)TritonContext_liftToLLVM,  METH_VARARGS | METH_KEYWORDS,  ""},
         {"liftToPython",                        (PyCFunction)TritonContext_liftToPython,                                METH_O,                        ""},
         {"liftToSMT",                           (PyCFunction)TritonContext_liftToSMT,                                   METH_VARARGS,                  ""},
         {"newSymbolicExpression",               (PyCFunction)TritonContext_newSymbolicExpression,                       METH_VARARGS,                  ""},

diff --git a/src/libtriton/engines/lifters/liftingToLLVM.cpp b/src/libtriton/engines/lifters/liftingToLLVM.cpp
@@ -22,21 +22,21 @@ namespace triton {
       }
 
 
-      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr) {
-        this->liftToLLVM(stream, expr->getAst());
+      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname) {
+        this->liftToLLVM(stream, expr->getAst(), fname);
         return stream;
       }
 
 
-      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node) {
+      std::ostream& LiftingToLLVM::liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname) {
         /* The LLVM context */
         llvm::LLVMContext context;
 
         /* The lifter Triton -> LLVM */
         triton::ast::TritonToLLVM lifter(context);
 
         /* Lift AST to LLVM IR */
-        auto llvmModule = lifter.convert(node);
+        auto llvmModule = lifter.convert(node, fname);
 
         /* Print the LLVM module into the stream */
         std::string dump;

diff --git a/src/libtriton/includes/triton/api.hpp b/src/libtriton/includes/triton/api.hpp
@@ -672,11 +672,11 @@ namespace triton {
 
         /* Lifters engine API ================================================================================= */
 
-        //! [**lifting api**] - Lifts an AST and all its references to LLVM format.
-        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node);
+        //! [**lifting api**] - Lifts an AST and all its references to LLVM format. `fname` represents the name of the LLVM function.
+        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
 
-        //! [**lifting api**] - Lifts a symbolic expression and all its references to LLVM format.
-        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr);
+        //! [**lifting api**] - Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
+        TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton");
 
         //! [**lifting api**] - Lifts a symbolic expression and all its references to Python format.
         TRITON_EXPORT std::ostream& liftToPython(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr);

diff --git a/src/libtriton/includes/triton/liftingToLLVM.hpp b/src/libtriton/includes/triton/liftingToLLVM.hpp
@@ -52,11 +52,11 @@ namespace triton {
           //! Constructor.
           TRITON_EXPORT LiftingToLLVM();
 
-          //! Lifts a symbolic expression and all its references to LLVM format.
-          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr);
+          //! Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
+          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, const char* fname="__triton");
 
-          //! Lifts a abstract node and all its references to LLVM format.
-          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node);
+          //! Lifts a abstract node and all its references to LLVM format. `fname` represents the name of the LLVM function.
+          TRITON_EXPORT std::ostream& liftToLLVM(std::ostream& stream, const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
       };
 
     /*! @} End of lifters namespace */

diff --git a/src/libtriton/includes/triton/tritonToLLVM.hpp b/src/libtriton/includes/triton/tritonToLLVM.hpp
@@ -42,29 +42,29 @@ namespace triton {
     /*! \brief Converts a Triton's AST to LVM IR. */
     class TritonToLLVM {
       private:
-        //! The LLVM context
+        //! The LLVM context.
         llvm::LLVMContext& llvmContext;
 
-        //! The LLVM module
+        //! The LLVM module.
         std::shared_ptr<llvm::Module> llvmModule;
 
-        //! The LLVM IR builder
+        //! The LLVM IR builder.
         llvm::IRBuilder<> llvmIR;
 
-        //! Map Triton variables to LLVM ones
+        //! Map Triton variables to LLVM ones.
         std::map<triton::ast::SharedAbstractNode, llvm::Value*> llvmVars;
 
-        //! Create a LLVM function
+        //! Create a LLVM function. `fname` represents the name of the LLVM function.
         void createFunction(const triton::ast::SharedAbstractNode& node, const char* fname);
 
-        //! Converts Triton AST to LLVM IR
+        //! Converts Triton AST to LLVM IR.
         llvm::Value* do_convert(const triton::ast::SharedAbstractNode& node, std::unordered_map<triton::ast::SharedAbstractNode, llvm::Value*>* results);
 
       public:
         //! Constructor.
         TRITON_EXPORT TritonToLLVM(llvm::LLVMContext& llvmContext);
 
-        //! Lifts a symbolic expression and all its references to LLVM format.
+        //! Lifts a symbolic expression and all its references to LLVM format. `fname` represents the name of the LLVM function.
         TRITON_EXPORT std::shared_ptr<llvm::Module> convert(const triton::ast::SharedAbstractNode& node, const char* fname="__triton");
     };
 

diff --git a/src/testers/unittests/test_ast_representation.py b/src/testers/unittests/test_ast_representation.py
@@ -269,4 +269,4 @@ def test_lifting(self):
             ]
 
             for n in nodes:
-                self.assertNotEqual(len(self.ctx.liftToLLVM(n)), 0)
+                self.assertNotEqual(len(self.ctx.liftToLLVM(n, fname="test")), 0)