vicis_core/ir/function/data/
mod.rs

1use crate::ir::{
2    function::{
3        basic_block::{BasicBlock, BasicBlockId},
4        instruction::{Instruction, InstructionId},
5    },
6    value::{Value, ValueId},
7};
8use id_arena::Arena;
9use rustc_hash::{FxHashMap, FxHashSet};
10
11pub struct Data {
12    pub values: Arena<Value>,
13    pub instructions: Arena<Instruction>,
14    pub basic_blocks: Arena<BasicBlock>,
15    pub users_map: FxHashMap<InstructionId, FxHashSet<InstructionId>>,
16}
17
18impl Default for Data {
19    fn default() -> Self {
20        Self {
21            values: Arena::new(),
22            instructions: Arena::new(),
23            basic_blocks: Arena::new(),
24            users_map: FxHashMap::default(),
25        }
26    }
27}
28
29impl Data {
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    pub fn create_block(&mut self) -> BasicBlockId {
35        self.basic_blocks.alloc(BasicBlock::new())
36    }
37
38    pub fn create_inst(&mut self, mut inst: Instruction) -> InstructionId {
39        let id = self.instructions.alloc_with_id(|id| {
40            inst.id = Some(id);
41            inst
42        });
43        self.users_map.insert(id, FxHashSet::default());
44        self.validate_inst_uses(id);
45        id
46    }
47
48    pub fn create_value(&mut self, inst: Value) -> ValueId {
49        self.values.alloc(inst)
50    }
51
52    pub fn replace_inst(&mut self, from: InstructionId, to: Instruction) {
53        self.instructions[from].replace(to);
54        self.validate_inst_uses(from)
55    }
56
57    pub fn block_ref(&self, id: BasicBlockId) -> &BasicBlock {
58        &self.basic_blocks[id]
59    }
60
61    pub fn block_ref_mut(&mut self, id: BasicBlockId) -> &mut BasicBlock {
62        &mut self.basic_blocks[id]
63    }
64
65    pub fn remove_block_pred(&mut self, id: BasicBlockId, pred: BasicBlockId) -> bool {
66        self.basic_blocks[id].preds.remove(&pred)
67    }
68
69    pub fn remove_block_succ(&mut self, id: BasicBlockId, succ: BasicBlockId) -> bool {
70        self.basic_blocks[id].succs.remove(&succ)
71    }
72
73    pub fn inst_ref(&self, id: InstructionId) -> &Instruction {
74        &self.instructions[id]
75    }
76
77    pub fn inst_ref_mut(&mut self, id: InstructionId) -> &mut Instruction {
78        &mut self.instructions[id]
79    }
80
81    pub fn value_ref(&self, id: ValueId) -> &Value {
82        &self.values[id]
83    }
84
85    pub fn value_ref_mut(&mut self, id: ValueId) -> &mut Value {
86        &mut self.values[id]
87    }
88
89    pub fn users_of(&self, id: InstructionId) -> &FxHashSet<InstructionId> {
90        &self.users_map[&id]
91    }
92
93    /// If an instruction with `id` has the only one user, return it.
94    /// Otherwise, return None.
95    pub fn only_one_user_of(&self, id: InstructionId) -> Option<InstructionId> {
96        if self.users_of(id).len() != 1 {
97            return None;
98        }
99        self.users_of(id).iter().next().copied()
100    }
101
102    pub fn replace_inst_arg(&mut self, inst_id: InstructionId, from: InstructionId, to: ValueId) {
103        let inst = &mut self.instructions[inst_id];
104        let mut replaced = false;
105
106        // Replace `from` args with `to`
107        for arg_id in inst.operand.args_mut() {
108            let arg = &self.values[*arg_id];
109            if matches!(arg, Value::Instruction(i) if i == &from) {
110                *arg_id = to;
111                replaced = true;
112            }
113            continue;
114        }
115
116        // Return if no args replaced
117        if !replaced {
118            return;
119        }
120
121        // Update users map
122
123        if let Some(map) = self.users_map.get_mut(&from) {
124            assert!(map.remove(&inst_id));
125        }
126
127        if let Value::Instruction(to) = self.values[to] {
128            self.users_map
129                .entry(to)
130                .or_insert_with(FxHashSet::default)
131                .insert(inst_id);
132        }
133    }
134
135    pub fn replace_all_uses(&mut self, inst_id: InstructionId, to: ValueId) {
136        for user_id in self.users_map[&inst_id].clone() {
137            self.replace_inst_arg(user_id, inst_id, to);
138        }
139    }
140
141    pub fn validate_inst_uses(&mut self, id: InstructionId) {
142        let args = self.instructions[id]
143            .operand
144            .args()
145            .iter()
146            .filter_map(|&arg| match &self.values[arg] {
147                Value::Instruction(id) => Some(*id),
148                _ => None,
149            })
150            .collect::<Vec<InstructionId>>();
151        for arg in args {
152            self.users_map
153                .entry(arg)
154                .or_insert_with(FxHashSet::default)
155                .insert(id);
156        }
157    }
158
159    pub fn remove_uses(&mut self, id: InstructionId) {
160        let args = self.instructions[id]
161            .operand
162            .args()
163            .iter()
164            .filter_map(|&arg| match &self.values[arg] {
165                Value::Instruction(id) => Some(*id),
166                _ => None,
167            })
168            .collect::<Vec<InstructionId>>();
169        for arg in args {
170            if let Some(users) = self.users_map.get_mut(&arg) {
171                users.remove(&id);
172            }
173        }
174    }
175}
vicis_core/ir/function/data/mod.rs

vicis_core/ir/function/data/
mod.rs
