rune/compile/v1/
scopes.rs

Help
1use core::cell::{Cell, RefCell};
2use core::fmt;
3
4use crate::alloc::prelude::*;
5use crate::alloc::{self, HashMap};
6use crate::ast::Spanned;
7use crate::compile::{self, Assembly, ErrorKind, WithSpan};
8use crate::hir;
9use crate::query::Query;
10use crate::runtime::{Inst, InstAddress, Output};
11use crate::SourceId;
12
13use super::{Address, Any, DisplayNamed, Linear, Slab, Slots};
14
15/// Root scope.
16const ROOT: ScopeId = ScopeId { index: 0, id: 0 };
17
18#[derive(Debug)]
19pub(super) struct Scope<'hir> {
20    /// Parent scope.
21    parent: ScopeId,
22    /// Scope.
23    id: ScopeId,
24    /// Named variables.
25    names: HashMap<hir::Variable, VarInner<'hir>>,
26    /// Slots owned by this scope.
27    locals: Dangling,
28}
29
30impl Scope<'_> {
31    /// Construct a new locals handlers.
32    fn new(parent: ScopeId, id: ScopeId) -> Self {
33        Self {
34            parent,
35            id,
36            names: HashMap::new(),
37            locals: Dangling::default(),
38        }
39    }
40
41    /// Get the parent scope.
42    fn parent(&self) -> Option<ScopeId> {
43        (self.parent != self.id).then_some(self.parent)
44    }
45}
46
47/// A scope handle which does not implement Copy to make it harder to misuse.
48#[must_use = "Scope handles must be handed back to Scopes to be freed"]
49pub(super) struct ScopeHandle {
50    pub(super) id: ScopeId,
51}
52
53/// A scope that has been popped but not freed.
54#[must_use = "Scope handles must be handed back to Scopes to be freed"]
55pub(super) struct DanglingScope {
56    pub(super) id: ScopeId,
57}
58
59/// A guard returned from [push][Scopes::push].
60///
61/// This should be provided to a subsequent [pop][Scopes::pop] to allow it to be
62/// sanity checked.
63#[derive(Clone, Copy, PartialEq, Eq, Hash)]
64#[must_use]
65pub(super) struct ScopeId {
66    index: usize,
67    id: usize,
68}
69
70impl fmt::Display for ScopeId {
71    #[inline]
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        write!(f, "{}:{}", self.index, self.id)
74    }
75}
76
77impl fmt::Debug for ScopeId {
78    #[inline]
79    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80        fmt::Display::fmt(self, f)
81    }
82}
83
84/// Dangling address set which keeps track of addresses used and the order in
85/// which they are inserted.
86#[derive(Default)]
87struct Dangling {
88    addresses: Vec<Option<InstAddress>>,
89    address_to_index: HashMap<InstAddress, usize>,
90}
91
92impl Dangling {
93    fn clear(&mut self) {
94        self.addresses.clear();
95        self.address_to_index.clear();
96    }
97
98    fn insert(&mut self, addr: InstAddress) -> alloc::Result<()> {
99        if self.address_to_index.contains_key(&addr) {
100            return Ok(());
101        }
102
103        self.address_to_index
104            .try_insert(addr, self.addresses.len())?;
105
106        self.addresses.try_push(Some(addr))?;
107        Ok(())
108    }
109
110    fn remove(&mut self, addr: InstAddress) -> bool {
111        if let Some(index) = self.address_to_index.remove(&addr) {
112            self.addresses[index] = None;
113            true
114        } else {
115            false
116        }
117    }
118
119    /// Iterate over addresses.
120    #[inline]
121    fn addresses(&self) -> impl Iterator<Item = InstAddress> + '_ {
122        self.addresses.iter().filter_map(|addr| *addr)
123    }
124}
125
126impl fmt::Debug for Dangling {
127    #[inline]
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        self.addresses.fmt(f)
130    }
131}
132
133pub(crate) struct Scopes<'hir> {
134    scopes: RefCell<Slab<Scope<'hir>>>,
135    /// Set of addresses that are dangling.
136    dangling: RefCell<Dangling>,
137    source_id: SourceId,
138    size: Cell<usize>,
139    slots: RefCell<Slots>,
140    id: Cell<usize>,
141    top: Cell<ScopeId>,
142}
143
144impl<'hir> Scopes<'hir> {
145    /// Construct a new collection of scopes.
146    pub(crate) fn new(source_id: SourceId) -> alloc::Result<Self> {
147        let mut scopes = Slab::new();
148        scopes.insert(Scope::new(ROOT, ROOT))?;
149
150        Ok(Self {
151            scopes: RefCell::new(scopes),
152            dangling: RefCell::new(Dangling::default()),
153            source_id,
154            size: Cell::new(0),
155            slots: RefCell::new(Slots::new()),
156            id: Cell::new(1),
157            top: Cell::new(ROOT),
158        })
159    }
160
161    /// Drain dangling addresses into a vector.
162    pub(crate) fn drain_dangling_into(&self, out: &mut Vec<InstAddress>) -> alloc::Result<()> {
163        let mut dangling = self.dangling.borrow_mut();
164
165        for addr in dangling.addresses.drain(..).flatten() {
166            out.try_push(addr)?;
167        }
168
169        dangling.clear();
170        Ok(())
171    }
172
173    /// Get the maximum total number of variables used in a function.
174    /// Effectively the required stack size.
175    #[inline]
176    pub(crate) fn size(&self) -> usize {
177        self.size.get()
178    }
179
180    /// Get the last scope guard.
181    #[inline]
182    pub(super) fn top_id(&self) -> ScopeId {
183        self.top.get()
184    }
185
186    /// Get the local with the given name.
187    #[tracing::instrument(skip(self, q, span))]
188    pub(super) fn get(
189        &self,
190        q: &mut Query<'_, '_>,
191        span: &dyn Spanned,
192        name: hir::Variable,
193    ) -> compile::Result<Var<'hir>> {
194        let scopes = self.scopes.borrow();
195        let mut current = Some(self.top.get());
196
197        while let Some(id) = current.take() {
198            let Some(scope) = scopes.get(id.index) else {
199                return Err(compile::Error::msg(span, format!("Missing scope {id}")));
200            };
201
202            current = scope.parent();
203
204            let Some(var) = scope.names.get(&name) else {
205                continue;
206            };
207
208            if let Some(_moved_at) = var.moved_at.get() {
209                return Err(compile::Error::new(
210                    span,
211                    ErrorKind::VariableMoved {
212                        #[cfg(feature = "emit")]
213                        moved_at: _moved_at.span(),
214                    },
215                ));
216            }
217
218            q.visitor
219                .visit_variable_use(self.source_id, var.span, span)
220                .with_span(span)?;
221
222            let var = Var {
223                span: var.span,
224                name: var.name,
225                addr: var.addr,
226            };
227
228            tracing::trace!(?scope, ?var);
229            return Ok(var);
230        }
231
232        Err(compile::Error::msg(
233            span,
234            try_format!("Missing variable `{name}`"),
235        ))
236    }
237
238    /// Take the local with the given name.
239    #[tracing::instrument(skip(self, q, span))]
240    pub(super) fn take(
241        &self,
242        q: &mut Query<'_, '_>,
243        span: &'hir dyn Spanned,
244        name: hir::Variable,
245    ) -> compile::Result<Var<'hir>> {
246        let scopes = self.scopes.borrow();
247        let mut current = Some(self.top.get());
248
249        while let Some(id) = current.take() {
250            let Some(scope) = scopes.get(id.index) else {
251                return Err(compile::Error::msg(span, format!("Missing scope {id}")));
252            };
253
254            current = scope.parent();
255
256            let Some(var) = scope.names.get(&name) else {
257                continue;
258            };
259
260            if let Some(_moved_at) = var.moved_at.get() {
261                return Err(compile::Error::new(
262                    span,
263                    ErrorKind::VariableMoved {
264                        #[cfg(feature = "emit")]
265                        moved_at: _moved_at.span(),
266                    },
267                ));
268            }
269
270            q.visitor
271                .visit_variable_use(self.source_id, var.span, span)
272                .with_span(span)?;
273
274            var.moved_at.set(Some(span));
275
276            let var = Var {
277                span: var.span,
278                name: var.name,
279                addr: var.addr,
280            };
281
282            tracing::trace!(?scope, ?var);
283            return Ok(var);
284        }
285
286        Err(compile::Error::msg(
287            span,
288            try_format!("Missing variable `{name}` to take"),
289        ))
290    }
291
292    /// Construct a new variable.
293    #[tracing::instrument(skip(self, span))]
294    pub(super) fn define(
295        &self,
296        span: &'hir dyn Spanned,
297        name: hir::Variable,
298        addr: &Address<'_, 'hir>,
299    ) -> compile::Result<()> {
300        let mut scopes = self.scopes.borrow_mut();
301
302        let Some(scope) = scopes.get_mut(self.top.get().index) else {
303            return Err(compile::Error::msg(span, "Missing top scope"));
304        };
305
306        let var = VarInner {
307            span,
308            name,
309            addr: addr.addr(),
310            moved_at: Cell::new(None),
311        };
312
313        scope.names.try_insert(name, var).with_span(span)?;
314        tracing::trace!(?scope, ?name);
315        Ok(())
316    }
317
318    /// Defer slot allocation.
319    #[tracing::instrument(skip_all)]
320    pub(super) fn defer(&self, span: &'hir dyn Spanned) -> Any<'_, 'hir> {
321        Any::defer(self, self.top.get(), span)
322    }
323
324    /// Explicitly allocate a slot.
325    #[tracing::instrument(skip_all)]
326    pub(super) fn alloc(&self, span: &'hir dyn Spanned) -> compile::Result<Address<'_, 'hir>> {
327        let mut scopes = self.scopes.borrow_mut();
328
329        let Some(scope) = scopes.get_mut(self.top.get().index) else {
330            return Err(compile::Error::msg(span, "Missing top scope"));
331        };
332
333        let addr = InstAddress::new(self.slots.borrow_mut().insert()?);
334        self.size.set(self.size.get().max(addr.offset() + 1));
335        scope.locals.insert(addr).with_span(span)?;
336        self.dangling.borrow_mut().remove(addr);
337
338        tracing::trace!(?scope, ?addr, size = self.size.get());
339        Ok(Address::local(span, self, addr))
340    }
341
342    /// Declare an anonymous variable.
343    #[tracing::instrument(skip(self, span))]
344    pub(super) fn alloc_in(&self, span: &dyn Spanned, id: ScopeId) -> compile::Result<InstAddress> {
345        let mut scopes = self.scopes.borrow_mut();
346
347        let Some(scope) = scopes.get_mut(id.index) else {
348            return Err(compile::Error::msg(
349                span,
350                format!("Missing scope {id} to allocate in"),
351            ));
352        };
353
354        if scope.id != id {
355            return Err(compile::Error::msg(
356                span,
357                try_format!(
358                    "Scope id mismatch, {} (actual) != {id} (expected)",
359                    scope.id
360                ),
361            ));
362        }
363
364        let addr = InstAddress::new(self.slots.borrow_mut().insert()?);
365        scope.locals.insert(addr).with_span(span)?;
366        self.size.set(self.size.get().max(addr.offset() + 1));
367        self.dangling.borrow_mut().remove(addr);
368
369        tracing::trace!(?scope, ?addr, size = self.size.get());
370        Ok(addr)
371    }
372
373    /// Perform a linear allocation.
374    #[tracing::instrument(ret(level = "trace"), skip(self, span))]
375    pub(super) fn linear(
376        &self,
377        span: &'hir dyn Spanned,
378        n: usize,
379    ) -> compile::Result<Linear<'_, 'hir>> {
380        let mut scopes = self.scopes.borrow_mut();
381
382        let Some(scope) = scopes.get_mut(self.top.get().index) else {
383            return Err(compile::Error::msg(span, "Missing top scope"));
384        };
385
386        let mut dangling = self.dangling.borrow_mut();
387
388        let linear = match n {
389            0 => Linear::empty(),
390            1 => {
391                let addr = InstAddress::new(self.slots.borrow_mut().insert()?);
392                scope.locals.insert(addr).with_span(span)?;
393                dangling.remove(addr);
394                self.size.set(self.size.get().max(addr.offset() + 1));
395                Linear::single(Address::local(span, self, addr))
396            }
397            n => {
398                let mut slots = self.slots.borrow_mut();
399                let mut addresses = Vec::try_with_capacity(n).with_span(span)?;
400
401                for _ in 0..n {
402                    let addr = InstAddress::new(slots.push().with_span(span)?);
403                    scope.locals.insert(addr).with_span(span)?;
404                    dangling.remove(addr);
405                    addresses.try_push(Address::dangling(span, self, addr))?;
406                    self.size.set(self.size.get().max(addr.offset() + 1));
407                }
408
409                Linear::new(addresses)
410            }
411        };
412
413        tracing::trace!(?scope, ?linear, size = self.size.get());
414        Ok(linear)
415    }
416
417    /// Free an address if it's in the specified scope.
418    #[tracing::instrument(skip(self, span))]
419    pub(super) fn free_addr(
420        &self,
421        span: &dyn Spanned,
422        addr: InstAddress,
423        name: Option<&'static str>,
424        dangling: bool,
425    ) -> compile::Result<()> {
426        let mut scopes = self.scopes.borrow_mut();
427
428        let Some(scope) = scopes.get_mut(self.top.get().index) else {
429            return Err(compile::Error::msg(
430                span,
431                format!(
432                    "Freed address {} does not have an implicit scope",
433                    DisplayNamed::new(addr, name)
434                ),
435            ));
436        };
437
438        tracing::trace!(?scope);
439
440        if !scope.locals.remove(addr) {
441            return Err(compile::Error::msg(
442                span,
443                format!(
444                    "Freed address {} is not defined in scope {}",
445                    DisplayNamed::new(addr, name),
446                    scope.id
447                ),
448            ));
449        }
450
451        if dangling {
452            self.dangling.borrow_mut().insert(addr).with_span(span)?;
453        }
454
455        let mut slots = self.slots.borrow_mut();
456
457        if !slots.remove(addr.offset()) {
458            return Err(compile::Error::msg(
459                span,
460                format!(
461                    "Freed adddress {} does not have a global slot in scope {}",
462                    DisplayNamed::new(addr, name),
463                    scope.id
464                ),
465            ));
466        }
467
468        Ok(())
469    }
470
471    #[tracing::instrument(skip(self, span, handle), fields(id = ?handle.id))]
472    pub(super) fn pop(&self, span: &dyn Spanned, handle: ScopeHandle) -> compile::Result<()> {
473        let ScopeHandle { id } = handle;
474
475        let Some(mut scope) = self.scopes.borrow_mut().try_remove(id.index) else {
476            return Err(compile::Error::msg(
477                span,
478                format!("Missing scope while expected {id}"),
479            ));
480        };
481
482        if scope.id != id {
483            return Err(compile::Error::msg(
484                span,
485                try_format!(
486                    "Scope id mismatch, {} (actual) != {id} (expected)",
487                    scope.id
488                ),
489            ));
490        }
491
492        tracing::trace!(?scope, "freeing locals");
493
494        let mut slots = self.slots.borrow_mut();
495        let mut dangling = self.dangling.borrow_mut();
496
497        // Free any locally defined variables associated with the scope.
498        for addr in scope.locals.addresses() {
499            if !slots.remove(addr.offset()) {
500                return Err(compile::Error::msg(
501                    span,
502                    format!(
503                        "Address {addr} is not globally allocated in {:?}",
504                        self.slots
505                    ),
506                ));
507            }
508
509            dangling.insert(addr).with_span(span)?;
510        }
511
512        scope.locals.clear();
513        self.top.set(scope.parent);
514        Ok(())
515    }
516
517    /// Pop the last of the scope.
518    #[tracing::instrument(skip(self, span))]
519    pub(super) fn pop_last(&self, span: &dyn Spanned) -> compile::Result<()> {
520        self.pop(span, ScopeHandle { id: ROOT })?;
521        Ok(())
522    }
523
524    /// Construct a new child scope and return its guard.
525    #[tracing::instrument(skip_all)]
526    pub(super) fn child(&self, span: &dyn Spanned) -> compile::Result<ScopeHandle> {
527        let mut scopes = self.scopes.borrow_mut();
528
529        let id = ScopeId {
530            index: scopes.vacant_key(),
531            id: self.id.replace(self.id.get().wrapping_add(1)),
532        };
533
534        let scope = Scope::new(self.top.replace(id), id);
535        tracing::trace!(?scope);
536        scopes.insert(scope).with_span(span)?;
537        Ok(ScopeHandle { id })
538    }
539
540    #[tracing::instrument(skip(self, span, handle), fields(id = ?handle.id))]
541    pub(super) fn dangle(
542        &self,
543        span: &dyn Spanned,
544        handle: ScopeHandle,
545    ) -> compile::Result<DanglingScope> {
546        let ScopeHandle { id } = handle;
547
548        let scopes = self.scopes.borrow();
549
550        let Some(scope) = scopes.get(id.index) else {
551            return Err(compile::Error::msg(
552                span,
553                format!("Missing scope while expected {id}"),
554            ));
555        };
556
557        if scope.id != id {
558            return Err(compile::Error::msg(
559                span,
560                try_format!(
561                    "Scope id mismatch, {} (actual) != {id} (expected)",
562                    scope.id
563                ),
564            ));
565        }
566
567        self.top.set(scope.parent);
568        tracing::trace!(?scope);
569        Ok(DanglingScope { id })
570    }
571
572    /// Push a dangling scope back onto the stack.
573    #[tracing::instrument(skip_all)]
574    pub(super) fn restore(&self, handle: DanglingScope) -> ScopeHandle {
575        let DanglingScope { id } = handle;
576        tracing::trace!(?id);
577        self.top.set(id);
578        ScopeHandle { id }
579    }
580}
581
582/// A locally declared variable, its calculated stack offset and where it was
583/// declared in its source file.
584pub(super) struct Var<'hir> {
585    /// The span where the variable was declared.
586    pub(super) span: &'hir dyn Spanned,
587    /// The name of the variable.
588    name: hir::Variable,
589    /// Address where the variable is currently live.
590    pub(super) addr: InstAddress,
591}
592
593impl fmt::Debug for Var<'_> {
594    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
595        f.debug_struct("Var")
596            .field("span", &self.span.span())
597            .field("name", &self.name)
598            .field("addr", &self.addr)
599            .finish()
600    }
601}
602
603impl fmt::Display for Var<'_> {
604    #[inline]
605    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
606        self.name.fmt(f)
607    }
608}
609
610impl Var<'_> {
611    /// Copy the declared variable.
612    pub(super) fn copy(
613        &self,
614        asm: &mut Assembly,
615        span: &dyn Spanned,
616        comment: Option<&dyn fmt::Display>,
617        out: Output,
618    ) -> compile::Result<()> {
619        asm.push_with_comment(
620            Inst::Copy {
621                addr: self.addr,
622                out,
623            },
624            span,
625            &format_args!("var `{}`{}", self.name, Append("; ", comment)),
626        )
627    }
628
629    /// Move the declared variable.
630    pub(super) fn move_(
631        &self,
632        asm: &mut Assembly,
633        span: &dyn Spanned,
634        comment: Option<&dyn fmt::Display>,
635        out: Output,
636    ) -> compile::Result<()> {
637        asm.push_with_comment(
638            Inst::Move {
639                addr: self.addr,
640                out,
641            },
642            span,
643            &format_args!("var `{}`{}", self.name, Append("; ", comment)),
644        )
645    }
646}
647
648struct Append<P, T>(P, T);
649
650impl<P, T> fmt::Display for Append<P, T>
651where
652    P: fmt::Display,
653    T: Copy + IntoIterator,
654    T::Item: fmt::Display,
655{
656    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
657        for item in self.1 {
658            self.0.fmt(f)?;
659            item.fmt(f)?;
660        }
661
662        Ok(())
663    }
664}
665
666/// A locally declared variable, its calculated stack offset and where it was
667/// declared in its source file.
668struct VarInner<'hir> {
669    /// Token assocaited with the variable.
670    span: &'hir dyn Spanned,
671    /// The name of the variable.
672    name: hir::Variable,
673    /// Offset from the current stack frame.
674    addr: InstAddress,
675    /// Variable has been taken at the given position.
676    moved_at: Cell<Option<&'hir dyn Spanned>>,
677}
678
679impl fmt::Debug for VarInner<'_> {
680    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
681        f.debug_struct("Var")
682            .field("span", &self.span.span())
683            .field("name", &self.name)
684            .field("addr", &self.addr)
685            .field("moved_at", &self.moved_at.get().map(|s| s.span()))
686            .finish()
687    }
688}
rune/compile/v1/scopes.rs

rune/compile/v1/
scopes.rs
