Support CastPtr in IR optimizations

sway-ir/src/irtype.rs

@@ -478,9 +478,21 @@ impl Type {
         })
     }
 
-    /// What's the offset, in bytes, of the indexed element?
-    /// Returns `None` on invalid indices.
-    /// Panics if `self` is not an aggregate (struct, union, or array).
+    /// Returns the offset, in bytes, of the indexed element of an aggregate `self`.
+    ///
+    /// Returns `None`:
+    /// - on invalid `indices`.
+    /// - if `self` is not an aggregate (struct, union, or array).
+    ///
+    /// Note that the function accepts both invalid `indices` and `self` not
+    /// being an aggregate as a valid input. It is up to the caller to decide
+    /// what to do in case of `None` being returned.
+    /// E.g.:
+    /// - SROA internally guarantees that `self` is an aggregate and that
+    ///   `indices` are valid and `expect`s a valid offset.
+    /// - Memcpyopt can try to index into a `slice` (not an aggregate) as a result of
+    ///   a valid inspection of a `{ ptr, u64 }` the `slice` got `cast_ptr`ed into,
+    ///   and simply bail out if `None` is return (as a semantically valid result).
     pub fn get_indexed_offset(&self, context: &Context, indices: &[u64]) -> Option<u64> {
         indices
             .iter()
@@ -505,12 +517,7 @@ impl Type {
                                 + (union_size_in_bytes - field_ty.size(context).in_bytes()),
                         )
                     })
-                } else {
-                    assert!(
-                        ty.is_array(context),
-                        "Expected aggregate type. Got {}.",
-                        ty.as_string(context)
-                    );
+                } else if ty.is_array(context) {
                     // size_of_element * idx will be the offset of idx.
                     ty.get_array_elem_type(context).map(|elm_ty| {
                         let prev_idxs_offset = ty
@@ -521,6 +528,8 @@ impl Type {
                             * idx;
                         (elm_ty, accum_offset + prev_idxs_offset)
                     })
+                } else {
+                    None
                 }
             })
             .map(|pair| pair.1)

sway-ir/src/optimize/memcpyopt.rs

@@ -457,6 +457,12 @@ fn local_copy_prop(
 
     struct ReplGep {
         base: Symbol,
+        // When the memcpy's source symbol has a different (but layout-identical)
+        // type than its destination symbol - i.e. the memcpy copies through a
+        // layout-preserving `cast_ptr` - we must re-interpret the base symbol as
+        // the destination's pointee type before indexing into it. This holds the
+        // pointer type to `cast_ptr` the base to or `None` when no cast is needed.
+        cast_base_to: Option<Type>,
         elem_ptr_ty: Type,
         indices: Vec<Value>,
     }
@@ -523,7 +529,44 @@ fn local_copy_prop(
                             .get_type(context)
                             .get_pointee_type(context)
                             .unwrap();
-                        if memcpy_src_sym_type == memcpy_dst_sym_type
+
+                        // The memcpy copies the whole destination symbol, iff the
+                        // source symbol has the same layout as the destination.
+                        // If the two symbol types are identical we can index the
+                        // source directly. If they only share a layout (the memcpy
+                        // went through a layout-preserving `cast_ptr`) we index it
+                        // after re-interpreting it as the destination's type.
+                        let same_type = memcpy_src_sym_type == memcpy_dst_sym_type;
+                        let same_layout = same_type
+                            || memory_utils::types_are_gep_equivalent(
+                                context,
+                                memcpy_src_sym_type,
+                                memcpy_dst_sym_type,
+                            );
+
+                        // In the `cast_ptr` crossing case we re-interpret the source as
+                        // `memcpy_dst_sym_type` and index it with `new_indices`.
+                        //
+                        // Those indices were computed against the access pointer,
+                        // whose type can differ from `memcpy_dst_sym_type` when the
+                        // symbol is reached through a layout-preserving cast (e.g.
+                        // via a block argument.
+                        //
+                        // E.g., let's say the original symbol is a slice that gets
+                        // passed as block argument via `cast_ptr` to `{ ptr, u64 }`.
+                        // In this case, the access sees the block parameter as `{ ptr, u64 }`
+                        // while the symbol passed as the block argument stays a
+                        // non-indexable `slice`.
+                        //
+                        // We want to only proceed if the indices are actually valid for the
+                        // `memcpy_dst_sym_type`.
+                        let indices_valid_for_dst = same_type
+                            || memcpy_dst_sym_type
+                                .get_value_indexed_offset(context, &new_indices)
+                                .is_some();
+
+                        if same_layout
+                            && indices_valid_for_dst
                             && memcpy_dst_sym_type.size(context).in_bytes() == copy_len
                         {
                             replacements.insert(
@@ -532,6 +575,8 @@ fn local_copy_prop(
                                     src_val_ptr,
                                     Replacement::NewGep(ReplGep {
                                         base: memcpy_src_sym,
+                                        cast_base_to: (!same_type)
+                                            .then(|| memcpy_dst_sym.get_type(context)),
                                         elem_ptr_ty: src_val_ptr.get_type(context).unwrap(),
                                         indices: new_indices,
                                     }),
@@ -770,11 +815,12 @@ fn local_copy_prop(
                             to_replace,
                             Replacement::NewGep(ReplGep {
                                 base,
+                                cast_base_to,
                                 elem_ptr_ty,
                                 indices,
                             }),
                         ) => {
-                            let base = match base {
+                            let mut base = match base {
                                 Symbol::Local(local) => {
                                     let base = Value::new_instruction(
                                         context,
@@ -788,6 +834,16 @@ fn local_copy_prop(
                                     block_arg.block.get_arg(context, block_arg.idx).unwrap()
                                 }
                             };
+                            // Re-interpret the base as the destination symbol's type
+                            // when the memcpy went through a layout-preserving cast.
+                            if let Some(cast_to) = cast_base_to {
+                                base = Value::new_instruction(
+                                    context,
+                                    block,
+                                    InstOp::CastPtr(base, cast_to),
+                                );
+                                new_insts.push(base);
+                            }
                             let v = Value::new_instruction(
                                 context,
                                 block,

sway-ir/src/optimize/sroa.rs

@@ -449,7 +449,7 @@ fn profitability(context: &Context, function: Function, candidates: &mut FxHashS
     }
 }
 
-/// Only the following aggregates can be scalarised:
+/// Only a following aggregate can be scalarised:
 /// 1. Does not escape.
 /// 2. Is always accessed via a scalar (register sized) field.
 ///    i.e., The entire aggregate or a sub-aggregate isn't loaded / stored.

sway-ir/tests/memcpyopt_castptr/collapse_through_layout_preserving_cast.ir

@@ -0,0 +1,38 @@
+// A chain of `mem_copy_val`s that alternates between a `{ ptr, u64 }` local and
+// a `slice` local (which have an identical memory layout) via `cast_ptr`.
+//
+// This mirrors the shape produced by usual `raw_slice` handlings.
+//
+// The expected result is:
+// - intermediate copies through the layout-preserving casts must collapse,
+// - the final read of the length field must come straight from `b` (through an inserted, layout-preserving cast),
+// - the redundant copy into `c` must be gone.
+
+script {
+    entry fn main(p: __ptr b256) -> u64 {
+        local { ptr, u64 } a
+        local slice b
+        local { ptr, u64 } c
+
+        entry(p: __ptr b256):
+        va = get_local __ptr { ptr, u64 }, a
+        i0 = const u64 0
+        pa0 = get_elem_ptr va, __ptr ptr, i0
+        pc = cast_ptr p to ptr
+        store pc to pa0
+        i1 = const u64 1
+        pa1 = get_elem_ptr va, __ptr u64, i1
+        c32 = const u64 32
+        store c32 to pa1
+        vcast = cast_ptr va to __ptr slice
+        vb = get_local __ptr slice, b
+        mem_copy_val vb, vcast
+        vcast2 = cast_ptr vb to __ptr { ptr, u64 }
+        vc = get_local __ptr { ptr, u64 }, c
+        mem_copy_val vc, vcast2
+        j1 = const u64 1
+        vlen_p = get_elem_ptr vc, __ptr u64, j1
+        vlen = load vlen_p
+        ret u64 vlen
+    }
+}

sway-ir/tests/memcpyopt_castptr/collapse_through_layout_preserving_cast.ir.snap

@@ -0,0 +1,37 @@
+---
+source: sway-ir/tests/tests.rs
+---
+Modified: true
+
+script {
+    entry fn main(p: __ptr b256) -> u64 {
+        local { ptr, u64 } a
+        local slice b
+-         local { ptr, u64 } c
+
+        entry(mut p: __ptr b256):
+        v2v1 = get_local __ptr { ptr, u64 }, a
+        v3v1 = const u64 0
+        v4v1 = get_elem_ptr v2v1, __ptr ptr, v3v1
+        v5v1 = cast_ptr p to ptr
+        store v5v1 to v4v1
+        v7v1 = const u64 1
+        v8v1 = get_elem_ptr v2v1, __ptr u64, v7v1
+        v9v1 = const u64 32
+        store v9v1 to v8v1
+        v11v1 = cast_ptr v2v1 to __ptr slice
+        v12v1 = get_local __ptr slice, b
+        mem_copy_val v12v1, v11v1
+-         v14v1 = cast_ptr v12v1 to __ptr { ptr, u64 }
+-         v15v1 = get_local __ptr { ptr, u64 }, c
+-         mem_copy_val v15v1, v14v1
++         v21v1 = get_local __ptr slice, b
++         v22v1 = cast_ptr v21v1 to __ptr { ptr, u64 }
+        v17v1 = const u64 1
+-         v18v1 = get_elem_ptr v15v1, __ptr u64, v17v1
+-         v19v1 = load v18v1
++         v23v1 = get_elem_ptr v22v1, __ptr u64, v17v1
++         v19v1 = load v23v1
+        ret u64 v19v1
+    }
+}

sway-ir/tests/memcpyopt_castptr/no_collapse_different_layout_cast.ir

@@ -0,0 +1,30 @@
+// The `cast_ptr`s here reinterpret between `[u64; 4]` and `b256`.
+//
+// The optimization must not fire:
+// - both `mem_copy_val`s are preserved,
+// - the length is still read out of `c`.
+
+script {
+    entry fn main() -> u64 {
+        local [u64; 4] a
+        local b256 b
+        local [u64; 4] c
+
+        entry():
+        va = get_local __ptr [u64; 4], a
+        i0 = const u64 0
+        e0 = get_elem_ptr va, __ptr u64, i0
+        c32 = const u64 32
+        store c32 to e0
+        vcast = cast_ptr va to __ptr b256
+        vb = get_local __ptr b256, b
+        mem_copy_val vb, vcast
+        vcast2 = cast_ptr vb to __ptr [u64; 4]
+        vc = get_local __ptr [u64; 4], c
+        mem_copy_val vc, vcast2
+        j0 = const u64 0
+        ep = get_elem_ptr vc, __ptr u64, j0
+        vlen = load ep
+        ret u64 vlen
+    }
+}

sway-ir/tests/memcpyopt_castptr/no_collapse_different_layout_cast.ir.snap

@@ -0,0 +1,29 @@
+---
+source: sway-ir/tests/tests.rs
+---
+Modified: false
+
+script {
+    entry fn main() -> u64 {
+        local [u64; 4] a
+        local b256 b
+        local [u64; 4] c
+
+        entry():
+        v1v1 = get_local __ptr [u64; 4], a
+        v2v1 = const u64 0
+        v3v1 = get_elem_ptr v1v1, __ptr u64, v2v1
+        v4v1 = const u64 32
+        store v4v1 to v3v1
+        v6v1 = cast_ptr v1v1 to __ptr b256
+        v7v1 = get_local __ptr b256, b
+        mem_copy_val v7v1, v6v1
+        v9v1 = cast_ptr v7v1 to __ptr [u64; 4]
+        v10v1 = get_local __ptr [u64; 4], c
+        mem_copy_val v10v1, v9v1
+        v12v1 = const u64 0
+        v13v1 = get_elem_ptr v10v1, __ptr u64, v12v1
+        v14v1 = load v13v1
+        ret u64 v14v1
+    }
+}

sway-ir/tests/memcpyopt_castptr/no_collapse_tree_mismatch_cast.ir

@@ -0,0 +1,32 @@
+// The `cast_ptr`s here reinterpret between `[u64; 4]` and `[[u64; 2]; 2]`.
+// The two types have the same byte layout four words,
+// but they are not GEP-equivalent.
+//
+// The layout-preserving `cast_ptr` must only be seen through
+// when the pointee types are GEP-equivalent, so the optimization must not fire:
+// - both `mem_copy_val`s are preserved,
+// - the length is still read out of `c`.
+
+script {
+    entry fn main() -> u64 {
+        local [u64; 4] a
+        local [[u64; 2]; 2] b
+        local [u64; 4] c
+
+        entry():
+        va = get_local __ptr [u64; 4], a
+        i3 = const u64 3
+        pa3 = get_elem_ptr va, __ptr u64, i3
+        c40 = const u64 40
+        store c40 to pa3
+        vcast = cast_ptr va to __ptr [[u64; 2]; 2]
+        vb = get_local __ptr [[u64; 2]; 2], b
+        mem_copy_val vb, vcast
+        vcast2 = cast_ptr vb to __ptr [u64; 4]
+        vc = get_local __ptr [u64; 4], c
+        mem_copy_val vc, vcast2
+        ep = get_elem_ptr vc, __ptr u64, i3
+        vlen = load ep
+        ret u64 vlen
+    }
+}

sway-ir/tests/memcpyopt_castptr/no_collapse_tree_mismatch_cast.ir.snap

@@ -0,0 +1,28 @@
+---
+source: sway-ir/tests/tests.rs
+---
+Modified: false
+
+script {
+    entry fn main() -> u64 {
+        local [u64; 4] a
+        local [[u64; 2]; 2] b
+        local [u64; 4] c
+
+        entry():
+        v1v1 = get_local __ptr [u64; 4], a
+        v2v1 = const u64 3
+        v3v1 = get_elem_ptr v1v1, __ptr u64, v2v1
+        v4v1 = const u64 40
+        store v4v1 to v3v1
+        v6v1 = cast_ptr v1v1 to __ptr [[u64; 2]; 2]
+        v7v1 = get_local __ptr [[u64; 2]; 2], b
+        mem_copy_val v7v1, v6v1
+        v9v1 = cast_ptr v7v1 to __ptr [u64; 4]
+        v10v1 = get_local __ptr [u64; 4], c
+        mem_copy_val v10v1, v9v1
+        v12v1 = get_elem_ptr v10v1, __ptr u64, v2v1
+        v13v1 = load v12v1
+        ret u64 v13v1
+    }
+}

sway-ir/tests/memcpyopt_castptr/no_collapse_when_escaped_through_asm.ir

@@ -0,0 +1,34 @@
+// The casts here are layout-preserving (`{ ptr, u64 }` <-> `slice`),
+// but the copied-into pointer `c` also escapes into an `asm` block. Escape
+// analysis must still treat `c` as escaped, so the copy propagation must not
+// fire and the `mem_copy_val` into `c` is preserved.
+
+script {
+    entry fn main() -> u64 {
+        local { ptr, u64 } a
+        local slice b
+        local { ptr, u64 } c
+
+        entry():
+        va = get_local __ptr { ptr, u64 }, a
+        i1 = const u64 1
+        pa1 = get_elem_ptr va, __ptr u64, i1
+        c32 = const u64 32
+        store c32 to pa1
+        vcast = cast_ptr va to __ptr slice
+        vb = get_local __ptr slice, b
+        mem_copy_val vb, vcast
+        vcast2 = cast_ptr vb to __ptr { ptr, u64 }
+        vc = get_local __ptr { ptr, u64 }, c
+        mem_copy_val vc, vcast2
+        i0 = const u64 0
+        vptr = get_elem_ptr vc, __ptr ptr, i0
+        vesc = asm(ptr: vptr) -> u64 {
+            lw result ptr i1
+        }
+        jl = const u64 1
+        vlen_p = get_elem_ptr vc, __ptr u64, jl
+        vlen = load vlen_p
+        ret u64 vlen
+    }
+}