Skip to main content

plexus_engine/capabilities/
types.rs

1use std::collections::BTreeSet;
2use std::str::FromStr;
3
4use plexus_serde::{
5    deserialize_engine_capability_decl, serialize_engine_capability_decl,
6    CapabilitySemver as WireSemver, CapabilityVersionRange as WireVersionRange,
7    EngineCapabilityDecl as WireCapabilityDecl, OpOrderingDecl as WireOpOrderingDecl, Version,
8};
9use serde::{Deserialize, Serialize};
10
11use crate::capabilities::ordering::{op_ordering_contract, OpOrderingContract};
12use crate::capabilities::wire::{
13    from_wire_ordering_contract, to_wire_ordering_contract, EngineCapabilityDocument,
14    OpOrderingDocument,
15};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
18pub struct PlanSemver {
19    pub major: u32,
20    pub minor: u32,
21    pub patch: u32,
22}
23
24impl PlanSemver {
25    pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
26        Self {
27            major,
28            minor,
29            patch,
30        }
31    }
32}
33
34impl From<&Version> for PlanSemver {
35    fn from(v: &Version) -> Self {
36        Self::new(v.major, v.minor, v.patch)
37    }
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
41pub struct VersionRange {
42    pub min_supported: PlanSemver,
43    pub max_supported: PlanSemver,
44}
45
46impl VersionRange {
47    pub const fn new(min_supported: PlanSemver, max_supported: PlanSemver) -> Self {
48        Self {
49            min_supported,
50            max_supported,
51        }
52    }
53
54    pub fn supports(&self, version: PlanSemver) -> bool {
55        self.min_supported <= version && version <= self.max_supported
56    }
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
60pub enum OpKind {
61    ScanNodes,
62    ScanRels,
63    Expand,
64    OptionalExpand,
65    SemiExpand,
66    ExpandVarLen,
67    Filter,
68    BlockMarker,
69    Project,
70    Aggregate,
71    Sort,
72    Limit,
73    Unwind,
74    PathConstruct,
75    Union,
76    CreateNode,
77    CreateRel,
78    Merge,
79    Delete,
80    SetProperty,
81    RemoveProperty,
82    VectorScan,
83    Rerank,
84    Return,
85    ConstRow,
86}
87
88impl OpKind {
89    pub fn as_str(self) -> &'static str {
90        match self {
91            Self::ScanNodes => "ScanNodes",
92            Self::ScanRels => "ScanRels",
93            Self::Expand => "Expand",
94            Self::OptionalExpand => "OptionalExpand",
95            Self::SemiExpand => "SemiExpand",
96            Self::ExpandVarLen => "ExpandVarLen",
97            Self::Filter => "Filter",
98            Self::BlockMarker => "BlockMarker",
99            Self::Project => "Project",
100            Self::Aggregate => "Aggregate",
101            Self::Sort => "Sort",
102            Self::Limit => "Limit",
103            Self::Unwind => "Unwind",
104            Self::PathConstruct => "PathConstruct",
105            Self::Union => "Union",
106            Self::CreateNode => "CreateNode",
107            Self::CreateRel => "CreateRel",
108            Self::Merge => "Merge",
109            Self::Delete => "Delete",
110            Self::SetProperty => "SetProperty",
111            Self::RemoveProperty => "RemoveProperty",
112            Self::VectorScan => "VectorScan",
113            Self::Rerank => "Rerank",
114            Self::Return => "Return",
115            Self::ConstRow => "ConstRow",
116        }
117    }
118}
119
120impl FromStr for OpKind {
121    type Err = ();
122
123    fn from_str(name: &str) -> Result<Self, Self::Err> {
124        match name {
125            "ScanNodes" => Ok(Self::ScanNodes),
126            "ScanRels" => Ok(Self::ScanRels),
127            "Expand" => Ok(Self::Expand),
128            "OptionalExpand" => Ok(Self::OptionalExpand),
129            "SemiExpand" => Ok(Self::SemiExpand),
130            "ExpandVarLen" => Ok(Self::ExpandVarLen),
131            "Filter" => Ok(Self::Filter),
132            "BlockMarker" => Ok(Self::BlockMarker),
133            "Project" => Ok(Self::Project),
134            "Aggregate" => Ok(Self::Aggregate),
135            "Sort" => Ok(Self::Sort),
136            "Limit" => Ok(Self::Limit),
137            "Unwind" => Ok(Self::Unwind),
138            "PathConstruct" => Ok(Self::PathConstruct),
139            "Union" => Ok(Self::Union),
140            "CreateNode" => Ok(Self::CreateNode),
141            "CreateRel" => Ok(Self::CreateRel),
142            "Merge" => Ok(Self::Merge),
143            "Delete" => Ok(Self::Delete),
144            "SetProperty" => Ok(Self::SetProperty),
145            "RemoveProperty" => Ok(Self::RemoveProperty),
146            "VectorScan" => Ok(Self::VectorScan),
147            "Rerank" => Ok(Self::Rerank),
148            "Return" => Ok(Self::Return),
149            "ConstRow" => Ok(Self::ConstRow),
150            _ => Err(()),
151        }
152    }
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
156pub enum ExprKind {
157    ColRef,
158    PropAccess,
159    IntLiteral,
160    FloatLiteral,
161    BoolLiteral,
162    StringLiteral,
163    NullLiteral,
164    Cmp,
165    And,
166    Or,
167    Not,
168    IsNull,
169    IsNotNull,
170    StartsWith,
171    EndsWith,
172    Contains,
173    In,
174    ListLiteral,
175    MapLiteral,
176    Exists,
177    ListComprehension,
178    Agg,
179    Arith,
180    Param,
181    Case,
182    VectorSimilarity,
183}
184
185impl ExprKind {
186    pub fn as_str(self) -> &'static str {
187        match self {
188            Self::ColRef => "ColRef",
189            Self::PropAccess => "PropAccess",
190            Self::IntLiteral => "IntLiteral",
191            Self::FloatLiteral => "FloatLiteral",
192            Self::BoolLiteral => "BoolLiteral",
193            Self::StringLiteral => "StringLiteral",
194            Self::NullLiteral => "NullLiteral",
195            Self::Cmp => "Cmp",
196            Self::And => "And",
197            Self::Or => "Or",
198            Self::Not => "Not",
199            Self::IsNull => "IsNull",
200            Self::IsNotNull => "IsNotNull",
201            Self::StartsWith => "StartsWith",
202            Self::EndsWith => "EndsWith",
203            Self::Contains => "Contains",
204            Self::In => "In",
205            Self::ListLiteral => "ListLiteral",
206            Self::MapLiteral => "MapLiteral",
207            Self::Exists => "Exists",
208            Self::ListComprehension => "ListComprehension",
209            Self::Agg => "Agg",
210            Self::Arith => "Arith",
211            Self::Param => "Param",
212            Self::Case => "Case",
213            Self::VectorSimilarity => "VectorSimilarity",
214        }
215    }
216}
217
218impl FromStr for ExprKind {
219    type Err = ();
220
221    fn from_str(name: &str) -> Result<Self, Self::Err> {
222        match name {
223            "ColRef" => Ok(Self::ColRef),
224            "PropAccess" => Ok(Self::PropAccess),
225            "IntLiteral" => Ok(Self::IntLiteral),
226            "FloatLiteral" => Ok(Self::FloatLiteral),
227            "BoolLiteral" => Ok(Self::BoolLiteral),
228            "StringLiteral" => Ok(Self::StringLiteral),
229            "NullLiteral" => Ok(Self::NullLiteral),
230            "Cmp" => Ok(Self::Cmp),
231            "And" => Ok(Self::And),
232            "Or" => Ok(Self::Or),
233            "Not" => Ok(Self::Not),
234            "IsNull" => Ok(Self::IsNull),
235            "IsNotNull" => Ok(Self::IsNotNull),
236            "StartsWith" => Ok(Self::StartsWith),
237            "EndsWith" => Ok(Self::EndsWith),
238            "Contains" => Ok(Self::Contains),
239            "In" => Ok(Self::In),
240            "ListLiteral" => Ok(Self::ListLiteral),
241            "MapLiteral" => Ok(Self::MapLiteral),
242            "Exists" => Ok(Self::Exists),
243            "ListComprehension" => Ok(Self::ListComprehension),
244            "Agg" => Ok(Self::Agg),
245            "Arith" => Ok(Self::Arith),
246            "Param" => Ok(Self::Param),
247            "Case" => Ok(Self::Case),
248            "VectorSimilarity" => Ok(Self::VectorSimilarity),
249            _ => Err(()),
250        }
251    }
252}
253
254pub const ALL_OP_KINDS: [OpKind; 25] = [
255    OpKind::ScanNodes,
256    OpKind::ScanRels,
257    OpKind::Expand,
258    OpKind::OptionalExpand,
259    OpKind::SemiExpand,
260    OpKind::ExpandVarLen,
261    OpKind::Filter,
262    OpKind::BlockMarker,
263    OpKind::Project,
264    OpKind::Aggregate,
265    OpKind::Sort,
266    OpKind::Limit,
267    OpKind::Unwind,
268    OpKind::PathConstruct,
269    OpKind::Union,
270    OpKind::CreateNode,
271    OpKind::CreateRel,
272    OpKind::Merge,
273    OpKind::Delete,
274    OpKind::SetProperty,
275    OpKind::RemoveProperty,
276    OpKind::VectorScan,
277    OpKind::Rerank,
278    OpKind::Return,
279    OpKind::ConstRow,
280];
281
282pub const ALL_EXPR_KINDS: [ExprKind; 26] = [
283    ExprKind::ColRef,
284    ExprKind::PropAccess,
285    ExprKind::IntLiteral,
286    ExprKind::FloatLiteral,
287    ExprKind::BoolLiteral,
288    ExprKind::StringLiteral,
289    ExprKind::NullLiteral,
290    ExprKind::Cmp,
291    ExprKind::And,
292    ExprKind::Or,
293    ExprKind::Not,
294    ExprKind::IsNull,
295    ExprKind::IsNotNull,
296    ExprKind::StartsWith,
297    ExprKind::EndsWith,
298    ExprKind::Contains,
299    ExprKind::In,
300    ExprKind::ListLiteral,
301    ExprKind::MapLiteral,
302    ExprKind::Exists,
303    ExprKind::ListComprehension,
304    ExprKind::Agg,
305    ExprKind::Arith,
306    ExprKind::Param,
307    ExprKind::Case,
308    ExprKind::VectorSimilarity,
309];
310
311#[derive(Debug, Clone, PartialEq, Eq)]
312pub struct RequiredCapabilities {
313    pub plan_version: PlanSemver,
314    pub required_ops: BTreeSet<OpKind>,
315    pub required_exprs: BTreeSet<ExprKind>,
316}
317
318#[derive(Debug, Clone, PartialEq, Eq)]
319pub struct EngineCapabilities {
320    pub version_range: VersionRange,
321    pub supported_ops: BTreeSet<OpKind>,
322    pub supported_exprs: BTreeSet<ExprKind>,
323    pub supports_graph_ref: bool,
324    pub supports_multi_graph: bool,
325    pub supports_graph_params: bool,
326}
327
328impl EngineCapabilities {
329    pub fn full(version_range: VersionRange) -> Self {
330        Self {
331            version_range,
332            supported_ops: BTreeSet::from_iter(ALL_OP_KINDS),
333            supported_exprs: BTreeSet::from_iter(ALL_EXPR_KINDS),
334            supports_graph_ref: false,
335            supports_multi_graph: false,
336            supports_graph_params: false,
337        }
338    }
339
340    pub fn to_document(&self) -> EngineCapabilityDocument {
341        EngineCapabilityDocument {
342            version_range: self.version_range,
343            supported_ops: self
344                .supported_ops
345                .iter()
346                .copied()
347                .map(OpKind::as_str)
348                .map(str::to_string)
349                .collect(),
350            supported_exprs: self
351                .supported_exprs
352                .iter()
353                .copied()
354                .map(ExprKind::as_str)
355                .map(str::to_string)
356                .collect(),
357            op_ordering_contracts: self
358                .supported_ops
359                .iter()
360                .copied()
361                .map(|op| OpOrderingDocument {
362                    op: op.as_str().to_string(),
363                    contract: op_ordering_contract(op),
364                })
365                .collect(),
366            supports_graph_ref: self.supports_graph_ref,
367            supports_multi_graph: self.supports_multi_graph,
368            supports_graph_params: self.supports_graph_params,
369        }
370    }
371
372    pub fn from_document(doc: EngineCapabilityDocument) -> Result<Self, CapabilityError> {
373        let mut supported_ops = BTreeSet::new();
374        for name in doc.supported_ops {
375            let Some(kind) = name.parse::<OpKind>().ok() else {
376                return Err(CapabilityError::InvalidCapabilityName { kind: "op", name });
377            };
378            supported_ops.insert(kind);
379        }
380
381        let mut supported_exprs = BTreeSet::new();
382        for name in doc.supported_exprs {
383            let Some(kind) = name.parse::<ExprKind>().ok() else {
384                return Err(CapabilityError::InvalidCapabilityName { kind: "expr", name });
385            };
386            supported_exprs.insert(kind);
387        }
388
389        let mut declared_ops = BTreeSet::new();
390        for decl in doc.op_ordering_contracts {
391            let Some(kind) = decl.op.parse::<OpKind>().ok() else {
392                return Err(CapabilityError::InvalidCapabilityName {
393                    kind: "op-ordering",
394                    name: decl.op,
395                });
396            };
397            if !supported_ops.contains(&kind) {
398                return Err(CapabilityError::OrderingDeclarationForUnsupportedOp {
399                    op: kind.as_str().to_string(),
400                });
401            }
402            let expected = op_ordering_contract(kind);
403            if decl.contract != expected {
404                return Err(CapabilityError::OrderingContractMismatch {
405                    op: kind.as_str().to_string(),
406                    expected,
407                    actual: decl.contract,
408                });
409            }
410            declared_ops.insert(kind);
411        }
412
413        if !declared_ops.is_empty() {
414            for op in &supported_ops {
415                if !declared_ops.contains(op) {
416                    return Err(CapabilityError::MissingOrderingDeclaration {
417                        op: op.as_str().to_string(),
418                    });
419                }
420            }
421        }
422
423        Ok(Self {
424            version_range: doc.version_range,
425            supported_ops,
426            supported_exprs,
427            supports_graph_ref: doc.supports_graph_ref,
428            supports_multi_graph: doc.supports_multi_graph,
429            supports_graph_params: doc.supports_graph_params,
430        })
431    }
432
433    pub fn to_json_pretty(&self) -> Result<String, CapabilityError> {
434        serde_json::to_string_pretty(&self.to_document()).map_err(CapabilityError::Serialize)
435    }
436
437    pub fn from_json(json: &str) -> Result<Self, CapabilityError> {
438        let doc: EngineCapabilityDocument =
439            serde_json::from_str(json).map_err(CapabilityError::Deserialize)?;
440        Self::from_document(doc)
441    }
442
443    fn to_wire_decl(&self) -> WireCapabilityDecl {
444        WireCapabilityDecl {
445            version_range: WireVersionRange {
446                min_supported: WireSemver {
447                    major: self.version_range.min_supported.major,
448                    minor: self.version_range.min_supported.minor,
449                    patch: self.version_range.min_supported.patch,
450                },
451                max_supported: WireSemver {
452                    major: self.version_range.max_supported.major,
453                    minor: self.version_range.max_supported.minor,
454                    patch: self.version_range.max_supported.patch,
455                },
456            },
457            supported_ops: self
458                .supported_ops
459                .iter()
460                .copied()
461                .map(OpKind::as_str)
462                .map(str::to_string)
463                .collect(),
464            supported_exprs: self
465                .supported_exprs
466                .iter()
467                .copied()
468                .map(ExprKind::as_str)
469                .map(str::to_string)
470                .collect(),
471            op_ordering: self
472                .supported_ops
473                .iter()
474                .copied()
475                .map(|op| WireOpOrderingDecl {
476                    op: op.as_str().to_string(),
477                    contract: to_wire_ordering_contract(op_ordering_contract(op)),
478                })
479                .collect(),
480            supports_graph_ref: self.supports_graph_ref,
481            supports_multi_graph: self.supports_multi_graph,
482            supports_graph_params: self.supports_graph_params,
483        }
484    }
485
486    fn from_wire_decl(doc: WireCapabilityDecl) -> Result<Self, CapabilityError> {
487        let mut supported_ops = BTreeSet::new();
488        for name in doc.supported_ops {
489            let Some(kind) = name.parse::<OpKind>().ok() else {
490                return Err(CapabilityError::InvalidCapabilityName { kind: "op", name });
491            };
492            supported_ops.insert(kind);
493        }
494
495        let mut supported_exprs = BTreeSet::new();
496        for name in doc.supported_exprs {
497            let Some(kind) = name.parse::<ExprKind>().ok() else {
498                return Err(CapabilityError::InvalidCapabilityName { kind: "expr", name });
499            };
500            supported_exprs.insert(kind);
501        }
502
503        let mut declared_ops = BTreeSet::new();
504        for decl in doc.op_ordering {
505            let Some(kind) = decl.op.parse::<OpKind>().ok() else {
506                return Err(CapabilityError::InvalidCapabilityName {
507                    kind: "op-ordering",
508                    name: decl.op,
509                });
510            };
511            if !supported_ops.contains(&kind) {
512                return Err(CapabilityError::OrderingDeclarationForUnsupportedOp {
513                    op: kind.as_str().to_string(),
514                });
515            }
516            let actual = from_wire_ordering_contract(decl.contract);
517            let expected = op_ordering_contract(kind);
518            if actual != expected {
519                return Err(CapabilityError::OrderingContractMismatch {
520                    op: kind.as_str().to_string(),
521                    expected,
522                    actual,
523                });
524            }
525            declared_ops.insert(kind);
526        }
527
528        if !declared_ops.is_empty() {
529            for op in &supported_ops {
530                if !declared_ops.contains(op) {
531                    return Err(CapabilityError::MissingOrderingDeclaration {
532                        op: op.as_str().to_string(),
533                    });
534                }
535            }
536        }
537
538        Ok(Self {
539            version_range: VersionRange {
540                min_supported: PlanSemver {
541                    major: doc.version_range.min_supported.major,
542                    minor: doc.version_range.min_supported.minor,
543                    patch: doc.version_range.min_supported.patch,
544                },
545                max_supported: PlanSemver {
546                    major: doc.version_range.max_supported.major,
547                    minor: doc.version_range.max_supported.minor,
548                    patch: doc.version_range.max_supported.patch,
549                },
550            },
551            supported_ops,
552            supported_exprs,
553            supports_graph_ref: doc.supports_graph_ref,
554            supports_multi_graph: doc.supports_multi_graph,
555            supports_graph_params: doc.supports_graph_params,
556        })
557    }
558
559    pub fn to_flatbuffer_bytes(&self) -> Result<Vec<u8>, CapabilityError> {
560        serialize_engine_capability_decl(&self.to_wire_decl()).map_err(CapabilityError::WireSerde)
561    }
562
563    pub fn from_flatbuffer_bytes(bytes: &[u8]) -> Result<Self, CapabilityError> {
564        let doc = deserialize_engine_capability_decl(bytes).map_err(CapabilityError::WireSerde)?;
565        Self::from_wire_decl(doc)
566    }
567}
568
569#[derive(Debug, thiserror::Error)]
570pub enum CapabilityError {
571    #[error(
572        "unsupported plan version {plan_major}.{plan_minor}.{plan_patch}; supported range {min_major}.{min_minor}.{min_patch}..={max_major}.{max_minor}.{max_patch}"
573    )]
574    UnsupportedPlanVersion {
575        plan_major: u32,
576        plan_minor: u32,
577        plan_patch: u32,
578        min_major: u32,
579        min_minor: u32,
580        min_patch: u32,
581        max_major: u32,
582        max_minor: u32,
583        max_patch: u32,
584    },
585    #[error("plan requires unsupported features")]
586    MissingFeatureSupport {
587        missing_ops: Vec<OpKind>,
588        missing_exprs: Vec<ExprKind>,
589    },
590    #[error("plan requires graph_ref support but engine declares supports_graph_ref=false")]
591    GraphRefUnsupported,
592    #[error("plan mixes multiple graph_ref values but engine declares supports_multi_graph=false")]
593    MultiGraphUnsupported,
594    #[error(
595        "plan uses graph parameter variables ($g) but engine declares supports_graph_params=false"
596    )]
597    GraphParamUnsupported,
598    #[error("unknown capability {kind} name `{name}`")]
599    InvalidCapabilityName { kind: &'static str, name: String },
600    #[error("failed to serialize capability JSON: {0}")]
601    Serialize(serde_json::Error),
602    #[error("failed to deserialize capability JSON: {0}")]
603    Deserialize(serde_json::Error),
604    #[error("failed to encode/decode capability flatbuffer: {0}")]
605    WireSerde(plexus_serde::SerdeError),
606    #[error("ordering declaration provided for unsupported op `{op}`")]
607    OrderingDeclarationForUnsupportedOp { op: String },
608    #[error("missing ordering declaration for supported op `{op}`")]
609    MissingOrderingDeclaration { op: String },
610    #[error("ordering contract mismatch for `{op}`: expected `{expected:?}`, got `{actual:?}`")]
611    OrderingContractMismatch {
612        op: String,
613        expected: OpOrderingContract,
614        actual: OpOrderingContract,
615    },
616}