// SPDX-License-Identifier: GPL-2.0 OR MIT
/**************************************************************************
 *
 * Copyright (c) 2019-2025 Broadcom. All Rights Reserved. The term
 * “Broadcom” refers to Broadcom Inc. and/or its subsidiaries.
 *
 **************************************************************************/
#include "vmwgfx_bo.h"
#include "vmwgfx_drv.h"

/*
 * Different methods for tracking dirty:
 * VMW_BO_DIRTY_PAGETABLE - Scan the pagetable for hardware dirty bits
 * VMW_BO_DIRTY_MKWRITE - Write-protect page table entries and record write-
 * accesses in the VM mkwrite() callback
 */
enum vmw_bo_dirty_method {
        VMW_BO_DIRTY_PAGETABLE,
        VMW_BO_DIRTY_MKWRITE,
};

/*
 * No dirtied pages at scan trigger a transition to the _MKWRITE method,
 * similarly a certain percentage of dirty pages trigger a transition to
 * the _PAGETABLE method. How many triggers should we wait for before
 * changing method?
 */
#define VMW_DIRTY_NUM_CHANGE_TRIGGERS 2

/* Percentage to trigger a transition to the _PAGETABLE method */
#define VMW_DIRTY_PERCENTAGE 10

/**
 * struct vmw_bo_dirty - Dirty information for buffer objects
 * @ref_count: Reference count for this structure. Must be first member!
 * @start: First currently dirty bit
 * @end: Last currently dirty bit + 1
 * @method: The currently used dirty method
 * @change_count: Number of consecutive method change triggers
 * @bitmap_size: The size of the bitmap in bits. Typically equal to the
 * nuber of pages in the bo.
 * @bitmap: A bitmap where each bit represents a page. A set bit means a
 * dirty page.
 */
struct vmw_bo_dirty {
        struct   kref ref_count;
        unsigned long start;
        unsigned long end;
        enum vmw_bo_dirty_method method;
        unsigned int change_count;
        unsigned long bitmap_size;
        unsigned long bitmap[];
};

bool vmw_bo_is_dirty(struct vmw_bo *vbo)
{
        return vbo->dirty && (vbo->dirty->start < vbo->dirty->end);
}

/**
 * vmw_bo_dirty_scan_pagetable - Perform a pagetable scan for dirty bits
 * @vbo: The buffer object to scan
 *
 * Scans the pagetable for dirty bits. Clear those bits and modify the
 * dirty structure with the results. This function may change the
 * dirty-tracking method.
 */
static void vmw_bo_dirty_scan_pagetable(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;
        pgoff_t offset = drm_vma_node_start(&vbo->tbo.base.vma_node);
        struct address_space *mapping = vbo->tbo.bdev->dev_mapping;
        pgoff_t num_marked;

        num_marked = clean_record_shared_mapping_range
                (mapping,
                 offset, dirty->bitmap_size,
                 offset, &dirty->bitmap[0],
                 &dirty->start, &dirty->end);
        if (num_marked == 0)
                dirty->change_count++;
        else
                dirty->change_count = 0;

        if (dirty->change_count > VMW_DIRTY_NUM_CHANGE_TRIGGERS) {
                dirty->change_count = 0;
                dirty->method = VMW_BO_DIRTY_MKWRITE;
                wp_shared_mapping_range(mapping,
                                        offset, dirty->bitmap_size);
                clean_record_shared_mapping_range(mapping,
                                                  offset, dirty->bitmap_size,
                                                  offset, &dirty->bitmap[0],
                                                  &dirty->start, &dirty->end);
        }
}

/**
 * vmw_bo_dirty_scan_mkwrite - Reset the mkwrite dirty-tracking method
 * @vbo: The buffer object to scan
 *
 * Write-protect pages written to so that consecutive write accesses will
 * trigger a call to mkwrite.
 *
 * This function may change the dirty-tracking method.
 */
static void vmw_bo_dirty_scan_mkwrite(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;
        unsigned long offset = drm_vma_node_start(&vbo->tbo.base.vma_node);
        struct address_space *mapping = vbo->tbo.bdev->dev_mapping;
        pgoff_t num_marked;

        if (dirty->end <= dirty->start)
                return;

        num_marked = wp_shared_mapping_range(vbo->tbo.bdev->dev_mapping,
                                             dirty->start + offset,
                                             dirty->end - dirty->start);

        if (100UL * num_marked / dirty->bitmap_size >
            VMW_DIRTY_PERCENTAGE)
                dirty->change_count++;
        else
                dirty->change_count = 0;

        if (dirty->change_count > VMW_DIRTY_NUM_CHANGE_TRIGGERS) {
                pgoff_t start = 0;
                pgoff_t end = dirty->bitmap_size;

                dirty->method = VMW_BO_DIRTY_PAGETABLE;
                clean_record_shared_mapping_range(mapping, offset, end, offset,
                                                  &dirty->bitmap[0],
                                                  &start, &end);
                bitmap_clear(&dirty->bitmap[0], 0, dirty->bitmap_size);
                if (dirty->start < dirty->end)
                        bitmap_set(&dirty->bitmap[0], dirty->start,
                                   dirty->end - dirty->start);
                dirty->change_count = 0;
        }
}

/**
 * vmw_bo_dirty_scan - Scan for dirty pages and add them to the dirty
 * tracking structure
 * @vbo: The buffer object to scan
 *
 * This function may change the dirty tracking method.
 */
void vmw_bo_dirty_scan(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;

        if (dirty->method == VMW_BO_DIRTY_PAGETABLE)
                vmw_bo_dirty_scan_pagetable(vbo);
        else
                vmw_bo_dirty_scan_mkwrite(vbo);
}

/**
 * vmw_bo_dirty_pre_unmap - write-protect and pick up dirty pages before
 * an unmap_mapping_range operation.
 * @vbo: The buffer object,
 * @start: First page of the range within the buffer object.
 * @end: Last page of the range within the buffer object + 1.
 *
 * If we're using the _PAGETABLE scan method, we may leak dirty pages
 * when calling unmap_mapping_range(). This function makes sure we pick
 * up all dirty pages.
 */
static void vmw_bo_dirty_pre_unmap(struct vmw_bo *vbo,
                                   pgoff_t start, pgoff_t end)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;
        unsigned long offset = drm_vma_node_start(&vbo->tbo.base.vma_node);
        struct address_space *mapping = vbo->tbo.bdev->dev_mapping;

        if (dirty->method != VMW_BO_DIRTY_PAGETABLE || start >= end)
                return;

        wp_shared_mapping_range(mapping, start + offset, end - start);
        clean_record_shared_mapping_range(mapping, start + offset,
                                          end - start, offset,
                                          &dirty->bitmap[0], &dirty->start,
                                          &dirty->end);
}

/**
 * vmw_bo_dirty_unmap - Clear all ptes pointing to a range within a bo
 * @vbo: The buffer object,
 * @start: First page of the range within the buffer object.
 * @end: Last page of the range within the buffer object + 1.
 *
 * This is similar to ttm_bo_unmap_virtual() except it takes a subrange.
 */
void vmw_bo_dirty_unmap(struct vmw_bo *vbo,
                        pgoff_t start, pgoff_t end)
{
        unsigned long offset = drm_vma_node_start(&vbo->tbo.base.vma_node);
        struct address_space *mapping = vbo->tbo.bdev->dev_mapping;

        vmw_bo_dirty_pre_unmap(vbo, start, end);
        unmap_shared_mapping_range(mapping, (offset + start) << PAGE_SHIFT,
                                   (loff_t) (end - start) << PAGE_SHIFT);
}

/**
 * vmw_bo_dirty_add - Add a dirty-tracking user to a buffer object
 * @vbo: The buffer object
 *
 * This function registers a dirty-tracking user to a buffer object.
 * A user can be for example a resource or a vma in a special user-space
 * mapping.
 *
 * Return: Zero on success, -ENOMEM on memory allocation failure.
 */
int vmw_bo_dirty_add(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;
        pgoff_t num_pages = PFN_UP(vbo->tbo.resource->size);
        size_t size;
        int ret;

        if (dirty) {
                kref_get(&dirty->ref_count);
                return 0;
        }

        size = sizeof(*dirty) + BITS_TO_LONGS(num_pages) * sizeof(long);
        dirty = kvzalloc(size, GFP_KERNEL);
        if (!dirty) {
                ret = -ENOMEM;
                goto out_no_dirty;
        }

        dirty->bitmap_size = num_pages;
        dirty->start = dirty->bitmap_size;
        dirty->end = 0;
        kref_init(&dirty->ref_count);
        if (num_pages < PAGE_SIZE / sizeof(pte_t)) {
                dirty->method = VMW_BO_DIRTY_PAGETABLE;
        } else {
                struct address_space *mapping = vbo->tbo.bdev->dev_mapping;
                pgoff_t offset = drm_vma_node_start(&vbo->tbo.base.vma_node);

                dirty->method = VMW_BO_DIRTY_MKWRITE;

                /* Write-protect and then pick up already dirty bits */
                wp_shared_mapping_range(mapping, offset, num_pages);
                clean_record_shared_mapping_range(mapping, offset, num_pages,
                                                  offset,
                                                  &dirty->bitmap[0],
                                                  &dirty->start, &dirty->end);
        }

        vbo->dirty = dirty;

        return 0;

out_no_dirty:
        return ret;
}

static void vmw_bo_dirty_free(struct kref *kref)
{
        struct vmw_bo_dirty *dirty = container_of(kref, struct vmw_bo_dirty, ref_count);

        kvfree(dirty);
}

/**
 * vmw_bo_dirty_release - Release a dirty-tracking user from a buffer object
 * @vbo: The buffer object
 *
 * This function releases a dirty-tracking user from a buffer object.
 * If the reference count reaches zero, then the dirty-tracking object is
 * freed and the pointer to it cleared.
 *
 * Return: Zero on success, -ENOMEM on memory allocation failure.
 */
void vmw_bo_dirty_release(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;

        if (dirty && kref_put(&dirty->ref_count, vmw_bo_dirty_free))
                vbo->dirty = NULL;
}

/**
 * vmw_bo_dirty_transfer_to_res - Pick up a resource's dirty region from
 * its backing mob.
 * @res: The resource
 *
 * This function will pick up all dirty ranges affecting the resource from
 * it's backup mob, and call vmw_resource_dirty_update() once for each
 * range. The transferred ranges will be cleared from the backing mob's
 * dirty tracking.
 */
void vmw_bo_dirty_transfer_to_res(struct vmw_resource *res)
{
        struct vmw_bo *vbo = res->guest_memory_bo;
        struct vmw_bo_dirty *dirty = vbo->dirty;
        pgoff_t start, cur, end;
        unsigned long res_start = res->guest_memory_offset;
        unsigned long res_end = res->guest_memory_offset + res->guest_memory_size;

        WARN_ON_ONCE(res_start & ~PAGE_MASK);
        res_start >>= PAGE_SHIFT;
        res_end = DIV_ROUND_UP(res_end, PAGE_SIZE);

        if (res_start >= dirty->end || res_end <= dirty->start)
                return;

        cur = max(res_start, dirty->start);
        res_end = max(res_end, dirty->end);
        while (cur < res_end) {
                unsigned long num;

                start = find_next_bit(&dirty->bitmap[0], res_end, cur);
                if (start >= res_end)
                        break;

                end = find_next_zero_bit(&dirty->bitmap[0], res_end, start + 1);
                cur = end + 1;
                num = end - start;
                bitmap_clear(&dirty->bitmap[0], start, num);
                vmw_resource_dirty_update(res, start, end);
        }

        if (res_start <= dirty->start && res_end > dirty->start)
                dirty->start = res_end;
        if (res_start < dirty->end && res_end >= dirty->end)
                dirty->end = res_start;
}

void vmw_bo_dirty_clear(struct vmw_bo *vbo)
{
        struct vmw_bo_dirty *dirty = vbo->dirty;
        pgoff_t start, cur, end;
        unsigned long res_start = 0;
        unsigned long res_end = vbo->tbo.base.size;

        WARN_ON_ONCE(res_start & ~PAGE_MASK);
        res_start >>= PAGE_SHIFT;
        res_end = DIV_ROUND_UP(res_end, PAGE_SIZE);

        if (res_start >= dirty->end || res_end <= dirty->start)
                return;

        cur = max(res_start, dirty->start);
        res_end = max(res_end, dirty->end);
        while (cur < res_end) {
                unsigned long num;

                start = find_next_bit(&dirty->bitmap[0], res_end, cur);
                if (start >= res_end)
                        break;

                end = find_next_zero_bit(&dirty->bitmap[0], res_end, start + 1);
                cur = end + 1;
                num = end - start;
                bitmap_clear(&dirty->bitmap[0], start, num);
        }

        if (res_start <= dirty->start && res_end > dirty->start)
                dirty->start = res_end;
        if (res_start < dirty->end && res_end >= dirty->end)
                dirty->end = res_start;
}

/**
 * vmw_bo_dirty_clear_res - Clear a resource's dirty region from
 * its backing mob.
 * @res: The resource
 *
 * This function will clear all dirty ranges affecting the resource from
 * it's backup mob's dirty tracking.
 */
void vmw_bo_dirty_clear_res(struct vmw_resource *res)
{
        unsigned long res_start = res->guest_memory_offset;
        unsigned long res_end = res->guest_memory_offset + res->guest_memory_size;
        struct vmw_bo *vbo = res->guest_memory_bo;
        struct vmw_bo_dirty *dirty = vbo->dirty;

        res_start >>= PAGE_SHIFT;
        res_end = DIV_ROUND_UP(res_end, PAGE_SIZE);

        if (res_start >= dirty->end || res_end <= dirty->start)
                return;

        res_start = max(res_start, dirty->start);
        res_end = min(res_end, dirty->end);
        bitmap_clear(&dirty->bitmap[0], res_start, res_end - res_start);

        if (res_start <= dirty->start && res_end > dirty->start)
                dirty->start = res_end;
        if (res_start < dirty->end && res_end >= dirty->end)
                dirty->end = res_start;
}

vm_fault_t vmw_bo_vm_mkwrite(struct vm_fault *vmf)
{
        struct vm_area_struct *vma = vmf->vma;
        struct ttm_buffer_object *bo = (struct ttm_buffer_object *)
            vma->vm_private_data;
        vm_fault_t ret;
        unsigned long page_offset;
        unsigned int save_flags;
        struct vmw_bo *vbo = to_vmw_bo(&bo->base);

        /*
         * mkwrite() doesn't handle the VM_FAULT_RETRY return value correctly.
         * So make sure the TTM helpers are aware.
         */
        save_flags = vmf->flags;
        vmf->flags &= ~FAULT_FLAG_ALLOW_RETRY;
        ret = ttm_bo_vm_reserve(bo, vmf);
        vmf->flags = save_flags;
        if (ret)
                return ret;

        page_offset = vmf->pgoff - drm_vma_node_start(&bo->base.vma_node);
        if (unlikely(page_offset >= PFN_UP(bo->resource->size))) {
                ret = VM_FAULT_SIGBUS;
                goto out_unlock;
        }

        if (vbo->dirty && vbo->dirty->method == VMW_BO_DIRTY_MKWRITE &&
            !test_bit(page_offset, &vbo->dirty->bitmap[0])) {
                struct vmw_bo_dirty *dirty = vbo->dirty;

                __set_bit(page_offset, &dirty->bitmap[0]);
                dirty->start = min(dirty->start, page_offset);
                dirty->end = max(dirty->end, page_offset + 1);
        }

out_unlock:
        dma_resv_unlock(bo->base.resv);
        return ret;
}

vm_fault_t vmw_bo_vm_fault(struct vm_fault *vmf)
{
        struct vm_area_struct *vma = vmf->vma;
        struct ttm_buffer_object *bo = (struct ttm_buffer_object *)
            vma->vm_private_data;
        struct vmw_bo *vbo = to_vmw_bo(&bo->base);
        pgoff_t num_prefault;
        pgprot_t prot;
        vm_fault_t ret;

        ret = ttm_bo_vm_reserve(bo, vmf);
        if (ret)
                return ret;

        num_prefault = (vma->vm_flags & VM_RAND_READ) ? 1 :
                TTM_BO_VM_NUM_PREFAULT;

        if (vbo->dirty) {
                pgoff_t allowed_prefault;
                unsigned long page_offset;

                page_offset = vmf->pgoff -
                        drm_vma_node_start(&bo->base.vma_node);
                if (page_offset >= PFN_UP(bo->resource->size) ||
                    vmw_resources_clean(vbo, page_offset,
                                        page_offset + PAGE_SIZE,
                                        &allowed_prefault)) {
                        ret = VM_FAULT_SIGBUS;
                        goto out_unlock;
                }

                num_prefault = min(num_prefault, allowed_prefault);
        }

        /*
         * If we don't track dirty using the MKWRITE method, make sure
         * sure the page protection is write-enabled so we don't get
         * a lot of unnecessary write faults.
         */
        if (vbo->dirty && vbo->dirty->method == VMW_BO_DIRTY_MKWRITE)
                prot = vm_get_page_prot(vma->vm_flags & ~VM_SHARED);
        else
                prot = vm_get_page_prot(vma->vm_flags);

        ret = ttm_bo_vm_fault_reserved(vmf, prot, num_prefault);
        if (ret == VM_FAULT_RETRY && !(vmf->flags & FAULT_FLAG_RETRY_NOWAIT))
                return ret;

out_unlock:
        dma_resv_unlock(bo->base.resv);

        return ret;
}