added support for collect

rssh · Apr 14, 2016 · 019c1b8 · 019c1b8
1 parent d32fa0f
commit 019c1b8
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git a/agent/src/main/scala/trackedfuture/agent/TFMethodAdapter.scala b/agent/src/main/scala/trackedfuture/agent/TFMethodAdapter.scala
@@ -40,7 +40,11 @@ class TFMethodAdapter(up: MethodVisitor) extends MethodVisitor(Opcodes.ASM5, up)
         (Opcodes.INVOKEINTERFACE,
           "scala/concurrent/Future", "withFilter", 
           "(Lscala/Function1;Lscala/concurrent/ExecutionContext;)Lscala/concurrent/Future;"
-        ) -> "rFilter"
+        ) -> "rFilter",
+        (Opcodes.INVOKEINTERFACE,
+          "scala/concurrent/Future", "collect", 
+          "(Lscala/PartialFunction;Lscala/concurrent/ExecutionContext;)Lscala/concurrent/Future;"
+        ) -> "collect"
   )
 
   def appendParam(ownerType: String, desc: String): String =

diff --git a/agent/src/main/scala/trackedfuture/runtime/TrackedFuture.scala b/agent/src/main/scala/trackedfuture/runtime/TrackedFuture.scala
@@ -54,7 +54,6 @@ object TrackedFuture
 
   def rFilter[A](future: Future[A], function: A => Boolean, executor: ExecutionContext):Future[A]=
   {
-    System.err.println("rFilter")
     val trace = Thread.currentThread.getStackTrace()
     val prevTrace = new StackTraces(trace, ThreadTrace.prevTraces.value)
     future.map { a => trackedCall(
@@ -65,6 +64,16 @@ object TrackedFuture
                }(executor)
   }
 
+  def collect[A,B](future: Future[A], pf: PartialFunction[A,B], executor: ExecutionContext):Future[B]=
+  {
+    val trace = Thread.currentThread.getStackTrace()
+    val prevTrace = new StackTraces(trace, ThreadTrace.prevTraces.value)
+    future.map { a => trackedCall({
+          pf.applyOrElse(a, (t: A) => throw new NoSuchElementException("Future.collect partial function is not defined at: " + t)) 
+    }, prevTrace) }(executor)
+  }
+
+
   private def trackedCall[A](body: =>A, prevTrace:StackTraces):A =
   {
     ThreadTrace.setPrev(prevTrace)

diff --git a/example/src/main/scala/tracedfuture/example/Main.scala b/example/src/main/scala/tracedfuture/example/Main.scala
@@ -60,5 +60,11 @@ object Main
     Future{ () } withFilter { _ => throw new Exception("AAA") }
   }
 
+  def fCollect0(f:PartialFunction[String,String]):Future[String] =
+    fCollect1(f)
+
+  def fCollect1(f:PartialFunction[String,String]):Future[String] =
+    Future{ "aaa" } collect f
+
 
 }
diff --git a/example/src/test/scala/tracedfuture/example/MainCallSpec.scala b/example/src/test/scala/tracedfuture/example/MainCallSpec.scala
@@ -33,6 +33,11 @@ class MainCallSpec extends FlatSpec with AsyncAssertions
     callAndCheckMethod( Main.withFilter0(), "withFilter0")
   }
 
+  "MainCall" should "show origin method with collect " in {
+    callAndCheckMethod( Main.fCollect0{case "bbb" => "ccc"}, "fCollect0")
+  }
+
+
   private def callAndCheckMethod(body: =>Future[_],method:String): Unit = {
     val f = body
     val w = new Waiter