1use std::{
2 collections::{HashMap, HashSet},
3 fmt::Display,
4};
5
6mod deal_ii;
7pub use deal_ii::deal_ii_factory;
8
9use indexmap::IndexSet;
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12
13use crate::codegen::input_schema::{FunctionDef, GenConfig};
14use symrs::{Equation, Expr};
15
16use super::input_schema::{FiniteElement, mesh::Mesh};
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct BuildingBlock {
20 pub includes: IndexSet<String>,
21 pub data: Vec<String>,
22 pub setup: Vec<String>,
23 pub additional_names: HashSet<String>,
24 pub constructor: Vec<String>,
25 pub methods_defs: Vec<String>,
26 pub methods_impls: Vec<String>,
27 pub main: Vec<String>,
28 pub main_setup: Vec<String>,
29 pub additional_vectors: HashSet<String>,
30 pub additional_matrixes: HashSet<String>,
31 pub global: Vec<String>,
32 pub output: Vec<String>,
33}
34
35impl BuildingBlock {
36 pub fn new() -> Self {
37 BuildingBlock {
38 includes: IndexSet::new(),
39 data: Vec::new(),
40 setup: Vec::new(),
41 additional_names: HashSet::new(),
42 constructor: Vec::new(),
43 methods_defs: Vec::new(),
44 methods_impls: Vec::new(),
45 main: Vec::new(),
46 main_setup: Vec::new(),
47 additional_vectors: HashSet::new(),
48 additional_matrixes: HashSet::new(),
49 global: Vec::new(),
50 output: Vec::new(),
51 }
52 }
53
54 fn add_global<T: ToString>(&mut self, global: T) {
55 self.global.push(global.to_string())
56 }
57
58 fn add_includes(&mut self, includes: &[&str]) {
59 for include in includes {
60 self.includes.insert(include.to_string());
61 }
62 }
63
64 fn add_data(&mut self, data: &str) {
65 self.data.push(data.to_string());
66 }
67
68 pub fn add_setup(&mut self, setup: &[&str]) {
69 for line in setup {
70 self.setup.push(line.to_string());
71 }
72 }
73
74 pub fn push_setup<I>(&mut self, setup: I)
75 where
76 I: IntoIterator<Item = String>,
77 {
78 for line in setup {
79 self.setup.push(line);
80 }
81 }
82
83 fn push_data(&mut self, data: String) {
84 self.data.push(data);
85 }
86
87 pub fn push_method_impl<I>(&mut self, method_impl: I)
88 where
89 I: IntoIterator<Item = String>,
90 {
91 for line in method_impl {
92 self.methods_impls.push(line);
93 }
94 }
95}
96
97pub type BlockRes = Result<BuildingBlock, BuildingBlockError>;
98type BlockGetter<'a, T> = &'a dyn Fn(&str, &T, &GenConfig) -> BlockRes;
99macro_rules! block_getter {
100 ($t:ty) => {
101 &'a dyn Fn(&str, &$t, &GenConfig) -> BlockRes
102 }
103}
104macro_rules! block_getter_no_context {
105 ($t:ty) => {
106 &'a dyn Fn(&str, &$t) -> BlockRes
107 };
108}
109
110macro_rules! block_accessers {
111 ($name:ident, $setter:ident, $config:ty) => {
112 pub fn $setter(&mut self, block: block_getter!($config)) {
113 self.$name = Some(block);
114 }
115
116 pub fn $name(&self, name: &str, config: &$config, gen_config: &GenConfig) -> BlockRes {
117 if self.$name.is_none() {
118 Err(BuildingBlockError::BlockMissing(
119 stringify!($name).to_string(),
120 self.name.clone(),
121 ))?
122 }
123 Ok(self.$name.unwrap()(name, config, gen_config)?)
124 }
125 };
126}
127
128pub struct VectorConfig<'a> {
129 pub dof_handler: &'a str,
130 pub is_unknown: bool,
131}
132
133pub struct MatrixConfig<'a> {
134 pub sparsity_pattern: &'a str,
135 pub dof_handler: &'a str,
137}
138
139#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140pub enum ShapeMatrix {
141 Laplace,
142 Mass,
143}
144
145impl Display for ShapeMatrix {
146 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147 let s = format!("{self:?}");
148 write!(f, "{}", s.to_lowercase())
149 }
150}
151
152pub struct ShapeMatrixConfig<'a> {
153 pub dof_handler: &'a str,
154 pub element: &'a str,
155 pub matrix_config: &'a MatrixConfig<'a>,
156 pub kind: ShapeMatrix,
157}
158
159pub struct DofHandlerConfig<'a> {
160 pub mesh: &'a str,
161 pub element: &'a str,
162}
163
164pub struct SparsityPatternConfig<'a> {
165 pub dof_handler: &'a str,
166}
167
168pub struct InitialConditionConfig<'a> {
169 pub dof_handler: &'a str,
170 pub function: &'a str,
171 pub element: &'a str,
172 pub target: &'a str,
173}
174
175#[derive(Clone)]
176pub enum Block<'a> {
177 Matrix(&'a MatrixConfig<'a>),
178 Vector(&'a VectorConfig<'a>),
179 SolveUnknown(&'a SolveUnknownConfig<'a>),
180 EquationSetup(&'a EquationSetupConfig<'a>),
181 Parameter(f64),
182 Function(&'a FunctionDef),
183 AppyBoundaryCondition(&'a ApplyBoundaryConditionConfig<'a>),
184 InitialCondition(&'a InitialConditionConfig<'a>),
185}
186
187pub struct SolveUnknownConfig<'a> {
188 pub dof_handler: &'a str,
189 pub rhs: &'a str,
190 pub unknown_vec: &'a str,
191 pub unknown_mat: &'a str,
192}
193
194pub struct EquationSetupConfig<'a> {
195 pub equation: &'a Equation,
196 pub unknown: &'a dyn Expr,
197 pub vectors: &'a [&'a dyn Expr],
198 pub matrixes: &'a [&'a dyn Expr],
199}
200
201pub struct VectorFromFnConfig<'a> {
202 pub function: &'a str,
203 pub dof_handler: &'a str,
204 pub element: &'a str,
205}
206
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
208pub enum TargetIteration {
209 Next,
210 Current,
211 Previous,
212}
213
214pub struct ApplyBoundaryConditionConfig<'a> {
215 pub function: &'a str,
216 pub dof_handler: &'a str,
217 pub matrix: &'a str,
218 pub solution: &'a str,
219 pub rhs: &'a str,
220}
221
222#[derive(Clone)]
223pub struct BuildingBlockFactory<'a> {
224 name: String,
225 mesh: HashMap<String, BlockGetter<'a, dyn Mesh>>,
226 vector: Option<block_getter!(VectorConfig)>,
227 matrix: Option<block_getter!(MatrixConfig)>,
228 dof_handler: Option<block_getter!(DofHandlerConfig)>,
229 finite_element: Option<block_getter!(FiniteElement)>,
230 sparsity_pattern: Option<block_getter!(SparsityPatternConfig)>,
231 shape_matrix: Option<&'a dyn Fn(&str, BuildingBlock, &ShapeMatrixConfig) -> BlockRes>,
232 solve_unknown: Option<block_getter!(SolveUnknownConfig)>,
233 equation_setup: Option<block_getter!(EquationSetupConfig)>,
234 call: Option<block_getter_no_context!([&str])>,
235 parameter: Option<block_getter!(f64)>,
236 function: Option<block_getter!(FunctionDef)>,
237 vector_from_function: Option<block_getter!(VectorFromFnConfig)>,
238 apply_boundary_condition: Option<block_getter!(ApplyBoundaryConditionConfig)>,
239 initial_condition: Option<block_getter!(InitialConditionConfig)>,
240 add_vector_output: Option<block_getter!(str)>,
241}
242
243impl<'a> BuildingBlockFactory<'a> {
244 pub fn new(name: &str) -> Self {
245 BuildingBlockFactory {
246 name: name.to_string(),
247 mesh: HashMap::new(),
248 vector: None,
249 matrix: None,
250 dof_handler: None,
251 finite_element: None,
252 sparsity_pattern: None,
253 shape_matrix: None,
254 solve_unknown: None,
255 equation_setup: None,
256 call: Some(&|name, args| {
257 let mut block = BuildingBlock::new();
258 block.main.push(format!("{}({});", name, args.join(", ")));
259
260 Ok(block)
261 }),
262 parameter: None,
263 function: None,
264 vector_from_function: None,
265 apply_boundary_condition: None,
266 initial_condition: None,
267 add_vector_output: None,
268 }
269 }
270
271 block_accessers!(vector, set_vector, VectorConfig);
272 block_accessers!(dof_handler, set_dof_handler, DofHandlerConfig);
273 block_accessers!(
274 sparsity_pattern,
275 set_sparsity_pattern,
276 SparsityPatternConfig
277 );
278 block_accessers!(
279 vector_from_function,
280 set_vector_from_function,
281 VectorFromFnConfig
282 );
283 block_accessers!(
284 apply_boundary_condition,
285 set_apply_boundary_condition,
286 ApplyBoundaryConditionConfig
287 );
288 block_accessers!(
289 initial_condition,
290 set_initial_condition,
291 InitialConditionConfig
292 );
293 block_accessers!(add_vector_output, set_add_vector_output, str);
294
295 block_accessers!(matrix, set_matrix, MatrixConfig);
296
297 pub fn add_mesh(&mut self, r#type: &str, block: BlockGetter<'a, dyn Mesh>) {
298 self.mesh.insert(r#type.to_string(), block);
299 }
300
301 pub fn mesh(&self, name: &str, config: &'a dyn Mesh, gen_config: &GenConfig) -> BlockRes {
302 let r#type = config.typetag_name();
303 Ok(self.mesh.get(r#type).ok_or(
304 BuildingBlockError::BlockMissing(r#type.to_string(), self.name.clone()),
305 )?(name, config, gen_config)?)
306 }
307
308 block_accessers!(finite_element, set_finite_element, FiniteElement);
309
310 pub fn shape_matrix<'b: 'a>(
311 &self,
312 name: &str,
313 config: &ShapeMatrixConfig<'b>,
314 gen_config: &GenConfig,
315 ) -> BlockRes {
316 if self.shape_matrix.is_none() {
317 Err(BuildingBlockError::BlockMissing(
318 "shape_matrix".to_string(),
319 self.name.clone(),
320 ))?
321 }
322 let matrix = self.matrix(name, &config.matrix_config, gen_config)?;
323 Ok(self.shape_matrix.unwrap()(name, matrix, config)?)
324 }
325
326 pub fn set_shape_matrix(
327 &mut self,
328 block: &'a dyn Fn(&str, BuildingBlock, &ShapeMatrixConfig) -> BlockRes,
329 ) {
330 self.shape_matrix = Some(block);
331 }
332
333 block_accessers!(solve_unknown, set_solve_unknown, SolveUnknownConfig);
334 block_accessers!(equation_setup, set_equation_setup, EquationSetupConfig);
335
336 pub fn call(&self, name: &str, args: &[&str]) -> BlockRes {
337 if self.call.is_none() {
338 Err(BuildingBlockError::BlockMissing(
339 "call".to_string(),
340 self.name.clone(),
341 ))?
342 }
343 Ok(self.call.unwrap()(name, args)?)
344 }
345
346 pub fn set_call(&mut self, block: block_getter_no_context!([&str])) {
347 self.call = Some(block);
348 }
349
350 pub(crate) fn newline(&self) -> BuildingBlock {
351 let mut block = BuildingBlock::new();
352 block.main.push("\n".to_string());
353 block
354 }
355
356 block_accessers!(parameter, set_parameter, f64);
357
358 pub(crate) fn comment(&self, content: &str) -> BuildingBlock {
359 let mut block = BuildingBlock::new();
360 block.main.push(format!("// {}", content));
361 block
362 }
363
364 block_accessers!(function, set_function, FunctionDef);
365}
366
367#[derive(Error, Debug)]
368pub enum BuildingBlockError {
369 #[error("block '{0}' missing in factory '{1}'")]
370 BlockMissing(String, String),
371 #[error("wrong input supplied to block {0}")]
372 WrongInput(String),
373 #[error("block {0} already exists")]
374 BlockAlreadyExists(String),
375 #[error("name {0} already exists")]
376 NameAlreadyExists(String),
377 #[error("failed to generate expression code")]
378 ExprCodeGen(#[from] ExprCodeGenError),
379}
380
381#[derive(Error, Debug)]
382pub enum ExprCodeGenError {
383 #[error("too many vectors in expr: {0}")]
384 TooManyVectors(Box<dyn Expr>),
385 #[error("too many matrixes in expr: {0}")]
386 TooManyMatrixes(Box<dyn Expr>),
387 #[error("unsupported expr: {0}")]
388 UnsupportedExpr(Box<dyn Expr>),
389 #[error("unsupported operand in multiplication: {0}")]
390 UnsupportedMulOperand(Box<dyn Expr>),
391 #[error("can't multiply two vectors")]
392 VectorMul,
393 #[error("operations resulted in a matrix when a vector was expected")]
394 MatResult,
395 #[error("operations resulted in a vector when a matrix was expected")]
396 VecResult,
397 #[error("unsupported equation : {reason}\nequation: {equation}")]
398 UnsupportedEquation { reason: String, equation: Equation },
399}