gistfile1.java

package com.tinkerpop.gremlin.computer;

import com.tinkerpop.blueprints.Direction;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.computer.GraphMemory;
import com.tinkerpop.blueprints.computer.VertexProgram;
import com.tinkerpop.gremlin.pipes.FilterPipe;
import com.tinkerpop.gremlin.pipes.FlatMapPipe;
import com.tinkerpop.gremlin.pipes.Gremlin;
import com.tinkerpop.gremlin.pipes.MapPipe;
import com.tinkerpop.gremlin.pipes.Pipe;
import com.tinkerpop.gremlin.pipes.util.Holder;
import com.tinkerpop.gremlin.pipes.util.PipeHelper;
import com.tinkerpop.gremlin.pipes.util.SingleIterator;

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

/**
 * @author Marko A. Rodriguez (http://markorodriguez.com)
 */
public class GremlinVertexProgram implements VertexProgram {

    private static final String GREMLINS = "gremlins";
    private static final String RESULT = "result";
    private final Gremlin gremlin;

    public GremlinVertexProgram(Gremlin gremlin) {
        this.gremlin = gremlin;
    }

    public void setup(final GraphMemory graphMemory) {
        graphMemory.setIfAbsent("gremlin", gremlin);
    }

    public void execute(final Vertex vertex, final GraphMemory graphMemory) {
        Pipe pipe = getNextStep(graphMemory);

        if (pipe instanceof FilterPipe) {
            ((FilterPipe) pipe).addStarts(new SingleIterator<Holder>(new Holder(Pipe.NONE, vertex)));
            if (!PipeHelper.hasNextIteration(pipe)) {
                vertex.setProperty(GREMLINS, Collections.emptyList());
            }
        } else if(pipe instanceof FlatMapPipe) {
            ((FlatMapPipe) pipe).addStarts(new SingleIterator<Holder>(new Holder(Pipe.NONE, vertex)));
            
        } else if(pipe instanceof MapPipe) {
            
        } else {
            throw new IllegalStateException("There are no other pipe types -- how did you get here?");
        }
    }

    public boolean terminate(final GraphMemory graphMemory) {
        Gremlin gremlin = graphMemory.get("gremlin");
        return !(graphMemory.getIteration() < gremlin.getPipes().size());
    }

    public Map<String, KeyType> getComputeKeys() {
        return VertexProgram.ofComputeKeys(GREMLINS, KeyType.VARIABLE);
    }

    private Pipe getNextStep(final GraphMemory graphMemory) {
        final Gremlin gremlin = graphMemory.get("gremlin");
        return (Pipe) gremlin.getPipes().get(graphMemory.getIteration());
    }

    public static Builder create() {
        return new Builder();
    }

    public static class Builder {
        private Gremlin gremlin;

        public Builder gremlin(final Gremlin gremlin) {
            this.gremlin = gremlin;
            return this;
        }

        public GremlinVertexProgram build() {
            return new GremlinVertexProgram(this.gremlin);
        }
    }
}