// SPDX-License-Identifier: GPL-2.0
#define pr_fmt(fmt) "%s() " fmt "\n", __func__

#include <linux/generic-radix-tree.h>
#include <linux/mm.h>
#include <linux/percpu.h>
#include <linux/slab.h>
#include <linux/srcu.h>
#include <linux/vmalloc.h>

#include "rcu_pending.h"
#include "darray.h"
#include "util.h"

#define static_array_for_each(_a, _i)   \
 for (typeof(&(_a)[0]) _i = _a;   \
      _i < (_a) + ARRAY_SIZE(_a);  \
      _i++)

enum rcu_pending_special {
 RCU_PENDING_KVFREE = 1,
 RCU_PENDING_CALL_RCU = 2,
};

#define RCU_PENDING_KVFREE_FN  ((rcu_pending_process_fn) (ulong) RCU_PENDING_KVFREE)
#define RCU_PENDING_CALL_RCU_FN  ((rcu_pending_process_fn) (ulong) RCU_PENDING_CALL_RCU)

#ifdef __KERNEL__
typedef unsigned long   rcu_gp_poll_state_t;

static inline bool rcu_gp_poll_cookie_eq(rcu_gp_poll_state_t l, rcu_gp_poll_state_t r)
{
 return l == r;
}
#else
typedef struct urcu_gp_poll_state rcu_gp_poll_state_t;

static inline bool rcu_gp_poll_cookie_eq(rcu_gp_poll_state_t l, rcu_gp_poll_state_t r)
{
 return l.grace_period_id == r.grace_period_id;
}
#endif

static inline rcu_gp_poll_state_t __get_state_synchronize_rcu(struct srcu_struct *ssp)
{
 return ssp
  ? get_state_synchronize_srcu(ssp)
  : get_state_synchronize_rcu();
}

static inline rcu_gp_poll_state_t __start_poll_synchronize_rcu(struct srcu_struct *ssp)
{
 return ssp
  ? start_poll_synchronize_srcu(ssp)
  : start_poll_synchronize_rcu();
}

static inline bool __poll_state_synchronize_rcu(struct srcu_struct *ssp, rcu_gp_poll_state_t cookie)
{
 return ssp
  ? poll_state_synchronize_srcu(ssp, cookie)
  : poll_state_synchronize_rcu(cookie);
}

static inline void __rcu_barrier(struct srcu_struct *ssp)
{
 return ssp
  ? srcu_barrier(ssp)
  : rcu_barrier();
}

static inline void __call_rcu(struct srcu_struct *ssp, struct rcu_head *rhp,
         rcu_callback_t func)
{
 if (ssp)
  call_srcu(ssp, rhp, func);
 else
  call_rcu(rhp, func);
}

struct rcu_pending_seq {
 /*
 * We're using a radix tree like a vector - we're just pushing elements
 * onto the end; we're using a radix tree instead of an actual vector to
 * avoid reallocation overhead
 */
 GENRADIX(struct rcu_head *) objs;
 size_t    nr;
 struct rcu_head   **cursor;
 rcu_gp_poll_state_t  seq;
};

struct rcu_pending_list {
 struct rcu_head   *head;
 struct rcu_head   *tail;
 rcu_gp_poll_state_t  seq;
};

struct rcu_pending_pcpu {
 struct rcu_pending  *parent;
 spinlock_t   lock;
 int    cpu;

 /*
 * We can't bound the number of unprocessed gp sequence numbers, and we
 * can't efficiently merge radix trees for expired grace periods, so we
 * need darray/vector:
 */
 DARRAY_PREALLOCATED(struct rcu_pending_seq, 4) objs;

 /* Third entry is for expired objects: */
 struct rcu_pending_list  lists[NUM_ACTIVE_RCU_POLL_OLDSTATE + 1];

 struct rcu_head   cb;
 bool    cb_armed;
 struct work_struct  work;
};

static bool __rcu_pending_has_pending(struct rcu_pending_pcpu *p)
{
 if (p->objs.nr)
  return true;

 static_array_for_each(p->lists, i)
  if (i->head)
   return true;

 return false;
}

static void rcu_pending_list_merge(struct rcu_pending_list *l1,
       struct rcu_pending_list *l2)
{
#ifdef __KERNEL__
 if (!l1->head)
  l1->head = l2->head;
 else
  l1->tail->next = l2->head;
#else
 if (!l1->head)
  l1->head = l2->head;
 else
  l1->tail->next.next = (void *) l2->head;
#endif

 l1->tail = l2->tail;
 l2->head = l2->tail = NULL;
}

static void rcu_pending_list_add(struct rcu_pending_list *l,
     struct rcu_head *n)
{
#ifdef __KERNEL__
 if (!l->head)
  l->head = n;
 else
  l->tail->next = n;
 l->tail = n;
 n->next = NULL;
#else
 if (!l->head)
  l->head = n;
 else
  l->tail->next.next = (void *) n;
 l->tail = n;
 n->next.next = NULL;
#endif
}

static void merge_expired_lists(struct rcu_pending_pcpu *p)
{
 struct rcu_pending_list *expired = &p->lists[NUM_ACTIVE_RCU_POLL_OLDSTATE];

 for (struct rcu_pending_list *i = p->lists; i < expired; i++)
  if (i->head && __poll_state_synchronize_rcu(p->parent->srcu, i->seq))
   rcu_pending_list_merge(expired, i);
}

#ifndef __KERNEL__
static inline void kfree_bulk(size_t nr, void ** p)
{
 while (nr--)
  kfree(*p);
}
#endif

static noinline void __process_finished_items(struct rcu_pending *pending,
           struct rcu_pending_pcpu *p,
           unsigned long flags)
{
 struct rcu_pending_list *expired = &p->lists[NUM_ACTIVE_RCU_POLL_OLDSTATE];
 struct rcu_pending_seq objs = {};
 struct rcu_head *list = NULL;

 if (p->objs.nr &&
     __poll_state_synchronize_rcu(pending->srcu, p->objs.data[0].seq)) {
  objs = p->objs.data[0];
  darray_remove_item(&p->objs, p->objs.data);
 }

 merge_expired_lists(p);

 list = expired->head;
 expired->head = expired->tail = NULL;

 spin_unlock_irqrestore(&p->lock, flags);

 switch ((ulong) pending->process) {
 case RCU_PENDING_KVFREE:
  for (size_t i = 0; i < objs.nr; ) {
   size_t nr_this_node = min(GENRADIX_NODE_SIZE / sizeof(void *), objs.nr - i);

   kfree_bulk(nr_this_node, (void **) genradix_ptr(&objs.objs, i));
   i += nr_this_node;
  }
  genradix_free(&objs.objs);

  while (list) {
   struct rcu_head *obj = list;
#ifdef __KERNEL__
   list = obj->next;
#else
   list = (void *) obj->next.next;
#endif

   /*
 * low bit of pointer indicates whether rcu_head needs
 * to be freed - kvfree_rcu_mightsleep()
 */
   BUILD_BUG_ON(ARCH_SLAB_MINALIGN == 0);

   void *ptr = (void *)(((unsigned long) obj->func) & ~1UL);
   bool free_head = ((unsigned long) obj->func) & 1UL;

   kvfree(ptr);
   if (free_head)
    kfree(obj);
  }

  break;

 case RCU_PENDING_CALL_RCU:
  for (size_t i = 0; i < objs.nr; i++) {
   struct rcu_head *obj = *genradix_ptr(&objs.objs, i);
   obj->func(obj);
  }
  genradix_free(&objs.objs);

  while (list) {
   struct rcu_head *obj = list;
#ifdef __KERNEL__
   list = obj->next;
#else
   list = (void *) obj->next.next;
#endif
   obj->func(obj);
  }
  break;

 default:
  for (size_t i = 0; i < objs.nr; i++)
   pending->process(pending, *genradix_ptr(&objs.objs, i));
  genradix_free(&objs.objs);

  while (list) {
   struct rcu_head *obj = list;
#ifdef __KERNEL__
   list = obj->next;
#else
   list = (void *) obj->next.next;
#endif
   pending->process(pending, obj);
  }
  break;
 }
}

static bool process_finished_items(struct rcu_pending *pending,
       struct rcu_pending_pcpu *p,
       unsigned long flags)
{
 /*
 * XXX: we should grab the gp seq once and avoid multiple function
 * calls, this is called from __rcu_pending_enqueue() fastpath in
 * may_sleep==true mode
 */
 if ((p->objs.nr && __poll_state_synchronize_rcu(pending->srcu, p->objs.data[0].seq)) ||
     (p->lists[0].head && __poll_state_synchronize_rcu(pending->srcu, p->lists[0].seq)) ||
     (p->lists[1].head && __poll_state_synchronize_rcu(pending->srcu, p->lists[1].seq)) ||
     p->lists[2].head) {
  __process_finished_items(pending, p, flags);
  return true;
 }

 return false;
}

static void rcu_pending_work(struct work_struct *work)
{
 struct rcu_pending_pcpu *p =
  container_of(work, struct rcu_pending_pcpu, work);
 struct rcu_pending *pending = p->parent;
 unsigned long flags;

 do {
  spin_lock_irqsave(&p->lock, flags);
 } while (process_finished_items(pending, p, flags));

 spin_unlock_irqrestore(&p->lock, flags);
}

static void rcu_pending_rcu_cb(struct rcu_head *rcu)
{
 struct rcu_pending_pcpu *p = container_of(rcu, struct rcu_pending_pcpu, cb);

 schedule_work_on(p->cpu, &p->work);

 unsigned long flags;
 spin_lock_irqsave(&p->lock, flags);
 if (__rcu_pending_has_pending(p)) {
  spin_unlock_irqrestore(&p->lock, flags);
  __call_rcu(p->parent->srcu, &p->cb, rcu_pending_rcu_cb);
 } else {
  p->cb_armed = false;
  spin_unlock_irqrestore(&p->lock, flags);
 }
}

static __always_inline struct rcu_pending_seq *
get_object_radix(struct rcu_pending_pcpu *p, rcu_gp_poll_state_t seq)
{
 darray_for_each_reverse(p->objs, objs)
  if (rcu_gp_poll_cookie_eq(objs->seq, seq))
   return objs;

 if (darray_push_gfp(&p->objs, ((struct rcu_pending_seq) { .seq = seq }), GFP_ATOMIC))
  return NULL;

 return &darray_last(p->objs);
}

static noinline bool
rcu_pending_enqueue_list(struct rcu_pending_pcpu *p, rcu_gp_poll_state_t seq,
    struct rcu_head *head, void *ptr,
    unsigned long *flags)
{
 if (ptr) {
  if (!head) {
   /*
 * kvfree_rcu_mightsleep(): we weren't passed an
 * rcu_head, but we need one: use the low bit of the
 * ponter to free to flag that the head needs to be
 * freed as well:
 */
   ptr = (void *)(((unsigned long) ptr)|1UL);
   head = kmalloc(sizeof(*head), __GFP_NOWARN);
   if (!head) {
    spin_unlock_irqrestore(&p->lock, *flags);
    head = kmalloc(sizeof(*head), GFP_KERNEL|__GFP_NOFAIL);
    /*
 * dropped lock, did GFP_KERNEL allocation,
 * check for gp expiration
 */
    if (unlikely(__poll_state_synchronize_rcu(p->parent->srcu, seq))) {
     kvfree(--ptr);
     kfree(head);
     spin_lock_irqsave(&p->lock, *flags);
     return false;
    }
   }
  }

  head->func = ptr;
 }
again:
 for (struct rcu_pending_list *i = p->lists;
      i < p->lists + NUM_ACTIVE_RCU_POLL_OLDSTATE; i++) {
  if (rcu_gp_poll_cookie_eq(i->seq, seq)) {
   rcu_pending_list_add(i, head);
   return false;
  }
 }

 for (struct rcu_pending_list *i = p->lists;
      i < p->lists + NUM_ACTIVE_RCU_POLL_OLDSTATE; i++) {
  if (!i->head) {
   i->seq = seq;
   rcu_pending_list_add(i, head);
   return true;
  }
 }

 merge_expired_lists(p);
 goto again;
}

/*
 * __rcu_pending_enqueue: enqueue a pending RCU item, to be processed (via
 * pending->pracess) once grace period elapses.
 *
 * Attempt to enqueue items onto a radix tree; if memory allocation fails, fall
 * back to a linked list.
 *
 * - If @ptr is NULL, we're enqueuing an item for a generic @pending with a
 *   process callback
 *
 * - If @ptr and @head are both not NULL, we're kvfree_rcu()
 *
 * - If @ptr is not NULL and @head is, we're kvfree_rcu_mightsleep()
 *
 * - If @may_sleep is true, will do GFP_KERNEL memory allocations and process
 *   expired items.
 */
static __always_inline void
__rcu_pending_enqueue(struct rcu_pending *pending, struct rcu_head *head,
        void *ptr, bool may_sleep)
{

 struct rcu_pending_pcpu *p;
 struct rcu_pending_seq *objs;
 struct genradix_node *new_node = NULL;
 unsigned long flags;
 bool start_gp = false;

 BUG_ON((ptr != NULL) != (pending->process == RCU_PENDING_KVFREE_FN));

 /* We could technically be scheduled before taking the lock and end up
 * using a different cpu's rcu_pending_pcpu: that's ok, it needs a lock
 * anyways
 *
 * And we have to do it this way to avoid breaking PREEMPT_RT, which
 * redefines how spinlocks work:
 */
 p = raw_cpu_ptr(pending->p);
 spin_lock_irqsave(&p->lock, flags);
 rcu_gp_poll_state_t seq = __get_state_synchronize_rcu(pending->srcu);
restart:
 if (may_sleep &&
     unlikely(process_finished_items(pending, p, flags)))
  goto check_expired;

 /*
 * In kvfree_rcu() mode, the radix tree is only for slab pointers so
 * that we can do kfree_bulk() - vmalloc pointers always use the linked
 * list:
 */
 if (ptr && unlikely(is_vmalloc_addr(ptr)))
  goto list_add;

 objs = get_object_radix(p, seq);
 if (unlikely(!objs))
  goto list_add;

 if (unlikely(!objs->cursor)) {
  /*
 * New radix tree nodes must be added under @p->lock because the
 * tree root is in a darray that can be resized (typically,
 * genradix supports concurrent unlocked allocation of new
 * nodes) - hence preallocation and the retry loop:
 */
  objs->cursor = genradix_ptr_alloc_preallocated_inlined(&objs->objs,
      objs->nr, &new_node, GFP_ATOMIC|__GFP_NOWARN);
  if (unlikely(!objs->cursor)) {
   if (may_sleep) {
    spin_unlock_irqrestore(&p->lock, flags);

    gfp_t gfp = GFP_KERNEL;
    if (!head)
     gfp |= __GFP_NOFAIL;

    new_node = genradix_alloc_node(gfp);
    if (!new_node)
     may_sleep = false;
    goto check_expired;
   }
list_add:
   start_gp = rcu_pending_enqueue_list(p, seq, head, ptr, &flags);
   goto start_gp;
  }
 }

 *objs->cursor++ = ptr ?: head;
 /* zero cursor if we hit the end of a radix tree node: */
 if (!(((ulong) objs->cursor) & (GENRADIX_NODE_SIZE - 1)))
  objs->cursor = NULL;
 start_gp = !objs->nr;
 objs->nr++;
start_gp:
 if (unlikely(start_gp)) {
  /*
 * We only have one callback (ideally, we would have one for
 * every outstanding graceperiod) - so if our callback is
 * already in flight, we may still have to start a grace period
 * (since we used get_state() above, not start_poll())
 */
  if (!p->cb_armed) {
   p->cb_armed = true;
   __call_rcu(pending->srcu, &p->cb, rcu_pending_rcu_cb);
  } else {
   __start_poll_synchronize_rcu(pending->srcu);
  }
 }
 spin_unlock_irqrestore(&p->lock, flags);
free_node:
 if (new_node)
  genradix_free_node(new_node);
 return;
check_expired:
 if (unlikely(__poll_state_synchronize_rcu(pending->srcu, seq))) {
  switch ((ulong) pending->process) {
  case RCU_PENDING_KVFREE:
   kvfree(ptr);
   break;
  case RCU_PENDING_CALL_RCU:
   head->func(head);
   break;
  default:
   pending->process(pending, head);
   break;
  }
  goto free_node;
 }

 p = raw_cpu_ptr(pending->p);
 spin_lock_irqsave(&p->lock, flags);
 goto restart;
}

void rcu_pending_enqueue(struct rcu_pending *pending, struct rcu_head *obj)
{
 __rcu_pending_enqueue(pending, obj, NULL, true);
}

static struct rcu_head *rcu_pending_pcpu_dequeue(struct rcu_pending_pcpu *p)
{
 struct rcu_head *ret = NULL;

 spin_lock_irq(&p->lock);
 darray_for_each(p->objs, objs)
  if (objs->nr) {
   ret = *genradix_ptr(&objs->objs, --objs->nr);
   objs->cursor = NULL;
   if (!objs->nr)
    genradix_free(&objs->objs);
   goto out;
  }

 static_array_for_each(p->lists, i)
  if (i->head) {
   ret = i->head;
#ifdef __KERNEL__
   i->head = ret->next;
#else
   i->head = (void *) ret->next.next;
#endif
   if (!i->head)
    i->tail = NULL;
   goto out;
  }
out:
 spin_unlock_irq(&p->lock);

 return ret;
}

struct rcu_head *rcu_pending_dequeue(struct rcu_pending *pending)
{
 return rcu_pending_pcpu_dequeue(raw_cpu_ptr(pending->p));
}

struct rcu_head *rcu_pending_dequeue_from_all(struct rcu_pending *pending)
{
 struct rcu_head *ret = rcu_pending_dequeue(pending);

 if (ret)
  return ret;

 int cpu;
 for_each_possible_cpu(cpu) {
  ret = rcu_pending_pcpu_dequeue(per_cpu_ptr(pending->p, cpu));
  if (ret)
   break;
 }
 return ret;
}

static bool rcu_pending_has_pending_or_armed(struct rcu_pending *pending)
{
 int cpu;
 for_each_possible_cpu(cpu) {
  struct rcu_pending_pcpu *p = per_cpu_ptr(pending->p, cpu);
  spin_lock_irq(&p->lock);
  if (__rcu_pending_has_pending(p) || p->cb_armed) {
   spin_unlock_irq(&p->lock);
   return true;
  }
  spin_unlock_irq(&p->lock);
 }

 return false;
}

void rcu_pending_exit(struct rcu_pending *pending)
{
 int cpu;

 if (!pending->p)
  return;

 while (rcu_pending_has_pending_or_armed(pending)) {
  __rcu_barrier(pending->srcu);

  for_each_possible_cpu(cpu) {
   struct rcu_pending_pcpu *p = per_cpu_ptr(pending->p, cpu);
   flush_work(&p->work);
  }
 }

 for_each_possible_cpu(cpu) {
  struct rcu_pending_pcpu *p = per_cpu_ptr(pending->p, cpu);
  flush_work(&p->work);
 }

 for_each_possible_cpu(cpu) {
  struct rcu_pending_pcpu *p = per_cpu_ptr(pending->p, cpu);

  static_array_for_each(p->lists, i)
   WARN_ON(i->head);
  WARN_ON(p->objs.nr);
  darray_exit(&p->objs);
 }
 free_percpu(pending->p);
}

/**
 * rcu_pending_init: - initialize a rcu_pending
 *
 * @pending: Object to init
 * @srcu: May optionally be used with an srcu_struct; if NULL, uses normal
 * RCU flavor
 * @process: Callback function invoked on objects once their RCU barriers
 * have completed; if NULL, kvfree() is used.
 */
int rcu_pending_init(struct rcu_pending *pending,
       struct srcu_struct *srcu,
       rcu_pending_process_fn process)
{
 pending->p = alloc_percpu(struct rcu_pending_pcpu);
 if (!pending->p)
  return -ENOMEM;

 int cpu;
 for_each_possible_cpu(cpu) {
  struct rcu_pending_pcpu *p = per_cpu_ptr(pending->p, cpu);
  p->parent = pending;
  p->cpu  = cpu;
  spin_lock_init(&p->lock);
  darray_init(&p->objs);
  INIT_WORK(&p->work, rcu_pending_work);
 }

 pending->srcu = srcu;
 pending->process = process;

 return 0;
}