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 #[must_use]
151 pub fn name(&self) -> &str {
152 &self.0
153 }
154
155 pub(crate) fn new(name: &str) -> Self {
156 FunctionId(name.to_string())
157 }
158}
159
160impl Borrow<str> for FunctionId {
161 fn borrow(&self) -> &str {
162 &self.0
163 }
164}
165
166pub(crate) struct Function {
167 pub(crate) module_path: Vec<ModuleId>,
168 pub(crate) name: FunctionId,
169}
170
171impl std::fmt::Display for Function {
172 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173 for m in &self.module_path {
174 write!(f, "{m}::")?;
175 }
176 write!(f, "{}", self.name)
177 }
178}
179
180pub struct Module {
182 pub(crate) name: ModuleId,
183 pub(crate) doc: &'static str,
184 pub(crate) align_functions: HashMap<FunctionId, AlignFunction>,
185 pub(crate) mapping_functions: HashMap<FunctionId, MapFunction>,
186 pub(crate) group_functions: HashMap<FunctionId, GroupFunction>,
187 pub(crate) bucket_functions: HashMap<FunctionId, BucketType>,
188 pub(crate) compute_functions: HashMap<FunctionId, ComputeFunction>,
189 pub(crate) submodules: HashMap<ModuleId, Module>,
190}
191
192impl std::fmt::Display for Module {
193 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
194 write!(f, "Module: {}", self.name)
195 }
196}
197
198struct FunctionSerializer<'a, F: FunctionTrait>(&'a F);
199
200impl<F> Serialize for FunctionSerializer<'_, F>
201where
202 F: FunctionTrait,
203{
204 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
205 where
206 S: serde::Serializer,
207 {
208 let mut state = serializer.serialize_struct("Function", 2)?;
209 state.serialize_field("doc", self.0.doc())?;
210 state.serialize_field("args", &self.0.args())?;
211 state.end()
212 }
213}
214
215struct FunctionMapSerializer<'a, F: FunctionTrait>(&'a HashMap<FunctionId, F>);
216
217impl<F> Serialize for FunctionMapSerializer<'_, F>
218where
219 F: FunctionTrait,
220{
221 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
222 where
223 S: serde::Serializer,
224 {
225 let mut state = serializer.serialize_map(Some(self.0.len()))?;
226 for (id, func) in self.0 {
227 state.serialize_entry(&id.to_string(), &FunctionSerializer(func))?;
228 }
229 state.end()
230 }
231}
232
233impl Serialize for Module {
234 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
235 where
236 S: serde::Serializer,
237 {
238 let mut state = serializer.serialize_struct("Module", 8)?;
239 state.serialize_field("name", &self.name)?;
240 state.serialize_field("doc", &self.doc)?;
241 state.serialize_field(
242 "align_functions",
243 &FunctionMapSerializer(&self.align_functions),
244 )?;
245 state.serialize_field(
246 "mapping_functions",
247 &FunctionMapSerializer(&self.mapping_functions),
248 )?;
249 state.serialize_field(
250 "group_functions",
251 &FunctionMapSerializer(&self.group_functions),
252 )?;
253 state.serialize_field(
254 "compute_functions",
255 &FunctionMapSerializer(&self.compute_functions),
256 )?;
257 state.serialize_field(
258 "bucket_functions",
259 &FunctionMapSerializer(&self.bucket_functions),
260 )?;
261 state.serialize_field("submodules", &self.submodules)?;
262 state.end()
263 }
264}
265
266impl Module {
267 pub fn submodule_iter(&self) -> impl Iterator<Item = (&str, &Module)> {
271 self.submodules.iter().map(|(k, v)| (k.0.as_str(), v))
272 }
273
274 #[must_use]
276 pub fn submodule(&self, name: &str) -> Option<&Module> {
277 self.submodules.get(name)
278 }
279
280 pub fn align_function_iter(&self) -> impl Iterator<Item = (&str, &AlignFunction)> {
282 self.align_functions.iter().map(|(k, v)| (k.0.as_str(), v))
283 }
284
285 pub fn mapping_function_iter(&self) -> impl Iterator<Item = (&str, &MapFunction)> {
287 self.mapping_functions
288 .iter()
289 .map(|(k, v)| (k.0.as_str(), v))
290 }
291
292 pub fn group_function_iter(&self) -> impl Iterator<Item = (&str, &GroupFunction)> {
294 self.group_functions.iter().map(|(k, v)| (k.0.as_str(), v))
295 }
296
297 pub fn bucket_function_iter(&self) -> impl Iterator<Item = (&str, &BucketType)> {
299 self.bucket_functions.iter().map(|(k, v)| (k.0.as_str(), v))
300 }
301
302 pub fn compute_function_iter(&self) -> impl Iterator<Item = (&str, &ComputeFunction)> {
304 self.compute_functions
305 .iter()
306 .map(|(k, v)| (k.0.as_str(), v))
307 }
308
309 #[must_use]
311 pub fn align_function(&self, name: &str) -> Option<&AlignFunction> {
312 self.align_functions.get(name)
313 }
314
315 #[must_use]
317 pub fn mapping_function(&self, name: &str) -> Option<&MapFunction> {
318 self.mapping_functions.get(name)
319 }
320
321 #[must_use]
323 pub fn group_function(&self, name: &str) -> Option<&GroupFunction> {
324 self.group_functions.get(name)
325 }
326
327 #[must_use]
329 pub fn bucket_function(&self, name: &str) -> Option<&BucketType> {
330 self.bucket_functions.get(name)
331 }
332
333 #[must_use]
335 pub fn compute_function(&self, name: &str) -> Option<&ComputeFunction> {
336 self.compute_functions.get(name)
337 }
338
339 pub fn documentation(&self, level: usize) -> Result<String, std::fmt::Error> {
341 let header = "#".repeat(level + 1);
342 let mut functions = String::new();
343 let mut align_functions: Vec<_> = self.align_functions.iter().collect();
344
345 align_functions.sort_by_key(|(i, _)| *i);
346 if !align_functions.is_empty() {
347 writeln!(&mut functions, "#{header} Align Functions")?;
348 }
349 for (n, f) in &align_functions {
350 writeln!(&mut functions, "{}", f.documentation(n))?;
351 writeln!(&mut functions)?;
352 }
353 let mut mapping_functions: Vec<_> = self.mapping_functions.iter().collect();
354 mapping_functions.sort_by_key(|(i, _)| *i);
355 if !mapping_functions.is_empty() {
356 writeln!(&mut functions, "#{header} Map Functions")?;
357 }
358 for (n, f) in &mapping_functions {
359 writeln!(&mut functions, "{}", f.documentation(n))?;
360 writeln!(&mut functions)?;
361 }
362 let mut group_functions: Vec<_> = self.group_functions.iter().collect();
363 group_functions.sort_by_key(|(i, _)| *i);
364 if !group_functions.is_empty() {
365 writeln!(&mut functions, "#{header} Group Functions")?;
366 }
367 for (n, f) in &group_functions {
368 writeln!(&mut functions, "{}", f.documentation(n))?;
369 writeln!(&mut functions)?;
370 }
371 let mut compute_functions: Vec<_> = self.compute_functions.iter().collect();
372 compute_functions.sort_by_key(|(i, _)| *i);
373 if !compute_functions.is_empty() {
374 writeln!(&mut functions, "#{header} Compute Functions")?;
375 }
376 for (n, f) in &compute_functions {
377 writeln!(&mut functions, "{}", f.documentation(n))?;
378 writeln!(&mut functions)?;
379 }
380 let mut bucket_functions: Vec<_> = self.bucket_functions.iter().collect();
381 bucket_functions.sort_by_key(|(i, _)| *i);
382 if !bucket_functions.is_empty() {
383 writeln!(&mut functions, "#{header} Bucket Functions")?;
384 }
385 for (n, f) in &bucket_functions {
386 writeln!(&mut functions, "{}", f.documentation(n))?;
387 writeln!(&mut functions)?;
388 }
389
390 let mut submodule_list: Vec<_> = self.submodules.iter().collect();
391 submodule_list.sort_by_key(|(i, _)| *i);
392 let mut submodules = String::new();
393 for (_, m) in submodule_list {
394 writeln!(&mut submodules, "{}", m.documentation(level + 1)?)?;
395 }
396 Ok(format!(
397 r"{header} {name}
398{doc}
399{functions}{submodules}",
400 name = self.name,
401 doc = self.doc,
402 ))
403 }
404 pub(crate) fn map_fn(&self, id: &Function) -> Option<&MapFunction> {
405 self.map_fn_(&id.module_path, &id.name)
406 }
407 fn map_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&MapFunction> {
408 if let Some((first, rest)) = modules.split_first() {
409 self.submodules.get(first)?.map_fn_(rest, id)
410 } else {
411 self.mapping_functions.get(id)
412 }
413 }
414
415 pub(crate) fn align_fn(&self, id: &Function) -> Option<&AlignFunction> {
416 self.align_fn_(&id.module_path, &id.name)
417 }
418 fn align_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&AlignFunction> {
419 if let Some((first, rest)) = modules.split_first() {
420 self.submodules.get(first)?.align_fn_(rest, id)
421 } else {
422 self.align_functions.get(id)
423 }
424 }
425 pub(crate) fn group_fn(&self, id: &Function) -> Option<&GroupFunction> {
426 self.group_fn_(&id.module_path, &id.name)
427 }
428 fn group_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&GroupFunction> {
429 if let Some((first, rest)) = modules.split_first() {
430 self.submodules.get(first)?.group_fn_(rest, id)
431 } else {
432 self.group_functions.get(id)
433 }
434 }
435
436 pub(crate) fn compute_fn(&self, id: &Function) -> Option<&ComputeFunction> {
437 self.compute_fn_(&id.module_path, &id.name)
438 }
439 fn compute_fn_(&self, modules: &[ModuleId], id: &FunctionId) -> Option<&ComputeFunction> {
440 if let Some((first, rest)) = modules.split_first() {
441 self.submodules.get(first)?.compute_fn_(rest, id)
442 } else {
443 self.compute_functions.get(id)
444 }
445 }
446}
447
448pub trait MapFunctionTrait:
450 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
451{
452 #[must_use]
454 fn call(&self, input: &str) -> String;
455 #[must_use]
457 fn box_clone(&self) -> Box<dyn MapFunctionTrait>;
458}
459
460#[derive(Debug, serde::Serialize, serde::Deserialize)]
461pub enum MapFunction {
463 Builtin(MapType),
465 #[serde(skip)]
466 UserDefined(Box<dyn MapFunctionTrait>),
468}
469
470impl FunctionTrait for MapFunction {
471 fn doc(&self) -> &str {
472 match self {
473 MapFunction::Builtin(t) => t.doc(),
474 MapFunction::UserDefined(func) => func.doc(),
475 }
476 }
477 fn args(&self) -> Vec<Arg> {
478 match self {
479 MapFunction::Builtin(t) => t.args(),
480 MapFunction::UserDefined(func) => func.args(),
481 }
482 }
483}
484
485impl From<MapType> for MapFunction {
486 fn from(t: MapType) -> Self {
487 MapFunction::Builtin(t)
488 }
489}
490impl std::fmt::Display for MapFunction {
491 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
492 match self {
493 MapFunction::Builtin(t) => write!(f, "{t}"),
494 MapFunction::UserDefined(func) => write!(f, "{func}"),
495 }
496 }
497}
498impl Clone for MapFunction {
499 fn clone(&self) -> Self {
500 match self {
501 MapFunction::Builtin(t) => MapFunction::Builtin(*t),
502 MapFunction::UserDefined(func) => MapFunction::UserDefined(func.box_clone()),
503 }
504 }
505}
506
507pub trait AlignFunctionTrait:
509 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
510{
511 #[must_use]
513 fn call(&self, input: &str) -> String;
514 #[must_use]
516 fn box_clone(&self) -> Box<dyn AlignFunctionTrait>;
517}
518
519#[derive(Debug, serde::Serialize, serde::Deserialize)]
520pub enum AlignFunction {
522 Builtin(TimeType),
524 #[serde(skip)]
525 UserDefined(Box<dyn AlignFunctionTrait>),
527}
528impl FunctionTrait for AlignFunction {
529 fn doc(&self) -> &str {
530 match self {
531 AlignFunction::Builtin(t) => t.doc(),
532 AlignFunction::UserDefined(func) => func.doc(),
533 }
534 }
535 fn args(&self) -> Vec<Arg> {
536 match self {
537 AlignFunction::Builtin(t) => t.args(),
538 AlignFunction::UserDefined(func) => func.args(),
539 }
540 }
541}
542impl From<TimeType> for AlignFunction {
543 fn from(t: TimeType) -> Self {
544 AlignFunction::Builtin(t)
545 }
546}
547impl std::fmt::Display for AlignFunction {
548 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
549 match self {
550 AlignFunction::Builtin(t) => write!(f, "{t}"),
551 AlignFunction::UserDefined(func) => write!(f, "{func}"),
552 }
553 }
554}
555impl Clone for AlignFunction {
556 fn clone(&self) -> Self {
557 match self {
558 AlignFunction::Builtin(t) => AlignFunction::Builtin(*t),
559 AlignFunction::UserDefined(func) => AlignFunction::UserDefined(func.box_clone()),
560 }
561 }
562}
563
564pub trait GroupFunctionTrait:
566 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
567{
568 #[must_use]
570 fn call(&self, input: &str) -> String;
571 #[must_use]
573 fn box_clone(&self) -> Box<dyn GroupFunctionTrait>;
574}
575
576#[derive(Debug, serde::Serialize, serde::Deserialize)]
577pub enum GroupFunction {
579 Builtin(TagsType),
581 #[serde(skip)]
582 UserDefined(Box<dyn GroupFunctionTrait>),
584}
585impl FunctionTrait for GroupFunction {
586 fn doc(&self) -> &str {
587 match self {
588 GroupFunction::Builtin(t) => t.doc(),
589 GroupFunction::UserDefined(func) => func.doc(),
590 }
591 }
592 fn args(&self) -> Vec<Arg> {
593 match self {
594 GroupFunction::Builtin(t) => t.args(),
595 GroupFunction::UserDefined(func) => func.args(),
596 }
597 }
598}
599impl From<TagsType> for GroupFunction {
600 fn from(t: TagsType) -> Self {
601 GroupFunction::Builtin(t)
602 }
603}
604impl std::fmt::Display for GroupFunction {
605 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
606 match self {
607 GroupFunction::Builtin(t) => write!(f, "{t}"),
608 GroupFunction::UserDefined(func) => write!(f, "{func}"),
609 }
610 }
611}
612impl Clone for GroupFunction {
613 fn clone(&self) -> Self {
614 match self {
615 GroupFunction::Builtin(t) => GroupFunction::Builtin(*t),
616 GroupFunction::UserDefined(func) => GroupFunction::UserDefined(func.box_clone()),
617 }
618 }
619}
620
621pub trait ComputeFunctionTrait:
623 Send + Sync + std::fmt::Debug + std::fmt::Display + FunctionTrait
624{
625 #[must_use]
627 fn call(&self, input: &str) -> String;
628 #[must_use]
630 fn box_clone(&self) -> Box<dyn ComputeFunctionTrait>;
631}
632#[derive(Debug, serde::Serialize, serde::Deserialize)]
633pub enum ComputeFunction {
635 Builtin(ComputeType),
637 #[serde(skip)]
638 UserDefined(Box<dyn ComputeFunctionTrait>),
640}
641impl FunctionTrait for ComputeFunction {
642 fn doc(&self) -> &str {
643 match self {
644 ComputeFunction::Builtin(t) => t.doc(),
645 ComputeFunction::UserDefined(func) => func.doc(),
646 }
647 }
648 fn args(&self) -> Vec<Arg> {
649 match self {
650 ComputeFunction::Builtin(t) => t.args(),
651 ComputeFunction::UserDefined(func) => func.args(),
652 }
653 }
654}
655impl From<ComputeType> for ComputeFunction {
656 fn from(c: ComputeType) -> Self {
657 ComputeFunction::Builtin(c)
658 }
659}
660impl std::fmt::Display for ComputeFunction {
661 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
662 match self {
663 ComputeFunction::Builtin(t) => write!(f, "{t}"),
664 ComputeFunction::UserDefined(func) => write!(f, "{func}"),
665 }
666 }
667}
668impl Clone for ComputeFunction {
669 fn clone(&self) -> Self {
670 match self {
671 ComputeFunction::Builtin(c) => ComputeFunction::Builtin(*c),
672 ComputeFunction::UserDefined(func) => ComputeFunction::UserDefined(func.box_clone()),
673 }
674 }
675}