rune/compile/v1/
needs.rs

1use core::fmt;
2use core::mem::replace;
3
4use crate::ast::Spanned;
5use crate::compile;
6use crate::runtime::{Inst, InstAddress, Output};
7use crate::shared::{rune_diagnose, Backtrace};
8
9use super::{Ctxt, DisplayNamed, ScopeId, Scopes};
10
11/// Trait used to abstract interactions over different needs.
12pub(super) trait Needs<'a, 'hir> {
13    /// Access the span for the needs.
14    fn span(&self) -> &'hir dyn Spanned;
15
16    /// Get output of the needs or error.
17    fn output(&self) -> compile::Result<Output>;
18
19    /// Get the need as an output.
20    ///
21    /// Returns `None` if `Needs::None` is set.
22    fn try_alloc_output(&mut self) -> compile::Result<Option<Output>>;
23
24    fn assign_addr(
25        &mut self,
26        cx: &mut Ctxt<'_, 'hir, '_>,
27        from: InstAddress,
28    ) -> compile::Result<()>;
29
30    /// Allocate an output falling back to discarding if one is not available.
31    fn alloc_output(&mut self) -> compile::Result<Output>;
32
33    /// Try to allocate an address from this needs.
34    fn try_alloc_addr(&mut self) -> compile::Result<Option<&mut Address<'a, 'hir>>>;
35
36    /// Get the need as an address without trying to allocate it if it's
37    /// missing.
38    fn try_as_addr(&self) -> compile::Result<Option<&Address<'a, 'hir>>>;
39
40    /// Get the populated address of the need.
41    fn addr(&self) -> compile::Result<&Address<'a, 'hir>> {
42        let Some(addr) = self.try_as_addr()? else {
43            return Err(compile::Error::msg(
44                self.span(),
45                "Expected need to be populated",
46            ));
47        };
48
49        Ok(addr)
50    }
51}
52
53impl<'a, 'hir> Needs<'a, 'hir> for Any<'a, 'hir> {
54    #[inline]
55    fn span(&self) -> &'hir dyn Spanned {
56        self.span
57    }
58
59    #[inline]
60    fn output(&self) -> compile::Result<Output> {
61        Any::output(self)
62    }
63
64    #[inline]
65    fn assign_addr(
66        &mut self,
67        cx: &mut Ctxt<'_, 'hir, '_>,
68        from: InstAddress,
69    ) -> compile::Result<()> {
70        Any::assign_addr(self, cx, from)
71    }
72
73    #[inline]
74    fn alloc_output(&mut self) -> compile::Result<Output> {
75        Any::alloc_output(self)
76    }
77
78    #[inline]
79    fn try_alloc_output(&mut self) -> compile::Result<Option<Output>> {
80        Any::try_alloc_output(self)
81    }
82
83    #[inline]
84    fn try_alloc_addr(&mut self) -> compile::Result<Option<&mut Address<'a, 'hir>>> {
85        Any::try_alloc_addr(self)
86    }
87
88    #[inline]
89    fn try_as_addr(&self) -> compile::Result<Option<&Address<'a, 'hir>>> {
90        Any::try_as_addr(self)
91    }
92}
93
94impl<'a, 'hir> Needs<'a, 'hir> for Address<'a, 'hir> {
95    #[inline]
96    fn span(&self) -> &'hir dyn Spanned {
97        self.span
98    }
99
100    #[inline]
101    fn output(&self) -> compile::Result<Output> {
102        Ok(Address::output(self))
103    }
104
105    #[inline]
106    fn assign_addr(
107        &mut self,
108        cx: &mut Ctxt<'_, 'hir, '_>,
109        from: InstAddress,
110    ) -> compile::Result<()> {
111        Address::assign_addr(self, cx, from)
112    }
113
114    #[inline]
115    fn alloc_output(&mut self) -> compile::Result<Output> {
116        Address::alloc_output(self)
117    }
118
119    #[inline]
120    fn try_alloc_output(&mut self) -> compile::Result<Option<Output>> {
121        Ok(Some(Address::alloc_output(self)?))
122    }
123
124    #[inline]
125    fn try_alloc_addr(&mut self) -> compile::Result<Option<&mut Address<'a, 'hir>>> {
126        Ok(Some(Address::alloc_addr(self)?))
127    }
128
129    #[inline]
130    fn try_as_addr(&self) -> compile::Result<Option<&Address<'a, 'hir>>> {
131        Ok(Some(self))
132    }
133}
134
135#[derive(Clone, Copy)]
136pub(super) enum AddressKind {
137    /// The value is locally allocated and should be freed in the immediate
138    /// scope.
139    Local,
140    /// The slot has been reserved, but has not been assigned to anything yet.
141    Dangling,
142    /// The address is assigned from elsewhere and *should not* be touched.
143    Assigned,
144    /// The address is allocated on behalf of the given scope, and we should
145    /// defer deallocating it until the given scope is deallocated.
146    Scope(ScopeId),
147    /// The address has been freed.
148    Freed,
149}
150
151impl fmt::Display for AddressKind {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        match self {
154            AddressKind::Local => write!(f, "local"),
155            AddressKind::Dangling => write!(f, "dangling"),
156            AddressKind::Assigned => write!(f, "assigned"),
157            AddressKind::Scope(scope) => write!(f, "scope({scope})"),
158            AddressKind::Freed => write!(f, "freed"),
159        }
160    }
161}
162
163#[must_use = "An allocated address must be freed"]
164pub(super) struct Address<'a, 'hir> {
165    span: &'hir dyn Spanned,
166    scopes: &'a Scopes<'hir>,
167    address: InstAddress,
168    kind: AddressKind,
169    /// A diagnostical name for the address.
170    name: Option<&'static str>,
171    #[cfg_attr(not(feature = "tracing"), allow(unused))]
172    backtrace: Backtrace,
173}
174
175impl<'a, 'hir> Address<'a, 'hir> {
176    /// A locally allocated address.
177    #[inline]
178    #[track_caller]
179    pub(super) fn local(
180        span: &'hir dyn Spanned,
181        scopes: &'a Scopes<'hir>,
182        addr: InstAddress,
183    ) -> Self {
184        Self {
185            span,
186            scopes,
187            address: addr,
188            kind: AddressKind::Local,
189            name: None,
190            backtrace: Backtrace::capture(),
191        }
192    }
193
194    /// A locally assigned address.
195    #[inline]
196    #[track_caller]
197    pub(super) fn assigned(
198        span: &'hir dyn Spanned,
199        scopes: &'a Scopes<'hir>,
200        addr: InstAddress,
201    ) -> Self {
202        Self {
203            span,
204            scopes,
205            address: addr,
206            kind: AddressKind::Assigned,
207            name: None,
208            backtrace: Backtrace::capture(),
209        }
210    }
211
212    /// A locally reserved address.
213    #[inline]
214    #[track_caller]
215    pub(super) fn dangling(
216        span: &'hir dyn Spanned,
217        scopes: &'a Scopes<'hir>,
218        addr: InstAddress,
219    ) -> Self {
220        Self {
221            span,
222            scopes,
223            address: addr,
224            kind: AddressKind::Dangling,
225            name: None,
226            backtrace: Backtrace::capture(),
227        }
228    }
229
230    /// Assign a name to the address.
231    #[inline]
232    pub(super) fn with_name(mut self, name: &'static str) -> Self {
233        self.name = Some(name);
234        self
235    }
236
237    #[inline]
238    pub(super) fn addr(&self) -> InstAddress {
239        self.address
240    }
241
242    #[inline]
243    pub(super) fn alloc_addr(&mut self) -> compile::Result<&mut Self> {
244        if matches!(self.kind, AddressKind::Dangling) {
245            self.kind = AddressKind::Local;
246        }
247
248        Ok(self)
249    }
250
251    #[inline]
252    pub(super) fn alloc_output(&mut self) -> compile::Result<Output> {
253        Ok(self.alloc_addr()?.output())
254    }
255
256    #[inline]
257    pub(super) fn output(&self) -> Output {
258        self.address.output()
259    }
260
261    pub(super) fn assign_addr(
262        &self,
263        cx: &mut Ctxt<'_, '_, '_>,
264        from: InstAddress,
265    ) -> compile::Result<()> {
266        if from != self.address {
267            cx.asm.push(
268                Inst::Copy {
269                    addr: from,
270                    out: self.address.output(),
271                },
272                self.span,
273            )?;
274        }
275
276        Ok(())
277    }
278
279    /// Forget the current address.
280    ///
281    /// This will be freed when the scope it's associated with is freed.
282    pub(super) fn forget(mut self) -> compile::Result<()> {
283        self.kind = AddressKind::Freed;
284        Ok(())
285    }
286
287    pub(super) fn free(self) -> compile::Result<()> {
288        self.free_inner(true)
289    }
290
291    pub(super) fn free_non_dangling(self) -> compile::Result<()> {
292        self.free_inner(false)
293    }
294
295    /// Free the current address.
296    fn free_inner(mut self, dangling: bool) -> compile::Result<()> {
297        match replace(&mut self.kind, AddressKind::Freed) {
298            AddressKind::Local | AddressKind::Dangling => {
299                self.scopes
300                    .free_addr(self.span, self.address, self.name, dangling)?;
301            }
302            AddressKind::Scope(scope) => {
303                if self.scopes.top_id() == scope {
304                    self.scopes
305                        .free_addr(self.span, self.address, self.name, dangling)?;
306                }
307            }
308            AddressKind::Freed => {
309                return Err(compile::Error::msg(
310                    self.span,
311                    "Address has already been freed",
312                ));
313            }
314            _ => {}
315        }
316
317        Ok(())
318    }
319}
320
321impl fmt::Display for Address<'_, '_> {
322    #[inline]
323    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
324        write!(
325            f,
326            "{} address {}",
327            self.kind,
328            DisplayNamed::new(self.address, self.name)
329        )
330    }
331}
332
333impl fmt::Debug for Address<'_, '_> {
334    #[inline]
335    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
336        fmt::Display::fmt(self, f)
337    }
338}
339
340impl Drop for Address<'_, '_> {
341    fn drop(&mut self) {
342        if matches!(self.kind, AddressKind::Freed) {
343            return;
344        }
345
346        rune_diagnose!("{self} was not freed:\nallocated at:\n{}", self.backtrace);
347    }
348}
349
350/// The kind of a needs.
351enum AnyKind<'a, 'hir> {
352    Defer {
353        scopes: &'a Scopes<'hir>,
354        scope: ScopeId,
355        name: Option<&'static str>,
356    },
357    Address {
358        address: Address<'a, 'hir>,
359    },
360    Ignore {
361        #[allow(unused)]
362        name: Option<&'static str>,
363    },
364    Freed,
365}
366
367impl fmt::Display for AnyKind<'_, '_> {
368    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
369        match self {
370            AnyKind::Defer { scope, name, .. } => {
371                write!(f, "defer({})", DisplayNamed::new(scope, *name))
372            }
373            AnyKind::Address { address } => address.fmt(f),
374            AnyKind::Ignore { name } => DisplayNamed::new("ignore", *name).fmt(f),
375            AnyKind::Freed => write!(f, "freed"),
376        }
377    }
378}
379
380impl fmt::Debug for AnyKind<'_, '_> {
381    #[inline]
382    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
383        fmt::Display::fmt(self, f)
384    }
385}
386
387/// A needs hint for an expression.
388/// This is used to contextually determine what an expression is expected to
389/// produce.
390#[must_use = "A need must be freed when its no longer in use"]
391pub(super) struct Any<'a, 'hir> {
392    span: &'hir dyn Spanned,
393    kind: AnyKind<'a, 'hir>,
394    #[cfg_attr(not(feature = "tracing"), allow(unused))]
395    backtrace: Backtrace,
396}
397
398impl<'a, 'hir> Any<'a, 'hir> {
399    /// A needs that should be ignored.
400    #[track_caller]
401    pub(super) fn ignore(span: &'hir dyn Spanned) -> Self {
402        Self {
403            span,
404            kind: AnyKind::Ignore { name: None },
405            backtrace: Backtrace::capture(),
406        }
407    }
408
409    /// Defer allocation of a slot until it is requested.
410    #[track_caller]
411    pub(super) fn defer(scopes: &'a Scopes<'hir>, scope: ScopeId, span: &'hir dyn Spanned) -> Self {
412        Self {
413            span,
414            kind: AnyKind::Defer {
415                scopes,
416                scope,
417                name: None,
418            },
419            backtrace: Backtrace::capture(),
420        }
421    }
422
423    /// An assigned address.
424    #[track_caller]
425    pub(super) fn assigned(
426        span: &'hir dyn Spanned,
427        scopes: &'a Scopes<'hir>,
428        addr: InstAddress,
429    ) -> Self {
430        Self {
431            span,
432            kind: AnyKind::Address {
433                address: Address::assigned(span, scopes, addr),
434            },
435            backtrace: Backtrace::capture(),
436        }
437    }
438
439    /// Assign a name to the request.
440    pub(super) fn with_name(mut self, new_name: &'static str) -> Self {
441        match &mut self.kind {
442            AnyKind::Defer { name, .. } => {
443                *name = Some(new_name);
444            }
445            AnyKind::Address { address } => {
446                address.name = Some(new_name);
447            }
448            AnyKind::Ignore { name, .. } => {
449                *name = Some(new_name);
450            }
451            AnyKind::Freed => {}
452        };
453
454        self
455    }
456
457    pub(super) fn assign_addr(
458        &mut self,
459        cx: &mut Ctxt<'_, 'hir, '_>,
460        from: InstAddress,
461    ) -> compile::Result<()> {
462        match &self.kind {
463            AnyKind::Defer { scopes, name, .. } => {
464                self.kind = AnyKind::Address {
465                    address: Address {
466                        span: self.span,
467                        scopes,
468                        address: from,
469                        kind: AddressKind::Assigned,
470                        name: *name,
471                        backtrace: Backtrace::capture(),
472                    },
473                };
474            }
475            AnyKind::Address { address } => {
476                address.assign_addr(cx, from)?;
477            }
478            _ => {}
479        }
480
481        Ok(())
482    }
483
484    /// Allocate an address even if it's a locally allocated one.
485    #[inline]
486    pub(super) fn try_alloc_addr(&mut self) -> compile::Result<Option<&mut Address<'a, 'hir>>> {
487        if let AnyKind::Defer {
488            scopes,
489            scope,
490            name,
491        } = self.kind
492        {
493            let address = Address {
494                span: self.span,
495                scopes,
496                address: scopes.alloc_in(self.span, scope)?,
497                kind: AddressKind::Scope(scope),
498                name,
499                backtrace: Backtrace::capture(),
500            };
501
502            self.kind = AnyKind::Address { address };
503        }
504
505        match &mut self.kind {
506            AnyKind::Address { address } => Ok(Some(address.alloc_addr()?)),
507            _ => Ok(None),
508        }
509    }
510
511    /// Get the needs as an output.
512    #[inline]
513    pub(super) fn try_alloc_output(&mut self) -> compile::Result<Option<Output>> {
514        let Some(addr) = self.try_alloc_addr()? else {
515            return Ok(None);
516        };
517
518        Ok(Some(addr.output()))
519    }
520
521    /// Test if any sort of value is needed.
522    #[inline(always)]
523    pub(super) fn alloc_output(&mut self) -> compile::Result<Output> {
524        let Some(addr) = self.try_alloc_addr()? else {
525            return Ok(Output::discard());
526        };
527
528        Ok(addr.output())
529    }
530
531    /// Try to treat as an address.
532    #[inline]
533    pub(super) fn addr(&self) -> compile::Result<&Address<'a, 'hir>> {
534        match &self.kind {
535            AnyKind::Address { address } => Ok(address),
536            kind => Err(compile::Error::msg(
537                self.span,
538                format!("No address for need {kind}"),
539            )),
540        }
541    }
542
543    /// Coerce into an owned address.
544    #[inline]
545    pub(super) fn into_addr(mut self) -> compile::Result<Address<'a, 'hir>> {
546        match replace(&mut self.kind, AnyKind::Freed) {
547            AnyKind::Address { address } => Ok(address),
548            kind => Err(compile::Error::msg(
549                self.span,
550                format!("No address for need {kind}"),
551            )),
552        }
553    }
554
555    /// Coerce into a output.
556    #[inline]
557    pub(super) fn output(&self) -> compile::Result<Output> {
558        match &self.kind {
559            AnyKind::Address { address } => Ok(Output::keep(address.address.offset())),
560            AnyKind::Ignore { .. } => Ok(Output::discard()),
561            kind => Err(compile::Error::msg(
562                self.span,
563                format!("Needs {kind} has not been allocated for output"),
564            )),
565        }
566    }
567
568    /// Gets the initialized address unless it is specifically set to `None`.
569    #[inline]
570    pub(super) fn try_as_addr(&self) -> compile::Result<Option<&Address<'a, 'hir>>> {
571        match &self.kind {
572            AnyKind::Address { address } => {
573                if matches!(address.kind, AddressKind::Dangling) {
574                    return Err(compile::Error::msg(
575                        self.span,
576                        "Expected address to be initialized",
577                    ));
578                }
579
580                Ok(Some(address))
581            }
582            _ => Ok(None),
583        }
584    }
585
586    /// Free the current needs.
587    pub(super) fn free(mut self) -> compile::Result<()> {
588        if let AnyKind::Address { address } = replace(&mut self.kind, AnyKind::Freed) {
589            address.free()?;
590        }
591
592        Ok(())
593    }
594}
595
596impl fmt::Display for Any<'_, '_> {
597    #[inline]
598    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
599        self.kind.fmt(f)
600    }
601}
602
603impl fmt::Debug for Any<'_, '_> {
604    #[inline]
605    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
606        fmt::Display::fmt(self, f)
607    }
608}
609
610impl Drop for Any<'_, '_> {
611    fn drop(&mut self) {
612        if matches!(self.kind, AnyKind::Ignore { .. } | AnyKind::Freed) {
613            return;
614        }
615
616        rune_diagnose!("{self} was not freed:\nallocated at:\n{}", self.backtrace);
617    }
618}