// Small-object allocation. See include/kmem.h for what this is for.
//
// A slab is one page from kalloc(), with a header at the front and objects
// packed after it. Free objects hold the free list in their own first eight
// bytes -- so an object may never be smaller than a pointer, which is what the
// alignment floor below is really enforcing.
//
// Freeing takes no search and no per-object header: an object's page base is
// its slab, and the slab names its cache. That is the entire design, and it is
// the reason an object cannot exceed a page.
//
// Every cache keeps two lists. `partial` holds slabs with at least one free
// object, `full` holds the rest. Allocation takes the head of `partial`;
// freeing into a full slab moves it back. A slab that empties completely is
// handed back to kalloc() straight away rather than kept: hoarding would make
// /proc/slab lie about what the kernel is using, and a subsystem that
// oscillates across a slab boundary is rare enough not to design around.

#ifdef KMEM_HOST
// Host harness: see tools/kmemtest.c. kalloc/kfree become page-aligned
// malloc/free and the locks vanish, because the point of the harness is the
// allocator's own bookkeeping rather than its concurrency.
#include <stdlib.h>
#include <string.h>
typedef unsigned int   uint;
typedef unsigned long  uint64;
#define PGSIZE 4096
static void *kalloc(void)      { void *p = 0;
                                 if (posix_memalign(&p, PGSIZE, PGSIZE)) return 0;
                                 return p; }
static void  kfree(char *p)    { free(p); }
struct spinlock { int x; };
static void initlock(struct spinlock *l, char *n) { (void)n; l->x = 0; }
static void acquire(struct spinlock *l) { (void)l; }
static void release(struct spinlock *l) { (void)l; }
#define PGROUNDDOWN(a) (((uint64)(a)) & ~((uint64)PGSIZE - 1))
#else
#include "types.h"
#include "defs.h"
#include "param.h"
#include "mmu.h"
#include "spinlock.h"
#include "sbuf.h"
#endif

#include "kmem.h"

// Objects are 16-byte aligned. Eight would be enough for the free list, but
// sixteen is what the ABI wants for anything that might hold a long double or
// be memcpy'd with SSE, and the waste at these sizes is a rounding error.
#define KMEM_ALIGN 16
#define ALIGN_UP(n) (((n) + KMEM_ALIGN - 1) & ~(uint)(KMEM_ALIGN - 1))

struct slab {
  struct slab       *next;
  struct slab       *prev;
  struct kmem_cache *cache;
  void              *freelist;   // first free object, or 0
  uint               inuse;
  uint               pad;
};

struct kmem_cache {
  const char        *name;
  uint               objsize;    // after rounding
  uint               per_slab;
  struct slab       *partial;
  struct slab       *full;
  uint               nslab;      // pages held
  uint               inuse;      // objects handed out
  uint64             nalloc, nfree;
  struct spinlock    lock;
  struct kmem_cache *next;       // the global list, for /proc/slab
};

// Caches are themselves fixed, and deliberately: a cache is created once per
// subsystem at init and there are a handful of them, so making them dynamic
// would mean bootstrapping an allocator out of itself for no benefit.
#define NCACHE 32
static struct kmem_cache caches[NCACHE];
static int ncache;
static struct spinlock caches_lock;
static int kmem_ready;

// Rounded *down* to the alignment, so that the value this returns is one
// kmem_cache_create() will actually accept. It used to return the raw
// remainder, which create() then rounded up past its own limit -- so the
// largest legal size was reported as illegal, and the only caller who would
// ever ask is one trying to use every byte.
uint
kmem_max_obj(void)
{
  return (PGSIZE - (uint)sizeof(struct slab)) & ~(uint)(KMEM_ALIGN - 1);
}

static void
kmem_boot(void)
{
  if (kmem_ready)
    return;
  initlock(&caches_lock, "kmem");
  kmem_ready = 1;
}

struct kmem_cache *
kmem_cache_create(const char *name, uint size)
{
  struct kmem_cache *c;

  kmem_boot();
  size = ALIGN_UP(size < KMEM_ALIGN ? KMEM_ALIGN : size);
  if (size > kmem_max_obj())
    return 0;

  acquire(&caches_lock);
  if (ncache >= NCACHE) {
    release(&caches_lock);
    return 0;
  }
  c = &caches[ncache++];
  release(&caches_lock);

  c->name     = name;
  c->objsize  = size;
  c->per_slab = (PGSIZE - (uint)sizeof(struct slab)) / size;
  c->partial  = 0;
  c->full     = 0;
  c->nslab    = c->inuse = 0;
  c->nalloc   = c->nfree = 0;
  initlock(&c->lock, "slab");
  return c;
}

// --- slab list surgery -----------------------------------------------------
//
// Doubly linked because a slab moves between `partial` and `full` from either
// end: allocation takes the head of partial, but a free arrives at whichever
// slab happens to own the object.

static void
list_remove(struct slab **head, struct slab *s)
{
  if (s->prev) s->prev->next = s->next;
  else         *head         = s->next;
  if (s->next) s->next->prev = s->prev;
  s->next = s->prev = 0;
}

static void
list_push(struct slab **head, struct slab *s)
{
  s->prev = 0;
  s->next = *head;
  if (*head) (*head)->prev = s;
  *head = s;
}

static struct slab *
slab_new(struct kmem_cache *c)
{
  char *page = kalloc();
  struct slab *s;
  uint i;

  if (!page)
    return 0;

  s = (struct slab *)page;
  s->next = s->prev = 0;
  s->cache = c;
  s->inuse = 0;

  // Thread the free list through the objects themselves, back to front so the
  // list comes out in address order. Nothing depends on that order; it just
  // makes a dump of the freelist readable when something has gone wrong.
  s->freelist = 0;
  for (i = c->per_slab; i-- > 0; ) {
    char *obj = page + sizeof(struct slab) + i * c->objsize;
    *(void **)obj = s->freelist;
    s->freelist = obj;
  }
  return s;
}

void *
kmem_cache_alloc(struct kmem_cache *c)
{
  struct slab *s;
  void *obj;

  if (!c)
    return 0;

  acquire(&c->lock);
  if (!c->partial) {
    release(&c->lock);
    // Outside the lock: kalloc() takes its own, and holding two locks in one
    // order here and the other order anywhere else is how a kernel deadlocks.
    s = slab_new(c);
    if (!s)
      return 0;
    acquire(&c->lock);
    list_push(&c->partial, s);
    c->nslab++;
  }
  s = c->partial;

  obj = s->freelist;
  s->freelist = *(void **)obj;
  s->inuse++;
  c->inuse++;
  c->nalloc++;

  if (!s->freelist) {
    list_remove(&c->partial, s);
    list_push(&c->full, s);
  }
  release(&c->lock);
  return obj;
}

void
kmem_cache_free(struct kmem_cache *c, void *obj)
{
  struct slab *s;
  int empty;

  if (!obj)
    return;
  s = (struct slab *)PGROUNDDOWN((uint64)obj);
  if (!c)
    c = s->cache;

  acquire(&c->lock);
  if (!s->freelist)                 // was full; it has room again
    { list_remove(&c->full, s); list_push(&c->partial, s); }
  *(void **)obj = s->freelist;
  s->freelist = obj;
  s->inuse--;
  c->inuse--;
  c->nfree++;

  empty = (s->inuse == 0);
  if (empty) {
    list_remove(&c->partial, s);
    c->nslab--;
  }
  release(&c->lock);

  // Outside the lock, for the same reason slab_new() is.
  if (empty)
    kfree((char *)s);
}

// --- the general case ------------------------------------------------------
//
// Power-of-two classes from 16 to 2048. Above that a caller is within a factor
// of two of a page and should take the page.

static const uint gp_size[] = { 16, 32, 64, 128, 256, 512, 1024, 2048 };
#define NGP ((int)(sizeof(gp_size) / sizeof(gp_size[0])))
static struct kmem_cache *gp[NGP];
static const char *gp_name[NGP] = {
  "kmalloc-16",  "kmalloc-32",  "kmalloc-64",   "kmalloc-128",
  "kmalloc-256", "kmalloc-512", "kmalloc-1024", "kmalloc-2048",
};

void *
kmalloc(uint n)
{
  int i;

  if (n == 0 || n > gp_size[NGP - 1])
    return 0;
  for (i = 0; i < NGP; i++)
    if (n <= gp_size[i]) {
      // Created on first use rather than at boot: a kernel that never calls
      // kmalloc for a given size should not hold a page for it.
      if (!gp[i] && !(gp[i] = kmem_cache_create(gp_name[i], gp_size[i])))
        return 0;
      return kmem_cache_alloc(gp[i]);
    }
  return 0;
}

void
kmfree(void *p)
{
  // The cache comes from the object's own page, which is the whole reason
  // kmalloc needs no size argument on the way back.
  if (p)
    kmem_cache_free(0, p);
}

// --- /proc/slab ------------------------------------------------------------

#ifndef KMEM_HOST
void
kmem_stats(struct sbuf *out)
{
  int i;

  sb_puts(out, "name               size  inuse  slabs  per     alloc      free\n");
  for (i = 0; i < ncache; i++) {
    struct kmem_cache *c = &caches[i];

    sb_pad(out, c->name, 18);
    sb_putu(out, c->objsize);      sb_puts(out, "  ");
    sb_putu(out, c->inuse);        sb_puts(out, "  ");
    sb_putu(out, c->nslab);        sb_puts(out, "  ");
    sb_putu(out, c->per_slab);     sb_puts(out, "  ");
    sb_putu(out, c->nalloc);       sb_puts(out, "  ");
    sb_putu(out, c->nfree);
    sb_putc(out, '\n');
  }
}

#endif