dbt-antlr4 1.2.3

Dbt fork of ANTLR4 runtime for Rust
Documentation
#![allow(clippy::mut_from_ref)]
use crate::{
    cast_unchecked,
    token::Token,
    tree::{NodeInner, NodeKindType, TreeNode},
};

#[derive(Debug)]
pub struct Arena {
    tokens: bumpalo::Bump,
    contexts: bumpalo::Bump,
    payloads: bumpalo::Bump,
}

impl Arena {
    /// Invoke the callback function with a reference to a local Arena.
    ///
    /// The arena will be dropped when the callback returns.
    pub fn with<F, R>(f: F) -> R
    where
        F: for<'a> FnOnce(&'a Arena) -> R,
    {
        let arena = Arena::new();
        f(&arena)
    }

    pub(crate) fn new() -> Self {
        Self {
            tokens: bumpalo::Bump::new(),
            contexts: bumpalo::Bump::new(),
            payloads: bumpalo::Bump::new(),
        }
    }

    pub(crate) fn children_arena(&self) -> &bumpalo::Bump {
        &self.payloads
    }

    pub fn alloc_token<T>(&self, value: T) -> &mut T
    where
        T: Token,
    {
        self.tokens.alloc(value)
    }

    pub fn alloc_string(&self, value: String) -> &str {
        self.tokens.alloc(value)
    }

    pub fn alloc_labeled_node<'input, 'a, N, T>(&self, value: T) -> *mut TreeNode<'input, 'a, N>
    where
        'input: 'a,
        N: NodeKindType<'a>,
        T: NodeInner<'input, 'a, N>,
    {
        self.contexts.alloc(value) as *mut _ as *mut TreeNode<'input, 'a, N>
    }

    pub fn alloc_zeroed_node<'input, 'a, N, I>(&'a self, body: I) -> *mut TreeNode<'input, 'a, N>
    where
        'input: 'a,
        N: NodeKindType<'a>,
        I: NodeInner<'input, 'a, N>,
    {
        let zeroed_header = unsafe { std::mem::zeroed::<TreeNode<'input, 'a, N>>() };
        let ptr = self.contexts.alloc((zeroed_header, body));
        ptr as *mut _ as *mut TreeNode<'input, 'a, N>
    }

    pub fn alloc_node<'input, 'a, N, I>(
        &'a self,
        header: TreeNode<'input, 'a, N>,
        body: I,
    ) -> &'a mut TreeNode<'input, 'a, N>
    where
        'input: 'a,
        N: NodeKindType<'a>,
        I: NodeInner<'input, 'a, N>,
    {
        let ptr = self.contexts.alloc((header, body));
        // Safety: casting to the header portion:
        cast_unchecked!(ptr => mut TreeNode<'input, 'a, N>)
    }

    pub fn alloc_exception<'a, T>(&'a self, value: T) -> bumpalo::boxed::Box<'a, T> {
        bumpalo::boxed::Box::new_in(value, &self.payloads)
    }

    pub fn alloc_payload<T>(&self, value: T) -> &mut T {
        self.payloads.alloc(value)
    }
}

pub(crate) fn is_ref_in_arena<T>(ptr: &T, arena: &bumpalo::Bump) -> bool {
    let p = ptr as *const T as usize;
    ptr_in_arena(p, arena)
}

pub(crate) fn is_slice_in_arena<T>(ptr: &[T], arena: &bumpalo::Bump) -> bool {
    let p = ptr.as_ptr() as usize;
    ptr_in_arena(p, arena)
}

fn ptr_in_arena(ptr: usize, arena: &bumpalo::Bump) -> bool {
    // SAFETY: We're not allocating from the arena while iterating over the
    // chunks
    unsafe {
        arena.iter_allocated_chunks_raw().any(|(start, size)| {
            let start = start as usize;
            let end = start + size;
            ptr >= start && ptr < end
        })
    }
}