1use std::{
3 borrow::Borrow,
4 collections::HashMap,
5 fmt::{Display, Write as _},
6};
7
8use serde::{
9 Serialize,
10 ser::{SerializeMap, SerializeStruct as _},
11};
12
13use crate::types::{BucketType, ComputeType, MapType, TagsType, TimeType};
14
15#[derive(Debug, Clone, serde::Serialize)]
16pub enum ArgType {
18 Float,
20 Enum(&'static [&'static str]),
22 Repeated {
24 typ: Box<ArgType>,
26 min: usize,
28 max: Option<usize>,
30 },
31 OneOf(Vec<ArgType>),
33 Optional(Box<ArgType>),
35}
36impl Display for ArgType {
37 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38 match self {
39 ArgType::Float => write!(f, "float"),
40 ArgType::Enum(values) => write!(f, "enum({})", values.join(", ")),
41 ArgType::Repeated { typ, min, max } => {
42 write!(f, "repeated({typ}")?;
43 if *min > 0 {
44 write!(f, ", min={min}")?;
45 }
46 if let Some(max) = max {
47 write!(f, ", max={max}")?;
48 }
49 write!(f, ")")
50 }
51 ArgType::OneOf(types) => write!(
52 f,
53 "one_of({})",
54 types
55 .iter()
56 .map(ToString::to_string)
57 .collect::<Vec<String>>()
58 .join(", ")
59 ),
60 ArgType::Optional(typ) => write!(f, "[{typ}]"),
61 }
62 }
63}
64
65#[derive(Debug, Clone, serde::Serialize)]
66pub struct Arg {
68 pub name: &'static str,
70 pub typ: ArgType,
72}
73impl Display for Arg {
74 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75 write!(f, "{}: {}", self.name, self.typ)
76 }
77}
78impl Arg {
79 #[must_use]
81 pub const fn new(name: &'static str, typ: ArgType) -> Self {
82 Self { name, typ }
83 }
84}
85
86pub trait FunctionTrait {
88 fn doc(&self) -> &str;
90 fn args(&self) -> Vec<Arg>;
92 fn documentation(&self, name: &FunctionId) -> String {
94 let args = self
95 .args()
96 .iter()
97 .map(ToString::to_string)
98 .collect::<Vec<_>>()
99 .join(", ");
100 let doc = self.doc();
101 let e = if name.0.contains('*') { "__" } else { "**" };
102 if args.is_empty() {
103 format!(
104 r"{e}{name}{e}:
105{doc}"
106 )
107 } else {
108 format!(
109 r"{e}{name}{e}({args}):
110
111{doc}"
112 )
113 }
114 }
115}
116
117#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize)]
119pub(crate) struct ModuleId(pub(crate) String);
120
121impl std::fmt::Display for ModuleId {
122 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123 write!(f, "{}", self.0)
124 }
125}
126
127impl ModuleId {
128 pub(crate) fn new(name: &str) -> Self {
129 ModuleId(name.to_string())
130 }
131}
132
133impl Borrow<str> for ModuleId {
134 fn borrow(&self) -> &str {
135 &self.0
136 }
137}
138
139#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize)]
141pub struct FunctionId(pub(crate) String);
142
143impl std::fmt::Display for FunctionId {
144 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145 write!(f, "{}", self.0)
146 }
147}
148impl FunctionId {
149 pub(crate) fn new(name: &str) -> Self {
150 FunctionId(name.to_string())
151 }
152}
153
154impl Borrow<str> for FunctionId {
155 fn borrow(&self) -> &str {
156 &self.0
157 }
158}
159
160pub(crate) struct Function {
161 pub(crate) module_path: Vec<ModuleId>,
162 pub(crate) name: FunctionId,
163}
164
165impl std::fmt::Display for Function {
166 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167 for m in &self.module_path {
168 write!(f, "{m}::")?;
169 }
170 write!(f, "{}", self.name)
171 }
172}
173
174pub struct Module {
176 pub(crate) name: ModuleId,
177 pub(crate) doc: &'static str,
178 pub(crate) align_functions: HashMap<FunctionId, AlignFunction>,
179 pub(crate) mapping_functions: HashMap<FunctionId, MapFunction>,
180 pub(crate) group_functions: HashMap<FunctionId, GroupFunction>,
181 pub(crate) bucket_functions: HashMap<FunctionId, BucketType>,
182 pub(crate) compute_functions: HashMap<FunctionId, ComputeFunction>,
183 pub(crate) submodules: HashMap<ModuleId, Module>,
184}
185
186impl std::fmt::Display for Module {
187 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188 write!(f, "Module: {}", self.name)
189 }
190}
191
192struct FunctionSerializer<'a, F: FunctionTrait>(&'a F);
193
194impl<F> Serialize for FunctionSerializer<'_, F>
195where
196 F: FunctionTrait,
197{
198 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
199 where
200 S: serde::Serializer,
201 {
202 let mut state = serializer.serialize_struct("Function", 2)?;
203 state.serialize_field("doc", self.0.doc())?;
204 state.serialize_field("args", &self.0.args())?;
205 state.end()
206 }
207}
208
209struct FunctionMapSerializer<'a, F: FunctionTrait>(&'a HashMap<FunctionId, F>);
210
211impl<F> Serialize for FunctionMapSerializer<'_, F>
212where
213 F: FunctionTrait,
214{
215 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
216 where
217 S: serde::Serializer,
218 {
219 let mut state = serializer.serialize_map(Some(self.0.len()))?;
220 for (id, func) in self.0 {
221 state.serialize_entry(&id.to_string(), &FunctionSerializer(func))?;
222 }
223 state.end()
224 }
225}
226
227impl Serialize for Module {
228 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
229 where
230 S: serde::Serializer,
231 {
232 let mut state = serializer.serialize_struct("Module", 8)?;
233 state.serialize_field("name", &self.name)?;
234 state.serialize_field("doc", &self.doc)?;
235 state.serialize_field(
236 "align_functions",
237 &FunctionMapSerializer(&self.align_functions),
238 )?;
239 state.serialize_field(
240 "mapping_functions",
241 &FunctionMapSerializer(&self.mapping_functions),
242 )?;
243 state.serialize_field(
244 "group_functions",
245 &FunctionMapSerializer(&self.group_functions),
246 )?;
247 state.serialize_field(
248 "compute_functions",
249 &FunctionMapSerializer(&self.compute_functions),
250 )?;
251 state.serialize_field(
252 "bucket_functions",
253 &FunctionMapSerializer(&self.bucket_functions),
254 )?;
255 state.serialize_field("submodules", &self.submodules)?;
256 state.end()
257 }
258}
259
260impl Module {
261 pub fn documentation(&self, level: usize) -> Result<String, std::fmt::Error> {
263 let header = "#".repeat(level + 1);
264 let mut functions = String::new();
265 let mut align_functions: Vec<_> = self.align_functions.iter().collect();
266
267 align_functions.sort_by_key(|(i, _)| *i);
268 if !align_functions.is_empty() {
269 writeln!(&mut functions, "#{header} Align Functions")?;
270 }
271 for (n, f) in &align_functions {
272 writeln!(&mut functions, "{}", f.documentation(n))?;
273 writeln!(&mut functions)?;
274 }
275 let mut mapping_functions: Vec<_> = self.mapping_functions.iter().collect();
276 mapping_functions.sort_by_key(|(i, _)| *i);
277 if !mapping_functions.is_empty() {
278 writeln!(&mut functions, "#{header} Map Functions")?;
279 }
280 for (n, f) in &mapping_functions {
281 writeln!(&mut functions, "{}", f.documentation(n))?;
282 writeln!(&mut functions)?;
283 }
284 let mut group_functions: Vec<_> = self.group_functions.iter().collect();
285 group_functions.sort_by_key(|(i, _)| *i);
286 if !group_functions.is_empty() {
287 writeln!(&mut functions, "#{header} Group Functions")?;
288 }
289 for (n, f) in &group_functions {
290 writeln!(&mut functions, "{}", f.documentation(n))?;
291 writeln!(&mut functions)?;
292 }
293 let mut compute_functions: Vec<_> = self.compute_functions.iter().collect();
294 compute_functions.sort_by_key(|(i, _)| *i);
295 if !compute_functions.is_empty() {
296 writeln!(&mut functions, "#{header} Compute Functions")?;
297 }
298 for (n, f) in &compute_functions {
299 writeln!(&mut functions, "{}", f.documentation(n))?;
300 writeln!(&mut functions)?;
301 }
302 let mut bucket_functions: Vec<_> = self.bucket_functions.iter().collect();
303 bucket_functions.sort_by_key(|(i, _)| *i);
304 if !bucket_functions.is_empty() {
305 writeln!(&mut functions, "#{header} Bucket Functions")?;
306 }
307 for (n, f) in &bucket_functions {
308 writeln!(&mut functions, "{}", f.documentation(n))?;
309 writeln!(&mut functions)?;
310 }
311
312 let mut submodule_list: Vec<_> = self.submodules.iter().collect();
313 submodule_list.sort_by_key(|(i, _)| *i);
314 let mut submodules = String::new();
315 for (_, m) in submodule_list {
316 writeln!(&mut submodules, "{}", m.documentation(level + 1)?)?;
317 }
318 Ok(format!(
319 r"{header} {name}
320{doc}
321{functions}{submodules}",
322 name = self.name,
323 doc = self.doc,
324 ))
325 }
326 pub(crate) fn map_fn(&self, id: &Function) -> Option<&MapFunction> {
327 self.map_fn_(&id.module_path, &id.name)
328 }
329 fn map_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&MapFunction> {
330 if let Some((first, rest)) = modules.split_first() {
331 self.submodules.get(first)?.map_fn_(rest, id)
332 } else {
333 self.mapping_functions.get(id)
334 }
335 }
336
337 pub(crate) fn align_fn(&self, id: &Function) -> Option<&AlignFunction> {
338 self.align_fn_(&id.module_path, &id.name)
339 }
340 fn align_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&AlignFunction> {
341 if let Some((first, rest)) = modules.split_first() {
342 self.submodules.get(first)?.align_fn_(rest, id)
343 } else {
344 self.align_functions.get(id)
345 }
346 }
347 pub(crate) fn group_fn(&self, id: &Function) -> Option<&GroupFunction> {
348 self.group_fn_(&id.module_path, &id.name)
349 }
350 fn group_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&GroupFunction> {
351 if let Some((first, rest)) = modules.split_first() {
352 self.submodules.get(first)?.group_fn_(rest, id)
353 } else {
354 self.group_functions.get(id)
355 }
356 }
357
358 pub(crate) fn compute_fn(&self, id: &Function) -> Option<&ComputeFunction> {
359 self.compute_fn_(&id.module_path, &id.name)
360 }
361 fn compute_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&ComputeFunction> {
362 if let Some((first, rest)) = modules.split_first() {
363 self.submodules.get(first)?.compute_fn_(rest, id)
364 } else {
365 self.compute_functions.get(id)
366 }
367 }
368}
369
370pub trait MapFunctionTrait:
372 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
373{
374 #[must_use]
376 fn call(&self, input: &str) -> String;
377 #[must_use]
379 fn box_clone(&self) -> Box<dyn MapFunctionTrait>;
380}
381
382#[derive(Debug, serde::Serialize, serde::Deserialize)]
383#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
384#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
385pub enum MapFunction {
387 Builtin(MapType),
389 #[serde(skip)]
390 UserDefined(Box<dyn MapFunctionTrait>),
392}
393
394impl FunctionTrait for MapFunction {
395 fn doc(&self) -> &str {
396 match self {
397 MapFunction::Builtin(t) => t.doc(),
398 MapFunction::UserDefined(func) => func.doc(),
399 }
400 }
401 fn args(&self) -> Vec<Arg> {
402 match self {
403 MapFunction::Builtin(t) => t.args(),
404 MapFunction::UserDefined(func) => func.args(),
405 }
406 }
407}
408
409impl From<MapType> for MapFunction {
410 fn from(t: MapType) -> Self {
411 MapFunction::Builtin(t)
412 }
413}
414impl std::fmt::Display for MapFunction {
415 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
416 match self {
417 MapFunction::Builtin(t) => write!(f, "{t}"),
418 MapFunction::UserDefined(func) => write!(f, "{func}"),
419 }
420 }
421}
422impl Clone for MapFunction {
423 fn clone(&self) -> Self {
424 match self {
425 MapFunction::Builtin(t) => MapFunction::Builtin(*t),
426 MapFunction::UserDefined(func) => MapFunction::UserDefined(func.box_clone()),
427 }
428 }
429}
430
431pub trait AlignFunctionTrait:
433 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
434{
435 #[must_use]
437 fn call(&self, input: &str) -> String;
438 #[must_use]
440 fn box_clone(&self) -> Box<dyn AlignFunctionTrait>;
441}
442
443#[derive(Debug, serde::Serialize, serde::Deserialize)]
444#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
445#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
446pub enum AlignFunction {
448 Builtin(TimeType),
450 #[serde(skip)]
451 UserDefined(Box<dyn AlignFunctionTrait>),
453}
454impl FunctionTrait for AlignFunction {
455 fn doc(&self) -> &str {
456 match self {
457 AlignFunction::Builtin(t) => t.doc(),
458 AlignFunction::UserDefined(func) => func.doc(),
459 }
460 }
461 fn args(&self) -> Vec<Arg> {
462 match self {
463 AlignFunction::Builtin(t) => t.args(),
464 AlignFunction::UserDefined(func) => func.args(),
465 }
466 }
467}
468impl From<TimeType> for AlignFunction {
469 fn from(t: TimeType) -> Self {
470 AlignFunction::Builtin(t)
471 }
472}
473impl std::fmt::Display for AlignFunction {
474 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
475 match self {
476 AlignFunction::Builtin(t) => write!(f, "{t}"),
477 AlignFunction::UserDefined(func) => write!(f, "{func}"),
478 }
479 }
480}
481impl Clone for AlignFunction {
482 fn clone(&self) -> Self {
483 match self {
484 AlignFunction::Builtin(t) => AlignFunction::Builtin(*t),
485 AlignFunction::UserDefined(func) => AlignFunction::UserDefined(func.box_clone()),
486 }
487 }
488}
489
490pub trait GroupFunctionTrait:
492 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
493{
494 #[must_use]
496 fn call(&self, input: &str) -> String;
497 #[must_use]
499 fn box_clone(&self) -> Box<dyn GroupFunctionTrait>;
500}
501
502#[derive(Debug, serde::Serialize, serde::Deserialize)]
503#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
504#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
505pub enum GroupFunction {
507 Builtin(TagsType),
509 #[serde(skip)]
510 UserDefined(Box<dyn GroupFunctionTrait>),
512}
513impl FunctionTrait for GroupFunction {
514 fn doc(&self) -> &str {
515 match self {
516 GroupFunction::Builtin(t) => t.doc(),
517 GroupFunction::UserDefined(func) => func.doc(),
518 }
519 }
520 fn args(&self) -> Vec<Arg> {
521 match self {
522 GroupFunction::Builtin(t) => t.args(),
523 GroupFunction::UserDefined(func) => func.args(),
524 }
525 }
526}
527impl From<TagsType> for GroupFunction {
528 fn from(t: TagsType) -> Self {
529 GroupFunction::Builtin(t)
530 }
531}
532impl std::fmt::Display for GroupFunction {
533 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
534 match self {
535 GroupFunction::Builtin(t) => write!(f, "{t}"),
536 GroupFunction::UserDefined(func) => write!(f, "{func}"),
537 }
538 }
539}
540impl Clone for GroupFunction {
541 fn clone(&self) -> Self {
542 match self {
543 GroupFunction::Builtin(t) => GroupFunction::Builtin(*t),
544 GroupFunction::UserDefined(func) => GroupFunction::UserDefined(func.box_clone()),
545 }
546 }
547}
548
549pub trait ComputeFunctionTrait:
551 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
552{
553 #[must_use]
555 fn call(&self, input: &str) -> String;
556 #[must_use]
558 fn box_clone(&self) -> Box<dyn ComputeFunctionTrait>;
559}
560#[derive(Debug, serde::Serialize, serde::Deserialize)]
561#[cfg_attr(feature = "wasm", derive(tsify::Tsify))]
562#[cfg_attr(feature = "wasm", tsify(into_wasm_abi, from_wasm_abi))]
563pub enum ComputeFunction {
565 Builtin(ComputeType),
567 #[serde(skip)]
568 UserDefined(Box<dyn ComputeFunctionTrait>),
570}
571impl FunctionTrait for ComputeFunction {
572 fn doc(&self) -> &str {
573 match self {
574 ComputeFunction::Builtin(t) => t.doc(),
575 ComputeFunction::UserDefined(func) => func.doc(),
576 }
577 }
578 fn args(&self) -> Vec<Arg> {
579 match self {
580 ComputeFunction::Builtin(t) => t.args(),
581 ComputeFunction::UserDefined(func) => func.args(),
582 }
583 }
584}
585impl From<ComputeType> for ComputeFunction {
586 fn from(c: ComputeType) -> Self {
587 ComputeFunction::Builtin(c)
588 }
589}
590impl std::fmt::Display for ComputeFunction {
591 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
592 match self {
593 ComputeFunction::Builtin(t) => write!(f, "{t}"),
594 ComputeFunction::UserDefined(func) => write!(f, "{func}"),
595 }
596 }
597}
598impl Clone for ComputeFunction {
599 fn clone(&self) -> Self {
600 match self {
601 ComputeFunction::Builtin(c) => ComputeFunction::Builtin(*c),
602 ComputeFunction::UserDefined(func) => ComputeFunction::UserDefined(func.box_clone()),
603 }
604 }
605}