use std::collections::BTreeSet;
use std::str::FromStr;
use plexus_serde::{
deserialize_engine_capability_decl, serialize_engine_capability_decl,
CapabilitySemver as WireSemver, CapabilityVersionRange as WireVersionRange,
EngineCapabilityDecl as WireCapabilityDecl, OpOrderingDecl as WireOpOrderingDecl, Version,
};
use serde::{Deserialize, Serialize};
use crate::capabilities::ordering::{op_ordering_contract, OpOrderingContract};
use crate::capabilities::wire::{
from_wire_ordering_contract, to_wire_ordering_contract, EngineCapabilityDocument,
OpOrderingDocument,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct PlanSemver {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl PlanSemver {
pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
Self {
major,
minor,
patch,
}
}
}
impl From<&Version> for PlanSemver {
fn from(v: &Version) -> Self {
Self::new(v.major, v.minor, v.patch)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct VersionRange {
pub min_supported: PlanSemver,
pub max_supported: PlanSemver,
}
impl VersionRange {
pub const fn new(min_supported: PlanSemver, max_supported: PlanSemver) -> Self {
Self {
min_supported,
max_supported,
}
}
pub fn supports(&self, version: PlanSemver) -> bool {
self.min_supported <= version && version <= self.max_supported
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum OpKind {
ScanNodes,
ScanRels,
Expand,
OptionalExpand,
SemiExpand,
ExpandVarLen,
Filter,
BlockMarker,
Project,
Aggregate,
Sort,
Limit,
Unwind,
PathConstruct,
Union,
CreateNode,
CreateRel,
Merge,
Delete,
SetProperty,
RemoveProperty,
VectorScan,
Rerank,
Return,
ConstRow,
}
impl OpKind {
pub fn as_str(self) -> &'static str {
match self {
Self::ScanNodes => "ScanNodes",
Self::ScanRels => "ScanRels",
Self::Expand => "Expand",
Self::OptionalExpand => "OptionalExpand",
Self::SemiExpand => "SemiExpand",
Self::ExpandVarLen => "ExpandVarLen",
Self::Filter => "Filter",
Self::BlockMarker => "BlockMarker",
Self::Project => "Project",
Self::Aggregate => "Aggregate",
Self::Sort => "Sort",
Self::Limit => "Limit",
Self::Unwind => "Unwind",
Self::PathConstruct => "PathConstruct",
Self::Union => "Union",
Self::CreateNode => "CreateNode",
Self::CreateRel => "CreateRel",
Self::Merge => "Merge",
Self::Delete => "Delete",
Self::SetProperty => "SetProperty",
Self::RemoveProperty => "RemoveProperty",
Self::VectorScan => "VectorScan",
Self::Rerank => "Rerank",
Self::Return => "Return",
Self::ConstRow => "ConstRow",
}
}
}
impl FromStr for OpKind {
type Err = ();
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name {
"ScanNodes" => Ok(Self::ScanNodes),
"ScanRels" => Ok(Self::ScanRels),
"Expand" => Ok(Self::Expand),
"OptionalExpand" => Ok(Self::OptionalExpand),
"SemiExpand" => Ok(Self::SemiExpand),
"ExpandVarLen" => Ok(Self::ExpandVarLen),
"Filter" => Ok(Self::Filter),
"BlockMarker" => Ok(Self::BlockMarker),
"Project" => Ok(Self::Project),
"Aggregate" => Ok(Self::Aggregate),
"Sort" => Ok(Self::Sort),
"Limit" => Ok(Self::Limit),
"Unwind" => Ok(Self::Unwind),
"PathConstruct" => Ok(Self::PathConstruct),
"Union" => Ok(Self::Union),
"CreateNode" => Ok(Self::CreateNode),
"CreateRel" => Ok(Self::CreateRel),
"Merge" => Ok(Self::Merge),
"Delete" => Ok(Self::Delete),
"SetProperty" => Ok(Self::SetProperty),
"RemoveProperty" => Ok(Self::RemoveProperty),
"VectorScan" => Ok(Self::VectorScan),
"Rerank" => Ok(Self::Rerank),
"Return" => Ok(Self::Return),
"ConstRow" => Ok(Self::ConstRow),
_ => Err(()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum ExprKind {
ColRef,
PropAccess,
IntLiteral,
FloatLiteral,
BoolLiteral,
StringLiteral,
NullLiteral,
Cmp,
And,
Or,
Not,
IsNull,
IsNotNull,
StartsWith,
EndsWith,
Contains,
In,
ListLiteral,
MapLiteral,
Exists,
ListComprehension,
Agg,
Arith,
Param,
Case,
VectorSimilarity,
}
impl ExprKind {
pub fn as_str(self) -> &'static str {
match self {
Self::ColRef => "ColRef",
Self::PropAccess => "PropAccess",
Self::IntLiteral => "IntLiteral",
Self::FloatLiteral => "FloatLiteral",
Self::BoolLiteral => "BoolLiteral",
Self::StringLiteral => "StringLiteral",
Self::NullLiteral => "NullLiteral",
Self::Cmp => "Cmp",
Self::And => "And",
Self::Or => "Or",
Self::Not => "Not",
Self::IsNull => "IsNull",
Self::IsNotNull => "IsNotNull",
Self::StartsWith => "StartsWith",
Self::EndsWith => "EndsWith",
Self::Contains => "Contains",
Self::In => "In",
Self::ListLiteral => "ListLiteral",
Self::MapLiteral => "MapLiteral",
Self::Exists => "Exists",
Self::ListComprehension => "ListComprehension",
Self::Agg => "Agg",
Self::Arith => "Arith",
Self::Param => "Param",
Self::Case => "Case",
Self::VectorSimilarity => "VectorSimilarity",
}
}
}
impl FromStr for ExprKind {
type Err = ();
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name {
"ColRef" => Ok(Self::ColRef),
"PropAccess" => Ok(Self::PropAccess),
"IntLiteral" => Ok(Self::IntLiteral),
"FloatLiteral" => Ok(Self::FloatLiteral),
"BoolLiteral" => Ok(Self::BoolLiteral),
"StringLiteral" => Ok(Self::StringLiteral),
"NullLiteral" => Ok(Self::NullLiteral),
"Cmp" => Ok(Self::Cmp),
"And" => Ok(Self::And),
"Or" => Ok(Self::Or),
"Not" => Ok(Self::Not),
"IsNull" => Ok(Self::IsNull),
"IsNotNull" => Ok(Self::IsNotNull),
"StartsWith" => Ok(Self::StartsWith),
"EndsWith" => Ok(Self::EndsWith),
"Contains" => Ok(Self::Contains),
"In" => Ok(Self::In),
"ListLiteral" => Ok(Self::ListLiteral),
"MapLiteral" => Ok(Self::MapLiteral),
"Exists" => Ok(Self::Exists),
"ListComprehension" => Ok(Self::ListComprehension),
"Agg" => Ok(Self::Agg),
"Arith" => Ok(Self::Arith),
"Param" => Ok(Self::Param),
"Case" => Ok(Self::Case),
"VectorSimilarity" => Ok(Self::VectorSimilarity),
_ => Err(()),
}
}
}
pub const ALL_OP_KINDS: [OpKind; 25] = [
OpKind::ScanNodes,
OpKind::ScanRels,
OpKind::Expand,
OpKind::OptionalExpand,
OpKind::SemiExpand,
OpKind::ExpandVarLen,
OpKind::Filter,
OpKind::BlockMarker,
OpKind::Project,
OpKind::Aggregate,
OpKind::Sort,
OpKind::Limit,
OpKind::Unwind,
OpKind::PathConstruct,
OpKind::Union,
OpKind::CreateNode,
OpKind::CreateRel,
OpKind::Merge,
OpKind::Delete,
OpKind::SetProperty,
OpKind::RemoveProperty,
OpKind::VectorScan,
OpKind::Rerank,
OpKind::Return,
OpKind::ConstRow,
];
pub const ALL_EXPR_KINDS: [ExprKind; 26] = [
ExprKind::ColRef,
ExprKind::PropAccess,
ExprKind::IntLiteral,
ExprKind::FloatLiteral,
ExprKind::BoolLiteral,
ExprKind::StringLiteral,
ExprKind::NullLiteral,
ExprKind::Cmp,
ExprKind::And,
ExprKind::Or,
ExprKind::Not,
ExprKind::IsNull,
ExprKind::IsNotNull,
ExprKind::StartsWith,
ExprKind::EndsWith,
ExprKind::Contains,
ExprKind::In,
ExprKind::ListLiteral,
ExprKind::MapLiteral,
ExprKind::Exists,
ExprKind::ListComprehension,
ExprKind::Agg,
ExprKind::Arith,
ExprKind::Param,
ExprKind::Case,
ExprKind::VectorSimilarity,
];
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequiredCapabilities {
pub plan_version: PlanSemver,
pub required_ops: BTreeSet<OpKind>,
pub required_exprs: BTreeSet<ExprKind>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EngineCapabilities {
pub version_range: VersionRange,
pub supported_ops: BTreeSet<OpKind>,
pub supported_exprs: BTreeSet<ExprKind>,
pub supports_graph_ref: bool,
pub supports_multi_graph: bool,
pub supports_graph_params: bool,
}
impl EngineCapabilities {
pub fn full(version_range: VersionRange) -> Self {
Self {
version_range,
supported_ops: BTreeSet::from_iter(ALL_OP_KINDS),
supported_exprs: BTreeSet::from_iter(ALL_EXPR_KINDS),
supports_graph_ref: false,
supports_multi_graph: false,
supports_graph_params: false,
}
}
pub fn to_document(&self) -> EngineCapabilityDocument {
EngineCapabilityDocument {
version_range: self.version_range,
supported_ops: self
.supported_ops
.iter()
.copied()
.map(OpKind::as_str)
.map(str::to_string)
.collect(),
supported_exprs: self
.supported_exprs
.iter()
.copied()
.map(ExprKind::as_str)
.map(str::to_string)
.collect(),
op_ordering_contracts: self
.supported_ops
.iter()
.copied()
.map(|op| OpOrderingDocument {
op: op.as_str().to_string(),
contract: op_ordering_contract(op),
})
.collect(),
supports_graph_ref: self.supports_graph_ref,
supports_multi_graph: self.supports_multi_graph,
supports_graph_params: self.supports_graph_params,
}
}
pub fn from_document(doc: EngineCapabilityDocument) -> Result<Self, CapabilityError> {
let mut supported_ops = BTreeSet::new();
for name in doc.supported_ops {
let Some(kind) = name.parse::<OpKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName { kind: "op", name });
};
supported_ops.insert(kind);
}
let mut supported_exprs = BTreeSet::new();
for name in doc.supported_exprs {
let Some(kind) = name.parse::<ExprKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName { kind: "expr", name });
};
supported_exprs.insert(kind);
}
let mut declared_ops = BTreeSet::new();
for decl in doc.op_ordering_contracts {
let Some(kind) = decl.op.parse::<OpKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName {
kind: "op-ordering",
name: decl.op,
});
};
if !supported_ops.contains(&kind) {
return Err(CapabilityError::OrderingDeclarationForUnsupportedOp {
op: kind.as_str().to_string(),
});
}
let expected = op_ordering_contract(kind);
if decl.contract != expected {
return Err(CapabilityError::OrderingContractMismatch {
op: kind.as_str().to_string(),
expected,
actual: decl.contract,
});
}
declared_ops.insert(kind);
}
if !declared_ops.is_empty() {
for op in &supported_ops {
if !declared_ops.contains(op) {
return Err(CapabilityError::MissingOrderingDeclaration {
op: op.as_str().to_string(),
});
}
}
}
Ok(Self {
version_range: doc.version_range,
supported_ops,
supported_exprs,
supports_graph_ref: doc.supports_graph_ref,
supports_multi_graph: doc.supports_multi_graph,
supports_graph_params: doc.supports_graph_params,
})
}
pub fn to_json_pretty(&self) -> Result<String, CapabilityError> {
serde_json::to_string_pretty(&self.to_document()).map_err(CapabilityError::Serialize)
}
pub fn from_json(json: &str) -> Result<Self, CapabilityError> {
let doc: EngineCapabilityDocument =
serde_json::from_str(json).map_err(CapabilityError::Deserialize)?;
Self::from_document(doc)
}
fn to_wire_decl(&self) -> WireCapabilityDecl {
WireCapabilityDecl {
version_range: WireVersionRange {
min_supported: WireSemver {
major: self.version_range.min_supported.major,
minor: self.version_range.min_supported.minor,
patch: self.version_range.min_supported.patch,
},
max_supported: WireSemver {
major: self.version_range.max_supported.major,
minor: self.version_range.max_supported.minor,
patch: self.version_range.max_supported.patch,
},
},
supported_ops: self
.supported_ops
.iter()
.copied()
.map(OpKind::as_str)
.map(str::to_string)
.collect(),
supported_exprs: self
.supported_exprs
.iter()
.copied()
.map(ExprKind::as_str)
.map(str::to_string)
.collect(),
op_ordering: self
.supported_ops
.iter()
.copied()
.map(|op| WireOpOrderingDecl {
op: op.as_str().to_string(),
contract: to_wire_ordering_contract(op_ordering_contract(op)),
})
.collect(),
supports_graph_ref: self.supports_graph_ref,
supports_multi_graph: self.supports_multi_graph,
supports_graph_params: self.supports_graph_params,
}
}
fn from_wire_decl(doc: WireCapabilityDecl) -> Result<Self, CapabilityError> {
let mut supported_ops = BTreeSet::new();
for name in doc.supported_ops {
let Some(kind) = name.parse::<OpKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName { kind: "op", name });
};
supported_ops.insert(kind);
}
let mut supported_exprs = BTreeSet::new();
for name in doc.supported_exprs {
let Some(kind) = name.parse::<ExprKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName { kind: "expr", name });
};
supported_exprs.insert(kind);
}
let mut declared_ops = BTreeSet::new();
for decl in doc.op_ordering {
let Some(kind) = decl.op.parse::<OpKind>().ok() else {
return Err(CapabilityError::InvalidCapabilityName {
kind: "op-ordering",
name: decl.op,
});
};
if !supported_ops.contains(&kind) {
return Err(CapabilityError::OrderingDeclarationForUnsupportedOp {
op: kind.as_str().to_string(),
});
}
let actual = from_wire_ordering_contract(decl.contract);
let expected = op_ordering_contract(kind);
if actual != expected {
return Err(CapabilityError::OrderingContractMismatch {
op: kind.as_str().to_string(),
expected,
actual,
});
}
declared_ops.insert(kind);
}
if !declared_ops.is_empty() {
for op in &supported_ops {
if !declared_ops.contains(op) {
return Err(CapabilityError::MissingOrderingDeclaration {
op: op.as_str().to_string(),
});
}
}
}
Ok(Self {
version_range: VersionRange {
min_supported: PlanSemver {
major: doc.version_range.min_supported.major,
minor: doc.version_range.min_supported.minor,
patch: doc.version_range.min_supported.patch,
},
max_supported: PlanSemver {
major: doc.version_range.max_supported.major,
minor: doc.version_range.max_supported.minor,
patch: doc.version_range.max_supported.patch,
},
},
supported_ops,
supported_exprs,
supports_graph_ref: doc.supports_graph_ref,
supports_multi_graph: doc.supports_multi_graph,
supports_graph_params: doc.supports_graph_params,
})
}
pub fn to_flatbuffer_bytes(&self) -> Result<Vec<u8>, CapabilityError> {
serialize_engine_capability_decl(&self.to_wire_decl()).map_err(CapabilityError::WireSerde)
}
pub fn from_flatbuffer_bytes(bytes: &[u8]) -> Result<Self, CapabilityError> {
let doc = deserialize_engine_capability_decl(bytes).map_err(CapabilityError::WireSerde)?;
Self::from_wire_decl(doc)
}
}
#[derive(Debug, thiserror::Error)]
pub enum CapabilityError {
#[error(
"unsupported plan version {plan_major}.{plan_minor}.{plan_patch}; supported range {min_major}.{min_minor}.{min_patch}..={max_major}.{max_minor}.{max_patch}"
)]
UnsupportedPlanVersion {
plan_major: u32,
plan_minor: u32,
plan_patch: u32,
min_major: u32,
min_minor: u32,
min_patch: u32,
max_major: u32,
max_minor: u32,
max_patch: u32,
},
#[error("plan requires unsupported features")]
MissingFeatureSupport {
missing_ops: Vec<OpKind>,
missing_exprs: Vec<ExprKind>,
},
#[error("plan requires graph_ref support but engine declares supports_graph_ref=false")]
GraphRefUnsupported,
#[error("plan mixes multiple graph_ref values but engine declares supports_multi_graph=false")]
MultiGraphUnsupported,
#[error(
"plan uses graph parameter variables ($g) but engine declares supports_graph_params=false"
)]
GraphParamUnsupported,
#[error("unknown capability {kind} name `{name}`")]
InvalidCapabilityName { kind: &'static str, name: String },
#[error("failed to serialize capability JSON: {0}")]
Serialize(serde_json::Error),
#[error("failed to deserialize capability JSON: {0}")]
Deserialize(serde_json::Error),
#[error("failed to encode/decode capability flatbuffer: {0}")]
WireSerde(plexus_serde::SerdeError),
#[error("ordering declaration provided for unsupported op `{op}`")]
OrderingDeclarationForUnsupportedOp { op: String },
#[error("missing ordering declaration for supported op `{op}`")]
MissingOrderingDeclaration { op: String },
#[error("ordering contract mismatch for `{op}`: expected `{expected:?}`, got `{actual:?}`")]
OrderingContractMismatch {
op: String,
expected: OpOrderingContract,
actual: OpOrderingContract,
},
}