Merge pull request #147 from jad-hamza/merge-calls

Implement @samarion's fix for mergeCalls code explosion

Author: samarion

build.sbt

@@ -85,7 +85,7 @@ script := {
                        |
                        |SCALACLASSPATH=$cp
                        |
-                       |java -Xmx2G -Xms512M -Xss64M -classpath "$${SCALACLASSPATH}" -Dscala.usejavacp=true inox.Main $$@ 2>&1
+                       |java -Xmx2G -Xms512M -Xss64M -classpath "$${SCALACLASSPATH}" -Dscala.usejavacp=true inox.Main "$$@" 2>&1
                        |""".stripMargin)
     file.setExecutable(true)
   } catch {

src/main/scala/inox/ast/SymbolOps.scala

@@ -1028,188 +1028,6 @@ trait SymbolOps { self: TypeOps =>
     case _ => None
   } (expr)
 
-  private def mergeFunctions(expr: Expr)(implicit opts: PurityOptions): Expr = {
-    type Bindings = Seq[(ValDef, Expr)]
-    implicit class BindingsWrapper(bindings: Bindings) {
-      def merge(that: Bindings): Bindings = (bindings ++ that).distinct
-      def wrap(that: Expr): Expr = bindings.foldRight(that) { case ((vd, e), b) =>
-        val freshVd = vd.freshen
-        val fb = replaceFromSymbols(Map(vd -> freshVd.toVariable), b)
-        let(freshVd, e, fb)
-      }
-    }
-
-    def liftCalls(e: Expr): Expr = {
-      def rec(e: Expr): (Expr, Bindings) = e match {
-        case Let(i, fi @ FunctionInvocation(id, tps, args), body) =>
-          val (recArgs, argsBindings) = args.map(rec).unzip
-          val (recBody, bodyBindings) = rec(body)
-          val recBindings = argsBindings.flatten ++ bodyBindings
-          (recBody, ((argsBindings.flatten :+ (i -> FunctionInvocation(id, tps, recArgs).copiedFrom(fi))) ++ recBindings).distinct)
-
-        case fi @ FunctionInvocation(id, tps, args) =>
-          val v = Variable.fresh("call", fi.tfd.getType, true)
-          val (recArgs, recBindings) = args.map(rec).unzip
-          (v, recBindings.flatten.distinct :+ (v.toVal -> FunctionInvocation(id, tps, recArgs).copiedFrom(fi)))
-
-        case Let(i, v, b) =>
-          val (vr, vBindings) = rec(v)
-          val (br, bBindings) = rec(b)
-          (br, vBindings merge ((i -> vr) +: bBindings))
-
-        case Assume(pred, body) =>
-          val (recPred, predBindings) = rec(pred)
-          val (recBody, bodyBindings) = rec(body)
-          (Assume(recPred, bodyBindings wrap recBody), predBindings)
-
-        case IfExpr(cond, thenn, elze) =>
-          val (recCond, condBindings) = rec(cond)
-          val (recThen, thenBindings) = rec(thenn)
-          val (recElse, elseBindings) = rec(elze)
-          (IfExpr(recCond, thenBindings wrap recThen, elseBindings wrap recElse), condBindings)
-
-        case And(e +: es) =>
-          val (recE, eBindings) = rec(e)
-          val newEs = es.map { e =>
-            val (recE, eBindings) = rec(e)
-            eBindings wrap recE
-          }
-          (And(recE +: newEs), eBindings)
-
-        case Or(e +: es) =>
-          val (recE, eBindings) = rec(e)
-          val newEs = es.map { e =>
-            val (recE, eBindings) = rec(e)
-            eBindings wrap recE
-          }
-          (Or(recE +: newEs), eBindings)
-
-        case Implies(lhs, rhs) =>
-          val (recLhs, lhsBindings) = rec(lhs)
-          val (recRhs, rhsBindings) = rec(rhs)
-          (Implies(recLhs, rhsBindings wrap recRhs), lhsBindings)
-
-        case v: Variable => (v, Seq.empty)
-
-        case ex @ VariableExtractor(vs) if vs.nonEmpty =>
-          val Operator(subs, recons) = ex
-          val recSubs = subs.map(rec)
-          (recons(recSubs.map { case (e, bindings) => bindings wrap e }), Seq.empty)
-
-        case Lambda(Seq(), body) =>
-          val (recBody, bodyBindings) = rec(body)
-          (Lambda(Seq(), bodyBindings wrap recBody), Seq.empty)
-
-        case Operator(es, recons) =>
-          val (recEs, esBindings) = es.map(rec).unzip
-          (recons(recEs), esBindings.flatten.distinct)
-      }
-
-      val (newE, bindings) = rec(e)
-      bindings wrap newE
-    }
-
-    def mergeCalls(e: Expr): Expr = {
-      def evCalls(e: Expr): Map[TypedFunDef, Set[(Path, Seq[Expr])]] = {
-
-        var inLambda = false
-        var pathFis: Seq[(Path, FunctionInvocation)] = Seq.empty
-        transformWithPC(e, false /* ignore calls within types */)((e, path, op) => e match {
-          case l: Lambda =>
-            val old = inLambda
-            inLambda = true
-            val nl = op.sup(l, path)
-            inLambda = old
-            nl
-
-          case fi: FunctionInvocation =>
-            def freeVars(elements: Seq[Path.Element]): Set[Variable] = {
-              val path = Path(elements)
-              path.freeVariables ++ variablesOf(fi) -- path.bound.map(_.toVariable)
-            }
-
-            val elements = path.elements.foldRight(Some(Seq[Path.Element]()): Option[Seq[Path.Element]]) {
-              case (_, None) => None
-              case (Path.OpenBound(vd), Some(elements)) =>
-                if (freeVars(elements) contains vd.toVariable) None
-                else Some(elements) // No need to keep open bounds as we'll flatten to a condition
-              case (cb @ Path.CloseBound(vd, e), Some(elements)) =>
-                if (exists { case fi: FunctionInvocation => true case _ => false }(e)) None
-                else if (freeVars(elements) contains vd.toVariable) Some(cb +: elements)
-                else Some(elements)
-              case (c @ Path.Condition(cond), Some(elements)) =>
-                if (exists { case fi: FunctionInvocation => true case _ => false }(cond)) None
-                else Some(c +: elements)
-            }
-
-            if ((!inLambda || isPure(fi)) && elements.isDefined) {
-              pathFis :+= Path(elements.get) -> fi
-            }
-            op.sup(fi, path)
-
-          case _ =>
-            op.sup(e, path)
-        })
-
-        pathFis.groupBy(_._2.tfd).mapValues(_.map(p => (p._1, p._2.args)).toSet)
-      }
-
-      def replace(path: Path, oldE: Expr, newE: Expr, body: Expr): Expr =
-        transformWithPC(body, false /* ignore calls within types */)((e, env, op) => {
-          if ((path.bindings.toSet subsetOf env.bindings.toSet) &&
-            (path.bound.toSet subsetOf env.bound.toSet) &&
-            (path.conditions == env.conditions) && e == oldE) {
-            newE
-          } else {
-            op.sup(e, env)
-          }
-        })
-
-      postMap {
-        case IfExpr(cond, thenn, elze) =>
-          val condVar = Variable.fresh("cond", BooleanType(), true)
-          val condPath = Path(condVar)
-
-          def rec(bindings: Bindings, thenn: Expr, elze: Expr): (Bindings, Expr, Expr) = {
-            val thenCalls = evCalls(thenn)
-            val elseCalls = evCalls(elze)
-
-            (thenCalls.keySet & elseCalls.keySet).headOption match {
-              case Some(tfd) =>
-                val (pathThen, argsThen) = thenCalls(tfd).toSeq.sortBy(_._1.elements.size).head
-                val (pathElse, argsElse) = elseCalls(tfd).toSeq.sortBy(_._1.elements.size).head
-
-                val v = Variable.fresh("res", tfd.getType, true)
-                val condThen = Variable.fresh("condThen", BooleanType(), true)
-
-                val result = IfExpr(Or(condThen, freshenLocals((condPath.negate merge pathElse).toClause)),
-                  tfd.applied((argsThen zip argsElse).map { case (argThen, argElse) =>
-                    ifExpr(condThen, pathThen.bindings wrap argThen, pathElse.bindings wrap argElse)
-                  }), Choose(Variable.fresh("res", tfd.getType).toVal, BooleanLiteral(true)))
-
-                val newBindings = bindings ++ Seq(
-                  condThen.toVal -> freshenLocals((condPath merge pathThen).toClause),
-                  v.toVal -> result
-                )
-
-                val newThen = replace(pathThen, tfd.applied(argsThen), v, thenn)
-                val newElse = replace(pathElse, tfd.applied(argsElse), v, elze)
-                rec(newBindings, newThen, newElse)
-
-              case None => (bindings, thenn, elze)
-            }
-          }
-
-          val (bindings, newThen, newElse) = rec(Seq.empty, thenn, elze)
-          Some(((condVar.toVal -> cond) +: bindings) wrap ifExpr(condVar, newThen, newElse))
-
-        case _ => None
-      } (e)
-    }
-
-    mergeCalls(liftCalls(expr))
-  }
-
   private[inox] def simplifyFormula(e: Expr)(implicit ctx: Context, sem: symbols.Semantics): Expr = {
     implicit val simpOpts = SimplificationOptions(ctx)
     implicit val purityOpts = PurityOptions(ctx)
@@ -1219,8 +1037,7 @@ trait SymbolOps { self: TypeOps =>
         ((e: Expr) => simplifyGround(e))      compose
         ((e: Expr) => simplifyExpr(e))        compose
         ((e: Expr) => simplifyForalls(e))     compose
-        ((e: Expr) => simplifyAssumptions(e)) compose
-        ((e: Expr) => mergeFunctions(e))
+        ((e: Expr) => simplifyAssumptions(e))
       simp(e)
     } else {
       e

src/main/scala/inox/solvers/unrolling/TemplateGenerator.scala

@@ -75,6 +75,81 @@ trait TemplateGenerator { self: Templates =>
     tmplClauses + (pathVar -> p)
   }
 
+  def mergeCalls(pathVar: Variable, condVar: Variable, substMap: Map[Variable, Encoded],
+                thenClauses: TemplateClauses, elseClauses: TemplateClauses): TemplateClauses = {
+    val builder = new Builder(pathVar, substMap)
+    builder ++= thenClauses
+    builder ++= elseClauses
+
+    // Clear all guardedExprs in builder since we're going to transform them by merging calls.
+    // The transformed guardedExprs will be added to builder at the end of the function.
+    builder.guardedExprs = Map.empty
+
+    def collectCalls(expr: Expr): Set[FunctionInvocation] =
+      exprOps.collect { case fi: FunctionInvocation => Set(fi) case _ => Set.empty[FunctionInvocation] }(expr)
+    def countCalls(expr: Expr): Int =
+      exprOps.count { case fi: FunctionInvocation => 1 case _ => 0}(expr)
+    def replaceCall(call: FunctionInvocation, newExpr: Expr)(e: Expr): Expr =
+      exprOps.replace(Map(call -> newExpr), e)
+
+    def getCalls(guardedExprs: Map[Variable, Seq[Expr]]): Map[TypedFunDef, Seq[(FunctionInvocation, Set[Variable])]] =
+      (for { (b, es) <- guardedExprs.toSeq; e <- es; fi <- collectCalls(e) } yield (b -> fi))
+      .groupBy(_._2)
+      .mapValues(_.map(_._1).toSet)
+      .toSeq
+      .groupBy(_._1.tfd)
+      .mapValues(_.toList.distinct.sortBy(p => countCalls(p._1)))  // place inner calls first
+      .toMap
+
+    var thenGuarded = thenClauses._4
+    var elseGuarded = elseClauses._4
+
+    val thenCalls = getCalls(thenGuarded)
+    val elseCalls = getCalls(elseGuarded)
+
+    // We sort common function calls in order to merge nested calls first.
+    var toMerge: Seq[((FunctionInvocation, Set[Variable]), (FunctionInvocation, Set[Variable]))] =
+      (thenCalls.keySet & elseCalls.keySet)
+        .flatMap(tfd => thenCalls(tfd) zip elseCalls(tfd))
+        .toSeq
+        .sortBy(p => countCalls(p._1._1) + countCalls(p._2._1))
+
+    while (toMerge.nonEmpty) {
+      val ((thenCall, thenBlockers), (elseCall, elseBlockers)) = toMerge.head
+      toMerge = toMerge.tail
+
+      val newExpr: Variable = Variable.fresh("call", thenCall.tfd.getType, true)
+      builder.storeExpr(newExpr)
+
+      val replaceThen = replaceCall(thenCall, newExpr) _
+      val replaceElse = replaceCall(elseCall, newExpr) _
+
+      thenGuarded = thenGuarded.mapValues(_.map(replaceThen))
+      elseGuarded = elseGuarded.mapValues(_.map(replaceElse))
+      toMerge = toMerge.map(p => (
+        (replaceThen(p._1._1).asInstanceOf[FunctionInvocation], p._1._2),
+        (replaceElse(p._2._1).asInstanceOf[FunctionInvocation], p._2._2)
+      ))
+
+      val newBlocker: Variable = Variable.fresh("bm", BooleanType(), true)
+      builder.storeConds(thenBlockers ++ elseBlockers, newBlocker)
+      builder.iff(orJoin((thenBlockers ++ elseBlockers).toSeq), newBlocker)
+
+      val newArgs = (thenCall.args zip elseCall.args).map { case (thenArg, elseArg) =>
+        val (newArg, argClauses) = mkExprClauses(newBlocker, ifExpr(condVar, thenArg, elseArg), builder.localSubst)
+        builder ++= argClauses
+        newArg
+      }
+
+      val newCall = thenCall.tfd.applied(newArgs)
+      builder.storeGuarded(newBlocker, Equals(newExpr, newCall))
+    }
+
+    for ((b, es) <- thenGuarded; e <- es) builder.storeGuarded(b, e)
+    for ((b, es) <- elseGuarded; e <- es) builder.storeGuarded(b, e)
+    builder.result
+  }
+
   protected def mkExprStructure(
     pathVar: Variable,
     expr: Expr,
@@ -157,6 +232,13 @@ trait TemplateGenerator { self: Templates =>
       condTree += pathVar -> (condTree.getOrElse(pathVar, Set.empty) + id)
     }
 
+    def storeConds(pathVars: Set[Variable], id: Variable): Unit = {
+      condVars += id -> encodeSymbol(id)
+      for (pathVar <- pathVars) {
+        condTree += pathVar -> (condTree.getOrElse(pathVar, Set.empty) + id)
+      }
+    }
+
     @inline def encodedCond(id: Variable): Encoded = substMap.getOrElse(id, condVars(id))
 
     var exprVars = Map[Variable, Encoded]()
@@ -367,15 +449,16 @@ trait TemplateGenerator { self: Templates =>
           storeExpr(condVar)
 
           val crec = rec(pathVar, cond, None)
-          val trec = rec(newBool1, thenn, pol)
-          val erec = rec(newBool2, elze, pol)
-
           storeGuarded(pathVar, Equals(condVar, crec))
           iff(and(pathVar, condVar), newBool1)
           iff(and(pathVar, not(condVar)), newBool2)
 
-          storeGuarded(newBool1, Equals(newExpr, trec))
-          storeGuarded(newBool2, Equals(newExpr, erec))
+          val (trec, tClauses) = mkExprClauses(newBool1, thenn, localSubst, pol)
+          val (erec, eClauses) = mkExprClauses(newBool2, elze, localSubst, pol)
+          builder ++=  mergeCalls(pathVar, condVar, localSubst,
+                                  tClauses + (newBool1 -> Equals(newExpr, trec)),
+                                  eClauses + (newBool2 -> Equals(newExpr, erec)))
+
           newExpr
         }
       }