Introduce index() and index_assign() to StorageMap and showcase nested maps

using both `get/insert` and the `[]` operator

Author: mohammadfawaz

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

@@ -1674,14 +1674,14 @@ impl ty::TyExpression {
                 errors,
             )
         } else {
-            // Otherwise convert into a method call 'index(self, index)' via the std::ops::Index trait.
+            // Otherwise convert into a method call `index(self, index)`
             let method_name = TypeBinding {
                 inner: MethodName::FromTrait {
                     call_path: CallPath {
-                        prefixes: vec![
-                            Ident::new_with_override("core".into(), span.clone()),
-                            Ident::new_with_override("ops".into(), span.clone()),
-                        ],
+                        // Eventually, this should be `core::ops::index` once we are able to implement
+                        // the `Index` trait, but for now, we assume that `index` is part of `impl`
+                        // self for the type of `prefix`.
+                        prefixes: vec![],
                         suffix: Ident::new_with_override("index".into(), span.clone()),
                         is_absolute: true,
                     },
@@ -1788,8 +1788,10 @@ impl ty::TyExpression {
             ReassignmentTarget::VariableExpression(var) => {
                 let mut expr = var;
                 let mut names_vec = Vec::new();
-                let (base_name, final_return_type) = loop {
-                    match expr.kind {
+                let mut projection_index = 0;
+                let mut first_array_index_expr = None;
+                let base_name_and_final_return_type = loop {
+                    match expr.kind.clone() {
                         ExpressionKind::Variable(name) => {
                             // check that the reassigned name exists
                             let unknown_decl = check!(
@@ -1808,7 +1810,7 @@ impl ty::TyExpression {
                                 errors.push(CompileError::AssignmentToNonMutable { name, span });
                                 return err(warnings, errors);
                             }
-                            break (name, variable_decl.body.return_type);
+                            break Some((name, variable_decl.body.return_type));
                         }
                         ExpressionKind::Subfield(SubfieldExpression {
                             prefix,
@@ -1830,6 +1832,19 @@ impl ty::TyExpression {
                             expr = prefix;
                         }
                         ExpressionKind::ArrayIndex(ArrayIndexExpression { prefix, index }) => {
+                            // If this is the right most project (i.e. the first, since the we
+                            // start counting from the right), and if this projection is an array
+                            // index projection, then keep track of the `prefix` and the `index` in
+                            // this case: we may need to call `index_assign()` later on if the
+                            // compiler does not offer an intrinsic way of dealing with this
+                            // reassignment
+                            if projection_index == 0 {
+                                first_array_index_expr = Some(ArrayIndexExpression {
+                                    prefix: prefix.clone(),
+                                    index: index.clone(),
+                                });
+                            }
+
                             let ctx = ctx.by_ref().with_help_text("");
                             let typed_index = check!(
                                 ty::TyExpression::type_check(ctx, index.as_ref().clone()),
@@ -1839,50 +1854,88 @@ impl ty::TyExpression {
                             );
                             names_vec.push(ty::ProjectionKind::ArrayIndex {
                                 index: Box::new(typed_index),
-                                index_span: index.span(),
+                                index_span: index.span().clone(),
                             });
                             expr = prefix;
                         }
-                        _ => {
-                            errors.push(CompileError::InvalidExpressionOnLhs { span });
-                            return err(warnings, errors);
-                        }
+                        _ => break None,
                     }
+                    projection_index += 1;
                 };
-                let names_vec = names_vec.into_iter().rev().collect::<Vec<_>>();
-                let (ty_of_field, _ty_of_parent) = check!(
-                    ctx.namespace
-                        .find_subfield_type(ctx.engines(), &base_name, &names_vec),
-                    return err(warnings, errors),
-                    warnings,
-                    errors
-                );
-                // type check the reassignment
-                let ctx = ctx.with_type_annotation(ty_of_field).with_help_text("");
-                let rhs_span = rhs.span();
-                let rhs = check!(
-                    ty::TyExpression::type_check(ctx, rhs),
-                    ty::TyExpression::error(rhs_span, engines),
-                    warnings,
-                    errors
-                );
 
-                ok(
-                    ty::TyExpression {
-                        expression: ty::TyExpressionVariant::Reassignment(Box::new(
-                            ty::TyReassignment {
-                                lhs_base_name: base_name,
-                                lhs_type: final_return_type,
-                                lhs_indices: names_vec,
-                                rhs,
+                match base_name_and_final_return_type {
+                    Some((base_name, final_return_type)) => {
+                        let names_vec = names_vec.into_iter().rev().collect::<Vec<_>>();
+                        let (ty_of_field, _ty_of_parent) = check!(
+                            ctx.namespace
+                                .find_subfield_type(ctx.engines(), &base_name, &names_vec),
+                            return err(warnings, errors),
+                            warnings,
+                            errors
+                        );
+                        // type check the reassignment
+                        let ctx = ctx.with_type_annotation(ty_of_field).with_help_text("");
+                        let rhs_span = rhs.span();
+                        let rhs = check!(
+                            ty::TyExpression::type_check(ctx, rhs),
+                            ty::TyExpression::error(rhs_span, engines),
+                            warnings,
+                            errors
+                        );
+
+                        ok(
+                            ty::TyExpression {
+                                expression: ty::TyExpressionVariant::Reassignment(Box::new(
+                                    ty::TyReassignment {
+                                        lhs_base_name: base_name,
+                                        lhs_type: final_return_type,
+                                        lhs_indices: names_vec,
+                                        rhs,
+                                    },
+                                )),
+                                return_type: type_engine
+                                    .insert(decl_engine, TypeInfo::Tuple(Vec::new())),
+                                span,
                             },
-                        )),
-                        return_type: type_engine.insert(decl_engine, TypeInfo::Tuple(Vec::new())),
-                        span,
+                            warnings,
+                            errors,
+                        )
+                    }
+                    None => match first_array_index_expr {
+                        Some(ArrayIndexExpression { prefix, index }) => {
+                            let method_name = TypeBinding {
+                                inner: MethodName::FromTrait {
+                                    call_path: CallPath {
+                                        // Eventually, this should be `core::ops::index_assign`
+                                        // once we are able to implement the `IndexAssign` trait,
+                                        // but for now, we assume that `index_assign` is part of
+                                        // `impl` self for the type of `prefix`.
+                                        prefixes: vec![],
+                                        suffix: Ident::new_with_override(
+                                            "index_assign".into(),
+                                            span.clone(),
+                                        ),
+                                        is_absolute: true,
+                                    },
+                                },
+                                type_arguments: TypeArgs::Regular(vec![]),
+                                span: span.clone(),
+                            };
+
+                            type_check_method_application(
+                                ctx,
+                                method_name,
+                                vec![],
+                                vec![*prefix, *index, rhs],
+                                span,
+                            )
+                        }
+                        None => {
+                            errors.push(CompileError::InvalidExpressionOnLhs { span });
+                            err(warnings, errors)
+                        }
                     },
-                    warnings,
-                    errors,
-                )
+                }
             }
             ReassignmentTarget::StorageField(storage_keyword_span, fields) => {
                 let ctx = ctx

sway-lib-std/src/experimental/storage.sw

@@ -24,6 +24,10 @@ use core::experimental::storage::StorageKey;
 /// ```
 #[storage(read, write)]
 pub fn write<T>(key: b256, offset: u64, value: T) {
+    if __size_of::<T>() == 0 {
+        return;
+    }
+
     // Get the number of storage slots needed based on the size of `T`
     let number_of_slots = (offset * 8 + __size_of::<T>() + 31) >> 5;
 
@@ -68,6 +72,10 @@ pub fn write<T>(key: b256, offset: u64, value: T) {
 /// ```
 #[storage(read)]
 pub fn read<T>(key: b256, offset: u64) -> Option<T> {
+    if __size_of::<T>() == 0 {
+        return Option::None;
+    }
+
     // NOTE: we are leaking this value on the heap.
     // Get the number of storage slots needed based on the size of `T`
     let number_of_slots = (offset * 8 + __size_of::<T>() + 31) >> 5;
@@ -103,7 +111,7 @@ pub fn read<T>(key: b256, offset: u64) -> Option<T> {
 /// assert(read::<u64>(ZERO_B256, 0).is_none());
 /// ```
 #[storage(write)]
-pub fn clear<T>(key: b256) -> bool {
+fn clear<T>(key: b256) -> bool {
     // Get the number of storage slots needed based on the size of `T` as the ceiling of 
     // `__size_of::<T>() / 32`
     let number_of_slots = (__size_of::<T>() + 31) >> 5;
@@ -116,13 +124,13 @@ impl<T> StorageKey<T> {
     /// Reads a value of type `T` starting at the location specified by `self`. If the value
     /// crosses the boundary of a storage slot, reading continues at the following slot.
     ///
-    /// Returns the value previously stored if a the storage slots read were
-    /// valid and contain `value`. Panics otherwise.
+    /// Returns the value previously stored if the storage slots read were valid and contain 
+    /// `value`. Reverts otherwise.
     ///
-    /// ### Arguments
-    ///
-    /// None
+    /// ### Reverts
     ///
+    /// Reverts if at least one of the storage slots needed to read a value of type `T` is not set.
+    /// 
     /// ### Examples
     ///
     /// ```sway
@@ -225,6 +233,38 @@ impl<K, V> StorageKey<StorageMap<K, V>> {
         write::<V>(key, 0, value);
     }
 
+    /// Retrieves the `StorageKey` that describes the raw location in storage of the value
+    /// Inserts a key-value pair into the map using the `[]` operator
+    ///
+    /// This is temporary until we are able to implement `trait IndexAssign`. The Sway compiler will
+    /// de-sugar the index operator `[]` in an assignment expression to a call to `index_assign()`.
+    ///
+    /// ### Arguments
+    ///
+    /// * `key` - The key to which the value is paired.
+    /// * `value` - The value to be stored.
+    ///
+    /// ### Examples
+    ///
+    /// ```sway
+    /// storage {
+    ///     map: StorageMap<u64, bool> = StorageMap {}
+    /// }
+    ///
+    /// fn foo() {
+    ///     let key = 5_u64;
+    ///     let value = true;
+    ///     storage.map[key] = value; // de-sugars to `storage.map.index_assign(key, value);`
+    ///     let retrieved = storage.map.get(key).read();
+    ///     assert(value == retrieved);
+    /// }
+    /// ```
+    #[storage(read, write)]
+    pub fn index_assign(self, key: K, value: V) {
+        let key = sha256((key, self.key));
+        write::<V>(key, 0, value);
+    }
+
     /// Retrieves the `StorageKey` that describes the raw location in storage of the value
     /// stored at `key`, regardless of whether a value is actually stored at that location or not.
     ///
@@ -243,8 +283,8 @@ impl<K, V> StorageKey<StorageMap<K, V>> {
     ///     let key = 5_u64;
     ///     let value = true;
     ///     storage.map.insert(key, value);
-    ///     let retrieved_value = storage.map.get(key).read();
-    ///     assert(value == retrieved_value);
+    ///     let retrieved = storage.map.get(key).read();
+    ///     assert(value == retrieved);
     /// }
     /// ```
     #[storage(read)]
@@ -255,6 +295,37 @@ impl<K, V> StorageKey<StorageMap<K, V>> {
         }
     }
 
+    /// Retrieves the `StorageKey` that describes the raw location in storage of the value 
+    /// stored at `key`, regardless of whether a value is actually stored at that location or not.
+    ///
+    /// This is temporary until we are able to implement `trait Index`. The Sway compiler will
+    /// de-sugar the index operator `[]` in an expression to a call to `index()`.
+    ///
+    /// ### Arguments
+    ///
+    /// * `key` - The key to which the value is paired.
+    ///
+    /// ### Examples
+    ///
+    /// ```sway
+    /// storage {
+    ///     map: StorageMap<u64, bool> = StorageMap {}
+    /// }
+    ///
+    /// fn foo() {
+    ///     let key = 5_u64;
+    ///     let value = true;
+    ///     storage.map.insert(key, value);
+    ///     let retrieved = storage.map[key].read(); // de-sugars to `storage.map.get(key).read()`
+    ///     assert(value == retrieved);
+    /// }
+    pub fn index(self, key: K) -> StorageKey<V> {
+        StorageKey {
+            key: sha256((key, self.key)),
+            offset: 0,
+        }
+    }
+
     /// Clears a value previously stored using a key
     ///
     /// Return a Boolean indicating whether there was a value previously stored at `key`.

test/src/sdk-harness/test_projects/experimental_storage_nested_maps/Forc.lock

@@ -0,0 +1,13 @@
+[[package]]
+name = 'core'
+source = 'path+from-root-2064A4F50B965AB3'
+
+[[package]]
+name = 'experimental_storage_nested_maps'
+source = 'member'
+dependencies = ['std']
+
+[[package]]
+name = 'std'
+source = 'path+from-root-2064A4F50B965AB3'
+dependencies = ['core']

test/src/sdk-harness/test_projects/experimental_storage_nested_maps/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "experimental_storage_nested_maps"
+
+[dependencies]
+std = { path = "../../../../../sway-lib-std" }

test/src/sdk-harness/test_projects/experimental_storage_nested_maps/mod.rs

@@ -0,0 +1,47 @@
+use fuels::prelude::*;
+
+abigen!(Contract(
+    name = "TestExperimentalStorageNestedMapsContract",
+    abi = "test_projects/experimental_storage_nested_maps/out/debug/experimental_storage_nested_maps-abi.json",
+));
+
+async fn test_experimental_storage_nested_maps_instance(
+) -> TestExperimentalStorageNestedMapsContract<WalletUnlocked> {
+    let wallet = launch_provider_and_get_wallet().await;
+    let id = Contract::deploy(
+        "test_projects/experimental_storage_nested_maps/out/debug/experimental_storage_nested_maps.bin",
+        &wallet,
+        DeployConfiguration::default(),
+    )
+    .await
+    .unwrap();
+
+    TestExperimentalStorageNestedMapsContract::new(id.clone(), wallet)
+}
+
+#[tokio::test]
+async fn nested_map_1_access() {
+    let methods = test_experimental_storage_nested_maps_instance()
+        .await
+        .methods();
+
+    methods.nested_map_1_access().call().await.unwrap();
+}
+
+#[tokio::test]
+async fn nested_map_2_access() {
+    let methods = test_experimental_storage_nested_maps_instance()
+        .await
+        .methods();
+
+    methods.nested_map_2_access().call().await.unwrap();
+}
+
+#[tokio::test]
+async fn nested_map_3_access() {
+    let methods = test_experimental_storage_nested_maps_instance()
+        .await
+        .methods();
+
+    methods.nested_map_3_access().call().await.unwrap();
+}