vicis_core/pass/
mod.rs

1pub mod analysis;
2pub mod transform;
3
4use rustc_hash::FxHashMap;
5use std::any::{Any, TypeId};
6
7pub trait AnalysisPass<T> {
8    fn run_on(&self, _: &T, _: &mut Box<dyn Any>) {}
9}
10
11pub trait TransformPass<T> {
12    fn run_on(&self, _: &mut T, _: &mut Box<dyn Any>) {}
13}
14
15pub enum Pass<T> {
16    Analysis(Box<dyn AnalysisPass<T>>),
17    Transform(Box<dyn TransformPass<T>>),
18}
19
20pub struct PassManager<T> {
21    passes: Vec<Pass<T>>,
22    results: FxHashMap<TypeId, Box<dyn Any>>,
23}
24
25impl<T> Default for PassManager<T> {
26    fn default() -> Self {
27        Self {
28            passes: vec![],
29            results: FxHashMap::default(),
30        }
31    }
32}
33
34impl<T> PassManager<T> {
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    pub fn add(&mut self, pass: Pass<T>) {
40        self.passes.push(pass)
41    }
42
43    pub fn add_analysis<P: 'static + AnalysisPass<T>>(&mut self, pass: P) {
44        self.passes.push(Pass::Analysis(Box::new(pass)))
45    }
46
47    pub fn add_transform<P: 'static + TransformPass<T>>(&mut self, pass: P) {
48        self.passes.push(Pass::Transform(Box::new(pass)))
49    }
50
51    pub fn run_on(&mut self, target: &mut T) {
52        self.results.clear();
53
54        for pass in &self.passes {
55            let mut result: Box<dyn Any> = Box::new(());
56            match pass {
57                Pass::Analysis(analysis) => analysis.run_on(target, &mut result),
58                Pass::Transform(transform) => transform.run_on(target, &mut result),
59            }
60            self.results.insert((*result).type_id(), result);
61        }
62    }
63
64    pub fn run_analyses_on(&mut self, target: &T) {
65        self.results.clear();
66
67        for pass in &self.passes {
68            let mut result: Box<dyn Any> = Box::new(());
69            match pass {
70                Pass::Analysis(analysis) => analysis.run_on(target, &mut result),
71                Pass::Transform(_) => {}
72            }
73            self.results.insert((*result).type_id(), result);
74        }
75    }
76
77    pub fn get_result<P: 'static>(&self) -> Option<&P> {
78        self.results
79            .get(&TypeId::of::<P>())
80            .and_then(|result| result.downcast_ref())
81    }
82}
83
84use crate::ir::{function::Function, module::Module};
85
86impl PassManager<Function> {
87    pub fn run_analyses_on_module(&mut self, module: &Module) {
88        for (_, func) in &module.functions {
89            self.run_analyses_on(func)
90        }
91    }
92
93    pub fn run_on_module(&mut self, module: &mut Module) {
94        for (_, func) in &mut module.functions {
95            self.run_on(func)
96        }
97    }
98}
99
100impl<T> Pass<T> {
101    pub fn analysis<P: 'static + AnalysisPass<T>>(pass: P) -> Self {
102        Self::Analysis(Box::new(pass))
103    }
104}
105
106#[cfg(test)]
107mod test {
108    use super::*;
109    use crate::ir::{
110        function::Function,
111        module::{parse_assembly, Module},
112    };
113
114    pub struct TestFunctionAnalysisPass {}
115    pub struct TestFunctionAnalysisResult(String);
116
117    impl AnalysisPass<Function> for TestFunctionAnalysisPass {
118        fn run_on(&self, func: &Function, result: &mut Box<dyn Any>) {
119            *result = Box::new(TestFunctionAnalysisResult(func.name.to_owned()));
120        }
121    }
122
123    pub struct TestFunctionTransformPass {}
124
125    impl TransformPass<Function> for TestFunctionTransformPass {
126        fn run_on(&self, _func: &mut Function, _result: &mut Box<dyn Any>) {}
127    }
128
129    fn test_module() -> Module {
130        parse_assembly(
131            r#"
132source_filename = "sample"                                                                            
133target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"       
134target triple = "x86_64-pc-linux-gnu"                                                              
135         
136define dso_local i32 @main() {                                            
137  ret i32 0
138}                                                                            
139        "#,
140        ).expect("failed to parse IR")
141    }
142
143    #[test]
144    fn analysis() {
145        let module = test_module();
146
147        let mut pm = PassManager::new();
148        pm.add_analysis(TestFunctionAnalysisPass {});
149
150        pm.run_analyses_on_module(&module);
151
152        assert_eq!(
153            pm.get_result::<TestFunctionAnalysisResult>().unwrap().0,
154            "main"
155        );
156    }
157
158    #[test]
159    fn analysis_transform() {
160        let mut module = test_module();
161
162        let mut pm = PassManager::new();
163        pm.add_analysis(TestFunctionAnalysisPass {});
164        pm.add_transform(TestFunctionTransformPass {});
165
166        pm.run_on_module(&mut module);
167
168        assert_eq!(
169            pm.get_result::<TestFunctionAnalysisResult>().unwrap().0,
170            "main"
171        );
172    }
173}
vicis_core/pass/mod.rs

vicis_core/pass/
mod.rs
