Support FuelVM wqxx instructions in ASM blocks (#6484)

## Description

This PR adds support for FuelVM `wqxx` instructions in ASM blocks. Some
of those instructions were up to now partially supported in IR and ASM
generation but not in ASM blocks. The immediate motivation for adding
support for these instructions is the need for `wqml` instruction in the
new implementation of `u256::pow()`.

Additionally, the PR adds expressive diagnostics for
`UnrecognizedOpCode` parsing error.

## Checklist

- [ ] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [ ] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [ ] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

sway-ast/src/expr/op_code.rs

@@ -200,6 +200,13 @@ define_op_codes!(
     (Srli, SrliOpcode, "srli", (ret: reg, lhs: reg, rhs: imm)),
     (Sub, SubOpcode, "sub", (ret: reg, lhs: reg, rhs: reg)),
     (Subi, SubiOpcode, "subi", (ret: reg, lhs: reg, rhs: imm)),
+    (Wqcm, WqcmOpcode, "wqcm", (ret: reg, lhs: reg, rhs: reg, op_mode: imm)),
+    (Wqop, WqopOpcode, "wqop", (ret: reg, lhs: reg, rhs: reg, op_mode: imm)),
+    (Wqml, WqmlOpcode, "wqml", (ret: reg, lhs: reg, rhs: reg, indirect: imm)),
+    (Wqdv, WqdvOpcode, "wqdv", (ret: reg, lhs: reg, rhs: reg, indirect: imm)),
+    (Wqmd, WqmdOpcode, "wqmd", (ret: reg, lhs_a: reg, lhs_b: reg, rhs: reg)),
+    (Wqam, WqamOpcode, "wqam", (ret: reg, lhs_a: reg, lhs_b: reg, rhs: reg)),
+    (Wqmm, WqmmOpcode, "wqmm", (ret: reg, lhs_a: reg, lhs_b: reg, rhs: reg)),
     (Xor, XorOpcode, "xor", (ret: reg, lhs: reg, rhs: reg)),
     (Xori, XoriOpcode, "xori", (ret: reg, lhs: reg, rhs: imm)),
     /* Control Flow Instructions */

sway-core/src/asm_lang/allocated_ops.rs

@@ -110,6 +110,12 @@ pub(crate) enum AllocatedOpcode {
         AllocatedRegister,
         VirtualImmediate06,
     ),
+    WQMD(
+        AllocatedRegister,
+        AllocatedRegister,
+        AllocatedRegister,
+        AllocatedRegister,
+    ),
     WQCM(
         AllocatedRegister,
         AllocatedRegister,
@@ -122,6 +128,12 @@ pub(crate) enum AllocatedOpcode {
         AllocatedRegister,
         AllocatedRegister,
     ),
+    WQMM(
+        AllocatedRegister,
+        AllocatedRegister,
+        AllocatedRegister,
+        AllocatedRegister,
+    ),
 
     /* Control Flow Instructions */
     JMP(AllocatedRegister),
@@ -303,8 +315,10 @@ impl AllocatedOpcode {
             WQOP(_, _, _, _) => vec![],
             WQML(_, _, _, _) => vec![],
             WQDV(_, _, _, _) => vec![],
+            WQMD(_, _, _, _) => vec![],
             WQCM(r1, _, _, _) => vec![r1],
             WQAM(_, _, _, _) => vec![],
+            WQMM(_, _, _, _) => vec![],
 
             /* Control Flow Instructions */
             JMP(_r1) => vec![],
@@ -428,8 +442,10 @@ impl fmt::Display for AllocatedOpcode {
             WQOP(a, b, c, d) => write!(fmtr, "wqop {a} {b} {c} {d}"),
             WQML(a, b, c, d) => write!(fmtr, "wqml {a} {b} {c} {d}"),
             WQDV(a, b, c, d) => write!(fmtr, "wqdv {a} {b} {c} {d}"),
+            WQMD(a, b, c, d) => write!(fmtr, "wqmd {a} {b} {c} {d}"),
             WQCM(a, b, c, d) => write!(fmtr, "wqcm {a} {b} {c} {d}"),
             WQAM(a, b, c, d) => write!(fmtr, "wqam {a} {b} {c} {d}"),
+            WQMM(a, b, c, d) => write!(fmtr, "wqmm {a} {b} {c} {d}"),
 
             /* Control Flow Instructions */
             JMP(a) => write!(fmtr, "jmp {a}"),
@@ -594,12 +610,18 @@ impl AllocatedOp {
             WQDV(a, b, c, d) => {
                 op::WQDV::new(a.to_reg_id(), b.to_reg_id(), c.to_reg_id(), d.value.into()).into()
             }
+            WQMD(a, b, c, d) => {
+                op::WQMD::new(a.to_reg_id(), b.to_reg_id(), c.to_reg_id(), d.to_reg_id()).into()
+            }
             WQCM(a, b, c, d) => {
                 op::WQCM::new(a.to_reg_id(), b.to_reg_id(), c.to_reg_id(), d.value.into()).into()
             }
             WQAM(a, b, c, d) => {
                 op::WQAM::new(a.to_reg_id(), b.to_reg_id(), c.to_reg_id(), d.to_reg_id()).into()
             }
+            WQMM(a, b, c, d) => {
+                op::WQMM::new(a.to_reg_id(), b.to_reg_id(), c.to_reg_id(), d.to_reg_id()).into()
+            }
 
             /* Control Flow Instructions */
             JMP(a) => op::JMP::new(a.to_reg_id()).into(),

sway-core/src/asm_lang/mod.rs

@@ -415,6 +415,34 @@ impl Op {
                 let (r1, r2, imm) = two_regs_imm_12(handler, args, immediate, whole_op_span)?;
                 VirtualOp::SUBI(r1, r2, imm)
             }
+            "wqcm" => {
+                let (r1, r2, r3, imm) = three_regs_imm_06(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQCM(r1, r2, r3, imm)
+            }
+            "wqop" => {
+                let (r1, r2, r3, imm) = three_regs_imm_06(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQOP(r1, r2, r3, imm)
+            }
+            "wqml" => {
+                let (r1, r2, r3, imm) = three_regs_imm_06(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQML(r1, r2, r3, imm)
+            }
+            "wqdv" => {
+                let (r1, r2, r3, imm) = three_regs_imm_06(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQDV(r1, r2, r3, imm)
+            }
+            "wqmd" => {
+                let (r1, r2, r3, r4) = four_regs(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQMD(r1, r2, r3, r4)
+            }
+            "wqam" => {
+                let (r1, r2, r3, r4) = four_regs(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQAM(r1, r2, r3, r4)
+            }
+            "wqmm" => {
+                let (r1, r2, r3, r4) = four_regs(handler, args, immediate, whole_op_span)?;
+                VirtualOp::WQMM(r1, r2, r3, r4)
+            }
             "xor" => {
                 let (r1, r2, r3) = three_regs(handler, args, immediate, whole_op_span)?;
                 VirtualOp::XOR(r1, r2, r3)
@@ -1148,8 +1176,10 @@ impl fmt::Display for VirtualOp {
             WQOP(a, b, c, d) => write!(fmtr, "wqop {a} {b} {c} {d}"),
             WQML(a, b, c, d) => write!(fmtr, "wqml {a} {b} {c} {d}"),
             WQDV(a, b, c, d) => write!(fmtr, "wqdv {a} {b} {c} {d}"),
+            WQMD(a, b, c, d) => write!(fmtr, "wqmd {a} {b} {c} {d}"),
             WQCM(a, b, c, d) => write!(fmtr, "wqcm {a} {b} {c} {d}"),
             WQAM(a, b, c, d) => write!(fmtr, "wqam {a} {b} {c} {d}"),
+            WQMM(a, b, c, d) => write!(fmtr, "wqmm {a} {b} {c} {d}"),
 
             /* Control Flow Instructions */
             JMP(a) => write!(fmtr, "jmp {a}"),

sway-core/src/asm_lang/virtual_ops.rs

@@ -75,6 +75,12 @@ pub(crate) enum VirtualOp {
         VirtualRegister,
         VirtualImmediate06,
     ),
+    WQMD(
+        VirtualRegister,
+        VirtualRegister,
+        VirtualRegister,
+        VirtualRegister,
+    ),
     WQCM(
         VirtualRegister,
         VirtualRegister,
@@ -87,6 +93,12 @@ pub(crate) enum VirtualOp {
         VirtualRegister,
         VirtualRegister,
     ),
+    WQMM(
+        VirtualRegister,
+        VirtualRegister,
+        VirtualRegister,
+        VirtualRegister,
+    ),
 
     /* Control Flow Instructions */
     JMP(VirtualRegister),
@@ -263,8 +275,10 @@ impl VirtualOp {
             WQOP(r1, r2, r3, _) => vec![r1, r2, r3],
             WQML(r1, r2, r3, _) => vec![r1, r2, r3],
             WQDV(r1, r2, r3, _) => vec![r1, r2, r3],
+            WQMD(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
             WQCM(r1, r2, r3, _) => vec![r1, r2, r3],
             WQAM(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
+            WQMM(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
 
             /* Control Flow Instructions */
             JMP(r1) => vec![r1],
@@ -399,8 +413,10 @@ impl VirtualOp {
             WQOP(_, _, _, _)
             | WQML(_, _, _, _)
             | WQDV(_, _, _, _)
+            | WQMD(_, _, _, _)
             | WQCM(_, _, _, _)
             | WQAM(_, _, _, _)
+            | WQMM(_, _, _, _)
             | JMP(_)
             | JI(_)
             | JNE(_, _, _)
@@ -493,8 +509,10 @@ impl VirtualOp {
             | WQOP(_, _, _, _)
             | WQML(_, _, _, _)
             | WQDV(_, _, _, _)
+            | WQMD(_, _, _, _)
             | WQCM(_, _, _, _)
             | WQAM(_, _, _, _)
+            | WQMM(_, _, _, _)
             // Cryptographic
             | ECK1(_, _, _)
             | ECR1(_, _, _)
@@ -596,11 +614,16 @@ impl VirtualOp {
             SUBI(_r1, r2, _i) => vec![r2],
             XOR(_r1, r2, r3) => vec![r2, r3],
             XORI(_r1, r2, _i) => vec![r2],
+            // Note that most of the `WQ..` instructions *read* from the `r1` result register,
+            // because the register itself does not contain the result, but provides the
+            // memory address at which the result will be stored.
             WQOP(r1, r2, r3, _) => vec![r1, r2, r3],
             WQML(r1, r2, r3, _) => vec![r1, r2, r3],
             WQDV(r1, r2, r3, _) => vec![r1, r2, r3],
+            WQMD(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
             WQCM(_, r2, r3, _) => vec![r2, r3],
             WQAM(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
+            WQMM(r1, r2, r3, r4) => vec![r1, r2, r3, r4],
 
             /* Control Flow Instructions */
             JMP(r1) => vec![r1],
@@ -718,8 +741,10 @@ impl VirtualOp {
             WQOP(_, _, _, _) => vec![],
             WQML(_, _, _, _) => vec![],
             WQDV(_, _, _, _) => vec![],
+            WQMD(_, _, _, _) => vec![],
             WQCM(r1, _, _, _) => vec![r1],
             WQAM(_, _, _, _) => vec![],
+            WQMM(_, _, _, _) => vec![],
 
             /* Control Flow Instructions */
             JMP(_r1) => vec![],
@@ -987,6 +1012,12 @@ impl VirtualOp {
                 update_reg(reg_to_reg_map, r3),
                 i.clone(),
             ),
+            WQMD(r1, r2, r3, r4) => Self::WQMD(
+                update_reg(reg_to_reg_map, r1),
+                update_reg(reg_to_reg_map, r2),
+                update_reg(reg_to_reg_map, r3),
+                update_reg(reg_to_reg_map, r4),
+            ),
             WQCM(r1, r2, r3, i) => Self::WQCM(
                 update_reg(reg_to_reg_map, r1),
                 update_reg(reg_to_reg_map, r2),
@@ -999,6 +1030,12 @@ impl VirtualOp {
                 update_reg(reg_to_reg_map, r3),
                 update_reg(reg_to_reg_map, r4),
             ),
+            WQMM(r1, r2, r3, r4) => Self::WQMM(
+                update_reg(reg_to_reg_map, r1),
+                update_reg(reg_to_reg_map, r2),
+                update_reg(reg_to_reg_map, r3),
+                update_reg(reg_to_reg_map, r4),
+            ),
 
             /* Control Flow Instructions */
             JMP(r1) => Self::JMP(update_reg(reg_to_reg_map, r1)),
@@ -1464,6 +1501,12 @@ impl VirtualOp {
                 map_reg(&mapping, reg3),
                 imm.clone(),
             ),
+            WQMD(reg1, reg2, reg3, reg4) => AllocatedOpcode::WQMD(
+                map_reg(&mapping, reg1),
+                map_reg(&mapping, reg2),
+                map_reg(&mapping, reg3),
+                map_reg(&mapping, reg4),
+            ),
             WQCM(reg1, reg2, reg3, imm) => AllocatedOpcode::WQCM(
                 map_reg(&mapping, reg1),
                 map_reg(&mapping, reg2),
@@ -1476,6 +1519,12 @@ impl VirtualOp {
                 map_reg(&mapping, reg3),
                 map_reg(&mapping, reg4),
             ),
+            WQMM(reg1, reg2, reg3, reg4) => AllocatedOpcode::WQMM(
+                map_reg(&mapping, reg1),
+                map_reg(&mapping, reg2),
+                map_reg(&mapping, reg3),
+                map_reg(&mapping, reg4),
+            ),
 
             /* Control Flow Instructions */
             JMP(reg1) => AllocatedOpcode::JMP(map_reg(&mapping, reg1)),

sway-error/Cargo.toml

@@ -13,6 +13,7 @@ either = "1.9.0"
 in_definite = "1.0.0"
 num-traits = "0.2.14"
 smallvec = "1.7"
+strsim = "0.11.1"
 sway-types = { version = "0.63.3", path = "../sway-types" }
 thiserror = "1.0"
 uwuify = { version = "^0.2", optional = true }

sway-error/src/error.rs

@@ -2646,6 +2646,31 @@ impl ToDiagnostic for CompileError {
                             format!("{}  E.g., `*ref_to_mutable_value` or `*max_mut(&mut x, &mut y)`.", Indent::Single),
                         ]
                     },
+                    ParseErrorKind::UnrecognizedOpCode { known_op_codes } => Diagnostic {
+                        reason: Some(Reason::new(code(1), "Assembly instruction is unknown".to_string())),
+                        issue: Issue::error(
+                            source_engine,
+                            error.span.clone(),
+                            format!("\"{}\" is not a known assembly instruction.",
+                                error.span.as_str()
+                            )
+                        ),
+                        hints: {
+                            let suggestions = &did_you_mean(error.span.as_str(), known_op_codes.iter(), 2);
+                            if suggestions.is_empty() {
+                                vec![]
+                            } else {
+                                vec![
+                                    Hint::help(
+                                        source_engine,
+                                        error.span.clone(),
+                                        format!("Did you mean {}?", sequence_to_str_or(suggestions, Enclosing::DoubleQuote, 2))
+                                    ),
+                                ]
+                            }
+                        },
+                        help: vec![]
+                    },
                     _ => Diagnostic {
                                 // TODO: Temporary we use self here to achieve backward compatibility.
                                 //       In general, self must not be used and will not be used once we