pipeline_script/context/
mod.rs

1use crate::ast::module::Module;
2use crate::ast::r#type::Type;
3use crate::ast::type_alias::TypeAlias;
4use crate::context::key::ContextKey;
5use crate::context::scope::Scope;
6use crate::context::value::ContextValue;
7use crate::llvm::builder::ir::IRBuilder;
8use crate::llvm::context::LLVMContext;
9use crate::llvm::module::LLVMModule;
10use crate::llvm::types::LLVMType;
11use crate::llvm::value::fucntion::FunctionValue;
12use crate::llvm::value::LLVMValue;
13use llvm_sys::prelude::LLVMBasicBlockRef;
14use slotmap::DefaultKey;
15use std::collections::HashMap;
16use std::rc::Rc;
17use std::sync::{Arc, Mutex, RwLock};
18
19pub mod key;
20pub mod scope;
21pub mod value;
22
23#[derive(Clone, Debug)]
24pub struct Context {
25    parent: Option<Box<Context>>,
26    key: ContextKey,
27    value: ContextValue,
28}
29
30impl Context {
31    pub fn background() -> Self {
32        Self {
33            parent: None,
34            key: ContextKey::Background,
35            value: ContextValue::Background,
36        }
37    }
38    pub fn with_builder(parent: &Context, builder: IRBuilder) -> Self {
39        Self {
40            parent: Some(Box::new(parent.clone())),
41            key: ContextKey::Builder,
42            value: ContextValue::Builder(Arc::new(builder)),
43        }
44    }
45    pub fn with_module_slot_map(parent: &Context, t: slotmap::SlotMap<DefaultKey, Module>) -> Self {
46        Self::with_value(
47            parent,
48            ContextKey::ModuleSlotMap,
49            ContextValue::ModuleSlotMap(Arc::new(RwLock::new(t))),
50        )
51    }
52    pub fn apply_mut_module(&self, key: DefaultKey, apply: impl Fn(&mut Module)) {
53        let slot_map = self.get(ContextKey::ModuleSlotMap).unwrap();
54        match slot_map {
55            ContextValue::ModuleSlotMap(slot_map) => {
56                let mut slot_map = slot_map.write().unwrap();
57                let module = slot_map.get_mut(key).unwrap();
58                apply(module)
59            }
60            _ => panic!("not a module slot map"),
61        }
62    }
63    pub fn get_module_slot_map(&self) -> Arc<RwLock<slotmap::SlotMap<DefaultKey, Module>>> {
64        match self.get(ContextKey::ModuleSlotMap) {
65            Some(ContextValue::ModuleSlotMap(slot_map)) => slot_map.clone(),
66            _ => Arc::new(RwLock::new(slotmap::SlotMap::new())),
67        }
68    }
69
70    pub fn apply_module(&self, key: DefaultKey, mut apply: impl FnMut(&Module)) {
71        let slot_map = self.get(ContextKey::ModuleSlotMap).unwrap();
72        match slot_map {
73            ContextValue::ModuleSlotMap(slot_map) => {
74                let slot_map = slot_map.read().unwrap();
75                let module = slot_map.get(key).unwrap();
76                apply(module)
77            }
78            _ => panic!("not a module slot map"),
79        }
80    }
81    pub fn register_module(&self, module: Module) -> DefaultKey {
82        let slot_map = self.get(ContextKey::ModuleSlotMap).unwrap();
83        match slot_map {
84            ContextValue::ModuleSlotMap(slot_map) => {
85                let mut slot_map = slot_map.write().unwrap();
86                slot_map.insert(module)
87            }
88            _ => panic!("not a module slot map"),
89        }
90    }
91    pub fn with_type(parent: &Context, name: String, ty: Type) -> Self {
92        Self::with_value(parent, ContextKey::Type(name), ContextValue::Type(ty))
93    }
94    pub fn get_builder(&self) -> Arc<IRBuilder> {
95        match self.get(ContextKey::Builder) {
96            Some(ContextValue::Builder(b)) => b.clone(),
97            _ => panic!("not a builder"),
98        }
99    }
100    pub fn get_current_function(&self) -> FunctionValue {
101        match self.get(ContextKey::Function) {
102            Some(ContextValue::Function(b)) => b.clone(),
103            _ => panic!("not a function"),
104        }
105    }
106    pub fn get_current_function_type(&self) -> &Type {
107        match self.get(ContextKey::Type("current_function".into())) {
108            Some(ContextValue::Type(ty)) => ty,
109            _ => panic!("not a function"),
110        }
111    }
112    pub fn with_scope(parent: &Context) -> Self {
113        Self::with_value(parent, ContextKey::Scope, ContextValue::Scope(Scope::new()))
114    }
115    pub fn with_type_table(parent: &Context, t: HashMap<Type, LLVMType>) -> Self {
116        Self::with_value(
117            parent,
118            ContextKey::TypeTable,
119            ContextValue::TypeTable(Rc::new(RwLock::new(t))),
120        )
121    }
122    pub fn create_llvm_context() -> Self {
123        let llvm_ctx = LLVMContext::new();
124        let module = llvm_ctx.create_module("main");
125        let ctx = Context {
126            parent: None,
127            key: ContextKey::LLVMContext,
128            value: ContextValue::LLVMContext(Rc::new(Mutex::new(LLVMContext::new()))),
129        };
130        Self {
131            parent: Some(Box::new(ctx)),
132            key: ContextKey::LLVMModule,
133            value: ContextValue::LLVMModule(Rc::new(RwLock::new(module))),
134        }
135    }
136    pub fn with_function(parent: &Context, f: FunctionValue) -> Self {
137        Self::with_value(parent, ContextKey::Function, ContextValue::Function(f))
138    }
139    pub fn with_flag(parent: &Context, key: impl Into<String>, flag: bool) -> Self {
140        Self::with_value(
141            parent,
142            ContextKey::Flag(key.into()),
143            ContextValue::Flag(Arc::new(RwLock::new(flag))),
144        )
145    }
146    pub fn get_llvm_module(&self) -> Rc<RwLock<LLVMModule>> {
147        match self.get(ContextKey::LLVMModule) {
148            Some(ContextValue::LLVMModule(m)) => m.clone(),
149            _ => panic!("not a llvm module"),
150        }
151    }
152    pub fn with_llvm_module(parent: &Context, m: LLVMModule) -> Self {
153        Self::with_value(
154            parent,
155            ContextKey::LLVMModule,
156            ContextValue::LLVMModule(Rc::new(RwLock::new(m))),
157        )
158    }
159    pub fn get_llvm_context(&self) -> Rc<Mutex<LLVMContext>> {
160        match self.get(ContextKey::LLVMContext) {
161            Some(ContextValue::LLVMContext(c)) => c.clone(),
162            _ => panic!("not a llvm context"),
163        }
164    }
165
166    pub fn with_local(parent: &Context, local: Vec<String>) -> Self {
167        Self::with_value(
168            parent,
169            ContextKey::LocalVariable,
170            ContextValue::LocalVariable(Arc::new(RwLock::new(local))),
171        )
172    }
173    pub fn get_type_binding_functions(&self, name: &str) -> Vec<crate::ast::function::Function> {
174        let slot_map = self.get_module_slot_map();
175        let slot_map = slot_map.read().unwrap();
176        let mut result = Vec::new();
177
178        for module in slot_map.values() {
179            for (_, function) in module.get_functions() {
180                if function.has_binding() && function.get_binding() == name {
181                    result.push(function);
182                }
183            }
184        }
185
186        result
187    }
188
189    pub fn with_capture(parent: &Context) -> Self {
190        Self::with_value(
191            parent,
192            ContextKey::CaptureVariable,
193            ContextValue::CaptureVariable(Arc::new(RwLock::new(vec![]))),
194        )
195    }
196    pub fn set_flag(&self, key: impl Into<String>, flag: bool) {
197        match self.get(ContextKey::Flag(key.into())) {
198            Some(ContextValue::Flag(f)) => {
199                let mut f = f.write().unwrap();
200                *f = flag;
201            }
202            _ => panic!("not a flag"),
203        }
204    }
205    pub fn get_flag(&self, key: impl Into<String>) -> Option<bool> {
206        match self.get(ContextKey::Flag(key.into())) {
207            Some(ContextValue::Flag(f)) => {
208                let f = f.read().unwrap();
209                Some(*f)
210            }
211            _ => None,
212        }
213    }
214    pub fn with_value(parent: &Context, key: ContextKey, value: ContextValue) -> Self {
215        Self {
216            parent: Some(Box::new(parent.clone())),
217            key,
218            value,
219        }
220    }
221
222    pub fn get(&self, key: ContextKey) -> Option<&ContextValue> {
223        if self.key == key {
224            return Some(&self.value);
225        }
226        match &self.parent {
227            None => None,
228            Some(parent) => parent.get(key),
229        }
230    }
231    pub fn get_scope(&self) -> Scope {
232        match self.get(ContextKey::Scope) {
233            Some(ContextValue::Scope(s)) => s.clone(),
234            _ => panic!("not a scope"),
235        }
236    }
237    pub fn set_symbol(&self, name: String, v: LLVMValue) {
238        let scope = self.get(ContextKey::Scope).unwrap().as_scope();
239        scope.set(name, v);
240    }
241    pub fn get_symbol(&self, name: impl AsRef<str>) -> Option<LLVMValue> {
242        let scope = self.get(ContextKey::Scope).unwrap().as_scope();
243        match scope.has(name.as_ref()) {
244            true => scope.get(name),
245            false => match &self.parent {
246                None => None,
247                Some(parent) => parent.get_symbol(name),
248            },
249        }
250    }
251    pub fn get_type(&self, t: &Type) -> Option<LLVMType> {
252        match self.get(ContextKey::TypeTable) {
253            Some(ContextValue::TypeTable(tt0)) => {
254                let tt = tt0.read().unwrap();
255                let r = tt.get(t);
256                r.cloned()
257            }
258            _ => panic!("not a type table"),
259        }
260    }
261    pub fn register_type(&self, ty: &Type, llvm_ty: &LLVMType) {
262        match self.get(ContextKey::TypeTable) {
263            Some(ContextValue::TypeTable(tt0)) => {
264                let mut tt = tt0.write().unwrap();
265                tt.insert(ty.clone(), llvm_ty.clone());
266            }
267            _ => panic!("not a type table"),
268        }
269    }
270    pub fn get_alias_type(&self, name: impl AsRef<str>) -> Option<Type> {
271        match self.get(ContextKey::AliasType) {
272            Some(ContextValue::AliasType(t)) => {
273                let t = t.read().unwrap();
274                t.get(name.as_ref()).cloned()
275            }
276            _ => panic!("not a symbol type"),
277        }
278    }
279    pub fn set_alias_type(&self, name: String, t0: Type) {
280        match self.get(ContextKey::AliasType) {
281            Some(ContextValue::AliasType(t)) => {
282                let mut t = t.write().unwrap();
283                t.insert(name, t0);
284            }
285            _ => panic!("not a symbol type"),
286        }
287    }
288    pub fn get_symbol_type(&self, name: impl AsRef<str>) -> Option<Type> {
289        match self.get(ContextKey::SymbolType) {
290            Some(ContextValue::SymbolType(t)) => {
291                let t = t.read().unwrap();
292                let r = t.get(name.as_ref());
293                match r {
294                    None => self.parent.clone()?.parent?.get_symbol_type(name),
295                    Some(s) => Some(s.clone()),
296                }
297            }
298            _ => None,
299        }
300    }
301    pub fn try_add_local(&self, name: String) {
302        if let Some(ContextValue::LocalVariable(local)) = self.get(ContextKey::LocalVariable) {
303            let mut local = local.write().unwrap();
304            local.push(name);
305        }
306    }
307    pub fn add_capture(&self, name: String, ty: Type) {
308        match self.get(ContextKey::CaptureVariable) {
309            Some(ContextValue::CaptureVariable(local)) => {
310                let mut local = local.write().unwrap();
311                local.push((name, ty));
312            }
313            _ => panic!("not a local variable"),
314        }
315    }
316    pub fn is_local_variable(&self, name: impl AsRef<str>) -> bool {
317        match self.get(ContextKey::LocalVariable) {
318            Some(ContextValue::LocalVariable(local)) => {
319                let local = local.read().unwrap();
320                local.contains(&name.as_ref().to_string())
321            }
322            _ => panic!("not a local variable"),
323        }
324    }
325    pub fn get_captures(&self) -> Option<Vec<(String, Type)>> {
326        match self.get(ContextKey::CaptureVariable) {
327            Some(ContextValue::CaptureVariable(local)) => {
328                let local = local.read().unwrap();
329                Some(local.clone())
330            }
331            _ => panic!("not a local variable"),
332        }
333    }
334    pub fn get_function(&self, name: impl AsRef<str>) -> Option<crate::ast::function::Function> {
335        let slot_map = self.get_module_slot_map();
336        let slot_map = slot_map.read().unwrap();
337        for module in slot_map.values() {
338            if let Some(fun) = module.get_function(name.as_ref()) {
339                return Some(fun.clone());
340            }
341        }
342        None
343    }
344    pub fn set_symbol_type(&self, name: String, t0: Type) {
345        match self.get(ContextKey::SymbolType) {
346            Some(ContextValue::SymbolType(t)) => {
347                let mut t = t.write().unwrap();
348                t.insert(name, t0);
349            }
350            _ => panic!("not a symbol type"),
351        }
352    }
353    pub fn get_type_alias(&self, name: impl AsRef<str>) -> Option<TypeAlias> {
354        let slot_map = self.get_module_slot_map();
355        let slot_map = slot_map.read().unwrap();
356
357        for module in slot_map.values() {
358            if let Some(ty) = module.get_type_alias(name.as_ref()) {
359                return Some(ty.clone());
360            }
361        }
362
363        None
364    }
365    pub fn with_loop_block(parent: &Context, name: String, block: LLVMBasicBlockRef) -> Self {
366        Self::with_value(
367            parent,
368            ContextKey::LoopBlock(name),
369            ContextValue::LoopBlock(block),
370        )
371    }
372
373    pub fn get_loop_block(&self, name: &str) -> Option<LLVMBasicBlockRef> {
374        match self.get(ContextKey::LoopBlock(name.to_string())) {
375            Some(ContextValue::LoopBlock(block)) => Some(*block),
376            _ => None,
377        }
378    }
379
380    pub fn with_default_expr(parent: &Context) -> Self {
381        Self::with_value(
382            parent,
383            ContextKey::DefaultExpr,
384            ContextValue::DefaultExpr(Arc::new(RwLock::new(HashMap::new()))),
385        )
386    }
387    pub fn set_default_expr(&self, name: String, expr: Box<crate::ast::expr::ExprNode>) {
388        match self.get(ContextKey::DefaultExpr) {
389            Some(ContextValue::DefaultExpr(map)) => {
390                let mut map = map.write().unwrap();
391                map.insert(name, expr);
392            }
393            _ => panic!("not a default expr"),
394        }
395    }
396
397    pub fn get_default_expr(&self, name: &str) -> Option<Box<crate::ast::expr::ExprNode>> {
398        if let Some(ContextValue::DefaultExpr(map)) = self.get(ContextKey::DefaultExpr) {
399            map.read().unwrap().get(name).cloned()
400        } else {
401            None
402        }
403    }
404}
pipeline_script/context/mod.rs

pipeline_script/context/
mod.rs
