// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2025 - Google LLC
 * Author: Vincent Donnefort <vdonnefort@google.com>
 */

#include <linux/atomic.h>
#include <linux/simple_ring_buffer.h>

#include <asm/barrier.h>
#include <asm/local.h>

enum simple_rb_link_type {
        SIMPLE_RB_LINK_NORMAL           = 0,
        SIMPLE_RB_LINK_HEAD             = 1,
        SIMPLE_RB_LINK_HEAD_MOVING
};

#define SIMPLE_RB_LINK_MASK ~(SIMPLE_RB_LINK_HEAD | SIMPLE_RB_LINK_HEAD_MOVING)

static void simple_bpage_set_head_link(struct simple_buffer_page *bpage)
{
        unsigned long link = (unsigned long)bpage->link.next;

        link &= SIMPLE_RB_LINK_MASK;
        link |= SIMPLE_RB_LINK_HEAD;

        /*
         * Paired with simple_rb_find_head() to order access between the head
         * link and overrun. It ensures we always report an up-to-date value
         * after swapping the reader page.
         */
        smp_store_release(&bpage->link.next, (struct list_head *)link);
}

static bool simple_bpage_unset_head_link(struct simple_buffer_page *bpage,
                                         struct simple_buffer_page *dst,
                                         enum simple_rb_link_type new_type)
{
        unsigned long *link = (unsigned long *)(&bpage->link.next);
        unsigned long old = (*link & SIMPLE_RB_LINK_MASK) | SIMPLE_RB_LINK_HEAD;
        unsigned long new = (unsigned long)(&dst->link) | new_type;

        return try_cmpxchg(link, &old, new);
}

static void simple_bpage_set_normal_link(struct simple_buffer_page *bpage)
{
        unsigned long link = (unsigned long)bpage->link.next;

        WRITE_ONCE(bpage->link.next, (struct list_head *)(link & SIMPLE_RB_LINK_MASK));
}

static struct simple_buffer_page *simple_bpage_from_link(struct list_head *link)
{
        unsigned long ptr = (unsigned long)link & SIMPLE_RB_LINK_MASK;

        return container_of((struct list_head *)ptr, struct simple_buffer_page, link);
}

static struct simple_buffer_page *simple_bpage_next_page(struct simple_buffer_page *bpage)
{
        return simple_bpage_from_link(bpage->link.next);
}

static void simple_bpage_reset(struct simple_buffer_page *bpage)
{
        bpage->write = 0;
        bpage->entries = 0;

        local_set(&bpage->page->commit, 0);
}

static void simple_bpage_init(struct simple_buffer_page *bpage, void *page)
{
        INIT_LIST_HEAD(&bpage->link);
        bpage->page = (struct buffer_data_page *)page;

        simple_bpage_reset(bpage);
}

#define simple_rb_meta_inc(__meta, __inc)               \
        WRITE_ONCE((__meta), (__meta + __inc))

static bool simple_rb_loaded(struct simple_rb_per_cpu *cpu_buffer)
{
        return !!cpu_buffer->bpages;
}

static int simple_rb_find_head(struct simple_rb_per_cpu *cpu_buffer)
{
        int retry = cpu_buffer->nr_pages * 2;
        struct simple_buffer_page *head;

        head = cpu_buffer->head_page;

        while (retry--) {
                unsigned long link;

spin:
                /* See smp_store_release in simple_bpage_set_head_link() */
                link = (unsigned long)smp_load_acquire(&head->link.prev->next);

                switch (link & ~SIMPLE_RB_LINK_MASK) {
                /* Found the head */
                case SIMPLE_RB_LINK_HEAD:
                        cpu_buffer->head_page = head;
                        return 0;
                /* The writer caught the head, we can spin, that won't be long */
                case SIMPLE_RB_LINK_HEAD_MOVING:
                        goto spin;
                }

                head = simple_bpage_next_page(head);
        }

        return -EBUSY;
}

/**
 * simple_ring_buffer_swap_reader_page - Swap ring-buffer head with the reader
 * @cpu_buffer: A simple_rb_per_cpu
 *
 * This function enables consuming reading. It ensures the current head page will not be overwritten
 * and can be safely read.
 *
 * Returns 0 on success, -ENODEV if @cpu_buffer was unloaded or -EBUSY if we failed to catch the
 * head page.
 */
int simple_ring_buffer_swap_reader_page(struct simple_rb_per_cpu *cpu_buffer)
{
        struct simple_buffer_page *last, *head, *reader;
        unsigned long overrun;
        int retry = 8;
        int ret;

        if (!simple_rb_loaded(cpu_buffer))
                return -ENODEV;

        reader = cpu_buffer->reader_page;

        do {
                /* Run after the writer to find the head */
                ret = simple_rb_find_head(cpu_buffer);
                if (ret)
                        return ret;

                head = cpu_buffer->head_page;

                /* Connect the reader page around the header page */
                reader->link.next = head->link.next;
                reader->link.prev = head->link.prev;

                /* The last page before the head */
                last = simple_bpage_from_link(head->link.prev);

                /* The reader page points to the new header page */
                simple_bpage_set_head_link(reader);

                overrun = cpu_buffer->meta->overrun;
        } while (!simple_bpage_unset_head_link(last, reader, SIMPLE_RB_LINK_NORMAL) && retry--);

        if (!retry)
                return -EINVAL;

        cpu_buffer->head_page = simple_bpage_from_link(reader->link.next);
        cpu_buffer->head_page->link.prev = &reader->link;
        cpu_buffer->reader_page = head;
        cpu_buffer->meta->reader.lost_events = overrun - cpu_buffer->last_overrun;
        cpu_buffer->meta->reader.id = cpu_buffer->reader_page->id;
        cpu_buffer->last_overrun = overrun;

        return 0;
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_swap_reader_page);

static struct simple_buffer_page *simple_rb_move_tail(struct simple_rb_per_cpu *cpu_buffer)
{
        struct simple_buffer_page *tail, *new_tail;

        tail = cpu_buffer->tail_page;
        new_tail = simple_bpage_next_page(tail);

        if (simple_bpage_unset_head_link(tail, new_tail, SIMPLE_RB_LINK_HEAD_MOVING)) {
                /*
                 * Oh no! we've caught the head. There is none anymore and
                 * swap_reader will spin until we set the new one. Overrun must
                 * be written first, to make sure we report the correct number
                 * of lost events.
                 */
                simple_rb_meta_inc(cpu_buffer->meta->overrun, new_tail->entries);
                simple_rb_meta_inc(cpu_buffer->meta->pages_lost, 1);

                simple_bpage_set_head_link(new_tail);
                simple_bpage_set_normal_link(tail);
        }

        simple_bpage_reset(new_tail);
        cpu_buffer->tail_page = new_tail;

        simple_rb_meta_inc(cpu_buffer->meta->pages_touched, 1);

        return new_tail;
}

static unsigned long rb_event_size(unsigned long length)
{
        struct ring_buffer_event *event;

        return length + RB_EVNT_HDR_SIZE + sizeof(event->array[0]);
}

static struct ring_buffer_event *
rb_event_add_ts_extend(struct ring_buffer_event *event, u64 delta)
{
        event->type_len = RINGBUF_TYPE_TIME_EXTEND;
        event->time_delta = delta & TS_MASK;
        event->array[0] = delta >> TS_SHIFT;

        return (struct ring_buffer_event *)((unsigned long)event + 8);
}

static struct ring_buffer_event *
simple_rb_reserve_next(struct simple_rb_per_cpu *cpu_buffer, unsigned long length, u64 timestamp)
{
        unsigned long ts_ext_size = 0, event_size = rb_event_size(length);
        struct simple_buffer_page *tail = cpu_buffer->tail_page;
        struct ring_buffer_event *event;
        u32 write, prev_write;
        u64 time_delta;

        time_delta = timestamp - cpu_buffer->write_stamp;

        if (test_time_stamp(time_delta))
                ts_ext_size = 8;

        prev_write = tail->write;
        write = prev_write + event_size + ts_ext_size;

        if (unlikely(write > (PAGE_SIZE - BUF_PAGE_HDR_SIZE)))
                tail = simple_rb_move_tail(cpu_buffer);

        if (!tail->entries) {
                tail->page->time_stamp = timestamp;
                time_delta = 0;
                ts_ext_size = 0;
                write = event_size;
                prev_write = 0;
        }

        tail->write = write;
        tail->entries++;

        cpu_buffer->write_stamp = timestamp;

        event = (struct ring_buffer_event *)(tail->page->data + prev_write);
        if (ts_ext_size) {
                event = rb_event_add_ts_extend(event, time_delta);
                time_delta = 0;
        }

        event->type_len = 0;
        event->time_delta = time_delta;
        event->array[0] = event_size - RB_EVNT_HDR_SIZE;

        return event;
}

/**
 * simple_ring_buffer_reserve - Reserve an entry in @cpu_buffer
 * @cpu_buffer: A simple_rb_per_cpu
 * @length:     Size of the entry in bytes
 * @timestamp:  Timestamp of the entry
 *
 * Returns the address of the entry where to write data or NULL
 */
void *simple_ring_buffer_reserve(struct simple_rb_per_cpu *cpu_buffer, unsigned long length,
                                 u64 timestamp)
{
        struct ring_buffer_event *rb_event;

        if (cmpxchg(&cpu_buffer->status, SIMPLE_RB_READY, SIMPLE_RB_WRITING) != SIMPLE_RB_READY)
                return NULL;

        rb_event = simple_rb_reserve_next(cpu_buffer, length, timestamp);

        return &rb_event->array[1];
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_reserve);

/**
 * simple_ring_buffer_commit - Commit the entry reserved with simple_ring_buffer_reserve()
 * @cpu_buffer: The simple_rb_per_cpu where the entry has been reserved
 */
void simple_ring_buffer_commit(struct simple_rb_per_cpu *cpu_buffer)
{
        local_set(&cpu_buffer->tail_page->page->commit,
                  cpu_buffer->tail_page->write);
        simple_rb_meta_inc(cpu_buffer->meta->entries, 1);

        /*
         * Paired with simple_rb_enable_tracing() to ensure data is
         * written to the ring-buffer before teardown.
         */
        smp_store_release(&cpu_buffer->status, SIMPLE_RB_READY);
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_commit);

static u32 simple_rb_enable_tracing(struct simple_rb_per_cpu *cpu_buffer, bool enable)
{
        u32 prev_status;

        if (enable)
                return cmpxchg(&cpu_buffer->status, SIMPLE_RB_UNAVAILABLE, SIMPLE_RB_READY);

        /* Wait for the buffer to be released */
        do {
                prev_status = cmpxchg_acquire(&cpu_buffer->status,
                                              SIMPLE_RB_READY,
                                              SIMPLE_RB_UNAVAILABLE);
        } while (prev_status == SIMPLE_RB_WRITING);

        return prev_status;
}

/**
 * simple_ring_buffer_reset - Reset @cpu_buffer
 * @cpu_buffer: A simple_rb_per_cpu
 *
 * This will not clear the content of the data, only reset counters and pointers
 *
 * Returns 0 on success or -ENODEV if @cpu_buffer was unloaded.
 */
int simple_ring_buffer_reset(struct simple_rb_per_cpu *cpu_buffer)
{
        struct simple_buffer_page *bpage;
        u32 prev_status;
        int ret;

        if (!simple_rb_loaded(cpu_buffer))
                return -ENODEV;

        prev_status = simple_rb_enable_tracing(cpu_buffer, false);

        ret = simple_rb_find_head(cpu_buffer);
        if (ret)
                return ret;

        bpage = cpu_buffer->tail_page = cpu_buffer->head_page;
        do {
                simple_bpage_reset(bpage);
                bpage = simple_bpage_next_page(bpage);
        } while (bpage != cpu_buffer->head_page);

        simple_bpage_reset(cpu_buffer->reader_page);

        cpu_buffer->last_overrun = 0;
        cpu_buffer->write_stamp = 0;

        cpu_buffer->meta->reader.read = 0;
        cpu_buffer->meta->reader.lost_events = 0;
        cpu_buffer->meta->entries = 0;
        cpu_buffer->meta->overrun = 0;
        cpu_buffer->meta->read = 0;
        cpu_buffer->meta->pages_lost = 0;
        cpu_buffer->meta->pages_touched = 0;

        if (prev_status == SIMPLE_RB_READY)
                simple_rb_enable_tracing(cpu_buffer, true);

        return 0;
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_reset);

int simple_ring_buffer_init_mm(struct simple_rb_per_cpu *cpu_buffer,
                               struct simple_buffer_page *bpages,
                               const struct ring_buffer_desc *desc,
                               void *(*load_page)(unsigned long va),
                               void (*unload_page)(void *va))
{
        struct simple_buffer_page *bpage = bpages;
        int ret = 0;
        void *page;
        int i;

        /* At least 1 reader page and two pages in the ring-buffer */
        if (desc->nr_page_va < 3)
                return -EINVAL;

        memset(cpu_buffer, 0, sizeof(*cpu_buffer));

        cpu_buffer->meta = load_page(desc->meta_va);
        if (!cpu_buffer->meta)
                return -EINVAL;

        memset(cpu_buffer->meta, 0, sizeof(*cpu_buffer->meta));
        cpu_buffer->meta->meta_page_size = PAGE_SIZE;
        cpu_buffer->meta->nr_subbufs = cpu_buffer->nr_pages;

        /* The reader page is not part of the ring initially */
        page = load_page(desc->page_va[0]);
        if (!page) {
                unload_page(cpu_buffer->meta);
                return -EINVAL;
        }

        simple_bpage_init(bpage, page);
        bpage->id = 0;

        cpu_buffer->nr_pages = 1;

        cpu_buffer->reader_page = bpage;
        cpu_buffer->tail_page = bpage + 1;
        cpu_buffer->head_page = bpage + 1;

        for (i = 1; i < desc->nr_page_va; i++) {
                page = load_page(desc->page_va[i]);
                if (!page) {
                        ret = -EINVAL;
                        break;
                }

                simple_bpage_init(++bpage, page);

                bpage->link.next = &(bpage + 1)->link;
                bpage->link.prev = &(bpage - 1)->link;
                bpage->id = i;

                cpu_buffer->nr_pages = i + 1;
        }

        if (ret) {
                for (i--; i >= 0; i--)
                        unload_page((void *)desc->page_va[i]);
                unload_page(cpu_buffer->meta);

                return ret;
        }

        /* Close the ring */
        bpage->link.next = &cpu_buffer->tail_page->link;
        cpu_buffer->tail_page->link.prev = &bpage->link;

        /* The last init'ed page points to the head page */
        simple_bpage_set_head_link(bpage);

        cpu_buffer->bpages = bpages;

        return 0;
}

static void *__load_page(unsigned long page)
{
        return (void *)page;
}

static void __unload_page(void *page) { }

/**
 * simple_ring_buffer_init - Init @cpu_buffer based on @desc
 * @cpu_buffer: A simple_rb_per_cpu buffer to init, allocated by the caller.
 * @bpages:     Array of simple_buffer_pages, with as many elements as @desc->nr_page_va
 * @desc:       A ring_buffer_desc
 *
 * Returns 0 on success or -EINVAL if the content of @desc is invalid
 */
int simple_ring_buffer_init(struct simple_rb_per_cpu *cpu_buffer, struct simple_buffer_page *bpages,
                            const struct ring_buffer_desc *desc)
{
        return simple_ring_buffer_init_mm(cpu_buffer, bpages, desc, __load_page, __unload_page);
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_init);

void simple_ring_buffer_unload_mm(struct simple_rb_per_cpu *cpu_buffer,
                                  void (*unload_page)(void *))
{
        int p;

        if (!simple_rb_loaded(cpu_buffer))
                return;

        simple_rb_enable_tracing(cpu_buffer, false);

        unload_page(cpu_buffer->meta);
        for (p = 0; p < cpu_buffer->nr_pages; p++)
                unload_page(cpu_buffer->bpages[p].page);

        cpu_buffer->bpages = NULL;
}

/**
 * simple_ring_buffer_unload - Prepare @cpu_buffer for deletion
 * @cpu_buffer: A simple_rb_per_cpu that will be deleted.
 */
void simple_ring_buffer_unload(struct simple_rb_per_cpu *cpu_buffer)
{
        return simple_ring_buffer_unload_mm(cpu_buffer, __unload_page);
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_unload);

/**
 * simple_ring_buffer_enable_tracing - Enable or disable writing to @cpu_buffer
 * @cpu_buffer: A simple_rb_per_cpu
 * @enable:     True to enable tracing, False to disable it
 *
 * Returns 0 on success or -ENODEV if @cpu_buffer was unloaded
 */
int simple_ring_buffer_enable_tracing(struct simple_rb_per_cpu *cpu_buffer, bool enable)
{
        if (!simple_rb_loaded(cpu_buffer))
                return -ENODEV;

        simple_rb_enable_tracing(cpu_buffer, enable);

        return 0;
}
EXPORT_SYMBOL_GPL(simple_ring_buffer_enable_tracing);