llvm_bitcode/
visitor.rs

1use crate::BitStreamReader;
2use crate::bitcode::{BitcodeElement, Block, Record, Signature};
3
4/// A visitor which receives callbacks while reading a bitstream.
5pub trait BitStreamVisitor {
6    /// Validate a bitstream's signature or "magic number".
7    fn validate(&self, _signature: Signature) -> bool {
8        true
9    }
10
11    /// Called when a new block is encountered. Return `true` to enter the block
12    /// and read its contents, or `false` to skip it.
13    fn should_enter_block(&mut self, block_id: u32) -> bool;
14
15    /// Called when a block is exited.
16    fn did_exit_block(&mut self, block_id: u32);
17
18    /// Called whenever a record is encountered.
19    fn visit(&mut self, block_id: u32, record: Record);
20}
21
22/// A basic visitor that collects all the blocks and records in a stream.
23pub struct CollectingVisitor {
24    stack: Vec<(u32, Vec<BitcodeElement>)>,
25}
26
27impl CollectingVisitor {
28    #[must_use]
29    pub fn new() -> Self {
30        Self {
31            stack: vec![(BitStreamReader::TOP_LEVEL_BLOCK_ID, Vec::new())],
32        }
33    }
34
35    #[must_use]
36    pub fn finalize_top_level_elements(mut self) -> Vec<BitcodeElement> {
37        assert_eq!(self.stack.len(), 1);
38        self.stack.pop().unwrap().1
39    }
40}
41
42impl BitStreamVisitor for CollectingVisitor {
43    fn should_enter_block(&mut self, id: u32) -> bool {
44        self.stack.push((id, Vec::new()));
45        true
46    }
47
48    fn did_exit_block(&mut self, block_id: u32) {
49        if let Some((id, elements)) = self.stack.pop() {
50            assert_eq!(id, block_id);
51            let block = Block { id, elements };
52            let last = self.stack.last_mut().unwrap();
53            last.1.push(BitcodeElement::Block(block));
54        }
55    }
56
57    fn visit(&mut self, _block_id: u32, record: Record) {
58        let last = self.stack.last_mut().unwrap();
59        last.1.push(BitcodeElement::Record(record));
60    }
61}