com.tinkerpop.gremlin.structure.strategy.StrategyWrappedTraversalStrategy Maven / Gradle / Ivy

Go to download
package com.tinkerpop.gremlin.structure.strategy;

import com.tinkerpop.gremlin.process.Traversal;
import com.tinkerpop.gremlin.process.TraversalStrategy;
import com.tinkerpop.gremlin.process.Traverser;
import com.tinkerpop.gremlin.process.graph.step.map.EdgeVertexStep;
import com.tinkerpop.gremlin.process.graph.step.map.GraphStep;
import com.tinkerpop.gremlin.process.graph.step.map.MapStep;
import com.tinkerpop.gremlin.process.graph.step.map.VertexStep;
import com.tinkerpop.gremlin.process.util.TraversalHelper;
import com.tinkerpop.gremlin.structure.Edge;
import com.tinkerpop.gremlin.structure.Graph;
import com.tinkerpop.gremlin.structure.Property;
import com.tinkerpop.gremlin.structure.Vertex;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * Force all {@link Vertex}, {@link Edge}, {@link Property}, and {@link Graph} objects in the {@link Traversal} to be wrapped
 * appropriately.
 *
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public class StrategyWrappedTraversalStrategy implements TraversalStrategy.NoDependencies {

    private final StrategyWrappedGraph graph;

    public StrategyWrappedTraversalStrategy(final StrategyWrappedGraph graph) {
        this.graph = graph;
    }

    public void apply(final Traversal traversal) {
        // MapStep after each GraphStep, VertexStep or EdgeVertexStep
        final List stepsToLookFor = Arrays.asList(GraphStep.class, VertexStep.class, EdgeVertexStep.class);
        final List 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((Traverser