Last Updated: February 25, 2016
·
595
· dkuppitz

Sorted Trees in Gremlin

// closure can't have the same name as the defined step
tree = { vertices ->

  def results = []

  vertices.each() {
    results << it
    unordered = label == null ? it."$direction"() : it."$direction"(label);
    ordered = order == null ? unordered : unordered.order(order)
    children = ordered.toList()
    if (children) {
      child_tree = tree(children)
      results << child_tree
    }
  }
  results
}

extractParams = {final Object... params ->
  label = null;
  order = null;
  switch (params.size()) {
    case 1:
      if (params[0] instanceof String) {
        label = params[0];
      }
      else {
        order = params[0];
      }
      break;
    case 2:
      label = params[0];
      order = params[1];
      break;
  }
} 

inClosure = {final Object... params -> 
  extractParams(params)
  results = []
  direction = "in"
  _().transform{ tree(it) }
}    

outClosure = {final Object... params -> 
  extractParams(params)
  results = []
  direction = "out"
  _().transform{ tree(it) }
}    

Gremlin.defineStep("inTree", [Vertex,Pipe], inClosure) 
Gremlin.defineStep("outTree", [Vertex,Pipe], outClosure)