<commit_before>package ir

object ExprSimplifier {

  private def primeFactors(n: Int, i: Int =2) : List[Int] = {
    if (i >= n)
      return List(n)

    if (n % i == 0)
      i :: primeFactors(n/i, i)
    else
      primeFactors(n, i+1)
  }

  private def simplifyPow(pow: Pow): ArithExpr = {
    pow match {
      case Pow(Cst(0), Cst(0)) => throw new NotEvaluableException(pow.toString())
      case Pow(Cst(b), Cst(e)) =>
        val powDbl = scala.math.pow(b,e)
        if (powDbl.isValidInt)
          Cst(powDbl.toInt)
        else {
          if (e < 0) {
            // x>0, y>0, pow(x,-y) = pow(pow(x,y), -1)
            val powDbl = scala.math.pow(b,-e)
            if (powDbl.isValidInt)
              Pow(powDbl.toInt, -1)
            else
              pow
          } else
            pow
        }
      case Pow(base, Cst(0)) => Cst(1)
      case Pow(base, Cst(1)) => base
      case Pow(Cst(0), _) => Cst(0)
      case Pow(Cst(1), _) => Cst(1)

      // This rule should not be used since prodFactorsSimplify turns x^m * x^n into x^(m+n)
      // x^(m+n) => x^m * x^n
      case Pow(base, Sum(terms)) =>
        val newProd = Prod(terms.map(t => simplifyPow(Pow(base, t))))
        val simpNewProd = simplifyProd(newProd)
        simpNewProd

      // (x*y)^(n) => x^n * y^n
      case Pow(Prod(terms), exp) => simplifyProd(Prod(terms.map(t => simplifyPow(Pow(t, exp)))))

      // (x^e1)^e2 => x^(e1*e2)
      case Pow(Pow(b, e1), e2) => simplifyPow(Pow(b, simplifyProd(e1 * e2)))

      // x^(a*log(x,b)*c) => b^(a*b)
      case Pow(x,Prod(factors)) =>
        var b : ArithExpr = null
        val idx = factors.indexWhere({
          case Log(_,b1) =>
            b = b1
            true
          case _ => false
        })
        if (idx != -1)
          simplifyPow(Pow(b, simplifyProd(Prod(factors.diff(List(factors(idx)))))))
        else
          pow

      // x^log(x,b) => b
      case Pow(x1,Log(x2,b)) if x1 == x2 => b



      case _ => pow
    }
  }


  private def flattenSum(sum: Sum) : Sum = {
    var result = List[ArithExpr]()

    sum.terms.foreach {
      case s: Sum => result = result ++ s.terms
      case t => result = t :: result
    }

    Sum(result)
  }

  private def simplifySum(sum: Sum): ArithExpr = {
    var resultSum = sum

    resultSum = flattenSum(resultSum)

    // constant folding
    resultSum = sumCstFolding(resultSum)

    if (resultSum.terms.length == 1)
      return resultSum.terms(0)

    resultSum
  }



  // constant folding
  private def cstFolding(l : List[ArithExpr], op : ((Double,Double) => Double), neutral: Int) : List[ArithExpr] = {

    // fixed point iteration until everything has been folded
    var change = true
    var curLen = l.length
    var curResult = l
    while (change) {

      var cstVal : Int = neutral
      var newResult = List[ArithExpr]()

      curResult.foreach(e =>
        try {
          val cstValDbl : Double = op(e.evalDbl(),cstVal)
          if (cstValDbl.isValidInt)
            cstVal = cstValDbl.toInt
          else
            newResult = e :: newResult<commit_msg>Added basic simplification for terms in a sum<commit_after>package ir

object ExprSimplifier {

  private def primeFactors(n: Int, i: Int =2) : List[Int] = {
    if (i >= n)
      return List(n)

    if (n % i == 0)
      i :: primeFactors(n/i, i)
    else
      primeFactors(n, i+1)
  }

  private def simplifyPow(pow: Pow): ArithExpr = {
    pow match {
      case Pow(Cst(0), Cst(0)) => throw new NotEvaluableException(pow.toString())
      case Pow(Cst(b), Cst(e)) =>
        val powDbl = scala.math.pow(b,e)
        if (powDbl.isValidInt)
          Cst(powDbl.toInt)
        else {
          if (e < 0) {
            // x>0, y>0, pow(x,-y) = pow(pow(x,y), -1)
            val powDbl = scala.math.pow(b,-e)
            if (powDbl.isValidInt)
              Pow(powDbl.toInt, -1)
            else
              pow
          } else
            pow
        }
      case Pow(base, Cst(0)) => Cst(1)
      case Pow(base, Cst(1)) => base
      case Pow(Cst(0), _) => Cst(0)
      case Pow(Cst(1), _) => Cst(1)

      // This rule should not be used since prodFactorsSimplify turns x^m * x^n into x^(m+n)
      // x^(m+n) => x^m * x^n
      case Pow(base, Sum(terms)) =>
        val newProd = Prod(terms.map(t => simplifyPow(Pow(base, t))))
        val simpNewProd = simplifyProd(newProd)
        simpNewProd

      // (x*y)^(n) => x^n * y^n
      case Pow(Prod(terms), exp) => simplifyProd(Prod(terms.map(t => simplifyPow(Pow(t, exp)))))

      // (x^e1)^e2 => x^(e1*e2)
      case Pow(Pow(b, e1), e2) => simplifyPow(Pow(b, simplifyProd(e1 * e2)))

      // x^(a*log(x,b)*c) => b^(a*b)
      case Pow(x,Prod(factors)) =>
        var b : ArithExpr = null
        val idx = factors.indexWhere({
          case Log(_,b1) =>
            b = b1
            true
          case _ => false
        })
        if (idx != -1)
          simplifyPow(Pow(b, simplifyProd(Prod(factors.diff(List(factors(idx)))))))
        else
          pow

      // x^log(x,b) => b
      case Pow(x1,Log(x2,b)) if x1 == x2 => b



      case _ => pow
    }
  }


  private def flattenSum(sum: Sum) : Sum = {
    var result = List[ArithExpr]()

    sum.terms.foreach {
      case s: Sum => result = result ++ s.terms
      case t => result = t :: result
    }

    Sum(result)
  }

  private def simplifySum(sum: Sum): ArithExpr = {
    var resultSum = sum

    resultSum = flattenSum(resultSum)

    resultSum = simplifySumTerms(resultSum)

    // constant folding
    resultSum = sumCstFolding(resultSum)

    if (resultSum.terms.length == 1)
      return resultSum.terms(0)

    resultSum
  }

  private def simplifySumTerms(sum: Sum): Sum = {

    val terms: List[ArithExpr] = sum.terms

    def simplifyTerm(i: Int, ae: ArithExpr): Option[Sum] = {
      val vars = Var.getVars(ae)

      for (j <- 0 until vars.length) {
        val v = vars(j)

        for (k <- i + 1 until terms.length) {
          val term = terms(k)

          if (ArithExpr.contains(term, v)) {
            val e: Sum = ae / v + term / v
            val simplified = ExprSimplifier.simplify(e)

            if (!simplified.isInstanceOf[Sum])
              return Some(simplifySumTerms(Sum(ExprSimplifier.simplify(v * simplified) :: terms.slice(0, i) ++ terms.slice(i + 1, k) ++ terms.slice(k + 1, terms.length))))

          }
        }
      }
      None
    }

    for (i <- 0 until terms.length) {
      terms(i) match {
        case p: Var =>
          simplifyTerm(i, p) match {
            case Some(toReturn) => return toReturn
            case None =>
          }

        case p: Product =>
          simplifyTerm(i, p) match {
            case Some(toReturn) => return toReturn
            case None =>
          }

        case t =>
      }
    }

    sum
  }

  // constant folding
  private def cstFolding(l : List[ArithExpr], op : ((Double,Double) => Double), neutral: Int) : List[ArithExpr] = {

    // fixed point iteration until everything has been folded
    var change = true
    var curLen = l.length
    var curResult = l
    while (change) {

      var cstVal : Int = neutral
      var newResult = List[ArithExpr]()

      curResult.foreach(e =>
        try {
          val cstValDbl : Double = op(e.evalDbl(),cstVal)
          if (cstValDbl.isValidInt)
            cstVal = cstValDbl.toInt
          else
            newResult = e :: newResult
{'fraction_non_alphanumeric': 0.14274128142741282, 'fraction_numerical': 0.010002703433360368, 'mean_word_length': 1.8947574334898278, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 66, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3598524', 'n_tokens_mistral': 2753, 'n_tokens_neox': 2594, 'n_words': 618}