Handle compile time case for vector element access using lane access

2 files changed, 47 insertions, 9 deletions

diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index c5376792a7..70cba5eb9f 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -4527,17 +4527,38 @@ fn airArrayElemVal(func: *CodeGen, inst: Air.Inst.Index) InnerError!void {
     } else {
         std.debug.assert(array_ty.zigTypeTag() == .Vector);
 
-        // TODO: Check if index is constant; if so, use a lane extract
+        switch (index) {
+            inline .imm32, .imm64 => |lane| {
+                const opcode: wasm.SimdOpcode = switch (elem_ty.bitSize(func.target)) {
+                    8 => if (elem_ty.isSignedInt()) .i8x16_extract_lane_s else .i8x16_extract_lane_u,
+                    16 => if (elem_ty.isSignedInt()) .i16x8_extract_lane_s else .i16x8_extract_lane_u,
+                    32 => if (elem_ty.isInt()) .i32x4_extract_lane else .f32x4_extract_lane,
+                    64 => if (elem_ty.isInt()) .i64x2_extract_lane else .f64x2_extract_lane,
+                    else => unreachable,
+                };
 
-        var stack_vec = try func.allocStack(array_ty);
-        try func.store(stack_vec, array, array_ty, 0);
+                var operands = [_]u32{ std.wasm.simdOpcode(opcode), @intCast(u8, lane) };
 
-        // Is a non-unrolled vector (v128)
-        try func.lowerToStack(stack_vec);
-        try func.emitWValue(index);
-        try func.addImm32(@bitCast(i32, @intCast(u32, elem_size)));
-        try func.addTag(.i32_mul);
-        try func.addTag(.i32_add);
+                try func.emitWValue(array);
+
+                const extra_index = @intCast(u32, func.mir_extra.items.len);
+                try func.mir_extra.appendSlice(func.gpa, &operands);
+                try func.addInst(.{ .tag = .simd_prefix, .data = .{ .payload = extra_index } });
+
+                return func.finishAir(inst, try WValue.toLocal(.stack, func, elem_ty), &.{ bin_op.lhs, bin_op.rhs });
+            },
+            else => {
+                var stack_vec = try func.allocStack(array_ty);
+                try func.store(stack_vec, array, array_ty, 0);
+
+                // Is a non-unrolled vector (v128)
+                try func.lowerToStack(stack_vec);
+                try func.emitWValue(index);
+                try func.addImm32(@bitCast(i32, @intCast(u32, elem_size)));
+                try func.addTag(.i32_mul);
+                try func.addTag(.i32_add);
+            },
+        }
     }
 
     const elem_result = val: {
diff --git a/src/arch/wasm/Emit.zig b/src/arch/wasm/Emit.zig
index 3d05ff7987..1d039f7495 100644
--- a/src/arch/wasm/Emit.zig
+++ b/src/arch/wasm/Emit.zig
@@ -492,6 +492,23 @@ fn emitSimd(emit: *Emit, inst: Mir.Inst.Index) !void {
             const simd_value = emit.mir.extra[extra_index + 1 ..][0..4];
             try writer.writeAll(std.mem.asBytes(simd_value));
         },
+        .i8x16_extract_lane_s,
+        .i8x16_extract_lane_u,
+        .i8x16_replace_lane,
+        .i16x8_extract_lane_s,
+        .i16x8_extract_lane_u,
+        .i16x8_replace_lane,
+        .i32x4_extract_lane,
+        .i32x4_replace_lane,
+        .i64x2_extract_lane,
+        .i64x2_replace_lane,
+        .f32x4_extract_lane,
+        .f32x4_replace_lane,
+        .f64x2_extract_lane,
+        .f64x2_replace_lane,
+        => {
+            try writer.writeByte(@intCast(u8, emit.mir.extra[extra_index + 1]));
+        },
         .i8x16_splat,
         .i16x8_splat,
         .i32x4_splat,