1use std::any::TypeId;
10use std::collections::{HashMap, HashSet};
11use std::sync::Arc;
12
13use crate::persistence::{PersistenceDefinition, PersistenceKind};
14use crate::rpc::RegisteredRpc;
15use crate::step::RegisteredStep;
16use crate::step_options::ErasedStepOptions;
17use crate::{Context, Flow, HandlerResult, SdkError, SdkResult, StepDecision, Wait};
18use dex_protocol::dex::Value as ProtoValue;
19
20#[derive(Clone, Default)]
21pub struct Registry {
39 inner: Arc<RegistryInner>,
40}
41
42#[derive(Clone, Default)]
43struct RegistryInner {
44 flows: HashMap<&'static str, RegisteredFlow>,
45 attribute_indexes: HashMap<String, i32>,
46}
47
48#[derive(Clone)]
49pub(crate) struct RegisteredFlow {
50 pub(crate) name: &'static str,
51 pub(crate) type_id: TypeId,
52 pub(crate) steps: HashMap<&'static str, RegisteredStep>,
53 pub(crate) start_step: Option<RegisteredStep>,
54 pub(crate) rpcs: HashMap<&'static str, RegisteredRpc>,
55 pub(crate) persistence: HashMap<String, PersistenceDefinition>,
56 pub(crate) handler: Arc<dyn ErasedFlow>,
57}
58
59impl Registry {
60 pub fn new() -> Self {
62 Self::default()
63 }
64
65 pub fn register<SomeFlow: Flow>(mut self, flow: SomeFlow) -> SdkResult<Self> {
72 let flow = Arc::new(flow);
73 let name = require_static_name(flow.flow_type(), "Flow type")?;
74 let registered = assemble_flow(name, Arc::clone(&flow))?;
75 let inner = Arc::make_mut(&mut self.inner);
76 let mut additions = Vec::new();
77 for definition in registered.persistence.values() {
78 let Some(index) = &definition.index else {
79 continue;
80 };
81 if definition.kind == PersistenceKind::AttributeMap
82 && index.key().is_none_or(str::is_empty)
83 {
84 return Err(definition_error(format!(
85 "Flow {name} indexed AttributeMap {} requires an index key",
86 definition.name
87 )));
88 }
89 let key = index.key().unwrap_or(&definition.name).to_string();
90 let index_type = index.proto_value();
91 if inner
92 .attribute_indexes
93 .get(&key)
94 .is_some_and(|existing| *existing != index_type)
95 {
96 return Err(definition_error(format!(
97 "Attribute index {key} has conflicting types"
98 )));
99 }
100 additions.push((key, index_type));
101 }
102 if inner.flows.insert(name, registered).is_some() {
103 return Err(definition_error(format!("duplicate Flow {name}")));
104 }
105 inner.attribute_indexes.extend(additions);
106 Ok(self)
107 }
108
109 pub(crate) fn flow(&self, name: &str) -> SdkResult<&RegisteredFlow> {
110 self.inner
111 .flows
112 .get(name)
113 .ok_or_else(|| SdkError::FlowDefinition {
114 message: format!("Flow is not registered: {name}"),
115 })
116 }
117
118 pub(crate) fn rpc(&self, name: &str) -> SdkResult<&RegisteredRpc> {
119 let mut matched = None;
120 for flow in self.inner.flows.values() {
121 let Some(rpc) = flow.rpcs.get(name) else {
122 continue;
123 };
124 if matched.is_some() {
125 return Err(SdkError::FlowDefinition {
126 message: format!("RPC name is ambiguous across registered Flows: {name}"),
127 });
128 }
129 matched = Some(rpc);
130 }
131 matched.ok_or_else(|| SdkError::FlowDefinition {
132 message: format!("RPC is not registered: {name}"),
133 })
134 }
135
136 pub(crate) fn attribute_indexes(&self) -> &HashMap<String, i32> {
137 &self.inner.attribute_indexes
138 }
139}
140
141fn assemble_flow<SomeFlow: Flow>(
142 name: &'static str,
143 flow: Arc<SomeFlow>,
144) -> SdkResult<RegisteredFlow> {
145 let mut steps = HashMap::new();
146 let mut start_step = None;
147 for step in flow.steps().into_definitions() {
148 require_static_name(step.name, "Step type")?;
149 if step.starting {
150 if start_step.is_some() {
151 return Err(definition_error(format!(
152 "Flow {name} must not have multiple start Steps"
153 )));
154 }
155 start_step = Some(step.clone());
156 }
157 if steps.insert(step.name, step).is_some() {
158 return Err(definition_error(format!(
159 "Flow {name} has a duplicate Step"
160 )));
161 }
162 }
163
164 let persistence = assemble_persistence(name, flow.persistence().definitions())?;
165 let mut rpcs = HashMap::new();
166 for rpc in flow.rpcs().bind(Arc::clone(&flow)) {
167 require_static_name(rpc.name, "RPC name")?;
168 validate_rpc(name, &rpc, &persistence)?;
169 if rpcs.insert(rpc.name, rpc).is_some() {
170 return Err(definition_error(format!("Flow {name} has a duplicate RPC")));
171 }
172 }
173 Ok(RegisteredFlow {
174 name,
175 type_id: TypeId::of::<SomeFlow>(),
176 steps,
177 start_step,
178 rpcs,
179 persistence,
180 handler: Arc::new(TypedFlow { flow }),
181 })
182}
183
184pub(crate) trait ErasedFlow: Send + Sync {
185 fn wait_for(
186 &self,
187 step_type: &str,
188 context: &mut Context,
189 input: &ProtoValue,
190 ) -> HandlerResult<Wait>;
191 fn execute(
192 &self,
193 step_type: &str,
194 context: &mut Context,
195 input: &ProtoValue,
196 ) -> HandlerResult<StepDecision>;
197 fn step_options(&self, step_type: &str) -> HandlerResult<ErasedStepOptions>;
198}
199
200struct TypedFlow<SomeFlow> {
201 flow: Arc<SomeFlow>,
202}
203
204impl<SomeFlow: Flow> ErasedFlow for TypedFlow<SomeFlow> {
205 fn wait_for(
206 &self,
207 step_type: &str,
208 context: &mut Context,
209 input: &ProtoValue,
210 ) -> HandlerResult<Wait> {
211 self.flow
212 .steps()
213 .find(step_type)
214 .ok_or_else(|| {
215 crate::HandlerError::new(format!("Step is not registered: {step_type}"))
216 })?
217 .wait_for(context, input)
218 }
219
220 fn execute(
221 &self,
222 step_type: &str,
223 context: &mut Context,
224 input: &ProtoValue,
225 ) -> HandlerResult<StepDecision> {
226 self.flow
227 .steps()
228 .find(step_type)
229 .ok_or_else(|| {
230 crate::HandlerError::new(format!("Step is not registered: {step_type}"))
231 })?
232 .execute(context, input)
233 }
234
235 fn step_options(&self, step_type: &str) -> HandlerResult<ErasedStepOptions> {
236 Ok(self
237 .flow
238 .steps()
239 .find(step_type)
240 .ok_or_else(|| {
241 crate::HandlerError::new(format!("Step is not registered: {step_type}"))
242 })?
243 .options())
244 }
245}
246
247fn assemble_persistence(
248 flow_name: &str,
249 definitions: &[PersistenceDefinition],
250) -> SdkResult<HashMap<String, PersistenceDefinition>> {
251 let mut persistence = HashMap::new();
252 for definition in definitions {
253 if definition.name.is_empty() {
254 return Err(definition_error(format!(
255 "Flow {flow_name} has an empty persistence definition name"
256 )));
257 }
258 if persistence
259 .insert(definition.name.clone(), definition.clone())
260 .is_some()
261 {
262 return Err(definition_error(format!(
263 "Flow {flow_name} has duplicate persistence definition {}",
264 definition.name
265 )));
266 }
267 }
268 Ok(persistence)
269}
270
271fn validate_rpc(
272 flow_name: &str,
273 rpc: &RegisteredRpc,
274 persistence: &HashMap<String, PersistenceDefinition>,
275) -> SdkResult<()> {
276 if rpc.timeout.is_some_and(|timeout| timeout.is_zero()) {
277 return Err(definition_error(format!(
278 "Flow {flow_name} RPC {} timeout must be positive",
279 rpc.name
280 )));
281 }
282 let mut locks = HashSet::new();
283 for lock in &rpc.locks {
284 let physical_name = lock.physical_name();
285 let logical_name = physical_name.split('/').next().unwrap_or(&physical_name);
286 let definition = persistence.get(logical_name).ok_or_else(|| {
287 definition_error(format!(
288 "Flow {flow_name} RPC {} locks an unregistered Attribute",
289 rpc.name
290 ))
291 })?;
292 if !matches!(
293 definition.kind,
294 PersistenceKind::Attribute | PersistenceKind::AttributeMap
295 ) {
296 return Err(definition_error(format!(
297 "Flow {flow_name} RPC {} locks a non-Attribute",
298 rpc.name
299 )));
300 }
301 if !locks.insert(physical_name) {
302 return Err(definition_error(format!(
303 "Flow {flow_name} RPC {} has a duplicate lock",
304 rpc.name
305 )));
306 }
307 }
308 Ok(())
309}
310
311fn require_static_name(name: &'static str, kind: &str) -> SdkResult<&'static str> {
312 if name.is_empty() {
313 Err(definition_error(format!("{kind} is required")))
314 } else {
315 Ok(name)
316 }
317}
318
319fn definition_error(message: impl Into<String>) -> SdkError {
320 SdkError::FlowDefinition {
321 message: message.into(),
322 }
323}
324
325pub(crate) fn physical_name(name: &str, instance: &str) -> String {
326 if instance.is_empty() {
327 panic!("persistence instance is required");
328 }
329 let mut encoded = String::new();
330 for byte in instance.bytes() {
331 if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') {
332 encoded.push(char::from(byte));
333 } else {
334 encoded.push_str(&format!("%{byte:02X}"));
335 }
336 }
337 format!("{name}/{encoded}")
338}
339
340pub(crate) fn decode_instance(physical_name: &str, prefix: &str) -> Result<String, String> {
341 let encoded = physical_name
342 .strip_prefix(prefix)
343 .ok_or_else(|| format!("map key {physical_name} does not start with {prefix}"))?;
344 let bytes = encoded.as_bytes();
345 let mut decoded = Vec::with_capacity(bytes.len());
346 let mut index = 0;
347 while index < bytes.len() {
348 if bytes[index] == b'%' {
349 if index + 2 >= bytes.len() {
350 return Err(format!("invalid percent encoding in {physical_name}"));
351 }
352 let high = hex_value(bytes[index + 1])?;
353 let low = hex_value(bytes[index + 2])?;
354 decoded.push(high * 16 + low);
355 index += 3;
356 } else {
357 decoded.push(bytes[index]);
358 index += 1;
359 }
360 }
361 String::from_utf8(decoded).map_err(|error| error.to_string())
362}
363
364fn hex_value(byte: u8) -> Result<u8, String> {
365 match byte {
366 b'0'..=b'9' => Ok(byte - b'0'),
367 b'a'..=b'f' => Ok(byte - b'a' + 10),
368 b'A'..=b'F' => Ok(byte - b'A' + 10),
369 _ => Err(format!(
370 "invalid percent-encoding digit {}",
371 char::from(byte)
372 )),
373 }
374}