Skip to main content

xc3_model/material/
assignments.rs

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/// Assignment information for the channels of each output.
14/// This includes channels from textures, material parameters, or shader constants.
15#[derive(Debug, Clone, PartialEq)]
16pub struct OutputAssignments {
17    pub output_assignments: [OutputAssignment; 6],
18
19    // TODO: make this the same type as normal intensity.
20    /// The parameter multiplied by vertex alpha to determine outline width.
21    pub outline_width: Option<AssignmentValue>,
22
23    /// Index into [assignments](#structfield.assignments) for the intensity map for normal mapping.
24    pub normal_intensity: Option<usize>,
25
26    /// Index into [assignments](#structfield.assignments) for the intensity for vValInf normal mapping.
27    pub val_inf_intensity: Option<usize>,
28
29    /// Unique values shared between all outputs.
30    pub assignments: Vec<Assignment>,
31}
32
33impl OutputAssignments {
34    /// Calculate the material ID from a hardcoded shader constant if present.
35    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            // TODO: Why is this sometimes 7?
40            Some((v.0 * 255.0 + 0.1) as u32 & 0x7)
41        } else {
42            None
43        }
44    }
45}
46
47// TODO: Come up with better names.
48#[derive(Debug, Clone, PartialEq, Default)]
49pub struct OutputAssignment {
50    /// Index into [assignments](struct.OutputAssignments.html#structfield.assignments) for the x value.
51    pub x: Option<usize>,
52    /// Index into [assignments](struct.OutputAssignments.html#structfield.assignments) for the y value.
53    pub y: Option<usize>,
54    /// Index into [assignments](struct.OutputAssignments.html#structfield.assignments) for the z value.
55    pub z: Option<usize>,
56    /// Index into [assignments](struct.OutputAssignments.html#structfield.assignments) for the w value.
57    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        /// Index into [assignments](struct.OutputAssignments.html#structfield.assignments)
66        /// for the function argument list `[arg0, arg1, ...]`.
67        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    /// The name of the texture like `s0` or `gTResidentTex09`.
85    pub name: SmolStr,
86    pub channel: Option<char>,
87    /// Indices into [assigmments](struct.OutputAssignments.html#structfield.assignments)
88    /// for the texture coordinates.
89    pub texcoords: Vec<usize>,
90}
91
92/// Assignment information for the channels of each output.
93/// This includes channels from textures, material parameters, or shader constants.
94#[derive(Debug, Clone, PartialEq)]
95pub struct OutputAssignmentXyz {
96    /// Index into [assignments](#structfield.assignments) for the most recent assignment.
97    pub assignment: usize, // TODO: option int?
98
99    /// Unique shared values.
100    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        /// Index into XYZ [assignments](struct.OutputAssignmentXyz.html#structfield.assignments)
109        /// for the function argument list `[arg0, arg1, ...]`.
110        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    /// The name of the texture like `s0` or `gTResidentTex09`.
127    pub name: SmolStr,
128    pub channel: Option<ChannelXyz>,
129    /// Indices into scalar [assigmments](struct.OutputAssignments.html#structfield.assignments)
130    /// for the texture coordinates.
131    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    // Avoid processing the same set of assignments more than once.
193    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                    // TODO: Check that channels are one of the supported channels.
236                    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(), // TODO: These should refer to the scalar assignments?
247                                };
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    // Skip any operations that involve vector arguments.
295    // TODO: Should all operations be supported?
296    // TODO: Combine xyz and scalar operations?
297    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    // Use the existing indices to avoid costly caching or recursion.
343    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    // No assignment data is available.
407    // Guess reasonable defaults based on the texture names or types.
408    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            // TODO: Why does this index out of range for xc2 legacy mxmd?
439            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    // This may only have two channels since BC5 is common.
447    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}