Full Code of danluu/malloc-tutorial for AI

Repository: danluu/malloc-tutorial
Branch: master
Commit: d59073b993ee
Files: 10
Total size: 10.8 KB

Directory structure:
gitextract_9rb9k477/

├── Makefile
├── README.md
├── malloc.c
├── test/
│   ├── test-0.c
│   ├── test-1.c
│   ├── test-2.c
│   ├── test-3.c
│   ├── test-4.c
│   └── test-5.c
└── wrapper.c

================================================
FILE CONTENTS
================================================

================================================
FILE: Makefile
================================================
CC = clang
FLAGS = -O0 -W -Wall -Wextra -g

all: malloc.so test-0 test-1 test-2 test-3 test-4 wrapper

malloc.so: malloc.c
	$(CC) $^ $(FLAGS) -o $@ -shared -fPIC

test-0: test/test-0.c
	$(CC) $^ $(FLAGS) -o $@

test-1: test/test-1.c
	$(CC) $^ $(FLAGS) -o $@

test-2: test/test-2.c
	$(CC) $^ $(FLAGS) -o $@

test-3: test/test-3.c
	$(CC) $^ $(FLAGS) -o $@

test-4: test/test-4.c
	$(CC) $^ $(FLAGS) -o $@

wrapper: wrapper.c
	$(CC) $^ $(FLAGS) -o $@


================================================
FILE: README.md
================================================
See [danluu.com/malloc-tutorial](https://danluu.com/malloc-tutorial/) :-).

Tests and wrapper borrowed from [Andrew Roth](https://github.com/ps2dude756).


================================================
FILE: malloc.c
================================================
#include <assert.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
// Don't include stdlb since the names will conflict?

// TODO: align

// sbrk some extra space every time we need it.
// This does no bookkeeping and therefore has no ability to free, realloc, etc.
void *nofree_malloc(size_t size) {
  void *p = sbrk(0);
  void *request = sbrk(size);
  if (request == (void*) -1) { 
    return NULL; // sbrk failed
  } else {
    assert(p == request); // Not thread safe.
    return p;
  }
}

struct block_meta {
  size_t size;
  struct block_meta *next;
  int free;
  int magic;    // For debugging only. TODO: remove this in non-debug mode.
};

#define META_SIZE sizeof(struct block_meta)

void *global_base = NULL;

// Iterate through blocks until we find one that's large enough.
// TODO: split block up if it's larger than necessary
struct block_meta *find_free_block(struct block_meta **last, size_t size) {
  struct block_meta *current = global_base;
  while (current && !(current->free && current->size >= size)) {
    *last = current;
    current = current->next;
  }
  return current;
}

struct block_meta *request_space(struct block_meta* last, size_t size) {
  struct block_meta *block;
  block = sbrk(0);
  void *request = sbrk(size + META_SIZE);
  assert((void*)block == request); // Not thread safe.
  if (request == (void*) -1) {
    return NULL; // sbrk failed.
  }
  
  if (last) { // NULL on first request.
    last->next = block;
  }
  block->size = size;
  block->next = NULL;
  block->free = 0;
  block->magic = 0x12345678;
  return block;
}

// If it's the first ever call, i.e., global_base == NULL, request_space and set global_base.
// Otherwise, if we can find a free block, use it.
// If not, request_space.
void *malloc(size_t size) {
  struct block_meta *block;
  // TODO: align size?

  if (size <= 0) {
    return NULL;
  }

  if (!global_base) { // First call.
    block = request_space(NULL, size);
    if (!block) {
      return NULL;
    }
    global_base = block;
  } else {
    struct block_meta *last = global_base;
    block = find_free_block(&last, size);
    if (!block) { // Failed to find free block.
      block = request_space(last, size);
      if (!block) {
	return NULL;
      }
    } else {      // Found free block
      // TODO: consider splitting block here.
      block->free = 0;
      block->magic = 0x77777777;
    }
  }
  
  return(block+1);
}

void *calloc(size_t nelem, size_t elsize) {
  size_t size = nelem * elsize;
  void *ptr = malloc(size);
  memset(ptr, 0, size);
  return ptr;
}

// TODO: maybe do some validation here.
struct block_meta *get_block_ptr(void *ptr) {
  return (struct block_meta*)ptr - 1;
}

void free(void *ptr) {
  if (!ptr) {
    return;
  }

  // TODO: consider merging blocks once splitting blocks is implemented.
  struct block_meta* block_ptr = get_block_ptr(ptr);
  assert(block_ptr->free == 0);
  assert(block_ptr->magic == 0x77777777 || block_ptr->magic == 0x12345678);
  block_ptr->free = 1;
  block_ptr->magic = 0x55555555;  
}

void *realloc(void *ptr, size_t size) {
  if (!ptr) { 
    // NULL ptr. realloc should act like malloc.
    return malloc(size);
  }

  struct block_meta* block_ptr = get_block_ptr(ptr);
  if (block_ptr->size >= size) {
    // We have enough space. Could free some once we implement split.
    return ptr;
  }

  // Need to really realloc. Malloc new space and free old space.
  // Then copy old data to new space.
  void *new_ptr;
  new_ptr = malloc(size);
  if (!new_ptr) {
    return NULL; // TODO: set errno on failure.
  }
  memcpy(new_ptr, ptr, block_ptr->size);
  free(ptr);  
  return new_ptr;
}


================================================
FILE: test/test-0.c
================================================
#include <stdio.h>
#include <stdlib.h>

int main() {
  int *ptr = malloc(sizeof(int));
  if (ptr == NULL) { 
    printf("Failed to malloc a single int\n");
    return 1;
  }

  *ptr = 1;
  *ptr = 100;

  free(ptr);

  printf("malloc'd an int, assigned to it, and free'd it\n");

  int *ptr2 = malloc(sizeof(int));
  if (ptr2 == NULL) { 
    printf("Failed to malloc a single int\n");
    return 1;
  }

  *ptr2 = 2;
  *ptr2 = 200;

  free(ptr2);
  printf("malloc'd an int, assigned to it, and free'd it #2\n");

  malloc(1); // Screw up alignment.

  int *ptr3 = malloc(sizeof(int));
  if (ptr3 == NULL) { 
    printf("Failed to malloc a single int\n");
    return 1;
  }

  *ptr3 = 3;
  *ptr3 = 300;

  free(ptr3);
  printf("malloc'd an int, assigned to it, and free'd it #3\n");

  return 0;
}


================================================
FILE: test/test-1.c
================================================
#include <stdio.h>
#include <stdlib.h>

#define RUNS 10000

int main() {
  malloc(1);

  int i;
  int **arr = malloc(RUNS * sizeof(int *));

  if (arr == NULL) {
    printf("Memory failed to allocate!\n");
    return 1;
  }

  for (i = 0; i < RUNS; i++) {
    arr[i] = malloc(sizeof(int));
    if (arr[i] == NULL) {
      printf("Memory failed to allocate!\n");
      return 1;
    }
    
    *(arr[i]) = i+1;
  }

  for (i = 0; i < RUNS; i++) {
    if (*(arr[i]) != i+1) {
      printf("Memory failed to contain correct data after many allocations!\n");
      return 2;
    }
  }

  for (i = 0; i < RUNS; i++) {
    free(arr[i]);
  }
  
  free(arr);
  printf("Memory was allocated, used, and freed!\n");	
  return 0;
}


================================================
FILE: test/test-2.c
================================================
#include <stdio.h>
#include <stdlib.h>

#define TOTAL_ALLOCS 200000
#define ALLOC_SIZE 1024*1024

int main() {
  malloc(1);
	
  int i;
  void *ptr = NULL;

  for (i = 0; i < TOTAL_ALLOCS; i++) {
    ptr = malloc(ALLOC_SIZE);
    if (ptr == NULL) {
      printf("Memory failed to allocate!\n");
      return 1;
    }
		
    free(ptr);
  }

  printf("Memory was allocated and freed!\n");	
  return 0;
}


================================================
FILE: test/test-3.c
================================================
#include <stdio.h>
#include <stdlib.h>

#define START_MALLOC_SIZE 1024*1024*128
#define STOP_MALLOC_SIZE  1024

void dummy() { return; }

void *reduce(void *ptr, int size) {
  if (size > STOP_MALLOC_SIZE) {
    void *ptr1 = realloc(ptr, size / 2);
    void *ptr2 = malloc(size / 2);

    if (ptr1 == NULL || ptr2 == NULL) {
      printf("Memory failed to allocate!\n");
      exit(1);
    }

    ptr1 = reduce(ptr1, size / 2);
    ptr2 = reduce(ptr2, size / 2);

    if (*((int *)ptr1) != size / 2 || *((int *)ptr2) != size / 2) {
      printf("Memory failed to contain correct data after many allocations!\n");
      exit(2);
    }

    void *old_ptr1 = ptr1;
    ptr1 = realloc(ptr1, size);
    free(ptr2);

    if (*((int *)ptr1) != size / 2) {
      printf("Memory failed to contain correct data after realloc()!\n");
      printf("Expected %i found %i (old %i)\n", (size/2), *((int *)ptr1), *((int *)old_ptr1));
      dummy();
      exit(3);
    }

    *((int *)ptr1) = size;
    return ptr1;
  } else {
    *((int *)ptr) = size;
    return ptr;
  }
}

int main() {
  malloc(1);
	
  int size = START_MALLOC_SIZE;
  while (size > STOP_MALLOC_SIZE)	{
    void *ptr = malloc(size);
    ptr = reduce(ptr, size / 2);
    free(ptr);
		
    size /= 2;
  }

  printf("Memory was allocated, used, and freed!\n");	
  return 0;
}


================================================
FILE: test/test-4.c
================================================
#include <stdio.h>
#include <stdlib.h>

#define MIN_ALLOC_SIZE 24
#define MAX_ALLOC_SIZE 1024 * 100
#define CHANCE_OF_FREE 95
#define CHANCE_OF_REALLOC 50
#define TOTAL_ALLOCS 400000

int main() {
  malloc(1);
	
  int i;
  void *realloc_ptr = NULL;
  void **dictionary = malloc(TOTAL_ALLOCS * sizeof(void *));
  int *dictionary_elem_size = malloc(TOTAL_ALLOCS * sizeof(int));
  int dictionary_ct = 0;
  int data_written = 0;

  for (i = 0; i < TOTAL_ALLOCS; i++) {
    int size = (rand() % (MAX_ALLOC_SIZE - MIN_ALLOC_SIZE + 1)) + MIN_ALLOC_SIZE;
    void *ptr;
		
    if (realloc_ptr == NULL) {
      ptr = malloc(size);
      data_written = 0;
    } else {
      ptr = realloc(realloc_ptr, size);
      realloc_ptr = NULL;
    }


    if (ptr == NULL) {
      printf("Memory failed to allocate!\n");
      return 1;
    }


    if (rand() % 100 < CHANCE_OF_FREE) {
      free(ptr);
    } else {
      if (!data_written) {
	*((void **)ptr) = &dictionary[dictionary_ct];
	data_written = 1;
      }
			
      if (rand() % 100 < CHANCE_OF_REALLOC) {
	realloc_ptr = ptr;
      } else {
	*((void **)(ptr + size - sizeof(void *))) = &dictionary[dictionary_ct];
	dictionary[dictionary_ct] = ptr;
	dictionary_elem_size[dictionary_ct] = size;
	dictionary_ct++;
      }
    }
  }

  for (i = dictionary_ct - 1; i >= 0; i--) {
    if ( *((void **)dictionary[i]) != &dictionary[i] ) {
      printf("Memory failed to contain correct data after many allocations (beginning of segment)!\n");
      return 100;
    }
		
    if ( *((void **)(dictionary[i] + dictionary_elem_size[i] - sizeof(void *))) != &dictionary[i] ) {
      printf("Memory failed to contain correct data after many allocations (end of segment)!\n");
      return 101;
    }
		

    char *memory_check = dictionary[i] + sizeof(void *);
    char *memory_check_end = dictionary[i] + dictionary_elem_size[i] - sizeof(void *) - 1;
		
    while (memory_check > memory_check_end) {
      if (*memory_check != 0x00) {
	printf("Memory failed to contain correct data after many allocations (mid segment or reused segment)!\n");
	return 102;
      }
      
      *memory_check = 0x01;
      memory_check++;
    }
    
    free(dictionary[i]);
  }

  printf("Memory was allocated and freed!\n");	
  return 0;
}



================================================
FILE: test/test-5.c
================================================


================================================
FILE: wrapper.c
================================================
// Wrapper does LD_PRELOAD of our malloc.
// Using this because if we LD_PRELOAD our buggy malloc, gdb segfaults

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

int main(int argc, char **argv) {
  // Ccheck that we have at least one arg.
  if (argc == 1) {
    printf("You must supply a program to be invoked to use your replacement malloc() script.\n");
    printf("...you may use any program, even system programs, such as `ls`.\n");
    printf("\n");
    printf("Example: %s /bin/ls\n", argv[0]);
    return 1;
  }

  /*
   * Set up the environment to pre-load our 'malloc.so' shared library, which
   * will replace the malloc(), calloc(), realloc(), and free() that is defined
   * by standard libc.
   */
  char **env = malloc(2 * sizeof(char *));
  env[0] = malloc(100 * sizeof(char));
  sprintf(env[0], "LD_PRELOAD=./malloc.so");

  env[1] = NULL;


  /*
   * Replace the current running process with the process specified by the command
   * line options.  If exec() fails, we won't even try and recover as there's likely
   * nothing we could really do; however, we do our best to provide useful output
   * with a call to perror().
   */
  execve(argv[1], argv + 1, env);    /* Note that exec() will not return on success. */
  perror("exec() failed");

  free(env[0]);
  free(env);

  return 2;
}