use super::{Step, StepResult, Traverser, TraverserRequirement, TraverserValue};
use std::any::Any;
pub trait FlatMapStep: Step {
fn flat_map(&self, traverser: &Traverser) -> Vec<Traverser>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Direction {
Out,
In,
Both,
}
#[derive(Debug, Clone)]
pub struct VertexStep {
id: String,
labels: Vec<String>,
direction: Direction,
edge_labels: Vec<String>,
return_edges: bool,
}
impl VertexStep {
pub fn out(edge_labels: Vec<String>) -> Self {
Self {
id: format!("out_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Out,
edge_labels,
return_edges: false,
}
}
pub fn in_(edge_labels: Vec<String>) -> Self {
Self {
id: format!("in_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::In,
edge_labels,
return_edges: false,
}
}
pub fn both(edge_labels: Vec<String>) -> Self {
Self {
id: format!("both_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Both,
edge_labels,
return_edges: false,
}
}
pub fn out_e(edge_labels: Vec<String>) -> Self {
Self {
id: format!("outE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Out,
edge_labels,
return_edges: true,
}
}
pub fn in_e(edge_labels: Vec<String>) -> Self {
Self {
id: format!("inE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::In,
edge_labels,
return_edges: true,
}
}
pub fn both_e(edge_labels: Vec<String>) -> Self {
Self {
id: format!("bothE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Both,
edge_labels,
return_edges: true,
}
}
pub fn direction(&self) -> Direction {
self.direction
}
pub fn edge_labels(&self) -> &[String] {
&self.edge_labels
}
pub fn returns_edges(&self) -> bool {
self.return_edges
}
pub fn with_id(mut self, id: String) -> Self {
self.id = id;
self
}
}
impl Step for VertexStep {
fn id(&self) -> &str {
&self.id
}
fn name(&self) -> &str {
match (self.direction, self.return_edges) {
(Direction::Out, false) => "OutStep",
(Direction::In, false) => "InStep",
(Direction::Both, false) => "BothStep",
(Direction::Out, true) => "OutEStep",
(Direction::In, true) => "InEStep",
(Direction::Both, true) => "BothEStep",
}
}
fn labels(&self) -> &[String] {
&self.labels
}
fn add_label(&mut self, label: String) {
if !self.labels.contains(&label) {
self.labels.push(label);
}
}
fn requirements(&self) -> &[TraverserRequirement] {
&[]
}
fn process_traverser(&self, traverser: Traverser) -> StepResult {
let new_traversers = self.flat_map(&traverser);
StepResult::emit_many(new_traversers)
}
fn reset(&mut self) {}
fn clone_step(&self) -> Box<dyn Step> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
impl FlatMapStep for VertexStep {
fn flat_map(&self, traverser: &Traverser) -> Vec<Traverser> {
Vec::new()
}
}
pub type OutStep = VertexStep;
pub type InStep = VertexStep;
pub type BothStep = VertexStep;
#[derive(Debug, Clone)]
pub struct EdgeStep {
id: String,
labels: Vec<String>,
direction: Direction,
edge_labels: Vec<String>,
}
impl EdgeStep {
pub fn out(edge_labels: Vec<String>) -> Self {
Self {
id: format!("outE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Out,
edge_labels,
}
}
pub fn in_(edge_labels: Vec<String>) -> Self {
Self {
id: format!("inE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::In,
edge_labels,
}
}
pub fn both(edge_labels: Vec<String>) -> Self {
Self {
id: format!("bothE_{}", edge_labels.join("_")),
labels: Vec::new(),
direction: Direction::Both,
edge_labels,
}
}
}
impl Step for EdgeStep {
fn id(&self) -> &str {
&self.id
}
fn name(&self) -> &str {
match self.direction {
Direction::Out => "OutEStep",
Direction::In => "InEStep",
Direction::Both => "BothEStep",
}
}
fn labels(&self) -> &[String] {
&self.labels
}
fn add_label(&mut self, label: String) {
if !self.labels.contains(&label) {
self.labels.push(label);
}
}
fn requirements(&self) -> &[TraverserRequirement] {
&[]
}
fn process_traverser(&self, traverser: Traverser) -> StepResult {
let new_traversers = self.flat_map(&traverser);
StepResult::emit_many(new_traversers)
}
fn reset(&mut self) {}
fn clone_step(&self) -> Box<dyn Step> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
impl FlatMapStep for EdgeStep {
fn flat_map(&self, _traverser: &Traverser) -> Vec<Traverser> {
Vec::new()
}
}
#[derive(Debug, Clone)]
pub struct EdgeVertexStep {
id: String,
labels: Vec<String>,
direction: Direction,
}
impl EdgeVertexStep {
pub fn out_v() -> Self {
Self {
id: "outV_0".to_string(),
labels: Vec::new(),
direction: Direction::Out,
}
}
pub fn in_v() -> Self {
Self {
id: "inV_0".to_string(),
labels: Vec::new(),
direction: Direction::In,
}
}
pub fn both_v() -> Self {
Self {
id: "bothV_0".to_string(),
labels: Vec::new(),
direction: Direction::Both,
}
}
pub fn other_v() -> Self {
Self {
id: "otherV_0".to_string(),
labels: Vec::new(),
direction: Direction::Both, }
}
}
impl Step for EdgeVertexStep {
fn id(&self) -> &str {
&self.id
}
fn name(&self) -> &str {
match self.direction {
Direction::Out => "OutVStep",
Direction::In => "InVStep",
Direction::Both => "BothVStep",
}
}
fn labels(&self) -> &[String] {
&self.labels
}
fn add_label(&mut self, label: String) {
if !self.labels.contains(&label) {
self.labels.push(label);
}
}
fn requirements(&self) -> &[TraverserRequirement] {
&[]
}
fn process_traverser(&self, traverser: Traverser) -> StepResult {
if let TraverserValue::Edge {
id: _,
source,
target,
label: _,
} = traverser.value()
{
let new_traversers = match self.direction {
Direction::Out => {
vec![traverser.clone_with_value(TraverserValue::Vertex(target.clone()))]
}
Direction::In => {
vec![traverser.clone_with_value(TraverserValue::Vertex(source.clone()))]
}
Direction::Both => vec![
traverser.clone_with_value(TraverserValue::Vertex(source.clone())),
traverser.clone_with_value(TraverserValue::Vertex(target.clone())),
],
};
StepResult::emit_many(new_traversers)
} else {
StepResult::Filter
}
}
fn reset(&mut self) {}
fn clone_step(&self) -> Box<dyn Step> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
impl FlatMapStep for EdgeVertexStep {
fn flat_map(&self, traverser: &Traverser) -> Vec<Traverser> {
match self.process_traverser(traverser.clone()) {
StepResult::Emit(t) => t,
_ => Vec::new(),
}
}
}
#[derive(Debug, Clone)]
pub struct PropertiesStep {
id: String,
labels: Vec<String>,
keys: Vec<String>,
}
impl PropertiesStep {
pub fn new() -> Self {
Self {
id: "properties_0".to_string(),
labels: Vec::new(),
keys: Vec::new(),
}
}
pub fn with_keys(keys: Vec<String>) -> Self {
Self {
id: format!("properties_{}", keys.join("_")),
labels: Vec::new(),
keys,
}
}
}
impl Default for PropertiesStep {
fn default() -> Self {
Self::new()
}
}
impl Step for PropertiesStep {
fn id(&self) -> &str {
&self.id
}
fn name(&self) -> &str {
"PropertiesStep"
}
fn labels(&self) -> &[String] {
&self.labels
}
fn add_label(&mut self, label: String) {
if !self.labels.contains(&label) {
self.labels.push(label);
}
}
fn requirements(&self) -> &[TraverserRequirement] {
&[]
}
fn process_traverser(&self, traverser: Traverser) -> StepResult {
let new_traversers = self.flat_map(&traverser);
StepResult::emit_many(new_traversers)
}
fn reset(&mut self) {}
fn clone_step(&self) -> Box<dyn Step> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
}
impl FlatMapStep for PropertiesStep {
fn flat_map(&self, traverser: &Traverser) -> Vec<Traverser> {
if let TraverserValue::Map(map) = traverser.value() {
let mut result = Vec::new();
for (key, value) in map {
if self.keys.is_empty() || self.keys.contains(key) {
result.push(
traverser
.clone_with_value(TraverserValue::Property(key.clone(), value.clone())),
);
}
}
result
} else {
Vec::new()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_out_step() {
let step = VertexStep::out(vec!["knows".to_string()]);
assert_eq!(step.direction(), Direction::Out);
assert_eq!(step.edge_labels(), &["knows"]);
assert!(!step.returns_edges());
}
#[test]
fn test_in_step() {
let step = VertexStep::in_(vec![]);
assert_eq!(step.direction(), Direction::In);
assert!(step.edge_labels().is_empty());
}
#[test]
fn test_both_step() {
let step = VertexStep::both(vec!["connects".to_string(), "knows".to_string()]);
assert_eq!(step.direction(), Direction::Both);
assert_eq!(step.edge_labels().len(), 2);
}
#[test]
fn test_out_e_step() {
let step = VertexStep::out_e(vec![]);
assert_eq!(step.direction(), Direction::Out);
assert!(step.returns_edges());
}
#[test]
fn test_edge_vertex_step() {
let edge = TraverserValue::Edge {
id: "e1".to_string(),
source: "v1".to_string(),
target: "v2".to_string(),
label: "knows".to_string(),
};
let traverser = Traverser::with_value(edge);
let out_v = EdgeVertexStep::out_v();
let result = out_v.process_traverser(traverser.clone());
if let StepResult::Emit(t) = result {
assert_eq!(t.len(), 1);
assert!(matches!(t[0].value(), TraverserValue::Vertex(id) if id == "v2"));
}
let in_v = EdgeVertexStep::in_v();
let result = in_v.process_traverser(traverser.clone());
if let StepResult::Emit(t) = result {
assert_eq!(t.len(), 1);
assert!(matches!(t[0].value(), TraverserValue::Vertex(id) if id == "v1"));
}
let both_v = EdgeVertexStep::both_v();
let result = both_v.process_traverser(traverser);
if let StepResult::Emit(t) = result {
assert_eq!(t.len(), 2);
}
}
#[test]
fn test_properties_step() {
let step = PropertiesStep::with_keys(vec!["name".to_string(), "age".to_string()]);
assert_eq!(step.name(), "PropertiesStep");
}
}