sfpool.h

/* SPDX-FileCopyrightText: 2025 Dyne.org foundation
 * SPDX-License-Identifier: GPL-3.0-or-later
 *
 * Copyright (C) 2025 Dyne.org foundation
 * designed, written and maintained by Denis Roio <jaromil@dyne.org>
 *
 * This program is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Affero General Public License for more details.
 *
 * You should have received a copy of the GNU General Public
 * License along with this program.  If not, see
 * <https://www.gnu.org/licenses/>.
 *
 */
 
#ifndef __SFPOOL_H__
#define __SFPOOL_H__
 
// This header carries both declarations and implementation.
// All functions use internal linkage, so it is safe to include it
// from multiple translation units in the same program.
 
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <assert.h>
#ifdef __EMSCRIPTEN__
#include <emscripten.h>
#elif defined(_WIN32)
#include <windows.h>
#else // lucky shot on POSIX
// defined(__unix__) || defined(__linux__) || defined(__APPLE__) ||  defined(__DragonFly__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__)
#include <unistd.h> // for geteuid to switch protected memory map
#include <sys/resource.h>
#include <sys/mman.h>
#endif
 
// Configuration
#define SECURE_ZERO // Enable secure zeroing
#define PROFILING // Profile most used sizes allocated
 
#if defined(__x86_64__) || defined(_M_X64) || defined(__ppc64__) || defined(__LP64__)
#define ptr_t uint64_t
#define ptr_align 8
#define struct_align 16
#else
#define ptr_t uint32_t
#define ptr_align 4
#define struct_align 8
#endif
 
// Memory pool structure
typedef struct __attribute__((aligned(struct_align))) sfpool_t {
  uint8_t *buffer; // raw
  uint8_t *data; // aligned
  uint8_t *free_list;
  uint32_t free_count;
  uint32_t total_blocks;
  uint32_t total_bytes;
  uint32_t block_size;
#ifdef PROFILING
  uint32_t *hits;
  uint32_t hits_total;
  size_t   hits_bytes;
  uint32_t miss_total;
  size_t   miss_bytes;
  size_t   alloc_total;
#endif
} sfpool_t;
 
 
#if !defined(__MUSL__)
static_assert(sizeof(ptr_t) == sizeof(void*), "Unknown memory pointer size detected");
#endif
static inline bool _is_in_pool(sfpool_t *pool, const void *ptr) {
  volatile ptr_t p = (ptr_t)ptr;
  return(p >= (ptr_t)pool->data
         && p < (ptr_t)(pool->data + pool->total_bytes));
}
 
static inline void sfutil_zero(void *ptr, uint32_t size) {
  volatile uint8_t *p = (volatile uint8_t*)ptr;
  while (size--) *p++ = 0;
}
 
static inline void *sfutil_memalign(const void* ptr) {
    register ptr_t mask = ptr_align - 1;
    ptr_t aligned = ((ptr_t)ptr + mask) & ~mask;
    return (void*)aligned;
}
 
static inline void *sfutil_secalloc(size_t size) {
    // add bytes to every allocation to support alignment
    size_t alloc_size = size + ptr_align;
    void *res = NULL;
#if defined(__EMSCRIPTEN__)
    res = (uint8_t *)malloc(alloc_size);
#elif defined(_WIN32)
    res = VirtualAlloc(NULL, alloc_size,
                       MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
#elif defined(__APPLE__)
    int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE;
    res = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE, flags, -1, 0);
    if (res == MAP_FAILED) return NULL;
    mlock(res, alloc_size);
#else // assume POSIX
    int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE;
    struct rlimit rl;
    if (getrlimit(RLIMIT_MEMLOCK, &rl) == 0)
        if(alloc_size<=rl.rlim_cur) flags |= MAP_LOCKED;
    res = mmap(NULL, alloc_size, PROT_READ | PROT_WRITE, flags, -1, 0);
    if (res == MAP_FAILED) return NULL;
#endif
    return res;
}
 
static inline void sfutil_secfree(void *ptr, size_t size) {
    size_t alloc_size = size + ptr_align;
#if defined(__EMSCRIPTEN__)
    free(ptr);
#elif defined(_WIN32)
    VirtualFree(ptr, 0, MEM_RELEASE);
#else // Posix
    munmap(ptr, alloc_size);
#endif
}
 
 
 
static inline size_t sfpool_init(sfpool_t *pool, size_t nmemb, size_t blocksize) {
  if (pool == NULL) return 0;
  memset(pool, 0, sizeof(sfpool_t));
  if (nmemb == 0) return 0;
  if (blocksize < sizeof(void*)) return 0;
  if((blocksize & (blocksize - 1)) != 0) return 0;
  if (nmemb > (SIZE_MAX / blocksize)) return 0;
  // SFPool blocksize must be a power of two
  size_t totalsize = nmemb * blocksize;
  pool->buffer = sfutil_secalloc(totalsize);
  if (pool->buffer == NULL) return 0;
  // Failed to allocate pool memory
  pool->data   = sfutil_memalign(pool->buffer);
  if (pool->data == NULL) return 0;
  // Failed to allocate pool memory
  pool->total_bytes  = totalsize;
  pool->total_blocks = nmemb;
  pool->block_size   = blocksize;
  // Initialize the embedded free list
  pool->free_list = pool->data;
  register uint32_t i, bi;
  for (i = 0; i < pool->total_blocks - 1; ++i) {
    bi = i*blocksize;
    *(uint8_t **)(pool->data + bi) =
      pool->data + bi + blocksize;
  }
  pool->free_count = pool->total_blocks;
  *(uint8_t **)
    (pool->data + (pool->total_blocks - 1) * blocksize) = NULL;
#ifdef PROFILING
  pool->miss_total = pool->miss_bytes = 0;
  pool->hits_total = pool->hits_bytes = 0;
  pool->alloc_total = 0;
#endif
  return totalsize;
}
 
 
static inline void sfpool_teardown(sfpool_t *restrict pool) {
  // Free pool memory
  sfutil_secfree(pool->buffer, pool->total_bytes);
#ifdef PROFILING
  pool->miss_total = pool->miss_bytes = 0;
  pool->hits_total = pool->hits_bytes = 0;
  pool->alloc_total = 0;
#endif
}
 
static inline void *sfpool_malloc(void *restrict opaque, const size_t size) {
  sfpool_t *pool = (sfpool_t*)opaque;
  void *ptr;
  if (size <= pool->block_size
      && pool->free_list != NULL) {
#ifdef PROFILING
    pool->hits_total++;
    pool->hits_bytes+=size;
    pool->alloc_total+=size;
#endif
    // Remove the first block from the free list
    uint8_t *block = pool->free_list;
    pool->free_list = *(uint8_t **)block;
    pool->free_count-- ;
    return block;
  }
  // Fallback to system malloc for large allocations
  ptr = malloc(size);
  if(ptr == NULL) perror("system malloc error");
#ifdef PROFILING
  pool->miss_total++;
  pool->miss_bytes+=size;
  pool->alloc_total+=size;
#endif
  return ptr;
}
 
 
static inline void sfpool_free(void *restrict opaque, void *ptr) {
  sfpool_t *pool = (sfpool_t*)opaque;
  if (ptr == NULL) return; // Freeing NULL is a no-op
  if (_is_in_pool(pool,ptr)) {
#ifdef SECURE_ZERO
    // Zero the user-visible contents before restoring the free-list link.
    sfutil_zero(ptr, pool->block_size);
#endif
    // Add the block back to the free list
    *(uint8_t **)ptr = pool->free_list;
    pool->free_list = (uint8_t *)ptr;
    pool->free_count++ ;
    return;
  } else {
    free(ptr);
  }
}
 
 
static inline void *sfpool_realloc(void *restrict opaque, void *ptr, const size_t size) {
  sfpool_t *pool = (sfpool_t*)opaque;
  if (ptr == NULL) {
    return sfpool_malloc(pool, size);
  }
  if (size == 0) {
    sfpool_free(pool, ptr);
    return NULL;
  }
  if (_is_in_pool((sfpool_t*)pool,ptr)) {
    if (size <= pool->block_size) {
#ifdef PROFILING
      pool->hits_total++;
      pool->hits_bytes+=size;
      pool->alloc_total+=size;
#endif
      return ptr; // No need to reallocate
    } else {
      void *new_ptr = malloc(size);
      if (new_ptr == NULL) return NULL;
      memcpy(new_ptr, ptr, pool->block_size); // Copy only BLOCK_SIZE bytes
#ifdef SECURE_ZERO
      // Zero the old pool block before relinking it into the free list.
      sfutil_zero(ptr, pool->block_size);
#endif
      // Add the block back to the free list
      *(uint8_t **)ptr = pool->free_list;
      pool->free_list = (uint8_t *)ptr;
      pool->free_count++ ;
#ifdef PROFILING
  pool->miss_total++;
  pool->miss_bytes+=size;
  pool->alloc_total+=size;
#endif
      return new_ptr;
    }
  } else {
    // Handle large allocations
    return realloc(ptr, size);
#ifdef PROFILING
    pool->miss_total++;
    pool->miss_bytes+=size;
    pool->alloc_total+=size;
#endif
  }
}
 
static inline int sfpool_contains(void *restrict opaque, const void *ptr) {
  sfpool_t *pool = (sfpool_t*)opaque;
  int res = 0;
  if( _is_in_pool(pool,ptr) ) res = 1;
  return res;
}
 
 
static inline void sfpool_status(sfpool_t *restrict p) {
  fprintf(stderr,"\n🌊 sfpool: %u blocks %u B each\n",
          p->total_blocks, p->block_size);
#ifdef PROFILING
  fprintf(stderr,"🌊 Total:  %lu K\n",
          p->alloc_total/1024);
  fprintf(stderr,"🌊 Misses: %lu K (%u calls)\n",p->miss_bytes/1024,p->miss_total);
  fprintf(stderr,"🌊 Hits:   %lu K (%u calls)\n",p->hits_bytes/1024,p->hits_total);
#endif
}
 
 
#endif