1use std::collections::{BTreeMap, BTreeSet};
2use std::fs;
3use std::path::{Path, PathBuf};
4
5use crate::disasm::ShaderKind;
6
7#[derive(Debug, Clone)]
8pub struct StructField {
9 pub ty: String,
10 pub name: String,
11 pub semantic: String,
12}
13
14#[derive(Debug, Clone)]
15struct ParsedHlsl {
16 samplers: Vec<String>,
17 consts: Vec<(String, String, String)>,
18 input_name: String,
19 output_name: String,
20 input_fields: Vec<StructField>,
21 output_fields: Vec<StructField>,
22 body_lines: Vec<String>,
23 register_count: usize,
24}
25
26pub fn discover_reference_hlsl_roots(input: &Path) -> Vec<PathBuf> {
27 let mut out = Vec::new();
28 let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
29 let candidates = [
30 cwd.join("reference_hlsl"),
31 cwd.join("output").join("hlsl"),
32 input
33 .parent()
34 .unwrap_or(Path::new("."))
35 .join("reference_hlsl"),
36 input
37 .parent()
38 .unwrap_or(Path::new("."))
39 .join("output")
40 .join("hlsl"),
41 ];
42 for c in candidates {
43 if c.is_dir() && !out.iter().any(|x| x == &c) {
44 out.push(c);
45 }
46 }
47 out
48}
49
50pub fn load_reference_hlsl(roots: &[PathBuf], prefix: &str) -> Option<String> {
51 for root in roots {
52 let path = root.join(format!("{prefix}.hlsl"));
53 if let Ok(s) = fs::read_to_string(&path) {
54 return Some(s);
55 }
56 }
57 None
58}
59
60pub fn transpile_reference_hlsl_to_wgsl(src: &str, kind: ShaderKind) -> Result<String, String> {
61 let parsed = parse_hlsl(src)?;
62 Ok(emit_wgsl(&parsed, kind))
63}
64
65fn parse_hlsl(src: &str) -> Result<ParsedHlsl, String> {
66 let mut samplers = Vec::new();
67 let mut consts = Vec::new();
68 for line in src.lines() {
69 let t = line.trim();
70 if let Some(name) = t
71 .strip_prefix("uniform sampler2D ")
72 .and_then(|s| s.strip_suffix(';'))
73 {
74 samplers.push(name.trim().to_string());
75 } else if let Some(rest) = t.strip_prefix("static const ")
76 && let Some((lhs, rhs)) = rest.split_once('=')
77 {
78 let lhs = lhs.trim();
79 let rhs = rhs.trim().trim_end_matches(';').trim();
80 let parts = lhs.split_whitespace().collect::<Vec<_>>();
81 if parts.len() == 2 {
82 consts.push((parts[1].to_string(), parts[0].to_string(), rhs.to_string()));
83 }
84 }
85 }
86
87 let (input_name, input_fields) = parse_struct(src, "VS_INPUT")
88 .or_else(|| parse_struct(src, "PS_INPUT"))
89 .ok_or("missing input struct")?;
90 let (output_name, output_fields) = parse_struct(src, "VS_OUTPUT")
91 .or_else(|| parse_struct(src, "PS_OUTPUT"))
92 .ok_or("missing output struct")?;
93
94 let body = extract_main_body(src)?;
95 let mut body_lines = Vec::new();
96 let const_names: BTreeSet<String> = consts.iter().map(|(n, _, _)| n.clone()).collect();
97 let mut max_reg = 0usize;
98 for raw in body.lines() {
99 let t = raw.trim();
100 if t.is_empty() {
101 continue;
102 }
103 if t == format!("{} output;", output_name)
104 || t.starts_with("output.")
105 && t.contains("= ")
106 && t.contains("float4(0.0, 0.0, 0.0, 0.0)")
107 {
108 continue;
109 }
110 let line = translate_statement(t, &const_names, &mut max_reg)?;
111 if !line.is_empty() {
112 body_lines.push(line);
113 }
114 }
115
116 Ok(ParsedHlsl {
117 samplers,
118 consts,
119 input_name,
120 output_name,
121 input_fields,
122 output_fields,
123 body_lines,
124 register_count: max_reg.max(1),
125 })
126}
127
128fn parse_struct(src: &str, name: &str) -> Option<(String, Vec<StructField>)> {
129 let marker = format!("struct {} {{", name);
130 let start = src.find(&marker)?;
131 let rest = &src[start + marker.len()..];
132 let end = rest.find("};")?;
133 let body = &rest[..end];
134 let mut fields = Vec::new();
135 for line in body.lines() {
136 let t = line.trim();
137 if t.is_empty() {
138 continue;
139 }
140 let t = t.trim_end_matches(';');
141 let (lhs, semantic) = t.split_once(':')?;
142 let parts = lhs.split_whitespace().collect::<Vec<_>>();
143 if parts.len() != 2 {
144 continue;
145 }
146 fields.push(StructField {
147 ty: parts[0].to_string(),
148 name: parts[1].to_string(),
149 semantic: semantic.trim().to_string(),
150 });
151 }
152 Some((name.to_string(), fields))
153}
154
155fn extract_main_body(src: &str) -> Result<String, String> {
156 let main_pos = src
157 .find(" main(")
158 .or_else(|| src.find("main("))
159 .ok_or("missing main")?;
160 let rest = &src[main_pos..];
161 let brace = rest.find('{').ok_or("missing main body")?;
162 let mut depth = 0i32;
163 let mut end_idx = None;
164 for (i, ch) in rest[brace..].char_indices() {
165 match ch {
166 '{' => depth += 1,
167 '}' => {
168 depth -= 1;
169 if depth == 0 {
170 end_idx = Some(brace + i);
171 break;
172 }
173 }
174 _ => {}
175 }
176 }
177 let end = end_idx.ok_or("unterminated main body")?;
178 Ok(rest[brace + 1..end].to_string())
179}
180
181fn translate_statement(
182 line: &str,
183 const_names: &BTreeSet<String>,
184 max_reg: &mut usize,
185) -> Result<String, String> {
186 let mut out = line.to_string();
187 out = out.trim_end_matches(';').to_string();
188 out = replace_types(&out);
189 out = replace_intrinsics(&out);
190 out = replace_sampler_calls(&out);
191 out = replace_registers(&out, const_names, max_reg);
192 if out.contains('?') {
193 out = replace_ternary_expr(&out)?;
194 }
195 Ok(format!(" {};;", out).replace(";;", ";"))
196}
197
198fn replace_types(s: &str) -> String {
199 let mut out = s.to_string();
200 for (a, b) in [
201 ("float4(", "vec4<f32>("),
202 ("float3(", "vec3<f32>("),
203 ("float2(", "vec2<f32>("),
204 ("float4 ", "vec4<f32> "),
205 ("float3 ", "vec3<f32> "),
206 ("float2 ", "vec2<f32> "),
207 ("float ", "f32 "),
208 ("int ", "i32 "),
209 ] {
210 out = out.replace(a, b);
211 }
212 out
213}
214
215fn replace_intrinsics(s: &str) -> String {
216 let mut out = s.to_string();
217 out = out.replace("lerp(", "mix(");
218 out = out.replace("rsqrt(", "inverseSqrt(");
219 out = out.replace("frac(", "fract(");
220 while let Some(pos) = out.find("saturate(") {
221 let start = pos + "saturate(".len();
222 let end = find_matching_paren(&out, start - 1).unwrap_or(out.len() - 1);
223 let inner = &out[start..end];
224 let repl = format!("clamp({}, 0.0, 1.0)", inner);
225 out.replace_range(pos..=end, &repl);
226 }
227 out
228}
229
230fn replace_sampler_calls(s: &str) -> String {
231 let mut out = String::new();
232 let bytes = s.as_bytes();
233 let mut i = 0usize;
234 while i < bytes.len() {
235 if s[i..].starts_with("tex2D(") {
236 let args_start = i + "tex2D(".len();
237 let end = find_matching_paren(s, args_start - 1).unwrap_or(s.len() - 1);
238 let args = &s[args_start..end];
239 if let Some((sampler, coord)) = split_top_level_once(args, ',') {
240 out.push_str(&format!(
241 "textureSample(tex_{}, samp_{}, {})",
242 sampler.trim(),
243 sampler.trim(),
244 coord.trim()
245 ));
246 i = end + 1;
247 continue;
248 }
249 }
250 out.push(bytes[i] as char);
251 i += 1;
252 }
253 out
254}
255
256fn replace_registers(s: &str, const_names: &BTreeSet<String>, max_reg: &mut usize) -> String {
257 let mut out = String::new();
258 let chars = s.chars().collect::<Vec<_>>();
259 let mut i = 0usize;
260 while i < chars.len() {
261 if chars[i] == 'c' {
262 let start = i;
263 let mut j = i + 1;
264 while j < chars.len() && chars[j].is_ascii_digit() {
265 j += 1;
266 }
267 if j > i + 1 {
268 let name = chars[start..j].iter().collect::<String>();
269 if !const_names.contains(&name)
270 && let Ok(idx) = name[1..].parse::<usize>()
271 {
272 *max_reg = (*max_reg).max(idx + 1);
273 out.push_str(&format!("u.c[{}]", idx));
274 i = j;
275 continue;
276 }
277 }
278 }
279 out.push(chars[i]);
280 i += 1;
281 }
282 out
283}
284
285fn replace_ternary_expr(s: &str) -> Result<String, String> {
286 let q = s.find('?').ok_or_else(|| format!("bad ternary: {s}"))?;
287 let c = s[q + 1..]
288 .find(':')
289 .map(|x| q + 1 + x)
290 .ok_or_else(|| format!("bad ternary colon: {s}"))?;
291 let open = s[..q]
292 .rfind('(')
293 .ok_or_else(|| format!("bad ternary open: {s}"))?;
294 let close = s[c + 1..]
295 .rfind(')')
296 .map(|x| c + 1 + x)
297 .ok_or_else(|| format!("bad ternary close: {s}"))?;
298 let prefix = &s[..open];
299 let cond = s[open + 1..q].trim();
300 let true_expr = s[q + 1..c].trim();
301 let false_expr = s[c + 1..close].trim();
302 let suffix = &s[close + 1..];
303 Ok(format!(
304 "{}select({}, {}, {}){}",
305 prefix, false_expr, true_expr, cond, suffix
306 ))
307}
308
309fn find_matching_paren(s: &str, open_idx: usize) -> Option<usize> {
310 let mut depth = 0i32;
311 for (i, ch) in s.char_indices().skip_while(|(i, _)| *i < open_idx) {
312 match ch {
313 '(' => depth += 1,
314 ')' => {
315 depth -= 1;
316 if depth == 0 {
317 return Some(i);
318 }
319 }
320 _ => {}
321 }
322 }
323 None
324}
325
326fn find_top_level_char(s: &str, target: char) -> Option<usize> {
327 let mut depth = 0i32;
328 for (i, ch) in s.char_indices() {
329 match ch {
330 '(' => depth += 1,
331 ')' => depth -= 1,
332 _ if ch == target && depth == 0 => return Some(i),
333 _ => {}
334 }
335 }
336 None
337}
338
339fn find_matching_colon(s: &str, q_index: usize) -> Option<usize> {
340 let mut depth = 0i32;
341 for (i, ch) in s.char_indices().skip_while(|(i, _)| *i <= q_index) {
342 match ch {
343 '(' => depth += 1,
344 ')' => depth -= 1,
345 ':' if depth == 0 => return Some(i),
346 _ => {}
347 }
348 }
349 None
350}
351
352fn split_top_level_once(s: &str, needle: char) -> Option<(&str, &str)> {
353 let mut depth = 0i32;
354 for (i, ch) in s.char_indices() {
355 match ch {
356 '(' => depth += 1,
357 ')' => depth -= 1,
358 _ if ch == needle && depth == 0 => return Some((&s[..i], &s[i + 1..])),
359 _ => {}
360 }
361 }
362 None
363}
364
365fn emit_wgsl(hlsl: &ParsedHlsl, kind: ShaderKind) -> String {
366 let mut out = String::new();
367 out.push_str(&format!(
368 "struct FloatRegs {{\n c: array<vec4<f32>, {}>,\n}};\n\n",
369 hlsl.register_count
370 ));
371 out.push_str("@group(0) @binding(0) var<uniform> u: FloatRegs;\n\n");
372
373 let mut sampler_idx = BTreeMap::new();
374 for name in &hlsl.samplers {
375 let idx = name.trim_start_matches('s').parse::<u32>().unwrap_or(0);
376 sampler_idx.insert(name.clone(), idx);
377 }
378 for name in &hlsl.samplers {
379 let idx = sampler_idx[name];
380 out.push_str(&format!(
381 "@group(1) @binding({}) var tex_{}: texture_2d<f32>;\n",
382 idx * 2,
383 name
384 ));
385 out.push_str(&format!(
386 "@group(1) @binding({}) var samp_{}: sampler;\n",
387 idx * 2 + 1,
388 name
389 ));
390 }
391 if !hlsl.samplers.is_empty() {
392 out.push('\n');
393 }
394
395 for (name, ty, value) in &hlsl.consts {
396 out.push_str(
397 &format!(
398 "const {}: {} = {};;\n",
399 name,
400 wgsl_type(ty),
401 replace_types(value)
402 )
403 .replace(";;", ";"),
404 );
405 }
406 if !hlsl.consts.is_empty() {
407 out.push('\n');
408 }
409
410 emit_struct(&mut out, &hlsl.input_name, &hlsl.input_fields, kind, true);
411 out.push('\n');
412 emit_struct(
413 &mut out,
414 &hlsl.output_name,
415 &hlsl.output_fields,
416 kind,
417 false,
418 );
419 out.push('\n');
420
421 match kind {
422 ShaderKind::Vertex => out.push_str("@vertex\n"),
423 ShaderKind::Pixel => out.push_str("@fragment\n"),
424 }
425 out.push_str(&format!(
426 "fn main(input: {}) -> {} {{\n",
427 hlsl.input_name, hlsl.output_name
428 ));
429 out.push_str(&format!(" var output: {};\n", hlsl.output_name));
430 for line in &hlsl.body_lines {
431 out.push_str(line);
432 out.push('\n');
433 }
434 out.push_str(" return output;\n}\n");
435 out
436}
437
438fn emit_struct(
439 out: &mut String,
440 name: &str,
441 fields: &[StructField],
442 kind: ShaderKind,
443 is_input: bool,
444) {
445 out.push_str(&format!("struct {} {{\n", name));
446 for field in fields {
447 let attr = semantic_to_wgsl_attribute(&field.semantic, kind, is_input);
448 out.push_str(&format!(
449 " {}{}: {},\n",
450 attr,
451 field.name,
452 wgsl_type(&field.ty)
453 ));
454 }
455 out.push_str("};\n");
456}
457
458fn wgsl_type(ty: &str) -> &'static str {
459 match ty {
460 "float4" | "vec4<f32>" => "vec4<f32>",
461 "float3" | "vec3<f32>" => "vec3<f32>",
462 "float2" | "vec2<f32>" => "vec2<f32>",
463 "float" | "f32" => "f32",
464 _ => "vec4<f32>",
465 }
466}
467
468fn semantic_to_wgsl_attribute(semantic: &str, kind: ShaderKind, is_input: bool) -> String {
469 if kind == ShaderKind::Pixel
470 && !is_input
471 && let Some(n) = semantic
472 .strip_prefix("COLOR")
473 .and_then(|s| s.parse::<u32>().ok())
474 {
475 return format!("@location({}) ", n);
476 }
477 if semantic.starts_with("POSITION") {
478 if kind == ShaderKind::Vertex && !is_input {
479 return "@builtin(position) ".to_string();
480 }
481 return "@location(0) ".to_string();
482 }
483 if semantic == "NORMAL" {
484 return "@location(1) ".to_string();
485 }
486 if let Some(n) = semantic
487 .strip_prefix("COLOR")
488 .and_then(|s| s.parse::<u32>().ok())
489 {
490 return format!("@location({}) ", 2 + n);
491 }
492 if let Some(n) = semantic
493 .strip_prefix("TEXCOORD")
494 .and_then(|s| s.parse::<u32>().ok())
495 {
496 return format!("@location({}) ", 4 + n);
497 }
498 "@location(15) ".to_string()
499}