u256 shifts operators (FuelLabs#4928)

## Description

This PR is part of FuelLabs#4794.

It implements left and right shifts. These have a little difference
because the `rhs` is always `u64`, which means that the WQOP used the
`rhs` is not indirect.

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] 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-core/src/asm_generation/fuel/fuel_asm_builder.rs

@@ -12,7 +12,7 @@ use crate::{
     },
     asm_lang::{
         virtual_register::*, Label, Op, VirtualImmediate06, VirtualImmediate12, VirtualImmediate18,
-        VirtualOp, WideCmp,
+        VirtualOp, WideCmp, WideOperations,
     },
     decl_engine::DeclRefFunction,
     metadata::MetadataManager,
@@ -580,13 +580,25 @@ impl<'ir, 'eng> FuelAsmBuilder<'ir, 'eng> {
                 result_reg,
                 val1_reg,
                 val2_reg,
-                VirtualImmediate06::wide_op(crate::asm_lang::WideOperations::Add, true),
+                VirtualImmediate06::wide_op(WideOperations::Add, true),
             ),
             BinaryOpKind::Sub => VirtualOp::WQOP(
                 result_reg,
                 val1_reg,
                 val2_reg,
-                VirtualImmediate06::wide_op(crate::asm_lang::WideOperations::Sub, true),
+                VirtualImmediate06::wide_op(WideOperations::Sub, true),
+            ),
+            BinaryOpKind::Lsh => VirtualOp::WQOP(
+                result_reg,
+                val1_reg,
+                val2_reg,
+                VirtualImmediate06::wide_op(WideOperations::Lsh, false),
+            ),
+            BinaryOpKind::Rsh => VirtualOp::WQOP(
+                result_reg,
+                val1_reg,
+                val2_reg,
+                VirtualImmediate06::wide_op(WideOperations::Rsh, false),
             ),
             BinaryOpKind::Mul => VirtualOp::WQML(
                 result_reg,

sway-core/src/asm_lang/virtual_immediate.rs

@@ -11,6 +11,8 @@ pub enum WideOperations {
     Add = 0,
     Sub = 1,
     Not = 2,
+    Lsh = 6,
+    Rsh = 7,
 }
 
 #[repr(u8)]

sway-core/src/semantic_analysis/ast_node/expression/intrinsic_function.rs

@@ -888,6 +888,7 @@ fn type_check_shift_binary_op(
             .with_type_annotation(return_type),
         lhs,
     )?;
+
     let rhs = arguments[1].clone();
     let rhs = ty::TyExpression::type_check(
         handler,

sway-ir/src/optimize/misc_demotion.rs

@@ -40,6 +40,7 @@ pub fn misc_demotion(
     let addrof_res = ptr_to_int_demotion(context, function)?;
 
     let wide_binary_op_res = wide_binary_op_demotion(context, function)?;
+    let wide_shifts_op_res = wide_shift_op_demotion(context, function)?;
     let wide_cmp_res = wide_cmp_demotion(context, function)?;
     let wide_unary_op_res = wide_unary_op_demotion(context, function)?;
 
@@ -49,6 +50,7 @@ pub fn misc_demotion(
         || addrof_res
         || wide_unary_op_res
         || wide_binary_op_res
+        || wide_shifts_op_res
         || wide_cmp_res)
 }
 
@@ -851,3 +853,153 @@ fn wide_unary_op_demotion(context: &mut Context, function: Function) -> Result<b
 
     Ok(true)
 }
+
+/// Find all shift operations on types bigger than 64 bits
+/// and demote them to fuel specific `wide binary ops`, that
+/// work only on pointers
+fn wide_shift_op_demotion(context: &mut Context, function: Function) -> Result<bool, IrError> {
+    // Find all intrinsics on wide operators
+    let candidates = function
+        .instruction_iter(context)
+        .filter_map(|(block, instr_val)| {
+            let instr = instr_val.get_instruction(context)?;
+
+            if let Instruction::BinaryOp { op, arg1, arg2 } = instr {
+                let arg1_type = arg1
+                    .get_type(context)
+                    .and_then(|x| x.get_uint_width(context));
+                let arg2_type = arg2
+                    .get_type(context)
+                    .and_then(|x| x.get_uint_width(context));
+
+                match (arg1_type, arg2_type) {
+                    (Some(256), Some(64)) => {
+                        use BinaryOpKind::*;
+                        match op {
+                            Lsh | Rsh => Some((block, instr_val)),
+                            _ => todo!(),
+                        }
+                    }
+                    _ => None,
+                }
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+
+    if candidates.is_empty() {
+        return Ok(false);
+    }
+
+    // Now create a local for the result
+    // get ptr to each arg
+    // and store the result after
+    for (block, binary_op_instr_val) in candidates {
+        let Instruction::BinaryOp { op, arg1, arg2 } = binary_op_instr_val
+            .get_instruction(context)
+            .cloned()
+            .unwrap()
+        else {
+            continue;
+        };
+
+        let binary_op_metadata = binary_op_instr_val.get_metadata(context);
+
+        let arg1_ty = arg1.get_type(context).unwrap();
+        let arg1_metadata = arg1.get_metadata(context);
+
+        let arg2_ty = arg2.get_type(context).unwrap();
+
+        let operand_ty = arg1.get_type(context).unwrap();
+
+        let result_local = function.new_unique_local_var(
+            context,
+            "__wide_result".to_owned(),
+            operand_ty,
+            None,
+            true,
+        );
+        let get_result_local = Value::new_instruction(context, Instruction::GetLocal(result_local))
+            .add_metadatum(context, binary_op_metadata);
+        let load_result_local =
+            Value::new_instruction(context, Instruction::Load(get_result_local))
+                .add_metadatum(context, binary_op_metadata);
+
+        // If arg1 is not a pointer, store it to a local
+        let lhs_store = if !arg1_ty.is_ptr(context) {
+            let lhs_local = function.new_unique_local_var(
+                context,
+                "__wide_lhs".to_owned(),
+                operand_ty,
+                None,
+                false,
+            );
+            let get_lhs_local = Value::new_instruction(context, Instruction::GetLocal(lhs_local))
+                .add_metadatum(context, arg1_metadata);
+            let store_lhs_local = Value::new_instruction(
+                context,
+                Instruction::Store {
+                    dst_val_ptr: get_lhs_local,
+                    stored_val: arg1,
+                },
+            )
+            .add_metadatum(context, arg1_metadata);
+            Some((get_lhs_local, store_lhs_local))
+        } else {
+            None
+        };
+
+        let (arg1_needs_insert, get_arg1) = if let Some((lhs_local, _)) = &lhs_store {
+            (false, *lhs_local)
+        } else {
+            (true, arg1)
+        };
+
+        // Assert result and lhs are pointers
+        // Assert rhs is u64
+        assert!(get_arg1.get_type(context).unwrap().is_ptr(context));
+        assert!(get_result_local.get_type(context).unwrap().is_ptr(context));
+        assert!(arg2_ty.is_uint64(context));
+
+        let wide_op = Value::new_instruction(
+            context,
+            Instruction::FuelVm(FuelVmInstruction::WideBinaryOp {
+                op,
+                arg1: get_arg1,
+                arg2,
+                result: get_result_local,
+            }),
+        )
+        .add_metadatum(context, binary_op_metadata);
+
+        // We don't have an actual instruction _inserter_ yet, just an appender, so we need to find
+        // the ptr_to_int instruction index and insert instructions manually.
+        let block_instrs = &mut context.blocks[block.0].instructions;
+        let idx = block_instrs
+            .iter()
+            .position(|&instr_val| instr_val == binary_op_instr_val)
+            .unwrap();
+
+        block
+            .replace_instruction(context, binary_op_instr_val, load_result_local)
+            .unwrap();
+
+        let block_instrs = &mut context.blocks[block.0].instructions;
+
+        block_instrs.insert(idx, wide_op);
+        block_instrs.insert(idx, get_result_local);
+
+        if arg1_needs_insert {
+            block_instrs.insert(idx, get_arg1);
+        }
+
+        // lhs
+        if let Some((get_lhs_local, store_lhs_local)) = lhs_store {
+            block_instrs.insert(idx, store_lhs_local);
+            block_instrs.insert(idx, get_lhs_local);
+        }
+    }
+
+    Ok(true)
+}

sway-ir/src/verify.rs

@@ -374,7 +374,7 @@ impl<'a, 'eng> InstructionVerifier<'a, 'eng> {
 
     fn verify_wide_binary_op(
         &self,
-        _op: &BinaryOpKind,
+        op: &BinaryOpKind,
         result: &Value,
         arg1: &Value,
         arg2: &Value,
@@ -389,11 +389,31 @@ impl<'a, 'eng> InstructionVerifier<'a, 'eng> {
             .get_type(self.context)
             .ok_or(IrError::VerifyBinaryOpIncorrectArgType)?;
 
-        if !arg1_ty.is_ptr(self.context)
-            || !arg2_ty.is_ptr(self.context)
-            || !result_ty.is_ptr(self.context)
-        {
-            return Err(IrError::VerifyBinaryOpIncorrectArgType);
+        match op {
+            // Shifts rhs are 64 bits
+            BinaryOpKind::Lsh | BinaryOpKind::Rsh => {
+                if !arg1_ty.is_ptr(self.context)
+                    || !arg2_ty.is_uint64(self.context)
+                    || !result_ty.is_ptr(self.context)
+                {
+                    return Err(IrError::VerifyBinaryOpIncorrectArgType);
+                }
+            }
+            BinaryOpKind::Add
+            | BinaryOpKind::Sub
+            | BinaryOpKind::Mul
+            | BinaryOpKind::Div
+            | BinaryOpKind::And
+            | BinaryOpKind::Or
+            | BinaryOpKind::Xor
+            | BinaryOpKind::Mod => {
+                if !arg1_ty.is_ptr(self.context)
+                    || !arg2_ty.is_ptr(self.context)
+                    || !result_ty.is_ptr(self.context)
+                {
+                    return Err(IrError::VerifyBinaryOpIncorrectArgType);
+                }
+            }
         }
 
         Ok(())
@@ -421,7 +441,7 @@ impl<'a, 'eng> InstructionVerifier<'a, 'eng> {
 
     fn verify_binary_op(
         &self,
-        _op: &BinaryOpKind,
+        op: &BinaryOpKind,
         arg1: &Value,
         arg2: &Value,
     ) -> Result<(), IrError> {
@@ -431,8 +451,26 @@ impl<'a, 'eng> InstructionVerifier<'a, 'eng> {
         let arg2_ty = arg2
             .get_type(self.context)
             .ok_or(IrError::VerifyBinaryOpIncorrectArgType)?;
-        if !arg1_ty.eq(self.context, &arg2_ty) || !arg1_ty.is_uint(self.context) {
-            return Err(IrError::VerifyBinaryOpIncorrectArgType);
+
+        match op {
+            // Shifts can have the rhs with different type
+            BinaryOpKind::Lsh | BinaryOpKind::Rsh => {
+                if !arg1_ty.is_uint(self.context) || !arg2_ty.is_uint(self.context) {
+                    return Err(IrError::VerifyBinaryOpIncorrectArgType);
+                }
+            }
+            BinaryOpKind::Add
+            | BinaryOpKind::Sub
+            | BinaryOpKind::Mul
+            | BinaryOpKind::Div
+            | BinaryOpKind::And
+            | BinaryOpKind::Or
+            | BinaryOpKind::Xor
+            | BinaryOpKind::Mod => {
+                if !arg1_ty.eq(self.context, &arg2_ty) || !arg1_ty.is_uint(self.context) {
+                    return Err(IrError::VerifyBinaryOpIncorrectArgType);
+                }
+            }
         }
 
         Ok(())

sway-ir/tests/demote_misc/demote_wide_binary_ops_constants.ir

@@ -4,6 +4,8 @@ script {
         v0 = const u256 0x0000000000000000000000000000000000000000000000000000000000000001
         v1 = const u256 0x0000000000000000000000000000000000000000000000000000000000000002
         v2 = add v0, v1
+        v3 = const u64 2
+        v4 = lsh v0, v3
         ret u256 v2
     }
 }

sway-ir/tests/serialize/wide_ops.ir

@@ -3,6 +3,7 @@ script {
         local u256 a
         entry():
         v0 = get_local ptr u256, a
+        v1 = const u64 2
 
         wide add v0, v0 to v0
 // check: wide add v0, v0 to v0
@@ -19,6 +20,12 @@ script {
         wide mod v0, v0 to v0
 // check: wide mod v0, v0 to v0
 
+        wide lsh v0, v1 to v0
+// check: wide lsh v0, v1 to v0
+
+        wide rsh v0, v1 to v0
+// check: wide rsh v0, v1 to v0
+
         wide not v0 to v0
 // check: wide not v0 to v0
 

sway-lib-core/src/ops.sw

@@ -543,6 +543,15 @@ pub trait Shift {
     fn rsh(self, other: u64) -> Self;
 }
 
+impl Shift for u256 {
+    fn lsh(self, other: u64) -> Self {
+        __lsh(self, other)
+    }
+    fn rsh(self, other: u64) -> Self {
+        __rsh(self, other)
+    }
+}
+
 impl Shift for u64 {
     fn lsh(self, other: u64) -> Self {
         __lsh(self, other)