vegafusion_core/task_graph/
scope.rs

1use crate::error::{Result, ResultWithContext, VegaFusionError};
2use crate::expression::supported::BUILT_IN_SIGNALS;
3use crate::proto::gen::tasks::{Variable, VariableNamespace};
4use std::collections::{HashMap, HashSet};
5
6#[derive(Clone, Debug, Default)]
7pub struct TaskScope {
8    pub signals: HashSet<String>,
9    pub data: HashSet<String>,
10    pub scales: HashSet<String>,
11    /// Tasks that have definition output variables (e.g. task that produces a signal)
12    pub output_var_defs: HashMap<Variable, Variable>,
13    pub children: Vec<TaskScope>,
14}
15
16impl TaskScope {
17    pub fn new() -> Self {
18        Self {
19            signals: Default::default(),
20            data: Default::default(),
21            output_var_defs: Default::default(),
22            scales: Default::default(),
23            children: Default::default(),
24        }
25    }
26
27    pub fn get_child(&self, scope: &[u32]) -> Result<&TaskScope> {
28        let mut child = self;
29        for index in scope {
30            child = child
31                .children
32                .get(*index as usize)
33                .with_context(|| format!("No group with scope {scope:?} found"))?;
34        }
35        Ok(child)
36    }
37
38    pub fn get_child_mut(&mut self, scope: &[u32]) -> Result<&mut TaskScope> {
39        let mut child = self;
40        for index in scope {
41            child = child
42                .children
43                .get_mut(*index as usize)
44                .with_context(|| format!("No group with scope {scope:?} found"))?;
45        }
46        Ok(child)
47    }
48
49    pub fn add_variable(&mut self, variable: &Variable, scope: &[u32]) -> Result<()> {
50        let child = self.get_child_mut(scope)?;
51
52        match variable.ns() {
53            VariableNamespace::Signal => {
54                child.signals.insert(variable.name.clone());
55            }
56            VariableNamespace::Data => {
57                child.data.insert(variable.name.clone());
58            }
59            VariableNamespace::Scale => {
60                child.scales.insert(variable.name.clone());
61            }
62        }
63
64        Ok(())
65    }
66
67    pub fn add_data_signal(&mut self, data: &str, signal: &str, scope: &[u32]) -> Result<()> {
68        let child = self.get_child_mut(scope)?;
69        child
70            .output_var_defs
71            .insert(Variable::new_signal(signal), Variable::new_data(data));
72        Ok(())
73    }
74
75    pub fn remove_data_signal(&mut self, signal: &str, scope: &[u32]) -> Result<Variable> {
76        let child = self.get_child_mut(scope)?;
77        child
78            .output_var_defs
79            .remove(&Variable::new_signal(signal))
80            .with_context(|| format!("No data signal named: {signal}"))
81    }
82
83    pub fn resolve_scope(&self, variable: &Variable, usage_scope: &[u32]) -> Result<Resolved> {
84        // Search for matching variable, start with full usage scope, then iterate up
85        for level in (0..=usage_scope.len()).rev() {
86            let curr_scope = &usage_scope[0..level];
87            let task_scope = self.get_child(curr_scope)?;
88
89            let found_it = match variable.ns() {
90                VariableNamespace::Signal => task_scope.signals.contains(&variable.name),
91                VariableNamespace::Data => task_scope.data.contains(&variable.name),
92                VariableNamespace::Scale => task_scope.scales.contains(&variable.name),
93            };
94            if found_it {
95                // Found it in the regular signal/data/scale
96                return Ok(Resolved {
97                    var: variable.clone(),
98                    scope: Vec::from(curr_scope),
99                    output_var: None,
100                });
101            }
102
103            // Check for output variable
104            if let Some(main_var) = task_scope.output_var_defs.get(variable) {
105                return Ok(Resolved {
106                    var: main_var.clone(),
107                    scope: Vec::from(curr_scope),
108                    output_var: Some(variable.clone()),
109                });
110            }
111
112            // Check for built-in signal
113            if matches!(variable.ns(), VariableNamespace::Signal)
114                && BUILT_IN_SIGNALS.contains(variable.name.as_str())
115            {
116                return Ok(Resolved {
117                    var: variable.clone(),
118                    scope: Vec::new(),
119                    output_var: None,
120                });
121            }
122        }
123
124        // Didn't find it
125        Err(VegaFusionError::internal(format!(
126            "Failed to resolve variable {variable:?} used in scope {usage_scope:?}"
127        )))
128    }
129}
130
131pub struct Resolved {
132    pub var: Variable,
133    pub scope: Vec<u32>,
134    pub output_var: Option<Variable>,
135}
vegafusion_core/task_graph/scope.rs

vegafusion_core/task_graph/
scope.rs
