<commit_before>//! \file
/*
**  Copyright (C) - Triton
**
**  This program is under the terms of the BSD License.
*/

#include <triton/cpuSize.hpp>
#include <triton/exceptions.hpp>
#include <triton/tritonToZ3Ast.hpp>

namespace triton {
  namespace ast {

    TritonToZ3Ast::TritonToZ3Ast(triton::engines::symbolic::SymbolicEngine* symbolicEngine, bool eval)
      : context() {
      if (symbolicEngine == nullptr)
        throw triton::exceptions::AstTranslations("TritonToZ3Ast::TritonToZ3Ast(): The symbolicEngine API cannot be null.");

      this->symbolicEngine = symbolicEngine;
      this->isEval = eval;
    }

    triton::__uint TritonToZ3Ast::getUintValue(const z3::expr& expr) {
      if (!expr.is_int())
        throw triton::exceptions::Exception("TritonToZ3Ast::getUintValue(): The ast is not a numerical value.");

      #if defined(__x86_64__) || defined(_M_X64)
      return expr.get_numeral_uint64();
      #endif
      #if defined(__i386) || defined(_M_IX86)
      return expr.get_numeral_uint();
      #endif
    }

    std::string TritonToZ3Ast::getStringValue(const z3::expr& expr) {
      return Z3_get_numeral_string(this->context, expr);
    }

    z3::expr TritonToZ3Ast::convert(const triton::ast::SharedAbstractNode& node) {
      if (node == nullptr)
        throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): node cannot be null.");

      switch (node->getKind()) {
        case BVADD_NODE:
          return to_expr(this->context, Z3_mk_bvadd(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVAND_NODE:
          return to_expr(this->context, Z3_mk_bvand(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVASHR_NODE:
          return to_expr(this->context, Z3_mk_bvashr(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVLSHR_NODE:
          return to_expr(this->context, Z3_mk_bvlshr(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVMUL_NODE:
          return to_expr(this->context, Z3_mk_bvmul(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNAND_NODE:
          return to_expr(this->context, Z3_mk_bvnand(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNEG_NODE:
          return to_expr(this->context, Z3_mk_bvneg(this->context, this->convert(node->getChildren()[0])));

        case BVNOR_NODE:
          return to_expr(this->context, Z3_mk_bvnor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNOT_NODE:
          return to_expr(this->context, Z3_mk_bvnot(this->context, this->convert(node->getChildren()[0])));

        case BVOR_NODE:
          return to_expr(this->context, Z3_mk_bvor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVROL_NODE: {
          triton::uint32 op1 = reinterpret_cast<triton::ast::DecimalNode*>(node->getChildren()[0].get())->getValue().convert_to<triton::uint32>();
          return to_expr(this->context, Z3_mk_rotate_left(this->context, op1, this->convert(node->getChildren()[1])));
        }

        case BVROR_NODE: {
          triton::uint32 op1 = reinterpret_cast<triton::ast::DecimalNode*>(node->getChildren()[0].get())->getValue().convert_to<triton::uint32>();
          return to_expr(this->context, Z3_mk_rotate_right(this->context, op1, this->convert(node->getChildren()[1])));
        }

        case BVSDIV_NODE:
          return to_expr(this->context, Z3_mk_bvsdiv(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSGE_NODE:
          return to_expr(this->context, Z3_mk_bvsge(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSGT_NODE:
          return to_expr(this->context, Z3_mk_bvsgt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSHL_NODE:
          return to_expr(this->context, Z3_mk_bvshl(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSLE_NODE:
          return to_expr(this->context, Z3_mk_bvsle(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSLT_NODE:
          return to_expr(this->context, Z3_mk_bvslt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSMOD_NODE:
          return to_expr(this->context, Z3_mk_bvsmod(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSREM_NODE:
          return to_expr(this->context, Z3_mk_bvsrem(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSUB_NODE:
          return to_expr(this->context, Z3_mk_bvsub(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUDIV_NODE:
          return to_expr(this->context, Z3_mk_bvudiv(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUGE_NODE:
          return to_expr(this->context, Z3_mk_bvuge(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUGT_NODE:
          return to_expr(this->context, Z3_mk_bvugt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVULE_NODE:
          return to_expr(this->context, Z3_mk_bvule(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVULT_NODE:
          return to_expr(this->context, Z3_mk_bvult(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUREM_NODE:
          return to_expr(this->context, Z3_mk_bvurem(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVXNOR_NODE:
          return to_expr(this->context, Z3_mk_bvxnor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVXOR_NODE:
          return to_expr(this->context, Z3_mk_bvxor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BV_NODE: {
          z3::expr value        = this->convert(node->getChildren()[0]);
          z3::expr size         = this->convert(node->getChildren()[1]);
          triton::uint32 bvsize = static_cast<triton::uint32>(this->getUintValue(size));
          return this->context.bv_val(this->getStringValue(value).c_str(), bvsize);
        }

        case CONCAT_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          z3::expr nextValue(this->context);

          // Child[0] is the LSB
          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            currentValue = to_expr(this->context, Z3_mk_concat(this->context, currentValue, nextValue));
          }

          return currentValue;
        }

        case DECIMAL_NODE: {
          std::string value(reinterpret_cast<triton::ast::DecimalNode*>(node.get())->getValue());
          return this->context.int_val(value.c_str());
        }

        case DISTINCT_NODE: {
          z3::expr op1 = this->convert(node->getChildren()[0]);
          z3::expr op2 = this->convert(node->getChildren()[1]);
          Z3_ast ops[] = {op1, op2};

          return to_expr(this->context, Z3_mk_distinct(this->context, 2, ops));
        }

        case EQUAL_NODE:
          return to_expr(this->context, Z3_mk_eq(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case EXTRACT_NODE: {
          z3::expr high     = this->convert(node->getChildren()[0]);
          z3::expr low      = this->convert(node->getChildren()[1]);
          z3::expr value    = this->convert(node->getChildren()[2]);
          triton::uint32 hv = static_cast<triton::uint32>(this->getUintValue(high));
          triton::uint32 lv = static_cast<triton::uint32>(this->getUintValue(low));

          return to_expr(this->context, Z3_mk_extract(this->context, hv, lv, value));
        }

        case ITE_NODE: {
          z3::expr op1 = this->convert(node->getChildren()[0]); // condition
          z3::expr op2 = this->convert(node->getChildren()[1]); // if true
          z3::expr op3 = this->convert(node->getChildren()[2]); // if false

          return to_expr(this->context, Z3_mk_ite(this->context, op1, op2, op3));
        }

        case LAND_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          if (!currentValue.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LandNode(): Land can be apply only on bool value.");
          }
          z3::expr nextValue(this->context);

          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            if (!nextValue.get_sort().is_bool()) {
              throw triton::exceptions::AstTranslations("TritonToZ3Ast::LandNode(): Land can be apply only on bool value.");
            }
            Z3_ast ops[] = {currentValue, nextValue};
            currentValue = to_expr(this->context, Z3_mk_and(this->context, 2, ops));
          }

          return currentValue;
        }

        case LET_NODE: {
          std::string symbol    = reinterpret_cast<triton::ast::StringNode*>(node->getChildren()[0].get())->getValue();
          this->symbols[symbol] = node->getChildren()[1];

          return this->convert(node->getChildren()[2]);
        }

        case LNOT_NODE: {
          z3::expr value = this->convert(node->getChildren()[0]);
          if (!value.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
          }
          return to_expr(this->context, Z3_mk_not(this->context, value));
        }

        case LOR_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          if (!currentValue.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
          }
          z3::expr nextValue(this->context);

          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            if (!nextValue.get_sort().is_bool()) {
              throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
            }
            Z3_ast ops[] = {currentValue, nextValue};
            currentValue = to_expr(this->context, Z3_mk_or(this->context, 2, ops));
          }

          return currentValue;
        }

        case REFERENCE_NODE:
          return this->convert(reinterpret_cast<triton::ast::ReferenceNode*>(node.get())->getSymbolicExpression()->getAst());

        case STRING_NODE: {
          std::string value = reinterpret_cast<triton::ast::StringNode*>(node.get())->getValue();

          if (this->symbols.find(value) == this->symbols.end())
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): [STRING_NODE] Symbols not found.");

          return this->convert(this->symbols[value]);
        }

        case SX_NODE: {
          z3::expr ext        = this->convert(node->getChildren()[0]);
          z3::expr value      = this->convert(node->getChildren()[1]);
          triton::uint32 extv = static_cast<triton::uint32>(this->getUintValue(ext));

          return to_expr(this->context, Z3_mk_sign_ext(this->context, extv, value));
        }

        case VARIABLE_NODE: {
          triton::usize varId = reinterpret_cast<triton::ast::VariableNode*>(node.get())->getVar().getId();
          triton::engines::symbolic::SymbolicVariable* symVar = this->symbolicEngine->getSymbolicVariableFromId(varId);

          if (symVar == nullptr)
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): [VARIABLE_NODE] Can't get the symbolic variable (nullptr).");

          /* If the conversion is used to evaluate a node, we concretize symbolic variables */
          if (this->isEval) {
            triton::uint512 value = reinterpret_cast<triton::ast::VariableNode*>(node.get())->evaluate();
            std::string strValue(value);
            return this->context.bv_val(strValue.c_str(), symVar->getSize());
          }

          /* Otherwise, we keep the symbolic variables for a real conversion */
          return this->context.bv_const(symVar->getName().c_str(), symVar->getSize());
        }

        case ZX_NODE: {
          z3::expr ext        = this->convert(node->getChildren()[0]);
          z3::expr value      = this->convert(node->getChildren()[1]);
          triton::uint32 extv = static_cast<triton::uint32>(this->getUintValue(ext));

          return to_expr(this->context, Z3_mk_zero_ext(this->context, extv, value));
        }

        default:
          throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): Invalid kind of node.");
      }
    }

  }; /* ast namespace */
}; /* triton namespace */
<commit_msg>Speedup into tritonToZ3Ast<commit_after>//! \file
/*
**  Copyright (C) - Triton
**
**  This program is under the terms of the BSD License.
*/

#include <triton/cpuSize.hpp>
#include <triton/exceptions.hpp>
#include <triton/tritonToZ3Ast.hpp>

namespace triton {
  namespace ast {

    TritonToZ3Ast::TritonToZ3Ast(triton::engines::symbolic::SymbolicEngine* symbolicEngine, bool eval)
      : context() {
      if (symbolicEngine == nullptr)
        throw triton::exceptions::AstTranslations("TritonToZ3Ast::TritonToZ3Ast(): The symbolicEngine API cannot be null.");

      this->symbolicEngine = symbolicEngine;
      this->isEval = eval;
    }

    triton::__uint TritonToZ3Ast::getUintValue(const z3::expr& expr) {
      if (!expr.is_int())
        throw triton::exceptions::Exception("TritonToZ3Ast::getUintValue(): The ast is not a numerical value.");

      #if defined(__x86_64__) || defined(_M_X64)
      return expr.get_numeral_uint64();
      #endif
      #if defined(__i386) || defined(_M_IX86)
      return expr.get_numeral_uint();
      #endif
    }

    std::string TritonToZ3Ast::getStringValue(const z3::expr& expr) {
      return Z3_get_numeral_string(this->context, expr);
    }

    z3::expr TritonToZ3Ast::convert(const triton::ast::SharedAbstractNode& node) {
      if (node == nullptr)
        throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): node cannot be null.");

      switch (node->getKind()) {
        case BVADD_NODE:
          return to_expr(this->context, Z3_mk_bvadd(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVAND_NODE:
          return to_expr(this->context, Z3_mk_bvand(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVASHR_NODE:
          return to_expr(this->context, Z3_mk_bvashr(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVLSHR_NODE:
          return to_expr(this->context, Z3_mk_bvlshr(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVMUL_NODE:
          return to_expr(this->context, Z3_mk_bvmul(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNAND_NODE:
          return to_expr(this->context, Z3_mk_bvnand(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNEG_NODE:
          return to_expr(this->context, Z3_mk_bvneg(this->context, this->convert(node->getChildren()[0])));

        case BVNOR_NODE:
          return to_expr(this->context, Z3_mk_bvnor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVNOT_NODE:
          return to_expr(this->context, Z3_mk_bvnot(this->context, this->convert(node->getChildren()[0])));

        case BVOR_NODE:
          return to_expr(this->context, Z3_mk_bvor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVROL_NODE: {
          triton::uint32 op1 = reinterpret_cast<triton::ast::DecimalNode*>(node->getChildren()[0].get())->getValue().convert_to<triton::uint32>();
          return to_expr(this->context, Z3_mk_rotate_left(this->context, op1, this->convert(node->getChildren()[1])));
        }

        case BVROR_NODE: {
          triton::uint32 op1 = reinterpret_cast<triton::ast::DecimalNode*>(node->getChildren()[0].get())->getValue().convert_to<triton::uint32>();
          return to_expr(this->context, Z3_mk_rotate_right(this->context, op1, this->convert(node->getChildren()[1])));
        }

        case BVSDIV_NODE:
          return to_expr(this->context, Z3_mk_bvsdiv(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSGE_NODE:
          return to_expr(this->context, Z3_mk_bvsge(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSGT_NODE:
          return to_expr(this->context, Z3_mk_bvsgt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSHL_NODE:
          return to_expr(this->context, Z3_mk_bvshl(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSLE_NODE:
          return to_expr(this->context, Z3_mk_bvsle(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSLT_NODE:
          return to_expr(this->context, Z3_mk_bvslt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSMOD_NODE:
          return to_expr(this->context, Z3_mk_bvsmod(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSREM_NODE:
          return to_expr(this->context, Z3_mk_bvsrem(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVSUB_NODE:
          return to_expr(this->context, Z3_mk_bvsub(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUDIV_NODE:
          return to_expr(this->context, Z3_mk_bvudiv(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUGE_NODE:
          return to_expr(this->context, Z3_mk_bvuge(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUGT_NODE:
          return to_expr(this->context, Z3_mk_bvugt(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVULE_NODE:
          return to_expr(this->context, Z3_mk_bvule(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVULT_NODE:
          return to_expr(this->context, Z3_mk_bvult(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVUREM_NODE:
          return to_expr(this->context, Z3_mk_bvurem(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVXNOR_NODE:
          return to_expr(this->context, Z3_mk_bvxnor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BVXOR_NODE:
          return to_expr(this->context, Z3_mk_bvxor(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case BV_NODE: {
          z3::expr value        = this->convert(node->getChildren()[0]);
          z3::expr size         = this->convert(node->getChildren()[1]);
          triton::uint32 bvsize = static_cast<triton::uint32>(this->getUintValue(size));
          return this->context.bv_val(this->getStringValue(value).c_str(), bvsize);
        }

        case CONCAT_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          z3::expr nextValue(this->context);

          // Child[0] is the LSB
          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            currentValue = to_expr(this->context, Z3_mk_concat(this->context, currentValue, nextValue));
          }

          return currentValue;
        }

        case DECIMAL_NODE: {
          std::string value(reinterpret_cast<triton::ast::DecimalNode*>(node.get())->getValue());
          return this->context.int_val(value.c_str());
        }

        case DISTINCT_NODE: {
          z3::expr op1 = this->convert(node->getChildren()[0]);
          z3::expr op2 = this->convert(node->getChildren()[1]);
          Z3_ast ops[] = {op1, op2};

          return to_expr(this->context, Z3_mk_distinct(this->context, 2, ops));
        }

        case EQUAL_NODE:
          return to_expr(this->context, Z3_mk_eq(this->context, this->convert(node->getChildren()[0]), this->convert(node->getChildren()[1])));

        case EXTRACT_NODE: {
          z3::expr high     = this->convert(node->getChildren()[0]);
          z3::expr low      = this->convert(node->getChildren()[1]);
          z3::expr value    = this->convert(node->getChildren()[2]);
          triton::uint32 hv = static_cast<triton::uint32>(this->getUintValue(high));
          triton::uint32 lv = static_cast<triton::uint32>(this->getUintValue(low));

          return to_expr(this->context, Z3_mk_extract(this->context, hv, lv, value));
        }

        case ITE_NODE: {
          z3::expr op1 = this->convert(node->getChildren()[0]); // condition
          z3::expr op2 = this->convert(node->getChildren()[1]); // if true
          z3::expr op3 = this->convert(node->getChildren()[2]); // if false

          return to_expr(this->context, Z3_mk_ite(this->context, op1, op2, op3));
        }

        case LAND_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          if (!currentValue.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LandNode(): Land can be apply only on bool value.");
          }
          z3::expr nextValue(this->context);

          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            if (!nextValue.get_sort().is_bool()) {
              throw triton::exceptions::AstTranslations("TritonToZ3Ast::LandNode(): Land can be apply only on bool value.");
            }
            Z3_ast ops[] = {currentValue, nextValue};
            currentValue = to_expr(this->context, Z3_mk_and(this->context, 2, ops));
          }

          return currentValue;
        }

        case LET_NODE: {
          std::string symbol    = reinterpret_cast<triton::ast::StringNode*>(node->getChildren()[0].get())->getValue();
          this->symbols[symbol] = node->getChildren()[1];

          return this->convert(node->getChildren()[2]);
        }

        case LNOT_NODE: {
          z3::expr value = this->convert(node->getChildren()[0]);
          if (!value.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
          }
          return to_expr(this->context, Z3_mk_not(this->context, value));
        }

        case LOR_NODE: {
          const std::vector<triton::ast::SharedAbstractNode>& children = node->getChildren();

          z3::expr currentValue = this->convert(node->getChildren()[0]);
          if (!currentValue.get_sort().is_bool()) {
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
          }
          z3::expr nextValue(this->context);

          for (triton::uint32 idx = 1; idx < children.size(); idx++) {
            nextValue = this->convert(children[idx]);
            if (!nextValue.get_sort().is_bool()) {
              throw triton::exceptions::AstTranslations("TritonToZ3Ast::LnotNode(): Lnot can be apply only on bool value.");
            }
            Z3_ast ops[] = {currentValue, nextValue};
            currentValue = to_expr(this->context, Z3_mk_or(this->context, 2, ops));
          }

          return currentValue;
        }

        case REFERENCE_NODE:
          return this->convert(reinterpret_cast<triton::ast::ReferenceNode*>(node.get())->getSymbolicExpression()->getAst());

        case STRING_NODE: {
          std::string value = reinterpret_cast<triton::ast::StringNode*>(node.get())->getValue();

          if (this->symbols.find(value) == this->symbols.end())
            throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): [STRING_NODE] Symbols not found.");

          return this->convert(this->symbols[value]);
        }

        case SX_NODE: {
          z3::expr ext        = this->convert(node->getChildren()[0]);
          z3::expr value      = this->convert(node->getChildren()[1]);
          triton::uint32 extv = static_cast<triton::uint32>(this->getUintValue(ext));

          return to_expr(this->context, Z3_mk_sign_ext(this->context, extv, value));
        }

        case VARIABLE_NODE: {
          const triton::engines::symbolic::SymbolicVariable& symVar = reinterpret_cast<triton::ast::VariableNode*>(node.get())->getVar();

          /* If the conversion is used to evaluate a node, we concretize symbolic variables */
          if (this->isEval) {
            triton::uint512 value = reinterpret_cast<triton::ast::VariableNode*>(node.get())->evaluate();
            std::string strValue(value);
            return this->context.bv_val(strValue.c_str(), symVar.getSize());
          }

          /* Otherwise, we keep the symbolic variables for a real conversion */
          return this->context.bv_const(symVar.getName().c_str(), symVar.getSize());
        }

        case ZX_NODE: {
          z3::expr ext        = this->convert(node->getChildren()[0]);
          z3::expr value      = this->convert(node->getChildren()[1]);
          triton::uint32 extv = static_cast<triton::uint32>(this->getUintValue(ext));

          return to_expr(this->context, Z3_mk_zero_ext(this->context, extv, value));
        }

        default:
          throw triton::exceptions::AstTranslations("TritonToZ3Ast::convert(): Invalid kind of node.");
      }
    }

  }; /* ast namespace */
}; /* triton namespace */
<|endoftext|>


<commit_before>/**
 * @file sr_parser.cpp
 * @author Chase Geigle
 */

#include <fstream>

#include "io/binary.h"
#include "logging/logger.h"
#include "parallel/parallel_for.h"
#include "parser/sr_parser.h"
#include "parser/state.h"
#include "parser/state_analyzer.h"
#include "parser/training_data.h"
#include "parser/trees/internal_node.h"
#include "parser/trees/leaf_node.h"
#include "util/progress.h"
#include "util/range.h"

namespace meta
{
namespace parser
{

sr_parser::sr_parser(const std::string& prefix) : trans_{prefix}
{
    load(prefix);
}

void sr_parser::train(std::vector<parse_tree>& trees, training_options options)
{
    training_data data{options, trees};
    trans_ = data.preprocess();

    LOG(info) << "Found " << trans_.size() << " transitions" << ENDLG;

    parallel::thread_pool pool{options.num_threads};

    for (uint64_t iter = 1; iter <= options.max_iterations; ++iter)
    {
        printing::progress progress{
            " > Iteration " + std::to_string(iter) + ": ", trees.size()};
        data.shuffle();

        for (size_t start = 0; start < data.size(); start += options.batch_size)
        {
            progress(start);
            auto end = std::min(start + options.batch_size, data.size());

            auto update = train_batch({data, start, end}, pool);

            for (const auto& feat_vec : update)
            {
                const auto& feat = feat_vec.first;
                auto& wv = weights_[feat];
                for (const auto& up : feat_vec.second)
                    wv[up.first] += up.second;
            }
        }
        progress.end();

        condense();
    }
}

auto sr_parser::train_batch(training_batch batch, parallel::thread_pool& pool)
    -> weight_vectors
{
    // TODO: real beam search

    weight_vectors update;
    auto range = util::range(batch.start, batch.end - 1); // inclusive range

    // Perform a reduction across threads: each thread stores its update in
    // a separate location, and we then add them all up after we join
    util::sparse_vector<std::thread::id, weight_vectors> updates;
    for (const auto& tid : pool.thread_ids())
        updates[tid] = {};

    state_analyzer analyzer;
    parallel::parallel_for(range.begin(), range.end(), pool, [&](size_t i)
                           {
        auto& tree = batch.data.tree(i);
        auto& transitions = batch.data.transitions(i);

        state state{tree};

        for (const auto& gold_trans : transitions)
        {
            auto feats = analyzer.featurize(state);
            auto trans = best_transition(feats);

            if (trans == gold_trans)
            {
                state.advance(trans_[trans]);
            }
            else
            {
                auto& update = updates[std::this_thread::get_id()];
                for (const auto& feat : feats)
                {
                    auto& wv = update[feat.first];
                    wv[gold_trans] += feat.second;
                    wv[trans] -= feat.second;
                }
                break;
            }
        }
    });

    // Reduce partial results down to final update vector
    for (const auto& thread_update : updates)
    {
        for (const auto& feat : thread_update.second)
        {
            auto& wv = update[feat.first];
            for (const auto& weight : feat.second)
                wv[weight.first] += weight.second;
        }
    }
    return update;
}

auto sr_parser::best_transition(
    const feature_vector& features) const -> trans_id
{
    weight_vector class_scores;
    for (const auto& feat : features)
    {
        const auto& name = feat.first;
        auto val = feat.second;

        auto it = weights_.find(name);
        if (it == weights_.end())
            continue;

        for (const auto& trans_weight : it->second)
        {
            auto tid = trans_weight.first;
            auto trans_w = trans_weight.second;

            class_scores[tid] += val * trans_w;
        }
    }

    auto best_score = std::numeric_limits<float>::lowest();
    trans_id best_trans{};
    for (const auto& score : class_scores)
    {
        if (score.second > best_score)
        {
            best_trans = score.first;
            best_score = score.second;
        }
    }

    return best_trans;
}

void sr_parser::condense()
{
    // build feature set
    std::vector<std::string> features;
    features.reserve(weights_.size());
    for (const auto& feat_vec : weights_)
        features.push_back(feat_vec.first);

    uint64_t nnz = 0;
    for (const auto& feat : features)
    {
        auto it = weights_.find(feat);
        it->second.condense();
        if (it->second.empty())
            weights_.erase(it);
        else
            nnz += it->second.size();
    }

    LOG(info) << "Number of nonzero weights: " << nnz << ENDLG;
}

void sr_parser::save(const std::string& prefix) const
{
    trans_.save(prefix);
    std::ofstream model{prefix + "/parser.model", std::ios::binary};

    io::write_binary(model, weights_.size());
    for (const auto& feat_vec : weights_)
    {
        const auto& feat = feat_vec.first;
        const auto& weights = feat_vec.second;

        io::write_binary(model, feat);
        io::write_binary(model, weights.size());

        for (const auto& weight : weights)
        {
            auto tid = weight.first;
            int val = weight.second;
            io::write_binary(model, tid);
            io::write_binary(model, val);
        }
    }
}

void sr_parser::load(const std::string& prefix)
{
    std::ifstream model{prefix + "/parser.model", std::ios::binary};

    if (!model)
        throw exception{"model file not found"};

    size_t num_feats;
    io::read_binary(model, num_feats);

    for (size_t i = 0; i < num_feats; ++i)
    {
        if (!model)
            throw exception{"malformed model file (too few features written)"};

        std::string feature_name;
        io::read_binary(model, feature_name);

        size_t num_trans;
        io::read_binary(model, num_trans);

        for (size_t j = 0; j < num_trans; ++j)
        {
            if (!model)
                throw exception{"malformed model file (too few transitions "
                                "written for feature)"};

            trans_id tid;
            float val;
            io::read_binary(model, tid);
            io::read_binary(model, val);

            weights_[feature_name][tid] = val;
        }
    }
}
}
}
<commit_msg>Print number of features after each iteration.<commit_after>/**
 * @file sr_parser.cpp
 * @author Chase Geigle
 */

#include <fstream>

#include "io/binary.h"
#include "logging/logger.h"
#include "parallel/parallel_for.h"
#include "parser/sr_parser.h"
#include "parser/state.h"
#include "parser/state_analyzer.h"
#include "parser/training_data.h"
#include "parser/trees/internal_node.h"
#include "parser/trees/leaf_node.h"
#include "util/progress.h"
#include "util/range.h"

namespace meta
{
namespace parser
{

sr_parser::sr_parser(const std::string& prefix) : trans_{prefix}
{
    load(prefix);
}

void sr_parser::train(std::vector<parse_tree>& trees, training_options options)
{
    training_data data{options, trees};
    trans_ = data.preprocess();

    LOG(info) << "Found " << trans_.size() << " transitions" << ENDLG;

    parallel::thread_pool pool{options.num_threads};

    for (uint64_t iter = 1; iter <= options.max_iterations; ++iter)
    {
        printing::progress progress{
            " > Iteration " + std::to_string(iter) + ": ", trees.size()};
        data.shuffle();

        for (size_t start = 0; start < data.size(); start += options.batch_size)
        {
            progress(start);
            auto end = std::min(start + options.batch_size, data.size());

            auto update = train_batch({data, start, end}, pool);

            for (const auto& feat_vec : update)
            {
                const auto& feat = feat_vec.first;
                auto& wv = weights_[feat];
                for (const auto& up : feat_vec.second)
                    wv[up.first] += up.second;
            }
        }
        progress.end();

        condense();
    }
}

auto sr_parser::train_batch(training_batch batch, parallel::thread_pool& pool)
    -> weight_vectors
{
    // TODO: real beam search

    weight_vectors update;
    auto range = util::range(batch.start, batch.end - 1); // inclusive range

    // Perform a reduction across threads: each thread stores its update in
    // a separate location, and we then add them all up after we join
    util::sparse_vector<std::thread::id, weight_vectors> updates;
    for (const auto& tid : pool.thread_ids())
        updates[tid] = {};

    state_analyzer analyzer;
    parallel::parallel_for(range.begin(), range.end(), pool, [&](size_t i)
                           {
        auto& tree = batch.data.tree(i);
        auto& transitions = batch.data.transitions(i);

        state state{tree};

        for (const auto& gold_trans : transitions)
        {
            auto feats = analyzer.featurize(state);
            auto trans = best_transition(feats);

            if (trans == gold_trans)
            {
                state.advance(trans_[trans]);
            }
            else
            {
                auto& update = updates[std::this_thread::get_id()];
                for (const auto& feat : feats)
                {
                    auto& wv = update[feat.first];
                    wv[gold_trans] += feat.second;
                    wv[trans] -= feat.second;
                }
                break;
            }
        }
    });

    // Reduce partial results down to final update vector
    for (const auto& thread_update : updates)
    {
        for (const auto& feat : thread_update.second)
        {
            auto& wv = update[feat.first];
            for (const auto& weight : feat.second)
                wv[weight.first] += weight.second;
        }
    }
    return update;
}

auto sr_parser::best_transition(
    const feature_vector& features) const -> trans_id
{
    weight_vector class_scores;
    for (const auto& feat : features)
    {
        const auto& name = feat.first;
        auto val = feat.second;

        auto it = weights_.find(name);
        if (it == weights_.end())
            continue;

        for (const auto& trans_weight : it->second)
        {
            auto tid = trans_weight.first;
            auto trans_w = trans_weight.second;

            class_scores[tid] += val * trans_w;
        }
    }

    auto best_score = std::numeric_limits<float>::lowest();
    trans_id best_trans{};
    for (const auto& score : class_scores)
    {
        if (score.second > best_score)
        {
            best_trans = score.first;
            best_score = score.second;
        }
    }

    return best_trans;
}

void sr_parser::condense()
{
    // build feature set
    std::vector<std::string> features;
    features.reserve(weights_.size());
    for (const auto& feat_vec : weights_)
        features.push_back(feat_vec.first);

    uint64_t nnz = 0;
    for (const auto& feat : features)
    {
        auto it = weights_.find(feat);
        it->second.condense();
        if (it->second.empty())
            weights_.erase(it);
        else
            nnz += it->second.size();
    }

    LOG(info) << "Number of total features: " << weights_.size() << ENDLG;
    LOG(info) << "Number of nonzero weights: " << nnz << ENDLG;
}

void sr_parser::save(const std::string& prefix) const
{
    trans_.save(prefix);
    std::ofstream model{prefix + "/parser.model", std::ios::binary};

    io::write_binary(model, weights_.size());
    for (const auto& feat_vec : weights_)
    {
        const auto& feat = feat_vec.first;
        const auto& weights = feat_vec.second;

        io::write_binary(model, feat);
        io::write_binary(model, weights.size());

        for (const auto& weight : weights)
        {
            auto tid = weight.first;
            int val = weight.second;
            io::write_binary(model, tid);
            io::write_binary(model, val);
        }
    }
}

void sr_parser::load(const std::string& prefix)
{
    std::ifstream model{prefix + "/parser.model", std::ios::binary};

    if (!model)
        throw exception{"model file not found"};

    size_t num_feats;
    io::read_binary(model, num_feats);

    for (size_t i = 0; i < num_feats; ++i)
    {
        if (!model)
            throw exception{"malformed model file (too few features written)"};

        std::string feature_name;
        io::read_binary(model, feature_name);

        size_t num_trans;
        io::read_binary(model, num_trans);

        for (size_t j = 0; j < num_trans; ++j)
        {
            if (!model)
                throw exception{"malformed model file (too few transitions "
                                "written for feature)"};

            trans_id tid;
            float val;
            io::read_binary(model, tid);
            io::read_binary(model, val);

            weights_[feature_name][tid] = val;
        }
    }
}
}
}
<|endoftext|>


<commit_before>///
/// @file  pi_lmo_parallel1.cpp
/// @brief Parallel implementation of the Lagarias-Miller-Odlyzko
///        prime counting algorithm using OpenMP. This implementation
///        is derived from pi_lmo4(x).
///
/// Copyright (C) 2014 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include "internal.hpp"
#include "PhiTiny.hpp"
#include "pmath.hpp"
#include "tos_counters.hpp"

#include <primesieve.hpp>
#include <stdint.h>
#include <algorithm>
#include <vector>

#ifdef _OPENMP
  #include <omp.h>
  #include "get_omp_threads.hpp"
#endif

using namespace std;

#ifdef _OPENMP

namespace {

template <typename T1, typename T2>
void cross_off(int64_t prime,
               int64_t low,
               int64_t high,
               int64_t& next_multiple,
               T1& sieve,
               T2& counters)
{
  int64_t segment_size = sieve.size();
  int64_t k = next_multiple + prime * (~next_multiple & 1);

  for (; k < high; k += prime * 2)
  {
    if (sieve[k - low])
    {
      sieve[k - low] = 0;
      cnt_update(counters, k - low, segment_size);
    }
  }

  next_multiple = k;
}

/// Calculate the contribution of the special leaves.
/// @pre y > 0 && c > 1
///
int64_t S2(int64_t x,
           int64_t y,
           int64_t pi_y,
           int64_t c,
           vector<int32_t>& primes,
           vector<int32_t>& lpf,
           vector<int32_t>& mu,
           int threads)
{
  threads = std::max(1, threads);
  int64_t limit = x / y + 1;
  int64_t segment_size = next_power_of_2(isqrt(limit));
  int64_t segments = (limit + segment_size - 1) / segment_size;
  int64_t segments_per_thread = (segments + threads - 1) / threads;
  int64_t S2_total = 0;

  vector<vector<int64_t> > phi(threads);
  vector<vector<int64_t> > mu_sum(threads);

  for (int i = 0; i < threads; i++)
  {
    phi[i].resize(primes.size(), 0);
    mu_sum[i].resize(primes.size(), 0);
  }

  #pragma omp parallel for num_threads(threads) reduction(+: S2_total)
  for (int i = 0; i < threads; i++)
  {
    vector<char> sieve(segment_size);
    vector<int32_t> counters(segment_size);
    vector<int32_t>::iterator iter = primes.begin() + 1;
    vector<int64_t> next;
    next.reserve(primes.size());
    next.push_back(0);

    int64_t start_idx = i * segments_per_thread;
    int64_t stop_idx = min(start_idx + segments_per_thread, segments);
    int64_t S2_thread = 0;

    for (int64_t j = start_idx; j < stop_idx; j++)
    {
      fill(sieve.begin(), sieve.end(), 1);

      // Current segment = interval [low, high[
      int64_t low = j * segment_size + 1;
      int64_t high = min(low + segment_size, limit);

      // Initialize sieve indices
      while (iter != primes.end() && isquare(*iter) < high)
      {
        int64_t prime = *iter++;
        int64_t next_multiple = ((low + prime - 1) / prime) * prime;
        next.push_back(next_multiple);
      }

      for (int64_t b = 1; b <= c; b++)
      {
        int64_t k = next[b];
        for (int64_t prime = primes[b]; k < high; k += prime)
          sieve[k - low] = 0;
        next[b] = k;
      }

      // Initialize special tree data structure from sieve
      cnt_finit(sieve, counters, segment_size);

      for (int64_t b = c + 1; b < pi_y; b++)
      {
        int64_t prime = primes[b];
        int64_t min_m = max(x / (prime * high), y / prime);
        int64_t max_m = min(x / (prime * low), y);

        if (prime >= max_m)
          break;

        for (int64_t m = max_m; m > min_m; m--)
        {
          if (mu[m] != 0 && prime < lpf[m])
          {
            // phi_xn is the number of currently unsieved elements in
            // the interval [start_idx * segment_size + 1, x / n].
            // This is not the entire contribution of the special leaf
            // n, the missing contribution is the number of unsieved
            // elements in the interval [1, start_idx * segment_size].
            // The missing contribution is reconstructed and added
            // once all threads have finished.
            //
            int64_t n = prime * m;
            int64_t count = cnt_query(counters, (x / n) - low);
            int64_t phi_xn = phi[i][b] + count;
            S2_thread -= mu[m] * phi_xn;
            mu_sum[i][b] -= mu[m];
          }
        }

        phi[i][b] += cnt_query(counters, high - 1 - low);
        cross_off(prime, low, high, next[b], sieve, counters);
      }
    }

    S2_total += S2_thread;
  }

  // Once all threads have finished reconstruct and add the
  // missing contribution of all special leaves. This must
  // be done in order as each thread (i) requires the sum of
  // the phi values from the previous threads.
  //
  for (int i = 1; i < threads; i++)
  {
    for (size_t j = 0; j < phi[i].size(); j++)
    {
      S2_total += phi[i - 1][j] * mu_sum[i][j];
      phi[i][j] += phi[i - 1][j];
    }
  }

  return S2_total;
}

} // namespace

#endif /* _OPENMP */

namespace primecount {

/// Calculate the number of primes below x using the
/// Lagarias-Miller-Odlyzko algorithm.
/// Run time: O(x^(2/3)) operations, O(x^(1/3) * log log x) space.
///
int64_t pi_lmo_parallel1(int64_t x, int threads)
{
#ifdef _OPENMP

  if (x < 2)
    return 0;

  double beta = 0.5;
  double alpha = max(1.0, log(log((double) x)) * beta);
  int64_t x13 = iroot<3>(x);
  int64_t y = (int64_t)(x13 * alpha);

  vector<int32_t> mu = make_moebius(y);
  vector<int32_t> lpf = make_least_prime_factor(y);
  vector<int32_t> primes;
  primes.push_back(0);
  primesieve::generate_primes(y, &primes);

  int64_t pi_y = primes.size() - 1;
  int64_t c = min(PhiTiny::MAX_A, pi_y);
  int64_t s1, s2, p2;

  threads = get_omp_threads(threads);
  omp_set_nested(true);

  #pragma omp parallel sections num_threads(threads)
  {
    #pragma omp section
    s2 = S2(x, y, pi_y, c, primes, lpf , mu, threads - 1);
    #pragma omp section
    {
      s1 = S1(x, y, c, primes, lpf , mu);
      p2 = P2(x, y);
    }
  }
  omp_set_nested(false);

  int64_t phi = s1 + s2;
  int64_t sum = phi + pi_y - 1 - p2;

  return sum;
#else
  return pi_lmo4(x);
#endif
}

} // namespace primecount
<commit_msg>Refactoring<commit_after>///
/// @file  pi_lmo_parallel1.cpp
/// @brief Parallel implementation of the Lagarias-Miller-Odlyzko
///        prime counting algorithm using OpenMP. This implementation
///        is derived from pi_lmo4(x).
///
/// Copyright (C) 2014 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include "internal.hpp"
#include "PhiTiny.hpp"
#include "pmath.hpp"
#include "tos_counters.hpp"

#include <primesieve.hpp>
#include <stdint.h>
#include <algorithm>
#include <vector>

#ifdef _OPENMP
  #include <omp.h>
  #include "get_omp_threads.hpp"
#endif

using namespace std;

#ifdef _OPENMP

namespace {

template <typename T1, typename T2>
void cross_off(int64_t prime,
               int64_t low,
               int64_t high,
               int64_t& next_multiple,
               T1& sieve,
               T2& counters)
{
  int64_t segment_size = sieve.size();
  int64_t k = next_multiple + prime * (~next_multiple & 1);

  for (; k < high; k += prime * 2)
  {
    if (sieve[k - low])
    {
      sieve[k - low] = 0;
      cnt_update(counters, k - low, segment_size);
    }
  }

  next_multiple = k;
}

/// Calculate the contribution of the special leaves.
/// @pre y > 0 && c > 1
///
int64_t S2(int64_t x,
           int64_t y,
           int64_t pi_y,
           int64_t c,
           vector<int32_t>& primes,
           vector<int32_t>& lpf,
           vector<int32_t>& mu,
           int threads)
{
  int64_t limit = x / y + 1;
  int64_t segment_size = next_power_of_2(isqrt(limit));
  int64_t segments = (limit + segment_size - 1) / segment_size;
  int64_t segments_per_thread = (segments + threads - 1) / threads;
  int64_t S2_total = 0;

  vector<vector<int64_t> > phi(threads);
  vector<vector<int64_t> > mu_sum(threads);

  for (int i = 0; i < threads; i++)
  {
    phi[i].resize(primes.size(), 0);
    mu_sum[i].resize(primes.size(), 0);
  }

  #pragma omp parallel for num_threads(threads) reduction(+: S2_total)
  for (int i = 0; i < threads; i++)
  {
    vector<char> sieve(segment_size);
    vector<int32_t> counters(segment_size);
    vector<int32_t>::iterator iter = primes.begin() + 1;
    vector<int64_t> next;
    next.reserve(primes.size());
    next.push_back(0);

    int64_t start_idx = i * segments_per_thread;
    int64_t stop_idx = min(start_idx + segments_per_thread, segments);
    int64_t S2_thread = 0;

    for (int64_t j = start_idx; j < stop_idx; j++)
    {
      fill(sieve.begin(), sieve.end(), 1);

      // Current segment = interval [low, high[
      int64_t low = j * segment_size + 1;
      int64_t high = min(low + segment_size, limit);

      // Initialize sieve indices
      while (iter != primes.end() && isquare(*iter) < high)
      {
        int64_t prime = *iter++;
        int64_t next_multiple = ((low + prime - 1) / prime) * prime;
        next.push_back(next_multiple);
      }

      for (int64_t b = 1; b <= c; b++)
      {
        int64_t k = next[b];
        for (int64_t prime = primes[b]; k < high; k += prime)
          sieve[k - low] = 0;
        next[b] = k;
      }

      // Initialize special tree data structure from sieve
      cnt_finit(sieve, counters, segment_size);

      for (int64_t b = c + 1; b < pi_y; b++)
      {
        int64_t prime = primes[b];
        int64_t min_m = max(x / (prime * high), y / prime);
        int64_t max_m = min(x / (prime * low), y);

        if (prime >= max_m)
          break;

        for (int64_t m = max_m; m > min_m; m--)
        {
          if (mu[m] != 0 && prime < lpf[m])
          {
            // phi_xn is the number of currently unsieved elements in
            // the interval [start_idx * segment_size + 1, x / n].
            // This is not the entire contribution of the special leaf
            // n, the missing contribution is the number of unsieved
            // elements in the interval [1, start_idx * segment_size].
            // The missing contribution is reconstructed and added
            // once all threads have finished.
            //
            int64_t n = prime * m;
            int64_t count = cnt_query(counters, (x / n) - low);
            int64_t phi_xn = phi[i][b] + count;
            S2_thread -= mu[m] * phi_xn;
            mu_sum[i][b] -= mu[m];
          }
        }

        phi[i][b] += cnt_query(counters, high - 1 - low);
        cross_off(prime, low, high, next[b], sieve, counters);
      }
    }

    S2_total += S2_thread;
  }

  // Once all threads have finished reconstruct and add the
  // missing contribution of all special leaves. This must
  // be done in order as each thread (i) requires the sum of
  // the phi values from the previous threads.
  //
  for (int i = 1; i < threads; i++)
  {
    for (size_t j = 0; j < phi[i].size(); j++)
    {
      S2_total += phi[i - 1][j] * mu_sum[i][j];
      phi[i][j] += phi[i - 1][j];
    }
  }

  return S2_total;
}

} // namespace

#endif /* _OPENMP */

namespace primecount {

/// Calculate the number of primes below x using the
/// Lagarias-Miller-Odlyzko algorithm.
/// Run time: O(x^(2/3)) operations, O(x^(1/3) * log log x) space.
///
int64_t pi_lmo_parallel1(int64_t x, int threads)
{
#ifdef _OPENMP

  if (x < 2)
    return 0;

  double beta = 0.5;
  double alpha = max(1.0, log(log((double) x)) * beta);
  int64_t x13 = iroot<3>(x);
  int64_t y = (int64_t)(x13 * alpha);

  vector<int32_t> mu = make_moebius(y);
  vector<int32_t> lpf = make_least_prime_factor(y);
  vector<int32_t> primes;
  primes.push_back(0);
  primesieve::generate_primes(y, &primes);

  int64_t pi_y = primes.size() - 1;
  int64_t c = min(PhiTiny::MAX_A, pi_y);
  int64_t s1, s2, p2;

  threads = get_omp_threads(threads);
  omp_set_nested(true);

  #pragma omp parallel sections num_threads(threads)
  {
    #pragma omp section
    s2 = S2(x, y, pi_y, c, primes, lpf , mu, max(1, threads - 1));
    #pragma omp section
    {
      s1 = S1(x, y, c, primes, lpf , mu);
      p2 = P2(x, y);
    }
  }
  omp_set_nested(false);

  int64_t phi = s1 + s2;
  int64_t sum = phi + pi_y - 1 - p2;

  return sum;
#else
  return pi_lmo4(x);
#endif
}

} // namespace primecount
<|endoftext|>


<commit_before>///
/// @file  pi_lmo_parallel3.cpp
/// @brief Parallel implementation of the Lagarias-Miller-Odlyzko
///        prime counting algorithm using OpenMP. This implementation
///        is based on pi_lmo_parallel2(x) but has an improved
///        load balancing.
///
/// Copyright (C) 2014 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include "internal.hpp"
#include "PhiTiny.hpp"
#include "pmath.hpp"
#include "tos_counters.hpp"
#include "pi_bsearch.hpp"

#include <primesieve.hpp>
#include <stdint.h>
#include <algorithm>
#include <vector>

#ifdef _OPENMP
  #include <omp.h>
  #include "get_omp_threads.hpp"
#endif

using namespace std;

#ifdef _OPENMP

namespace {

template <typename T1, typename T2>
void cross_off(int64_t prime,
               int64_t low,
               int64_t high,
               int64_t& next_multiple,
               T1& sieve,
               T2& counters)
{
  int64_t segment_size = sieve.size();
  int64_t k = next_multiple + prime * (~next_multiple & 1);

  for (; k < high; k += prime * 2)
  {
    if (sieve[k - low])
    {
      sieve[k - low] = 0;
      cnt_update(counters, k - low, segment_size);
    }
  }
  next_multiple = k;
}

/// Compute the S2 contribution for the interval
/// [low_process, low_process + segments * segment_size[.
/// The missing special leaf contributions for the interval
/// [1, low_process[ are later reconstructed and added in
/// the calling (parent) S2 function.
///
int64_t S2_thread(int64_t x,
                  int64_t y,
                  int64_t pi_sqrty,
                  int64_t pi_y,
                  int64_t c,
                  int64_t limit,
                  int64_t low_process,
                  int64_t segments,
                  int64_t segment_size,
                  int64_t segments_per_thread,
                  int64_t thread_num,
                  vector<int32_t>& pi,
                  vector<int32_t>& primes,
                  vector<int32_t>& lpf,
                  vector<int32_t>& mu,
                  vector<int64_t>& mu_sum,
                  vector<int64_t>& phi)
{
  vector<char> sieve(segment_size);
  vector<int32_t> counters(segment_size);
  vector<int64_t> next;
  phi.resize(primes.size(), 0);
  mu_sum.resize(primes.size(), 0);
  next.reserve(primes.size());
  next.push_back(0);

  int64_t start_idx = segments_per_thread * thread_num;
  int64_t stop_idx = min(segments_per_thread * (thread_num + 1), segments);
  int64_t low_thread = low_process + segment_size * start_idx;
  int64_t S2_thread = 0;

  // Initialize next multiples
  for (size_t j = 1; j < primes.size(); j++)
  {
    int64_t prime = primes[j];
    int64_t next_multiple = ((low_thread + prime - 1) / prime) * prime;
    next.push_back(next_multiple);
  }

  // Process the segments corresponding to the current thread
  for (int64_t j = start_idx; j < stop_idx; j++)
  {
    fill(sieve.begin(), sieve.end(), 1);

    // Current segment = interval [low, high[
    int64_t low = low_process + segment_size * j;
    int64_t high = min(low + segment_size, limit);
    int64_t special_leaf_threshold = max(x / high, y);
    int64_t b = 1;

    for (; b <= c; b++)
    {
      int64_t k = next[b];
      for (int64_t prime = primes[b]; k < high; k += prime)
        sieve[k - low] = 0;
      next[b] = k;
    }

    // Initialize special tree data structure from sieve
    cnt_finit(sieve, counters, segment_size);

    // For c + 1 <= b < pi_sqrty
    // Find all special leaves: n = primes[b] * m, with mu[m] != 0 and primes[b] < lpf[m]
    // Such that: low <= x / n < high
    for (; b < pi_sqrty; b++)
    {
      int64_t prime = primes[b];
      int64_t m_min = max(x / (prime * high), y / prime);
      int64_t m_max = min(x / (prime * low), y);

      if (prime >= m_max)
        break;

      for (int64_t m = m_max; m > m_min; m--)
      {
        if (mu[m] != 0 && prime < lpf[m])
        {
          int64_t n = prime * m;
          int64_t count = cnt_query(counters, (x / n) - low);
          int64_t phi_xn = phi[b] + count;
          S2_thread -= mu[m] * phi_xn;
          mu_sum[b] -= mu[m];
        }
      }

      phi[b] += cnt_query(counters, (high - 1) - low);
      cross_off(prime, low, high, next[b], sieve, counters);
    }

    // For pi_sqrty <= b < pi_y
    // Find all special leaves: n = primes[b] * prime2
    // Such that: low <= x / n < high
    for (; b < pi_y; b++)
    {
      int64_t prime = primes[b];
      int64_t l = pi[min(x / (prime * low), y)];
      if (prime >= primes[l])
        break;

      special_leaf_threshold = max(prime * prime, special_leaf_threshold);

      for (; prime * primes[l] > special_leaf_threshold; l--)
      {
        int64_t n = prime * primes[l];
        int64_t count = cnt_query(counters, (x / n) - low);
        int64_t phi_xn = phi[b] + count;
        S2_thread += phi_xn;
        mu_sum[b]++;
      }

      phi[b] += cnt_query(counters, (high - 1) - low);
      cross_off(prime, low, high, next[b], sieve, counters);
    }
  }

  return S2_thread;
}

/// Calculate the contribution of the special leaves.
/// This is a parallel implementation with advanced load balancing.
/// As most special leaves tend to be in the first segments we
/// start off with a small segment size and few segments
/// per thread, after each iteration we dynamically increase
/// the segment size and the segments per thread.
/// @pre y > 0 && c > 1
///
int64_t S2(int64_t x,
           int64_t y,
           int64_t pi_y,
           int64_t c,
           vector<int32_t>& primes,
           vector<int32_t>& lpf,
           vector<int32_t>& mu,
           int threads)
{
  int64_t S2_result = 0;
  int64_t low = 1;
  int64_t limit = x / y + 1;
  int64_t logx = (int64_t) max(1.0, log((double) x));
  int64_t segment_size_max = next_power_of_2(isqrt(limit));
  int64_t segment_size = next_power_of_2(segment_size_max / logx);
  int64_t segments_per_thread = 1;
  int64_t pi_sqrty = pi_bsearch(primes, isqrt(y));

  vector<int32_t> pi = make_pi(y);
  vector<int64_t> phi_total(primes.size(), 0);

  while (low < limit)
  {
    double time = omp_get_wtime();
    int64_t segments = (limit - low + segment_size - 1) / segment_size;
    threads = in_between(1, threads, segments);
    segments_per_thread = in_between(1, segments_per_thread, (segments + threads - 1) / threads);

    vector<vector<int64_t> > phi(threads);
    vector<vector<int64_t> > mu_sum(threads);

    #pragma omp parallel for num_threads(threads) reduction(+: S2_result)
    for (int i = 0; i < threads; i++)
    {
      S2_result += S2_thread(x, y, pi_sqrty, pi_y, c, limit, low, segments,
          segment_size, segments_per_thread, i, pi, primes, lpf, mu, mu_sum[i], phi[i]);
    }

    // Once all threads have finished reconstruct and add the
    // missing contribution of all special leaves. This must
    // be done in order as each thread (i) requires the sum of
    // the phi values from the previous threads.
    //
    for (size_t j = 0; j < phi[0].size(); j++)
    {
      S2_result += phi_total[j] * mu_sum[0][j];
      phi[0][j] += phi_total[j];
    }

    for (int i = 1; i < threads; i++)
    {
      for (size_t j = 1; j < phi[i].size(); j++)
      {
        S2_result += phi[i - 1][j] * mu_sum[i][j];
        phi[i][j] += phi[i - 1][j];
      }
    }

    phi_total = phi.back();

    low += segments_per_thread * threads * segment_size;

    // Dynamically increase segment_size and segments_per_thread
    // if the running time is less than a certain threshold.
    // We start off with a small segment size and few segments
    // per thread as most special leaves are in the first segments
    // whereas later on there are very few special leaves.
    //
    double seconds = omp_get_wtime() - time;
    if (seconds < 10)
    {
      segment_size = min(segment_size * 2, segment_size_max);
      segments_per_thread *= 2;
    }
  }

  return S2_result;
}

} // namespace

#endif /* _OPENMP */

namespace primecount {

/// Calculate the number of primes below x using the
/// Lagarias-Miller-Odlyzko algorithm.
/// Run time: O(x^(2/3)) operations, O(x^(1/3) * log log x) space.
///
int64_t pi_lmo_parallel3(int64_t x, int threads)
{
#ifdef _OPENMP

  if (x < 2)
    return 0;

  double beta = 1.0;
  double alpha = max(1.0, log(log((double) x)) * beta);
  int64_t x13 = iroot<3>(x);
  int64_t y = (int64_t)(x13 * alpha);

  vector<int32_t> mu = make_moebius(y);
  vector<int32_t> lpf = make_least_prime_factor(y);
  vector<int32_t> primes;
  primes.push_back(0);
  primesieve::generate_primes(y, &primes);

  int64_t pi_y = primes.size() - 1;
  int64_t c = min(PhiTiny::MAX_A, pi_y);
  int64_t s1, s2, p2;

  threads = get_omp_threads(threads);
  omp_set_nested(true);

  #pragma omp parallel sections num_threads(threads)
  {
    #pragma omp section
    s2 = S2(x, y, pi_y, c, primes, lpf , mu, threads - 1);
    #pragma omp section
    {
      s1 = S1(x, y, c, primes, lpf , mu);
      p2 = P2(x, y);
    }
  }
  omp_set_nested(false);

  int64_t phi = s1 + s2;
  int64_t sum = phi + pi_y - 1 - p2;

  return sum;
#else
  return pi_lmo5(x);
#endif
}

} // namespace primecount
<commit_msg>Refactoring<commit_after>///
/// @file  pi_lmo_parallel3.cpp
/// @brief Parallel implementation of the Lagarias-Miller-Odlyzko
///        prime counting algorithm using OpenMP. This implementation
///        is based on pi_lmo_parallel2(x) but has an improved
///        load balancing.
///
/// Copyright (C) 2014 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include "internal.hpp"
#include "PhiTiny.hpp"
#include "pmath.hpp"
#include "tos_counters.hpp"
#include "pi_bsearch.hpp"

#include <primesieve.hpp>
#include <stdint.h>
#include <algorithm>
#include <vector>

#ifdef _OPENMP
  #include <omp.h>
  #include "get_omp_threads.hpp"
#endif

using namespace std;

#ifdef _OPENMP

namespace {

template <typename T1, typename T2>
void cross_off(int64_t prime,
               int64_t low,
               int64_t high,
               int64_t& next_multiple,
               T1& sieve,
               T2& counters)
{
  int64_t segment_size = sieve.size();
  int64_t k = next_multiple + prime * (~next_multiple & 1);

  for (; k < high; k += prime * 2)
  {
    if (sieve[k - low])
    {
      sieve[k - low] = 0;
      cnt_update(counters, k - low, segment_size);
    }
  }
  next_multiple = k;
}

/// Compute the S2 contribution for the interval
/// [low_process, low_process + segments * segment_size[.
/// The missing special leaf contributions for the interval
/// [1, low_process[ are later reconstructed and added in
/// the calling (parent) S2 function.
///
int64_t S2_thread(int64_t x,
                  int64_t y,
                  int64_t pi_sqrty,
                  int64_t pi_y,
                  int64_t c,
                  int64_t limit,
                  int64_t low_process,
                  int64_t segments,
                  int64_t segment_size,
                  int64_t segments_per_thread,
                  int64_t thread_num,
                  vector<int32_t>& pi,
                  vector<int32_t>& primes,
                  vector<int32_t>& lpf,
                  vector<int32_t>& mu,
                  vector<int64_t>& mu_sum,
                  vector<int64_t>& phi)
{
  vector<char> sieve(segment_size);
  vector<int32_t> counters(segment_size);
  vector<int64_t> next;
  phi.resize(primes.size(), 0);
  mu_sum.resize(primes.size(), 0);
  next.reserve(primes.size());
  next.push_back(0);

  int64_t start_idx = segments_per_thread * thread_num;
  int64_t stop_idx = min(segments_per_thread * (thread_num + 1), segments);
  int64_t low_thread = low_process + segment_size * start_idx;
  int64_t S2_thread = 0;

  // Initialize next multiples
  for (size_t j = 1; j < primes.size(); j++)
  {
    int64_t prime = primes[j];
    int64_t next_multiple = ((low_thread + prime - 1) / prime) * prime;
    next.push_back(next_multiple);
  }

  // Process the segments corresponding to the current thread
  for (int64_t j = start_idx; j < stop_idx; j++)
  {
    fill(sieve.begin(), sieve.end(), 1);

    // Current segment = interval [low, high[
    int64_t low = low_process + segment_size * j;
    int64_t high = min(low + segment_size, limit);
    int64_t special_leaf_threshold = max(x / high, y);
    int64_t b = 1;

    for (; b <= c; b++)
    {
      int64_t k = next[b];
      for (int64_t prime = primes[b]; k < high; k += prime)
        sieve[k - low] = 0;
      next[b] = k;
    }

    // Initialize special tree data structure from sieve
    cnt_finit(sieve, counters, segment_size);

    // For c + 1 <= b < pi_sqrty
    // Find all special leaves: n = primes[b] * m, with mu[m] != 0 and primes[b] < lpf[m]
    // Such that: low <= x / n < high
    for (; b < pi_sqrty; b++)
    {
      int64_t prime = primes[b];
      int64_t m_min = max(x / (prime * high), y / prime);
      int64_t m_max = min(x / (prime * low), y);

      if (prime >= m_max)
        break;

      for (int64_t m = m_max; m > m_min; m--)
      {
        if (mu[m] != 0 && prime < lpf[m])
        {
          int64_t n = prime * m;
          int64_t count = cnt_query(counters, (x / n) - low);
          int64_t phi_xn = phi[b] + count;
          S2_thread -= mu[m] * phi_xn;
          mu_sum[b] -= mu[m];
        }
      }

      phi[b] += cnt_query(counters, (high - 1) - low);
      cross_off(prime, low, high, next[b], sieve, counters);
    }

    // For pi_sqrty <= b < pi_y
    // Find all special leaves: n = primes[b] * prime2
    // Such that: low <= x / n < high
    for (; b < pi_y; b++)
    {
      int64_t prime = primes[b];
      int64_t l = pi[min(x / (prime * low), y)];
      if (prime >= primes[l])
        break;

      special_leaf_threshold = max(prime * prime, special_leaf_threshold);

      for (; prime * primes[l] > special_leaf_threshold; l--)
      {
        int64_t n = prime * primes[l];
        int64_t count = cnt_query(counters, (x / n) - low);
        int64_t phi_xn = phi[b] + count;
        S2_thread += phi_xn;
        mu_sum[b]++;
      }

      phi[b] += cnt_query(counters, (high - 1) - low);
      cross_off(prime, low, high, next[b], sieve, counters);
    }
  }

  return S2_thread;
}

/// Calculate the contribution of the special leaves.
/// This is a parallel implementation with advanced load balancing.
/// As most special leaves tend to be in the first segments we
/// start off with a small segment size and few segments
/// per thread, after each iteration we dynamically increase
/// the segment size and the segments per thread.
/// @pre y > 0 && c > 1
///
int64_t S2(int64_t x,
           int64_t y,
           int64_t pi_y,
           int64_t c,
           vector<int32_t>& primes,
           vector<int32_t>& lpf,
           vector<int32_t>& mu,
           int threads)
{
  int64_t S2_total = 0;
  int64_t low = 1;
  int64_t limit = x / y + 1;
  int64_t sqrt_limit = isqrt(limit);
  int64_t logx = max(1, ilog(x));
  int64_t min_segment_size = 1 << 6;
  int64_t segment_size = next_power_of_2(sqrt_limit / (logx * threads));
  int64_t segments_per_thread = 1;
  int64_t pi_sqrty = pi_bsearch(primes, isqrt(y));

  vector<int32_t> pi = make_pi(y);
  vector<int64_t> phi_total(primes.size(), 0);
  segment_size = max(segment_size, min_segment_size);

  while (low < limit)
  {
    double seconds = omp_get_wtime();
    int64_t segments = (limit - low + segment_size - 1) / segment_size;
    threads = in_between(1, threads, segments);
    segments_per_thread = in_between(1, segments_per_thread, (segments + threads - 1) / threads);

    vector<vector<int64_t> > phi(threads);
    vector<vector<int64_t> > mu_sum(threads);

    #pragma omp parallel for num_threads(threads) reduction(+: S2_total)
    for (int i = 0; i < threads; i++)
    {
      S2_total += S2_thread(x, y, pi_sqrty, pi_y, c, limit, low, segments,
          segment_size, segments_per_thread, i, pi, primes, lpf, mu, mu_sum[i], phi[i]);
    }

    low += segments_per_thread * threads * segment_size;
    seconds = omp_get_wtime() - seconds;

    // Once all threads have finished reconstruct and add the
    // missing contribution of all special leaves. This must
    // be done in order as each thread (i) requires the sum of
    // the phi values from the previous threads.
    //
    for (int i = 0; i < threads; i++)
    {
      for (size_t j = 1; j < phi[i].size(); j++)
      {
        S2_total += phi_total[j] * mu_sum[i][j];
        phi_total[j] += phi[i][j];
      }
    }

    // Dynamically increase segment_size or segments_per_thread
    // if the running time is less than a certain threshold.
    // We start off with a small segment size and few segments
    // per thread as most special leaves are in the first segments
    // whereas later on there are very few special leaves.
    //
    if (low > sqrt_limit && seconds < 10)
    {
      if (segment_size < sqrt_limit)
        segment_size = next_power_of_2(segment_size * 2);
      else
        segments_per_thread *= 2;
    }
  }

  return S2_total;
}

} // namespace

#endif /* _OPENMP */

namespace primecount {

/// Calculate the number of primes below x using the
/// Lagarias-Miller-Odlyzko algorithm.
/// Run time: O(x^(2/3)) operations, O(x^(1/3) * log log x) space.
///
int64_t pi_lmo_parallel3(int64_t x, int threads)
{
#ifdef _OPENMP

  if (x < 2)
    return 0;

  double beta = 1.0;
  double alpha = max(1.0, log(log((double) x)) * beta);
  int64_t x13 = iroot<3>(x);
  int64_t y = (int64_t)(x13 * alpha);

  vector<int32_t> mu = make_moebius(y);
  vector<int32_t> lpf = make_least_prime_factor(y);
  vector<int32_t> primes;
  primes.push_back(0);
  primesieve::generate_primes(y, &primes);

  int64_t pi_y = primes.size() - 1;
  int64_t c = min(PhiTiny::MAX_A, pi_y);
  int64_t s1, s2, p2;

  threads = get_omp_threads(threads);
  omp_set_nested(true);

  #pragma omp parallel sections num_threads(threads)
  {
    #pragma omp section
    s2 = S2(x, y, pi_y, c, primes, lpf , mu, max(1, threads - 1));
    #pragma omp section
    {
      s1 = S1(x, y, c, primes, lpf , mu);
      p2 = P2(x, y);
    }
  }
  omp_set_nested(false);

  int64_t phi = s1 + s2;
  int64_t sum = phi + pi_y - 1 - p2;

  return sum;
#else
  return pi_lmo5(x);
#endif
}

} // namespace primecount
<|endoftext|>


<commit_before><commit_msg>Fix getNodePath load now correctly data<commit_after><|endoftext|>


<commit_before>/** \file   rss_aggregator.cc
 *  \brief  Downloads and aggregates RSS feeds.
 *  \author Dr. Johannes Ruscheinski (johannes.ruscheinski@uni-tuebingen.de)
 *
 *  \copyright 2018-2021 Universitätsbibliothek Tübingen.  All rights reserved.
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <algorithm>
#include <map>
#include <unordered_map>
#include <unordered_set>
#include <cinttypes>
#include <cstring>
#include "Compiler.h"
#include "DbConnection.h"
#include "DbResultSet.h"
#include "DnsUtil.h"
#include "Downloader.h"
#include "EmailSender.h"
#include "FileUtil.h"
#include "HtmlUtil.h"
#include "IniFile.h"
#include "RegexMatcher.h"
#include "SqlUtil.h"
#include "StringUtil.h"
#include "SyndicationFormat.h"
#include "UBTools.h"
#include "util.h"
#include "XmlWriter.h"

namespace {

[[noreturn]] void Usage() {
    ::Usage("[--use-web-proxy] subsystem_type email_address xml_output_path\n"
            "where subsystem_type must be one of {ixtheo,relbib,krimdok}");
}

// These must be in sync with the sizes in the vufind.rss_aggregator table!
const size_t MAX_ITEM_ID_LENGTH(768);
const size_t MAX_ITEM_URL_LENGTH(1000);
const size_t MAX_ITEM_TITLE_LENGTH(1000);
const size_t MAX_SERIAL_NAME_LENGTH(200);

struct HarvestedRSSItem {
    SyndicationFormat::Item item_;
    std::string feed_title_;
    std::string feed_url_;

    HarvestedRSSItem(const SyndicationFormat::Item item, const std::string feed_title, const std::string feed_url)
        : item_(item), feed_title_(feed_title), feed_url_(feed_url) {}
};

struct ChannelDesc {
    std::string title_;
    std::string link_;
public:
    ChannelDesc(const std::string &title, const std::string &link)
        : title_(title), link_(link) { }
};

const std::map<std::string, ChannelDesc> subsystem_type_to_channel_desc_map = {
    { "relbib",  ChannelDesc("RelBib Aggregator",  "https://relbib.de/")                },
    { "ixtheo",  ChannelDesc("IxTheo Aggregator",  "https://itheo.de/")                 },
    { "krimdok", ChannelDesc("KrimDok Aggregator", "https://krimdok.uni-tuebingen.de/") },
};

std::string GetChannelDescEntry(const std::string &subsystem_type, const std::string &entry) {
    const auto subsystem_type_and_channel_desc(subsystem_type_to_channel_desc_map.find(subsystem_type));
    if (subsystem_type_and_channel_desc == subsystem_type_to_channel_desc_map.cend())
        LOG_ERROR("unknown subsystem type \"" + subsystem_type + "\"!");

    if (entry == "title")
        return subsystem_type_and_channel_desc->second.title_;
    if (entry == "link")
        return subsystem_type_and_channel_desc->second.link_;
    LOG_ERROR("unknown entry name \"" + entry + "\"!");
}

void WriteRSSFeedXMLOutput(const std::string &subsystem_type, const std::vector<HarvestedRSSItem> &harvested_items,
                           XmlWriter * const xml_writer)
{
    xml_writer->openTag("rss", { { "version", "2.0" }, { "xmlns:tuefind", "https://github.com/ubtue/tuefind" } });
    xml_writer->openTag("channel");
    xml_writer->writeTagsWithData("title", GetChannelDescEntry(subsystem_type, "title"));
    xml_writer->writeTagsWithData("link", GetChannelDescEntry(subsystem_type, "link"));
    xml_writer->writeTagsWithData("description", "RSS Aggregator");

    for (const auto &harvested_item : harvested_items) {
        xml_writer->openTag("item");

        const auto title(harvested_item.item_.getTitle());
        if (not title.empty())
            xml_writer->writeTagsWithData("title", harvested_item.item_.getTitle());

        xml_writer->writeTagsWithData("link", harvested_item.item_.getLink());

        const auto description(HtmlUtil::ShortenText(harvested_item.item_.getDescription(), /*max_length = */500));
        if (not description.empty())
            xml_writer->writeTagsWithData("description", description);

        xml_writer->writeTagsWithData("pubDate",
                                      TimeUtil::TimeTToString(harvested_item.item_.getPubDate(), TimeUtil::RFC822_FORMAT,
                                                                     TimeUtil::UTC));
        xml_writer->writeTagsWithData("guid", harvested_item.item_.getId());
        xml_writer->writeTagsWithData("tuefind:rss_title", harvested_item.feed_title_);
        xml_writer->writeTagsWithData("tuefind:rss_url", harvested_item.feed_url_);
        xml_writer->closeTag("item", /* suppress_indent */ false);
    }

    xml_writer->closeTag("channel");
    xml_writer->closeTag("rss");
}

// \return true if the item was new, else false.
bool ProcessRSSItem(const std::string &feed_id, const SyndicationFormat::Item &item, DbConnection * const db_connection) {
    const std::string item_id(item.getId());
    db_connection->queryOrDie("SELECT insertion_time FROM rss_aggregator WHERE item_id='"
                              + db_connection->escapeString(item_id) + "'");
    const DbResultSet result_set(db_connection->getLastResultSet());
    if (not result_set.empty())
        return false;

    const std::string item_url(item.getLink());
    if (item_url.empty()) {
        LOG_WARNING("got an item w/o a URL, ID is \"" + item.getId() + "\"");
        return false;
    }

    db_connection->insertIntoTableOrDie("rss_aggregator",
                                        {
                                            { "rss_feeds_id",     StringUtil::Truncate(MAX_SERIAL_NAME_LENGTH, feed_id)        },
                                            { "item_id",          StringUtil::Truncate(MAX_ITEM_ID_LENGTH, item_id)            },
                                            { "item_url",         StringUtil::Truncate(MAX_ITEM_URL_LENGTH, item_url)          },
                                            { "item_title",       StringUtil::Truncate(MAX_ITEM_TITLE_LENGTH, item.getTitle()) },
                                            { "item_description", item.getDescription()                                        },
                                            { "pub_date",         SqlUtil::TimeTToDatetime(item.getPubDate())                  }
                                        },
                                        DbConnection::DuplicateKeyBehaviour::DKB_REPLACE);

    return true;
}

// \return the number of new items.
unsigned ProcessFeed(const std::string &feed_id, const std::string &feed_name, const std::string &feed_url,
                     const std::string &title_suppression_regex_str, const std::string &strptime_format,
                     Downloader * const downloader, DbConnection * const db_connection, const unsigned downloader_time_limit)
{
    SyndicationFormat::AugmentParams augment_params;
    augment_params.strptime_format_ = strptime_format;

    const auto title_suppression_regex(
        title_suppression_regex_str.empty() ? nullptr : RegexMatcher::RegexMatcherFactoryOrDie(title_suppression_regex_str));

    unsigned new_item_count(0);
    if (not downloader->newUrl(feed_url, downloader_time_limit))
        LOG_WARNING(feed_name + ": failed to download the feed: " + downloader->getLastErrorMessage());
    else {
        std::string error_message;
        std::unique_ptr<SyndicationFormat> syndication_format(
            SyndicationFormat::Factory(downloader->getMessageBody(), augment_params, &error_message));
        if (unlikely(syndication_format == nullptr))
            LOG_WARNING("failed to parse feed: " + error_message);
        else {
            for (const auto &item : *syndication_format) {
                if (title_suppression_regex != nullptr and title_suppression_regex->matched(item.getTitle())) {
                    LOG_INFO("Suppressed item because of title: \"" + StringUtil::ShortenText(item.getTitle(), 40) + "\".");
                    continue; // Skip suppressed item.
                }

                if (ProcessRSSItem(feed_id, item, db_connection))
                    ++new_item_count;
            }
        }
    }

    return new_item_count;
}

const unsigned HARVEST_TIME_WINDOW(60); // days

size_t SelectItems(const std::string &subsystem_type, DbConnection * const db_connection,
                   std::vector<HarvestedRSSItem> * const harvested_items)
{
    db_connection->queryOrDie("SELECT id,feed_name FROM tuefind_rss_feeds WHERE subsystem_types LIKE '%"
                              + subsystem_type + "%'");
    DbResultSet feeds_result_set(db_connection->getLastResultSet());

    const std::string NOW_AS_SQL_DATETIME(SqlUtil::TimeTToDatetime(std::time(nullptr) - HARVEST_TIME_WINDOW * 86400));
    while (const auto feed_row = feeds_result_set.getNextRow()) {
        db_connection->queryOrDie("SELECT item_title,item_description,item_url,item_id,pub_date FROM rss_aggregator "
                                  "WHERE pub_date >= '" + NOW_AS_SQL_DATETIME + "' AND rss_feeds_id = "
                                  + feed_row["id"] + " ORDER BY pub_date DESC");
        DbResultSet result_set(db_connection->getLastResultSet());
        while (const DbRow row = result_set.getNextRow())
            harvested_items->emplace_back(SyndicationFormat::Item(row["item_title"], row["item_description"], row["item_url"],
                                                                  row["item_id"], SqlUtil::DatetimeToTimeT(row["pub_date"])),
                                                                  feed_row["feed_name"], row["feed_url"]);
    }

    return harvested_items->size();
}

const unsigned DEFAULT_XML_INDENT_AMOUNT(2);

int ProcessFeeds(const std::string &subsystem_type, const std::string &xml_output_filename,
                 DbConnection * const db_connection, Downloader * const downloader)
{
    db_connection->queryOrDie("SELECT * FROM tuefind_rss_feeds WHERE subsystem_type LIKE '%" + subsystem_type + "%'");
    auto result_set(db_connection->getLastResultSet());
    while (const auto row = result_set.getNextRow()) {
        LOG_INFO("Processing feed \"" + row["feed_name"] + "\".");
        const unsigned new_item_count(ProcessFeed(row["id"], row["feed_name"], row["feed_url"],
                                                  row.getValue("title_suppression_regex"),
                                                  row.getValue("strptime_format"), downloader, db_connection,
                                                  StringUtil::ToUnsigned(row["downloader_time_limit"])));
        LOG_INFO("Downloaded " + std::to_string(new_item_count) + " new items.");
    }

    std::vector<HarvestedRSSItem> harvested_items;
    const auto feed_item_count(SelectItems(subsystem_type, db_connection, &harvested_items));

    // scoped here so that we flush and close the output file right away
    {
        XmlWriter xml_writer(FileUtil::OpenOutputFileOrDie(xml_output_filename).release(),
                             XmlWriter::WriteTheXmlDeclaration, DEFAULT_XML_INDENT_AMOUNT);
        WriteRSSFeedXMLOutput(subsystem_type, harvested_items, &xml_writer);
    }
    LOG_INFO("Created our feed with " + std::to_string(feed_item_count) + " items from the last "
             + std::to_string(HARVEST_TIME_WINDOW) + " days.");

    return EXIT_SUCCESS;
}

} // unnamed namespace

int Main(int argc, char *argv[]) {
    if (argc != 4 and argc != 5)
        Usage();

    Downloader::Params params;
    if (argc == 5) {
        if (std::strcmp(argv[0], "--use-web-proxy") != 0)
            Usage();
        --argc, ++argv;
        params.proxy_host_and_port_ = UBTools::GetUBWebProxyURL();
    }
    Downloader downloader(params);

    const std::string subsystem_type(argv[1]);
    if (subsystem_type != "ixtheo" and subsystem_type != "relbib" and subsystem_type != "krimdok")
        LOG_ERROR("subsystem_type must be one of {ixtheo,relbib,krimdok}!");

    const std::string email_address(argv[2]);
    const std::string xml_output_filename(argv[3]);

    DbConnection db_connection;

    try {
        return ProcessFeeds(subsystem_type, xml_output_filename, &db_connection, &downloader);
    } catch (const std::runtime_error &x) {
        const auto program_basename(FileUtil::GetBasename(::progname));
        const auto subject(program_basename + " failed on " + DnsUtil::GetHostname()
                           + " (subsystem_type: " + subsystem_type + ")");
        const auto message_body("caught exception: " + std::string(x.what()));
        if (EmailSender::SimplerSendEmail("no_reply@ub.uni-tuebingen.de", { email_address }, subject, message_body,
                                          EmailSender::VERY_HIGH) < 299)
            return EXIT_FAILURE;
        else
            LOG_ERROR("failed to send an email error report!");
    }
}
<commit_msg>Various fixes.<commit_after>/** \file   rss_aggregator.cc
 *  \brief  Downloads and aggregates RSS feeds.
 *  \author Dr. Johannes Ruscheinski (johannes.ruscheinski@uni-tuebingen.de)
 *
 *  \copyright 2018-2021 Universitätsbibliothek Tübingen.  All rights reserved.
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <algorithm>
#include <map>
#include <unordered_map>
#include <unordered_set>
#include <cinttypes>
#include <cstring>
#include "Compiler.h"
#include "DbConnection.h"
#include "DbResultSet.h"
#include "DnsUtil.h"
#include "Downloader.h"
#include "EmailSender.h"
#include "FileUtil.h"
#include "HtmlUtil.h"
#include "IniFile.h"
#include "RegexMatcher.h"
#include "SqlUtil.h"
#include "StringUtil.h"
#include "SyndicationFormat.h"
#include "UBTools.h"
#include "util.h"
#include "VuFind.h"
#include "XmlWriter.h"

namespace {

[[noreturn]] void Usage() {
    ::Usage("[--use-web-proxy] subsystem_type email_address xml_output_path\n"
            "where subsystem_type must be one of {ixtheo,relbib,krimdok}");
}

// These must be in sync with the sizes in the vufind.rss_aggregator table!
const size_t MAX_ITEM_ID_LENGTH(768);
const size_t MAX_ITEM_URL_LENGTH(1000);
const size_t MAX_ITEM_TITLE_LENGTH(1000);
const size_t MAX_SERIAL_NAME_LENGTH(200);

struct HarvestedRSSItem {
    SyndicationFormat::Item item_;
    std::string feed_title_;
    std::string feed_url_;

    HarvestedRSSItem(const SyndicationFormat::Item item, const std::string feed_title, const std::string feed_url)
        : item_(item), feed_title_(feed_title), feed_url_(feed_url) {}
};

struct ChannelDesc {
    std::string title_;
    std::string link_;
public:
    ChannelDesc(const std::string &title, const std::string &link)
        : title_(title), link_(link) { }
};

const std::map<std::string, ChannelDesc> subsystem_type_to_channel_desc_map = {
    { "relbib",  ChannelDesc("RelBib Aggregator",  "https://relbib.de/")                },
    { "ixtheo",  ChannelDesc("IxTheo Aggregator",  "https://itheo.de/")                 },
    { "krimdok", ChannelDesc("KrimDok Aggregator", "https://krimdok.uni-tuebingen.de/") },
};

std::string GetChannelDescEntry(const std::string &subsystem_type, const std::string &entry) {
    const auto subsystem_type_and_channel_desc(subsystem_type_to_channel_desc_map.find(subsystem_type));
    if (subsystem_type_and_channel_desc == subsystem_type_to_channel_desc_map.cend())
        LOG_ERROR("unknown subsystem type \"" + subsystem_type + "\"!");

    if (entry == "title")
        return subsystem_type_and_channel_desc->second.title_;
    if (entry == "link")
        return subsystem_type_and_channel_desc->second.link_;
    LOG_ERROR("unknown entry name \"" + entry + "\"!");
}

void WriteRSSFeedXMLOutput(const std::string &subsystem_type, const std::vector<HarvestedRSSItem> &harvested_items,
                           XmlWriter * const xml_writer)
{
    xml_writer->openTag("rss", { { "version", "2.0" }, { "xmlns:tuefind", "https://github.com/ubtue/tuefind" } });
    xml_writer->openTag("channel");
    xml_writer->writeTagsWithData("title", GetChannelDescEntry(subsystem_type, "title"));
    xml_writer->writeTagsWithData("link", GetChannelDescEntry(subsystem_type, "link"));
    xml_writer->writeTagsWithData("description", "RSS Aggregator");

    for (const auto &harvested_item : harvested_items) {
        xml_writer->openTag("item");

        const auto title(harvested_item.item_.getTitle());
        if (not title.empty())
            xml_writer->writeTagsWithData("title", harvested_item.item_.getTitle());

        xml_writer->writeTagsWithData("link", harvested_item.item_.getLink());

        const auto description(HtmlUtil::ShortenText(harvested_item.item_.getDescription(), /*max_length = */500));
        if (not description.empty())
            xml_writer->writeTagsWithData("description", description);

        xml_writer->writeTagsWithData("pubDate",
                                      TimeUtil::TimeTToString(harvested_item.item_.getPubDate(), TimeUtil::RFC822_FORMAT,
                                                                     TimeUtil::UTC));
        xml_writer->writeTagsWithData("guid", harvested_item.item_.getId());
        xml_writer->writeTagsWithData("tuefind:rss_title", harvested_item.feed_title_);
        xml_writer->writeTagsWithData("tuefind:rss_url", harvested_item.feed_url_);
        xml_writer->closeTag("item", /* suppress_indent */ false);
    }

    xml_writer->closeTag("channel");
    xml_writer->closeTag("rss");
}

// \return true if the item was new, else false.
bool ProcessRSSItem(const std::string &feed_id, const SyndicationFormat::Item &item, DbConnection * const db_connection) {
    const std::string item_id(item.getId());
    db_connection->queryOrDie("SELECT insertion_time FROM tuefind_rss_items WHERE item_id='"
                              + db_connection->escapeString(item_id) + "'");
    const DbResultSet result_set(db_connection->getLastResultSet());
    if (not result_set.empty())
        return false;

    const std::string item_url(item.getLink());
    if (item_url.empty()) {
        LOG_WARNING("got an item w/o a URL, ID is \"" + item.getId() + "\"");
        return false;
    }

    db_connection->insertIntoTableOrDie("tuefind_rss_items",
                                        {
                                            { "rss_feeds_id",     StringUtil::Truncate(MAX_SERIAL_NAME_LENGTH, feed_id)        },
                                            { "item_id",          StringUtil::Truncate(MAX_ITEM_ID_LENGTH, item_id)            },
                                            { "item_url",         StringUtil::Truncate(MAX_ITEM_URL_LENGTH, item_url)          },
                                            { "item_title",       StringUtil::Truncate(MAX_ITEM_TITLE_LENGTH, item.getTitle()) },
                                            { "item_description", item.getDescription()                                        },
                                            { "pub_date",         SqlUtil::TimeTToDatetime(item.getPubDate())                  }
                                        },
                                        DbConnection::DuplicateKeyBehaviour::DKB_REPLACE);

    return true;
}

// \return the number of new items.
unsigned ProcessFeed(const std::string &feed_id, const std::string &feed_name, const std::string &feed_url,
                     const std::string &title_suppression_regex_str, const std::string &strptime_format,
                     Downloader * const downloader, DbConnection * const db_connection, const unsigned downloader_time_limit)
{
    SyndicationFormat::AugmentParams augment_params;
    augment_params.strptime_format_ = strptime_format;

    const auto title_suppression_regex(
        title_suppression_regex_str.empty() ? nullptr : RegexMatcher::RegexMatcherFactoryOrDie(title_suppression_regex_str));

    unsigned new_item_count(0);
    if (not downloader->newUrl(feed_url, downloader_time_limit))
        LOG_WARNING(feed_name + ": failed to download the feed: " + downloader->getLastErrorMessage());
    else {
        std::string error_message;
        std::unique_ptr<SyndicationFormat> syndication_format(
            SyndicationFormat::Factory(downloader->getMessageBody(), augment_params, &error_message));
        if (unlikely(syndication_format == nullptr))
            LOG_WARNING("failed to parse feed: " + error_message);
        else {
            for (const auto &item : *syndication_format) {
                if (title_suppression_regex != nullptr and title_suppression_regex->matched(item.getTitle())) {
                    LOG_INFO("Suppressed item because of title: \"" + StringUtil::ShortenText(item.getTitle(), 40) + "\".");
                    continue; // Skip suppressed item.
                }

                if (ProcessRSSItem(feed_id, item, db_connection))
                    ++new_item_count;
            }
        }
    }

    return new_item_count;
}

const unsigned HARVEST_TIME_WINDOW(60); // days

struct FeedNameAndURL {
    std::string name_;
    std::string url_;
public:
    FeedNameAndURL() = default;
    FeedNameAndURL(const FeedNameAndURL &other) = default;
    FeedNameAndURL(const std::string &name, const std::string &url)
        : name_(name), url_(url) { }
};

size_t SelectItems(const std::string &subsystem_type, DbConnection * const db_connection,
                   std::vector<HarvestedRSSItem> * const harvested_items)
{
    db_connection->queryOrDie("SELECT id,feed_name,feed_url FROM tuefind_rss_feeds WHERE subsystem_types LIKE '%"
                              + subsystem_type + "%'");
    DbResultSet feeds_result_set(db_connection->getLastResultSet());
    std::unordered_map<std::string, FeedNameAndURL> feed_ids_to_names_and_urls_map;
    while (const auto row = feeds_result_set.getNextRow())
        feed_ids_to_names_and_urls_map[row["id"]] = FeedNameAndURL(row["feed_name"], row["feed_url"]);

    const std::string NOW_AS_SQL_DATETIME(SqlUtil::TimeTToDatetime(std::time(nullptr) - HARVEST_TIME_WINDOW * 86400));
    for (const auto &[feed_id, feed_name_and_url] : feed_ids_to_names_and_urls_map) {
        db_connection->queryOrDie("SELECT item_title,item_description,item_url,item_id,pub_date FROM tuefind_rss_items "
                                  "WHERE pub_date >= '" + NOW_AS_SQL_DATETIME + "' AND rss_feeds_id = "
                                  + feed_id + " ORDER BY pub_date DESC");
        DbResultSet result_set(db_connection->getLastResultSet());
        while (const DbRow row = result_set.getNextRow())
            harvested_items->emplace_back(SyndicationFormat::Item(row["item_title"], row["item_description"], row["item_url"],
                                                                  row["item_id"], SqlUtil::DatetimeToTimeT(row["pub_date"])),
                                                                  feed_name_and_url.name_, feed_name_and_url.url_);
    }

    return harvested_items->size();
}

const unsigned DEFAULT_XML_INDENT_AMOUNT(2);

int ProcessFeeds(const std::string &subsystem_type, const std::string &xml_output_filename,
                 DbConnection * const db_connection, Downloader * const downloader)
{
    db_connection->queryOrDie("SELECT * FROM tuefind_rss_feeds WHERE subsystem_types LIKE '%" + subsystem_type + "%'");
    auto result_set(db_connection->getLastResultSet());
    while (const auto row = result_set.getNextRow()) {
        LOG_INFO("Processing feed \"" + row["feed_name"] + "\".");
        const unsigned new_item_count(ProcessFeed(row["id"], row["feed_name"], row["feed_url"],
                                                  row.getValue("title_suppression_regex"),
                                                  row.getValue("strptime_format"), downloader, db_connection,
                                                  StringUtil::ToUnsigned(row["downloader_time_limit"])));
        LOG_INFO("Downloaded " + std::to_string(new_item_count) + " new items.");
    }

    std::vector<HarvestedRSSItem> harvested_items;
    const auto feed_item_count(SelectItems(subsystem_type, db_connection, &harvested_items));

    // scoped here so that we flush and close the output file right away
    {
        XmlWriter xml_writer(FileUtil::OpenOutputFileOrDie(xml_output_filename).release(),
                             XmlWriter::WriteTheXmlDeclaration, DEFAULT_XML_INDENT_AMOUNT);
        WriteRSSFeedXMLOutput(subsystem_type, harvested_items, &xml_writer);
    }
    LOG_INFO("Created our feed with " + std::to_string(feed_item_count) + " items from the last "
             + std::to_string(HARVEST_TIME_WINDOW) + " days.");

    return EXIT_SUCCESS;
}

} // unnamed namespace

int Main(int argc, char *argv[]) {
    if (argc != 4 and argc != 5)
        Usage();

    Downloader::Params params;
    if (argc == 5) {
        if (std::strcmp(argv[0], "--use-web-proxy") != 0)
            Usage();
        --argc, ++argv;
        params.proxy_host_and_port_ = UBTools::GetUBWebProxyURL();
    }
    Downloader downloader(params);

    const std::string subsystem_type(argv[1]);
    if (subsystem_type != "ixtheo" and subsystem_type != "relbib" and subsystem_type != "krimdok")
        LOG_ERROR("subsystem_type must be one of {ixtheo,relbib,krimdok}!");

    const std::string email_address(argv[2]);
    const std::string xml_output_filename(argv[3]);

    const auto db_connection(VuFind::GetDbConnection());

    try {
        return ProcessFeeds(subsystem_type, xml_output_filename, db_connection.get(), &downloader);
    } catch (const std::runtime_error &x) {
        const auto program_basename(FileUtil::GetBasename(::progname));
        const auto subject(program_basename + " failed on " + DnsUtil::GetHostname()
                           + " (subsystem_type: " + subsystem_type + ")");
        const auto message_body("caught exception: " + std::string(x.what()));
        if (EmailSender::SimplerSendEmail("no_reply@ub.uni-tuebingen.de", { email_address }, subject, message_body,
                                          EmailSender::VERY_HIGH) < 299)
            return EXIT_FAILURE;
        else
            LOG_ERROR("failed to send an email error report!");
    }
}
<|endoftext|>