// Kernel DMA core. This file owns device adapters, common buffers, pinned memory, streaming mappings, and bounce copies.

#include <ntlite/atomic.h>
#include <stdbool.h>
#include <intrin.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>

#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddk/dma.h>
#include <kernel/io_dma.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_process.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_pte.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pool.h>

#define IO_DMA_ADAPTER_TAG MM_POOL_TAG('d', 'm', 'a', 'a')
#define IO_DMA_BUFFER_TAG  MM_POOL_TAG('d', 'm', 'a', 'b')
#define IO_DMA_PIN_TAG     MM_POOL_TAG('d', 'm', 'a', 'p')
#define IO_DMA_MAP_TAG     MM_POOL_TAG('d', 'm', 'a', 'm')

struct IO_DMA_ADAPTER {
    volatile int32_t RefCount;
    volatile int32_t Removed;
    SPINLOCK Lock;
    DEVICE_OBJECT* Device;
    IO_DMA_REQUIREMENTS Requirements;
};

struct IO_DMA_BUFFER {
    IO_DMA_ADAPTER* Adapter;
    void* CpuAddress;
    PHYSICAL_ADDRESS PhysicalAddress;
    uint64_t DeviceAddress;
    size_t Size;
    size_t AllocationSize;
    size_t PageCount;
};

typedef enum MM_PINNED_MEMORY_ORIGIN {
    MmPinnedKernel = 1,
    MmPinnedUser
} MM_PINNED_MEMORY_ORIGIN;

struct MM_PINNED_MEMORY {
    volatile int32_t RefCount;
    MM_PINNED_MEMORY_ORIGIN Origin;
    IO_DMA_DIRECTION Direction;
    void* Address;
    size_t Size;
    size_t PageCount;
    size_t FirstPageOffset;
    MM_LOCKED_USER_BUFFER* UserBuffer;
    PHYSICAL_ADDRESS Frames[1];
};

struct IO_DMA_MAPPING {
    IO_DMA_ADAPTER* Adapter;
    MM_PINNED_MEMORY* Memory;
    IO_DMA_BUFFER* BounceBuffer;
    size_t MemoryOffset;
    size_t Length;
    IO_DMA_DIRECTION Direction;
    uint32_t SegmentCount;
    bool DeviceOwnsMapping;
    IO_DMA_SEGMENT Segments[1];
};

typedef struct IO_DMA_INTERNAL_STATS {
    volatile int64_t AdaptersCreated;
    volatile int64_t AdaptersDestroyed;
    volatile int64_t AdaptersInUse;
    volatile int64_t CommonBuffersAllocated;
    volatile int64_t CommonBuffersFreed;
    volatile int64_t CommonBuffersInUse;
    volatile int64_t CommonBufferBytesInUse;
    volatile int64_t PinnedKernelRanges;
    volatile int64_t PinnedUserRanges;
    volatile int64_t PinnedRangesReleased;
    volatile int64_t PinnedRangesInUse;
    volatile int64_t PinnedPagesInUse;
    volatile int64_t StreamingMappings;
    volatile int64_t StreamingMappingsUnmapped;
    volatile int64_t StreamingMappingsInUse;
    volatile int64_t DirectMappings;
    volatile int64_t BounceMappings;
    volatile int64_t BounceBytesCopiedToDevice;
    volatile int64_t BounceBytesCopiedToCpu;
    volatile int64_t FailedAllocations;
    volatile int64_t FailedMappings;
} IO_DMA_INTERNAL_STATS;

static IO_DMA_INTERNAL_STATS gIoDmaStats;

// Returns whether a size is a nonzero power of two.
static bool IoDmaIsPowerOfTwo(size_t value) {
    return value && ((value & (value - 1)) == 0);
}

// Normalizes and validates one adapter requirement set.
static KSTATUS IoDmaNormalizeRequirements(const IO_DMA_REQUIREMENTS* source,
                                           IO_DMA_REQUIREMENTS* destination) {
    if (!destination)
        return NTLSTATUS_INVALID_PARAMETER;

    memset(destination, 0, sizeof(*destination));
    if (source)
        *destination = *source;

    if (!destination->MaximumDeviceAddress)
        destination->MaximumDeviceAddress = (uint64_t)-1;
    if (!destination->Alignment)
        destination->Alignment = 1;
    if (!destination->MaximumSegmentLength)
        destination->MaximumSegmentLength = (size_t)-1;
    if (!destination->MaximumSegments)
        destination->MaximumSegments = UINT32_MAX;

    if (!IoDmaIsPowerOfTwo(destination->Alignment))
        return NTLSTATUS_INVALID_PARAMETER;
    if (destination->Boundary && !IoDmaIsPowerOfTwo(destination->Boundary))
        return NTLSTATUS_INVALID_PARAMETER;
    return NTLSTATUS_SUCCESS;
}

// Returns whether two normalized requirement sets are identical.
static bool IoDmaRequirementsEqual(const IO_DMA_REQUIREMENTS* left,
                                  const IO_DMA_REQUIREMENTS* right) {
    return left && right &&
           left->MaximumDeviceAddress == right->MaximumDeviceAddress &&
           left->Alignment == right->Alignment &&
           left->Boundary == right->Boundary &&
           left->MaximumSegmentLength == right->MaximumSegmentLength &&
           left->MaximumSegments == right->MaximumSegments;
}

// Returns whether device removal has closed the adapter to new ownership.
static bool IoDmaAdapterRemoved(const IO_DMA_ADAPTER* adapter) {
    if (!adapter)
        return true;
    return NtlAtomicLoad32(&adapter->Removed) != 0;
}

// Adds one reference to an adapter that is already stabilized by its owner.
static void IoDmaReferenceAdapter(IO_DMA_ADAPTER* adapter) {
    if (adapter)
        NtlAtomicIncrement32(&adapter->RefCount);
}

// Adds a reference only while the adapter still accepts new DMA ownership.
static bool IoDmaTryReferenceAdapter(IO_DMA_ADAPTER* adapter) {
    if (!adapter)
        return false;
    uintptr_t flags = KeAcquireSpinlock(&adapter->Lock);
    if (adapter->Removed) {
        KeReleaseSpinlock(&adapter->Lock, flags);
        return false;
    }
    NtlAtomicIncrement32(&adapter->RefCount);
    KeReleaseSpinlock(&adapter->Lock, flags);
    return true;
}

// Releases one adapter reference and destroys the detached adapter at zero.
static void IoDmaPutAdapterImpl(IO_DMA_ADAPTER* adapter) {
    if (!adapter)
        return;
    int32_t remaining = NtlAtomicDecrement32(&adapter->RefCount);
    if (remaining > 0)
        return;
    if (remaining < 0)
        KePanic(KE_PANIC_CODE_GENERIC,
                "DMA: adapter reference underflow adapter=0x%016llX\n",
                (unsigned long long)(uintptr_t)adapter);
    if (!IoDmaAdapterRemoved(adapter))
        KePanic(KE_PANIC_CODE_GENERIC,
                "DMA: attached adapter reached zero references adapter=0x%016llX\n",
                (unsigned long long)(uintptr_t)adapter);
    NtlAtomicIncrement64(&gIoDmaStats.AdaptersDestroyed);
    NtlAtomicDecrement64(&gIoDmaStats.AdaptersInUse);
    MmFreePool(adapter);
}

// Returns the one canonical DMA adapter owned by a device.
static KSTATUS IoDmaGetAdapterImpl(DEVICE_OBJECT* device,
                                   const IO_DMA_REQUIREMENTS* requirements,
                                   IO_DMA_ADAPTER** outAdapter) {
    if (!device || !outAdapter)
        return NTLSTATUS_INVALID_PARAMETER;
    *outAdapter = (IO_DMA_ADAPTER*)0;

    IO_DMA_REQUIREMENTS normalized;
    KSTATUS status = IoDmaNormalizeRequirements(requirements, &normalized);
    if (!KSUCCESS(status))
        return status;

    uintptr_t flags = KeAcquireSpinlock(&device->Lock);
    if (device->DeletePending || device->PnpState == DevicePnpDeleted) {
        KeReleaseSpinlock(&device->Lock, flags);
        return NTLSTATUS_OBJECT_DEAD;
    }

    IO_DMA_ADAPTER* adapter = device->DmaAdapter;
    if (adapter) {
        if (!IoDmaRequirementsEqual(&adapter->Requirements, &normalized)) {
            KeReleaseSpinlock(&device->Lock, flags);
            return NTLSTATUS_ALREADY_EXISTS;
        }
        IoDmaReferenceAdapter(adapter);
        KeReleaseSpinlock(&device->Lock, flags);
        *outAdapter = adapter;
        return NTLSTATUS_SUCCESS;
    }
    KeReleaseSpinlock(&device->Lock, flags);

    IO_DMA_ADAPTER* candidate = (IO_DMA_ADAPTER*)MmAllocatePoolTag(sizeof(*candidate),
                                                                   IO_DMA_ADAPTER_TAG);
    if (!candidate) {
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(candidate, 0, sizeof(*candidate));
    candidate->RefCount = 2;
    candidate->Device = device;
    candidate->Requirements = normalized;

    flags = KeAcquireSpinlock(&device->Lock);
    if (device->DeletePending || device->PnpState == DevicePnpDeleted) {
        KeReleaseSpinlock(&device->Lock, flags);
        MmFreePool(candidate);
        return NTLSTATUS_OBJECT_DEAD;
    }

    adapter = device->DmaAdapter;
    if (adapter) {
        if (!IoDmaRequirementsEqual(&adapter->Requirements, &normalized)) {
            KeReleaseSpinlock(&device->Lock, flags);
            MmFreePool(candidate);
            return NTLSTATUS_ALREADY_EXISTS;
        }
        IoDmaReferenceAdapter(adapter);
        KeReleaseSpinlock(&device->Lock, flags);
        MmFreePool(candidate);
        *outAdapter = adapter;
        return NTLSTATUS_SUCCESS;
    }

    device->DmaAdapter = candidate;
    KeReleaseSpinlock(&device->Lock, flags);

    NtlAtomicIncrement64(&gIoDmaStats.AdaptersCreated);
    NtlAtomicIncrement64(&gIoDmaStats.AdaptersInUse);
    *outAdapter = candidate;
    return NTLSTATUS_SUCCESS;
}

// Detaches a device-owned adapter so no new mappings can be created.
void IoDmaDetachDevice(DEVICE_OBJECT* device) {
    if (!device)
        return;

    uintptr_t flags = KeAcquireSpinlock(&device->Lock);
    IO_DMA_ADAPTER* adapter = device->DmaAdapter;
    device->DmaAdapter = (IO_DMA_ADAPTER*)0;
    if (adapter) {
        uintptr_t adapterFlags = KeAcquireSpinlock(&adapter->Lock);
        adapter->Removed = 1;
        adapter->Device = (DEVICE_OBJECT*)0;
        KeReleaseSpinlock(&adapter->Lock, adapterFlags);
    }
    KeReleaseSpinlock(&device->Lock, flags);

    if (adapter)
        IoDmaPutAdapterImpl(adapter);
}

// Converts an inclusive device-address limit to the allocator's exclusive limit.
static PHYSICAL_ADDRESS IoDmaMaximumExclusive(const IO_DMA_ADAPTER* adapter) {
    if (!adapter || adapter->Requirements.MaximumDeviceAddress == (uint64_t)-1)
        return (PHYSICAL_ADDRESS)-1;
    return adapter->Requirements.MaximumDeviceAddress + 1;
}

// Allocates one contiguous backing run while applying address and alignment requirements.
static KSTATUS IoDmaAllocateCommonBufferInternal(IO_DMA_ADAPTER* adapter,
                                                  size_t size,
                                                  size_t alignment,
                                                  int enforceAllocationBoundary,
                                                  IO_DMA_BUFFER** outBuffer) {
    if (!adapter || !outBuffer || !size)
        return NTLSTATUS_INVALID_PARAMETER;
    *outBuffer = (IO_DMA_BUFFER*)0;
    if (IoDmaAdapterRemoved(adapter))
        return NTLSTATUS_OBJECT_DEAD;

    if (!alignment)
        alignment = 1;
    if (!IoDmaIsPowerOfTwo(alignment))
        return NTLSTATUS_INVALID_PARAMETER;
    if (alignment < adapter->Requirements.Alignment)
        alignment = adapter->Requirements.Alignment;
    if (alignment < PAGE_SIZE)
        alignment = PAGE_SIZE;

    if (size > ((size_t)-1) - (PAGE_SIZE - 1))
        return NTLSTATUS_NO_MEMORY;
    size_t allocationSize = (size + PAGE_SIZE - 1) & PAGE_MASK;
    size_t pageCount = allocationSize / PAGE_SIZE;

    size_t allocationBoundary = 0;
    if (enforceAllocationBoundary && adapter->Requirements.Boundary) {
        if (size > adapter->Requirements.Boundary)
            return NTLSTATUS_NOT_SUPPORTED;
        if (adapter->Requirements.Boundary >= PAGE_SIZE)
            allocationBoundary = adapter->Requirements.Boundary;
    }

    PHYSICAL_ADDRESS physical = MmPfnAllocateContiguousInRange(
        pageCount,
        PAGE_SIZE,
        IoDmaMaximumExclusive(adapter),
        alignment,
        allocationBoundary,
        MmPfnUseDma);
    if (!physical) {
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_NO_MEMORY;
    }

    IO_DMA_BUFFER* buffer = (IO_DMA_BUFFER*)MmAllocatePoolTag(sizeof(*buffer), IO_DMA_BUFFER_TAG);
    if (!buffer) {
        for (size_t i = 0; i < pageCount; ++i)
            (void)MmPfnFree(physical + i * PAGE_SIZE);
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_NO_MEMORY;
    }

    memset(buffer, 0, sizeof(*buffer));
    buffer->Adapter = adapter;
    buffer->CpuAddress = MmPfnToVirtual(physical);
    buffer->PhysicalAddress = physical;
    buffer->DeviceAddress = physical;
    buffer->Size = size;
    buffer->AllocationSize = allocationSize;
    buffer->PageCount = pageCount;
    if (!buffer->CpuAddress) {
        MmFreePool(buffer);
        for (size_t i = 0; i < pageCount; ++i)
            (void)MmPfnFree(physical + i * PAGE_SIZE);
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_FAILURE;
    }

    memset(buffer->CpuAddress, 0, allocationSize);
    if (!IoDmaTryReferenceAdapter(adapter)) {
        MmFreePool(buffer);
        for (size_t i = 0; i < pageCount; ++i)
            (void)MmPfnFree(physical + i * PAGE_SIZE);
        return NTLSTATUS_OBJECT_DEAD;
    }
    NtlAtomicIncrement64(&gIoDmaStats.CommonBuffersAllocated);
    NtlAtomicIncrement64(&gIoDmaStats.CommonBuffersInUse);
    NtlAtomicExchangeAdd64(&gIoDmaStats.CommonBufferBytesInUse, (int64_t)allocationSize);
    *outBuffer = buffer;
    return NTLSTATUS_SUCCESS;
}

// Allocates one coherent common buffer represented by one device segment.
static KSTATUS IoDmaAllocateCommonBufferImpl(IO_DMA_ADAPTER* adapter,
                                              size_t size,
                                              size_t alignment,
                                              IO_DMA_BUFFER** outBuffer) {
    return IoDmaAllocateCommonBufferInternal(adapter, size, alignment, 1, outBuffer);
}

// Frees one common buffer and its owned physical pages.
static void IoDmaFreeCommonBufferImpl(IO_DMA_BUFFER* buffer) {
    if (!buffer)
        return;

    IO_DMA_ADAPTER* adapter = buffer->Adapter;
    PHYSICAL_ADDRESS physical = buffer->PhysicalAddress;
    size_t pageCount = buffer->PageCount;
    size_t allocationSize = buffer->AllocationSize;
    memset(buffer, 0, sizeof(*buffer));
    MmFreePool(buffer);

    for (size_t i = 0; i < pageCount; ++i)
        (void)MmPfnFree(physical + i * PAGE_SIZE);

    NtlAtomicIncrement64(&gIoDmaStats.CommonBuffersFreed);
    NtlAtomicDecrement64(&gIoDmaStats.CommonBuffersInUse);
    NtlAtomicExchangeAdd64(&gIoDmaStats.CommonBufferBytesInUse, -(int64_t)allocationSize);
    IoDmaPutAdapterImpl(adapter);
}

// Returns the CPU address of a common buffer.
static void* IoDmaGetCommonBufferAddressImpl(const IO_DMA_BUFFER* buffer) {
    return buffer ? buffer->CpuAddress : (void*)0;
}

// Returns the first device address of a common buffer.
static uint64_t IoDmaGetCommonBufferDeviceAddressImpl(const IO_DMA_BUFFER* buffer) {
    return buffer ? buffer->DeviceAddress : 0;
}

// Returns the caller-requested size of a common buffer.
static size_t IoDmaGetCommonBufferSizeImpl(const IO_DMA_BUFFER* buffer) {
    return buffer ? buffer->Size : 0;
}

// Returns whether a pinned access mode can satisfy a mapping direction.
static bool IoDmaDirectionAllows(IO_DMA_DIRECTION pinned, IO_DMA_DIRECTION mapped) {
    if (pinned == IoDmaBidirectional)
        return mapped == IoDmaToDevice || mapped == IoDmaFromDevice || mapped == IoDmaBidirectional;
    return pinned == mapped;
}

// Allocates a pinned-memory descriptor with room for every backing frame.
static MM_PINNED_MEMORY* IoDmaAllocatePinnedMemory(size_t pageCount) {
    if (!pageCount || pageCount > (((size_t)-1) - sizeof(MM_PINNED_MEMORY)) / sizeof(PHYSICAL_ADDRESS))
        return (MM_PINNED_MEMORY*)0;
    size_t bytes = sizeof(MM_PINNED_MEMORY) + (pageCount - 1) * sizeof(PHYSICAL_ADDRESS);
    MM_PINNED_MEMORY* memory = (MM_PINNED_MEMORY*)MmAllocatePoolTag(bytes, IO_DMA_PIN_TAG);
    if (memory)
        memset(memory, 0, bytes);
    return memory;
}

// Releases a pinned-memory descriptor after its last mapping leaves.
static void IoDmaReleasePinnedMemory(MM_PINNED_MEMORY* memory) {
    if (!memory)
        return;
    int32_t remaining = NtlAtomicDecrement32(&memory->RefCount);
    if (remaining > 0)
        return;
    if (remaining < 0)
        KePanic(KE_PANIC_CODE_GENERIC,
                "DMA: pinned-memory reference underflow memory=0x%016llX\n",
                (unsigned long long)(uintptr_t)memory);

    if (memory->Origin == MmPinnedUser) {
        MmUnlockUserBuffer(memory->UserBuffer);
    } else {
        for (size_t i = 0; i < memory->PageCount; ++i)
            (void)MmPfnUnpin(memory->Frames[i]);
    }
    NtlAtomicIncrement64(&gIoDmaStats.PinnedRangesReleased);
    NtlAtomicDecrement64(&gIoDmaStats.PinnedRangesInUse);
    NtlAtomicExchangeAdd64(&gIoDmaStats.PinnedPagesInUse, -(int64_t)memory->PageCount);
    MmFreePool(memory);
}

// Adds one mapping reference to pinned memory.
static void IoDmaReferencePinnedMemory(MM_PINNED_MEMORY* memory) {
    if (memory)
        NtlAtomicIncrement32(&memory->RefCount);
}

// Pins a kernel virtual range for later DMA mapping.
static KSTATUS IoDmaPinKernelMemoryImpl(void* address,
                                        size_t size,
                                        IO_DMA_DIRECTION direction,
                                        MM_PINNED_MEMORY** outMemory) {
    if (!outMemory)
        return NTLSTATUS_INVALID_PARAMETER;
    *outMemory = (MM_PINNED_MEMORY*)0;
    if (!address || !size)
        return NTLSTATUS_INVALID_PARAMETER;
    if (direction < IoDmaToDevice || direction > IoDmaBidirectional)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t start = (uintptr_t)address;
    if (start < KERNEL_BASE || start + size < start)
        return NTLSTATUS_INVALID_PARAMETER;
    uintptr_t firstPage = start & PAGE_MASK;
    uintptr_t lastByte = start + size - 1;
    size_t pageCount = ((lastByte & PAGE_MASK) - firstPage) / PAGE_SIZE + 1;

    MM_PINNED_MEMORY* memory = IoDmaAllocatePinnedMemory(pageCount);
    if (!memory) {
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_NO_MEMORY;
    }

    memory->RefCount = 1;
    memory->Origin = MmPinnedKernel;
    memory->Direction = direction;
    memory->Address = address;
    memory->Size = size;
    memory->PageCount = pageCount;
    memory->FirstPageOffset = start & (PAGE_SIZE - 1);

    size_t pinned = 0;
    KSTATUS status = NTLSTATUS_SUCCESS;
    for (size_t i = 0; i < pageCount; ++i) {
        PHYSICAL_ADDRESS physical = MmPteTranslate(firstPage + i * PAGE_SIZE);
        PHYSICAL_ADDRESS frame = physical & PAGE_MASK;
        if (!physical || !MmPfnGet(frame)) {
            status = NTLSTATUS_NOT_FOUND;
            break;
        }
        status = MmPfnPin(frame);
        if (!KSUCCESS(status))
            break;
        memory->Frames[i] = frame;
        pinned++;
    }

    if (!KSUCCESS(status)) {
        for (size_t i = 0; i < pinned; ++i)
            (void)MmPfnUnpin(memory->Frames[i]);
        MmFreePool(memory);
        return status;
    }

    NtlAtomicIncrement64(&gIoDmaStats.PinnedKernelRanges);
    NtlAtomicIncrement64(&gIoDmaStats.PinnedRangesInUse);
    NtlAtomicExchangeAdd64(&gIoDmaStats.PinnedPagesInUse, (int64_t)pageCount);
    *outMemory = memory;
    return NTLSTATUS_SUCCESS;
}

// Pins a current-process user range through the VAD locking contract.
static KSTATUS IoDmaPinUserMemoryImpl(struct PROCESS* process,
                                      void* address,
                                      size_t size,
                                      IO_DMA_DIRECTION direction,
                                      MM_PINNED_MEMORY** outMemory) {
    if (!outMemory)
        return NTLSTATUS_INVALID_PARAMETER;
    *outMemory = (MM_PINNED_MEMORY*)0;
    if (!address || !size)
        return NTLSTATUS_INVALID_PARAMETER;
    if (direction < IoDmaToDevice || direction > IoDmaBidirectional)
        return NTLSTATUS_INVALID_PARAMETER;

    if (!process)
        return NTLSTATUS_ACCESS_DENIED;

    MM_LOCKED_USER_BUFFER* locked = (MM_LOCKED_USER_BUFFER*)0;
    int writeAccess = direction == IoDmaFromDevice || direction == IoDmaBidirectional;
    KSTATUS status = MmLockUserBuffer(process, address, size, writeAccess, &locked);
    if (!KSUCCESS(status))
        return status;

    size_t pageCount = MmLockedUserBufferPageCount(locked);
    MM_PINNED_MEMORY* memory = IoDmaAllocatePinnedMemory(pageCount);
    if (!memory) {
        MmUnlockUserBuffer(locked);
        NtlAtomicIncrement64(&gIoDmaStats.FailedAllocations);
        return NTLSTATUS_NO_MEMORY;
    }

    memory->RefCount = 1;
    memory->Origin = MmPinnedUser;
    memory->Direction = direction;
    memory->Address = address;
    memory->Size = size;
    memory->PageCount = pageCount;
    memory->FirstPageOffset = (uintptr_t)address & (PAGE_SIZE - 1);
    memory->UserBuffer = locked;
    for (size_t i = 0; i < pageCount; ++i)
        memory->Frames[i] = MmLockedUserBufferPageFrame(locked, i);

    NtlAtomicIncrement64(&gIoDmaStats.PinnedUserRanges);
    NtlAtomicIncrement64(&gIoDmaStats.PinnedRangesInUse);
    NtlAtomicExchangeAdd64(&gIoDmaStats.PinnedPagesInUse, (int64_t)pageCount);
    *outMemory = memory;
    return NTLSTATUS_SUCCESS;
}

// Releases the caller's pinned-memory reference.
static void IoDmaUnpinMemoryImpl(MM_PINNED_MEMORY* memory) {
    IoDmaReleasePinnedMemory(memory);
}

// Returns the physical address and bytes remaining in one pinned page span.
static KSTATUS IoDmaGetPinnedSpan(const MM_PINNED_MEMORY* memory,
                                  size_t offset,
                                  PHYSICAL_ADDRESS* address,
                                  size_t* length) {
    if (!memory || !address || !length || offset >= memory->Size)
        return NTLSTATUS_INVALID_PARAMETER;

    size_t absolute = memory->FirstPageOffset + offset;
    size_t pageIndex = absolute / PAGE_SIZE;
    size_t pageOffset = absolute & (PAGE_SIZE - 1);
    if (pageIndex >= memory->PageCount || !memory->Frames[pageIndex])
        return NTLSTATUS_NOT_FOUND;

    *address = memory->Frames[pageIndex] + pageOffset;
    *length = PAGE_SIZE - pageOffset;
    size_t remaining = memory->Size - offset;
    if (*length > remaining)
        *length = remaining;
    return NTLSTATUS_SUCCESS;
}

// Copies a pinned range into one linear kernel buffer.
static KSTATUS IoDmaCopyFromPinned(const MM_PINNED_MEMORY* memory,
                                   size_t offset,
                                   void* destination,
                                   size_t length) {
    if (!memory || (!destination && length) || offset > memory->Size || length > memory->Size - offset)
        return NTLSTATUS_INVALID_PARAMETER;

    uint8_t* output = (uint8_t*)destination;
    size_t copied = 0;
    while (copied < length) {
        size_t absolute = memory->FirstPageOffset + offset + copied;
        size_t pageIndex = absolute / PAGE_SIZE;
        size_t pageOffset = absolute & (PAGE_SIZE - 1);
        size_t chunk = PAGE_SIZE - pageOffset;
        if (chunk > length - copied)
            chunk = length - copied;
        void* page = MmPfnToVirtual(memory->Frames[pageIndex]);
        if (!page)
            return NTLSTATUS_FAILURE;
        memcpy(output + copied, (uint8_t*)page + pageOffset, chunk);
        copied += chunk;
    }
    return NTLSTATUS_SUCCESS;
}

// Copies one linear kernel buffer back into a pinned range.
static KSTATUS IoDmaCopyToPinned(MM_PINNED_MEMORY* memory,
                                 size_t offset,
                                 const void* source,
                                 size_t length) {
    if (!memory || (!source && length) || offset > memory->Size || length > memory->Size - offset)
        return NTLSTATUS_INVALID_PARAMETER;

    const uint8_t* input = (const uint8_t*)source;
    size_t copied = 0;
    while (copied < length) {
        size_t absolute = memory->FirstPageOffset + offset + copied;
        size_t pageIndex = absolute / PAGE_SIZE;
        size_t pageOffset = absolute & (PAGE_SIZE - 1);
        size_t chunk = PAGE_SIZE - pageOffset;
        if (chunk > length - copied)
            chunk = length - copied;
        void* page = MmPfnToVirtual(memory->Frames[pageIndex]);
        if (!page)
            return NTLSTATUS_FAILURE;
        memcpy((uint8_t*)page + pageOffset, input + copied, chunk);
        copied += chunk;
    }
    return NTLSTATUS_SUCCESS;
}

// Returns the maximum bytes one segment may cover from a given address.
static size_t IoDmaSegmentCapacity(const IO_DMA_REQUIREMENTS* requirements,
                                   uint64_t address,
                                   size_t currentLength) {
    size_t capacity = requirements->MaximumSegmentLength - currentLength;
    if (requirements->Boundary) {
        size_t boundaryRemaining = requirements->Boundary -
                                   (size_t)(address & (requirements->Boundary - 1));
        if (boundaryRemaining > currentLength)
            boundaryRemaining -= currentLength;
        else
            boundaryRemaining = 0;
        if (capacity > boundaryRemaining)
            capacity = boundaryRemaining;
    }
    return capacity;
}

// Adds physical spans to a constrained scatter/gather list.
static KSTATUS IoDmaBuildSegments(const MM_PINNED_MEMORY* memory,
                                  size_t offset,
                                  size_t length,
                                  const IO_DMA_REQUIREMENTS* requirements,
                                  IO_DMA_SEGMENT* segments,
                                  uint32_t* segmentCount) {
    if (!memory || !requirements || !segmentCount || !length)
        return NTLSTATUS_INVALID_PARAMETER;

    uint32_t count = 0;
    uint64_t currentAddress = 0;
    size_t currentLength = 0;
    size_t consumed = 0;

    while (consumed < length) {
        PHYSICAL_ADDRESS address = 0;
        size_t spanLength = 0;
        KSTATUS status = IoDmaGetPinnedSpan(memory, offset + consumed, &address, &spanLength);
        if (!KSUCCESS(status))
            return status;
        if (spanLength > length - consumed)
            spanLength = length - consumed;

        if (address > requirements->MaximumDeviceAddress ||
            spanLength - 1 > requirements->MaximumDeviceAddress - address)
            return NTLSTATUS_NOT_SUPPORTED;

        size_t spanConsumed = 0;
        while (spanConsumed < spanLength) {
            uint64_t pieceAddress = address + spanConsumed;
            if (currentLength && currentAddress + currentLength == pieceAddress) {
                size_t capacity = IoDmaSegmentCapacity(requirements, currentAddress, currentLength);
                if (capacity) {
                    size_t take = spanLength - spanConsumed;
                    if (take > capacity)
                        take = capacity;
                    currentLength += take;
                    if (segments)
                        segments[count - 1].Length = currentLength;
                    spanConsumed += take;
                    continue;
                }
            }

            if ((pieceAddress & (requirements->Alignment - 1)) != 0)
                return NTLSTATUS_NOT_SUPPORTED;
            if (count == requirements->MaximumSegments)
                return NTLSTATUS_NOT_SUPPORTED;

            size_t capacity = requirements->MaximumSegmentLength;
            if (requirements->Boundary) {
                size_t boundaryRemaining = requirements->Boundary -
                                           (size_t)(pieceAddress & (requirements->Boundary - 1));
                if (capacity > boundaryRemaining)
                    capacity = boundaryRemaining;
            }
            if (!capacity)
                return NTLSTATUS_NOT_SUPPORTED;

            size_t take = spanLength - spanConsumed;
            if (take > capacity)
                take = capacity;
            currentAddress = pieceAddress;
            currentLength = take;
            if (segments) {
                segments[count].DeviceAddress = currentAddress;
                segments[count].Length = currentLength;
            }
            count++;
            spanConsumed += take;
        }
        consumed += spanLength;
    }

    *segmentCount = count;
    return NTLSTATUS_SUCCESS;
}

// Builds a constrained scatter/gather list for one contiguous bounce allocation.
static KSTATUS IoDmaBuildContiguousSegments(uint64_t address,
                                            size_t length,
                                            const IO_DMA_REQUIREMENTS* requirements,
                                            IO_DMA_SEGMENT* segments,
                                            uint32_t* segmentCount) {
    if (!requirements || !segmentCount || !address || !length)
        return NTLSTATUS_INVALID_PARAMETER;
    if (address > requirements->MaximumDeviceAddress ||
        length - 1 > requirements->MaximumDeviceAddress - address)
        return NTLSTATUS_NOT_SUPPORTED;

    uint32_t count = 0;
    size_t consumed = 0;
    while (consumed < length) {
        uint64_t pieceAddress = address + consumed;
        if ((pieceAddress & (requirements->Alignment - 1)) != 0)
            return NTLSTATUS_NOT_SUPPORTED;
        if (count == requirements->MaximumSegments)
            return NTLSTATUS_NOT_SUPPORTED;

        size_t capacity = requirements->MaximumSegmentLength;
        if (requirements->Boundary) {
            size_t boundaryRemaining = requirements->Boundary -
                                       (size_t)(pieceAddress & (requirements->Boundary - 1));
            if (capacity > boundaryRemaining)
                capacity = boundaryRemaining;
        }
        if (!capacity)
            return NTLSTATUS_NOT_SUPPORTED;

        size_t take = length - consumed;
        if (take > capacity)
            take = capacity;
        if (segments) {
            segments[count].DeviceAddress = pieceAddress;
            segments[count].Length = take;
        }
        count++;
        consumed += take;
    }

    *segmentCount = count;
    return NTLSTATUS_SUCCESS;
}

// Allocates one mapping object with a fixed segment array.
static IO_DMA_MAPPING* IoDmaAllocateMapping(uint32_t segmentCount) {
    if (!segmentCount)
        return (IO_DMA_MAPPING*)0;
    size_t extraBytes = ((size_t)segmentCount - 1) * sizeof(IO_DMA_SEGMENT);
    if (extraBytes / sizeof(IO_DMA_SEGMENT) != (size_t)segmentCount - 1)
        return (IO_DMA_MAPPING*)0;
    size_t bytes = sizeof(IO_DMA_MAPPING) + extraBytes;
    if (bytes < extraBytes)
        return (IO_DMA_MAPPING*)0;
    IO_DMA_MAPPING* mapping = (IO_DMA_MAPPING*)MmAllocatePoolTag(bytes, IO_DMA_MAP_TAG);
    if (mapping)
        memset(mapping, 0, bytes);
    return mapping;
}

// Maps pinned memory directly or through an automatic contiguous bounce buffer.
static KSTATUS IoDmaMapPinnedMemoryImpl(IO_DMA_ADAPTER* adapter,
                                         MM_PINNED_MEMORY* memory,
                                         size_t offset,
                                         size_t length,
                                         IO_DMA_DIRECTION direction,
                                         uint32_t flags,
                                         IO_DMA_MAPPING** outMapping) {
    if (!outMapping)
        return NTLSTATUS_INVALID_PARAMETER;
    *outMapping = (IO_DMA_MAPPING*)0;
    if (!adapter || !memory || !length || offset > memory->Size || length > memory->Size - offset)
        return NTLSTATUS_INVALID_PARAMETER;
    if (IoDmaAdapterRemoved(adapter))
        return NTLSTATUS_OBJECT_DEAD;
    if (direction < IoDmaToDevice || direction > IoDmaBidirectional ||
        !IoDmaDirectionAllows(memory->Direction, direction))
        return NTLSTATUS_ACCESS_DENIED;
    if (flags & ~IO_DMA_MAP_NO_BOUNCE)
        return NTLSTATUS_INVALID_PARAMETER;

    uint32_t segmentCount = 0;
    KSTATUS directStatus = IoDmaBuildSegments(memory,
                                               offset,
                                               length,
                                               &adapter->Requirements,
                                               (IO_DMA_SEGMENT*)0,
                                               &segmentCount);
    IO_DMA_BUFFER* bounce = (IO_DMA_BUFFER*)0;
    KSTATUS status = directStatus;

    if (!KSUCCESS(status)) {
        if (status != NTLSTATUS_NOT_SUPPORTED || (flags & IO_DMA_MAP_NO_BOUNCE)) {
            NtlAtomicIncrement64(&gIoDmaStats.FailedMappings);
            return status;
        }

        status = IoDmaAllocateCommonBufferInternal(adapter,
                                                     length,
                                                     adapter->Requirements.Alignment,
                                                     0,
                                                     &bounce);
        if (!KSUCCESS(status)) {
            NtlAtomicIncrement64(&gIoDmaStats.FailedMappings);
            return status;
        }
        status = IoDmaBuildContiguousSegments(bounce->DeviceAddress,
                                               length,
                                               &adapter->Requirements,
                                               (IO_DMA_SEGMENT*)0,
                                               &segmentCount);
        if (!KSUCCESS(status)) {
            IoDmaFreeCommonBufferImpl(bounce);
            NtlAtomicIncrement64(&gIoDmaStats.FailedMappings);
            return status;
        }
    }

    IO_DMA_MAPPING* mapping = IoDmaAllocateMapping(segmentCount);
    if (!mapping) {
        IoDmaFreeCommonBufferImpl(bounce);
        NtlAtomicIncrement64(&gIoDmaStats.FailedMappings);
        return NTLSTATUS_NO_MEMORY;
    }

    if (bounce) {
        status = IoDmaBuildContiguousSegments(bounce->DeviceAddress,
                                               length,
                                               &adapter->Requirements,
                                               mapping->Segments,
                                               &segmentCount);
    } else {
        status = IoDmaBuildSegments(memory,
                                    offset,
                                    length,
                                    &adapter->Requirements,
                                    mapping->Segments,
                                    &segmentCount);
    }
    if (!KSUCCESS(status)) {
        MmFreePool(mapping);
        IoDmaFreeCommonBufferImpl(bounce);
        NtlAtomicIncrement64(&gIoDmaStats.FailedMappings);
        return status;
    }

    mapping->Adapter = adapter;
    mapping->Memory = memory;
    mapping->BounceBuffer = bounce;
    mapping->MemoryOffset = offset;
    mapping->Length = length;
    mapping->Direction = direction;
    mapping->SegmentCount = segmentCount;
    if (!IoDmaTryReferenceAdapter(adapter)) {
        MmFreePool(mapping);
        IoDmaFreeCommonBufferImpl(bounce);
        return NTLSTATUS_OBJECT_DEAD;
    }
    IoDmaReferencePinnedMemory(memory);

    NtlAtomicIncrement64(&gIoDmaStats.StreamingMappings);
    NtlAtomicIncrement64(&gIoDmaStats.StreamingMappingsInUse);
    if (bounce)
        NtlAtomicIncrement64(&gIoDmaStats.BounceMappings);
    else
        NtlAtomicIncrement64(&gIoDmaStats.DirectMappings);
    *outMapping = mapping;
    return NTLSTATUS_SUCCESS;
}

// Synchronizes CPU-written data before the device consumes a mapping.
static KSTATUS IoDmaSynchronizeForDeviceImpl(IO_DMA_MAPPING* mapping) {
    if (!mapping)
        return NTLSTATUS_INVALID_PARAMETER;

    if (mapping->DeviceOwnsMapping)
        return NTLSTATUS_FAILURE;

    if (mapping->Direction != IoDmaFromDevice && mapping->BounceBuffer) {
        KSTATUS status = IoDmaCopyFromPinned(mapping->Memory,
                                             mapping->MemoryOffset,
                                             mapping->BounceBuffer->CpuAddress,
                                             mapping->Length);
        if (!KSUCCESS(status))
            return status;
        NtlAtomicExchangeAdd64(&gIoDmaStats.BounceBytesCopiedToDevice,
                                  (int64_t)mapping->Length);
    }
    _ReadWriteBarrier();
    _mm_mfence();
    _ReadWriteBarrier();
    mapping->DeviceOwnsMapping = 1;
    return NTLSTATUS_SUCCESS;
}

// Synchronizes device-written data before the CPU consumes a mapping.
static KSTATUS IoDmaSynchronizeForCpuImpl(IO_DMA_MAPPING* mapping) {
    if (!mapping)
        return NTLSTATUS_INVALID_PARAMETER;
    if (!mapping->DeviceOwnsMapping)
        return NTLSTATUS_SUCCESS;

    _ReadWriteBarrier();
    _mm_mfence();
    _ReadWriteBarrier();
    if (mapping->Direction != IoDmaToDevice && mapping->BounceBuffer) {
        KSTATUS status = IoDmaCopyToPinned(mapping->Memory,
                                           mapping->MemoryOffset,
                                           mapping->BounceBuffer->CpuAddress,
                                           mapping->Length);
        if (!KSUCCESS(status))
            return status;
        NtlAtomicExchangeAdd64(&gIoDmaStats.BounceBytesCopiedToCpu,
                                  (int64_t)mapping->Length);
    }
    mapping->DeviceOwnsMapping = 0;
    return NTLSTATUS_SUCCESS;
}

// Releases a streaming mapping and performs the final receive-side bounce copy.
static void IoDmaUnmapMemoryImpl(IO_DMA_MAPPING* mapping) {
    if (!mapping)
        return;

    if (mapping->DeviceOwnsMapping) {
        KSTATUS status = IoDmaSynchronizeForCpuImpl(mapping);
        if (!KSUCCESS(status))
            KePanic(KE_PANIC_CODE_GENERIC,
                    "DMA: final receive synchronization failed mapping=0x%016llX status=%u\n",
                    (unsigned long long)(uintptr_t)mapping,
                    status);
    }
    IO_DMA_BUFFER* bounce = mapping->BounceBuffer;
    MM_PINNED_MEMORY* memory = mapping->Memory;
    IO_DMA_ADAPTER* adapter = mapping->Adapter;
    memset(mapping, 0, sizeof(*mapping));
    MmFreePool(mapping);
    IoDmaFreeCommonBufferImpl(bounce);
    IoDmaReleasePinnedMemory(memory);
    IoDmaPutAdapterImpl(adapter);
    NtlAtomicIncrement64(&gIoDmaStats.StreamingMappingsUnmapped);
    NtlAtomicDecrement64(&gIoDmaStats.StreamingMappingsInUse);
}

// Returns the number of device segments in a streaming mapping.
static uint32_t IoDmaGetSegmentCountImpl(const IO_DMA_MAPPING* mapping) {
    return mapping ? mapping->SegmentCount : 0;
}

// Returns the immutable device-segment array owned by a mapping.
static const IO_DMA_SEGMENT* IoDmaGetSegmentsImpl(const IO_DMA_MAPPING* mapping) {
    return mapping ? mapping->Segments : (const IO_DMA_SEGMENT*)0;
}

// Copies the current DMA counters into caller-owned storage.
static void IoDmaQueryStatsImpl(IO_DMA_STATS* outStats) {
    if (!outStats)
        return;
    outStats->AdaptersCreated = (uint64_t)gIoDmaStats.AdaptersCreated;
    outStats->AdaptersDestroyed = (uint64_t)gIoDmaStats.AdaptersDestroyed;
    outStats->AdaptersInUse = (uint64_t)gIoDmaStats.AdaptersInUse;
    outStats->CommonBuffersAllocated = (uint64_t)gIoDmaStats.CommonBuffersAllocated;
    outStats->CommonBuffersFreed = (uint64_t)gIoDmaStats.CommonBuffersFreed;
    outStats->CommonBuffersInUse = (uint64_t)gIoDmaStats.CommonBuffersInUse;
    outStats->CommonBufferBytesInUse = (uint64_t)gIoDmaStats.CommonBufferBytesInUse;
    outStats->PinnedKernelRanges = (uint64_t)gIoDmaStats.PinnedKernelRanges;
    outStats->PinnedUserRanges = (uint64_t)gIoDmaStats.PinnedUserRanges;
    outStats->PinnedRangesReleased = (uint64_t)gIoDmaStats.PinnedRangesReleased;
    outStats->PinnedRangesInUse = (uint64_t)gIoDmaStats.PinnedRangesInUse;
    outStats->PinnedPagesInUse = (uint64_t)gIoDmaStats.PinnedPagesInUse;
    outStats->StreamingMappings = (uint64_t)gIoDmaStats.StreamingMappings;
    outStats->StreamingMappingsUnmapped = (uint64_t)gIoDmaStats.StreamingMappingsUnmapped;
    outStats->StreamingMappingsInUse = (uint64_t)gIoDmaStats.StreamingMappingsInUse;
    outStats->DirectMappings = (uint64_t)gIoDmaStats.DirectMappings;
    outStats->BounceMappings = (uint64_t)gIoDmaStats.BounceMappings;
    outStats->BounceBytesCopiedToDevice = (uint64_t)gIoDmaStats.BounceBytesCopiedToDevice;
    outStats->BounceBytesCopiedToCpu = (uint64_t)gIoDmaStats.BounceBytesCopiedToCpu;
    outStats->FailedAllocations = (uint64_t)gIoDmaStats.FailedAllocations;
    outStats->FailedMappings = (uint64_t)gIoDmaStats.FailedMappings;
}

static const NTL_DDK_DMA gIoDmaDdk = {
    .GetAdapter = IoDmaGetAdapterImpl,
    .PutAdapter = IoDmaPutAdapterImpl,
    .AllocateCommonBuffer = IoDmaAllocateCommonBufferImpl,
    .FreeCommonBuffer = IoDmaFreeCommonBufferImpl,
    .GetCommonBufferAddress = IoDmaGetCommonBufferAddressImpl,
    .GetCommonBufferDeviceAddress = IoDmaGetCommonBufferDeviceAddressImpl,
    .GetCommonBufferSize = IoDmaGetCommonBufferSizeImpl,
    .PinKernelMemory = IoDmaPinKernelMemoryImpl,
    .PinUserMemory = IoDmaPinUserMemoryImpl,
    .UnpinMemory = IoDmaUnpinMemoryImpl,
    .MapPinnedMemory = IoDmaMapPinnedMemoryImpl,
    .UnmapMemory = IoDmaUnmapMemoryImpl,
    .SynchronizeForDevice = IoDmaSynchronizeForDeviceImpl,
    .SynchronizeForCpu = IoDmaSynchronizeForCpuImpl,
    .GetSegmentCount = IoDmaGetSegmentCountImpl,
    .GetSegments = IoDmaGetSegmentsImpl,
    .QueryStats = IoDmaQueryStatsImpl,
};

// Returns the kernel-owned DDK DMA table.
const NTL_DDK_DMA* IoGetDmaDdk(void) {
    return &gIoDmaDdk;
}