Fixed related bugs with fold and rewriting

translation/src/main/scala/org/opencypher/gremlin/translation/ir/rewrite/SimplifySingleProjections.scala

@@ -61,9 +61,7 @@ object SimplifySingleProjections extends GremlinRewriter {
       case SelectK(key1) :: rest if key == key1 && rest.isEmpty =>
         Identity :: Nil
       case _ =>
-        steps.map({ step =>
-          step.mapTraversals(subSteps => mapTraversals(removeSelect(key, _))(subSteps))
-        })
+        removeSelectFromSubtraversals(key, steps)
     }
   }
 
@@ -77,6 +75,15 @@ object SimplifySingleProjections extends GremlinRewriter {
     } ::: rest
   }
 
+  private def removeSelectFromSubtraversals(key: String, steps: Seq[GremlinStep]): Seq[GremlinStep] = {
+    val (replaceable, rest) =
+      steps.span(s => !s.isInstanceOf[Constant] || s.isInstanceOf[SelectC] || s.isInstanceOf[SelectK])
+
+    replaceable.map({ step =>
+      step.mapTraversals(subSteps => mapTraversals(removeSelect(key, _))(subSteps))
+    }) ++ rest
+  }
+
   private def eqUnused(stepLabel: String, key2: String): Boolean = {
     stepLabel == UNUSED && key2 == UNUSED
   }

translation/src/main/scala/org/opencypher/gremlin/translation/walker/ExpressionWalker.scala

@@ -221,9 +221,9 @@ private class ExpressionWalker[T, P](context: WalkerContext[T, P], g: GremlinSte
           case "exists"           => traversals.head.flatMap(anyMatch(__.is(p.neq(NULL))))
           case "head"             => traversals.head.flatMap(emptyToNull(__.limit(Scope.local, 1), context))
           case "id"               => traversals.head.flatMap(notNull(__.id(), context))
-          case "keys" if onEntity => traversals.head.properties().key().fold()
+          case "keys" if onEntity => traversals.head.map(__.properties().key().fold())
           case "keys"             => traversals.head.select(Column.keys)
-          case "labels"           => traversals.head.label().is(p.neq(Vertex.DEFAULT_LABEL)).fold()
+          case "labels"           => traversals.head.map(__.label().is(p.neq(Vertex.DEFAULT_LABEL)).fold())
           case "length"           => traversals.head.count(Scope.local).math("(_-1)/2")
           case "last"             => traversals.head.flatMap(emptyToNull(__.tail(Scope.local, 1), context))
           case "nodes"            => traversals.head.flatMap(filterElements(args, includeNodes = true))
@@ -258,7 +258,7 @@ private class ExpressionWalker[T, P](context: WalkerContext[T, P], g: GremlinSte
         val functionT = walkLocal(function, maybeAlias)
 
         val Variable(dependencyName) = function.dependencies.head
-        targetT.unfold().as(dependencyName).flatMap(functionT).fold()
+        targetT.map(__.unfold().as(dependencyName).flatMap(functionT).fold())
 
       case PatternComprehension(_, RelationshipsPattern(relationshipChain), maybePredicate, PathExpression(_), _) =>
         val select = __
@@ -287,10 +287,10 @@ private class ExpressionWalker[T, P](context: WalkerContext[T, P], g: GremlinSte
 
         val functionT = walkLocal(projection, maybeAlias)
         if (projection.dependencies.isEmpty) {
-          traversal.flatMap(functionT).fold()
+          __.map(traversal.flatMap(functionT).fold())
         } else if (projection.dependencies.size == 1) {
           val Variable(dependencyName) = projection.dependencies.head
-          traversal.select(dependencyName).flatMap(functionT).fold()
+          __.map(traversal.select(dependencyName).flatMap(functionT).fold())
         } else {
           context.unsupported("pattern comprehension with multiple arguments", expression)
         }

translation/src/test/scala/org/opencypher/gremlin/translation/ir/rewrite/RemoveUselessStepsTest.scala

@@ -25,17 +25,14 @@ class RemoveUselessStepsTest {
 
   val flavor = new TranslatorFlavor(
     rewriters = Seq(
-      InlineFlatMapTraversal
+      InlineFlatMapTraversal,
+      SimplifySingleProjections
     ),
     postConditions = Nil
   )
 
   @Test def foldUnfold(): Unit = {
-    assertThat(parse("""
-        |MATCH (n)
-        |UNWIND labels(n) as l
-        |RETURN l
-      """.stripMargin))
+    assertThat(parse("MATCH (n) RETURN count(*)"))
       .withFlavor(flavor)
       .rewritingWith(RemoveUselessSteps)
       .removes(__.fold().unfold())