Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place.

1
#include "byakugan.h"
2
#include "heapStructs.h"
3

4
void initializeHeapModel(struct HeapState *heapModel) {
5
	
6
	memset(heapModel, 0, sizeof (struct HeapState));
7
	heapModel->hPoolListLen = 16;
8
	heapModel->heaps = (struct HPool *) calloc(heapModel->hPoolListLen, sizeof (struct HPool));
9
}
10

11
void heapAllocate(struct HeapState *heapModel, struct AllocateStruct *aStruct) {
12
	struct	HPool		*myHeap;
13
	struct	HeapChunk	*myChunk;
14

15
	if (aStruct->heapHandle == 0 || aStruct->ret == 0)
16
		return;
17

18
	myHeap	= getHeap(heapModel, aStruct->heapHandle);
19
	myChunk = getChunk(myHeap, aStruct->ret);
20
	
21
	if (myChunk == NULL)
22
		return;
23

24
	myChunk->size		= aStruct->size;
25
	myChunk->flags		= aStruct->flags;
26
	myChunk->free		= FALSE;
27
}
28

29
void heapReallocate(struct HeapState *heapModel, struct ReallocateStruct *rStruct) {
30
    struct  HPool   *myHeap; 
31
	struct	HeapChunk	*oChunk, *nChunk;
32
	    
33
    if (rStruct->heapHandle == 0 || rStruct->ret == 0)
34
        return;
35

36
	myHeap = getHeap(heapModel, rStruct->heapHandle);
37
	nChunk = getChunk(myHeap, rStruct->ret);
38
	if (nChunk == NULL)
39
		return;
40

41
	if (rStruct->ret != rStruct->memoryPointer && rStruct->memoryPointer != 0) {
42
		oChunk = getChunk(myHeap, rStruct->memoryPointer);
43
		if (oChunk != NULL)
44
			oChunk->free = TRUE;
45
	}
46

47
	nChunk->size	= rStruct->size;
48
	nChunk->flags	= rStruct->flags;
49
	nChunk->free	= FALSE;
50
}
51

52
void heapFree(struct HeapState *heapModel, struct FreeStruct *fStruct) {
53
    struct  HPool   *myHeap; 
54
	struct	HeapChunk	*myChunk;
55
	   
56
	if (fStruct->heapHandle == 0 || fStruct->memoryPointer == 0x00000000) { 
57
		// So many of these that it slows us down :(
58
		//dprintf("[T] Program attempted to free a NULL pointer.\n\n");
59
		return;
60
	}
61

62
	myHeap	= getHeap(heapModel, fStruct->heapHandle);
63
	myChunk	= getChunk(myHeap, fStruct->memoryPointer);
64
	if (myChunk == NULL)
65
		return;
66

67
	if (myChunk->free == TRUE)
68
		dprintf("[T] Possible 'double free' of chunk @ 0x%08x:0x%08x\n\n", 
69
				myChunk->addr, fStruct->memoryPointer);
70

71
	myChunk->flags	= fStruct->flags;
72
	myChunk->free	= TRUE;
73
}
74

75
void heapCreate(struct HeapState *heapModel, struct CreateStruct *cStruct) {
76
	struct HPool	*newHeap;
77
	ULONG			i;
78

79
    if (cStruct->base == 0)
80
		return;
81

82
	newHeap = getHeap(heapModel, cStruct->base);
83
    newHeap->commit		= cStruct->commit;
84
    newHeap->reserve	= cStruct->reserve;
85
    newHeap->lock		= cStruct->lock;
86
    newHeap->flags		= cStruct->flags;
87
}
88

89
void heapDestroy(struct HeapState *heapModel, struct DestroyStruct *dStruct) {
90
	// Only place we mark heaps as not in use
91
	struct HPool	*deadHeap;
92

93
	if (dStruct->heapHandle == 0)
94
		return;
95

96
	deadHeap = getHeap(heapModel, dStruct->heapHandle);
97
	free(deadHeap->chunks);
98
	memset(deadHeap, 0, sizeof(struct HPool));
99
	heapModel->numHeaps--;
100
}
101

102
void heapCoalesce(struct HeapState *heapModel, struct CoalesceStruct *cfbStruct) {
103
	// Only place we mark chunks as not in use
104
}
105

106
struct HPool *getHeap(struct HeapState *heapModel, PVOID heapHandle) {
107
	ULONG					i;
108
	struct HPool			*newHeap = NULL;
109

110
	for (i = 0; i < heapModel->hPoolListLen; i++) {
111
		if ((heapModel->heaps[i]).base == heapHandle)
112
			return (&(heapModel->heaps[i]));
113
		if (heapModel->numHeaps == heapModel->hPoolListLen)	// Add some short circuits here
114
			continue;
115
		if (newHeap != NULL)
116
			continue;
117
		if ((heapModel->heaps[i]).inUse == FALSE)
118
			newHeap = &(heapModel->heaps[i]);
119
	}
120

121
	// If we get here we haven't found a heap, so lets make one
122
	//dprintf("[T] Creating entry for heap: 0x%08x\n", heapHandle);
123

124
	if (newHeap == NULL) {
125
		heapModel->hPoolListLen = heapModel->hPoolListLen * 2;
126
		newHeap    = heapModel->heaps;
127
		heapModel->heaps = (struct HPool *) realloc(heapModel->heaps,
128
				            heapModel->hPoolListLen * sizeof(struct HPool));
129
		if (heapModel->heaps == NULL) {
130
            heapModel->heaps = newHeap;
131
            //crushModel(heapModel);    ADD ME XXX
132
            dprintf("[T] OOM getting new heaps!\n\n");
133
            return (NULL);
134
        }
135
		newHeap = &(heapModel->heaps[heapModel->numHeaps]);
136
		// Clean the newly allocated memory
137
		memset(newHeap, 0, (sizeof (struct HPool) * (heapModel->hPoolListLen - heapModel->numHeaps)));
138
	}
139

140
	heapModel->numHeaps++;
141
	newHeap->base			= heapHandle;
142
    newHeap->inUse			= TRUE;
143
	newHeap->chunkListLen	= 0x007fffff/8;
144
	newHeap->chunks			= (struct HeapChunk *) calloc(newHeap->chunkListLen, 
145
								sizeof (struct HeapChunk));
146

147
	return (newHeap);
148
}
149

150
// CRRRRAAAAZZYYY splay tree action on top of our chunk code...
151
// Many thanks to Sedgewick and Sleator :)
152
struct HeapChunk *splay(PVOID addr, struct HeapChunk *t) {
153
    struct HeapChunk N, *l, *r, *y;
154
    
155
	if (t == NULL) { 
156
		return (t);
157
	}
158
    memset(&N, 0, sizeof (HeapChunk));
159
	l = r = &N;
160

161
    for (;;) {
162
		if (addr < t->addr) {
163
			if (t->left == NULL) break;
164
			if (addr < (t->left)->addr) {
165
				y			= t->left; 
166
				t->left		= y->right; 
167
				y->right	= t; 
168
				t			= y;
169
				if (t->left == NULL) break;
170
			}
171
			r->left		= t; 
172
			r			= t; 
173
			t			= t->left;
174
		} else if (addr > t->addr) {
175
			if (t->right == NULL) break;
176
			if (addr > (t->right)->addr) {
177
				y			= t->right; 
178
				t->right	= y->left; 
179
				y->left		= t; 
180
				t			= y;
181
				if (t->right == NULL) break;
182
			}
183
			l->right	= t; 
184
			l			= t; 
185
			t			= t->right;
186
		} else {
187
				break;
188
		}
189
    }
190
    l->right	= t->left; 
191
	r->left		= t->right; 
192
	t->left		= N.right; 
193
	t->right	= N.left;
194
    
195
	return (t);
196
}
197

198
struct HeapChunk *insert(struct HeapChunk *newC, struct HPool *heap) {
199
/* Insert i into the tree t, unless it's already there.    */
200
/* Return a pointer to the resulting tree.                 */
201
	struct HeapChunk *t;
202

203
    heap->numChunks++;
204
    newC->inUse = TRUE;
205

206
	t = heap->t;
207
    if (t == NULL) {
208
		heap->t = newC;
209
        return (newC);
210
    }
211
    t = splay(newC->addr,t);
212
    if (newC->addr < t->addr) {
213
        newC->left = t->left;
214
        newC->right = t;
215
        t->left = NULL;
216
		heap->t = newC;
217
        return (newC);
218
    } else if (newC->addr > t->addr) {
219
        newC->right = t->right;
220
        newC->left = t;
221
        t->right = NULL;
222
		heap->t = newC;
223
        return (newC);
224
    } else { /* We get here if it's already in the tree */
225
             /* Don't add it again                      */
226
		dprintf("[XXX] WTF?? Inserting the same address twice!\nBase: 0x%08x\tAddr: 0x%08x:0x%08x\n\n",
227
				heap->base, newC->addr, t->addr);
228
        t->inUse = TRUE;
229
		heap->t = t;
230
		return (t);
231
    }
232
}
233

234

235
struct HeapChunk *getChunk(struct HPool *heap, PVOID memoryPointer) {
236
	ULONG					i;
237
	struct HeapChunk		*newChunk, *niuChunk = NULL;
238

239
// Old and slow - O(n)
240
#if 0
241
	for (i = 0; i < heap->chunkListLen; i++) {
242
		newChunk = &(heap->chunks[i]);
243
		if (newChunk->addr == memoryPointer)
244
			return (newChunk);
245
		if (heap->numChunks == heap->chunkListLen)
246
			continue;
247
		if (niuChunk != NULL)
248
			continue;
249
		if (newChunk->inUse == FALSE)
250
			niuChunk = newChunk;
251
	}
252

253
// New splayness - O(log n)
254
	newChunk = heap->t;
255
	i = 0;
256
	while (newChunk != NULL) {
257
		if (newChunk->addr == memoryPointer) {
258
			return (newChunk);
259
		}
260

261
		if (newChunk->inUse == FALSE) {
262
			dprintf("[XXX] Unused chunk in the tree!! L%d 0x%08x\n\n", 
263
					i, newChunk->addr);
264
			niuChunk = newChunk;
265
		}
266
		
267
		if (memoryPointer < newChunk->addr) {
268
			newChunk = newChunk->left;
269
			i++;
270
		} else {
271
			newChunk = newChunk->right;
272
			i++;
273
		}
274
	}
275
#endif
276
	heap->t = splay(memoryPointer, heap->t);
277
	if (heap->t != NULL && (heap->t)->addr == memoryPointer)
278
		return (heap->t);
279

280
	if (niuChunk == NULL) {
281
		if (heap->numChunks < heap->chunkListLen) {
282
			if ((heap->chunks[heap->numChunks]).inUse == FALSE) {
283
				//dprintf("[XXX] FOUND CHUNK WHERE EXPECTED! :)\n\n");
284
				niuChunk = &(heap->chunks[heap->numChunks]);
285
    			niuChunk->addr			= memoryPointer;
286
				niuChunk->heapHandle	= heap->base;
287
				insert(niuChunk, heap);
288
				return (niuChunk);
289
			}
290
			// Worst muddahfuggin case...
291
			dprintf("[XXX] Unused chunks should be at the END of the list!\n\n");
292
			for (i = 0; i < heap->chunkListLen; i++) {
293
				if ((heap->chunks[i]).inUse == FALSE) {
294
					niuChunk = &(heap->chunks[i]);
295
    				niuChunk->addr  		= memoryPointer;
296
					niuChunk->heapHandle	= heap->base;
297
					insert(niuChunk, heap);
298
					return (niuChunk);
299
				}
300
			}
301
		}
302

303
		heap->chunkListLen = heap->chunkListLen * 2;
304
		niuChunk = heap->chunks;
305
		heap->chunks = (struct HeapChunk *) realloc(heap->chunks,
306
						heap->chunkListLen * sizeof (struct HeapChunk));
307
		if (heap->chunks == NULL) {
308
			heap->chunks = niuChunk;
309
			//crushHeap(heap);
310
			heap->chunkListLen = heap->chunkListLen/2;
311
			dprintf("[T] OOM getting new chunks! %d -> %d\n\n",
312
					heap->chunkListLen, heap->chunkListLen*2);
313
			return (NULL);
314
		}
315
		niuChunk = &(heap->chunks[heap->numChunks]);
316
		memset(niuChunk, 0, (sizeof (struct HeapChunk) * (heap->chunkListLen - heap->numChunks)));
317
	} 
318

319
	niuChunk->addr			= memoryPointer;
320
	niuChunk->heapHandle	= heap->base;
321
	insert(niuChunk, heap);
322
	return (niuChunk);
323
}
324

325