std: mem span functions can take an optional pointer

Author: daurnimator

lib/std/mem.zig

@@ -512,36 +512,54 @@ pub fn toSlice(comptime T: type, ptr: [*:0]T) [:0]T {
 /// the constness of the input type. `[*c]` pointers are assumed to be 0-terminated,
 /// and assumed to not allow null.
 pub fn Span(comptime T: type) type {
-    var ptr_info = @typeInfo(T).Pointer;
-    switch (ptr_info.size) {
-        .One => switch (@typeInfo(ptr_info.child)) {
-            .Array => |info| {
-                ptr_info.child = info.child;
-                ptr_info.sentinel = info.sentinel;
-            },
-            else => @compileError("invalid type given to std.mem.Span"),
+    switch(@typeInfo(T)) {
+        .Optional => |optional_info| {
+            return ?Span(optional_info.child);
         },
-        .C => {
-            ptr_info.sentinel = 0;
-            ptr_info.is_allowzero = false;
+        .Pointer => |ptr_info| {
+            var new_ptr_info = ptr_info;
+            switch (ptr_info.size) {
+                .One => switch (@typeInfo(ptr_info.child)) {
+                    .Array => |info| {
+                        new_ptr_info.child = info.child;
+                        new_ptr_info.sentinel = info.sentinel;
+                    },
+                    else => @compileError("invalid type given to std.mem.Span"),
+                },
+                .C => {
+                    new_ptr_info.sentinel = 0;
+                    new_ptr_info.is_allowzero = false;
+                },
+                .Many, .Slice => {},
+            }
+            new_ptr_info.size = .Slice;
+            return @Type(std.builtin.TypeInfo{ .Pointer = new_ptr_info });
         },
-        .Many, .Slice => {},
+        else => @compileError("invalid type given to std.mem.Span"),
     }
-    ptr_info.size = .Slice;
-    return @Type(std.builtin.TypeInfo{ .Pointer = ptr_info });
 }
 
 test "Span" {
     testing.expect(Span(*[5]u16) == []u16);
+    testing.expect(Span(?*[5]u16) == ?[]u16);
     testing.expect(Span(*const [5]u16) == []const u16);
+    testing.expect(Span(?*const [5]u16) == ?[]const u16);
     testing.expect(Span([]u16) == []u16);
+    testing.expect(Span(?[]u16) == ?[]u16);
     testing.expect(Span([]const u8) == []const u8);
+    testing.expect(Span(?[]const u8) == ?[]const u8);
     testing.expect(Span([:1]u16) == [:1]u16);
+    testing.expect(Span(?[:1]u16) == ?[:1]u16);
     testing.expect(Span([:1]const u8) == [:1]const u8);
+    testing.expect(Span(?[:1]const u8) == ?[:1]const u8);
     testing.expect(Span([*:1]u16) == [:1]u16);
+    testing.expect(Span(?[*:1]u16) == ?[:1]u16);
     testing.expect(Span([*:1]const u8) == [:1]const u8);
+    testing.expect(Span(?[*:1]const u8) == ?[:1]const u8);
     testing.expect(Span([*c]u16) == [:0]u16);
+    testing.expect(Span(?[*c]u16) == ?[:0]u16);
     testing.expect(Span([*c]const u8) == [:0]const u8);
+    testing.expect(Span(?[*c]const u8) == ?[:0]const u8);
 }
 
 /// Takes a pointer to an array, a sentinel-terminated pointer, or a slice, and
@@ -553,6 +571,9 @@ test "Span" {
 /// length value is used instead of the sentinel.
 pub fn span(ptr: var) Span(@TypeOf(ptr)) {
     const Result = Span(@TypeOf(ptr));
+    if (@typeInfo(@TypeOf(ptr)) == .Optional and ptr == null) {
+        return null;
+    }
     const l = len(ptr);
     if (@typeInfo(Result).Pointer.sentinel) |s| {
         return ptr[0..l :s];
@@ -573,6 +594,9 @@ test "span" {
 /// rather than using the length.
 pub fn spanZ(ptr: var) Span(@TypeOf(ptr)) {
     const Result = Span(@TypeOf(ptr));
+    if (@typeInfo(@TypeOf(ptr)) == .Optional and ptr == null) {
+        return null;
+    }
     const l = lenZ(ptr);
     if (@typeInfo(Result).Pointer.sentinel) |s| {
         return ptr[0..l :s];