From 3791e889c84f9948565240b6ec0b0e3273955e1c Mon Sep 17 00:00:00 2001
From: gsemaj <gentsemaj@live.com>
Date: Thu, 11 Aug 2022 11:53:38 -0400
Subject: [PATCH] Add support for multiple subprograms

---
 dx2cg/README.md       |  3 +-
 dx2cg/disassembler.py | 66 ++++++++++++++++++++++++++-----------------
 dx2cg/swapper.py      | 22 +++++++++------
 3 files changed, 55 insertions(+), 36 deletions(-)

diff --git a/dx2cg/README.md b/dx2cg/README.md
index f89623e..fb02690 100644
--- a/dx2cg/README.md
+++ b/dx2cg/README.md
@@ -6,6 +6,5 @@ Tools for converting d3d9 shader assembly to HLSL/Cg.
 
 ## Known issues
 - Only vertex shaders with profile `vs_1_1` are supported
-- No fragment shaders are supported yet
-- Only one subprogram in a subshader will be converted (for now)
+- Only fragment shaders with profile `ps_2_0` are supported
 - Only a limited set of instructions (those used by FF and Unity 2.6) are supported
diff --git a/dx2cg/disassembler.py b/dx2cg/disassembler.py
index 8e3430a..a76a4a5 100644
--- a/dx2cg/disassembler.py
+++ b/dx2cg/disassembler.py
@@ -115,7 +115,6 @@ fragment_func = """float4 frag(v2f pdat) {{
 
 def process_header(prog):
     keywords = []
-    header = []
     loctab = {}
     locdecl = []
     binds = []
@@ -192,19 +191,10 @@ def process_header(prog):
             del prog[i]
             i = i - 1
         i = i + 1
-        
-    if len(binds) > 0:
-        header.append("BindChannels {")
-        for b in binds:
-            header.append(f"\t{b}")
-        header.append("}")
-    
-    if lighting:
-        header.append("Lighting On")
 
     # print(loctab)
     
-    return (keywords, header, loctab, locdecl)
+    return (keywords, loctab, locdecl, binds, lighting)
 
 def resolve_args(args, loctab, consts):
     for a in range(0, len(args)):
@@ -306,28 +296,52 @@ def process_asm(asm, loctab):
     
     return (shadertype, translated)
 
-def disassemble(text):
-    asm = text.split('\n')[1:-1]
-    (keywords, header, loctab, locdecl) = process_header(asm)
-    (shadertype, disasm) = process_asm(asm, loctab)
+def disassemble(blocks):
+    shaders = {}
+    keywords = set()
+    locdecl = set()
+    binds = set()
+    lighting = False
+    for block in blocks:
+        asm = block.split('\n')[1:-1]
+
+        (kw, ltab, ldecl, bds, light) = process_header(asm)
+        keywords.update(kw)
+        locdecl.update(ldecl)
+        binds.update(bds)
+        lighting |= light
+
+        (shadertype, disasm) = process_asm(asm, ltab)
+        shaders[shadertype] = disasm
+
+    text = ""
+    if len(binds) > 0:
+        text += "BindChannels {\n"
+        for b in binds:
+            text += f"\t{b}\n"
+        text += "}\n"
     
-    text = "\n".join(header) + "\n" if len(header) > 0 else ""
+    if lighting:
+        text += "Lighting On\n"
+
     text += cg_header
-    if keywords:
+    if len(keywords) > 0:
         text += "#pragma multi_compile " + " ".join(keywords)
-    if shadertype == "vertex":
+    if "vertex" in shaders:
         text += "#pragma vertex vert\n"
-    if shadertype == "fragment":
+    if "fragment" in shaders:
         text += "#pragma fragment frag\n"
     text += "\n"
-    text += struct_a2v + "\n"
+    if "vertex" in shaders:
+        text += struct_a2v + "\n"
     text += struct_v2f + "\n"
-    text += struct_f2a + "\n"
+    if "fragment" in shaders:
+        text += struct_f2a + "\n"
     text += "\n".join(locdecl) + "\n"
-    if shadertype == "vertex":
-        text += vertex_func.format("\t" + "\n\t".join(disasm))
-    if shadertype == "fragment":
-        text += fragment_func.format("\t" + "\n\t".join(disasm))
+    if "vertex" in shaders:
+        text += "\n" + vertex_func.format("\t" + "\n\t".join(shaders["vertex"]))
+    if "fragment" in shaders:
+        text += "\n" + fragment_func.format("\t" + "\n\t".join(shaders["fragment"]))
     text += cg_footer
     return text
 
@@ -337,5 +351,5 @@ if __name__ == "__main__":
     else:
         with open(sys.argv[1], "r") as fi:
             buf = fi.read()
-        disasm = disassemble(buf)
+        disasm = disassemble(buf.split('~'))
         print(disasm)
diff --git a/dx2cg/swapper.py b/dx2cg/swapper.py
index 555aca1..89e63d5 100644
--- a/dx2cg/swapper.py
+++ b/dx2cg/swapper.py
@@ -27,18 +27,24 @@ def find_closing_bracket(block, i):
     raise ValueError(f"Block at {i} has no closing bracket")
 
 def process_program(prog):
+    # print("processing:\n" + prog)
+    subprogs = []
     subprog_index = prog.find("SubProgram \"d3d9")
-    if subprog_index == -1:
-        raise ValueError(f"Program has no d3d9 subprogram")
-    subprog_end_index = find_closing_bracket(prog, subprog_index)
-    subprog = prog[subprog_index:subprog_end_index+1]
-    processed = disassemble(subprog) + "\n"
+    while subprog_index > -1:
+        subprog_end_index = find_closing_bracket(prog, subprog_index)
+        subprog = prog[subprog_index:subprog_end_index+1]
+        subprogs.append(subprog)
+        prog = prog[subprog_end_index+1:]
+        subprog_index = prog.find("SubProgram \"d3d9")
+    if len(subprogs) < 1:
+        raise ValueError(f"Program has no d3d9 subprograms")
+    processed = disassemble(subprogs) + "\n"
     return indent(processed)
 
 def process_shader(shader):
     buf = shader
     processed = ''
-    program_index = buf.find('Program')
+    program_index = buf.find("Program \"\"")
     while program_index > -1:
         processed = processed + buf[:program_index]
         buf = buf[program_index:]
@@ -46,11 +52,11 @@ def process_shader(shader):
         if not line:
             raise ValueError(f"Program at {program_index} has no #LINE marker")
         end_index = line.end() + 1
-        program_section = buf[program_index:end_index+1]
+        program_section = buf[:end_index+1]
         processed = processed + process_program(program_section)
         buf = buf[end_index+1:]
         
-        program_index = buf.find('Program')
+        program_index = buf.find("Program \"\"")
     processed = processed + buf
     return processed