thalir_core/builder/
function_builder.rs

1use super::{BlockBuilder, IRContext, IRRegistry};
2use crate::{
3    block::BlockId,
4    function::{Function, FunctionSignature, Mutability, Parameter, Visibility},
5    types::Type,
6    values::{ParamId, Value},
7    Result,
8};
9
10pub struct FunctionBuilder<'a> {
11    contract_name: String,
12    function: Function,
13    context: &'a mut IRContext,
14    registry: &'a mut IRRegistry,
15    current_block: Option<BlockId>,
16    created_blocks: std::collections::HashSet<BlockId>,
17}
18
19impl<'a> FunctionBuilder<'a> {
20    pub fn new(
21        contract_name: String,
22        name: String,
23        context: &'a mut IRContext,
24        registry: &'a mut IRRegistry,
25    ) -> Self {
26        let signature = FunctionSignature {
27            name: name.clone(),
28            params: Vec::new(),
29            returns: Vec::new(),
30            is_payable: false,
31        };
32
33        let function = Function::new(signature);
34
35        Self {
36            contract_name,
37            function,
38            context,
39            registry,
40            current_block: None,
41            created_blocks: std::collections::HashSet::new(),
42        }
43    }
44
45    pub fn param(&mut self, name: &str, ty: Type) -> &mut Self {
46        self.function
47            .signature
48            .params
49            .push(Parameter::new(name, ty));
50        self
51    }
52
53    pub fn get_params(&self) -> &Vec<Parameter> {
54        &self.function.signature.params
55    }
56
57    pub fn returns(&mut self, ty: Type) -> &mut Self {
58        self.function.signature.returns = vec![ty];
59        self
60    }
61
62    pub fn returns_multiple(&mut self, types: Vec<Type>) -> &mut Self {
63        self.function.signature.returns = types;
64        self
65    }
66
67    pub fn visibility(&mut self, vis: Visibility) -> &mut Self {
68        self.function.visibility = vis;
69        self
70    }
71
72    pub fn mutability(&mut self, mut_: Mutability) -> &mut Self {
73        self.function.mutability = mut_;
74        self
75    }
76
77    pub fn modifier(&mut self, _name: &str) -> &mut Self {
78        self.function.modifiers.push(crate::contract::ModifierRef {
79            id: crate::contract::ModifierId(0),
80            arguments: Vec::new(),
81        });
82        self
83    }
84
85    pub fn create_block_id(&mut self) -> BlockId {
86        let block_id = self.function.body.create_block();
87        self.created_blocks.insert(block_id);
88        block_id
89    }
90
91    pub fn switch_to_block(&mut self, block_id: BlockId) -> Result<BlockBuilder<'_>> {
92        if !self.created_blocks.contains(&block_id) {
93            return Err(crate::IrError::BuilderError(format!(
94                "Block {} does not exist in function",
95                block_id
96            )));
97        }
98
99        self.current_block = Some(block_id);
100        self.context.set_current_block(block_id);
101
102        let qualified_func = format!("{}::{}", self.contract_name, self.function.signature.name);
103
104        Ok(BlockBuilder::new(
105            block_id,
106            qualified_func,
107            self.context,
108            self.registry,
109        ))
110    }
111
112    pub fn current_block(&self) -> Option<BlockId> {
113        self.current_block
114    }
115
116    pub fn entry_block(&mut self) -> BlockBuilder<'_> {
117        let block_id = self.function.body.entry_block;
118        self.current_block = Some(block_id);
119        self.context.set_current_block(block_id);
120        self.created_blocks.insert(block_id);
121
122        let qualified_func = format!("{}::{}", self.contract_name, self.function.signature.name);
123
124        BlockBuilder::new(block_id, qualified_func, self.context, self.registry)
125    }
126
127    pub fn new_block(&mut self, _name: &str) -> BlockBuilder<'_> {
128        self.block(_name)
129    }
130
131    pub fn block(&mut self, _name: &str) -> BlockBuilder<'_> {
132        let block_id = self.function.body.create_block();
133        self.current_block = Some(block_id);
134        self.context.set_current_block(block_id);
135        self.created_blocks.insert(block_id);
136
137        let qualified_func = format!("{}::{}", self.contract_name, self.function.signature.name);
138
139        BlockBuilder::new(block_id, qualified_func, self.context, self.registry)
140    }
141
142    pub fn block_with_id(&mut self, block_id: BlockId) -> BlockBuilder<'_> {
143        self.current_block = Some(block_id);
144        self.context.set_current_block(block_id);
145        self.created_blocks.insert(block_id);
146
147        let qualified_func = format!("{}::{}", self.contract_name, self.function.signature.name);
148
149        BlockBuilder::new(block_id, qualified_func, self.context, self.registry)
150    }
151
152    pub fn local(&mut self, name: &str, ty: Type) -> Value {
153        let var_id = self.context.ssa().get_or_create_var(name);
154        self.function
155            .body
156            .locals
157            .push(crate::function::LocalVariable {
158                id: crate::function::LocalId(self.function.body.locals.len() as u32),
159                name: name.to_string(),
160                var_type: ty,
161                location: crate::function::DataLocation::Memory,
162            });
163        Value::Variable(var_id)
164    }
165
166    pub fn get_param(&self, index: usize) -> Value {
167        Value::Param(ParamId(index as u32))
168    }
169
170    pub fn current_function(&self) -> &Function {
171        &self.function
172    }
173
174    pub fn build(self) -> Result<Function> {
175        let qualified_name = format!("{}::{}", self.contract_name, self.function.signature.name);
176        self.registry
177            .add_function(self.contract_name.clone(), self.function)?;
178
179        self.registry
180            .get_function(&qualified_name)
181            .ok_or_else(|| {
182                crate::IrError::BuilderError(format!(
183                    "Function {} not found after registration",
184                    qualified_name
185                ))
186            })
187            .map(|f| f.clone())
188    }
189}
thalir_core/builder/function_builder.rs

thalir_core/builder/
function_builder.rs
