hugr_model/v0/scope/
vars.rs

1use indexmap::IndexSet;
2use rustc_hash::FxHasher;
3use std::hash::BuildHasherDefault;
4use thiserror::Error;
5
6use crate::v0::table::{NodeId, VarId};
7
8type FxIndexSet<K> = IndexSet<K, BuildHasherDefault<FxHasher>>;
9
10/// Table for keeping track of node parameters.
11///
12/// Variables refer to the parameters of a node which introduces a symbol.
13/// Variables have an associated name and are scoped via nodes. The types of
14/// parameters of a node may only refer to earlier parameters in the same node
15/// in the order they are defined. A variable name must be unique within a
16/// single node. Each node that introduces a symbol introduces a new isolated
17/// scope for variables.
18///
19/// # Examples
20///
21/// ```
22/// # pub use hugr_model::v0::table::{NodeId, VarId};
23/// # pub use hugr_model::v0::scope::VarTable;
24/// let mut vars = VarTable::new();
25/// vars.enter(NodeId(0));
26/// vars.insert("foo").unwrap();
27/// assert_eq!(vars.resolve("foo").unwrap(), VarId(NodeId(0), 0));
28/// assert!(!vars.is_visible(VarId(NodeId(0), 1)));
29/// vars.insert("bar").unwrap();
30/// assert!(vars.is_visible(VarId(NodeId(0), 1)));
31/// assert_eq!(vars.resolve("bar").unwrap(), VarId(NodeId(0), 1));
32/// vars.enter(NodeId(1));
33/// assert!(vars.resolve("foo").is_err());
34/// assert!(!vars.is_visible(VarId(NodeId(0), 0)));
35/// vars.exit();
36/// assert_eq!(vars.resolve("foo").unwrap(), VarId(NodeId(0), 0));
37/// assert!(vars.is_visible(VarId(NodeId(0), 0)));
38/// ```
39#[derive(Debug, Clone, Default)]
40pub struct VarTable<'a> {
41    /// The set of variables in the currently active node and all its parent nodes.
42    ///
43    /// The order in this index set is the order in which variables were added to the table.
44    /// This is used to efficiently remove all variables from the current node when exiting a scope.
45    vars: FxIndexSet<(NodeId, &'a str)>,
46    /// The stack of scopes that are currently open.
47    scopes: Vec<VarScope>,
48}
49
50impl<'a> VarTable<'a> {
51    /// Create a new empty variable table.
52    #[must_use]
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Enter a new scope for the given node.
58    pub fn enter(&mut self, node: NodeId) {
59        self.scopes.push(VarScope {
60            node,
61            var_count: 0,
62            var_stack: self.vars.len(),
63        });
64    }
65
66    /// Exit a previously entered scope.
67    ///
68    /// # Panics
69    ///
70    /// Panics if there are no open scopes.
71    pub fn exit(&mut self) {
72        let scope = self.scopes.pop().unwrap();
73        self.vars.drain(scope.var_stack..);
74    }
75
76    /// Resolve a variable name to a node and variable index.
77    ///
78    /// # Errors
79    ///
80    /// Returns an error if the variable is not defined in the current scope.
81    pub fn resolve(&self, name: &'a str) -> Result<VarId, UnknownVarError<'a>> {
82        let scope = self.scopes.last().ok_or(UnknownVarError::Root(name))?;
83        let set_index = self
84            .vars
85            .get_index_of(&(scope.node, name))
86            .ok_or(UnknownVarError::WithinNode(scope.node, name))?;
87        let var_index = (set_index - scope.var_stack) as u16;
88        Ok(VarId(scope.node, var_index))
89    }
90
91    /// Check if a variable is visible in the current scope.
92    #[must_use]
93    pub fn is_visible(&self, var: VarId) -> bool {
94        let Some(scope) = self.scopes.last() else {
95            return false;
96        };
97
98        scope.node == var.0 && var.1 < scope.var_count
99    }
100
101    /// Insert a new variable into the current scope.
102    ///
103    /// # Errors
104    ///
105    /// Returns an error if the variable is already defined in the current scope.
106    ///
107    /// # Panics
108    ///
109    /// Panics if there are no open scopes.
110    pub fn insert(&mut self, name: &'a str) -> Result<VarId, DuplicateVarError<'a>> {
111        let scope = self.scopes.last_mut().unwrap();
112        let inserted = self.vars.insert((scope.node, name));
113
114        if !inserted {
115            return Err(DuplicateVarError(scope.node, name));
116        }
117
118        let var_index = scope.var_count;
119        scope.var_count += 1;
120        Ok(VarId(scope.node, var_index))
121    }
122
123    /// Reset the variable table to an empty state while preserving the allocations.
124    pub fn clear(&mut self) {
125        self.vars.clear();
126        self.scopes.clear();
127    }
128}
129
130#[derive(Debug, Clone)]
131struct VarScope {
132    /// The node that introduces this scope.
133    node: NodeId,
134    /// The number of variables in this scope.
135    var_count: u16,
136    /// The length of `VarTable::vars` when the scope was opened.
137    var_stack: usize,
138}
139
140/// Error that occurs when a node defines two parameters with the same name.
141#[derive(Debug, Clone, Error)]
142#[error("node {0} already has a variable named `{1}`")]
143pub struct DuplicateVarError<'a>(NodeId, &'a str);
144
145/// Error that occurs when a variable is not defined in the current scope.
146#[derive(Debug, Clone, Error)]
147pub enum UnknownVarError<'a> {
148    /// Failed to resolve a variable when in scope of a node.
149    #[error("can not resolve variable `{1}` in node {0}")]
150    WithinNode(NodeId, &'a str),
151    /// Failed to resolve a variable when in the root scope.
152    #[error("can not resolve variable `{0}` in the root scope")]
153    Root(&'a str),
154}