1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
/*!
Frontend parsers that consume binary and text shaders and load them into [`Module`](super::Module)s.
*/

mod interpolator;
mod type_gen;

#[cfg(feature = "glsl-in")]
pub mod glsl;
#[cfg(feature = "spv-in")]
pub mod spv;
#[cfg(feature = "wgsl-in")]
pub mod wgsl;

use crate::{
    arena::{Arena, Handle, UniqueArena},
    proc::{ResolveContext, ResolveError, TypeResolution},
    FastHashMap,
};
use std::ops;

/// A table of types for an `Arena<Expression>`.
///
/// A front end can use a `Typifier` to get types for an arena's expressions
/// while it is still contributing expressions to it. At any point, you can call
/// [`typifier.grow(expr, arena, ctx)`], where `expr` is a `Handle<Expression>`
/// referring to something in `arena`, and the `Typifier` will resolve the types
/// of all the expressions up to and including `expr`. Then you can write
/// `typifier[handle]` to get the type of any handle at or before `expr`.
///
/// Note that `Typifier` does *not* build an `Arena<Type>` as a part of its
/// usual operation. Ideally, a module's type arena should only contain types
/// actually needed by `Handle<Type>`s elsewhere in the module — functions,
/// variables, [`Compose`] expressions, other types, and so on — so we don't
/// want every little thing that occurs as the type of some intermediate
/// expression to show up there.
///
/// Instead, `Typifier` accumulates a [`TypeResolution`] for each expression,
/// which refers to the `Arena<Type>` in the [`ResolveContext`] passed to `grow`
/// as needed. [`TypeResolution`] is a lightweight representation for
/// intermediate types like this; see its documentation for details.
///
/// If you do need to register a `Typifier`'s conclusion in an `Arena<Type>`
/// (say, for a [`LocalVariable`] whose type you've inferred), you can use
/// [`register_type`] to do so.
///
/// [`typifier.grow(expr, arena)`]: Typifier::grow
/// [`register_type`]: Typifier::register_type
/// [`Compose`]: crate::Expression::Compose
/// [`LocalVariable`]: crate::LocalVariable
#[derive(Debug, Default)]
pub struct Typifier {
    resolutions: Vec<TypeResolution>,
}

impl Typifier {
    pub const fn new() -> Self {
        Typifier {
            resolutions: Vec::new(),
        }
    }

    pub fn reset(&mut self) {
        self.resolutions.clear()
    }

    pub fn get<'a>(
        &'a self,
        expr_handle: Handle<crate::Expression>,
        types: &'a UniqueArena<crate::Type>,
    ) -> &'a crate::TypeInner {
        self.resolutions[expr_handle.index()].inner_with(types)
    }

    /// Add an expression's type to an `Arena<Type>`.
    ///
    /// Add the type of `expr_handle` to `types`, and return a `Handle<Type>`
    /// referring to it.
    ///
    /// # Note
    ///
    /// If you just need a [`TypeInner`] for `expr_handle`'s type, consider
    /// using `typifier[expression].inner_with(types)` instead. Calling
    /// [`TypeResolution::inner_with`] often lets us avoid adding anything to
    /// the arena, which can significantly reduce the number of types that end
    /// up in the final module.
    ///
    /// [`TypeInner`]: crate::TypeInner
    pub fn register_type(
        &self,
        expr_handle: Handle<crate::Expression>,
        types: &mut UniqueArena<crate::Type>,
    ) -> Handle<crate::Type> {
        match self[expr_handle].clone() {
            TypeResolution::Handle(handle) => handle,
            TypeResolution::Value(inner) => {
                types.insert(crate::Type { name: None, inner }, crate::Span::UNDEFINED)
            }
        }
    }

    /// Grow this typifier until it contains a type for `expr_handle`.
    pub fn grow(
        &mut self,
        expr_handle: Handle<crate::Expression>,
        expressions: &Arena<crate::Expression>,
        ctx: &ResolveContext,
    ) -> Result<(), ResolveError> {
        if self.resolutions.len() <= expr_handle.index() {
            for (eh, expr) in expressions.iter().skip(self.resolutions.len()) {
                //Note: the closure can't `Err` by construction
                let resolution = ctx.resolve(expr, |h| Ok(&self.resolutions[h.index()]))?;
                log::debug!("Resolving {:?} = {:?} : {:?}", eh, expr, resolution);
                self.resolutions.push(resolution);
            }
        }
        Ok(())
    }

    /// Recompute the type resolution for `expr_handle`.
    ///
    /// If the type of `expr_handle` hasn't yet been calculated, call
    /// [`grow`](Self::grow) to ensure it is covered.
    ///
    /// In either case, when this returns, `self[expr_handle]` should be an
    /// updated type resolution for `expr_handle`.
    pub fn invalidate(
        &mut self,
        expr_handle: Handle<crate::Expression>,
        expressions: &Arena<crate::Expression>,
        ctx: &ResolveContext,
    ) -> Result<(), ResolveError> {
        if self.resolutions.len() <= expr_handle.index() {
            self.grow(expr_handle, expressions, ctx)
        } else {
            let expr = &expressions[expr_handle];
            //Note: the closure can't `Err` by construction
            let resolution = ctx.resolve(expr, |h| Ok(&self.resolutions[h.index()]))?;
            self.resolutions[expr_handle.index()] = resolution;
            Ok(())
        }
    }
}

impl ops::Index<Handle<crate::Expression>> for Typifier {
    type Output = TypeResolution;
    fn index(&self, handle: Handle<crate::Expression>) -> &Self::Output {
        &self.resolutions[handle.index()]
    }
}

/// Type representing a lexical scope, associating a name to a single variable
///
/// The scope is generic over the variable representation and name representation
/// in order to allow larger flexibility on the frontends on how they might
/// represent them.
type Scope<Name, Var> = FastHashMap<Name, Var>;

/// Structure responsible for managing variable lookups and keeping track of
/// lexical scopes
///
/// The symbol table is generic over the variable representation and its name
/// to allow larger flexibility on the frontends on how they might represent them.
///
/// ```
/// use naga::front::SymbolTable;
///
/// // Create a new symbol table with `u32`s representing the variable
/// let mut symbol_table: SymbolTable<&str, u32> = SymbolTable::default();
///
/// // Add two variables named `var1` and `var2` with 0 and 2 respectively
/// symbol_table.add("var1", 0);
/// symbol_table.add("var2", 2);
///
/// // Check that `var1` exists and is `0`
/// assert_eq!(symbol_table.lookup("var1"), Some(&0));
///
/// // Push a new scope and add a variable to it named `var1` shadowing the
/// // variable of our previous scope
/// symbol_table.push_scope();
/// symbol_table.add("var1", 1);
///
/// // Check that `var1` now points to the new value of `1` and `var2` still
/// // exists with its value of `2`
/// assert_eq!(symbol_table.lookup("var1"), Some(&1));
/// assert_eq!(symbol_table.lookup("var2"), Some(&2));
///
/// // Pop the scope
/// symbol_table.pop_scope();
///
/// // Check that `var1` now refers to our initial variable with value `0`
/// assert_eq!(symbol_table.lookup("var1"), Some(&0));
/// ```
///
/// Scopes are ordered as a LIFO stack so a variable defined in a later scope
/// with the same name as another variable defined in a earlier scope will take
/// precedence in the lookup. Scopes can be added with [`push_scope`] and
/// removed with [`pop_scope`].
///
/// A root scope is added when the symbol table is created and must always be
/// present. Trying to pop it will result in a panic.
///
/// Variables can be added with [`add`] and looked up with [`lookup`]. Adding a
/// variable will do so in the currently active scope and as mentioned
/// previously a lookup will search from the current scope to the root scope.
///
/// [`push_scope`]: Self::push_scope
/// [`pop_scope`]: Self::push_scope
/// [`add`]: Self::add
/// [`lookup`]: Self::lookup
pub struct SymbolTable<Name, Var> {
    /// Stack of lexical scopes. Not all scopes are active; see [`cursor`].
    ///
    /// [`cursor`]: Self::cursor
    scopes: Vec<Scope<Name, Var>>,
    /// Limit of the [`scopes`] stack (exclusive). By using a separate value for
    /// the stack length instead of `Vec`'s own internal length, the scopes can
    /// be reused to cache memory allocations.
    ///
    /// [`scopes`]: Self::scopes
    cursor: usize,
}

impl<Name, Var> SymbolTable<Name, Var> {
    /// Adds a new lexical scope.
    ///
    /// All variables declared after this point will be added to this scope
    /// until another scope is pushed or [`pop_scope`] is called, causing this
    /// scope to be removed along with all variables added to it.
    ///
    /// [`pop_scope`]: Self::pop_scope
    pub fn push_scope(&mut self) {
        // If the cursor is equal to the scope's stack length then we need to
        // push another empty scope. Otherwise we can reuse the already existing
        // scope.
        if self.scopes.len() == self.cursor {
            self.scopes.push(FastHashMap::default())
        } else {
            self.scopes[self.cursor].clear();
        }

        self.cursor += 1;
    }

    /// Removes the current lexical scope and all its variables
    ///
    /// # PANICS
    /// - If the current lexical scope is the root scope
    pub fn pop_scope(&mut self) {
        // Despite the method title, the variables are only deleted when the
        // scope is reused. This is because while a clear is inevitable if the
        // scope needs to be reused, there are cases where the scope might be
        // popped and not reused, i.e. if another scope with the same nesting
        // level is never pushed again.
        assert!(self.cursor != 1, "Tried to pop the root scope");

        self.cursor -= 1;
    }
}

impl<Name, Var> SymbolTable<Name, Var>
where
    Name: std::hash::Hash + Eq,
{
    /// Perform a lookup for a variable named `name`.
    ///
    /// As stated in the struct level documentation the lookup will proceed from
    /// the current scope to the root scope, returning `Some` when a variable is
    /// found or `None` if there doesn't exist a variable with `name` in any
    /// scope.
    pub fn lookup<Q: ?Sized>(&self, name: &Q) -> Option<&Var>
    where
        Name: std::borrow::Borrow<Q>,
        Q: std::hash::Hash + Eq,
    {
        // Iterate backwards trough the scopes and try to find the variable
        for scope in self.scopes[..self.cursor].iter().rev() {
            if let Some(var) = scope.get(name) {
                return Some(var);
            }
        }

        None
    }

    /// Adds a new variable to the current scope.
    ///
    /// Returns the previous variable with the same name in this scope if it
    /// exists, so that the frontend might handle it in case variable shadowing
    /// is disallowed.
    pub fn add(&mut self, name: Name, var: Var) -> Option<Var> {
        self.scopes[self.cursor - 1].insert(name, var)
    }

    /// Adds a new variable to the root scope.
    ///
    /// This is used in GLSL for builtins which aren't known in advance and only
    /// when used for the first time, so there must be a way to add those
    /// declarations to the root unconditionally from the current scope.
    ///
    /// Returns the previous variable with the same name in the root scope if it
    /// exists, so that the frontend might handle it in case variable shadowing
    /// is disallowed.
    pub fn add_root(&mut self, name: Name, var: Var) -> Option<Var> {
        self.scopes[0].insert(name, var)
    }
}

impl<Name, Var> Default for SymbolTable<Name, Var> {
    /// Constructs a new symbol table with a root scope
    fn default() -> Self {
        Self {
            scopes: vec![FastHashMap::default()],
            cursor: 1,
        }
    }
}

use std::fmt;

impl<Name: fmt::Debug, Var: fmt::Debug> fmt::Debug for SymbolTable<Name, Var> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("SymbolTable ")?;
        f.debug_list()
            .entries(self.scopes[..self.cursor].iter())
            .finish()
    }
}