/*-
 * Copyright (c) 2023, Netflix, Inc.
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include "stand.h"
#include "seg.h"

#include <sys/param.h>

static struct memory_segments *segs;
static int nr_seg = 0;
static int segalloc = 0;

void
init_avail(void)
{
        if (segs)
                free(segs);
        nr_seg = 0;
        segalloc = 16;
        free(segs);
        segs = malloc(sizeof(*segs) * segalloc);
        if (segs == NULL)
                panic("not enough memory to get memory map\n");
}

/*
 * Make sure at least n items can be accessed in the segs array.  Note the
 * realloc here will invalidate cached pointers (potentially), so addresses
 * into the segs array must be recomputed after this call.
 */
void
need_avail(int n)
{
        if (n <= segalloc)
                return;

        while (n > segalloc)
                segalloc *= 2;
        segs = realloc(segs, segalloc * sizeof(*segs));
        if (segs == NULL)
                panic("not enough memory to get memory map\n");
}

/*
 * Always called for a new range, so always just append a range,
 * unless it's continuous with the prior range.
 */
void
add_avail(uint64_t start, uint64_t end, uint64_t type)
{
        /*
         * This range is contiguous with the previous range, and is
         * the same type: we can collapse the two.
         */
        if (nr_seg >= 1 &&
            segs[nr_seg - 1].end + 1 == start &&
            segs[nr_seg - 1].type == type) {
                segs[nr_seg - 1].end = end;
                return;
        }

        /*
         * Otherwise we need to add a new range at the end, but don't need to
         * adjust the current end.
         */
        need_avail(nr_seg + 1);
        segs[nr_seg].start = start;
        segs[nr_seg].end = end;
        segs[nr_seg].type = type;
        nr_seg++;
}

/*
 * All or part of a prior entry needs to be modified. Given the structure of the
 * code, we know that it will always be modifying the last time and/or extending
 * the one before it if its contiguous.
 */
void
remove_avail(uint64_t start, uint64_t end, uint64_t type)
{
        struct memory_segments *s;

        /*
         * simple case: we are extending a previously removed item.
         */
        if (nr_seg >= 2) {
                s = &segs[nr_seg - 2];
                if (s->end + 1 == start &&
                    s->type == type) {
                        s->end = end;
                        /* Now adjust the ending element */
                        s++;
                        if (s->end == end) {
                                /* we've used up the 'free' space */
                                nr_seg--;
                                return;
                        }
                        /* Otherwise adjust the 'free' space */
                        s->start = end + 1;
                        return;
                }
        }

        /*
         * OK, we have four cases:
         * (1) The new chunk is at the start of the free space, but didn't catch the above
         *     folding for whatever reason (different type, start of space). In this case,
         *     we allocate 1 additional item. The current end is copied to the new end. The
         *     current end is set to <start, end, type> and the new end's start is set to end + 1.
         * (2) The new chunk is in the middle of the free space. In this case we allocate 2
         *     additional items. We copy the current end to the new end, set the new end's start
         *     to end + 1, the old end's end to start - 1 and the new item is <start, end, type>
         * (3) The new chunk is at the end of the current end. In this case we allocate 1 more
         *     and adjust the current end's end to start - 1 and set the new end to <start, end, type>.
         * (4) The new chunk is exactly the current end, except for type. In this case, we just adjust
         *     the type.
         * We can assume we always have at least one chunk since that's created with new_avail() above
         * necessarily before we are called to subset it.
         */
        s = &segs[nr_seg - 1];
        if (s->start == start) {
                if (s->end == end) { /* (4) */
                        s->type = type;
                        return;
                }
                /* chunk at start of old chunk -> (1) */
                need_avail(nr_seg + 1);
                s = &segs[nr_seg - 1];  /* Realloc may change pointers */
                s[1] = s[0];
                s->start = start;
                s->end = end;
                s->type = type;
                s[1].start = end + 1;
                nr_seg++;
                return;
        }
        if (s->end == end) {    /* At end of old chunk (3) */
                need_avail(nr_seg + 1);
                s = &segs[nr_seg - 1];  /* Realloc may change pointers */
                s[1] = s[0];
                s->end = start - 1;
                s[1].start = start;
                s[1].type = type;
                nr_seg++;
                return;
        }
        /* In the middle, need to split things up (2) */
        need_avail(nr_seg + 2);
        s = &segs[nr_seg - 1];  /* Realloc may change pointers */
        s[2] = s[1] = s[0];
        s->end = start - 1;
        s[1].start = start;
        s[1].end = end;
        s[1].type = type;
        s[2].start = end + 1;
        nr_seg += 2;
}

void
print_avail(void)
{
        printf("Found %d RAM segments:\n", nr_seg);

        for (int i = 0; i < nr_seg; i++) {
                printf("%#jx-%#jx type %lu\n",
                    (uintmax_t)segs[i].start,
                    (uintmax_t)segs[i].end,
                    (u_long)segs[i].type);
        }
}

uint64_t
first_avail(uint64_t align, uint64_t min_size, uint64_t memtype)
{
        uint64_t s, len;

        for (int i = 0; i < nr_seg; i++) {
                if (segs[i].type != memtype)    /* Not candidate */
                        continue;
                s = roundup(segs[i].start, align);
                if (s >= segs[i].end)           /* roundup past end */
                        continue;
                len = segs[i].end - s + 1;
                if (len >= min_size) {
                        printf("Found a big enough hole at in seg %d at %#jx (%#jx-%#jx)\n",
                            i,
                            (uintmax_t)s,
                            (uintmax_t)segs[i].start,
                            (uintmax_t)segs[i].end);
                        return (s);
                }
        }

        return (0);
}

enum types {
        system_ram = SYSTEM_RAM,
        firmware_reserved,
        linux_code,
        linux_data,
        linux_bss,
        unknown,
};

static struct kv
{
        uint64_t        type;
        char *          name;
        int             flags;
#define KV_KEEPER 1
} str2type_kv[] = {
        { linux_code,           "Kernel code", KV_KEEPER },
        { linux_data,           "Kernel data", KV_KEEPER },
        { linux_bss,            "Kernel bss", KV_KEEPER },
        { firmware_reserved,    "Reserved" },
};

static const char *
parse_line(const char *line, uint64_t *startp, uint64_t *endp)
{
        const char *walker;
        char *next;
        uint64_t start, end;

        /*
         * Each line is a range followed by a description of the form:
         * <hex-number><dash><hex-number><space><colon><space><string>
         * Bail if we have any parsing errors.
         */
        walker = line;
        start = strtoull(walker, &next, 16);
        if (start == ULLONG_MAX || walker == next)
                return (NULL);
        walker = next;
        if (*walker != '-')
                return (NULL);
        walker++;
        end = strtoull(walker, &next, 16);
        if (end == ULLONG_MAX || walker == next)
                return (NULL);
        walker = next;
        /* Now eat the ' : ' in front of the string we want to return */
        if (strncmp(walker, " : ", 3) != 0)
                return (NULL);
        *startp = start;
        *endp = end;
        return (walker + 3);
}

static struct kv *
kvlookup(const char *str, struct kv *kvs, size_t nkv)
{
        for (int i = 0; i < nkv; i++)
                if (strcmp(kvs[i].name, str) == 0)
                        return (&kvs[i]);

        return (NULL);
}

/* Trim trailing whitespace */
static void
chop(char *line)
{
        char *ep = line + strlen(line) - 1;

        while (ep >= line && isspace(*ep))
                *ep-- = '\0';
}

#define SYSTEM_RAM_STR "System RAM"
#define RESERVED "reserved"

bool
populate_avail_from_iomem(void)
{
        int fd;
        char buf[128];
        const char *str;
        uint64_t start, end;
        struct kv *kv;

        fd = open("host:/proc/iomem", O_RDONLY);
        if (fd == -1) {
                printf("Can't get memory map\n");
                init_avail();
                // Hack: 32G of RAM starting at 4G
                add_avail(4ull << 30, 36ull << 30, system_ram);
                return false;
        }

        if (fgetstr(buf, sizeof(buf), fd) < 0)
                goto out;       /* Nothing to do ???? */
        init_avail();
        chop(buf);
        while (true) {
                /*
                 * Look for top level items we understand.  Skip anything that's
                 * a continuation, since we don't care here. If we care, we'll
                 * consume them all when we recognize that top level item.
                 */
                if (buf[0] == ' ')      /* Continuation lines? Ignore */
                        goto next_line;
                str = parse_line(buf, &start, &end);
                if (str == NULL)        /* Malformed -> ignore */
                        goto next_line;
                /*
                 * All we care about is System RAM
                 */
                if (strncmp(str, SYSTEM_RAM_STR, sizeof(SYSTEM_RAM_STR) - 1) == 0)
                        add_avail(start, end, system_ram);
                else if (strncmp(str, RESERVED, sizeof(RESERVED) - 1) == 0)
                        add_avail(start, end, firmware_reserved);
                else
                        goto next_line; /* Ignore hardware */
                while (fgetstr(buf, sizeof(buf), fd) >= 0 && buf[0] == ' ') {
                        chop(buf);
                        str = parse_line(buf, &start, &end);
                        if (str == NULL)
                                break;
                        kv = kvlookup(str, str2type_kv, nitems(str2type_kv));
                        if (kv == NULL) /* failsafe for new types: igonre */
                                remove_avail(start, end, unknown);
                        else if ((kv->flags & KV_KEEPER) == 0)
                                remove_avail(start, end, kv->type);
                        /* Else no need to adjust since it's a keeper */
                }

                /*
                 * if buf[0] == ' ' then we know that the fgetstr failed and we
                 * should break. Otherwise fgetstr succeeded and we have a
                 * buffer we need to examine for being a top level item.
                 */
                if (buf[0] == ' ')
                        break;
                chop(buf);
                continue; /* buf has next top level line to parse */
next_line:
                if (fgetstr(buf, sizeof(buf), fd) < 0)
                        break;
        }

out:
        close(fd);
        return true;
}

/*
 * Return the amount of space available in the segment that @start@ lives in,
 * from @start@ to the end of the segment.
 */
uint64_t
space_avail(uint64_t start)
{
        for (int i = 0; i < nr_seg; i++) {
                if (start >= segs[i].start && start <= segs[i].end)
                        return segs[i].end - start;
        }

        /*
         * Properly used, we should never get here. Unsure if this should be a
         * panic or not.
         */
        return 0;
}