rune/hir/
scopes.rs

1#![allow(unused)]
2
3use core::cell::RefCell;
4use core::fmt;
5use core::num::NonZeroUsize;
6
7use crate::alloc::prelude::*;
8use crate::alloc::{self, BTreeSet, HashMap, Vec};
9use crate::ast::Spanned;
10use crate::compile::error::{MissingScope, PopError};
11use crate::compile::{self, HasSpan};
12use crate::hir;
13use crate::parse::NonZeroId;
14use crate::shared::Gen;
15
16#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
17#[repr(transparent)]
18pub(crate) struct Scope(usize);
19
20/// The kind of a layer.
21#[derive(Default)]
22enum LayerKind {
23    #[default]
24    Default,
25    Loop,
26    Captures,
27}
28
29#[derive(Default)]
30pub(crate) struct Layer<'hir> {
31    /// Scope identifier of the layer.
32    scope: Scope,
33    /// The parent layer.
34    parent: Option<NonZeroUsize>,
35    ///  The kind of the layer.
36    kind: LayerKind,
37    /// Variables defined in this layer.
38    variables: HashMap<hir::Name<'hir>, hir::Variable>,
39    /// Order of variable definitions.
40    order: Vec<hir::Variable>,
41    /// Captures inside of this layer.
42    captures: BTreeSet<(hir::Name<'hir>, hir::Variable)>,
43    /// An optional layer label.
44    label: Option<&'hir str>,
45}
46
47impl<'hir> Layer<'hir> {
48    fn parent(&self) -> Option<usize> {
49        Some(self.parent?.get().wrapping_sub(1))
50    }
51
52    /// Convert layer into variable drop order.
53    #[inline(always)]
54    pub(crate) fn into_drop_order(self) -> impl ExactSizeIterator<Item = hir::Variable> {
55        self.order.into_iter().rev()
56    }
57
58    /// Variables captured by the layer.
59    pub(crate) fn captures(
60        &self,
61    ) -> impl ExactSizeIterator<Item = (hir::Name<'hir>, hir::Variable)> + '_ {
62        self.captures.iter().copied()
63    }
64}
65
66pub(crate) struct Scopes<'hir, 'a> {
67    scope: Scope,
68    scopes: Vec<Layer<'hir>>,
69    gen: &'a Gen,
70}
71
72impl<'hir, 'a> Scopes<'hir, 'a> {
73    /// Root scope.
74    pub const ROOT: Scope = Scope(0);
75
76    #[inline]
77    pub(crate) fn new(gen: &'a Gen) -> alloc::Result<Self> {
78        let mut scopes = Vec::new();
79        scopes.try_push(Layer::default())?;
80
81        Ok(Self {
82            scope: Scopes::ROOT,
83            scopes,
84            gen,
85        })
86    }
87
88    /// Push a scope.
89    pub(crate) fn push(&mut self, label: Option<&'hir str>) -> alloc::Result<()> {
90        self.push_kind(LayerKind::Default, label)
91    }
92
93    /// Push an async block.
94    pub(crate) fn push_captures(&mut self) -> alloc::Result<()> {
95        self.push_kind(LayerKind::Captures, None)
96    }
97
98    /// Push a loop.
99    pub(crate) fn push_loop(&mut self, label: Option<&'hir str>) -> alloc::Result<()> {
100        self.push_kind(LayerKind::Loop, label)
101    }
102
103    fn push_kind(&mut self, kind: LayerKind, label: Option<&'hir str>) -> alloc::Result<()> {
104        let scope = Scope(self.scopes.len());
105
106        let layer = Layer {
107            scope,
108            parent: Some(NonZeroUsize::new(self.scope.0.wrapping_add(1)).expect("ran out of ids")),
109            variables: HashMap::new(),
110            order: Vec::new(),
111            kind,
112            captures: BTreeSet::new(),
113            label,
114        };
115
116        self.scopes.try_push(layer)?;
117        self.scope = scope;
118        Ok(())
119    }
120
121    /// Pop the given scope.
122    #[tracing::instrument(skip_all, fields(?self.scope))]
123    pub(crate) fn pop(&mut self) -> Result<Layer<'hir>, PopError> {
124        let Some(layer) = self.scopes.pop() else {
125            return Err(PopError::MissingScope(self.scope.0));
126        };
127
128        if layer.scope.0 != self.scope.0 {
129            return Err(PopError::MissingScope(self.scope.0));
130        }
131
132        let Some(parent) = layer.parent() else {
133            return Err(PopError::MissingParentScope(self.scope.0));
134        };
135
136        let to = Scope(parent);
137        tracing::trace!(from = ?self.scope, ?to);
138        self.scope = to;
139        Ok(layer)
140    }
141
142    /// Insert a variable with the specified id.
143    #[tracing::instrument(skip_all, fields(?self.scope, ?name))]
144    pub(crate) fn insert(
145        &mut self,
146        name: hir::Name<'hir>,
147        id: hir::Variable,
148        span: &dyn Spanned,
149    ) -> compile::Result<hir::Variable> {
150        tracing::trace!(?self.scope, ?name, "define");
151
152        let Some(layer) = self.scopes.get_mut(self.scope.0) else {
153            return Err(HasSpan::new(span, MissingScope(self.scope.0)).into());
154        };
155
156        layer.variables.try_insert(name, id)?;
157        layer.order.try_push(id)?;
158        Ok(id)
159    }
160
161    /// Define the given variable.
162    #[tracing::instrument(skip_all, fields(?self.scope, ?name))]
163    pub(crate) fn define(
164        &mut self,
165        name: hir::Name<'hir>,
166        span: &dyn Spanned,
167    ) -> compile::Result<hir::Variable> {
168        tracing::trace!(?self.scope, ?name, "define");
169
170        let Some(layer) = self.scopes.get_mut(self.scope.0) else {
171            return Err(HasSpan::new(span, MissingScope(self.scope.0)).into());
172        };
173
174        let id = hir::Variable(self.gen.next());
175        layer.variables.try_insert(name, id)?;
176        layer.order.try_push(id)?;
177        Ok(id)
178    }
179
180    /// Try to lookup the given variable.
181    #[tracing::instrument(skip_all, fields(?self.scope, ?name))]
182    pub(crate) fn get(
183        &mut self,
184        name: hir::Name<'hir>,
185    ) -> alloc::Result<Option<(hir::Variable, Scope)>> {
186        tracing::trace!("get");
187
188        let mut blocks = Vec::new();
189        let mut scope = self.scopes.get(self.scope.0);
190
191        let (scope, id) = 'ok: {
192            loop {
193                let Some(layer) = scope.take() else {
194                    return Ok(None);
195                };
196
197                if let Some(id) = layer.variables.get(&name) {
198                    break 'ok (layer.scope, *id);
199                }
200
201                if let LayerKind::Captures = layer.kind {
202                    blocks.try_push(layer.scope)?;
203                }
204
205                tracing::trace!(parent = ?layer.parent());
206
207                let Some(parent) = layer.parent() else {
208                    return Ok(None);
209                };
210
211                scope = self.scopes.get(parent);
212            }
213        };
214
215        for s in blocks {
216            let Some(layer) = self.scopes.get_mut(s.0) else {
217                continue;
218            };
219
220            layer.captures.try_insert((name, id))?;
221        }
222
223        Ok(Some((id, scope)))
224    }
225
226    /// Walk the loop and construct captures for it.
227    #[tracing::instrument(skip_all, fields(?self.scope, ?label))]
228    pub(crate) fn loop_drop(
229        &self,
230        label: Option<&str>,
231    ) -> alloc::Result<Option<Vec<hir::Variable>>> {
232        let mut captures = Vec::new();
233        let mut scope = self.scopes.get(self.scope.0);
234
235        while let Some(layer) = scope.take() {
236            if let Some(label) = label {
237                if layer.label == Some(label) {
238                    return Ok(Some(captures));
239                }
240            } else if matches!(layer.kind, LayerKind::Loop) {
241                return Ok(Some(captures));
242            }
243
244            captures.try_extend(layer.order.iter().rev().copied())?;
245            tracing::trace!(parent = ?layer.parent());
246
247            let Some(parent) = layer.parent() else {
248                return Ok(None);
249            };
250
251            scope = self.scopes.get(parent);
252        }
253
254        Ok(None)
255    }
256}