Handle extended instruction length for dcl_function_table

Author: asavonic

mlir/lib/Target/DXSA/BinaryParser.cpp

@@ -1270,6 +1270,18 @@ class Parser {
     return success();
   }
 
+  FailureOr<std::optional<uint32_t>> parseExtendedLength(uint32_t opcode) {
+    std::optional<uint32_t> result;
+    if (opcode != D3D11_SB_OPCODE_DCL_FUNCTION_TABLE) {
+      return success(result);
+    }
+
+    Token lengthToken = parseToken();
+    FAILURE_IF_FAILED(lengthToken);
+    result = *lengthToken;
+    return success(result);
+  }
+
   FailureOr<Instruction> parseInstruction() {
     size_t beginOffset = currentTokenOffset;
     Token token = parseToken();
@@ -1284,11 +1296,14 @@ class Parser {
     modifier.saturate = DECODE_IS_D3D10_SB_INSTRUCTION_SATURATE_ENABLED(*token);
 
     uint32_t length = DECODE_D3D10_SB_TOKENIZED_INSTRUCTION_LENGTH(*token);
+    if (DECODE_IS_D3D10_SB_OPCODE_EXTENDED(*token)) {
+      auto extLength = parseExtendedLength(opcode);
+      FAILURE_IF_FAILED(extLength);
+      length = extLength->value();
+    }
 
     // TODO: extended instructions:
     // BOOL b51PlusShader =
-    // BOOL bExtended = DECODE_IS_D3D10_SB_OPCODE_EXTENDED(Token)
-    // ...
 
     if (opcode >= D3D10_SB_NUM_OPCODES) {
       emitError(getLocation(), "unknown opcode");

mlir/test/Target/DXSA/dcl_function_table.mlir

@@ -2,3 +2,6 @@
 
 // CHECK-LABEL: module
 // CHECK-NEXT: dxsa.dcl_function_table 1, <functions = [0, 1]>
+
+// Test extended instruction length:
+// CHECK-NEXT: dxsa.dcl_function_table 2, <functions = [0, 1]>