Provide __jmpb_ssp intrinsic

docs/book/src/reference/compiler_intrinsics.md

@@ -315,4 +315,15 @@ __not(op: T) -> T
 **Description:** Bitwise NOT of `op`
 
 **Constraints:** `T` is an integer type, i.e. `u8`, `u16`, `u32`, `u64`, `u256`, `b256`.
+
 ___
+
+```sway
+__jmpb_ssp(offset: u64)
+```
+
+**Description:** Jumps to `$ssp - offset`. When the offset is the growth
+of `$ssp` after an `ldc` call, this transfers control to the newly loaded
+contract.
+
+**Constraints:** offset must have type `u64`.

sway-ast/src/intrinsics.rs

@@ -35,6 +35,7 @@ pub enum Intrinsic {
     PtrSub,
     Smo,
     Not,
+    JmpbSsp,
 }
 
 impl fmt::Display for Intrinsic {
@@ -73,6 +74,7 @@ impl fmt::Display for Intrinsic {
             Intrinsic::PtrSub => "ptr_sub",
             Intrinsic::Smo => "smo",
             Intrinsic::Not => "not",
+            Intrinsic::JmpbSsp => "jmpb_ssp",
         };
         write!(f, "{s}")
     }
@@ -115,6 +117,7 @@ impl Intrinsic {
             "__ptr_sub" => PtrSub,
             "__smo" => Smo,
             "__not" => Not,
+            "__jmpb_ssp" => JmpbSsp,
             _ => return None,
         })
     }

sway-core/src/asm_generation/fuel/fuel_asm_builder.rs

@@ -305,6 +305,7 @@ impl<'ir, 'eng> FuelAsmBuilder<'ir, 'eng> {
                         arg2,
                         arg3,
                     } => self.compile_wide_modular_op(instr_val, op, result, arg1, arg2, arg3),
+                    FuelVmInstruction::JmpbSsp(offset) => self.compile_jmpb_ssp(instr_val, offset),
                 },
                 InstOp::GetElemPtr {
                     base,
@@ -1451,6 +1452,50 @@ impl<'ir, 'eng> FuelAsmBuilder<'ir, 'eng> {
         Ok(())
     }
 
+    fn compile_jmpb_ssp(&mut self, instr_val: &Value, offset: &Value) -> Result<(), CompileError> {
+        let owning_span = self.md_mgr.val_to_span(self.context, *instr_val);
+        let offset_reg = self.value_to_register(offset)?;
+        let is_offset_reg = self.reg_seqr.next();
+        let prev_ssp_reg = self.reg_seqr.next();
+        let by4_reg = self.reg_seqr.next();
+
+        self.cur_bytecode.push(Op {
+            owning_span: owning_span.clone(),
+            opcode: Either::Left(VirtualOp::SUB(
+                prev_ssp_reg.clone(),
+                VirtualRegister::Constant(ConstantRegister::StackStartPointer),
+                offset_reg,
+            )),
+            comment: "jmpb_ssp: Compute $ssp - offset".into(),
+        });
+        self.cur_bytecode.push(Op {
+            owning_span: owning_span.clone(),
+            opcode: Either::Left(VirtualOp::SUB(
+                is_offset_reg.clone(),
+                prev_ssp_reg,
+                VirtualRegister::Constant(ConstantRegister::InstructionStart),
+            )),
+            comment: "jmpb_ssp: Subtract $is since $jmp adds it back.".into(),
+        });
+        self.cur_bytecode.push(Op {
+            owning_span: owning_span.clone(),
+            opcode: Either::Left(VirtualOp::DIVI(
+                by4_reg.clone(),
+                is_offset_reg.clone(),
+                VirtualImmediate12::new(4, Span::dummy()).unwrap(),
+            )),
+            comment: "jmpb_ssp: Divide by 4 since Jmp multiplies by 4.".into(),
+        });
+
+        self.cur_bytecode.push(Op {
+            owning_span,
+            opcode: Either::Left(VirtualOp::JMP(by4_reg)),
+            comment: "jmpb_ssp: Jump to computed value".into(),
+        });
+
+        Ok(())
+    }
+
     fn compile_smo(
         &mut self,
         instr_val: &Value,

sway-core/src/ir_generation/const_eval.rs

@@ -1111,6 +1111,7 @@ fn const_eval_intrinsic(
         | Intrinsic::StateStoreQuad
         | Intrinsic::Log
         | Intrinsic::Revert
+        | Intrinsic::JmpbSsp
         | Intrinsic::Smo => Err(ConstEvalError::CannotBeEvaluatedToConst {
             span: intrinsic.span.clone(),
         }),

sway-core/src/ir_generation/function.rs

@@ -1095,6 +1095,19 @@ impl<'eng> FnCompiler<'eng> {
                     .add_metadatum(context, span_md_idx);
                 Ok(TerminatorValue::new(val, context))
             }
+            Intrinsic::JmpbSsp => {
+                let offset_val = return_on_termination_or_extract!(
+                    self.compile_expression_to_value(context, md_mgr, &arguments[0])?
+                );
+
+                let span_md_idx = md_mgr.span_to_md(context, &span);
+                let val = self
+                    .current_block
+                    .append(context)
+                    .jmpb_ssp(offset_val)
+                    .add_metadatum(context, span_md_idx);
+                Ok(TerminatorValue::new(val, context))
+            }
             Intrinsic::PtrAdd | Intrinsic::PtrSub => {
                 let op = match kind {
                     Intrinsic::PtrAdd => BinaryOpKind::Add,

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

@@ -83,6 +83,9 @@ impl ty::TyIntrinsicFunctionKind {
             }
             Intrinsic::Smo => type_check_smo(handler, ctx, kind, arguments, type_arguments, span),
             Intrinsic::Not => type_check_not(handler, ctx, kind, arguments, type_arguments, span),
+            Intrinsic::JmpbSsp => {
+                type_check_jmpb_ssp(handler, ctx, kind, arguments, type_arguments, span)
+            }
         }
     }
 }
@@ -1170,6 +1173,55 @@ fn type_check_revert(
     ))
 }
 
+/// Signature: `__jmpb_ssp(offset: u64) -> !`
+/// Description: Jumps to `$ssp - offset`.
+/// Constraints: offset has type `u64`.
+fn type_check_jmpb_ssp(
+    handler: &Handler,
+    mut ctx: TypeCheckContext,
+    kind: sway_ast::Intrinsic,
+    arguments: Vec<Expression>,
+    type_arguments: Vec<TypeArgument>,
+    span: Span,
+) -> Result<(ty::TyIntrinsicFunctionKind, TypeId), ErrorEmitted> {
+    let type_engine = ctx.engines.te();
+    let engines = ctx.engines();
+
+    if arguments.len() != 1 {
+        return Err(handler.emit_err(CompileError::IntrinsicIncorrectNumArgs {
+            name: kind.to_string(),
+            expected: 0,
+            span,
+        }));
+    }
+
+    if !type_arguments.is_empty() {
+        return Err(handler.emit_err(CompileError::IntrinsicIncorrectNumTArgs {
+            name: kind.to_string(),
+            expected: 0,
+            span,
+        }));
+    }
+
+    // Type check the argument which is the jmpb_ssp offset
+    let mut ctx = ctx.by_ref().with_type_annotation(type_engine.insert(
+        engines,
+        TypeInfo::UnsignedInteger(IntegerBits::SixtyFour),
+        None,
+    ));
+    let offset = ty::TyExpression::type_check(handler, ctx.by_ref(), arguments[0].clone())?;
+
+    Ok((
+        ty::TyIntrinsicFunctionKind {
+            kind,
+            arguments: vec![offset],
+            type_arguments: vec![],
+            span,
+        },
+        type_engine.insert(engines, TypeInfo::Never, None),
+    ))
+}
+
 /// Signature: `__ptr_add(ptr: raw_ptr, offset: u64)`
 /// Description: Adds `offset` to the raw value of pointer `ptr`.
 /// Constraints: None.

sway-core/src/semantic_analysis/cei_pattern_analysis.rs

@@ -608,7 +608,7 @@ fn effects_of_intrinsic(intr: &sway_ast::Intrinsic) -> HashSet<Effect> {
         StateClear | StateStoreWord | StateStoreQuad => HashSet::from([Effect::StorageWrite]),
         StateLoadWord | StateLoadQuad => HashSet::from([Effect::StorageRead]),
         Smo => HashSet::from([Effect::OutputMessage]),
-        Revert | IsReferenceType | IsStrArray | SizeOfType | SizeOfVal | SizeOfStr
+        Revert | JmpbSsp | IsReferenceType | IsStrArray | SizeOfType | SizeOfVal | SizeOfStr
         | AssertIsStrArray | ToStrArray | Eq | Gt | Lt | Gtf | AddrOf | Log | Add | Sub | Mul
         | Div | And | Or | Xor | Mod | Rsh | Lsh | PtrAdd | PtrSub | Not => HashSet::new(),
     }

sway-ir/src/analysis/memory_utils.rs

@@ -201,7 +201,7 @@ pub fn get_loaded_ptr_values(context: &Context, val: Value) -> Vec<Value> {
         InstOp::Store { dst_val_ptr: _, .. } => vec![],
         InstOp::FuelVm(FuelVmInstruction::Gtf { .. })
         | InstOp::FuelVm(FuelVmInstruction::ReadRegister(_))
-        | InstOp::FuelVm(FuelVmInstruction::Revert(_)) => vec![],
+        | InstOp::FuelVm(FuelVmInstruction::Revert(_) | FuelVmInstruction::JmpbSsp(_)) => vec![],
         InstOp::FuelVm(FuelVmInstruction::WideUnaryOp { arg, .. }) => vec![*arg],
         InstOp::FuelVm(FuelVmInstruction::WideBinaryOp { arg1, arg2, .. })
         | InstOp::FuelVm(FuelVmInstruction::WideCmpOp { arg1, arg2, .. }) => {
@@ -252,6 +252,7 @@ pub fn get_stored_ptr_values(context: &Context, val: Value) -> Vec<Value> {
             | FuelVmInstruction::Log { .. }
             | FuelVmInstruction::ReadRegister(_)
             | FuelVmInstruction::Revert(_)
+            | FuelVmInstruction::JmpbSsp(_)
             | FuelVmInstruction::Smo { .. }
             | FuelVmInstruction::StateClear { .. } => vec![],
             FuelVmInstruction::StateLoadQuadWord { load_val, .. } => vec![*load_val],

sway-ir/src/block.rs

@@ -368,7 +368,10 @@ impl Block {
             matches!(
                 i,
                 Instruction {
-                    op: InstOp::Ret(..) | InstOp::FuelVm(FuelVmInstruction::Revert(..)),
+                    op: InstOp::Ret(..)
+                        | InstOp::FuelVm(
+                            FuelVmInstruction::Revert(..) | FuelVmInstruction::JmpbSsp(..)
+                        ),
                     ..
                 }
             )

sway-ir/src/instruction.rs

@@ -194,6 +194,7 @@ pub enum FuelVmInstruction {
         arg1: Value,
         arg2: Value,
     },
+    JmpbSsp(Value),
 }
 
 /// Comparison operations.
@@ -298,7 +299,7 @@ impl InstOp {
             // These are all terminators which don't return, essentially.  No type.
             InstOp::Branch(_)
             | InstOp::ConditionalBranch { .. }
-            | InstOp::FuelVm(FuelVmInstruction::Revert(..))
+            | InstOp::FuelVm(FuelVmInstruction::Revert(..) | FuelVmInstruction::JmpbSsp(..))
             | InstOp::Ret(..) => None,
 
             // No-op is also no-type.
@@ -407,7 +408,7 @@ impl InstOp {
                     log_val, log_id, ..
                 } => vec![*log_val, *log_id],
                 FuelVmInstruction::ReadRegister(_) => vec![],
-                FuelVmInstruction::Revert(v) => vec![*v],
+                FuelVmInstruction::Revert(v) | FuelVmInstruction::JmpbSsp(v) => vec![*v],
                 FuelVmInstruction::Smo {
                     recipient,
                     message,
@@ -544,6 +545,7 @@ impl InstOp {
                 }
                 FuelVmInstruction::ReadRegister { .. } => (),
                 FuelVmInstruction::Revert(revert_val) => replace(revert_val),
+                FuelVmInstruction::JmpbSsp(contr_id) => replace(contr_id),
                 FuelVmInstruction::Smo {
                     recipient,
                     message,
@@ -629,7 +631,7 @@ impl InstOp {
             | InstOp::FuelVm(FuelVmInstruction::StateLoadQuadWord { .. })
             | InstOp::FuelVm(FuelVmInstruction::StateStoreQuadWord { .. })
             | InstOp::FuelVm(FuelVmInstruction::StateStoreWord { .. })
-            | InstOp::FuelVm(FuelVmInstruction::Revert(..))
+            | InstOp::FuelVm(FuelVmInstruction::Revert(..) | FuelVmInstruction::JmpbSsp(..))
             | InstOp::MemCopyBytes { .. }
             | InstOp::MemCopyVal { .. }
             | InstOp::Store { .. }
@@ -664,7 +666,7 @@ impl InstOp {
             InstOp::Branch(_)
                 | InstOp::ConditionalBranch { .. }
                 | InstOp::Ret(..)
-                | InstOp::FuelVm(FuelVmInstruction::Revert(..))
+                | InstOp::FuelVm(FuelVmInstruction::Revert(..) | FuelVmInstruction::JmpbSsp(..))
         )
     }
 }
@@ -1072,6 +1074,18 @@ impl<'a, 'eng> InstructionInserter<'a, 'eng> {
         revert_val
     }
 
+    pub fn jmpb_ssp(self, offset: Value) -> Value {
+        let ldc_exec = Value::new_instruction(
+            self.context,
+            self.block,
+            InstOp::FuelVm(FuelVmInstruction::JmpbSsp(offset)),
+        );
+        self.context.blocks[self.block.0]
+            .instructions
+            .push(ldc_exec);
+        ldc_exec
+    }
+
     pub fn smo(self, recipient: Value, message: Value, message_size: Value, coins: Value) -> Value {
         insert_instruction!(
             self,

sway-ir/src/optimize/fn_dedup.rs

@@ -154,6 +154,7 @@ fn hash_fn(context: &Context, function: Function, eq_class: &mut EqClass) -> u64
                     crate::FuelVmInstruction::Log { log_ty, .. } => log_ty.hash(state),
                     crate::FuelVmInstruction::ReadRegister(reg) => reg.hash(state),
                     crate::FuelVmInstruction::Revert(_)
+                    | crate::FuelVmInstruction::JmpbSsp(_)
                     | crate::FuelVmInstruction::Smo { .. }
                     | crate::FuelVmInstruction::StateClear { .. }
                     | crate::FuelVmInstruction::StateLoadQuadWord { .. }

sway-ir/src/optimize/inline.rs

@@ -535,6 +535,9 @@ fn inline_instruction(
                     new_block.append(context).read_register(reg)
                 }
                 FuelVmInstruction::Revert(val) => new_block.append(context).revert(map_value(val)),
+                FuelVmInstruction::JmpbSsp(offset) => {
+                    new_block.append(context).jmpb_ssp(map_value(offset))
+                }
                 FuelVmInstruction::Smo {
                     recipient,
                     message,

sway-ir/src/parser.rs

@@ -203,6 +203,7 @@ mod ir_builder {
                 / op_read_register()
                 / op_ret()
                 / op_revert()
+                / op_jmpb_ssp()
                 / op_smo()
                 / op_state_load_quad_word()
                 / op_state_load_word()
@@ -362,6 +363,11 @@ mod ir_builder {
                     IrAstOperation::Revert(vn)
                 }
 
+            rule op_jmpb_ssp() -> IrAstOperation
+                = "jmpb_ssp" _ vn:id() {
+                    IrAstOperation::JmpbSsp(vn)
+                }
+
             rule op_smo() -> IrAstOperation
                 = "smo" _
                 recipient_and_message:id() comma() message_size:id() comma() output_index:id() comma() coins:id() _ {
@@ -722,6 +728,7 @@ mod ir_builder {
         ReadRegister(String),
         Ret(IrAstTy, String),
         Revert(String),
+        JmpbSsp(String),
         Smo(String, String, String, String),
         StateClear(String, String),
         StateLoadQuadWord(String, String, String),
@@ -1344,6 +1351,10 @@ mod ir_builder {
                         .append(context)
                         .revert(*val_map.get(&ret_val_name).unwrap())
                         .add_metadatum(context, opt_metadata),
+                    IrAstOperation::JmpbSsp(offset_name) => block
+                        .append(context)
+                        .jmpb_ssp(*val_map.get(&offset_name).unwrap())
+                        .add_metadatum(context, opt_metadata),
                     IrAstOperation::Smo(recipient, message, message_size, coins) => block
                         .append(context)
                         .smo(

sway-ir/src/printer.rs

@@ -656,6 +656,12 @@ fn instruction_to_doc<'a>(
                         Doc::text(format!("revert {}", namer.name(context, v),))
                             .append(md_namer.md_idx_to_doc(context, metadata)),
                     )),
+                FuelVmInstruction::JmpbSsp(offset) => {
+                    maybe_constant_to_doc(context, md_namer, namer, offset).append(Doc::line(
+                        Doc::text(format!("jmpb_ssp {}", namer.name(context, offset),))
+                            .append(md_namer.md_idx_to_doc(context, metadata)),
+                    ))
+                }
                 FuelVmInstruction::Smo {
                     recipient,
                     message,

sway-ir/src/value.rs

@@ -122,7 +122,7 @@ impl Value {
                 InstOp::Branch(_)
                     | InstOp::ConditionalBranch { .. }
                     | InstOp::Ret(_, _)
-                    | InstOp::FuelVm(FuelVmInstruction::Revert(_))
+                    | InstOp::FuelVm(FuelVmInstruction::Revert(_) | FuelVmInstruction::JmpbSsp(_))
             ),
             ValueDatum::Argument(..) | ValueDatum::Configurable(..) | ValueDatum::Constant(..) => {
                 false

sway-ir/src/verify.rs

@@ -230,6 +230,7 @@ impl<'a, 'eng> InstructionVerifier<'a, 'eng> {
                             log_id,
                         } => self.verify_log(log_val, log_ty, log_id)?,
                         FuelVmInstruction::ReadRegister(_) => (),
+                        FuelVmInstruction::JmpbSsp(_) => (),
                         FuelVmInstruction::Revert(val) => self.verify_revert(val)?,
                         FuelVmInstruction::Smo {
                             recipient,