rune/cli/
visitor.rs

1use crate::alloc::prelude::*;
2use crate::compile::meta;
3use crate::compile::{CompileVisitor, MetaError, MetaRef};
4use crate::{Hash, ItemBuf};
5
6/// Attribute to collect.
7#[derive(Debug, Clone, Copy)]
8pub(super) enum Attribute {
9    /// Do not collect any functions.
10    None,
11    /// Collect `#[test]` functions.
12    Test,
13    /// Collect `#[bench]` functions.
14    Bench,
15}
16
17/// A compile visitor that collects functions with a specific attribute.
18pub(super) struct FunctionVisitor {
19    attribute: Attribute,
20    functions: Vec<(Hash, ItemBuf)>,
21}
22
23impl FunctionVisitor {
24    pub(super) fn new(kind: Attribute) -> Self {
25        Self {
26            attribute: kind,
27            functions: Vec::default(),
28        }
29    }
30
31    /// Convert visitor into test functions.
32    pub(super) fn into_functions(self) -> Vec<(Hash, ItemBuf)> {
33        self.functions
34    }
35}
36
37impl CompileVisitor for FunctionVisitor {
38    fn register_meta(&mut self, meta: MetaRef<'_>) -> Result<(), MetaError> {
39        let type_hash = match (self.attribute, &meta.kind) {
40            (Attribute::Test, meta::Kind::Function { is_test, .. }) if *is_test => meta.hash,
41            (Attribute::Bench, meta::Kind::Function { is_bench, .. }) if *is_bench => meta.hash,
42            _ => return Ok(()),
43        };
44
45        self.functions
46            .try_push((type_hash, meta.item.try_to_owned()?))?;
47        Ok(())
48    }
49}