#1074: Allow analysis on sub-expressions

src/libtriton/api/api.cpp

@@ -1320,10 +1320,10 @@ namespace triton {
 
   /* Synthesizer engine API ============================================================================= */
 
-  triton::engines::synthesis::SynthesisResult API::synthesize(const triton::ast::SharedAbstractNode& node, bool constant) {
+  triton::engines::synthesis::SynthesisResult API::synthesize(const triton::ast::SharedAbstractNode& node, bool constant, bool subexpr) {
     this->checkSymbolic();
     triton::engines::synthesis::Synthesizer synth(this->symbolic);
-    return synth.synthesize(node, constant);
+    return synth.synthesize(node, constant, subexpr);
   }
 
 }; /* triton namespace */

src/libtriton/ast/ast.cpp

@@ -3269,7 +3269,7 @@ namespace triton {
         auto current = worklist.top();
         worklist.pop();
 
-        // This means that node is already in work_stack and we will not need to convert it second time
+        // This means that node is already visited and we will not need to visited it second time
         if (visited.find(current) != visited.end()) {
           continue;
         }
@@ -3279,10 +3279,10 @@ namespace triton {
           result.push_front(current->shared_from_this());
 
         if (current->getType() == REFERENCE_NODE) {
-          worklist.push(reinterpret_cast<triton::ast::ReferenceNode *>(current)->getSymbolicExpression()->getAst().get());
+          worklist.push(reinterpret_cast<triton::ast::ReferenceNode*>(current)->getSymbolicExpression()->getAst().get());
         }
         else {
-          for (const auto &child : current->getChildren()) {
+          for (const auto& child : current->getChildren()) {
             worklist.push(child.get());
           }
         }

src/libtriton/bindings/python/objects/pyTritonContext.cpp

@@ -376,8 +376,8 @@ Converts a symbolic memory expression to a symbolic variable. This function retu
 - <b>\ref py_SymbolicVariable_page symbolizeRegister(\ref py_Register_page reg, string symVarAlias)</b><br>
 Converts a symbolic register expression to a symbolic variable. This function returns the new symbolic variable created.
 
-- <b>\ref py_AstNode_page synthesize(\ref py_AstNode_page node, bool constant=True)</b><br>
-Synthesizes a given node. If constant is defined to True, do constant synthesizing.
+- <b>\ref py_AstNode_page synthesize(\ref py_AstNode_page node, bool constant=True, bool subexpr=True)</b><br>
+Synthesizes a given node. If `constant` is defined to True, performs a constant synthesis. If `subexpr` is defined to True, performs synthesis on sub-expressions.
 
 - <b>bool taintAssignment(\ref py_MemoryAccess_page memDst, \ref py_Immediate_page immSrc)</b><br>
 Taints `memDst` from `immSrc` with an assignment - `memDst` is untained. Returns true if the `memDst` is still tainted.
@@ -2999,15 +2999,17 @@ namespace triton {
       static PyObject* TritonContext_synthesize(PyObject* self, PyObject* args, PyObject* kwargs) {
         PyObject* node     = nullptr;
         PyObject* constant = nullptr;
+        PyObject* subexpr  = nullptr;
 
         static char* keywords[] = {
           (char*)"node",
           (char*)"constant",
+          (char*)"subexpr",
           nullptr
         };
 
         /* Extract keywords */
-        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OO", keywords, &node, &constant) == false) {
+        if (PyArg_ParseTupleAndKeywords(args, kwargs, "|OOO", keywords, &node, &constant, &subexpr) == false) {
           return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Invalid number of arguments");
         }
 
@@ -3017,11 +3019,17 @@ namespace triton {
         if (constant != nullptr && !PyBool_Check(constant))
           return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as second argument.");
 
+        if (subexpr != nullptr && !PyBool_Check(subexpr))
+          return PyErr_Format(PyExc_TypeError, "TritonContext::synthesize(): Expects a boolean as second argument.");
+
         if (constant == nullptr)
           constant = PyLong_FromUint32(true);
 
+        if (subexpr == nullptr)
+          subexpr = PyLong_FromUint32(true);
+
         try {
-          auto result = PyTritonContext_AsTritonContext(self)->synthesize(PyAstNode_AsAstNode(node), PyLong_AsBool(constant));
+          auto result = PyTritonContext_AsTritonContext(self)->synthesize(PyAstNode_AsAstNode(node), PyLong_AsBool(constant), PyLong_AsBool(subexpr));
           if (result.successful()) {
             return PyAstNode(result.getOutput());
           }

src/libtriton/engines/synthesis/synthesizer.cpp

@@ -6,6 +6,8 @@
 */
 
 #include <chrono>
+#include <stack>
+#include <unordered_set>
 #include <vector>
 
 #include <triton/ast.hpp>
@@ -28,7 +30,7 @@ namespace triton {
       }
 
 
-      SynthesisResult Synthesizer::synthesize(const triton::ast::SharedAbstractNode& input, bool constant) {
+      SynthesisResult Synthesizer::synthesize(const triton::ast::SharedAbstractNode& input, bool constant, bool subexpr) {
         SynthesisResult result;
 
         // Save the input node
@@ -37,21 +39,16 @@ namespace triton {
         // Start to record the time of the synthesizing
         auto start = std::chrono::system_clock::now();
 
-        auto actx = input->getContext();
-        auto vars = triton::ast::search(input, triton::ast::VARIABLE_NODE);
+        // Do not alter original input
+        auto node = triton::ast::newInstance(input.get(), true);
 
-        if (vars.size() == 1 && input->getLevel() > 2) {
-          bool success = this->unaryOperatorSynthesis(actx, vars, input, result);
-
-          if (success == false && constant == true) {
-            this->constantSynthesis(actx, vars, input, result);
+        // Do the synthesize and if nothing has been synthesized, try on children expression
+        if (this->do_synthesize(node, constant, result) == false) {
+          if (subexpr == true) {
+            this->childrenSynthesis(node, constant, result);
           }
         }
 
-        else if (vars.size() == 2 && input->getLevel() > 2) {
-          this->binaryOperatorSynthesis(actx, vars, input, result);
-        }
-
         // Stop to record the time of the synthesizing
         auto end = std::chrono::system_clock::now();
 
@@ -62,13 +59,86 @@ namespace triton {
       }
 
 
-      bool Synthesizer::constantSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                          const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
+      bool Synthesizer::do_synthesize(const triton::ast::SharedAbstractNode& node, bool constant, SynthesisResult& result) {
+        bool ret = false;
+
+        // How many variables in the expression?
+        auto vars = triton::ast::search(node, triton::ast::VARIABLE_NODE);
+
+        // First, find if unary operators can be synthesized
+        if (vars.size() == 1 && node->getLevel() > 2) {
+          ret = this->unaryOperatorSynthesis(vars, node, result);
+
+          // Second, find if constant can be synthesized
+          if (ret == false && constant == true) {
+            ret = this->constantSynthesis(vars, node, result);
+          }
+        }
+
+        // Third, find if binary operators can be synthesized
+        else if (vars.size() == 2 && node->getLevel() > 2) {
+          ret = this->binaryOperatorSynthesis(vars, node, result);
+        }
+
+        return ret;
+      }
+
+
+      bool Synthesizer::childrenSynthesis(const triton::ast::SharedAbstractNode& node, bool constant, SynthesisResult& result) {
+        std::stack<triton::ast::AbstractNode*>                worklist;
+        std::unordered_set<const triton::ast::AbstractNode*>  visited;
+
+        worklist.push(node.get());
+        while (!worklist.empty()) {
+          auto current = worklist.top();
+          worklist.pop();
+
+          // This means that node is already visited and we will not need to visited it second time
+          if (visited.find(current) != visited.end()) {
+            continue;
+          }
+          visited.insert(current);
+
+          if (current->getType() == triton::ast::REFERENCE_NODE) {
+            worklist.push(reinterpret_cast<triton::ast::ReferenceNode*>(current)->getSymbolicExpression()->getAst().get());
+          }
+          else {
+            triton::usize index = 0;
+            for (const auto& child : current->getChildren()) {
+              SynthesisResult tmp;
+              if (this->do_synthesize(child, constant, tmp)) {
+                /* Replace the child node on the fly */
+                current->setChild(index++, tmp.getOutput());
+                /* Set true because we synthesized at least one child */
+                result.setSuccess(true);
+                continue;
+              }
+              worklist.push(child.get());
+              index++;
+            }
+          }
+        }
+
+        /*
+         * If we synthesized at least one child, we set the output as 'node'
+         * because it has been modified on the fly
+         */
+        if (result.successful()) {
+          result.setOutput(node);
+          return true;
+        }
+
+        return false;
+      }
+
+
+      bool Synthesizer::constantSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
         /* We need Z3 solver in order to use quantifier logic */
         #ifdef TRITON_Z3_INTERFACE
 
         /* We start by getting the symbolic variable of the expression */
         auto var_x = reinterpret_cast<triton::ast::VariableNode*>(vars[0].get())->getSymbolicVariable();
+        auto actx  = node->getContext();
 
         triton::uint32 bits    = var_x->getSize();
         triton::uint32 insize  = node->getBitvectorSize();
@@ -114,10 +184,10 @@ namespace triton {
       }
 
 
-      bool Synthesizer::unaryOperatorSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                               const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
+      bool Synthesizer::unaryOperatorSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
         /* We start by saving orignal value of symbolic variable */
         auto var_x = reinterpret_cast<triton::ast::VariableNode*>(vars[0].get())->getSymbolicVariable();
+        auto actx  = node->getContext();
 
         triton::uint512 save_x = actx->getVariableValue(var_x->getName());
         triton::uint32  bits   = var_x->getSize();
@@ -183,11 +253,11 @@ namespace triton {
       }
 
 
-      bool Synthesizer::binaryOperatorSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                                const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
+      bool Synthesizer::binaryOperatorSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result) {
         /* We start by saving orignal value of symbolic variables */
         auto var_x = reinterpret_cast<triton::ast::VariableNode*>(vars[0].get())->getSymbolicVariable();
         auto var_y = reinterpret_cast<triton::ast::VariableNode*>(vars[1].get())->getSymbolicVariable();
+        auto actx  = node->getContext();
 
         triton::uint512 save_x = actx->getVariableValue(var_x->getName());
         triton::uint512 save_y = actx->getVariableValue(var_y->getName());

src/libtriton/includes/triton/api.hpp

@@ -661,8 +661,8 @@ namespace triton {
 
         /* Synthesizer engine API ============================================================================== */
 
-        //! [**synthesizer api**] - Returns the synthesized AST.
-        TRITON_EXPORT triton::engines::synthesis::SynthesisResult synthesize(const triton::ast::SharedAbstractNode& node, bool constant=true);
+        //! [**synthesizer api**] - Synthesizes a given node. If `constant` is true, performa a constant synthesis. If `subexpr` is true, analyze children AST.
+        TRITON_EXPORT triton::engines::synthesis::SynthesisResult synthesize(const triton::ast::SharedAbstractNode& node, bool constant=true, bool subexpr=true);
     };
 
 /*! @} End of triton namespace */

src/libtriton/includes/triton/synthesizer.hpp

@@ -73,23 +73,26 @@ namespace triton {
           triton::engines::symbolic::SymbolicEngine* symbolic;
 
           //! Synthesize a given node that contains one variable (constant synthesizing)
-          bool constantSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                 const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
+          bool constantSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
 
           //! Synthesize a given node that contains one variable with one operator
-          bool unaryOperatorSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                      const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
+          bool unaryOperatorSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
 
           //! Synthesize a given node that contains two variables with one operator
-          bool binaryOperatorSynthesis(const triton::ast::SharedAstContext& actx, const std::deque<triton::ast::SharedAbstractNode>& vars,
-                                       const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
+          bool binaryOperatorSynthesis(const std::deque<triton::ast::SharedAbstractNode>& vars, const triton::ast::SharedAbstractNode& node, SynthesisResult& result);
+
+          //! Synthesize children expression
+          bool childrenSynthesis(const triton::ast::SharedAbstractNode& node, bool constant, SynthesisResult& result);
+
+          //! Do the synthesis
+          bool do_synthesize(const triton::ast::SharedAbstractNode& node, bool constant, SynthesisResult& result);
 
         public:
           //! Constructor.
           TRITON_EXPORT Synthesizer(triton::engines::symbolic::SymbolicEngine* symbolic);
 
-          //! Synthesize a given node
-          TRITON_EXPORT SynthesisResult synthesize(const triton::ast::SharedAbstractNode& node, bool constant=true);
+          //! Synthesizes a given node. If `constant` is true, perform a constant synthesis. If `subexpr` is true, analyze children AST.
+          TRITON_EXPORT SynthesisResult synthesize(const triton::ast::SharedAbstractNode& node, bool constant=true, bool subexpr=true);
       };
 
     /*! @} End of synthesis namespace */

src/testers/unittests/test_synthesizer.py

@@ -56,16 +56,6 @@ def setUp(self):
         self.CODE += b"\xC1\x8D\x14\x06\x8B\x45\xEC\x01\xD0\x83\xC0\x01\x89\x45\xFC\x8B"
         self.CODE += b"\x45\xFC\x5D\xC3"
 
-        self.ctx = TritonContext(ARCH.X86_64)
-        self.ctx.setMode(MODE.CONSTANT_FOLDING, True)
-        self.ctx.setMode(MODE.ALIGNED_MEMORY, True)
-        self.ctx.setMode(MODE.AST_OPTIMIZATIONS, True)
-        self.ctx.addCallback(CALLBACK.SYMBOLIC_SIMPLIFICATION, self.cb)
-
-        self.ctx.symbolizeRegister(self.ctx.registers.edi, "a")
-        self.ctx.symbolizeRegister(self.ctx.registers.esi, "b")
-        self.ctx.setConcreteMemoryAreaValue(0x1000, self.CODE)
-
     def emulate(self, ctx, pc):
         while pc:
             opcode = ctx.getConcreteMemoryAreaValue(pc, 16)
@@ -75,13 +65,36 @@ def emulate(self, ctx, pc):
         return
 
     def cb(self, ctx, node):
-        synth = ctx.synthesize(node, constant=False)
+        synth = ctx.synthesize(node, constant=False, subexpr=False)
         if synth:
             return synth
         return node
 
+    def init(self):
+        self.ctx = TritonContext(ARCH.X86_64)
+        self.ctx.setMode(MODE.CONSTANT_FOLDING, True)
+        self.ctx.setMode(MODE.ALIGNED_MEMORY, True)
+        self.ctx.setMode(MODE.AST_OPTIMIZATIONS, True)
+
+        self.ctx.symbolizeRegister(self.ctx.registers.edi, "a")
+        self.ctx.symbolizeRegister(self.ctx.registers.esi, "b")
+        self.ctx.setConcreteMemoryAreaValue(0x1000, self.CODE)
+
     def test_2(self):
+        # First test using on the fly synthesis
+        self.init()
+        self.ctx.addCallback(CALLBACK.SYMBOLIC_SIMPLIFICATION, self.cb)
         self.emulate(self.ctx, 0x1000)
-        eax = self.ctx.getRegisterAst(self.ctx.registers.eax)
-        ast = self.ctx.getAstContext()
-        self.assertLessEqual(eax.getLevel(), 20)
+        eax  = self.ctx.getRegisterAst(self.ctx.registers.eax)
+        ast  = self.ctx.getAstContext()
+        res1 = str(ast.unroll(eax))
+
+        # Second test using subexpr
+        self.init()
+        self.emulate(self.ctx, 0x1000)
+        eax  = self.ctx.getRegisterAst(self.ctx.registers.eax)
+        ast  = self.ctx.getAstContext()
+        res2 = str(ast.unroll(self.ctx.synthesize(eax, constant=False, subexpr=True)))
+
+        # Both results must be equal
+        self.assertLessEqual(res1, res2)