llvm-bitcode 0.4.0

use crate::BitStreamReader;
use crate::bitcode::{BitcodeElement, Block, Record, Signature};

/// A visitor which receives callbacks while reading a bitstream.
pub trait BitStreamVisitor {
    /// Validate a bitstream's signature or "magic number".
    fn validate(&self, _signature: Signature) -> bool {
        true
    }

    /// Called when a new block is encountered. Return `true` to enter the block
    /// and read its contents, or `false` to skip it.
    fn should_enter_block(&mut self, block_id: u32) -> bool;

    /// Called when a block is exited.
    fn did_exit_block(&mut self, block_id: u32);

    /// Called whenever a record is encountered.
    fn visit(&mut self, block_id: u32, record: Record);
}

/// A basic visitor that collects all the blocks and records in a stream.
pub struct CollectingVisitor {
    stack: Vec<(u32, Vec<BitcodeElement>)>,
}

impl CollectingVisitor {
    #[must_use]
    pub fn new() -> Self {
        Self {
            stack: vec![(BitStreamReader::TOP_LEVEL_BLOCK_ID, Vec::new())],
        }
    }

    #[must_use]
    pub fn finalize_top_level_elements(mut self) -> Vec<BitcodeElement> {
        assert_eq!(self.stack.len(), 1);
        self.stack.pop().unwrap().1
    }
}

impl BitStreamVisitor for CollectingVisitor {
    fn should_enter_block(&mut self, id: u32) -> bool {
        self.stack.push((id, Vec::new()));
        true
    }

    fn did_exit_block(&mut self, block_id: u32) {
        if let Some((id, elements)) = self.stack.pop() {
            assert_eq!(id, block_id);
            let block = Block { id, elements };
            let last = self.stack.last_mut().unwrap();
            last.1.push(BitcodeElement::Block(block));
        }
    }

    fn visit(&mut self, _block_id: u32, record: Record) {
        let last = self.stack.last_mut().unwrap();
        last.1.push(BitcodeElement::Record(record));
    }
}