snapshot: refactor the fork-based stack initialization
[model-checker.git] / mymemory.cc
1 #include "mymemory.h"
2 #include "snapshot.h"
3 #include "snapshotimp.h"
4 #include <stdio.h>
5 #include <dlfcn.h>
6 #include <unistd.h>
7 #include <cstring>
8 #define REQUESTS_BEFORE_ALLOC 1024
9 size_t allocatedReqs[ REQUESTS_BEFORE_ALLOC ] = { 0 };
10 int nextRequest = 0;
11 int howManyFreed = 0;
12 #if !USE_MPROTECT_SNAPSHOT
13 static mspace sStaticSpace = NULL;
14 #endif
15
16 /** Non-snapshotting malloc for our use. */
17
18 void *MYMALLOC(size_t size) {
19 #if USE_MPROTECT_SNAPSHOT
20         static void *(*mallocp)(size_t size);
21         char *error;
22         void *ptr;
23   
24         /* get address of libc malloc */
25         if (!mallocp) {
26                 mallocp = ( void * ( * )( size_t ) )dlsym(RTLD_NEXT, "malloc");
27                 if ((error = dlerror()) != NULL) {
28                         fputs(error, stderr);
29                         exit(EXIT_FAILURE);
30                 }
31         }
32         ptr = mallocp(size);     
33         return ptr;
34 #else
35         if( !sTheRecord ){
36                 createSharedLibrary();
37         }
38         if( NULL == sStaticSpace )
39                 sStaticSpace = create_mspace_with_base( ( void * )( sTheRecord->mSharedMemoryBase ), SHARED_MEMORY_DEFAULT -sizeof( struct Snapshot ), 1 );
40         return mspace_malloc( sStaticSpace, size );
41 #endif
42 }
43
44 void *system_malloc( size_t size ){
45         static void *(*mallocp)(size_t size);
46         char *error;
47         void *ptr;
48
49         /* get address of libc malloc */
50         if (!mallocp) {
51                 mallocp = ( void * ( * )( size_t ) )dlsym(RTLD_NEXT, "malloc");
52                 if ((error = dlerror()) != NULL) {
53                         fputs(error, stderr);
54                         exit(EXIT_FAILURE);
55                 }
56         }
57         ptr = mallocp(size);
58         return ptr;
59 }
60
61 void system_free( void * ptr ){
62         static void (*freep)(void *);
63         char *error;
64
65         /* get address of libc free */
66         if (!freep) {
67                 freep = ( void  ( * )( void * ) )dlsym(RTLD_NEXT, "free");
68                 if ((error = dlerror()) != NULL) {
69                         fputs(error, stderr);
70                         exit(EXIT_FAILURE);
71                 }
72         }
73         freep(ptr);
74 }
75
76 /** Non-snapshotting free for our use. */
77 void MYFREE(void *ptr) {
78 #if USE_MPROTECT_SNAPSHOT
79         static void (*freep)(void *);
80         char *error;
81
82         /* get address of libc free */
83         if (!freep) {
84                 freep = ( void  ( * )( void * ) )dlsym(RTLD_NEXT, "free");
85                 if ((error = dlerror()) != NULL) {
86                         fputs(error, stderr);
87                         exit(EXIT_FAILURE);
88                 }
89         }
90         freep(ptr);
91 #else
92         mspace_free( sStaticSpace, ptr );
93 #endif
94 }
95
96
97 /** This global references the mspace for the snapshotting heap */
98 mspace mySpace = NULL;
99
100 /** This global references the unaligned memory address that was malloced for the snapshotting heap */
101 void * basemySpace = NULL;
102
103 /** Adding the fix for not able to allocate through a reimplemented calloc at the beginning before instantiating our allocator
104 A bit circumspect about adding an sbrk. linux docs say to avoid using it... */
105
106 void * HandleEarlyAllocationRequest( size_t sz ){
107         if( 0 == mySpace ){
108                 void * returnAddress = sbrk( sz );
109                 if( nextRequest >= REQUESTS_BEFORE_ALLOC ){
110                         exit( EXIT_FAILURE );
111                 }
112                 allocatedReqs[ nextRequest++ ] = ( size_t )returnAddress;
113                 return returnAddress;
114         }
115         return NULL;
116 }
117
118 /** The fact that I am not expecting more than a handful requests is implicit in my not using a binary search here*/
119
120 bool DontFree( void * ptr ){
121         if( howManyFreed == nextRequest ) return false; //a minor optimization to reduce the number of instructions executed on each free call....
122         if( NULL == ptr ) return true;
123         for( int i =  nextRequest - 1; i >= 0; --i ){
124                 if( allocatedReqs[ i ] ==  ( size_t )ptr ) {
125                         ++howManyFreed;
126                         return true;
127                 }
128         }
129         return false;
130 }
131
132 /** Snapshotting malloc implementation for user programs. */
133
134 void *malloc( size_t size ) {
135         void * earlyReq = HandleEarlyAllocationRequest( size );
136         if( earlyReq ) return earlyReq;
137         return mspace_malloc( mySpace, size );
138 }
139
140 /** Snapshotting free implementation for user programs. */
141
142 void free( void * ptr ){
143         if( DontFree( ptr ) ) return;
144         mspace_free( mySpace, ptr );
145 }
146
147 /** Snapshotting realloc implementation for user programs. */
148
149 void *realloc( void *ptr, size_t size ){
150         return mspace_realloc( mySpace, ptr, size );
151 }
152
153 /** Snapshotting calloc implementation for user programs. */
154
155 void * calloc( size_t num, size_t size ){
156         void * earlyReq = HandleEarlyAllocationRequest( size * num );
157         if( earlyReq ) {
158                 std::memset( earlyReq, 0, size * num );
159                 return earlyReq;
160         }
161         return mspace_calloc( mySpace, num, size );
162 }
163
164 /** Snapshotting new operator for user programs. */
165
166 void * operator new(size_t size) throw(std::bad_alloc) {
167         return malloc(size);
168 }
169
170 /** Snapshotting delete operator for user programs. */
171
172 void operator delete(void *p) throw() {
173         free(p);
174 }
175
176 /** Snapshotting new[] operator for user programs. */
177
178 void * operator new[](size_t size) throw(std::bad_alloc) {
179         return malloc(size);
180 }
181
182 /** Snapshotting delete[] operator for user programs. */
183
184 void operator delete[](void *p, size_t size) {
185         free(p);
186 }