Remove elements from std.ArrayList while iterating over it

std/array_list.zig

@@ -1,4 +1,5 @@
 const std = @import("std.zig");
+const builtin = @import("builtin");
 const debug = std.debug;
 const assert = debug.assert;
 const testing = std.testing;
@@ -201,27 +202,78 @@ pub fn AlignedArrayList(comptime T: type, comptime alignment: ?u29) type {
             return self.pop();
         }
 
+        // an iterator which allows you to remove the current element
+        // while iterating through it.
         pub const Iterator = struct {
-            list: *const Self,
-            // how many items have we returned
-            count: usize,
+            list: *Self,
+
+            // index of the last element we returned, or null
+            // if we haven't called `next` yet.
+            cursor: ?usize = null,
+
+            // the last index we removed
+            removed_index: ?usize = null,
 
             pub fn next(it: *Iterator) ?T {
-                if (it.count >= it.list.len) return null;
-                const val = it.list.at(it.count);
-                it.count += 1;
-                return val;
+                var idx: usize = 0;
+                if (it.cursor) |cursor| {
+                    idx = cursor + 1;
+                }
+                if (idx >= it.list.len) return null;
+                it.cursor = idx;
+                return it.list.at(idx);
             }
 
             pub fn reset(it: *Iterator) void {
-                it.count = 0;
+                it.cursor = null;
+                it.removed_index = null;
+            }
+
+            pub fn swapRemove(it: *Iterator) T {
+                var cursor = it.cursor.?; // must call .next() at least once
+                if (it.removed_index) |ri| assert(ri != cursor); // removed the same element more than once
+                it.removed_index = cursor;
+                return it.list.swapRemove(cursor);
+            }
+
+            pub fn orderedRemove(it: *Iterator) T {
+                var cursor = it.cursor.?; // must call .next() at least once
+                if (it.removed_index) |ri| assert(ri != cursor); // removed the same element more than once
+                it.removed_index = cursor;
+                return it.list.orderedRemove(cursor);
             }
         };
 
-        pub fn iterator(self: *const Self) Iterator {
+        pub fn iterator(self: *Self) Iterator {
             return Iterator{
                 .list = self,
-                .count = 0,
+            };
+        }
+
+        // an iterator that only allows you to iterate
+        // through the elements; you cannot remove the current
+        // element.
+        pub const IteratorConst = struct {
+            list: *const Self,
+
+            // index of the next element to return
+            next_index: usize = 0,
+
+            pub fn next(it: *IteratorConst) ?T {
+                if (it.next_index >= it.list.len) return null;
+                const val = it.list.at(it.next_index);
+                it.next_index += 1;
+                return val;
+            }
+
+            pub fn reset(it: *IteratorConst) void {
+                it.next_index = 0;
+            }
+        };
+
+        pub fn iteratorConst(self: *const Self) IteratorConst {
+            return IteratorConst{
+                .list = self,
             };
         }
     };
@@ -409,6 +461,115 @@ test "std.ArrayList.iterator" {
     testing.expect(it.next().? == 1);
 }
 
+test "std.ArrayList.iterator.swapRemove" {
+    var bytes: [1024]u8 = undefined;
+    const allocator = &std.heap.FixedBufferAllocator.init(bytes[0..]).allocator;
+
+    var list = ArrayList(i32).init(allocator);
+    defer list.deinit();
+
+    try list.append(1);
+    try list.append(2);
+    try list.append(3);
+    try list.append(4);
+
+    {
+        var last_elem = list.at(list.len-1);
+        var itr = list.iterator();
+        _ = itr.next();
+        _ = itr.swapRemove();
+        testing.expectEqual(list.at(0), last_elem);
+        try list.insert(0, 1); // put the number back in.
+    }
+
+    var it = list.iterator();
+    while (it.next()) |next| {
+        if (next == 2) {
+            var removed = it.swapRemove();
+            testing.expect(removed == 2);
+            break;
+        }
+    }
+
+    it.reset();
+    testing.expect(it.next().? == 1);
+    testing.expect(it.next().? == 4);
+    testing.expect(it.next().? == 3);
+    testing.expect(it.next() == null);
+}
+
+test "std.ArrayList.iterator.orderedRemove" {
+    var bytes: [1024]u8 = undefined;
+    const allocator = &std.heap.FixedBufferAllocator.init(bytes[0..]).allocator;
+
+    var list = ArrayList(i32).init(allocator);
+    defer list.deinit();
+
+    try list.append(1);
+    try list.append(2);
+    try list.append(3);
+    try list.append(4);
+
+    // check we can remove the first element
+    {
+        var second_elem = list.at(1);
+        var itr = list.iterator();
+        _ = itr.next();
+        _ = itr.orderedRemove();
+        testing.expectEqual(list.at(0), second_elem);
+        try list.insert(0, 1); // put the number back in.
+    }
+
+    var it = list.iterator();
+    while (it.next()) |next| {
+        if (next == 2) {
+            var removed = it.orderedRemove();
+            testing.expect(removed == 2);
+            break;
+        }
+    }
+
+    it.reset();
+    testing.expect(it.next().? == 1);
+    testing.expect(it.next().? == 3);
+    testing.expect(it.next().? == 4);
+    testing.expect(it.next() == null);
+}
+
+test "std.ArrayList.iteratorConst" {
+    var bytes: [1024]u8 = undefined;
+    const allocator = &std.heap.FixedBufferAllocator.init(bytes[0..]).allocator;
+
+    var list = ArrayList(i32).init(allocator);
+    defer list.deinit();
+
+    try list.append(1);
+    try list.append(2);
+    try list.append(3);
+
+    const clist = list;
+
+    var count: i32 = 0;
+    var it = clist.iteratorConst();
+    while (it.next()) |next| {
+        testing.expect(next == count + 1);
+        count += 1;
+    }
+
+    testing.expect(count == 3);
+    testing.expect(it.next() == null);
+    it.reset();
+    count = 0;
+    while (it.next()) |next| {
+        testing.expect(next == count + 1);
+        count += 1;
+        if (count == 2) break;
+    }
+
+    it.reset();
+    testing.expect(it.next().? == 1);
+}
+
 test "std.ArrayList.insert" {
     var list = ArrayList(i32).init(debug.global_allocator);
     defer list.deinit();

std/http/headers.zig

@@ -133,7 +133,7 @@ pub const Headers = struct {
             self.index.deinit();
         }
         {
-            var it = self.data.iterator();
+            var it = self.data.iteratorConst();
             while (it.next()) |entry| {
                 entry.deinit();
             }
@@ -146,7 +146,7 @@ pub const Headers = struct {
         errdefer other.deinit();
         try other.data.ensureCapacity(self.data.count());
         try other.index.initCapacity(self.index.entries.len);
-        var it = self.data.iterator();
+        var it = self.data.iteratorConst();
         while (it.next()) |entry| {
             try other.append(entry.name, entry.value, entry.never_index);
         }
@@ -157,10 +157,10 @@ pub const Headers = struct {
         return self.data.count();
     }
 
-    pub const Iterator = HeaderList.Iterator;
+    pub const Iterator = HeaderList.IteratorConst;
 
     pub fn iterator(self: Self) Iterator {
-        return self.data.iterator();
+        return self.data.iteratorConst();
     }
 
     pub fn append(self: *Self, name: []const u8, value: []const u8, never_index: ?bool) !void {
@@ -290,7 +290,7 @@ pub const Headers = struct {
         const dex = self.getIndices(name) orelse return null;
 
         const buf = try allocator.alloc(HeaderEntry, dex.count());
-        var it = dex.iterator();
+        var it = dex.iteratorConst();
         var n: usize = 0;
         while (it.next()) |idx| {
             buf[n] = self.data.at(idx);
@@ -315,7 +315,7 @@ pub const Headers = struct {
         // adapted from mem.join
         const total_len = blk: {
             var sum: usize = dex.count() - 1; // space for separator(s)
-            var it = dex.iterator();
+            var it = dex.iteratorConst();
             while (it.next()) |idx|
                 sum += self.data.at(idx).value.len;
             break :blk sum;
@@ -351,7 +351,7 @@ pub const Headers = struct {
             var it = self.data.iterator();
             while (it.next()) |entry| {
                 var dex = &self.index.get(entry.name).?.value;
-                dex.appendAssumeCapacity(it.count);
+                dex.appendAssumeCapacity(it.cursor.?+1);
             }
         }
     }