llvm_bitcode/
visitor.rs

1use crate::bitcode::{BitcodeElement, Block, Record, Signature};
2use crate::BitStreamReader;
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    /// Called when a new block is encountered. Return `true` to enter the block
11    /// and read its contents, or `false` to skip it.
12    fn should_enter_block(&mut self, id: u64) -> bool;
13    /// Called when a block is exited.
14    fn did_exit_block(&mut self);
15    /// Called whenever a record is encountered.
16    fn visit(&mut self, record: Record);
17}
18
19/// A basic visitor that collects all the blocks and records in a stream.
20pub struct CollectingVisitor {
21    stack: Vec<(u64, Vec<BitcodeElement>)>,
22}
23
24impl CollectingVisitor {
25    pub fn new() -> Self {
26        Self {
27            stack: vec![(BitStreamReader::TOP_LEVEL_BLOCK_ID, Vec::new())],
28        }
29    }
30
31    pub fn finalize_top_level_elements(mut self) -> Vec<BitcodeElement> {
32        assert_eq!(self.stack.len(), 1);
33        self.stack.pop().unwrap().1
34    }
35}
36
37impl BitStreamVisitor for CollectingVisitor {
38    fn should_enter_block(&mut self, id: u64) -> bool {
39        self.stack.push((id, Vec::new()));
40        true
41    }
42
43    fn did_exit_block(&mut self) {
44        if let Some((id, elements)) = self.stack.pop() {
45            let block = Block { id, elements };
46            let last = self.stack.last_mut().unwrap();
47            last.1.push(BitcodeElement::Block(block));
48        }
49    }
50
51    fn visit(&mut self, record: Record) {
52        let last = self.stack.last_mut().unwrap();
53        last.1.push(BitcodeElement::Record(record));
54    }
55}