rxgraph 0.6.0

High-performance graph traversal engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

//! Example traversal kernel using a Rust kernel as a plugin.
//!
//! [`WeightedBudget`] accumulates an edge weight along a path and accepts an
//! edge only while the running total stays within a budget, stopping once a
//! target node is reached. It is registered under the name `"weighted_budget"`
//! so it can also be built by name through [`build_kernel`](crate::build_kernel).

use anyhow::{Result, anyhow};

use crate::{
    dsl::{StateRow, Value},
    graph::{Graph, NodeId, OwnedGraphId},
    traversal::{EdgeCtx, Kernel},
};

/// Per-path state for [`WeightedBudget`]: the weight accumulated so far.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BudgetState {
    /// Total edge weight spent along the current path.
    pub spent: u64,
}

/// A budgeted-weight kernel.
///
/// Traversal accumulates `weight_col` along each path. An edge is accepted only
/// while `spent + weight <= budget`, and a path stops once it reaches `target`.
///
/// A null/missing weight rejects the edge (returns `Ok(false)` from
/// [`visit`](Kernel::visit)): a missing cost is treated as "cannot price this
/// hop", which is safer for a budget than silently charging zero.
#[derive(Clone, Debug)]
pub struct WeightedBudget {
    /// Edge payload column holding the `u64` weight.
    pub weight_col: String,
    /// Maximum total weight a returned path may spend.
    pub budget: u64,
    /// Destination node that ends a path.
    pub target: OwnedGraphId,
}

impl Kernel for WeightedBudget {
    type State = BudgetState;

    fn initial_state(&self, _graph: &Graph, _start: NodeId) -> Self::State {
        BudgetState { spent: 0 }
    }

    fn visit(&self, cx: &EdgeCtx<'_, Self::State>) -> Result<bool> {
        // Null/missing weight rejects the edge (see type docs).
        let Some(weight) = cx.edge_u64(&self.weight_col)? else {
            return Ok(false);
        };
        Ok(cx.state().spent.saturating_add(weight) <= self.budget)
    }

    fn next_state(&self, cx: &EdgeCtx<'_, Self::State>) -> Result<Self::State> {
        // `next_state` runs only after `visit` accepted the edge, so the weight
        // is present; default to 0 defensively rather than re-erroring.
        let weight = cx.edge_u64(&self.weight_col)?.unwrap_or(0);
        Ok(BudgetState {
            spent: cx.state().spent.saturating_add(weight),
        })
    }

    fn stop(&self, cx: &EdgeCtx<'_, Self::State>) -> Result<bool> {
        Ok(cx.dest_id() == Some(self.target.as_ref()))
    }

    fn state_row(&self, state: &Self::State) -> StateRow {
        vec![("spent".to_string(), Value::U64(state.spent))]
    }
}

impl WeightedBudget {
    /// Builds a [`WeightedBudget`] from JSON params.
    ///
    /// Accepts something like:
    /// ```json
    /// { "weight_col": "cost", "budget": 10, "target": 3 }
    /// ```
    fn from_params(params: &serde_json::Value) -> Result<Self> {
        let weight_col = params
            .get("weight_col")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow!("weighted_budget: missing string param 'weight_col'"))?
            .to_string();

        let budget = params
            .get("budget")
            .and_then(serde_json::Value::as_u64)
            .ok_or_else(|| anyhow!("weighted_budget: missing u64 param 'budget'"))?;

        let target = match params.get("target") {
            Some(serde_json::Value::Number(n)) => n
                .as_u64()
                .map(OwnedGraphId::U64)
                .ok_or_else(|| anyhow!("weighted_budget: 'target' number must be a u64"))?,
            Some(serde_json::Value::String(s)) => OwnedGraphId::Str(s.clone()),
            _ => {
                return Err(anyhow!(
                    "weighted_budget: 'target' must be a u64 or string node id"
                ));
            }
        };

        Ok(Self {
            weight_col,
            budget,
            target,
        })
    }
}

// Link-time registration: makes `build_kernel("weighted_budget", ..)` work.
crate::inventory::submit! {
    crate::KernelEntry {
        name: "weighted_budget",
        make: |params| Ok(crate::boxed_run(WeightedBudget::from_params(params)?)),
    }
}