From 1cef4fa54b5eed93aa1da25796df81129585b0f1 Mon Sep 17 00:00:00 2001
From: Etienne Kneuss <etienne.kneuss@epfl.ch>
Date: Wed, 4 Feb 2015 14:50:27 +0100
Subject: [PATCH] --timeout is for entire repair now

---
 runTests.sh                                   |  62 +++---
 src/main/scala/leon/repair/RepairPhase.scala  |   4 +-
 src/main/scala/leon/repair/RepairResult.scala |   5 +-
 src/main/scala/leon/repair/Repairman.scala    | 186 +++++++++---------
 .../solvers/TimeoutAssumptionSolver.scala     |   2 +-
 .../scala/leon/solvers/TimeoutSolver.scala    |  43 +---
 .../synthesis/utils/ExpressionGrammar.scala   |   2 +-
 .../scala/leon/utils/InterruptManager.scala   |   2 +-
 src/main/scala/leon/utils/TimeoutFor.scala    |  55 ++++++
 9 files changed, 192 insertions(+), 169 deletions(-)
 create mode 100644 src/main/scala/leon/utils/TimeoutFor.scala

diff --git a/runTests.sh b/runTests.sh
index eaa74b8c0..8f571b316 100755
--- a/runTests.sh
+++ b/runTests.sh
@@ -10,40 +10,40 @@ echo "#" `hostname` >> $fullLog
 echo "#" `date +"%d.%m.%Y %T"` >> $fullLog
 echo "#" `git log -1 --pretty=format:"%h - %cd"` >> $fullLog
 echo "################################" >> $fullLog
-echo "#           Category,                 File,             function,   S, f.S,   Tms,   Fms,   Rms, verif?" >> $fullLog
+echo "#           Category,                 File,             function, p.S, fuS, foS,   Tms,   Fms,   Rms, verif?" >> $fullLog
 
 #All benchmarks:
-./leon --repair --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar1.scala
-./leon --repair --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar2.scala
-./leon --repair --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar3.scala
-./leon --repair --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar4.scala
-
-./leon --repair --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort3.scala
-./leon --repair --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort4.scala
-./leon --repair --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort5.scala
-./leon --repair --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort6.scala
-./leon --repair --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort7.scala
-./leon --repair --solvers=fairz3:enum --functions=insert   testcases/repair/HeapSort/HeapSort8.scala
-./leon --repair --solvers=fairz3:enum --functions=makeN    testcases/repair/HeapSort/HeapSort9.scala
-
-./leon --repair --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic1.scala 
-./leon --repair --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic2.scala 
-./leon --repair --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic3.scala 
-./leon --repair --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic4.scala 
-./leon --repair --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic5.scala 
-
-./leon --repair --solvers=fairz3:enum --functions=_pad     testcases/repair/List/List1.scala
-./leon --repair --solvers=fairz3:enum --functions=_ap      testcases/repair/List/List3.scala
-./leon --repair                       --functions=_drop    testcases/repair/List/List4.scala
-./leon --repair                       --functions=_replace testcases/repair/List/List5.scala
-./leon --repair --solvers=fairz3:enum --functions=_count   testcases/repair/List/List6.scala
-./leon --repair --solvers=fairz3:enum --functions=_find    testcases/repair/List/List7.scala
-./leon --repair --solvers=fairz3:enum --functions=_find    testcases/repair/List/List8.scala
-./leon --repair                       --functions=_find    testcases/repair/List/List9.scala
-./leon --repair --solvers=fairz3:enum --functions=_size    testcases/repair/List/List10.scala
-./leon --repair --solvers=fairz3:enum --functions=sum      testcases/repair/List/List11.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar1.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar2.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar3.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=desugar  testcases/repair/Desugar/Desugar4.scala
+
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort3.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort4.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort5.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort6.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=merge    testcases/repair/HeapSort/HeapSort7.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=insert   testcases/repair/HeapSort/HeapSort8.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=makeN    testcases/repair/HeapSort/HeapSort9.scala
+
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic1.scala 
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic2.scala 
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic3.scala 
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic4.scala 
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=nnf      testcases/repair/PropLogic/PropLogic5.scala 
+
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_pad     testcases/repair/List/List1.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_ap      testcases/repair/List/List3.scala
+./leon --repair --timeout=30                       --functions=_drop    testcases/repair/List/List4.scala
+./leon --repair --timeout=30                       --functions=_replace testcases/repair/List/List5.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_count   testcases/repair/List/List6.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_find    testcases/repair/List/List7.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_find    testcases/repair/List/List8.scala
+./leon --repair --timeout=30                       --functions=_find    testcases/repair/List/List9.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=_size    testcases/repair/List/List10.scala
+./leon --repair --timeout=30 --solvers=fairz3:enum --functions=sum      testcases/repair/List/List11.scala
 
 # Average results
 cat $log >> $fullLog
-awk '{ total1 += $6; total2 += $7; total3 += $8; count++ } END { printf "#%69s Avg: %5d, %5d, %5d\n\n", "", total1/count, total2/count, total3/count }' $log >> $fullLog
+awk '{ total1 += $7; total2 += $8; total3 += $9; count++ } END { printf "#%74s Avg: %5d, %5d, %5d\n\n", "", total1/count, total2/count, total3/count }' $log >> $fullLog
 
diff --git a/src/main/scala/leon/repair/RepairPhase.scala b/src/main/scala/leon/repair/RepairPhase.scala
index 7714b679b..59bc87321 100644
--- a/src/main/scala/leon/repair/RepairPhase.scala
+++ b/src/main/scala/leon/repair/RepairPhase.scala
@@ -49,9 +49,9 @@ object RepairPhase extends LeonPhase[Program, Program] {
     if (toRepair.isEmpty) reporter.warning("No functions found with the given names")
     
     for (fd <- toRepair) {
-      val res = new Repairman(ctx, program, fd, verifTimeoutMs).repair()
+      val res = new Repairman(ctx, program, fd, verifTimeoutMs, verifTimeoutMs).repair()
       val fw = new FileWriter("repairs.last", true)
-      try { 
+      try {
         fw.write(res.toLine)
       } finally {
         fw.close
diff --git a/src/main/scala/leon/repair/RepairResult.scala b/src/main/scala/leon/repair/RepairResult.scala
index 6d0dba0b1..7bc8f256d 100644
--- a/src/main/scala/leon/repair/RepairResult.scala
+++ b/src/main/scala/leon/repair/RepairResult.scala
@@ -4,7 +4,8 @@ package repair
 import java.io.File
 
 case class RepairResult(f: File,
-                        size: Int = -1,
+                        psize: Int,
+                        fsize: Int = -1,
                         name: String = "?",
                         initTime: Option[Long] = None,
                         focusTime: Option[Long] = None,
@@ -17,7 +18,7 @@ case class RepairResult(f: File,
     val benchName = f.getName()
     val benchFun  = name
 
-    f"$benchCat%20s, $benchName%20s, $benchFun%20s, $size%3s, ${focusSize.getOrElse("")}%3s, ${initTime.getOrElse("")}%5s, ${focusTime.getOrElse("")}%5s, ${repairTime.getOrElse("")}%5s, ${repairTrusted.getOrElse("")}%5s\n"
+    f"$benchCat%20s, $benchName%20s, $benchFun%20s, $psize%3s, $fsize%3s, ${focusSize.getOrElse("")}%3s, ${initTime.getOrElse("")}%5s, ${focusTime.getOrElse("")}%5s, ${repairTime.getOrElse("")}%5s, ${repairTrusted.getOrElse("")}%5s\n"
   }
 
 }
diff --git a/src/main/scala/leon/repair/Repairman.scala b/src/main/scala/leon/repair/Repairman.scala
index 5989941f3..1ee872e64 100644
--- a/src/main/scala/leon/repair/Repairman.scala
+++ b/src/main/scala/leon/repair/Repairman.scala
@@ -24,7 +24,7 @@ import synthesis.Witnesses._
 import graph.DotGenerator
 import leon.utils.ASCIIHelpers.title
 
-class Repairman(ctx: LeonContext, initProgram: Program, fd: FunDef, verifTimeoutMs: Option[Long]) {
+class Repairman(ctx: LeonContext, initProgram: Program, fd: FunDef, verifTimeoutMs: Option[Long], repairTimeoutMs: Option[Long]) {
   val reporter = ctx.reporter
 
   var program = initProgram
@@ -37,104 +37,110 @@ class Repairman(ctx: LeonContext, initProgram: Program, fd: FunDef, verifTimeout
     case class NotValid(passing : List[Example], failing : List[Example]) extends VerificationResult
   }
 
-  var result = RepairResult(ctx.files.head, 0);
+  var result = RepairResult(ctx.files.head, visibleFunDefsFromMain(initProgram).foldLeft(0) {
+    case (s, f) => formulaSize(f.fullBody) + s
+  });
 
   import VRes._
   
   def repair(): RepairResult = {
-    reporter.info(ASCIIHelpers.title("1. Discovering tests for "+fd.id))
-    val t1 = new Timer().start
-    discoverTests match {
-      case Valid => 
-        reporter.info(s"Function ${fd.id} is found valid, no repair needed!")
-      case NotValid(passingTests, failingTests) =>
-        reporter.info(f" - Passing: ${passingTests.size}%3d")
-        reporter.info(f" - Failing: ${failingTests.size}%3d")
-    
-        val initTime = t1.stop
-        result = result.copy(initTime = Some(initTime))
-        reporter.info("Finished in "+initTime+"ms")
-    
-    
-        reporter.info(ASCIIHelpers.title("2. Locating/Focusing synthesis problem"))
-        val t2    = new Timer().start
-        val synth = getSynthesizer(passingTests, failingTests)
-        val ci    = synth.ci
-        val p     = synth.problem
-    
-        var solutions = List[Solution]()
-        val focusTime = t2.stop
-        result = result.copy(focusTime = Some(focusTime))
-    
-        reporter.info("Finished in "+focusTime+"ms")
-        reporter.info(ASCIIHelpers.title("3. Synthesizing"))
-        reporter.info(p)
-    
-    
-        val t3    = new Timer().start
-        synth.synthesize() match {
-          case (search, sols) =>
-            for (sol <- sols) {
-    
-              // Validate solution if not trusted
-              if (!sol.isTrusted) {
-                reporter.info("Found untrusted solution! Verifying...")
-                val (npr, fds) = synth.solutionToProgram(sol)
-    
-                getVerificationCounterExamples(fds.head, npr) match {
-                  case NotValid(_, ces) if !ces.isEmpty =>
-                    reporter.error("I ended up finding this counter example:\n"+ces.mkString("  |  "))
-    
-                  case NotValid(_, _) =>
-                    solutions ::= sol
-                    reporter.warning("Solution is not trusted!")
-                    result = result.copy(repairTrusted = Some(false))
-
-                  case Valid =>
-                    solutions ::= sol
-                    reporter.info("Solution was not trusted but post-validation passed!")
-                    result = result.copy(repairTrusted = Some(true))
+    val to = new TimeoutFor(ctx.interruptManager)
+
+    to.interruptAfter(repairTimeoutMs) {
+      reporter.info(ASCIIHelpers.title("1. Discovering tests for "+fd.id))
+      val t1 = new Timer().start
+      discoverTests match {
+        case Valid => 
+          reporter.info(s"Function ${fd.id} is found valid, no repair needed!")
+        case NotValid(passingTests, failingTests) =>
+          reporter.info(f" - Passing: ${passingTests.size}%3d")
+          reporter.info(f" - Failing: ${failingTests.size}%3d")
+      
+          val initTime = t1.stop
+          result = result.copy(initTime = Some(initTime))
+          reporter.info("Finished in "+initTime+"ms")
+      
+      
+          reporter.info(ASCIIHelpers.title("2. Locating/Focusing synthesis problem"))
+          val t2    = new Timer().start
+          val synth = getSynthesizer(passingTests, failingTests)
+          val ci    = synth.ci
+          val p     = synth.problem
+      
+          var solutions = List[Solution]()
+          val focusTime = t2.stop
+          result = result.copy(focusTime = Some(focusTime))
+      
+          reporter.info("Finished in "+focusTime+"ms")
+          reporter.info(ASCIIHelpers.title("3. Synthesizing"))
+          reporter.info(p)
+      
+      
+          val t3    = new Timer().start
+          synth.synthesize() match {
+            case (search, sols) =>
+              for (sol <- sols) {
+      
+                // Validate solution if not trusted
+                if (!sol.isTrusted) {
+                  reporter.info("Found untrusted solution! Verifying...")
+                  val (npr, fds) = synth.solutionToProgram(sol)
+      
+                  getVerificationCounterExamples(fds.head, npr) match {
+                    case NotValid(_, ces) if !ces.isEmpty =>
+                      reporter.error("I ended up finding this counter example:\n"+ces.mkString("  |  "))
+      
+                    case NotValid(_, _) =>
+                      solutions ::= sol
+                      reporter.warning("Solution is not trusted!")
+                      result = result.copy(repairTrusted = Some(false))
+
+                    case Valid =>
+                      solutions ::= sol
+                      reporter.info("Solution was not trusted but post-validation passed!")
+                      result = result.copy(repairTrusted = Some(true))
+                  }
+                } else {
+                  reporter.info("Found trusted solution!")
+                  solutions ::= sol
+                  result = result.copy(repairTrusted = Some(true))
                 }
-              } else {
-                reporter.info("Found trusted solution!")
-                solutions ::= sol
-                result = result.copy(repairTrusted = Some(true))
               }
-            }
-    
-            if (synth.settings.generateDerivationTrees) {
-              val dot = new DotGenerator(search.g)
-              dot.writeFile("derivation"+DotGenerator.nextId()+".dot")
-            }
-    
-            if (solutions.isEmpty) {
-              reporter.error(ASCIIHelpers.title("Failed to repair!"))
-            } else {
-              result = result.copy(repairTime = Some(t3.stop))
-
-              reporter.info(ASCIIHelpers.title("Repair successful:"))
-              for ((sol, i) <- solutions.reverse.zipWithIndex) {
-                reporter.info(ASCIIHelpers.subTitle("Solution "+(i+1)+":"))
-                val expr = sol.toSimplifiedExpr(ctx, program)
-                reporter.info(ScalaPrinter(expr));
+      
+              if (synth.settings.generateDerivationTrees) {
+                val dot = new DotGenerator(search.g)
+                dot.writeFile("derivation"+DotGenerator.nextId()+".dot")
               }
-              reporter.info(ASCIIHelpers.title("In context:"))
-    
-    
-              for ((sol, i) <- solutions.reverse.zipWithIndex) {
-                reporter.info(ASCIIHelpers.subTitle("Solution "+(i+1)+":"))
-                val expr = sol.toSimplifiedExpr(ctx, program)
-                val nfd = fd.duplicate;
-    
-                nfd.body = fd.body.map { b =>
-                  replace(Map(ci.source -> expr), b)
+      
+              if (solutions.isEmpty) {
+                reporter.error(ASCIIHelpers.title("Failed to repair!"))
+              } else {
+                result = result.copy(repairTime = Some(t3.stop))
+
+                reporter.info(ASCIIHelpers.title("Repair successful:"))
+                for ((sol, i) <- solutions.reverse.zipWithIndex) {
+                  reporter.info(ASCIIHelpers.subTitle("Solution "+(i+1)+":"))
+                  val expr = sol.toSimplifiedExpr(ctx, program)
+                  reporter.info(ScalaPrinter(expr));
                 }
-    
-                reporter.info(ScalaPrinter(nfd));
+                reporter.info(ASCIIHelpers.title("In context:"))
+      
+      
+                for ((sol, i) <- solutions.reverse.zipWithIndex) {
+                  reporter.info(ASCIIHelpers.subTitle("Solution "+(i+1)+":"))
+                  val expr = sol.toSimplifiedExpr(ctx, program)
+                  val nfd = fd.duplicate;
+      
+                  nfd.body = fd.body.map { b =>
+                    replace(Map(ci.source -> expr), b)
+                  }
+      
+                  reporter.info(ScalaPrinter(nfd));
+                }
+      
               }
-    
             }
-          }
+      }
     }
 
     result
@@ -154,7 +160,7 @@ class Repairman(ctx: LeonContext, initProgram: Program, fd: FunDef, verifTimeout
     reporter.info("Original body size: "+size)
     reporter.info("Focused expr size : "+focusSize)
 
-    result = result.copy(name = fd.id.name, size = size, focusSize = Some(focusSize))
+    result = result.copy(name = fd.id.name, fsize = size, focusSize = Some(focusSize))
 
     val guide = Guide(replacedExpr)
 
diff --git a/src/main/scala/leon/solvers/TimeoutAssumptionSolver.scala b/src/main/scala/leon/solvers/TimeoutAssumptionSolver.scala
index 958260d18..7fae46a9c 100644
--- a/src/main/scala/leon/solvers/TimeoutAssumptionSolver.scala
+++ b/src/main/scala/leon/solvers/TimeoutAssumptionSolver.scala
@@ -10,7 +10,7 @@ trait TimeoutAssumptionSolver extends TimeoutSolver with AssumptionSolver {
   abstract override def checkAssumptions(assumptions: Set[Expr]): Option[Boolean] = {
     optTimeout match {
       case Some(to) =>
-        interruptAfter(to) {
+        ti.interruptAfter(to) {
           super.checkAssumptions(assumptions)
         }
       case None =>
diff --git a/src/main/scala/leon/solvers/TimeoutSolver.scala b/src/main/scala/leon/solvers/TimeoutSolver.scala
index b4e4aece3..1198e0d72 100644
--- a/src/main/scala/leon/solvers/TimeoutSolver.scala
+++ b/src/main/scala/leon/solvers/TimeoutSolver.scala
@@ -7,48 +7,9 @@ import utils._
 
 trait TimeoutSolver extends Solver with Interruptible {
 
-  private class Countdown(timeout: Long, onTimeout: => Unit) extends Thread {
-    private var keepRunning = true
-    override def run : Unit = {
-      val startTime : Long = System.currentTimeMillis
-      var exceeded : Boolean = false
-
-      while(!exceeded && keepRunning) {
-        if(timeout < (System.currentTimeMillis - startTime)) {
-          exceeded = true
-        }
-        Thread.sleep(10)
-      }
-
-      if(exceeded && keepRunning) {
-        onTimeout
-      }
-    }
-
-    def finishedRunning : Unit = {
-      keepRunning = false
-    }
-  }
+  val ti = new TimeoutFor(this)
 
   protected var optTimeout: Option[Long] = None;
-  protected def interruptAfter[T](timeout: Long)(body: => T): T = {
-    var reachedTimeout = false
-
-    val timer = new Countdown(timeout, {
-      interrupt()
-      reachedTimeout = true
-    })
-
-    timer.start
-    val res = body
-    timer.finishedRunning
-
-    if (reachedTimeout) {
-      recoverInterrupt()
-    }
-
-    res
-  }
 
   def setTimeout(timeout: Long): this.type = {
     optTimeout = Some(timeout)
@@ -58,7 +19,7 @@ trait TimeoutSolver extends Solver with Interruptible {
   abstract override def check: Option[Boolean] = {
     optTimeout match {
       case Some(to) =>
-        interruptAfter(to) {
+        ti.interruptAfter(to) {
           super.check
         }
       case None =>
diff --git a/src/main/scala/leon/synthesis/utils/ExpressionGrammar.scala b/src/main/scala/leon/synthesis/utils/ExpressionGrammar.scala
index d555318ba..2c8be3cb2 100644
--- a/src/main/scala/leon/synthesis/utils/ExpressionGrammar.scala
+++ b/src/main/scala/leon/synthesis/utils/ExpressionGrammar.scala
@@ -189,7 +189,7 @@ object ExpressionGrammars {
     
     val normalGrammar = BoundedGrammar(EmbeddedGrammar(
         BaseGrammar ||
-        OneOf(terminals.toSeq) ||
+        OneOf(terminals.toSeq :+ e) ||
         FunctionCalls(sctx.program, sctx.functionContext, p.as.map(_.getType), excludeFCalls) ||
         SafeRecCalls(sctx.program, p.ws, p.pc),
       { (t: TypeTree)      => Label(t, "B", None)},
diff --git a/src/main/scala/leon/utils/InterruptManager.scala b/src/main/scala/leon/utils/InterruptManager.scala
index abe0d927b..e4aa86b3a 100644
--- a/src/main/scala/leon/utils/InterruptManager.scala
+++ b/src/main/scala/leon/utils/InterruptManager.scala
@@ -9,7 +9,7 @@ import java.util.concurrent.atomic.AtomicBoolean
 import sun.misc.{Signal, SignalHandler}
 import java.util.WeakHashMap
 
-class InterruptManager(reporter: Reporter) {
+class InterruptManager(reporter: Reporter) extends Interruptible {
   private[this] val interruptibles = new WeakHashMap[Interruptible, Boolean]()
   private[this] val sigINT = new Signal("INT")
   private[this] var oldHandler: SignalHandler = null
diff --git a/src/main/scala/leon/utils/TimeoutFor.scala b/src/main/scala/leon/utils/TimeoutFor.scala
new file mode 100644
index 000000000..a92f011a9
--- /dev/null
+++ b/src/main/scala/leon/utils/TimeoutFor.scala
@@ -0,0 +1,55 @@
+package leon
+package utils
+
+class TimeoutFor(it: Interruptible) {
+  private class Countdown(timeout: Long, onTimeout: => Unit) extends Thread {
+    private var keepRunning = true
+    override def run : Unit = {
+      val startTime : Long = System.currentTimeMillis
+      var exceeded : Boolean = false
+
+      while(!exceeded && keepRunning) {
+        if(timeout < (System.currentTimeMillis - startTime)) {
+          exceeded = true
+        }
+        Thread.sleep(10)
+      }
+
+      if(exceeded && keepRunning) {
+        onTimeout
+      }
+    }
+
+    def finishedRunning : Unit = {
+      keepRunning = false
+    }
+  }
+
+  def interruptAfter[T](timeout: Long)(body: => T): T = {
+    var reachedTimeout = false
+
+    val timer = new Countdown(timeout, {
+      it.interrupt()
+      reachedTimeout = true
+    })
+
+    timer.start
+    val res = body
+    timer.finishedRunning
+
+    if (reachedTimeout) {
+      it.recoverInterrupt()
+    }
+
+    res
+  }
+
+  def interruptAfter[T](timeout: Option[Long])(body: => T): T = {
+    timeout match {
+      case Some(to) =>
+        interruptAfter(to)(body)
+      case None =>
+        body
+    }
+  }
+}
-- 
GitLab