beamr 0.16.2

//! Heap allocators for native BIFs.
//!
//! Allocators that take term arguments root them for the duration of the
//! allocation (a reserve may trigger a collection that moves them); the
//! `_prereserved` variants assume `ensure_heap_space` was already called and
//! cannot collect.

#[cfg(feature = "threads")]
use std::sync::Arc;

#[cfg(feature = "threads")]
use crate::io::resource::{FD_RESOURCE_WORDS, FdInner, write_fd_resource};
use crate::term::Term;
use crate::term::boxed::{
    write_bigint, write_cons, write_external_pid, write_external_reference, write_float, write_map,
    write_reference, write_tuple,
};
use crate::term::shared_binary::{alloc_binary, alloc_binary_word_count};

use super::ProcessContext;

impl ProcessContext<'_> {
    /// Allocate a tuple on the calling process heap.
    pub fn alloc_tuple(&mut self, elements: &[Term]) -> Result<Term, Term> {
        let words = 1 + elements.len();
        if self.process.is_none() {
            // Detached contexts allocate owned blocks and never collect.
            let heap = self.alloc_words(words)?;
            return write_tuple(heap, elements)
                .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG));
        }
        // Root the inputs: the reserve below may collect and move them.
        self.with_rooted(elements, |context, roots| {
            context.ensure_heap_space(words)?;
            let mut current = Vec::with_capacity(roots.len);
            for index in 0..roots.len {
                current.push(context.rooted(roots, index)?);
            }
            let heap = context.alloc_words_prereserved(words)?;
            write_tuple(heap, &current).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
        })
    }

    /// Allocate a tuple using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_tuple_prereserved(&mut self, elements: &[Term]) -> Result<Term, Term> {
        let words = 1 + elements.len();
        let heap = self.alloc_words_prereserved(words)?;
        write_tuple(heap, elements).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a reference on the calling process heap.
    pub fn alloc_reference(&mut self, id: u64) -> Result<Term, Term> {
        let heap = self.alloc_words(2)?;
        write_reference(heap, id).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a reference using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_reference_prereserved(&mut self, id: u64) -> Result<Term, Term> {
        let heap = self.alloc_words_prereserved(2)?;
        write_reference(heap, id).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a remote PID on the calling process heap.
    pub fn alloc_external_pid(
        &mut self,
        node: crate::atom::Atom,
        pid_number: u64,
        serial: u64,
    ) -> Result<Term, Term> {
        let heap = self.alloc_words(4)?;
        write_external_pid(heap, node, pid_number, serial)
            .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a remote PID using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_external_pid_prereserved(
        &mut self,
        node: crate::atom::Atom,
        pid_number: u64,
        serial: u64,
    ) -> Result<Term, Term> {
        let heap = self.alloc_words_prereserved(4)?;
        write_external_pid(heap, node, pid_number, serial)
            .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a remote reference on the calling process heap.
    pub fn alloc_external_reference(
        &mut self,
        node: crate::atom::Atom,
        id: u64,
    ) -> Result<Term, Term> {
        let heap = self.alloc_words(3)?;
        write_external_reference(heap, node, id)
            .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a remote reference using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_external_reference_prereserved(
        &mut self,
        node: crate::atom::Atom,
        id: u64,
    ) -> Result<Term, Term> {
        let heap = self.alloc_words_prereserved(3)?;
        write_external_reference(heap, node, id)
            .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a cons cell on the calling process heap.
    pub fn alloc_cons(&mut self, head: Term, tail: Term) -> Result<Term, Term> {
        if self.process.is_none() {
            // Detached contexts allocate owned blocks and never collect.
            let heap = self.alloc_words(2)?;
            return write_cons(heap, head, tail)
                .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG));
        }
        // Root the inputs: the reserve below may collect and move them.
        self.with_rooted(&[head, tail], |context, roots| {
            context.ensure_heap_space(2)?;
            let head = context.rooted(roots, 0)?;
            let tail = context.rooted(roots, 1)?;
            let heap = context.alloc_words_prereserved(2)?;
            write_cons(heap, head, tail).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
        })
    }

    /// Allocate a float on the calling process heap.
    pub fn alloc_float(&mut self, value: f64) -> Result<Term, Term> {
        let heap = self.alloc_words(2)?;
        write_float(heap, value).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a binary on the calling process heap, promoting large binaries to ProcBin.
    pub fn alloc_binary(&mut self, bytes: &[u8]) -> Result<Term, Term> {
        let words = alloc_binary_word_count(bytes.len());
        let heap = self.alloc_words(words)?;
        let term =
            alloc_binary(heap, bytes).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))?;
        // Large binaries land as a refcounted ProcBin; mark the allocation so the
        // GC release walk drops its Arc. (Small inline binaries are marked too;
        // the walk reads their Binary tag and skips them.) See `AllocKind`.
        self.mark_last_allocation_maybe_refcounted();
        Ok(term)
    }

    /// Allocate an FdResource on the calling process heap.
    #[cfg(feature = "threads")]
    pub fn alloc_fd_resource(&mut self, fd_inner: Arc<FdInner>) -> Result<Term, Term> {
        let heap = self.alloc_words(FD_RESOURCE_WORDS)?;
        let term = write_fd_resource(heap, fd_inner)
            .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))?;
        self.mark_last_allocation_maybe_refcounted();
        Ok(term)
    }

    /// Mark the most recent process-heap allocation as possibly holding a
    /// refcounted resource. No-op for detached contexts, whose owned-block
    /// allocations are not part of the young region the release walk scans.
    fn mark_last_allocation_maybe_refcounted(&mut self) {
        if let Some(process) = self.process.as_deref_mut() {
            process
                .heap_mut()
                .mark_last_young_allocation_maybe_refcounted();
        }
    }

    /// Deep-copy an ETS-owned term onto the calling process heap.
    ///
    /// Reserves the exact word budget first, so the copy itself cannot
    /// collect. ETS-owned source memory is unaffected by process GC, so no
    /// rooting is needed.
    pub fn copy_owned_term(&mut self, owned: &crate::ets::OwnedTerm) -> Result<Term, Term> {
        self.ensure_heap_space(owned.total_words())?;
        self.copy_owned_term_prereserved(owned)
    }

    /// Deep-copy an ETS-owned term using pre-reserved heap space (no GC
    /// trigger). Caller must have called `ensure_heap_space` for the total
    /// budget (`OwnedTerm::total_words`).
    ///
    /// The copy materialises binaries inline (never as refcounted ProcBins),
    /// so no allocation marking is needed. Detached contexts adopt a fresh
    /// owned copy into their allocation list instead.
    pub fn copy_owned_term_prereserved(
        &mut self,
        owned: &crate::ets::OwnedTerm,
    ) -> Result<Term, Term> {
        if let Some(process) = self.process.as_deref_mut() {
            return owned
                .copy_to_heap(process.heap_mut())
                .map_err(|_| Term::atom(crate::atom::Atom::BADARG));
        }
        let copied = crate::ets::copy_term_to_ets(owned.root())
            .map_err(|_| Term::atom(crate::atom::Atom::BADARG))?;
        let (root, allocations) = copied.into_raw_parts();
        self.detached_allocations.extend(allocations);
        Ok(root)
    }

    /// Allocate a big integer on the calling process heap.
    pub fn alloc_bigint(&mut self, negative: bool, limbs: &[u64]) -> Result<Term, Term> {
        let words = 3 + limbs.len();
        let heap = self.alloc_words(words)?;
        write_bigint(heap, negative, limbs).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a proper list on the calling process heap.
    pub fn alloc_list(&mut self, elements: &[Term]) -> Result<Term, Term> {
        self.alloc_list_with_tail(elements, Term::NIL)
    }

    /// Allocate list cells for `elements`, ending in `tail`.
    ///
    /// The inputs are registered as GC roots for the duration of the call, so
    /// boxed terms stay valid across the collection the initial reserve may
    /// trigger.
    pub fn alloc_list_with_tail(&mut self, elements: &[Term], tail: Term) -> Result<Term, Term> {
        if self.process.is_none() {
            // Detached contexts allocate owned blocks and never collect.
            let mut tail = tail;
            for element in elements.iter().rev().copied() {
                let heap = self.alloc_words_prereserved(2)?;
                tail = write_cons(heap, element, tail)
                    .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))?;
            }
            return Ok(tail);
        }
        let mut rooted = Vec::with_capacity(elements.len() + 1);
        rooted.extend_from_slice(elements);
        rooted.push(tail);
        self.with_rooted(&rooted, |context, roots| {
            context.ensure_heap_space(elements.len() * 2)?;
            let mut tail = context.rooted(roots, elements.len())?;
            for index in (0..elements.len()).rev() {
                let element = context.rooted(roots, index)?;
                let heap = context.alloc_words_prereserved(2)?;
                tail = write_cons(heap, element, tail)
                    .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))?;
            }
            Ok(tail)
        })
    }

    /// Allocate a cons cell using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_cons_prereserved(&mut self, head: Term, tail: Term) -> Result<Term, Term> {
        let heap = self.alloc_words_prereserved(2)?;
        write_cons(heap, head, tail).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }

    /// Allocate a flatmap on the calling process heap.
    ///
    /// Keys and values are registered as GC roots for the duration of the
    /// call.
    pub fn alloc_map(&mut self, keys: &[Term], values: &[Term]) -> Result<Term, Term> {
        let words = 2 + keys.len() + values.len();
        if self.process.is_none() {
            // Detached contexts allocate owned blocks and never collect.
            let heap = self.alloc_words(words)?;
            return write_map(heap, keys, values)
                .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG));
        }
        let mut rooted = Vec::with_capacity(keys.len() + values.len());
        rooted.extend_from_slice(keys);
        rooted.extend_from_slice(values);
        self.with_rooted(&rooted, |context, roots| {
            context.ensure_heap_space(words)?;
            let mut current_keys = Vec::with_capacity(keys.len());
            let mut current_values = Vec::with_capacity(values.len());
            for index in 0..keys.len() {
                current_keys.push(context.rooted(roots, index)?);
            }
            for index in 0..values.len() {
                current_values.push(context.rooted(roots, keys.len() + index)?);
            }
            let heap = context.alloc_words_prereserved(words)?;
            write_map(heap, &current_keys, &current_values)
                .ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
        })
    }

    /// Allocate a flatmap using pre-reserved heap space (no GC trigger).
    /// Caller must have called `ensure_heap_space` for the total budget.
    pub fn alloc_map_prereserved(&mut self, keys: &[Term], values: &[Term]) -> Result<Term, Term> {
        let words = 2 + keys.len() + values.len();
        let heap = self.alloc_words_prereserved(words)?;
        write_map(heap, keys, values).ok_or_else(|| Term::atom(crate::atom::Atom::BADARG))
    }
}