Intrinsics for the VM's storage opcodes. (FuelLabs#2508)

* Add storage load/store word intrinsics

* Storage load/store quad intrinsics

* Add asm gen test

* Update storage.sw and add exhaustive tests

sway-ast/src/intrinsics.rs

@@ -9,6 +9,10 @@ pub enum Intrinsic {
     Eq,
     Gtf,
     AddrOf,
+    StateLoadWord,
+    StateStoreWord,
+    StateLoadQuad,
+    StateStoreQuad,
 }
 
 impl fmt::Display for Intrinsic {
@@ -21,6 +25,10 @@ impl fmt::Display for Intrinsic {
             Intrinsic::Eq => "eq",
             Intrinsic::Gtf => "gtf",
             Intrinsic::AddrOf => "addr_of",
+            Intrinsic::StateLoadWord => "state_load_word",
+            Intrinsic::StateStoreWord => "state_store_word",
+            Intrinsic::StateLoadQuad => "state_load_quad",
+            Intrinsic::StateStoreQuad => "state_store_quad",
         };
         write!(f, "{}", s)
     }
@@ -37,6 +45,10 @@ impl Intrinsic {
             "__eq" => Eq,
             "__gtf" => Gtf,
             "__addr_of" => AddrOf,
+            "__state_load_word" => StateLoadWord,
+            "__state_store_word" => StateStoreWord,
+            "__state_load_quad" => StateLoadQuad,
+            "__state_store_quad" => StateStoreQuad,
             _ => return None,
         })
     }

sway-core/src/asm_generation/from_ir.rs

@@ -1598,6 +1598,7 @@ impl<'ir> AsmBuilder<'ir> {
             key.get_stripped_ptr_type(self.context),
             Some(Type::B256)
         ));
+        let owning_span = self.md_mgr.val_to_span(self.context, *instr_val);
 
         let key_ptr = self.resolve_ptr(key);
         if key_ptr.value.is_none() {
@@ -1609,39 +1610,50 @@ impl<'ir> AsmBuilder<'ir> {
         assert!(offset == 0);
         assert!(ptr_ty.get_type(self.context).eq(self.context, &Type::B256));
 
-        // Expect ptr_ty here to also be b256 and offset to be whatever...
-        let val_ptr = self.resolve_ptr(val);
-        if val_ptr.value.is_none() {
-            return val_ptr.map(|_| ());
-        }
-        let (val_ptr, ptr_ty, offset) = val_ptr.value.unwrap();
-
-        // Expect the ptr_ty for val to also be B256
-        assert!(ptr_ty.get_type(self.context).eq(self.context, &Type::B256));
+        let val_reg = if matches!(
+            &self.context.values[val.0].value,
+            ValueDatum::Instruction(Instruction::IntToPtr(..))
+        ) {
+            match self.reg_map.get(val) {
+                Some(vreg) => vreg.clone(),
+                None => unreachable!("int_to_ptr instruction doesn't have vreg mapped"),
+            }
+        } else {
+            // Expect ptr_ty here to also be b256 and offset to be whatever...
+            let val_ptr = self.resolve_ptr(val);
+            if val_ptr.value.is_none() {
+                return val_ptr.map(|_| ());
+            }
+            let (val_ptr, ptr_ty, offset) = val_ptr.value.unwrap();
+            // Expect the ptr_ty for val to also be B256
+            assert!(ptr_ty.get_type(self.context).eq(self.context, &Type::B256));
+            match self.ptr_map.get(&val_ptr) {
+                Some(Storage::Stack(val_offset)) => {
+                    let base_reg = self.stack_base_reg.as_ref().unwrap().clone();
+                    let val_offset_in_bytes = val_offset * 8 + offset * 32;
+                    self.offset_reg(&base_reg, val_offset_in_bytes, owning_span.clone())
+                }
+                _ => unreachable!("Unexpected storage locations for key and val"),
+            }
+        };
 
-        let owning_span = self.md_mgr.val_to_span(self.context, *instr_val);
-        match (self.ptr_map.get(&val_ptr), self.ptr_map.get(&key_ptr)) {
-            (Some(Storage::Stack(val_offset)), Some(Storage::Stack(key_offset))) => {
+        let key_reg = match self.ptr_map.get(&key_ptr) {
+            Some(Storage::Stack(key_offset)) => {
                 let base_reg = self.stack_base_reg.as_ref().unwrap().clone();
-                let val_offset_in_bytes = val_offset * 8 + offset * 32;
                 let key_offset_in_bytes = key_offset * 8;
-
-                let val_reg = self.offset_reg(&base_reg, val_offset_in_bytes, owning_span.clone());
-
-                let key_reg = self.offset_reg(&base_reg, key_offset_in_bytes, owning_span.clone());
-
-                self.bytecode.push(Op {
-                    opcode: Either::Left(match access_type {
-                        StateAccessType::Read => VirtualOp::SRWQ(val_reg, key_reg),
-                        StateAccessType::Write => VirtualOp::SWWQ(key_reg, val_reg),
-                    }),
-                    comment: "quad word state access".into(),
-                    owning_span,
-                });
+                self.offset_reg(&base_reg, key_offset_in_bytes, owning_span.clone())
             }
             _ => unreachable!("Unexpected storage locations for key and val"),
-        }
+        };
 
+        self.bytecode.push(Op {
+            opcode: Either::Left(match access_type {
+                StateAccessType::Read => VirtualOp::SRWQ(val_reg, key_reg),
+                StateAccessType::Write => VirtualOp::SWWQ(key_reg, val_reg),
+            }),
+            comment: "quad word state access".into(),
+            owning_span,
+        });
         ok((), Vec::new(), Vec::new())
     }
 

sway-core/src/ir_generation/function.rs

@@ -8,14 +8,14 @@ use super::{
 use crate::{
     asm_generation::from_ir::ir_type_size_in_bytes,
     constants,
-    error::CompileError,
+    error::{CompileError, Hint},
     ir_generation::const_eval::{
         compile_constant_expression, compile_constant_expression_to_constant,
     },
     metadata::MetadataManager,
     parse_tree::{AsmOp, AsmRegister, LazyOp, Literal},
     semantic_analysis::*,
-    type_system::{look_up_type_id, resolve_type, TypeId, TypeInfo},
+    type_system::{look_up_type_id, resolve_type, IntegerBits, TypeId, TypeInfo},
 };
 use sway_ast::intrinsics::Intrinsic;
 use sway_ir::{Context, *};
@@ -391,6 +391,41 @@ impl FnCompiler {
         }: TypedIntrinsicFunctionKind,
         span: Span,
     ) -> Result<Value, CompileError> {
+        fn store_key_in_local_mem(
+            compiler: &mut FnCompiler,
+            context: &mut Context,
+            value: Value,
+            span_md_idx: Option<MetadataIndex>,
+        ) -> Result<Value, CompileError> {
+            // New name for the key
+            let key_name = "key_for_storage".to_string();
+            let alias_key_name = compiler.lexical_map.insert(key_name.as_str().to_owned());
+
+            // Local pointer for the key
+            let key_ptr = compiler
+                .function
+                .new_local_ptr(context, alias_key_name, Type::B256, true, None)
+                .map_err(|ir_error| {
+                    CompileError::InternalOwned(ir_error.to_string(), Span::dummy())
+                })?;
+
+            // Convert the key pointer to a value using get_ptr
+            let key_ptr_ty = *key_ptr.get_type(context);
+            let key_ptr_val = compiler
+                .current_block
+                .ins(context)
+                .get_ptr(key_ptr, key_ptr_ty, 0)
+                .add_metadatum(context, span_md_idx);
+
+            // Store the value to the key pointer value
+            compiler
+                .current_block
+                .ins(context)
+                .store(key_ptr_val, value)
+                .add_metadatum(context, span_md_idx);
+            Ok(key_ptr_val)
+        }
+
         // We safely index into arguments and type_arguments arrays below
         // because the type-checker ensures that the arguments are all there.
         match kind {
@@ -497,6 +532,77 @@ impl FnCompiler {
                     .addr_of(value)
                     .add_metadatum(context, span_md_idx))
             }
+            Intrinsic::StateLoadWord => {
+                let exp = arguments[0].clone();
+                let value = self.compile_expression(context, md_mgr, exp)?;
+                let span_md_idx = md_mgr.span_to_md(context, &span);
+                let key_ptr_val = store_key_in_local_mem(self, context, value, span_md_idx)?;
+                Ok(self
+                    .current_block
+                    .ins(context)
+                    .state_load_word(key_ptr_val)
+                    .add_metadatum(context, span_md_idx))
+            }
+            Intrinsic::StateStoreWord => {
+                let key_exp = arguments[0].clone();
+                let val_exp = arguments[1].clone();
+                // Validate that the val_exp is of the right type. We couldn't do it
+                // earlier during type checking as the type arguments may not have been resolved.
+                let val_ty = resolve_type(val_exp.return_type, &span).unwrap();
+                if !val_ty.is_copy_type() {
+                    return Err(CompileError::IntrinsicUnsupportedArgType {
+                        name: kind.to_string(),
+                        span,
+                        hint: Hint::new("This argument must be a copy type".to_string()),
+                    });
+                }
+                let key_value = self.compile_expression(context, md_mgr, key_exp)?;
+                let val_value = self.compile_expression(context, md_mgr, val_exp)?;
+                let span_md_idx = md_mgr.span_to_md(context, &span);
+                let key_ptr_val = store_key_in_local_mem(self, context, key_value, span_md_idx)?;
+                Ok(self
+                    .current_block
+                    .ins(context)
+                    .state_store_word(val_value, key_ptr_val)
+                    .add_metadatum(context, span_md_idx))
+            }
+            Intrinsic::StateLoadQuad | Intrinsic::StateStoreQuad => {
+                let key_exp = arguments[0].clone();
+                let val_exp = arguments[1].clone();
+                // Validate that the val_exp is of the right type. We couldn't do it
+                // earlier during type checking as the type arguments may not have been resolved.
+                let val_ty = resolve_type(val_exp.return_type, &span).unwrap();
+                if val_ty != TypeInfo::UnsignedInteger(IntegerBits::SixtyFour) {
+                    return Err(CompileError::IntrinsicUnsupportedArgType {
+                        name: kind.to_string(),
+                        span,
+                        hint: Hint::new("This argument must be u64".to_string()),
+                    });
+                }
+                let key_value = self.compile_expression(context, md_mgr, key_exp)?;
+                let val_value = self.compile_expression(context, md_mgr, val_exp)?;
+                let span_md_idx = md_mgr.span_to_md(context, &span);
+                let key_ptr_val = store_key_in_local_mem(self, context, key_value, span_md_idx)?;
+                // For quad word, the IR instructions take in a pointer rather than a raw u64.
+                let val_ptr = self
+                    .current_block
+                    .ins(context)
+                    .int_to_ptr(val_value, Type::B256)
+                    .add_metadatum(context, span_md_idx);
+                match kind {
+                    Intrinsic::StateLoadQuad => Ok(self
+                        .current_block
+                        .ins(context)
+                        .state_load_quad_word(val_ptr, key_ptr_val)
+                        .add_metadatum(context, span_md_idx)),
+                    Intrinsic::StateStoreQuad => Ok(self
+                        .current_block
+                        .ins(context)
+                        .state_store_quad_word(val_ptr, key_ptr_val)
+                        .add_metadatum(context, span_md_idx)),
+                    _ => unreachable!(),
+                }
+            }
         }
     }
 

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

@@ -375,6 +375,119 @@ impl TypedIntrinsicFunctionKind {
                 let return_type = insert_type(TypeInfo::UnsignedInteger(IntegerBits::SixtyFour));
                 (intrinsic_function, return_type)
             }
+            Intrinsic::StateLoadWord => {
+                if arguments.len() != 1 {
+                    errors.push(CompileError::IntrinsicIncorrectNumArgs {
+                        name: kind.to_string(),
+                        expected: 1,
+                        span,
+                    });
+                    return err(warnings, errors);
+                }
+                let ctx = ctx
+                    .with_help_text("")
+                    .with_type_annotation(insert_type(TypeInfo::Unknown));
+                let exp = check!(
+                    TypedExpression::type_check(ctx, arguments[0].clone()),
+                    return err(warnings, errors),
+                    warnings,
+                    errors
+                );
+                let key_ty = resolve_type(exp.return_type, &span).unwrap();
+                if key_ty != TypeInfo::B256 {
+                    errors.push(CompileError::IntrinsicUnsupportedArgType {
+                        name: kind.to_string(),
+                        span,
+                        hint: Hint::new(
+                            "Argument type must be B256, a key into the state storage".to_string(),
+                        ),
+                    });
+                    return err(warnings, errors);
+                }
+                let intrinsic_function = TypedIntrinsicFunctionKind {
+                    kind,
+                    arguments: vec![exp],
+                    type_arguments: vec![],
+                    span,
+                };
+                let return_type = insert_type(TypeInfo::UnsignedInteger(IntegerBits::SixtyFour));
+                (intrinsic_function, return_type)
+            }
+            Intrinsic::StateStoreWord | Intrinsic::StateLoadQuad | Intrinsic::StateStoreQuad => {
+                if arguments.len() != 2 {
+                    errors.push(CompileError::IntrinsicIncorrectNumArgs {
+                        name: kind.to_string(),
+                        expected: 2,
+                        span,
+                    });
+                    return err(warnings, errors);
+                }
+                if type_arguments.len() > 1 {
+                    errors.push(CompileError::IntrinsicIncorrectNumTArgs {
+                        name: kind.to_string(),
+                        expected: 1,
+                        span,
+                    });
+                    return err(warnings, errors);
+                }
+                let mut ctx = ctx
+                    .with_help_text("")
+                    .with_type_annotation(insert_type(TypeInfo::Unknown));
+                let key_exp = check!(
+                    TypedExpression::type_check(ctx.by_ref(), arguments[0].clone()),
+                    return err(warnings, errors),
+                    warnings,
+                    errors
+                );
+                let key_ty = resolve_type(key_exp.return_type, &span).unwrap();
+                if key_ty != TypeInfo::B256 {
+                    errors.push(CompileError::IntrinsicUnsupportedArgType {
+                        name: kind.to_string(),
+                        span,
+                        hint: Hint::new(
+                            "Argument type must be B256, a key into the state storage".to_string(),
+                        ),
+                    });
+                    return err(warnings, errors);
+                }
+                let mut ctx = ctx
+                    .with_help_text("")
+                    .with_type_annotation(insert_type(TypeInfo::Unknown));
+                let val_exp = check!(
+                    TypedExpression::type_check(ctx.by_ref(), arguments[1].clone()),
+                    return err(warnings, errors),
+                    warnings,
+                    errors
+                );
+                let type_argument = type_arguments.get(0).map(|targ| {
+                    let mut ctx = ctx
+                        .with_help_text("")
+                        .with_type_annotation(insert_type(TypeInfo::Unknown));
+                    let type_id = check!(
+                        ctx.resolve_type_with_self(
+                            insert_type(resolve_type(targ.type_id, &targ.span).unwrap()),
+                            &targ.span,
+                            EnforceTypeArguments::Yes,
+                            None
+                        ),
+                        insert_type(TypeInfo::ErrorRecovery),
+                        warnings,
+                        errors,
+                    );
+                    TypeArgument {
+                        type_id,
+                        span: span.clone(),
+                    }
+                });
+                let intrinsic_function = TypedIntrinsicFunctionKind {
+                    kind,
+                    arguments: vec![key_exp, val_exp],
+                    type_arguments: type_argument.map_or(vec![], |ta| vec![ta]),
+                    span,
+                };
+                let return_type = insert_type(TypeInfo::Tuple(vec![]));
+                (intrinsic_function, return_type)
+            }
         };
         ok((intrinsic_function, return_type), warnings, errors)
     }