Node

pub const Node = struct {
    index: OptionalIndex,

    pub const none: Node = .{ .index = .none };

    pub const max_name_len = 120;

    const Storage = extern struct {
        /// Little endian.
        completed_count: u32,
        /// 0 means unknown.
        /// Little endian.
        estimated_total_count: u32,
        name: [max_name_len]u8 align(@alignOf(usize)),

        /// Not thread-safe.
        fn getIpcIndex(s: Storage) ?Ipc.Index {
            return if (s.estimated_total_count == std.math.maxInt(u32)) @bitCast(s.completed_count) else null;
        }

        /// Thread-safe.
        fn setIpcIndex(s: *Storage, ipc_index: Ipc.Index) void {
            // `estimated_total_count` max int indicates the special state that
            // causes `completed_count` to be treated as a file descriptor, so
            // the order here matters.
            @atomicStore(u32, &s.completed_count, @bitCast(ipc_index), .monotonic);
            @atomicStore(u32, &s.estimated_total_count, std.math.maxInt(u32), .release); // synchronizes with acquire in `serialize`
        }

        /// Not thread-safe.
        fn byteSwap(s: *Storage) void {
            s.completed_count = @byteSwap(s.completed_count);
            s.estimated_total_count = @byteSwap(s.estimated_total_count);
        }

        fn copyRoot(dest: *Node.Storage, src: *align(1) const Node.Storage) void {
            dest.* = .{
                .completed_count = src.completed_count,
                .estimated_total_count = src.estimated_total_count,
                .name = if (src.name[0] == 0) dest.name else src.name,
            };
        }

        comptime {
            assert((@sizeOf(Storage) % 4) == 0);
        }
    };

    const Parent = enum(u8) {
        /// Unallocated storage.
        unused = std.math.maxInt(u8) - 1,
        /// Indicates root node.
        none = std.math.maxInt(u8),
        /// Index into `node_storage`.
        _,

        fn unwrap(i: @This()) ?Index {
            return switch (i) {
                .unused, .none => return null,
                else => @fromBackingInt(@intCast(@backingInt(i))),
            };
        }
    };

    pub const OptionalIndex = enum(u8) {
        none = std.math.maxInt(u8),
        /// Index into `node_storage`.
        _,

        pub fn unwrap(i: @This()) ?Index {
            if (i == .none) return null;
            return @fromBackingInt(@intCast(@backingInt(i)));
        }

        fn toParent(i: @This()) Parent {
            assert(@backingInt(i) != @backingInt(Parent.unused));
            return @fromBackingInt(@intCast(@backingInt(i)));
        }
    };

    /// Index into `node_storage`.
    pub const Index = enum(u8) {
        _,

        fn toParent(i: @This()) Parent {
            assert(@backingInt(i) != @backingInt(Parent.unused));
            assert(@backingInt(i) != @backingInt(Parent.none));
            return @fromBackingInt(@intCast(@backingInt(i)));
        }

        pub fn toOptional(i: @This()) OptionalIndex {
            return @fromBackingInt(@intCast(@backingInt(i)));
        }
    };

    /// Create a new child progress node. Thread-safe.
    ///
    /// Passing 0 for `estimated_total_items` means unknown.
    pub fn start(node: Node, name: []const u8, estimated_total_items: usize) Node {
        if (noop_impl) {
            assert(node.index == .none);
            return Node.none;
        }
        const node_index = node.index.unwrap() orelse return Node.none;
        const parent = node_index.toParent();

        const freelist = &global_progress.node_freelist;
        var old_freelist = @atomicLoad(Freelist, freelist, .acquire); // acquire to ensure we have the correct "next" entry
        while (old_freelist.head.unwrap()) |free_index| {
            const next_ptr = freelistNextByIndex(free_index);
            const new_freelist: Freelist = .{
                .head = @atomicLoad(Node.OptionalIndex, next_ptr, .monotonic),
                // We don't need to increment the generation when removing nodes from the free list,
                // only when adding them. (This choice is arbitrary; the opposite would also work.)
                .generation = old_freelist.generation,
            };
            old_freelist = @cmpxchgWeak(
                Freelist,
                freelist,
                old_freelist,
                new_freelist,
                .acquire, // not theoretically necessary, but not allowed to be weaker than the failure order
                .acquire, // ensure we have the correct `node_freelist_next` entry on the next iteration
            ) orelse {
                // We won the allocation race.
                return init(free_index, parent, name, estimated_total_items);
            };
        }

        const free_index = @atomicRmw(u32, &global_progress.node_end_index, .Add, 1, .monotonic);
        if (free_index >= node_storage_buffer_len) {
            // Ran out of node storage memory. Progress for this node will not be tracked.
            _ = @atomicRmw(u32, &global_progress.node_end_index, .Sub, 1, .monotonic);
            return Node.none;
        }

        return init(@fromBackingInt(@intCast(free_index)), parent, name, estimated_total_items);
    }

    pub fn startFmt(node: Node, estimated_total_items: usize, comptime format: []const u8, args: anytype) Node {
        var buffer: [max_name_len]u8 = undefined;
        const name = std.mem.print(&buffer, format, args) catch &buffer;
        return Node.start(node, name, estimated_total_items);
    }

    /// This is the same as calling `start` and then `end` on the returned `Node`. Thread-safe.
    pub fn completeOne(n: Node) void {
        const index = n.index.unwrap() orelse return;
        const storage = storageByIndex(index);
        _ = @atomicRmw(u32, &storage.completed_count, .Add, 1, .monotonic);
    }

    /// Thread-safe. Bytes after '0' in `new_name` are ignored.
    pub fn setName(n: Node, new_name: []const u8) void {
        const index = n.index.unwrap() orelse return;
        const storage = storageByIndex(index);

        const name_len = @min(max_name_len, std.mem.findScalar(u8, new_name, 0) orelse new_name.len);

        copyAtomicStore(storage.name[0..name_len], new_name[0..name_len]);
        if (name_len < storage.name.len)
            @atomicStore(u8, &storage.name[name_len], 0, .monotonic);
    }

    /// Gets the name of this `Node`.
    /// A pointer to this array can later be passed to `setName` to restore the name.
    pub fn getName(n: Node) [max_name_len]u8 {
        var dest: [max_name_len]u8 align(@alignOf(usize)) = undefined;
        if (n.index.unwrap()) |index| {
            copyAtomicLoad(&dest, &storageByIndex(index).name);
        }
        return dest;
    }

    /// Thread-safe.
    pub fn setCompletedItems(n: Node, completed_items: usize) void {
        const index = n.index.unwrap() orelse return;
        const storage = storageByIndex(index);
        @atomicStore(u32, &storage.completed_count, std.math.lossyCast(u32, completed_items), .monotonic);
    }

    /// Thread-safe. 0 means unknown.
    pub fn setEstimatedTotalItems(n: Node, count: usize) void {
        const index = n.index.unwrap() orelse return;
        const storage = storageByIndex(index);
        // Avoid u32 max int which is used to indicate a special state.
        const saturated_total_count = @min(std.math.maxInt(u32) - 1, count);
        @atomicStore(u32, &storage.estimated_total_count, saturated_total_count, .monotonic);
    }

    /// Thread-safe.
    pub fn increaseEstimatedTotalItems(n: Node, count: usize) void {
        const index = n.index.unwrap() orelse return;
        const storage = storageByIndex(index);
        // Avoid u32 max int which is used to indicate a special state.
        const saturated_total_count = @min(std.math.maxInt(u32) - 1, count);
        _ = @atomicRmw(u32, &storage.estimated_total_count, .Add, saturated_total_count, .monotonic);
    }

    /// Finish a started `Node`. Thread-safe.
    pub fn end(n: Node) void {
        if (noop_impl) {
            assert(n.index == .none);
            return;
        }
        const index = n.index.unwrap() orelse return;
        const io = global_progress.io;
        const parent_ptr = parentByIndex(index);
        if (@atomicLoad(Node.Parent, parent_ptr, .monotonic).unwrap()) |parent_index| {
            _ = @atomicRmw(u32, &storageByIndex(parent_index).completed_count, .Add, 1, .monotonic);
            @atomicStore(Node.Parent, parent_ptr, .unused, .monotonic);

            if (storageByIndex(index).getIpcIndex()) |ipc_index| {
                const file = global_progress.ipc_files[ipc_index.slot];
                const ipc = @atomicRmw(
                    Ipc,
                    &global_progress.ipc[ipc_index.slot],
                    .And,
                    .{ .locked = true, .valid = false, .generation = std.math.maxInt(Ipc.Generation) },
                    .release,
                );
                assert(ipc.valid and ipc.generation == ipc_index.generation);
                if (!ipc.locked) file.close(io);
            }

            const freelist = &global_progress.node_freelist;
            var old_freelist = @atomicLoad(Freelist, freelist, .monotonic);
            while (true) {
                @atomicStore(Node.OptionalIndex, freelistNextByIndex(index), old_freelist.head, .monotonic);
                old_freelist = @cmpxchgWeak(
                    Freelist,
                    freelist,
                    old_freelist,
                    .{ .head = index.toOptional(), .generation = old_freelist.generation +% 1 },
                    .release, // ensure a matching `start` sees the freelist link written above
                    .monotonic, // our write above is irrelevant if we need to retry
                ) orelse {
                    // We won the race.
                    return;
                };
            }
        } else {
            if (global_progress.update_worker) |*worker| worker.cancel(io) catch {};
            for (&global_progress.ipc, &global_progress.ipc_files) |ipc, ipc_file| {
                assert(!ipc.locked or !ipc.valid); // missing call to end()
                if (ipc.locked or ipc.valid) ipc_file.close(io);
            }
        }
    }

    /// Used by `std.process.Child`. Thread-safe.
    pub fn setIpcFile(node: Node, expected_io_userdata: ?*anyopaque, file: Io.File) void {
        const index = node.index.unwrap() orelse return;
        const io = global_progress.io;
        assert(io.userdata == expected_io_userdata);
        for (0..ipc_storage_buffer_len) |_| {
            const slot: Ipc.Slot = @truncate(
                @atomicRmw(Ipc.SlotAtomic, &global_progress.ipc_next, .Add, 1, .monotonic),
            );
            if (slot >= ipc_storage_buffer_len) continue;
            const ipc_ptr = &global_progress.ipc[slot];
            const ipc = @atomicLoad(Ipc, ipc_ptr, .monotonic);
            if (ipc.locked or ipc.valid) continue;
            const generation = ipc.generation +% 1;
            if (@cmpxchgWeak(
                Ipc,
                ipc_ptr,
                ipc,
                .{ .locked = false, .valid = true, .generation = generation },
                .acquire,
                .monotonic,
            )) |_| continue;
            global_progress.ipc_files[slot] = file;
            storageByIndex(index).setIpcIndex(.{ .slot = slot, .generation = generation });
            break;
        } else {
            // There was no IPC slot available, so we'll drop this node's IPC info and just close
            // the fd. To avoid an old `estimated_total_items` or `completed_count` value still
            // being rendered for the node, we'll zero that field out (and the user is not allowed
            // to change it because they think we're doing IPC).
            file.close(io);
            @atomicStore(u32, &storageByIndex(index).completed_count, 0, .monotonic);
            @atomicStore(u32, &storageByIndex(index).estimated_total_count, 0, .monotonic);
        }
    }

    pub fn setIpcIndex(node: Node, ipc_index: Ipc.Index) void {
        storageByIndex(node.index.unwrap() orelse return).setIpcIndex(ipc_index);
    }

    /// Not thread-safe.
    pub fn takeIpcIndex(node: Node) ?Ipc.Index {
        const storage = storageByIndex(node.index.unwrap() orelse return null);
        switch (storage.estimated_total_count) {
            std.math.maxInt(u32) => {}, // indicates that there is an IPC index in `completed_count`
            0 => return null, // `setIpcFile` failed so we don't have an IPC index for this node
            else => unreachable, // not an IPC node
        }
        @atomicStore(u32, &storage.estimated_total_count, 0, .monotonic);
        return @bitCast(storage.completed_count);
    }

    fn storageByIndex(index: Node.Index) *Node.Storage {
        return &global_progress.node_storage[@backingInt(index)];
    }

    fn parentByIndex(index: Node.Index) *Node.Parent {
        return &global_progress.node_parents[@backingInt(index)];
    }

    fn freelistNextByIndex(index: Node.Index) *Node.OptionalIndex {
        return &global_progress.node_freelist_next[@backingInt(index)];
    }

    fn init(free_index: Index, parent: Parent, name: []const u8, estimated_total_items: usize) Node {
        assert(parent == .none or @backingInt(parent) < node_storage_buffer_len);

        const storage = storageByIndex(free_index);
        @atomicStore(u32, &storage.completed_count, 0, .monotonic);
        // Avoid u32 max int which is used to indicate a special state.
        const saturated_total_count = @min(std.math.maxInt(u32) - 1, estimated_total_items);
        @atomicStore(u32, &storage.estimated_total_count, saturated_total_count, .monotonic);
        const name_len = @min(max_name_len, name.len);
        copyAtomicStore(storage.name[0..name_len], name[0..name_len]);
        if (name_len < storage.name.len)
            @atomicStore(u8, &storage.name[name_len], 0, .monotonic);

        const parent_ptr = parentByIndex(free_index);
        if (std.debug.runtime_safety) {
            assert(@atomicLoad(Node.Parent, parent_ptr, .monotonic) == .unused);
        }
        @atomicStore(Node.Parent, parent_ptr, parent, .monotonic);

        return .{ .index = free_index.toOptional() };
    }
}