AllocFooter

const AllocFooter = struct {
    /// Hash of `data` with the seed as the hash of its address so that memcpys of allocation
    /// metadata are detected or are at least caught across runs.
    ///
    /// This stored value is xored with the canary value so that canary mismatches are detected.
    checksum: u32,
    /// Accesed atomically with `.monotonic` ordering to catch operation races.
    ///
    /// This stored value is xored with the hash of its address so that memcpys of allocation
    /// metadata are detected or are at least caught across runs.
    modify: Modify,
    data: Data,

    /// `8`: minimum alignment for `Bucket` allocations
    /// `@alignOf(usize)`: so that the offset of trailing data is at `@sizeOf(@This())`
    _: void align(@max(8, @alignOf(usize))) = {},

    comptime {
        assert(@alignOf(@This()) >= @max(8, @alignOf(usize)));
    }

    const Data = packed struct(u16) {
        len: Len,
        /// Low bits of the alignment.
        ///
        /// For non-extended headers, this is the entire alignment. The location of the previous
        /// header is directly before this allocation since footers in `Bucket` are gauraunteed to
        /// have at least 8-byte alignment.
        alignment: u2,
        /// Used only for bucket allocations.
        prev_extended: bool,

        const Len = enum(u13) {
            _,

            /// This footer is trailed (before the traces) by `Extended`.
            /// The high bits of the alignment are encoded as the offset from `extended_start`.
            ///
            /// This may be set even if `Extended` is not strictly necesary
            /// as a result of resizes and remaps.
            const extended_start: u13 = math.maxInt(u13) - ((@bitSizeOf(usize) - 1) >> 2);
        };
    };

    const Extended = struct {
        len: usize,
        container: Container,

        const Container = union {
            bucket_prev: ?*AllocFooter,
            large_entry: *Allocs.Entry,
        };

        comptime {
            // Exactly `usize` so this is directly after the regular footer
            // and so that traces start directly after `@sizeOf(@This())`.
            assert(@alignOf(@This()) == @alignOf(usize));
        }
    };

    const Modify = enum(u16) {
        // Random non-linear enum values to decrease the chance of undetected corruption.
        none = 0x2962,
        resized = 0x0030,
        remaped = 0x9068,
        freeing = 0x7f3d,
        freed = 0xb98b,
        _,

        fn setNone(m: *Modify) void {
            _ = @atomicRmw(Modify, m, .Xchg, .storedXor(.none, m), .monotonic);
        }

        fn opName(m: Modify) []const u8 {
            return switch (m) {
                .resized => "resize",
                .remaped => "remap",
                .freeing => "free",
                _, .none, .freed => unreachable,
            };
        }

        fn stateName(m: Modify) []const u8 {
            return switch (m) {
                .resized => "after resize",
                .remaped => "after remap",
                .freeing => "during free",
                .freed => "after free",
                _, .none => unreachable,
            };
        }

        fn storedXor(m: Modify, ptr: *Modify) Modify {
            const addr_hash: u16 = @truncate(std.hash.int(@intFromPtr(ptr)));
            return @fromBackingInt(@intCast(@backingInt(m) ^ addr_hash));
        }
    };

    fn isExtended(f: *AllocFooter) bool {
        return @backingInt(f.data.len) >= Data.Len.extended_start;
    }

    fn extended(f: *AllocFooter) *Extended {
        assert(f.isExtended());
        return @ptrFromInt(@intFromPtr(f) + @sizeOf(AllocFooter));
    }

    fn userMemory(f: *AllocFooter) []u8 {
        const memory_addr = @intFromPtr(f) - allocOffset(f.userLen());
        assert(f.userAlign().check(memory_addr));
        const memory_ptr: [*]u8 = @ptrFromInt(memory_addr);
        return memory_ptr[0..f.userLen()];
    }

    fn userLen(f: *AllocFooter) usize {
        const len_int = @backingInt(f.data.len);
        return if (len_int < Data.Len.extended_start) len_int else f.extended().len;
    }

    fn userAlign(f: *AllocFooter) Alignment {
        const high = (@backingInt(f.data.len) -| Data.Len.extended_start) << 2;
        return @fromBackingInt(@intCast(high | f.data.alignment));
    }

    fn bucketPrev(f: *AllocFooter, b: *Bucket, s: *SafeAllocator) ?*AllocFooter {
        if (f.isExtended()) return f.extended().container.bucket_prev;
        return b.allocFooterBefore(s, Bucket.fillAt(s, f.userMemory().ptr), f.data.prev_extended);
    }

    fn tracesPtr(f: *AllocFooter) [*]usize {
        const off_footer = @divExact(@sizeOf(AllocFooter), @sizeOf(usize));
        const off_extended = @as(usize, @divExact(@sizeOf(Extended), @sizeOf(usize))) *
            @intFromBool(f.isExtended());
        return @as([*]usize, @ptrCast(f))[off_footer + off_extended ..];
    }

    fn allocTrace(f: *AllocFooter, s: *SafeAllocator) []usize {
        return f.tracesPtr()[0..s.stack_trace_size];
    }

    fn freeTrace(f: *AllocFooter, s: *SafeAllocator) []usize {
        const trace_size = s.stack_trace_size;
        return f.tracesPtr()[trace_size..][0..trace_size];
    }

    fn actualChecksum(f: *AllocFooter, s: *SafeAllocator) u32 {
        if (f.isExtended()) {
            const len = f.extended().len;
            const addr: usize = if (s.isLarge(len, f.userAlign()))
                @intFromPtr(f.extended().container.large_entry)
            else
                @intFromPtr(f.extended().container.bucket_prev);

            const len_bytes: [@sizeOf(usize)]u8 = @bitCast(len);
            const container: [@sizeOf(usize)]u8 = @bitCast(addr);
            const regular_bytes: [2]u8 = @bitCast(f.data);
            const data_bytes = len_bytes ++ container ++ regular_bytes;

            return @truncate(std.hash.Wyhash.hash(@truncate(@intFromPtr(f)), &data_bytes));
        }
        return @truncate(std.hash.int(@as(u16, @bitCast(f.data)) ^ @intFromPtr(f)));
    }

    fn allocOffset(len: usize) usize {
        return Alignment.of(AllocFooter).forward(len);
    }

    fn allocAlign(a: Alignment) Alignment {
        return a.max(.of(AllocFooter));
    }

    /// Assumes the footer is in a bucket allocation; all
    /// large allocations require an extended header.
    fn requiresExtended(len: usize, alignment: Alignment) bool {
        return len >= Data.Len.extended_start or @backingInt(alignment) > math.maxInt(u2);
    }

    fn lenBucket(s: *SafeAllocator, is_extended: bool) usize {
        return Alignment.forward(.@"8", @sizeOf(AllocFooter) +
            @as(usize, @sizeOf(Extended)) * @intFromBool(is_extended) +
            s.stack_trace_size * @sizeOf(usize) * 2);
    }

    fn lenLarge(s: *SafeAllocator) usize {
        return @sizeOf(AllocFooter) + @sizeOf(Extended) + s.stack_trace_size * @sizeOf(usize);
    }

    fn allocLenBucket(s: *SafeAllocator, len: usize, is_extended: bool) usize {
        return allocOffset(len) + lenBucket(s, is_extended);
    }

    fn allocLenLarge(s: *SafeAllocator, len: usize) usize {
        return allocOffset(len) + lenLarge(s);
    }

    fn allocOffsetOrOom(len: usize) error{OutOfMemory}!usize {
        return alignForwardOrOom(.of(AllocFooter), len);
    }

    fn allocLenBucketOrOom(
        s: *SafeAllocator,
        len: usize,
        is_extended: bool,
    ) error{OutOfMemory}!usize {
        return addOrOom(try allocOffsetOrOom(len), lenBucket(s, is_extended));
    }

    fn of(user_memory: []u8) *AllocFooter {
        // Avoid panicing now if `memory.ptr` is not correctly aligned since a more
        // useful panic will be provided later by a mismatch or invalid footer.
        const aligned_start = Alignment.backward(.of(AllocFooter), @intFromPtr(user_memory.ptr));
        return @ptrFromInt(aligned_start + allocOffset(user_memory.len));
    }

    fn startModify(f: *AllocFooter, m: Modify, s: *SafeAllocator, mem_fmt: FormatMemory) void {
        const prev = @atomicRmw(
            Modify,
            &f.modify,
            .Xchg,
            .storedXor(m, &f.modify),
            .monotonic,
        ).storedXor(&f.modify);

        if (prev != .none) {
            @branchHint(.cold);
            const op_name = m.opName();
            switch (prev) {
                .none => unreachable,
                .resized, .remaped => panic(
                    \\{s} {s} of {f}
                    \\alloc: {f}
                    \\{s}:
                    // (panic stack trace)
                , .{
                    op_name,
                    prev.stateName(),
                    mem_fmt,
                    // The stack trace may have been overwritten, but at least give it a try
                    formatStackTrace(f.allocTrace(s)),
                    op_name,
                }),
                .freeing, .freed => {
                    if (prev == .freeing) {
                        // Wait for trace to become available
                        const complete: Modify = .storedXor(.freed, &f.modify);
                        while (@atomicLoad(Modify, &f.modify, .monotonic) != complete) {}
                        const b: *Bucket = .of(s, @ptrCast(f));
                        b.alloc_count.fenceAcqRel();
                    }
                    if (m == .freeing) {
                        panic(
                            \\double free of {f}
                            \\alloc: {f}
                            \\first free: {f}
                            \\second free:
                            // (panic stack trace)
                        , .{
                            mem_fmt,
                            formatStackTrace(f.allocTrace(s)),
                            formatStackTrace(f.freeTrace(s)),
                        });
                    } else {
                        panic(
                            \\{s} {s} of {f}
                            \\alloc: {f}
                            \\free: {f}
                            \\{s}:
                            // (panic stack trace)
                        , .{
                            op_name,
                            prev.stateName(),
                            mem_fmt,
                            formatStackTrace(f.allocTrace(s)),
                            formatStackTrace(f.freeTrace(s)),
                            op_name,
                        });
                    }
                },
                _ => panic(
                    "{s} of invalid memory {f} or corrupted metadata",
                    .{ m.opName(), mem_fmt },
                ),
            }
            comptime unreachable;
        }

        const expected_checksum = f.actualChecksum(s);
        if (f.checksum ^ s.canary != expected_checksum) {
            @branchHint(.cold);
            const other_canary = f.checksum ^ expected_checksum;
            panic(
                "{s} of invalid memory {f}, corrupted metadata, or foreign allocation from canary 0x{x}",
                .{ m.opName(), mem_fmt, other_canary },
            );
        }

        if (f.userLen() != mem_fmt.memory.len or f.userAlign() != mem_fmt.alignment) {
            const op_name = m.opName();
            panic(
                \\{s} of {f} mismatches allocation of {f}
                \\alloc: {f}
                \\{s}:
                // (panic stack trace)
            , .{ op_name, mem_fmt, FormatMemory{
                .memory = f.userMemory(),
                .alignment = f.userAlign(),
            }, formatStackTrace(f.allocTrace(s)), op_name });
        }
    }

    /// It is the caller's responsibility to `.acquire` fence the respective `Bucket.alloc_count`.
    fn populate(
        memory: []align(@alignOf(AllocFooter)) u8,
        len: usize,
        alignment: Alignment,
        ra: usize,
        /// `true` for large allocations
        is_extended: bool,
        /// `false` for large allocations
        prev_extended: bool,
        container: Extended.Container,
        s: *SafeAllocator,
    ) *AllocFooter {
        const footer: *AllocFooter = @ptrCast(@alignCast(memory[allocOffset(len)..].ptr));

        if (!is_extended) {
            footer.data = .{
                .len = @fromBackingInt(@intCast(len)),
                .alignment = @intCast(@backingInt(alignment)),
                .prev_extended = prev_extended,
            };
            assert(!footer.isExtended());
        } else {
            footer.data = .{
                .len = @fromBackingInt(@intCast(Data.Len.extended_start + (@backingInt(alignment) >> 2))),
                .alignment = @truncate(@backingInt(alignment)),
                .prev_extended = prev_extended,
            };
            assert(footer.isExtended());
            footer.extended().* = .{
                .len = len,
                .container = container,
            };
        }

        captureStackTrace(footer.allocTrace(s), ra);
        footer.checksum = footer.actualChecksum(s) ^ s.canary;
        footer.modify.setNone();

        return footer;
    }
}