include/afl-prealloc.h


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

/* If we know we'll reuse small elements often, we'll just preallocate a buffer,
 * then fall back to malloc */
// TODO: Replace free status check with bitmask+CLZ

#ifndef AFL_PREALLOC_H
#define AFL_PREALLOC_H

#include <stdio.h>
#include <stdbool.h>
#include <string.h>

#include "debug.h"

typedef enum prealloc_status {

  PRE_STATUS_UNUSED = 0,                                     /* free in buf */
  PRE_STATUS_USED,                                           /* used in buf */
  PRE_STATUS_MALLOC                                        /* system malloc */

} pre_status_t;

/* Adds the entry used for prealloc bookkeeping to this struct */

#define PREALLOCABLE \
  ;                  \
  pre_status_t pre_status;              /* prealloc status of this instance */

/* allocate an element of type *el_ptr, to this variable.
    Uses (and reuses) the given prealloc_buf before hitting libc's malloc.
    prealloc_buf must be the pointer to an array with type `type`.
    `type` must be a struct with uses PREALLOCABLE (a pre_status_t pre_status
   member). prealloc_size must be the array size. prealloc_counter must be a
   variable initialized with 0 (of any name).
    */

#define PRE_ALLOC(el_ptr, prealloc_buf, prealloc_size, prealloc_counter)       \
  do {                                                                         \
                                                                               \
    if ((prealloc_counter) >= (prealloc_size)) {                               \
                                                                               \
      el_ptr = malloc(sizeof(*el_ptr));                                        \
      el_ptr->pre_status = PRE_STATUS_MALLOC;                                  \
                                                                               \
    } else {                                                                   \
                                                                               \
      /* Find one of our preallocated elements */                              \
      u32 i;                                                                   \
      for (i = 0; i < (prealloc_size); i++) {                                  \
                                                                               \
        el_ptr = &((prealloc_buf)[i]);                                         \
        if (el_ptr->pre_status == PRE_STATUS_UNUSED) {                         \
                                                                               \
          (prealloc_counter)++;                                                \
          el_ptr->pre_status = PRE_STATUS_USED;                                \
          break;                                                               \
                                                                               \
        }                                                                      \
                                                                               \
      }                                                                        \
                                                                               \
    }                                                                          \
                                                                               \
    if (!el_ptr) { FATAL("BUG in list.h -> no element found or allocated!"); } \
                                                                               \
  } while (0);

/* Take a chosen (free) element from the prealloc_buf directly */

#define PRE_ALLOC_FORCE(el_ptr, prealloc_counter)         \
  do {                                                    \
                                                          \
    if ((el_ptr)->pre_status != PRE_STATUS_UNUSED) {      \
                                                          \
      FATAL("PRE_ALLOC_FORCE element already allocated"); \
                                                          \
    }                                                     \
    (el_ptr)->pre_status = PRE_STATUS_USED;               \
    (prealloc_counter)++;                                 \
                                                          \
  } while (0);

/* free an preallocated element */

#define PRE_FREE(el_ptr, prealloc_counter)        \
  do {                                            \
                                                  \
    switch ((el_ptr)->pre_status) {               \
                                                  \
      case PRE_STATUS_USED: {                     \
                                                  \
        (el_ptr)->pre_status = PRE_STATUS_UNUSED; \
        (prealloc_counter)--;                     \
        if ((prealloc_counter) < 0) {             \
                                                  \
          FATAL("Inconsistent data in PRE_FREE"); \
                                                  \
        }                                         \
        break;                                    \
                                                  \
      }                                           \
      case PRE_STATUS_MALLOC: {                   \
                                                  \
        (el_ptr)->pre_status = PRE_STATUS_UNUSED; \
        free((el_ptr));                           \
        break;                                    \
                                                  \
      }                                           \
      default: {                                  \
                                                  \
        FATAL("Double Free Detected");            \
        break;                                    \
                                                  \
      }                                           \
                                                  \
    }                                             \
                                                  \
  } while (0);

#endif