From a6ecbb747da311e38867884b35a1d440b9d8c913 Mon Sep 17 00:00:00 2001
From: Justine Tunney <jtunney@gmail.com>
Date: Tue, 7 May 2024 00:37:41 -0700
Subject: [PATCH] Introduce libc/mem/tinymalloc.inc

This allocator shaves ~20kb off single-threaded tool programs and is
slightly faster than proper malloc for simple non-demanding programs
---
 libc/mem/BUILD.mk       |   2 +
 libc/mem/tinymalloc.inc | 160 ++++++++++++++++++++++++++++++++++++++++
 tool/build/apelink.c    |   2 +
 tool/build/assimilate.c |   2 +
 tool/build/compile.c    |   2 +
 tool/build/cp.c         |   2 +
 tool/build/elf2pe.c     |   2 +
 tool/build/fixupobj.c   |   2 +
 tool/build/gzip.c       |   2 +
 tool/build/killall.c    |   2 +
 tool/build/mkdeps.c     |   2 +
 tool/build/mv.c         |   2 +
 tool/build/package.c    |   4 +-
 tool/build/pecheck.c    |   2 +-
 tool/build/resymbol.c   |  46 +++---------
 tool/build/rm.c         |   2 +
 tool/build/symtab.c     |   2 +
 17 files changed, 201 insertions(+), 37 deletions(-)
 create mode 100644 libc/mem/tinymalloc.inc

diff --git a/libc/mem/BUILD.mk b/libc/mem/BUILD.mk
index 5f9bd5f97..bc9ab42cb 100644
--- a/libc/mem/BUILD.mk
+++ b/libc/mem/BUILD.mk
@@ -8,6 +8,7 @@ LIBC_MEM = $(LIBC_MEM_A_DEPS) $(LIBC_MEM_A)
 LIBC_MEM_A = o/$(MODE)/libc/mem/mem.a
 LIBC_MEM_A_FILES := $(wildcard libc/mem/*)
 LIBC_MEM_A_HDRS = $(filter %.h,$(LIBC_MEM_A_FILES))
+LIBC_MEM_A_INCS = $(filter %.inc,$(LIBC_MEM_A_FILES))
 LIBC_MEM_A_SRCS = $(filter %.c,$(LIBC_MEM_A_FILES))
 LIBC_MEM_A_OBJS = $(LIBC_MEM_A_SRCS:%.c=o/$(MODE)/%.o)
 
@@ -46,6 +47,7 @@ $(LIBC_MEM_A_OBJS): private				\
 LIBC_MEM_LIBS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)))
 LIBC_MEM_SRCS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_SRCS))
 LIBC_MEM_HDRS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_HDRS))
+LIBC_MEM_INCS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_INCS))
 LIBC_MEM_BINS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_BINS))
 LIBC_MEM_CHECKS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_CHECKS))
 LIBC_MEM_OBJS = $(foreach x,$(LIBC_MEM_ARTIFACTS),$($(x)_OBJS))
diff --git a/libc/mem/tinymalloc.inc b/libc/mem/tinymalloc.inc
new file mode 100644
index 000000000..f214248aa
--- /dev/null
+++ b/libc/mem/tinymalloc.inc
@@ -0,0 +1,160 @@
+// Copyright 2024 Justine Alexandra Roberts Tunney
+//
+// Permission to use, copy, modify, and/or distribute this software for
+// any purpose with or without fee is hereby granted, provided that the
+// above copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
+// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
+// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
+// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
+// DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
+// PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
+// TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+#include "libc/assert.h"
+#include "libc/errno.h"
+#include "libc/mem/mem.h"
+#include "libc/stdckdint.h"
+#include "libc/str/str.h"
+
+#ifndef MODE_DBG /* don't interfere with asan dlmalloc hooking */
+
+_Alignas(65536) static struct {
+  char memory[1024 * 1024 * 1024];
+  unsigned used, last, free;
+} heap;
+
+static inline bool isheap(char *mem) {
+  return heap.memory <= mem && mem < heap.memory + heap.used;
+}
+
+void free(void *ptr) {
+  char *mem;
+  unsigned base;
+  if (ptr) {
+    mem = (char *)ptr;
+    unassert(isheap(mem));
+    base = mem - heap.memory;
+    *(unsigned *)mem = heap.free;
+    heap.free = base;
+  }
+}
+
+size_t malloc_usable_size(void *ptr) {
+  char *mem = (char *)ptr;
+  unassert(isheap(mem));
+  return ((unsigned *)mem)[-1];
+}
+
+void *memalign(size_t align, size_t need) {
+  char *res;
+  unsigned next, next2, base, toto, *link, *link2;
+
+  // normalize arguments
+  while (align & (align - 1))
+    ++align;
+  if (need < sizeof(unsigned))
+    need = sizeof(unsigned);
+  if (align < sizeof(unsigned))
+    align = sizeof(unsigned);
+  if (align > 65536)
+    goto InvalidArgument;
+  if (ckd_add(&need, need, sizeof(unsigned) - 1))
+    goto OutOfMemory;
+  need &= -sizeof(unsigned);
+
+  // allocate from free list
+  next = heap.free;
+  link = &heap.free;
+  while (next) {
+    next2 = *(unsigned *)(heap.memory + next);
+    link2 = (unsigned *)(heap.memory + next);
+    if (need <= ((unsigned *)(heap.memory + next))[-1]) {
+      *link = next2;
+      return (void *)(heap.memory + next);
+    }
+    next = next2;
+    link = link2;
+  }
+
+  // allocate new static memory
+  base = heap.used;
+  base += sizeof(unsigned);
+  base += align - 1;
+  base &= -align;
+  if (ckd_add(&toto, base, need))
+    goto OutOfMemory;
+  if (toto > sizeof(heap.memory))
+    goto OutOfMemory;
+  res = heap.memory + base;
+  ((unsigned *)res)[-1] = need;
+  heap.used = toto;
+  heap.last = base;
+  return res;
+
+  // we require more vespene gas
+OutOfMemory:
+  errno = ENOMEM;
+  return 0;
+InvalidArgument:
+  errno = EINVAL;
+  return 0;
+}
+
+void *malloc(size_t need) {
+  return memalign(sizeof(max_align_t), need);
+}
+
+void *calloc(size_t count, size_t size) {
+  char *res;
+  unsigned need, used;
+  if (ckd_mul(&need, count, size))
+    need = -1;
+  used = heap.used;
+  if ((res = (char *)malloc(need)))
+    if (res - heap.memory < used)
+      bzero(res, need);
+  return res;
+}
+
+void *realloc(void *ptr, size_t need) {
+  char *res, *mem;
+  unsigned base, have, toto;
+  if (!ptr) {
+    res = (char *)malloc(need);
+  } else {
+    mem = (char *)ptr;
+    unassert(isheap(mem));
+    have = ((unsigned *)mem)[-1];
+    base = mem - heap.memory;
+    if (need < have) {
+      res = mem;
+    } else if (base == heap.last) {
+      if (need < sizeof(unsigned))
+        need = sizeof(unsigned);
+      if (ckd_add(&need, need, sizeof(unsigned) - 1))
+        goto OutOfMemory;
+      need &= -sizeof(unsigned);
+      if (ckd_add(&toto, base, need))
+        goto OutOfMemory;
+      if (toto > sizeof(heap.memory))
+        goto OutOfMemory;
+      ((unsigned *)mem)[-1] = need;
+      heap.used = toto;
+      res = mem;
+    } else if ((res = (char *)malloc(need))) {
+      if (have > need)
+        have = need;
+      memcpy(res, mem, have);
+      free(mem);
+    }
+  }
+  return res;
+OutOfMemory:
+  errno = ENOMEM;
+  return 0;
+}
+
+#endif /* MODE_DBG */
diff --git a/tool/build/apelink.c b/tool/build/apelink.c
index 25bbbd365..cfdcaea61 100644
--- a/tool/build/apelink.c
+++ b/tool/build/apelink.c
@@ -259,6 +259,8 @@ static Elf64_Xword notesize;
 
 static char *r_off32_e_lfanew;
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void Die(const char *thing, const char *reason) {
   tinyprint(2, thing, ": ", reason, "\n", NULL);
   exit(1);
diff --git a/tool/build/assimilate.c b/tool/build/assimilate.c
index d294b8ec2..8c6d4c1cf 100644
--- a/tool/build/assimilate.c
+++ b/tool/build/assimilate.c
@@ -67,6 +67,8 @@
 #define FORMAT_MACHO  2
 #define FORMAT_PE     3
 
+#include "libc/mem/tinymalloc.inc"
+
 static int g_arch;
 static int g_format;
 static bool g_force;
diff --git a/tool/build/compile.c b/tool/build/compile.c
index 3fbecc55c..c493c396c 100644
--- a/tool/build/compile.c
+++ b/tool/build/compile.c
@@ -226,6 +226,8 @@ const char *const kSafeEnv[] = {
     "SYSTEMROOT",  // needed by socket()
 };
 
+#include "libc/mem/tinymalloc.inc"
+
 void OnAlrm(int sig) {
   ++gotalrm;
 }
diff --git a/tool/build/cp.c b/tool/build/cp.c
index 5aed44442..afd8e84cc 100644
--- a/tool/build/cp.c
+++ b/tool/build/cp.c
@@ -69,6 +69,8 @@ char linkbuf[PATH_MAX];
 
 void Cp(char *, char *);
 
+#include "libc/mem/tinymalloc.inc"
+
 bool IsDirectory(const char *path) {
   int e;
   bool res;
diff --git a/tool/build/elf2pe.c b/tool/build/elf2pe.c
index bebbde1d1..9403b0bef 100644
--- a/tool/build/elf2pe.c
+++ b/tool/build/elf2pe.c
@@ -158,6 +158,8 @@ static const char *stubpath;
 static long FLAG_SizeOfStackCommit = 64 * 1024;
 static long FLAG_SizeOfStackReserve = 8 * 1024 * 1024;
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void Die(const char *thing, const char *reason) {
   tinyprint(2, thing, ": ", reason, "\n", NULL);
   exit(1);
diff --git a/tool/build/fixupobj.c b/tool/build/fixupobj.c
index fd3cafdf8..fc2d53cf5 100644
--- a/tool/build/fixupobj.c
+++ b/tool/build/fixupobj.c
@@ -67,6 +67,8 @@ static Elf64_Ehdr *elf;
 static const char *epath;
 static Elf64_Xword symcount;
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void Die(const char *reason) {
   tinyprint(2, epath, ": ", reason, "\n", NULL);
   exit(1);
diff --git a/tool/build/gzip.c b/tool/build/gzip.c
index 9ebbfa5c1..c26d65a73 100644
--- a/tool/build/gzip.c
+++ b/tool/build/gzip.c
@@ -71,6 +71,8 @@ const char *prog;
 char databuf[32768];
 char pathbuf[PATH_MAX];
 
+#include "libc/mem/tinymalloc.inc"
+
 wontreturn void PrintUsage(int rc, FILE *f) {
   fputs("usage: ", f);
   fputs(prog, f);
diff --git a/tool/build/killall.c b/tool/build/killall.c
index 5d347c825..7a250983b 100644
--- a/tool/build/killall.c
+++ b/tool/build/killall.c
@@ -50,6 +50,8 @@ static const char *prog;
 static char16_t **filters;
 static uint32_t pids[10000];
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void PrintUsage(int rc, FILE *f) {
   fprintf(f,
           "Usage: %s [-nshv] NAME...\n"
diff --git a/tool/build/mkdeps.c b/tool/build/mkdeps.c
index 8298ede11..baa6ba843 100644
--- a/tool/build/mkdeps.c
+++ b/tool/build/mkdeps.c
@@ -146,6 +146,8 @@ static const char *buildroot;
 static const char *genroot;
 static const char *outpath;
 
+#include "libc/mem/tinymalloc.inc"
+
 static inline bool IsBlank(int c) {
   return c == ' ' || c == '\t';
 }
diff --git a/tool/build/mv.c b/tool/build/mv.c
index d1cc31b13..503a2855f 100644
--- a/tool/build/mv.c
+++ b/tool/build/mv.c
@@ -62,6 +62,8 @@ char linkbuf[PATH_MAX];
 
 void Mv(char *, char *);
 
+#include "libc/mem/tinymalloc.inc"
+
 wontreturn void Die(const char *path, const char *reason) {
   tinyprint(2, path, ": ", reason, "\n", NULL);
   exit(1);
diff --git a/tool/build/package.c b/tool/build/package.c
index c18b18248..57b5de82a 100644
--- a/tool/build/package.c
+++ b/tool/build/package.c
@@ -43,8 +43,6 @@
 #include "third_party/xed/x86.h"
 #include "tool/build/lib/getargs.h"
 
-__static_yoink("realloc");
-
 /**
  * @fileoverview Build Package Script.
  *
@@ -153,6 +151,8 @@ struct Relas {
   } *p;
 } prtu;
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void Die(const char *path, const char *reason) {
   tinyprint(2, path, ": ", reason, "\n", NULL);
   exit(1);
diff --git a/tool/build/pecheck.c b/tool/build/pecheck.c
index f7586b88e..3585f26e3 100644
--- a/tool/build/pecheck.c
+++ b/tool/build/pecheck.c
@@ -17,7 +17,6 @@
 │ PERFORMANCE OF THIS SOFTWARE.                                                │
 ╚─────────────────────────────────────────────────────────────────────────────*/
 #include "libc/calls/calls.h"
-#include "libc/serialize.h"
 #include "libc/limits.h"
 #include "libc/nt/struct/imageimportbyname.internal.h"
 #include "libc/nt/struct/imageimportdescriptor.internal.h"
@@ -25,6 +24,7 @@
 #include "libc/nt/struct/imageoptionalheader.internal.h"
 #include "libc/nt/struct/imagesectionheader.internal.h"
 #include "libc/runtime/runtime.h"
+#include "libc/serialize.h"
 #include "libc/stdckdint.h"
 #include "libc/stdio/stdio.h"
 #include "libc/str/str.h"
diff --git a/tool/build/resymbol.c b/tool/build/resymbol.c
index a913e3c61..f525324a2 100644
--- a/tool/build/resymbol.c
+++ b/tool/build/resymbol.c
@@ -20,6 +20,7 @@
 #include "libc/elf/struct/shdr.h"
 #include "libc/elf/struct/sym.h"
 #include "libc/errno.h"
+#include "libc/mem/mem.h"
 #include "libc/runtime/runtime.h"
 #include "libc/serialize.h"
 #include "libc/stdio/sysparam.h"
@@ -32,6 +33,8 @@ const char *FLAG_prefix;
 const char *FLAG_suffix;
 const char *path;
 
+#include "libc/mem/tinymalloc.inc"
+
 wontreturn void PrintUsage(int fd, int exitcode) {
   tinyprint(fd, "\n\
 NAME\n\
@@ -74,42 +77,17 @@ wontreturn void DieOom(void) {
   Die("out of memory");
 }
 
-struct {
-  char *last;
-  size_t used;
-  union {
-    char memory[1024 * 1024 * 1024];
-    size_t align;
-  };
-} heap;
-
-void *Malloc(size_t need) {
-  if (need <= sizeof(heap.memory)) {
-    int align = sizeof(size_t);
-    size_t base = heap.used;
-    base += align - 1;
-    base &= -align;
-    size_t toto = base + sizeof(size_t) + need;
-    if (toto >= heap.used && toto <= sizeof(heap.memory)) {
-      char *res = heap.memory + base;
-      *(size_t *)res = need;
-      heap.used = toto;
-      return res + sizeof(size_t);
-    }
-  }
-  DieOom();
+static void *Malloc(size_t n) {
+  void *p;
+  if (!(p = malloc(n)))
+    DieOom();
+  return p;
 }
 
-void *Realloc(void *ptr, size_t need) {
-  if (ptr == heap.last) {
-    heap.used = (char *)ptr - heap.memory;
-    return Malloc(need);
-  } else {
-    void *res = Malloc(need);
-    size_t size = *(size_t *)((char *)ptr - sizeof(size_t));
-    memcpy(res, ptr, MIN(need, size));
-    return res;
-  }
+static void *Realloc(void *p, size_t n) {
+  if (!(p = realloc(p, n)))
+    DieOom();
+  return p;
 }
 
 void ProcessFile(void) {
diff --git a/tool/build/rm.c b/tool/build/rm.c
index c5a041dbc..e4a3a5077 100644
--- a/tool/build/rm.c
+++ b/tool/build/rm.c
@@ -48,6 +48,8 @@ static bool recursive;
 static bool doemptydirs;
 static const char *prog;
 
+#include "libc/mem/tinymalloc.inc"
+
 static wontreturn void PrintUsage(int rc, int fd) {
   tinyprint(fd, "USAGE\n\n  ", prog, USAGE, NULL);
   exit(rc);
diff --git a/tool/build/symtab.c b/tool/build/symtab.c
index 3c3584be4..b372d6f4d 100644
--- a/tool/build/symtab.c
+++ b/tool/build/symtab.c
@@ -30,6 +30,8 @@
  * @fileoverview elf to symbol table file dump tool
  */
 
+#include "libc/mem/tinymalloc.inc"
+
 void PrintUsage(FILE *f) {
   fprintf(f, "%s%s%s\n", "usage: ", program_invocation_name,
           " [-?h] -o PATH COMDBG");