stage2: elemPtr for slices

* Restructure elemPtr a bit
* New AIR instruction: slice_elem_ptr, which returns a pointer to an element of a slice
* Value: adapt elemPtr to work on slices

10 files changed, 157 insertions, 32 deletions

diff --git a/src/Air.zig b/src/Air.zig
index ee93f51096..380e7fbbc2 100644
--- a/src/Air.zig
+++ b/src/Air.zig
@@ -384,6 +384,10 @@ pub const Inst = struct {
         /// Result type is the element type of the slice operand.
         /// Uses the `bin_op` field.
         slice_elem_val,
+        /// Given a slice value and element index, return a pointer to the element value at that index.
+        /// Result type is a pointer to the element type of the slice operand.
+        /// Uses the `ty_pl` field with payload `Bin`.
+        slice_elem_ptr,
         /// Given a pointer value, and element index, return the element value at that index.
         /// Result type is the element type of the pointer operand.
         /// Uses the `bin_op` field.
@@ -685,6 +689,7 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
         .constant,
         .struct_field_ptr,
         .struct_field_val,
+        .slice_elem_ptr,
         .ptr_elem_ptr,
         .cmpxchg_weak,
         .cmpxchg_strong,
diff --git a/src/Liveness.zig b/src/Liveness.zig
index 1b29b46ad8..054c8e38dc 100644
--- a/src/Liveness.zig
+++ b/src/Liveness.zig
@@ -360,7 +360,7 @@ fn analyzeInst(
             const extra = a.air.extraData(Air.StructField, inst_datas[inst].ty_pl.payload).data;
             return trackOperands(a, new_set, inst, main_tomb, .{ extra.struct_operand, .none, .none });
         },
-        .ptr_elem_ptr => {
+        .ptr_elem_ptr, .slice_elem_ptr => {
             const extra = a.air.extraData(Air.Bin, inst_datas[inst].ty_pl.payload).data;
             return trackOperands(a, new_set, inst, main_tomb, .{ extra.lhs, extra.rhs, .none });
         },
diff --git a/src/Sema.zig b/src/Sema.zig
index d0a94d3fd8..5a77f81aa0 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -333,6 +333,24 @@ pub const Block = struct {
         });
     }
 
+    pub fn addSliceElemPtr(
+        block: *Block,
+        slice: Air.Inst.Ref,
+        elem_index: Air.Inst.Ref,
+        elem_ptr_ty: Type,
+    ) !Air.Inst.Ref {
+        return block.addInst(.{
+            .tag = .slice_elem_ptr,
+            .data = .{ .ty_pl = .{
+                .ty = try block.sema.addType(elem_ptr_ty),
+                .payload = try block.sema.addExtra(Air.Bin{
+                    .lhs = slice,
+                    .rhs = elem_index,
+                }),
+            } },
+        });
+    }
+
     pub fn addInst(block: *Block, inst: Air.Inst) error{OutOfMemory}!Air.Inst.Ref {
         return Air.indexToRef(try block.addInstAsIndex(inst));
     }
@@ -11476,18 +11494,39 @@ fn elemPtr(
         else => return sema.fail(block, array_ptr_src, "expected pointer, found '{}'", .{array_ptr_ty}),
     };
     if (!array_ty.isIndexable()) {
-        return sema.fail(block, src, "array access of non-array type '{}'", .{array_ty});
-    }
-    if (array_ty.isSinglePointer() and array_ty.elemType().zigTypeTag() == .Array) {
-        // we have to deref the ptr operand to get the actual array pointer
-        const array_ptr_deref = try sema.analyzeLoad(block, src, array_ptr, array_ptr_src);
-        return sema.elemPtrArray(block, src, array_ptr_deref, elem_index, elem_index_src);
-    }
-    if (array_ty.zigTypeTag() == .Array) {
-        return sema.elemPtrArray(block, src, array_ptr, elem_index, elem_index_src);
+        return sema.fail(block, src, "array access of non-indexable type '{}'", .{array_ty});
     }
 
-    return sema.fail(block, src, "TODO implement more analyze elemptr", .{});
+    switch (array_ty.zigTypeTag()) {
+        .Pointer => {
+            // In all below cases, we have to deref the ptr operand to get the actual array pointer.
+            const array = try sema.analyzeLoad(block, src, array_ptr, array_ptr_src);
+            switch (array_ty.ptrSize()) {
+                .Slice => {
+                    const result_ty = try array_ty.elemPtrType(sema.arena);
+                    const maybe_slice_val = try sema.resolveDefinedValue(block, array_ptr_src, array);
+                    const maybe_index_val = try sema.resolveDefinedValue(block, elem_index_src, elem_index);
+                    const runtime_src = if (maybe_slice_val) |slice_val| rs: {
+                        const index_val = maybe_index_val orelse break :rs elem_index_src;
+                        const index = @intCast(usize, index_val.toUnsignedInt());
+                        const elem_ptr = try slice_val.elemPtr(sema.arena, index);
+                        return sema.addConstant(result_ty, elem_ptr);
+                    } else array_ptr_src;
+
+                    try sema.requireRuntimeBlock(block, runtime_src);
+                    return block.addSliceElemPtr(array, elem_index, result_ty);
+                },
+                .Many, .C => return sema.fail(block, src, "TODO implement Sema for elemPtr for many/c pointer", .{}),
+                .One => {
+                    assert(array_ty.childType().zigTypeTag() == .Array); // Guaranteed by isIndexable
+                    return sema.elemPtrArray(block, src, array, elem_index, elem_index_src);
+                },
+            }
+        },
+        .Array => return sema.elemPtrArray(block, src, array_ptr, elem_index, elem_index_src),
+        .Vector => return sema.fail(block, src, "TODO implement Sema for elemPtr for vector", .{}),
+        else => unreachable,
+    }
 }
 
 fn elemVal(
@@ -11529,7 +11568,7 @@ fn elemVal(
                 return block.addBinOp(.ptr_elem_val, array, elem_index);
             },
             .One => {
-                assert(array_ty.childType().zigTypeTag() == .Array);
+                assert(array_ty.childType().zigTypeTag() == .Array); // Guaranteed by isIndexable
                 const elem_ptr = try sema.elemPtr(block, src, array, elem_index, elem_index_src);
                 return sema.analyzeLoad(block, src, elem_ptr, elem_index_src);
             },
@@ -11562,12 +11601,7 @@ fn elemPtrArray(
     elem_index_src: LazySrcLoc,
 ) CompileError!Air.Inst.Ref {
     const array_ptr_ty = sema.typeOf(array_ptr);
-    const pointee_type = array_ptr_ty.elemType().elemType();
-    const result_ty = try Type.ptr(sema.arena, .{
-        .pointee_type = pointee_type,
-        .mutable = array_ptr_ty.ptrIsMutable(),
-        .@"addrspace" = array_ptr_ty.ptrAddressSpace(),
-    });
+    const result_ty = try array_ptr_ty.elemPtrType(sema.arena);
 
     if (try sema.resolveDefinedValue(block, src, array_ptr)) |array_ptr_val| {
         if (try sema.resolveDefinedValue(block, elem_index_src, elem_index)) |index_val| {
diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index 5d1856b2fa..8dd4526f4d 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -500,6 +500,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
 
                     .array_elem_val      => try self.airArrayElemVal(inst),
                     .slice_elem_val      => try self.airSliceElemVal(inst),
+                    .slice_elem_ptr      => try self.airSliceElemPtr(inst),
                     .ptr_elem_val        => try self.airPtrElemVal(inst),
                     .ptr_elem_ptr        => try self.airPtrElemPtr(inst),
 
@@ -1084,6 +1085,13 @@ fn airSliceElemVal(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ bin_op.lhs, bin_op.rhs, .none });
 }
 
+fn airSliceElemPtr(self: *Self, inst: Air.Inst.Index) !void {
+    const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+    const extra = self.air.extraData(Air.Bin, ty_pl.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement slice_elem_ptr for {}", .{self.target.cpu.arch});
+    return self.finishAir(inst, result, .{ extra.lhs, extra.rhs, .none });
+}
+
 fn airArrayElemVal(self: *Self, inst: Air.Inst.Index) !void {
     const bin_op = self.air.instructions.items(.data)[inst].bin_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement array_elem_val for {}", .{self.target.cpu.arch});
diff --git a/src/codegen.zig b/src/codegen.zig
index 873260329c..5963cdd376 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -848,6 +848,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
 
                     .array_elem_val      => try self.airArrayElemVal(inst),
                     .slice_elem_val      => try self.airSliceElemVal(inst),
+                    .slice_elem_ptr      => try self.airSliceElemPtr(inst),
                     .ptr_elem_val        => try self.airPtrElemVal(inst),
                     .ptr_elem_ptr        => try self.airPtrElemPtr(inst),
 
@@ -1533,6 +1534,15 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             return self.finishAir(inst, result, .{ bin_op.lhs, bin_op.rhs, .none });
         }
 
+        fn airSliceElemPtr(self: *Self, inst: Air.Inst.Index) !void {
+            const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+            const extra = self.air.extraData(Air.Bin, ty_pl.payload).data;
+            const result: MCValue = if (self.liveness.isUnused(inst)) .dead else switch (arch) {
+                else => return self.fail("TODO implement slice_elem_ptr for {}", .{self.target.cpu.arch}),
+            };
+            return self.finishAir(inst, result, .{ extra.lhs, extra.rhs, .none });
+        }
+
         fn airArrayElemVal(self: *Self, inst: Air.Inst.Index) !void {
             const bin_op = self.air.instructions.items(.data)[inst].bin_op;
             const result: MCValue = if (self.liveness.isUnused(inst)) .dead else switch (arch) {
diff --git a/src/codegen/c.zig b/src/codegen/c.zig
index f09bb574af..9050499fce 100644
--- a/src/codegen/c.zig
+++ b/src/codegen/c.zig
@@ -1083,6 +1083,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
             .ptr_elem_val       => try airPtrElemVal(f, inst, "["),
             .ptr_elem_ptr       => try airPtrElemPtr(f, inst),
             .slice_elem_val     => try airSliceElemVal(f, inst, "["),
+            .slice_elem_ptr     => try airSliceElemPtr(f, inst),
             .array_elem_val     => try airArrayElemVal(f, inst),
 
             .unwrap_errunion_payload     => try airUnwrapErrUnionPay(f, inst),
@@ -1165,6 +1166,24 @@ fn airSliceElemVal(f: *Function, inst: Air.Inst.Index, prefix: []const u8) !CVal
     return local;
 }
 
+fn airSliceElemPtr(f: *Function, inst: Air.Inst.Index) !CValue {
+    if (f.liveness.isUnused(inst))
+        return CValue.none;
+    const ty_pl = f.air.instructions.items(.data)[inst].ty_pl;
+    const bin_op = f.air.extraData(Air.Bin, ty_pl.payload).data;
+
+    const slice = try f.resolveInst(bin_op.lhs);
+    const index = try f.resolveInst(bin_op.rhs);
+    const writer = f.object.writer();
+    const local = try f.allocLocal(f.air.typeOfIndex(inst), .Const);
+    try writer.writeAll(" = &");
+    try f.writeCValue(writer, slice);
+    try writer.writeByte('[');
+    try f.writeCValue(writer, index);
+    try writer.writeAll("];\n");
+    return local;
+}
+
 fn airArrayElemVal(f: *Function, inst: Air.Inst.Index) !CValue {
     if (f.liveness.isUnused(inst)) return CValue.none;
 
diff --git a/src/codegen/llvm.zig b/src/codegen/llvm.zig
index e5356582e2..a607b7a753 100644
--- a/src/codegen/llvm.zig
+++ b/src/codegen/llvm.zig
@@ -1760,6 +1760,7 @@ pub const FuncGen = struct {
 
                 .array_elem_val     => try self.airArrayElemVal(inst),
                 .slice_elem_val     => try self.airSliceElemVal(inst),
+                .slice_elem_ptr     => try self.airSliceElemPtr(inst),
                 .ptr_elem_val       => try self.airPtrElemVal(inst),
                 .ptr_elem_ptr       => try self.airPtrElemPtr(inst),
 
@@ -2157,12 +2158,20 @@ pub const FuncGen = struct {
 
         const slice = try self.resolveInst(bin_op.lhs);
         const index = try self.resolveInst(bin_op.rhs);
-        const base_ptr = self.builder.buildExtractValue(slice, 0, "");
-        const indices: [1]*const llvm.Value = .{index};
-        const ptr = self.builder.buildInBoundsGEP(base_ptr, &indices, indices.len, "");
+        const ptr = self.sliceElemPtr(slice, index);
         return self.load(ptr, slice_ty);
     }
 
+    fn airSliceElemPtr(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
+        if (self.liveness.isUnused(inst)) return null;
+        const ty_pl = self.air.instructions.items(.data)[inst].ty_pl;
+        const bin_op = self.air.extraData(Air.Bin, ty_pl.payload).data;
+
+        const slice = try self.resolveInst(bin_op.lhs);
+        const index = try self.resolveInst(bin_op.rhs);
+        return self.sliceElemPtr(slice, index);
+    }
+
     fn airArrayElemVal(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
         if (self.liveness.isUnused(inst)) return null;
 
@@ -3575,6 +3584,16 @@ pub const FuncGen = struct {
         return self.builder.buildBitCast(union_field_ptr, result_llvm_ty, "");
     }
 
+    fn sliceElemPtr(
+        self: *FuncGen,
+        slice: *const llvm.Value,
+        index: *const llvm.Value,
+    ) *const llvm.Value {
+        const base_ptr = self.builder.buildExtractValue(slice, 0, "");
+        const indices: [1]*const llvm.Value = .{index};
+        return self.builder.buildInBoundsGEP(base_ptr, &indices, indices.len, "");
+    }
+
     fn getIntrinsic(self: *FuncGen, name: []const u8) *const llvm.Value {
         const id = llvm.lookupIntrinsicID(name.ptr, name.len);
         assert(id != 0);
diff --git a/src/print_air.zig b/src/print_air.zig
index 5231bf0055..a6d99eee9f 100644
--- a/src/print_air.zig
+++ b/src/print_air.zig
@@ -200,6 +200,7 @@ const Writer = struct {
             .loop,
             => try w.writeBlock(s, inst),
 
+            .slice_elem_ptr => try w.writeSliceElemPtr(s, inst),
             .ptr_elem_ptr => try w.writePtrElemPtr(s, inst),
             .struct_field_ptr => try w.writeStructField(s, inst),
             .struct_field_val => try w.writeStructField(s, inst),
@@ -281,6 +282,15 @@ const Writer = struct {
         try s.print(", {d}", .{extra.field_index});
     }
 
+    fn writeSliceElemPtr(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
+        const ty_pl = w.air.instructions.items(.data)[inst].ty_pl;
+        const extra = w.air.extraData(Air.Bin, ty_pl.payload).data;
+
+        try w.writeOperand(s, inst, 0, extra.lhs);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 1, extra.rhs);
+    }
+
     fn writePtrElemPtr(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
         const ty_pl = w.air.instructions.items(.data)[inst].ty_pl;
         const extra = w.air.extraData(Air.Bin, ty_pl.payload).data;
diff --git a/src/type.zig b/src/type.zig
index fceb3c4dfd..72347f5adc 100644
--- a/src/type.zig
+++ b/src/type.zig
@@ -2520,6 +2520,20 @@ pub const Type = extern union {
         };
     }
 
+    /// Returns the type of a pointer to an element.
+    /// Asserts that the type is a pointer, and that the element type is indexable.
+    /// For *[N]T, return *T
+    /// For [*]T, returns *T
+    /// For []T, returns *T
+    /// Handles const-ness and address spaces in particular.
+    pub fn elemPtrType(ptr_ty: Type, arena: *Allocator) !Type {
+        return try Type.ptr(arena, .{
+            .pointee_type = ptr_ty.elemType2(),
+            .mutable = ptr_ty.ptrIsMutable(),
+            .@"addrspace" = ptr_ty.ptrAddressSpace(),
+        });
+    }
+
     fn shallowElemType(child_ty: Type) Type {
         return switch (child_ty.zigTypeTag()) {
             .Array, .Vector => child_ty.childType(),
diff --git a/src/value.zig b/src/value.zig
index f1ca384b75..bb904cea30 100644
--- a/src/value.zig
+++ b/src/value.zig
@@ -114,7 +114,7 @@ pub const Value = extern union {
         /// This Tag will never be seen by machine codegen backends. It is changed into a
         /// `decl_ref` when a comptime variable goes out of scope.
         decl_ref_mut,
-        /// Pointer to a specific element of an array.
+        /// Pointer to a specific element of an array, vector or slice.
         elem_ptr,
         /// Pointer to a specific field of a struct or union.
         field_ptr,
@@ -1792,17 +1792,23 @@ pub const Value = extern union {
 
     /// Returns a pointer to the element value at the index.
     pub fn elemPtr(self: Value, allocator: *Allocator, index: usize) !Value {
-        if (self.castTag(.elem_ptr)) |elem_ptr| {
-            return Tag.elem_ptr.create(allocator, .{
-                .array_ptr = elem_ptr.data.array_ptr,
-                .index = elem_ptr.data.index + index,
-            });
+        switch (self.tag()) {
+            .elem_ptr => {
+                const elem_ptr = self.castTag(.elem_ptr).?.data;
+                return Tag.elem_ptr.create(allocator, .{
+                    .array_ptr = elem_ptr.array_ptr,
+                    .index = elem_ptr.index + index,
+                });
+            },
+            .slice => return Tag.elem_ptr.create(allocator, .{
+                .array_ptr = self.castTag(.slice).?.data.ptr,
+                .index = index,
+            }),
+            else => return Tag.elem_ptr.create(allocator, .{
+                .array_ptr = self,
+                .index = index,
+            }),
         }
-
-        return Tag.elem_ptr.create(allocator, .{
-            .array_ptr = self,
-            .index = index,
-        });
     }
 
     pub fn isUndef(self: Value) bool {