#1085: Simplify a node using LLVM

JonathanSalwan · Feb 15, 2022 · 96104b1 · 96104b1
1 parent fa049b7
commit 96104b1
Show file tree

Hide file tree

Showing 7 changed files with 81 additions and 27 deletions.
diff --git a/.build_number b/.build_number
@@ -1 +1 @@
-1520
+1521
diff --git a/src/libtriton/api/api.cpp b/src/libtriton/api/api.cpp
@@ -881,13 +881,15 @@ namespace triton {
   }
 
 
-  triton::ast::SharedAbstractNode API::simplify(const triton::ast::SharedAbstractNode& node, bool usingSolver) const {
-    this->checkSymbolic();
-
+  triton::ast::SharedAbstractNode API::simplify(const triton::ast::SharedAbstractNode& node, bool usingSolver, bool usingLLVM) const {
     if (usingSolver) {
       return this->simplifyAstViaSolver(node);
     }
+    else if (usingLLVM) {
+      return this->simplifyAstViaLLVM(node);
+    }
     else {
+      this->checkSymbolic();
       return this->symbolic->simplify(node);
     }
   }
@@ -1133,7 +1135,7 @@ namespace triton {
     #endif
     #ifdef TRITON_BITWUZLA_INTERFACE
     if (this->getSolver() == triton::engines::solver::SOLVER_BITWUZLA) {
-        return reinterpret_cast<const triton::engines::solver::BitwuzlaSolver*>(this->getSolverInstance())->evaluate(node);
+      return reinterpret_cast<const triton::engines::solver::BitwuzlaSolver*>(this->getSolverInstance())->evaluate(node);
     }
     #endif
     throw triton::exceptions::API("API::evaluateAstViaZ3(): Solver instance must be a SOLVER_Z3 or SOLVER_BITWUZLA.");
@@ -1398,4 +1400,13 @@ namespace triton {
     return this->lifting->liftToSMT(stream, expr, assert_);
   }
 
+
+  triton::ast::SharedAbstractNode API::simplifyAstViaLLVM(const triton::ast::SharedAbstractNode& node) const {
+    this->checkLifting();
+    #ifdef TRITON_LLVM_INTERFACE
+    return this->lifting->simplifyAstViaLLVM(node);
+    #endif
+    throw triton::exceptions::API("API::simplifyAstViaLLVM(): Triton not built with LLVM");
+  }
+
 }; /* triton namespace */
diff --git a/src/libtriton/bindings/python/objects/pyTritonContext.cpp b/src/libtriton/bindings/python/objects/pyTritonContext.cpp
@@ -369,9 +369,10 @@ Sets the targeted register as tainted or not. Returns true if the register is st
 - <b>void setThumb(bool state)</b><br>
 Sets CPU state to Thumb mode (only valid for ARM32).
 
-- <b>\ref py_AstNode_page simplify(\ref py_AstNode_page node, bool solver=False)</b><br>
+- <b>\ref py_AstNode_page simplify(\ref py_AstNode_page node, bool solver=False, bool llvm=False)</b><br>
 Calls all simplification callbacks recorded and returns a new simplified node. If the `solver` flag is
-set to True, Triton will use the current solver instance to simplify the given `node`.
+set to True, Triton will use the current solver instance to simplify the given `node`. If `llvm` is true,
+we use LLVM to simplify node.
 
 - <b>dict sliceExpressions(\ref py_SymbolicExpression_page expr)</b><br>
 Slices expressions from a given one (backward slicing) and returns all symbolic expressions as a dictionary of {integer SymExprId : \ref py_SymbolicExpression_page expr}.
@@ -1403,7 +1404,7 @@ namespace triton {
         }
 
         if (node == nullptr || !PyAstNode_Check(node)) {
-          return PyErr_Format(PyExc_TypeError, "TritonContext::getModel(): Expects a AstNode as argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::getModel(): Expects a AstNode as node argument.");
         }
 
         if (wb != nullptr && !PyBool_Check(wb)) {
@@ -1469,11 +1470,11 @@ namespace triton {
         }
 
         if (node == nullptr || !PyAstNode_Check(node)) {
-          return PyErr_Format(PyExc_TypeError, "TritonContext::getModels(): Expects a AstNode as first argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::getModels(): Expects a AstNode as node argument.");
         }
 
         if (limit == nullptr || (!PyLong_Check(limit) && !PyInt_Check(limit))) {
-          return PyErr_Format(PyExc_TypeError, "TritonContext::getModels(): Expects an integer as second argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::getModels(): Expects an integer as limit argument.");
         }
 
         if (wb != nullptr && !PyBool_Check(wb)) {
@@ -2249,13 +2250,13 @@ namespace triton {
         }
 
         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.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a SymbolicExpression or a AstNode as node argument.");
 
         if (fname != nullptr && !PyStr_Check(fname))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a string as fname parameter.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a string as fname argument.");
 
         if (optimize != nullptr && !PyBool_Check(optimize))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a boolean as optimize parameter.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::liftToLLVM(): Expects a boolean as optimize argument.");
 
         if (fname == nullptr)
           fname = PyStr_FromString("__triton");
@@ -2927,29 +2928,37 @@ namespace triton {
       static PyObject* TritonContext_simplify(PyObject* self, PyObject* args, PyObject* kwargs) {
         PyObject* node   = nullptr;
         PyObject* solver = nullptr;
+        PyObject* llvm   = nullptr;
 
         static char* keywords[] = {
           (char*)"node",
           (char*)"solver",
+          (char*)"llvm",
           nullptr
         };
 
         /* Extract keywords */
-        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OO", keywords, &node, &solver) == false) {
+        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OOO", keywords, &node, &solver, &llvm) == false) {
           return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Invalid number of arguments");
         }
 
         if (node == nullptr || !PyAstNode_Check(node))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Expects a AstNode as first argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Expects a AstNode as node argument.");
 
         if (solver != nullptr && !PyBool_Check(solver))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Expects a boolean as second argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Expects a boolean as solver argument.");
+
+        if (llvm != nullptr && !PyBool_Check(llvm))
+          return PyErr_Format(PyExc_TypeError, "TritonContext::simplify(): Expects a boolean as llvm argument.");
 
         if (solver == nullptr)
           solver = PyLong_FromUint32(false);
 
+        if (llvm == nullptr)
+          llvm = PyLong_FromUint32(false);
+
         try {
-          return PyAstNode(PyTritonContext_AsTritonContext(self)->simplify(PyAstNode_AsAstNode(node), PyLong_AsBool(solver)));
+          return PyAstNode(PyTritonContext_AsTritonContext(self)->simplify(PyAstNode_AsAstNode(node), PyLong_AsBool(solver), PyLong_AsBool(llvm)));
         }
         catch (const triton::exceptions::PyCallbacks&) {
           return nullptr;
@@ -3101,16 +3110,16 @@ namespace triton {
         }
 
         if (node == nullptr || !PyAstNode_Check(node))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a AstNode as first argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a AstNode as node argument.");
 
         if (constant != nullptr && !PyBool_Check(constant))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as second argument.");
+      
8000
    return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as constant argument.");
 
         if (subexpr != nullptr && !PyBool_Check(subexpr))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as third argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as subexpr argument.");
 
         if (opaque != nullptr && !PyBool_Check(opaque))
-          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as fourth argument.");
+          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as opaque argument.");
 
         if (constant == nullptr)
           constant = PyLong_FromUint32(true);

diff --git a/src/libtriton/engines/lifters/liftingToLLVM.cpp b/src/libtriton/engines/lifters/liftingToLLVM.cpp
@@ -9,9 +9,14 @@
 #include <ostream>
 #include <string>
 
+#include <triton/astContext.hpp>
 #include <triton/liftingToLLVM.hpp>
+#include <triton/llvmToTriton.hpp>
 #include <triton/tritonToLLVM.hpp>
 
+#include <llvm/IR/LLVMContext.h>
+#include <llvm/IR/Module.h>
+
 
 
 namespace triton {
@@ -47,6 +52,17 @@ namespace triton {
         return stream;
       }
 
+
+      triton::ast::SharedAbstractNode LiftingToLLVM::simplifyAstViaLLVM(const triton::ast::SharedAbstractNode& node) const {
+        llvm::LLVMContext context;
+
+        triton::ast::TritonToLLVM ttllvm(context);
+        triton::ast::LLVMToTriton llvmtt(node->getContext());
+
+        auto llvmModule = ttllvm.convert(node, "__tmp", true);  /* from triton to llvm */
+        return llvmtt.convert(llvmModule.get(), "__tmp");       /* from llvm to triton */
+      }
+
     }; /* lifters namespace */
   }; /* engines namespace */
 }; /* triton namespace */
diff --git a/src/libtriton/includes/triton/api.hpp b/src/libtriton/includes/triton/api.hpp
@@ -407,8 +407,8 @@ namespace triton {
         //! [**symbolic api**] - Assigns a symbolic expression to a register.
         TRITON_EXPORT void assignSymbolicExpressionToRegister(const triton::engines::symbolic::SharedSymbolicExpression& se, const triton::arch::Register& reg);
    
-        //! [**symbolic api**] - Processes all recorded AST simplifications or uses solver's simplifications if `usingSolver` is true. Returns the simplified AST.
-        TRITON_EXPORT triton::ast::SharedAbstractNode simplify(const triton::ast::SharedAbstractNode& node, bool usingSolver = false) const;
+        //! [**symbolic api**] - Processes all recorded AST simplifications, uses solver's simplifications if `usingSolver` is true or LLVM is `usingLLVM` is true. Returns the simplified AST.
+        TRITON_EXPORT triton::ast::SharedAbstractNode simplify(const triton::ast::SharedAbstractNode& node, bool usingSolver=false, bool usingLLVM=false) const;
 
         //! [**symbolic api**] - Returns the shared symbolic expression corresponding to an id.
         TRITON_EXPORT triton::engines::symbolic::SharedSymbolicExpression getSymbolicExpression(triton::usize symExprId) const;
@@ -683,6 +683,9 @@ namespace triton {
 
         //! [**lifting api**] - Lifts a symbolic expression and all its references to SMT format. If `assert_` is true, then (assert <expr>).
         TRITON_EXPORT std::ostream& liftToSMT(std::ostream& stream, const triton::engines::symbolic::SharedSymbolicExpression& expr, bool assert_);
+
+        //! [**lifting api**] - Lifts and simplify an AST using LLVM
+        TRITON_EXPORT triton::ast::SharedAbstractNode simplifyAstViaLLVM(const triton::ast::SharedAbstractNode& node) const;
     };
 
 /*! @} End of triton namespace */

diff --git a/src/libtriton/includes/triton/liftingToLLVM.hpp b/src/libtriton/includes/triton/liftingToLLVM.hpp
@@ -16,10 +16,6 @@
 #include <triton/dllexport.hpp>
 #include <triton/symbolicExpression.hpp>
 
-#include <llvm/IR/IRBuilder.h>
-#include <llvm/IR/LLVMContext.h>
-#include <llvm/IR/Module.h>
-
 
 
 //! The Triton namespace
@@ -57,6 +53,9 @@ namespace triton {
 
           //! 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", bool optimize=false);
+
+          //! Lifts and simplify an AST using LLVM
+          TRITON_EXPORT triton::ast::SharedAbstractNode simplifyAstViaLLVM(const triton::ast::SharedAbstractNode& node) const;
       };
 
     /*! @} End of lifters namespace */

diff --git a/src/testers/unittests/test_ast_simplification.py b/src/testers/unittests/test_ast_simplification.py
@@ -724,3 +724,19 @@ def test_issue_1002_2(self):
              "(define-fun ref!13 () (_ BitVec 1) ((_ extract 63 63) ref!8)) ; Sign flag - 0x0: sub qword ptr [rdx], rcx\n"
              "(define-fun ref!14 () (_ BitVec 1) (_ bv1 1)) ; Zero flag - 0x0: sub qword ptr [rdx], rcx\n"
              "(define-fun ref!15 () (_ BitVec 64) (_ bv3 64)) ; Program Counter - 0x0: sub qword ptr [rdx], rcx"))
+
+
+class TestAstSimplificationLLVM(unittest.TestCase):
+    def setUp(self):
+        self.ctx = TritonContext(ARCH.X86_64)
+        self.ast = self.ctx.getAstContext()
+
+    def test_1(self):
+        if VERSION.LLVM_INTERFACE is True:
+            x = self.ast.variable(self.ctx.newSymbolicVariable(8, 'x'))
+            y = self.ast.variable(self.ctx.newSymbolicVariable(8, 'y'))
+            n = (x & ~y) | (~x & y)
+            o = self.ctx.simplify(n, llvm=True)
+            r = str(o) == "(bvxor y x)" or str(o) == "(bvxor x y)"
+            self.assertTrue(r)
+        return