stainless_script/
node.rs

1use crate::{
2    class::Class,
3    module::{ModulePath, ModulePathParseError},
4    program::ProgramId,
5    socket::{InputSocket, OutputSocket},
6    ExecutionContext,
7};
8use serde::{Deserialize, Serialize};
9use std::{
10    borrow::Cow,
11    collections::BTreeMap,
12    fmt::{Debug, Display},
13    num::ParseIntError,
14    rc::Rc,
15    str::FromStr,
16};
17use thiserror::Error;
18
19/// ID of a node
20pub type NodeId = u32;
21
22/// An ID to point to a node in other program
23#[derive(Debug, Clone, Hash, PartialEq, Eq, Deserialize, Serialize)]
24pub struct AbsoluteNodeId(pub ProgramId, pub NodeId);
25
26impl Display for AbsoluteNodeId {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        write!(f, "{}@{}", self.0, self.1)
29    }
30}
31
32impl FromStr for AbsoluteNodeId {
33    type Err = AbsoluteNodeIdParseError;
34
35    fn from_str(s: &str) -> Result<Self, Self::Err> {
36        let mut seq: Vec<String> = s.split('@').map(String::from).collect();
37        let node_id: NodeId = seq
38            .pop()
39            .ok_or(AbsoluteNodeIdParseError::IdNotFound)?
40            .parse()?;
41        let path: ProgramId = seq[0].parse()?;
42        Ok(Self(path, node_id))
43    }
44}
45
46#[derive(Debug, Clone, Error)]
47pub enum AbsoluteNodeIdParseError {
48    #[error("Node ID not found in string")]
49    IdNotFound,
50    #[error("Failed to parse Node ID: {0}")]
51    NodeIdParseError(ParseIntError),
52    #[error("Failed to parse program ID path: {0}")]
53    ProgramIdParseError(ModulePathParseError),
54}
55
56impl From<ParseIntError> for AbsoluteNodeIdParseError {
57    fn from(e: ParseIntError) -> Self {
58        Self::NodeIdParseError(e)
59    }
60}
61
62impl From<ModulePathParseError> for AbsoluteNodeIdParseError {
63    fn from(e: ModulePathParseError) -> Self {
64        Self::ProgramIdParseError(e)
65    }
66}
67
68/// ID of a branch of node
69#[derive(Debug, Clone, Hash, PartialEq, Eq)]
70pub struct NodeBranchId(pub NodeId, pub usize);
71
72impl From<&NodeBranchId> for u64 {
73    fn from(s: &NodeBranchId) -> Self {
74        (s.0 as u64) << 32 | s.1 as u64
75    }
76}
77
78impl From<u64> for NodeBranchId {
79    fn from(n: u64) -> Self {
80        let branch_idx: u32 = (((1 << 33) - 1) & n) as u32;
81        let node_id: NodeId = ((((1 << 33) - 1) << 32) & n) as NodeId;
82        Self(node_id, branch_idx as usize)
83    }
84}
85
86impl Serialize for NodeBranchId {
87    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
88    where
89        S: serde::Serializer,
90    {
91        u64::from(self).serialize(serializer)
92    }
93}
94
95impl<'de> Deserialize<'de> for NodeBranchId {
96    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
97    where
98        D: serde::Deserializer<'de>,
99    {
100        Ok(NodeBranchId::from(u64::deserialize(deserializer)?))
101    }
102}
103
104pub trait Node: Debug {
105    /// Execution of the node's code. Returns a branch index.
106    fn execute(&self, context: &mut ExecutionContext) -> usize;
107
108    /// The class of the node
109    fn class(&self) -> Class;
110
111    /// Variants of a node. Internally can be anythingg that can be converted to string
112    fn variants(&self) -> Vec<Cow<'_, str>>;
113
114    /// Current selected variant of the node
115    fn current_variant(&self) -> Cow<'_, str>;
116
117    /// Set a specific variant of a node
118    fn set_variant(&mut self, variant: &str);
119
120    /// Whether variation can be set as a custom string (not listed in `variants`) or not
121    fn accepts_arbitrary_variants(&self) -> bool {
122        false
123    }
124
125    /// Get information about node's inputs
126    fn inputs(&self) -> Vec<InputSocket>;
127
128    /// Get information about node's outputs
129    fn outputs(&self) -> Vec<OutputSocket>;
130
131    /// How many branches this node has
132    fn branches(&self) -> u32 {
133        1
134    }
135
136    /// Clone the node itself instead of it wrapped in Rc
137    fn clone_node(&self) -> Rc<dyn Node>;
138}
139
140#[derive(Debug, Clone, Default)]
141pub struct NodeStorage {
142    pub nodes: BTreeMap<NodeId, Rc<dyn Node>>,
143    next_vacant: NodeId,
144}
145
146impl NodeStorage {
147    pub fn get_node(&self, node_id: NodeId) -> Option<Rc<dyn Node>> {
148        self.nodes.get(&node_id).cloned()
149    }
150
151    pub fn remove_node(&mut self, node_id: NodeId) -> Option<Rc<dyn Node>> {
152        let node = self.nodes.remove(&node_id);
153        if node_id < self.next_vacant {
154            self.next_vacant = node_id
155        }
156        node
157    }
158
159    pub fn insert_node(&mut self, node: Rc<dyn Node>) -> NodeId {
160        let mut node_id = self.next_vacant;
161        self.nodes.insert(node_id, node);
162        while self.nodes.get(&node_id).is_some() {
163            node_id += 1;
164        }
165        self.next_vacant = node_id;
166        node_id
167    }
168
169    pub fn insert_node_at(&mut self, node_id: NodeId, node: Rc<dyn Node>) {
170        self.nodes.insert(node_id, node);
171        while self.nodes.get(&self.next_vacant).is_some() {
172            self.next_vacant += 1;
173        }
174    }
175}
176
177/// Information about a node stored in the program
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct NodeInfo {
180    pub class: ModulePath,
181    pub idx: usize,
182    pub variant: String,
183}