pass assignment ptr

Signed-off-by: Cyrill Leutwiler <bigcyrill@hotmail.com>

crates/llvm-context/src/polkavm/context/pointer/storage.rs

@@ -2,6 +2,8 @@
 
 use inkwell::values::BasicValueEnum;
 
+use crate::polkavm::context::address_space::AddressSpace;
+use crate::polkavm::context::pointer::Pointer;
 use crate::polkavm::context::runtime::RuntimeFunction;
 use crate::polkavm::context::Context;
 use crate::polkavm::Dependency;
@@ -17,20 +19,27 @@ where
     const NAME: &'static str = "__revive_load_storage_word";
 
     fn r#type<'ctx>(context: &Context<'ctx, D>) -> inkwell::types::FunctionType<'ctx> {
-        context
-            .word_type()
-            .fn_type(&[context.llvm().ptr_type(Default::default()).into()], false)
+        context.void_type().fn_type(
+            &[
+                context.llvm().ptr_type(Default::default()).into(),
+                context.llvm().ptr_type(Default::default()).into(),
+            ],
+            false,
+        )
     }
 
     fn emit_body<'ctx>(
         &self,
         context: &mut Context<'ctx, D>,
     ) -> anyhow::Result<Option<BasicValueEnum<'ctx>>> {
-        Ok(Some(emit_load(
-            context,
-            Self::paramater(context, 0),
-            false,
-        )?))
+        let key = Self::paramater(context, 0);
+        let assignment_pointer = Self::paramater(context, 1).into_pointer_value();
+        let value = emit_load(context, key, false)?;
+        context.build_store(
+            Pointer::new(context.word_type(), AddressSpace::Stack, assignment_pointer),
+            value,
+        )?;
+        Ok(None)
     }
 }
 

crates/llvm-context/src/polkavm/evm/storage.rs

@@ -13,16 +13,32 @@ use crate::PolkaVMStoreTransientStorageWordFunction;
 pub fn load<'ctx, D>(
     context: &mut Context<'ctx, D>,
     position: &PolkaVMArgument<'ctx>,
-) -> anyhow::Result<inkwell::values::BasicValueEnum<'ctx>>
+    assignment_pointer: Option<inkwell::values::PointerValue<'ctx>>,
+) -> anyhow::Result<Option<inkwell::values::BasicValueEnum<'ctx>>>
 where
     D: Dependency + Clone,
 {
-    let name = <PolkaVMLoadStorageWordFunction as RuntimeFunction<D>>::NAME;
+    let _name = <PolkaVMLoadStorageWordFunction as RuntimeFunction<D>>::NAME;
     let declaration = <PolkaVMLoadStorageWordFunction as RuntimeFunction<D>>::declaration(context);
-    let arguments = [position.as_pointer(context)?.value.into()];
-    Ok(context
-        .build_call(declaration, &arguments, "storage_load")
-        .unwrap_or_else(|| panic!("runtime function {name} should return a value")))
+    match assignment_pointer {
+        Some(assignment_pointer) => {
+            let arguments = [
+                position.as_pointer(context)?.value.into(),
+                assignment_pointer.into(),
+            ];
+            context.build_call(declaration, &arguments, "storage_load");
+            Ok(None)
+        }
+        None => {
+            let pointer = context.build_alloca_at_entry(context.word_type(), "pointer");
+            let arguments = [
+                position.as_pointer(context)?.value.into(),
+                pointer.value.into(),
+            ];
+            context.build_call(declaration, &arguments, "storage_load");
+            Ok(Some(context.build_load(pointer, "storage_value")?))
+        }
+    }
 }
 
 /// Translates the storage store.

crates/solidity/src/yul/parser/statement/assignment.rs

@@ -121,7 +121,7 @@ where
     ) -> anyhow::Result<()> {
         context.set_debug_location(self.location.line, 0, None)?;
 
-        let value = match self.initializer.into_llvm(context)? {
+        let value = match self.initializer.into_llvm(context, None)? {
             Some(value) => value,
             None => return Ok(()),
         };
@@ -142,7 +142,6 @@ where
             context.build_store(pointer, value.access(context)?)?;
             return Ok(());
         }
-
         let value = value.access(context)?;
         let llvm_type = value.into_struct_value().get_type();
         let tuple_pointer = context.build_alloca(llvm_type, "assignment_pointer");

crates/solidity/src/yul/parser/statement/block.rs

@@ -184,7 +184,7 @@ where
                     block.into_llvm(context)?;
                 }
                 Statement::Expression(expression) => {
-                    expression.into_llvm(context)?;
+                    expression.into_llvm(context, None)?;
                 }
                 Statement::VariableDeclaration(statement) => statement.into_llvm(context)?,
                 Statement::Assignment(statement) => statement.into_llvm(context)?,

crates/solidity/src/yul/parser/statement/expression/function_call/mod.rs

@@ -118,6 +118,7 @@ impl FunctionCall {
     pub fn into_llvm<'ctx, D>(
         mut self,
         context: &mut revive_llvm_context::PolkaVMContext<'ctx, D>,
+        assignment_pointer: Option<inkwell::values::PointerValue<'ctx>>,
     ) -> anyhow::Result<Option<inkwell::values::BasicValueEnum<'ctx>>>
     where
         D: revive_llvm_context::PolkaVMDependency + Clone,
@@ -129,7 +130,7 @@ impl FunctionCall {
                 let mut values = Vec::with_capacity(self.arguments.len());
                 for argument in self.arguments.into_iter().rev() {
                     let value = argument
-                        .into_llvm(context)?
+                        .into_llvm(context, None)?
                         .expect("Always exists")
                         .access(context)?;
                     values.push(value);
@@ -465,7 +466,11 @@ impl FunctionCall {
 
             Name::SLoad => {
                 let arguments = self.pop_arguments::<D, 1>(context)?;
-                revive_llvm_context::polkavm_evm_storage::load(context, &arguments[0]).map(Some)
+                revive_llvm_context::polkavm_evm_storage::load(
+                    context,
+                    &arguments[0],
+                    assignment_pointer,
+                )
             }
             Name::SStore => {
                 let arguments = self.pop_arguments::<D, 2>(context)?;
@@ -989,7 +994,7 @@ impl FunctionCall {
         for expression in self.arguments.drain(0..N).rev() {
             arguments.push(
                 expression
-                    .into_llvm(context)?
+                    .into_llvm(context, None)?
                     .expect("Always exists")
                     .access(context)?,
             );
@@ -1009,7 +1014,7 @@ impl FunctionCall {
     {
         let mut arguments = Vec::with_capacity(N);
         for expression in self.arguments.drain(0..N).rev() {
-            arguments.push(expression.into_llvm(context)?.expect("Always exists"));
+            arguments.push(expression.into_llvm(context, None)?.expect("Always exists"));
         }
         arguments.reverse();
 

crates/solidity/src/yul/parser/statement/expression/mod.rs

@@ -101,6 +101,7 @@ impl Expression {
     pub fn into_llvm<'ctx, D>(
         self,
         context: &mut revive_llvm_context::PolkaVMContext<'ctx, D>,
+        assignment_pointer: Option<inkwell::values::PointerValue<'ctx>>,
     ) -> anyhow::Result<Option<revive_llvm_context::PolkaVMArgument<'ctx>>>
     where
         D: revive_llvm_context::PolkaVMDependency + Clone,
@@ -139,7 +140,7 @@ impl Expression {
                 }))
             }
             Self::FunctionCall(call) => Ok(call
-                .into_llvm(context)?
+                .into_llvm(context, assignment_pointer)?
                 .map(revive_llvm_context::PolkaVMArgument::value)),
         }
     }

crates/solidity/src/yul/parser/statement/for_loop.rs

@@ -76,7 +76,7 @@ where
         context.set_basic_block(condition_block);
         let condition = self
             .condition
-            .into_llvm(context)?
+            .into_llvm(context, None)?
             .expect("Always exists")
             .access(context)?
             .into_int_value();

crates/solidity/src/yul/parser/statement/if_conditional.rs

@@ -55,7 +55,7 @@ where
     fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         let condition = self
             .condition
-            .into_llvm(context)?
+            .into_llvm(context, None)?
             .expect("Always exists")
             .access(context)?
             .into_int_value();

crates/solidity/src/yul/parser/statement/switch/mod.rs

@@ -123,7 +123,7 @@ where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
     fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
-        let scrutinee = self.expression.into_llvm(context)?;
+        let scrutinee = self.expression.into_llvm(context, None)?;
 
         if self.cases.is_empty() {
             if let Some(block) = self.default {

crates/solidity/src/yul/parser/statement/variable_declaration.rs

@@ -111,7 +111,7 @@ where
                 .insert_stack_pointer(identifier.inner.clone(), pointer);
 
             let value = if let Some(expression) = self.expression {
-                match expression.into_llvm(context)? {
+                match expression.into_llvm(context, None)? {
                     Some(mut value) => {
                         if let Some(constant) = value.constant.take() {
                             context
@@ -156,7 +156,7 @@ where
             None => return Ok(()),
         };
         let location = expression.location();
-        let expression = match expression.into_llvm(context)? {
+        let expression = match expression.into_llvm(context, None)? {
             Some(expression) => expression,
             None => return Ok(()),
         };