// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2026 Meta.  All rights reserved.
 */

#include <linux/sizes.h>
#include "btrfs-tests.h"
#include "../volumes.h"
#include "../disk-io.h"
#include "../extent-io-tree.h"

/*
 * Tests for chunk allocator pending extent internals.
 * These two functions form the core of searching the chunk allocation pending
 * extent bitmap and have relatively easily definable semantics, so unit
 * testing them can help ensure the correctness of chunk allocation.
 */

/*
 * Describes the inputs to the system and expected results
 * when testing btrfs_find_hole_in_pending_extents().
 */
struct pending_extent_test_case {
        const char *name;
        /* Input range to search. */
        u64 hole_start;
        u64 hole_len;
        /* The size of hole we are searching for. */
        u64 min_hole_size;
        /*
         * Pending extents to set up (up to 2 for up to 3 holes)
         * If len == 0, then it is skipped.
         */
        struct {
                u64 start;
                u64 len;
        } pending_extents[2];
        /* Expected outputs. */
        bool expected_found;
        u64 expected_start;
        u64 expected_len;
};

static const struct pending_extent_test_case find_hole_tests[] = {
        {
                .name = "no pending extents",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = { },
                .expected_found = true,
                .expected_start = 0,
                .expected_len = 10ULL * SZ_1G,
        },
        {
                .name = "pending extent at start of range",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = 0, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = SZ_1G,
                .expected_len = 9ULL * SZ_1G,
        },
        {
                .name = "pending extent overlapping start of range",
                .hole_start = SZ_1G,
                .hole_len = 9ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = 0, .len = SZ_2G },
                },
                .expected_found = true,
                .expected_start = SZ_2G,
                .expected_len = 8ULL * SZ_1G,
        },
        {
                .name = "two holes; first hole is exactly big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = 0,
                .expected_len = SZ_1G,
        },
        {
                .name = "two holes; first hole is big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = SZ_2G, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = 0,
                .expected_len = SZ_2G,
        },
        {
                .name = "two holes; second hole is big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_2G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = SZ_2G,
                .expected_len = 8ULL * SZ_1G,
        },
        {
                .name = "three holes; first hole big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_2G,
                .pending_extents = {
                        { .start = SZ_2G, .len = SZ_1G },
                        { .start = 4ULL * SZ_1G, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = 0,
                .expected_len = SZ_2G,
        },
        {
                .name = "three holes; second hole big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_2G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                        { .start = 5ULL * SZ_1G, .len = SZ_1G },
                },
                .expected_found = true,
                .expected_start = SZ_2G,
                .expected_len = 3ULL * SZ_1G,
        },
        {
                .name = "three holes; third hole big enough",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_2G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                        { .start = 3ULL * SZ_1G, .len = 5ULL * SZ_1G },
                },
                .expected_found = true,
                .expected_start = 8ULL * SZ_1G,
                .expected_len = SZ_2G,
        },
        {
                .name = "three holes; all holes too small",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_2G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                        { .start = 3ULL * SZ_1G, .len = 6ULL * SZ_1G },
                },
                .expected_found = false,
                .expected_start = 0,
                .expected_len = SZ_1G,
        },
        {
                .name = "three holes; all holes too small; first biggest",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = 3ULL * SZ_1G,
                .pending_extents = {
                        { .start = SZ_2G, .len = SZ_1G },
                        { .start = 4ULL * SZ_1G, .len = 5ULL * SZ_1G },
                },
                .expected_found = false,
                .expected_start = 0,
                .expected_len = SZ_2G,
        },
        {
                .name = "three holes; all holes too small; second biggest",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = 3ULL * SZ_1G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                        { .start = 4ULL * SZ_1G, .len = 5ULL * SZ_1G },
                },
                .expected_found = false,
                .expected_start = SZ_2G,
                .expected_len = SZ_2G,
        },
        {
                .name = "three holes; all holes too small; third biggest",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = 3ULL * SZ_1G,
                .pending_extents = {
                        { .start = SZ_1G, .len = SZ_1G },
                        { .start = 3ULL * SZ_1G, .len = 5ULL * SZ_1G },
                },
                .expected_found = false,
                .expected_start = 8ULL * SZ_1G,
                .expected_len = SZ_2G,
        },
        {
                .name = "hole entirely allocated by pending",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = 0, .len = 10ULL * SZ_1G },
                },
                .expected_found = false,
                .expected_start = 10ULL * SZ_1G,
                .expected_len = 0,
        },
        {
                .name = "pending extent at end of range",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .min_hole_size = SZ_1G,
                .pending_extents = {
                        { .start = 9ULL * SZ_1G, .len = SZ_2G },
                },
                .expected_found = true,
                .expected_start = 0,
                .expected_len = 9ULL * SZ_1G,
        },
        {
                .name = "zero length input",
                .hole_start = SZ_1G,
                .hole_len = 0,
                .min_hole_size = SZ_1G,
                .pending_extents = { },
                .expected_found = false,
                .expected_start = SZ_1G,
                .expected_len = 0,
        },
};

static int test_find_hole_in_pending(u32 sectorsize, u32 nodesize)
{
        struct btrfs_fs_info *fs_info;
        struct btrfs_device *device;
        int ret = 0;

        test_msg("running find_hole_in_pending_extents tests");

        fs_info = btrfs_alloc_dummy_fs_info(nodesize, sectorsize);
        if (!fs_info) {
                test_std_err(TEST_ALLOC_FS_INFO);
                return -ENOMEM;
        }

        device = btrfs_alloc_dummy_device(fs_info);
        if (IS_ERR(device)) {
                test_err("failed to allocate dummy device");
                ret = PTR_ERR(device);
                goto out_free_fs_info;
        }
        device->fs_info = fs_info;

        for (int i = 0; i < ARRAY_SIZE(find_hole_tests); i++) {
                const struct pending_extent_test_case *test_case = &find_hole_tests[i];
                u64 hole_start = test_case->hole_start;
                u64 hole_len = test_case->hole_len;
                bool found;

                for (int j = 0; j < ARRAY_SIZE(test_case->pending_extents); j++) {
                        u64 start = test_case->pending_extents[j].start;
                        u64 len = test_case->pending_extents[j].len;

                        if (!len)
                                continue;
                        btrfs_set_extent_bit(&device->alloc_state,
                                             start, start + len - 1,
                                             CHUNK_ALLOCATED, NULL);
                }

                mutex_lock(&fs_info->chunk_mutex);
                found = btrfs_find_hole_in_pending_extents(device, &hole_start, &hole_len,
                                                           test_case->min_hole_size);
                mutex_unlock(&fs_info->chunk_mutex);

                if (found != test_case->expected_found) {
                        test_err("%s: expected found=%d, got found=%d",
                                 test_case->name, test_case->expected_found, found);
                        ret = -EINVAL;
                        goto out_clear_pending_extents;
                }
                if (hole_start != test_case->expected_start ||
                    hole_len != test_case->expected_len) {
                        test_err("%s: expected [%llu, %llu), got [%llu, %llu)",
                                 test_case->name, test_case->expected_start,
                                 test_case->expected_start +
                                         test_case->expected_len,
                                 hole_start, hole_start + hole_len);
                        ret = -EINVAL;
                        goto out_clear_pending_extents;
                }
out_clear_pending_extents:
                btrfs_clear_extent_bit(&device->alloc_state, 0, (u64)-1,
                                       CHUNK_ALLOCATED, NULL);
                if (ret)
                        break;
        }

out_free_fs_info:
        btrfs_free_dummy_fs_info(fs_info);
        return ret;
}

/*
 * Describes the inputs to the system and expected results
 * when testing btrfs_first_pending_extent().
 */
struct first_pending_test_case {
        const char *name;
        /* The range to look for a pending extent in. */
        u64 hole_start;
        u64 hole_len;
        /* The pending extent to look for. */
        struct {
                u64 start;
                u64 len;
        } pending_extent;
        /* Expected outputs. */
        bool expected_found;
        u64 expected_pending_start;
        u64 expected_pending_end;
};

static const struct first_pending_test_case first_pending_tests[] = {
        {
                .name = "no pending extent",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .pending_extent = { 0, 0 },
                .expected_found = false,
        },
        {
                .name = "pending extent at search start",
                .hole_start = SZ_1G,
                .hole_len = 9ULL * SZ_1G,
                .pending_extent = { SZ_1G, SZ_1G },
                .expected_found = true,
                .expected_pending_start = SZ_1G,
                .expected_pending_end = SZ_2G - 1,
        },
        {
                .name = "pending extent overlapping search start",
                .hole_start = SZ_1G,
                .hole_len = 9ULL * SZ_1G,
                .pending_extent = { 0, SZ_2G },
                .expected_found = true,
                .expected_pending_start = 0,
                .expected_pending_end = SZ_2G - 1,
        },
        {
                .name = "pending extent inside search range",
                .hole_start = 0,
                .hole_len = 10ULL * SZ_1G,
                .pending_extent = { SZ_2G, SZ_1G },
                .expected_found = true,
                .expected_pending_start = SZ_2G,
                .expected_pending_end = 3ULL * SZ_1G - 1,
        },
        {
                .name = "pending extent outside search range",
                .hole_start = 0,
                .hole_len = SZ_1G,
                .pending_extent = { SZ_2G, SZ_1G },
                .expected_found = false,
        },
        {
                .name = "pending extent overlapping end of search range",
                .hole_start = 0,
                .hole_len = SZ_2G,
                .pending_extent = { SZ_1G, SZ_2G },
                .expected_found = true,
                .expected_pending_start = SZ_1G,
                .expected_pending_end = 3ULL * SZ_1G - 1,
        },
};

static int test_first_pending_extent(u32 sectorsize, u32 nodesize)
{
        struct btrfs_fs_info *fs_info;
        struct btrfs_device *device;
        int ret = 0;

        test_msg("running first_pending_extent tests");

        fs_info = btrfs_alloc_dummy_fs_info(nodesize, sectorsize);
        if (!fs_info) {
                test_std_err(TEST_ALLOC_FS_INFO);
                return -ENOMEM;
        }

        device = btrfs_alloc_dummy_device(fs_info);
        if (IS_ERR(device)) {
                test_err("failed to allocate dummy device");
                ret = PTR_ERR(device);
                goto out_free_fs_info;
        }

        device->fs_info = fs_info;

        for (int i = 0; i < ARRAY_SIZE(first_pending_tests); i++) {
                const struct first_pending_test_case *test_case = &first_pending_tests[i];
                u64 start = test_case->pending_extent.start;
                u64 len = test_case->pending_extent.len;
                u64 pending_start, pending_end;
                bool found;

                if (len) {
                        btrfs_set_extent_bit(&device->alloc_state,
                                             start, start + len - 1,
                                             CHUNK_ALLOCATED, NULL);
                }

                mutex_lock(&fs_info->chunk_mutex);
                found = btrfs_first_pending_extent(device, test_case->hole_start,
                                                   test_case->hole_len,
                                                   &pending_start, &pending_end);
                mutex_unlock(&fs_info->chunk_mutex);

                if (found != test_case->expected_found) {
                        test_err("%s: expected found=%d, got found=%d",
                                 test_case->name, test_case->expected_found, found);
                        ret = -EINVAL;
                        goto out_clear_pending_extents;
                }
                if (!found)
                        goto out_clear_pending_extents;

                if (pending_start != test_case->expected_pending_start ||
                    pending_end != test_case->expected_pending_end) {
                        test_err("%s: expected pending [%llu, %llu], got [%llu, %llu]",
                                 test_case->name,
                                 test_case->expected_pending_start,
                                 test_case->expected_pending_end,
                                 pending_start, pending_end);
                        ret = -EINVAL;
                        goto out_clear_pending_extents;
                }

out_clear_pending_extents:
                btrfs_clear_extent_bit(&device->alloc_state, 0, (u64)-1,
                                       CHUNK_ALLOCATED, NULL);
                if (ret)
                        break;
        }

out_free_fs_info:
        btrfs_free_dummy_fs_info(fs_info);
        return ret;
}

int btrfs_test_chunk_allocation(u32 sectorsize, u32 nodesize)
{
        int ret;

        test_msg("running chunk allocation tests");

        ret = test_first_pending_extent(sectorsize, nodesize);
        if (ret)
                return ret;

        ret = test_find_hole_in_pending(sectorsize, nodesize);
        if (ret)
                return ret;

        return 0;
}