1#![allow(dead_code)]
10
11use alloc::{collections::BTreeMap, string::String, vec::Vec};
12
13const FP_SHIFT: i32 = 16;
19
20const FP_ONE: i32 = 1 << FP_SHIFT;
22
23fn fp_mul(a: i32, b: i32) -> i32 {
25 ((a as i64 * b as i64) >> FP_SHIFT) as i32
26}
27
28fn fp_from_int(v: i32) -> i32 {
30 v << FP_SHIFT
31}
32
33fn fp_to_int(v: i32) -> i32 {
35 v >> FP_SHIFT
36}
37
38fn fp_rcp(x: i32) -> i32 {
40 if x == 0 {
41 return 0;
42 }
43 ((FP_ONE as i64 * FP_ONE as i64) / x as i64) as i32
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum ShaderType {
53 Vertex,
55 Fragment,
57}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum UniformValue {
62 Int(i32),
64 Vec2(i32, i32),
66 Vec4(i32, i32, i32, i32),
68 Mat4([i32; 16]),
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum SrcOperand {
79 Reg(u8),
81 Uniform(u8),
83 Immediate(i32),
85}
86
87impl Default for SrcOperand {
88 fn default() -> Self {
89 Self::Immediate(0)
90 }
91}
92
93#[derive(Debug, Clone, PartialEq, Eq)]
95pub enum TgsiInstruction {
96 MOV { dst: u8, src: SrcOperand },
98 ADD {
100 dst: u8,
101 src0: SrcOperand,
102 src1: SrcOperand,
103 },
104 MUL {
106 dst: u8,
107 src0: SrcOperand,
108 src1: SrcOperand,
109 },
110 MAD {
112 dst: u8,
113 src0: SrcOperand,
114 src1: SrcOperand,
115 src2: SrcOperand,
116 },
117 DP3 {
119 dst: u8,
120 src0: SrcOperand,
121 src1: SrcOperand,
122 },
123 DP4 {
125 dst: u8,
126 src0: SrcOperand,
127 src1: SrcOperand,
128 },
129 TEX { dst: u8, coord: SrcOperand },
131 SAMPLE {
133 dst: u8,
134 coord: SrcOperand,
135 sampler: u8,
136 },
137 RCP { dst: u8, src: SrcOperand },
139 RSQ { dst: u8, src: SrcOperand },
141}
142
143#[derive(Debug, Clone)]
149pub struct ShaderProgram {
150 pub shader_type: ShaderType,
152 pub instructions: Vec<TgsiInstruction>,
154 pub uniforms: BTreeMap<String, u8>,
156 pub label: String,
158}
159
160impl ShaderProgram {
161 pub fn new(shader_type: ShaderType, label: &str) -> Self {
163 Self {
164 shader_type,
165 instructions: Vec::new(),
166 uniforms: BTreeMap::new(),
167 label: String::from(label),
168 }
169 }
170
171 pub fn push(&mut self, instr: TgsiInstruction) {
173 self.instructions.push(instr);
174 }
175
176 pub fn bind_uniform(&mut self, name: &str, slot: u8) {
178 self.uniforms.insert(String::from(name), slot);
179 }
180
181 pub fn instruction_count(&self) -> usize {
183 self.instructions.len()
184 }
185}
186
187#[derive(Debug, Clone, PartialEq, Eq)]
193pub enum ShaderOp {
194 SetColor { r: u8, g: u8, b: u8, a: u8 },
196 TintByUniform { uniform_name: String },
198 SampleTexture { sampler: u8 },
200 VerticalGradient { top: u32, bottom: u32 },
202 Passthrough,
204}
205
206pub struct ShaderCompiler;
208
209impl ShaderCompiler {
210 pub fn compile_fragment(label: &str, ops: &[ShaderOp]) -> ShaderProgram {
212 let mut prog = ShaderProgram::new(ShaderType::Fragment, label);
213
214 for op in ops {
215 match op {
216 ShaderOp::SetColor { r, g, b, a } => {
217 prog.push(TgsiInstruction::MOV {
218 dst: 0,
219 src: SrcOperand::Immediate(fp_from_int(*r as i32)),
220 });
221 prog.push(TgsiInstruction::MOV {
222 dst: 1,
223 src: SrcOperand::Immediate(fp_from_int(*g as i32)),
224 });
225 prog.push(TgsiInstruction::MOV {
226 dst: 2,
227 src: SrcOperand::Immediate(fp_from_int(*b as i32)),
228 });
229 prog.push(TgsiInstruction::MOV {
230 dst: 3,
231 src: SrcOperand::Immediate(fp_from_int(*a as i32)),
232 });
233 }
234 ShaderOp::TintByUniform { uniform_name } => {
235 let slot = prog.uniforms.len() as u8;
236 prog.bind_uniform(uniform_name, slot);
237 prog.push(TgsiInstruction::MUL {
238 dst: 0,
239 src0: SrcOperand::Reg(0),
240 src1: SrcOperand::Uniform(slot),
241 });
242 }
243 ShaderOp::SampleTexture { sampler } => {
244 prog.push(TgsiInstruction::SAMPLE {
245 dst: 0,
246 coord: SrcOperand::Reg(4), sampler: *sampler,
248 });
249 }
250 ShaderOp::VerticalGradient { top, bottom } => {
251 let tr = fp_from_int(((*top >> 16) & 0xFF) as i32);
252 let br = fp_from_int(((*bottom >> 16) & 0xFF) as i32);
253 prog.push(TgsiInstruction::MOV {
255 dst: 0,
256 src: SrcOperand::Immediate(tr),
257 });
258 prog.push(TgsiInstruction::MOV {
259 dst: 5,
260 src: SrcOperand::Immediate(br - tr),
261 });
262 prog.push(TgsiInstruction::MAD {
263 dst: 0,
264 src0: SrcOperand::Reg(5),
265 src1: SrcOperand::Reg(6), src2: SrcOperand::Reg(0),
267 });
268 }
269 ShaderOp::Passthrough => {
270 }
272 }
273 }
274
275 prog
276 }
277
278 pub fn compile_vertex_passthrough(label: &str) -> ShaderProgram {
280 let mut prog = ShaderProgram::new(ShaderType::Vertex, label);
281 prog.push(TgsiInstruction::MOV {
283 dst: 0,
284 src: SrcOperand::Reg(0),
285 });
286 prog
287 }
288}
289
290const MAX_REGISTERS: usize = 16;
296
297pub struct ShaderExecutor {
299 registers: [i32; MAX_REGISTERS],
301 uniform_values: Vec<UniformValue>,
303}
304
305impl Default for ShaderExecutor {
306 fn default() -> Self {
307 Self::new()
308 }
309}
310
311impl ShaderExecutor {
312 pub fn new() -> Self {
314 Self {
315 registers: [0; MAX_REGISTERS],
316 uniform_values: Vec::new(),
317 }
318 }
319
320 pub fn set_uniform(&mut self, slot: usize, value: UniformValue) {
322 while self.uniform_values.len() <= slot {
323 self.uniform_values.push(UniformValue::Int(0));
324 }
325 self.uniform_values[slot] = value;
326 }
327
328 pub fn set_register(&mut self, reg: u8, value: i32) {
330 if (reg as usize) < MAX_REGISTERS {
331 self.registers[reg as usize] = value;
332 }
333 }
334
335 pub fn get_register(&self, reg: u8) -> i32 {
337 if (reg as usize) < MAX_REGISTERS {
338 self.registers[reg as usize]
339 } else {
340 0
341 }
342 }
343
344 fn read_src(&self, src: &SrcOperand) -> i32 {
346 match src {
347 SrcOperand::Reg(r) => self.get_register(*r),
348 SrcOperand::Uniform(slot) => {
349 let s = *slot as usize;
350 if s < self.uniform_values.len() {
351 match &self.uniform_values[s] {
352 UniformValue::Int(v) => *v,
353 UniformValue::Vec2(x, _) => *x,
354 UniformValue::Vec4(x, _, _, _) => *x,
355 UniformValue::Mat4(m) => m[0],
356 }
357 } else {
358 0
359 }
360 }
361 SrcOperand::Immediate(v) => *v,
362 }
363 }
364
365 pub fn execute(&mut self, program: &ShaderProgram) {
370 for instr in &program.instructions {
371 match instr {
372 TgsiInstruction::MOV { dst, src } => {
373 let v = self.read_src(src);
374 self.registers[*dst as usize % MAX_REGISTERS] = v;
375 }
376 TgsiInstruction::ADD { dst, src0, src1 } => {
377 let a = self.read_src(src0);
378 let b = self.read_src(src1);
379 self.registers[*dst as usize % MAX_REGISTERS] = a.saturating_add(b);
380 }
381 TgsiInstruction::MUL { dst, src0, src1 } => {
382 let a = self.read_src(src0);
383 let b = self.read_src(src1);
384 self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
385 }
386 TgsiInstruction::MAD {
387 dst,
388 src0,
389 src1,
390 src2,
391 } => {
392 let a = self.read_src(src0);
393 let b = self.read_src(src1);
394 let c = self.read_src(src2);
395 self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b).saturating_add(c);
396 }
397 TgsiInstruction::DP3 { dst, src0, src1 } => {
398 let a = self.read_src(src0);
400 let b = self.read_src(src1);
401 self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
402 }
403 TgsiInstruction::DP4 { dst, src0, src1 } => {
404 let a = self.read_src(src0);
405 let b = self.read_src(src1);
406 self.registers[*dst as usize % MAX_REGISTERS] = fp_mul(a, b);
407 }
408 TgsiInstruction::TEX { dst, coord } => {
409 let c = self.read_src(coord);
411 self.registers[*dst as usize % MAX_REGISTERS] = c;
412 }
413 TgsiInstruction::SAMPLE {
414 dst,
415 coord,
416 sampler: _,
417 } => {
418 let c = self.read_src(coord);
419 self.registers[*dst as usize % MAX_REGISTERS] = c;
420 }
421 TgsiInstruction::RCP { dst, src } => {
422 let v = self.read_src(src);
423 self.registers[*dst as usize % MAX_REGISTERS] = fp_rcp(v);
424 }
425 TgsiInstruction::RSQ { dst, src } => {
426 let v = self.read_src(src);
427 if v <= 0 {
430 self.registers[*dst as usize % MAX_REGISTERS] = 0;
431 } else {
432 let mut guess = FP_ONE;
434 let mut test = fp_to_int(v);
435 while test > 1 {
436 test >>= 2;
437 guess >>= 1;
438 }
439 if guess == 0 {
440 guess = 1;
441 }
442 self.registers[*dst as usize % MAX_REGISTERS] = fp_rcp(guess);
443 }
444 }
445 }
446 }
447 }
448
449 pub fn execute_fragment(&mut self, program: &ShaderProgram) -> u32 {
453 self.execute(program);
454
455 let r = fp_to_int(self.registers[0]).clamp(0, 255) as u32;
456 let g = fp_to_int(self.registers[1]).clamp(0, 255) as u32;
457 let b = fp_to_int(self.registers[2]).clamp(0, 255) as u32;
458 let a = fp_to_int(self.registers[3]).clamp(0, 255) as u32;
459
460 (a << 24) | (r << 16) | (g << 8) | b
461 }
462}
463
464#[cfg(test)]
469mod tests {
470 use super::*;
471
472 #[test]
473 fn test_fp_mul() {
474 let a = fp_from_int(3);
475 let b = fp_from_int(4);
476 assert_eq!(fp_to_int(fp_mul(a, b)), 12);
477 }
478
479 #[test]
480 fn test_fp_rcp() {
481 let v = fp_from_int(2);
482 let r = fp_rcp(v);
483 assert!((r - (FP_ONE / 2)).abs() < 2);
485 }
486
487 #[test]
488 fn test_shader_set_color() {
489 let prog = ShaderCompiler::compile_fragment(
490 "test",
491 &[ShaderOp::SetColor {
492 r: 128,
493 g: 64,
494 b: 32,
495 a: 255,
496 }],
497 );
498 assert!(prog.instruction_count() >= 4);
499 let mut exec = ShaderExecutor::new();
500 let pixel = exec.execute_fragment(&prog);
501 let r = (pixel >> 16) & 0xFF;
502 let g = (pixel >> 8) & 0xFF;
503 let b = pixel & 0xFF;
504 assert_eq!(r, 128);
505 assert_eq!(g, 64);
506 assert_eq!(b, 32);
507 }
508
509 #[test]
510 fn test_shader_passthrough() {
511 let prog = ShaderCompiler::compile_fragment("pass", &[ShaderOp::Passthrough]);
512 let mut exec = ShaderExecutor::new();
513 exec.set_register(0, fp_from_int(200));
514 exec.set_register(1, fp_from_int(100));
515 exec.set_register(2, fp_from_int(50));
516 exec.set_register(3, fp_from_int(255));
517 let pixel = exec.execute_fragment(&prog);
518 let r = (pixel >> 16) & 0xFF;
519 assert_eq!(r, 200);
520 }
521
522 #[test]
523 fn test_vertex_passthrough() {
524 let prog = ShaderCompiler::compile_vertex_passthrough("vtx");
525 assert_eq!(prog.shader_type, ShaderType::Vertex);
526 assert_eq!(prog.instruction_count(), 1);
527 }
528
529 #[test]
530 fn test_mov_instruction() {
531 let mut exec = ShaderExecutor::new();
532 let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
533 prog.push(TgsiInstruction::MOV {
534 dst: 0,
535 src: SrcOperand::Immediate(fp_from_int(42)),
536 });
537 exec.execute(&prog);
538 assert_eq!(fp_to_int(exec.get_register(0)), 42);
539 }
540
541 #[test]
542 fn test_add_instruction() {
543 let mut exec = ShaderExecutor::new();
544 exec.set_register(0, fp_from_int(10));
545 let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
546 prog.push(TgsiInstruction::ADD {
547 dst: 1,
548 src0: SrcOperand::Reg(0),
549 src1: SrcOperand::Immediate(fp_from_int(5)),
550 });
551 exec.execute(&prog);
552 assert_eq!(fp_to_int(exec.get_register(1)), 15);
553 }
554
555 #[test]
556 fn test_uniform_binding() {
557 let mut exec = ShaderExecutor::new();
558 exec.set_uniform(0, UniformValue::Int(fp_from_int(2)));
559 exec.set_register(0, fp_from_int(100));
560 let mut prog = ShaderProgram::new(ShaderType::Fragment, "test");
561 prog.bind_uniform("scale", 0);
562 prog.push(TgsiInstruction::MUL {
563 dst: 0,
564 src0: SrcOperand::Reg(0),
565 src1: SrcOperand::Uniform(0),
566 });
567 exec.execute(&prog);
568 assert_eq!(fp_to_int(exec.get_register(0)), 200);
569 }
570}