PartitionGraphStrategyOptimizer

gistfile1.java

public void optimize(final Traversal traversal) {
    // inject a HasStep after each GraphStep, VertexStep or EdgeVertexStep
    final List<Class> stepsToLookFor = Arrays.<Class>asList(GraphStep.class, VertexStep.class, EdgeVertexStep.class);
    final List<Integer> positions = new ArrayList<>();
    final List<?> traversalSteps = traversal.getSteps();
    for (int ix = 0; ix < traversalSteps.size(); ix++) {
        final int pos = ix;
        if (stepsToLookFor.stream().anyMatch(c -> c.isAssignableFrom(traversalSteps.get(pos).getClass()))) positions.add(ix);
    }

    Collections.reverse(positions);
    for (int pos : positions) {
        final MapStep transformToStrategy = new MapStep(traversal);
        transformToStrategy.setFunction((t) -> {
            if (t instanceof Vertex)
                return new StrategyWrappedVertex((Vertex) t, graph);
            else if (t instanceof Edge)
                return new StrategyWrappedEdge((Edge) t, graph);
            else
                return t;
        });

        //TraversalHelper.insertStep(transformToStrategy, pos + 1, traversal);
        TraversalHelper.insertStep(new HasStep(traversal, new HasContainer(this.partitionKey, T.convert(T.in), readPartitions)), pos + 1, traversal);

        System.out.println(traversal);
    }
}