fix counter-intuitive union choose coalesce and repeat path tracking (#362)

* fix counter-intuitive path tracking of union, choose, coalesce and repeat.

Author: bbrehm

traversal-tests/src/test/scala/overflowdb/traversal/PathTraversalTests.scala

@@ -159,7 +159,7 @@ class PathTraversalTests extends AnyWordSpec with ExampleGraphSetup {
             case "L1" => _.out // -> L2
             case "R1" => _.repeat(_.out)(_.maxDepth(3)) // -> R4
           }.property(Name).path.toSetMutable shouldBe Set(
-          Seq(r1, r4, "R4"),
+          Seq(r1, r2, r3, r4, "R4"),
           Seq(l1, l2, "L2")
         )
       }
@@ -177,6 +177,18 @@ class PathTraversalTests extends AnyWordSpec with ExampleGraphSetup {
         traversalInvoked shouldBe false
       }
 
+      "union" in {
+         centerTrav.enablePathTracking.union(_.out.out).path.toSetMutable shouldBe Set(
+          Seq(center, l1, l2),
+          Seq(center, r1, r2)
+         )
+        //we can hide internal steps from path-tracking
+        centerTrav.enablePathTracking.union(t => Traversal.from(t.out.out).iterator).path.toSetMutable shouldBe Set(
+          Seq(center, l2),
+          Seq(center, r2)
+        )
+      }
+
     }
   }
 

traversal-tests/src/test/scala/overflowdb/traversal/RepeatTraversalTests.scala

@@ -381,6 +381,13 @@ class RepeatTraversalTests extends AnyWordSpec with ExampleGraphSetup {
       r1.start.enablePathTracking.repeat(_.out.out)(_.maxDepth(2)).l shouldBe Seq(r5)
       r1.start.enablePathTracking.repeat(_.out.out)(_.maxDepth(2)).path.head shouldBe List(r1, r2, r3, r4, r5)
     }
+    "should not forget steps preceding the repeat" in {
+      centerTrav.enablePathTracking.followedBy.repeat(_.followedBy.followedBy)(_.maxDepth(1)).path.toSetMutable shouldBe Set(
+        Seq(center, l1, l2, l3),
+        Seq(center, r1, r2, r3)
+      )
+    }
+
   }
 
 }

traversal/src/main/scala/overflowdb/traversal/PathAwareTraversal.scala

@@ -46,6 +46,42 @@ class PathAwareTraversal[+A](val elementsWithPath: IterableOnce[(A, Vector[Any])
   override def dedupBy(fun: A => Any): Traversal[A] =
     new PathAwareTraversal(elementsWithPathIterator.distinctBy(x => fun(x._1)))
 
+  override def union[B](traversals: (Traversal[A] => Traversal[B])*): Traversal[B] =
+    new PathAwareTraversal(elementsWithPathIterator.flatMap { case (a, p) =>
+      traversals.iterator.flatMap { inner =>
+        inner(new PathAwareTraversal(Iterator.single((a, p)))) match {
+          case stillPathAware: PathAwareTraversal[B] => stillPathAware.elementsWithPathIterator
+          //do we really want to allow the following, or is it an error?
+          case notPathAware => notPathAware.iterator.map { (b: B) => (b, p.appended(a)) }
+        }
+      }
+    })
+
+
+  override def choose[BranchOn >: Null, NewEnd]
+  (on: Traversal[A] => Traversal[BranchOn])
+  (options: PartialFunction[BranchOn, Traversal[A] => Traversal[NewEnd]]): Traversal[NewEnd] =
+    new PathAwareTraversal(elementsWithPathIterator.flatMap { case (a, p) =>
+      val branchOnValue: BranchOn = on(Traversal.fromSingle(a)).headOption.getOrElse(null)
+      options.applyOrElse(branchOnValue, (failState: BranchOn) => ((unused: Traversal[A]) => Traversal.empty[NewEnd])).apply(new PathAwareTraversal(Iterator.single((a, p)))) match {
+        case stillPathAware: PathAwareTraversal[NewEnd] => stillPathAware.elementsWithPathIterator
+        //do we really want to allow the following, or is it an error?
+        case notPathAware => notPathAware.iterator.map { (b: NewEnd) => (b, p.appended(a)) }
+      }
+    })
+
+  override def coalesce[NewEnd](options: (Traversal[A] => Traversal[NewEnd])*): Traversal[NewEnd] =
+    new PathAwareTraversal(elementsWithPathIterator.flatMap { case (a, p) =>
+      options.iterator.map { inner =>
+        inner(new PathAwareTraversal(Iterator.single((a, p)))) match {
+          case stillPathAware: PathAwareTraversal[NewEnd] => stillPathAware.elementsWithPathIterator
+          //do we really want to allow the following, or is it an error?
+          case notPathAware => notPathAware.iterator.map { (b: NewEnd) => (b, p.appended(a)) }
+        }
+      }.find(_.nonEmpty).getOrElse(Iterator.empty)
+    })
+
+
   // TODO add type safety once we're on dotty, similar to gremlin-scala's as/label steps with typelevel append?
   override def path: Traversal[Vector[Any]] =
     new Traversal(elementsWithPathIterator.map {
@@ -82,8 +118,8 @@ class PathAwareTraversal[+A](val elementsWithPath: IterableOnce[(A, Vector[Any])
     val behaviour = behaviourBuilder(new RepeatBehaviour.Builder[B]).build
     val _repeatTraversal = repeatTraversal.asInstanceOf[Traversal[B] => Traversal[B]] //this cast usually :tm: safe, because `B` is a supertype of `A`
     val repeat0: B => PathAwareTraversal[B] = PathAwareRepeatStep(_repeatTraversal, behaviour)
-    new PathAwareTraversal(iterator.flatMap { a =>
-      repeat0(a).elementsWithPath
+    new PathAwareTraversal(elementsWithPathIterator.flatMap { case (a,p) =>
+      repeat0(a).elementsWithPathIterator.map{case (b, pp) => (b, p ++ pp)}
     })
   }
 

traversal/src/main/scala/overflowdb/traversal/Traversal.scala

@@ -231,15 +231,11 @@ class Traversal[+A](elements: IterableOnce[A])
    */
   @Doc(info = "allows to implement conditional semantics: if, if/else, if/elseif, if/elseif/else, ...")
   def choose[BranchOn >: Null, NewEnd]
-    (on: Traversal[A] => Traversal[BranchOn])
-    (options: PartialFunction[BranchOn, Traversal[A] => Traversal[NewEnd]]): Traversal[NewEnd] =
+  (on: Traversal[A] => Traversal[BranchOn])
+  (options: PartialFunction[BranchOn, Traversal[A] => Traversal[NewEnd]]): Traversal[NewEnd] =
     flatMap { (a: A) =>
       val branchOnValue: BranchOn = on(Traversal.fromSingle(a)).headOption.getOrElse(null)
-      if (options.isDefinedAt(branchOnValue)) {
-        options(branchOnValue)(Traversal.fromSingle(a))
-      } else {
-        Traversal.empty
-      }
+      options.applyOrElse(branchOnValue, (fail: BranchOn) => ((unused: Traversal[A]) => Traversal.empty[NewEnd])).apply(Traversal.fromSingle(a))
     }
 
   /** Branch step: evaluates the provided traversals in order and returns the first traversal that emits at least one element.