1use indexmap::{IndexMap, IndexSet};
2use ordered_float::OrderedFloat;
3use smol_str::SmolStr;
4use xc3_lib::mxmd::TextureUsage;
5
6use crate::{
7 ImageTexture,
8 shader_database::{Dependency, Operation, OutputExpr, ShaderProgram, TextureDependency},
9};
10
11use super::{MaterialParameters, Texture};
12
13#[derive(Debug, Clone, PartialEq)]
16pub struct OutputAssignments {
17 pub output_assignments: [OutputAssignment; 6],
18
19 pub outline_width: Option<AssignmentValue>,
22
23 pub normal_intensity: Option<usize>,
25
26 pub val_inf_intensity: Option<usize>,
28
29 pub assignments: Vec<Assignment>,
31}
32
33impl OutputAssignments {
34 pub fn mat_id(&self) -> Option<u32> {
36 if let Assignment::Value(Some(AssignmentValue::Float(v))) =
37 self.assignments.get(self.output_assignments[1].w?)?
38 {
39 Some((v.0 * 255.0 + 0.1) as u32 & 0x7)
41 } else {
42 None
43 }
44 }
45}
46
47#[derive(Debug, Clone, PartialEq, Default)]
49pub struct OutputAssignment {
50 pub x: Option<usize>,
52 pub y: Option<usize>,
54 pub z: Option<usize>,
56 pub w: Option<usize>,
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Hash)]
61pub enum Assignment {
62 Value(Option<AssignmentValue>),
63 Func {
64 op: Operation,
65 args: Vec<usize>,
68 },
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Hash)]
72pub enum AssignmentValue {
73 Texture(TextureAssignment),
74 Attribute {
75 name: SmolStr,
76 channel: Option<char>,
77 },
78 Float(OrderedFloat<f32>),
79 Int(i32),
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Hash)]
83pub struct TextureAssignment {
84 pub name: SmolStr,
86 pub channel: Option<char>,
87 pub texcoords: Vec<usize>,
90}
91
92#[derive(Debug, Clone, PartialEq)]
95pub struct OutputAssignmentXyz {
96 pub assignment: usize, pub assignments: Vec<AssignmentXyz>,
101}
102
103#[derive(Debug, Clone, PartialEq, Eq, Hash)]
104pub enum AssignmentXyz {
105 Value(Option<AssignmentValueXyz>),
106 Func {
107 op: Operation,
108 args: Vec<usize>,
111 },
112}
113
114#[derive(Debug, Clone, PartialEq, Eq, Hash)]
115pub enum AssignmentValueXyz {
116 Texture(TextureAssignmentXyz),
117 Attribute {
118 name: SmolStr,
119 channel: Option<ChannelXyz>,
120 },
121 Float([OrderedFloat<f32>; 3]),
122}
123
124#[derive(Debug, Clone, PartialEq, Eq, Hash)]
125pub struct TextureAssignmentXyz {
126 pub name: SmolStr,
128 pub channel: Option<ChannelXyz>,
129 pub texcoords: Vec<usize>,
132}
133
134#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
135pub enum ChannelXyz {
136 Xyz,
137 X,
138 Y,
139 Z,
140 W,
141}
142
143impl Default for Assignment {
144 fn default() -> Self {
145 Self::Value(None)
146 }
147}
148
149impl AssignmentValue {
150 pub fn from_dependency(d: &Dependency, parameters: &MaterialParameters) -> Option<Self> {
151 match d {
152 Dependency::Int(i) => Some(Self::Int(*i)),
153 Dependency::Float(f) => Some(Self::Float(f.0.into())),
154 Dependency::Buffer(b) => parameters.get_dependency(b).map(|f| Self::Float(f.into())),
155 Dependency::Texture(texture) => Some(Self::Texture(texture_assignment(texture))),
156 Dependency::Attribute(a) => Some(Self::Attribute {
157 name: a.name.clone(),
158 channel: a.channel,
159 }),
160 }
161 }
162}
163
164impl OutputAssignment {
165 pub fn merge_xyz(&self, assignments: &[Assignment]) -> Option<OutputAssignmentXyz> {
166 let mut assignments_xyz_index = IndexMap::new();
167 let mut assignments_xyz = IndexSet::new();
168 let i = merge_xyz_assignments(
169 self.x?,
170 self.y?,
171 self.z?,
172 assignments,
173 &mut assignments_xyz_index,
174 &mut assignments_xyz,
175 )?;
176
177 Some(OutputAssignmentXyz {
178 assignment: i,
179 assignments: assignments_xyz.into_iter().collect(),
180 })
181 }
182}
183
184fn merge_xyz_assignments(
185 x: usize,
186 y: usize,
187 z: usize,
188 assignments: &[Assignment],
189 assignments_xyz_index: &mut IndexMap<(usize, usize, usize), usize>,
190 assignments_xyz: &mut IndexSet<AssignmentXyz>,
191) -> Option<usize> {
192 match assignments_xyz_index.get(&(x, y, z)) {
194 Some(index) => Some(*index),
195 None => {
196 let x_assignment = assignments.get(x)?;
197 let y_assignment = assignments.get(y)?;
198 let z_assignment = assignments.get(z)?;
199
200 let assignment_xyz = match (x_assignment, y_assignment, z_assignment) {
201 (
202 Assignment::Func {
203 op: op_x,
204 args: args_x,
205 },
206 Assignment::Func {
207 op: op_y,
208 args: args_y,
209 },
210 Assignment::Func {
211 op: op_z,
212 args: args_z,
213 },
214 ) => {
215 let op = op_xyz(*op_x, *op_y, *op_z)?;
216 if args_x.len() == args_y.len() && args_y.len() == args_z.len() {
217 let mut args = Vec::new();
218 for ((x, y), z) in args_x.iter().zip(args_y.iter()).zip(args_z.iter()) {
219 let arg = merge_xyz_assignments(
220 *x,
221 *y,
222 *z,
223 assignments,
224 assignments_xyz_index,
225 assignments_xyz,
226 )?;
227 args.push(arg);
228 }
229 Some(AssignmentXyz::Func { op, args })
230 } else {
231 None
232 }
233 }
234 (Assignment::Value(vx), Assignment::Value(vy), Assignment::Value(vz)) => {
235 match (vx, vy, vz) {
237 (
238 Some(AssignmentValue::Texture(tx)),
239 Some(AssignmentValue::Texture(ty)),
240 Some(AssignmentValue::Texture(tz)),
241 ) => {
242 if tx.texcoords == ty.texcoords && ty.texcoords == tz.texcoords {
243 let t_xyz = TextureAssignmentXyz {
244 name: name_xyz(&tx.name, &ty.name, &tz.name)?,
245 channel: channel_xyz(tx.channel, ty.channel, tz.channel)?,
246 texcoords: tx.texcoords.clone(), };
248 Some(AssignmentXyz::Value(Some(AssignmentValueXyz::Texture(
249 t_xyz,
250 ))))
251 } else {
252 None
253 }
254 }
255 (
256 Some(AssignmentValue::Attribute {
257 name: n_x,
258 channel: c_x,
259 }),
260 Some(AssignmentValue::Attribute {
261 name: n_y,
262 channel: c_y,
263 }),
264 Some(AssignmentValue::Attribute {
265 name: n_z,
266 channel: c_z,
267 }),
268 ) => Some(AssignmentXyz::Value(Some(AssignmentValueXyz::Attribute {
269 name: name_xyz(n_x, n_y, n_z)?,
270 channel: channel_xyz(*c_x, *c_y, *c_z)?,
271 }))),
272 (
273 Some(AssignmentValue::Float(fx)),
274 Some(AssignmentValue::Float(fy)),
275 Some(AssignmentValue::Float(fz)),
276 ) => Some(AssignmentXyz::Value(Some(AssignmentValueXyz::Float([
277 *fx, *fy, *fz,
278 ])))),
279 (None, None, None) => Some(AssignmentXyz::Value(None)),
280 _ => None,
281 }
282 }
283 _ => None,
284 }?;
285
286 let index = assignments_xyz.insert_full(assignment_xyz).0;
287 assignments_xyz_index.insert((x, y, z), index);
288 Some(index)
289 }
290 }
291}
292
293fn op_xyz(x: Operation, y: Operation, z: Operation) -> Option<Operation> {
294 if x == y
298 && y == z
299 && !matches!(
300 x,
301 Operation::AddNormalX
302 | Operation::AddNormalY
303 | Operation::ReflectX
304 | Operation::ReflectY
305 | Operation::ReflectZ
306 | Operation::Dot4
307 | Operation::NormalMapX
308 | Operation::NormalMapY
309 | Operation::NormalMapZ,
310 )
311 {
312 Some(x)
313 } else {
314 None
315 }
316}
317
318fn name_xyz(x: &SmolStr, y: &SmolStr, z: &SmolStr) -> Option<SmolStr> {
319 if x == y && y == z {
320 Some(x.clone())
321 } else {
322 None
323 }
324}
325
326fn channel_xyz(x: Option<char>, y: Option<char>, z: Option<char>) -> Option<Option<ChannelXyz>> {
327 match (x, y, z) {
328 (Some('x'), Some('y'), Some('z')) => Some(Some(ChannelXyz::Xyz)),
329 (Some('x'), Some('x'), Some('x')) => Some(Some(ChannelXyz::X)),
330 (Some('y'), Some('y'), Some('y')) => Some(Some(ChannelXyz::Y)),
331 (Some('z'), Some('z'), Some('z')) => Some(Some(ChannelXyz::Z)),
332 (Some('w'), Some('w'), Some('w')) => Some(Some(ChannelXyz::W)),
333 (None, None, None) => Some(None),
334 _ => None,
335 }
336}
337
338pub(crate) fn output_assignments(
339 shader: &ShaderProgram,
340 parameters: &MaterialParameters,
341) -> OutputAssignments {
342 let assignments = shader
344 .exprs
345 .iter()
346 .map(|e| assignment_value(parameters, e))
347 .collect();
348
349 OutputAssignments {
350 output_assignments: [0, 1, 2, 3, 4, 5].map(|i| output_assignment(shader, i)),
351 outline_width: shader
352 .outline_width
353 .as_ref()
354 .and_then(|d| AssignmentValue::from_dependency(d, parameters)),
355 normal_intensity: shader.normal_intensity,
356 val_inf_intensity: shader.val_inf_intensity,
357 assignments,
358 }
359}
360
361fn output_assignment(shader: &ShaderProgram, output_index: usize) -> OutputAssignment {
362 OutputAssignment {
363 x: output_channel_assignment(shader, output_index, 'x'),
364 y: output_channel_assignment(shader, output_index, 'y'),
365 z: output_channel_assignment(shader, output_index, 'z'),
366 w: output_channel_assignment(shader, output_index, 'w'),
367 }
368}
369
370fn output_channel_assignment(
371 shader: &ShaderProgram,
372 output_index: usize,
373 channel: char,
374) -> Option<usize> {
375 let output = format!("o{output_index}.{channel}");
376 shader
377 .output_dependencies
378 .get(&SmolStr::from(output))
379 .copied()
380}
381
382fn assignment_value(parameters: &MaterialParameters, value: &OutputExpr) -> Assignment {
383 match value {
384 crate::shader_database::OutputExpr::Value(d) => {
385 Assignment::Value(AssignmentValue::from_dependency(d, parameters))
386 }
387 crate::shader_database::OutputExpr::Func { op, args } => Assignment::Func {
388 op: *op,
389 args: args.clone(),
390 },
391 }
392}
393
394fn texture_assignment(texture: &TextureDependency) -> TextureAssignment {
395 TextureAssignment {
396 name: texture.name.clone(),
397 channel: texture.channel,
398 texcoords: texture.texcoords.clone(),
399 }
400}
401
402pub(crate) fn infer_assignment_from_textures(
403 textures: &[Texture],
404 image_textures: &[ImageTexture],
405) -> OutputAssignments {
406 let mut assignments = IndexSet::new();
409
410 let mut assignment = |i: Option<usize>, c: usize| {
411 let u = assignments
412 .insert_full(Assignment::Value(Some(AssignmentValue::Attribute {
413 name: "vTex0".into(),
414 channel: Some('x'),
415 })))
416 .0;
417 let v = assignments
418 .insert_full(Assignment::Value(Some(AssignmentValue::Attribute {
419 name: "vTex0".into(),
420 channel: Some('y'),
421 })))
422 .0;
423 Some(
424 assignments
425 .insert_full(Assignment::Value(i.map(|i| {
426 AssignmentValue::Texture(TextureAssignment {
427 name: format!("s{i}").into(),
428 channel: Some(['x', 'y', 'z', 'w'][c]),
429 texcoords: vec![u, v],
430 })
431 })))
432 .0,
433 )
434 };
435
436 let color_index = textures.iter().position(|t| {
437 matches!(
438 image_textures
440 .get(t.image_texture_index)
441 .and_then(|t| t.usage),
442 Some(TextureUsage::Col | TextureUsage::Col2 | TextureUsage::Col3 | TextureUsage::Col4)
443 )
444 });
445
446 let normal_index = textures.iter().position(|t| {
448 matches!(
449 image_textures
450 .get(t.image_texture_index)
451 .and_then(|t| t.usage),
452 Some(TextureUsage::Nrm | TextureUsage::Nrm2)
453 )
454 });
455
456 let spm_index = textures.iter().position(|t| {
457 matches!(
458 image_textures.get(t.image_texture_index).and_then(|t| t.name.as_ref()),
459 Some(name) if name.ends_with("_SPM")
460 )
461 });
462
463 OutputAssignments {
464 output_assignments: [
465 OutputAssignment {
466 x: assignment(color_index, 0),
467 y: assignment(color_index, 1),
468 z: assignment(color_index, 2),
469 w: assignment(color_index, 3),
470 },
471 OutputAssignment::default(),
472 OutputAssignment {
473 x: assignment(normal_index, 0),
474 y: assignment(normal_index, 1),
475 ..Default::default()
476 },
477 OutputAssignment::default(),
478 OutputAssignment::default(),
479 OutputAssignment {
480 x: assignment(spm_index, 0),
481 y: assignment(spm_index, 1),
482 z: assignment(spm_index, 2),
483 ..Default::default()
484 },
485 ],
486 outline_width: None,
487 normal_intensity: None,
488 val_inf_intensity: None,
489 assignments: assignments.into_iter().collect(),
490 }
491}
492
493#[cfg(test)]
494mod tests {
495 use super::*;
496
497 #[test]
498 fn merge_xyz_empty() {
499 assert_eq!(
500 None,
501 OutputAssignment {
502 x: Some(0),
503 y: Some(0),
504 z: Some(0),
505 w: None,
506 }
507 .merge_xyz(&[])
508 );
509 }
510
511 #[test]
512 fn merge_xyz_invalid_channels() {
513 let assignments = [
514 Assignment::Value(Some(AssignmentValue::Float(0.0.into()))),
515 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
516 name: "s0".into(),
517 channel: Some('z'),
518 texcoords: vec![0, 0],
519 }))),
520 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
521 name: "s0".into(),
522 channel: Some('y'),
523 texcoords: vec![0, 0],
524 }))),
525 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
526 name: "s0".into(),
527 channel: Some('x'),
528 texcoords: vec![0, 0],
529 }))),
530 ];
531 assert_eq!(
532 None,
533 OutputAssignment {
534 x: Some(1),
535 y: Some(2),
536 z: Some(3),
537 w: None,
538 }
539 .merge_xyz(&assignments)
540 );
541 }
542
543 #[test]
544 fn merge_xyz_single_channel() {
545 let assignments = [
546 Assignment::Value(Some(AssignmentValue::Float(0.0.into()))),
547 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
548 name: "s0".into(),
549 channel: Some('w'),
550 texcoords: vec![0, 0],
551 }))),
552 ];
553 assert_eq!(
554 Some(OutputAssignmentXyz {
555 assignment: 0,
556 assignments: vec![AssignmentXyz::Value(Some(AssignmentValueXyz::Texture(
557 TextureAssignmentXyz {
558 name: "s0".into(),
559 channel: Some(ChannelXyz::W),
560 texcoords: vec![0, 0]
561 }
562 )))]
563 }),
564 OutputAssignment {
565 x: Some(1),
566 y: Some(1),
567 z: Some(1),
568 w: None,
569 }
570 .merge_xyz(&assignments)
571 );
572 }
573
574 #[test]
575 fn merge_xyz_multiple_channels() {
576 let assignments = [
577 Assignment::Value(Some(AssignmentValue::Float(0.0.into()))),
578 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
579 name: "s0".into(),
580 channel: Some('x'),
581 texcoords: vec![0, 0],
582 }))),
583 Assignment::Value(Some(AssignmentValue::Float(1.0.into()))),
584 Assignment::Value(Some(AssignmentValue::Attribute {
585 name: "vColor".into(),
586 channel: Some('x'),
587 })),
588 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
589 name: "s0".into(),
590 channel: Some('y'),
591 texcoords: vec![0, 0],
592 }))),
593 Assignment::Value(Some(AssignmentValue::Float(2.0.into()))),
594 Assignment::Value(Some(AssignmentValue::Attribute {
595 name: "vColor".into(),
596 channel: Some('y'),
597 })),
598 Assignment::Value(Some(AssignmentValue::Texture(TextureAssignment {
599 name: "s0".into(),
600 channel: Some('z'),
601 texcoords: vec![0, 0],
602 }))),
603 Assignment::Value(Some(AssignmentValue::Float(3.0.into()))),
604 Assignment::Value(Some(AssignmentValue::Attribute {
605 name: "vColor".into(),
606 channel: Some('z'),
607 })),
608 Assignment::Func {
609 op: Operation::Fma,
610 args: vec![1, 2, 3],
611 },
612 Assignment::Func {
613 op: Operation::Fma,
614 args: vec![4, 5, 6],
615 },
616 Assignment::Func {
617 op: Operation::Fma,
618 args: vec![7, 8, 9],
619 },
620 ];
621 assert_eq!(
622 Some(OutputAssignmentXyz {
623 assignment: 3,
624 assignments: vec![
625 AssignmentXyz::Value(Some(AssignmentValueXyz::Texture(TextureAssignmentXyz {
626 name: "s0".into(),
627 channel: Some(ChannelXyz::Xyz),
628 texcoords: vec![0, 0]
629 }))),
630 AssignmentXyz::Value(Some(AssignmentValueXyz::Float([
631 1.0.into(),
632 2.0.into(),
633 3.0.into()
634 ]))),
635 AssignmentXyz::Value(Some(AssignmentValueXyz::Attribute {
636 name: "vColor".into(),
637 channel: Some(ChannelXyz::Xyz)
638 })),
639 AssignmentXyz::Func {
640 op: Operation::Fma,
641 args: vec![0, 1, 2]
642 }
643 ]
644 }),
645 OutputAssignment {
646 x: Some(10),
647 y: Some(11),
648 z: Some(12),
649 w: None,
650 }
651 .merge_xyz(&assignments)
652 );
653 }
654}