diff --git a/lib/std/mem.zig b/lib/std/mem.zig index 88d0e877b8..eec9d4d61f 100644 --- a/lib/std/mem.zig +++ b/lib/std/mem.zig @@ -1941,13 +1941,29 @@ test "joinZ" { /// Copies each T from slices into a new slice that exactly holds all the elements. pub fn concat(allocator: Allocator, comptime T: type, slices: []const []const T) ![]T { - if (slices.len == 0) return &[0]T{}; + return concatMaybeSentinel(allocator, T, slices, null); +} + +/// Copies each T from slices into a new slice that exactly holds all the elements. +pub fn concatWithSentinel(allocator: Allocator, comptime T: type, slices: []const []const T, comptime s: T) ![:s]T { + const ret = try concatMaybeSentinel(allocator, T, slices, s); + return ret[0 .. ret.len - 1 :s]; +} + +/// Copies each T from slices into a new slice that exactly holds all the elements as well as the sentinel. +pub fn concatMaybeSentinel(allocator: Allocator, comptime T: type, slices: []const []const T, comptime s: ?T) ![]T { + if (slices.len == 0) return if (s) |sentinel| try allocator.dupe(T, &[1]T{sentinel}) else &[0]T{}; const total_len = blk: { var sum: usize = 0; for (slices) |slice| { sum += slice.len; } + + if (s) |_| { + sum += 1; + } + break :blk sum; }; @@ -1960,6 +1976,10 @@ pub fn concat(allocator: Allocator, comptime T: type, slices: []const []const T) buf_index += slice.len; } + if (s) |sentinel| { + buf[buf.len - 1] = sentinel; + } + // No need for shrink since buf is exactly the correct size. return buf; } @@ -1980,6 +2000,26 @@ test "concat" { defer testing.allocator.free(str); try testing.expect(eql(u32, str, &[_]u32{ 0, 1, 2, 3, 4, 5 })); } + { + const str = try concatWithSentinel(testing.allocator, u8, &[_][]const u8{ "abc", "def", "ghi" }, 0); + defer testing.allocator.free(str); + try testing.expectEqualSentinel(u8, 0, str, "abcdefghi"); + } + { + const slice = try concatWithSentinel(testing.allocator, u8, &[_][]const u8{}, 0); + defer testing.allocator.free(slice); + try testing.expectEqualSentinel(u8, 0, slice, &[_:0]u8{}); + } + { + const slice = try concatWithSentinel(testing.allocator, u32, &[_][]const u32{ + &[_]u32{ 0, 1 }, + &[_]u32{ 2, 3, 4 }, + &[_]u32{}, + &[_]u32{5}, + }, 2); + defer testing.allocator.free(slice); + try testing.expectEqualSentinel(u32, 2, slice, &[_:2]u32{ 0, 1, 2, 3, 4, 5 }); + } } test "testStringEquality" {