feat: No more kernel page table switching upon interrupt.
[lunaix-os.git] / lunaix-os / kernel / mm / kalloc.c
1 /**
2  * @file kalloc.c
3  * @author Lunaixsky
4  * @brief Implicit free list implementation of malloc family, for kernel use.
5  * 
6  * This version of code is however the simplest and yet insecured, thread unsafe
7  * it just to demonstrate how the malloc/free works behind the curtain
8  * @version 0.1
9  * @date 2022-03-05
10  * 
11  * @copyright Copyright (c) 2022
12  * 
13  */
14 #include <lunaix/mm/kalloc.h>
15 #include <lunaix/mm/dmm.h>
16
17 #include <lunaix/common.h>
18 #include <lunaix/spike.h>
19
20 #include <klibc/string.h>
21
22 #include <stdint.h>
23
24 extern uint8_t __kernel_heap_start;
25
26 void*
27 lx_malloc_internal(heap_context_t* heap, size_t size);
28
29 void
30 place_chunk(uint8_t* ptr, size_t size);
31
32 void
33 lx_free_internal(void* ptr);
34
35 void*
36 coalesce(uint8_t* chunk_ptr);
37
38 void*
39 lx_grow_heap(heap_context_t* heap, size_t sz);
40
41 /*
42     At the beginning, we allocate an empty page and put our initial marker
43     
44     | 4/1 | 0/1 |
45     ^     ^ brk
46     start
47
48     Then, expand the heap further, with HEAP_INIT_SIZE (evaluated to 4096, i.e., 1 pg size)
49     This will allocate as much pages and override old epilogue marker with a free region hdr
50         and put new epilogue marker. These are handled by lx_grow_heap which is internally used
51         by alloc to expand the heap at many moment when needed.
52     
53     | 4/1 | 4096/0 |   .......   | 4096/0 | 0/1 |
54     ^     ^ brk_old                       ^
55     start                                 brk
56
57     Note: the brk always point to the beginning of epilogue.
58 */
59
60 static heap_context_t kheap;
61
62 int
63 kalloc_init() {
64     kheap.start = &__kernel_heap_start;
65     kheap.brk = NULL;
66     kheap.max_addr = (void*)KSTACK_START;
67
68     if (!dmm_init(&kheap)) {
69         return 0;
70     }
71
72     SW(kheap.start, PACK(4, M_ALLOCATED));
73     SW(kheap.start + WSIZE, PACK(0, M_ALLOCATED));
74     kheap.brk += WSIZE;
75
76     return lx_grow_heap(&kheap, HEAP_INIT_SIZE) != NULL;
77 }
78
79 void*
80 lxmalloc(size_t size) {
81     mutex_lock(&kheap.lock);
82     void* r = lx_malloc_internal(&kheap, size);
83     mutex_unlock(&kheap.lock);
84
85     return r;
86 }
87
88 void*
89 lxcalloc(size_t n, size_t elem) {
90     size_t pd = n * elem;
91
92     // overflow detection
93     if (pd < elem || pd < n) {
94         return NULL;
95     }
96
97     void* ptr = lxmalloc(pd);
98     if (!ptr) {
99         return NULL;
100     }
101
102     return memset(ptr, 0, pd);
103 }
104
105 void
106 lxfree(void* ptr) {
107     if (!ptr) {
108         return;
109     }
110     mutex_lock(&kheap.lock);
111
112     uint8_t* chunk_ptr = (uint8_t*)ptr - WSIZE;
113     uint32_t hdr = LW(chunk_ptr);
114     size_t sz = CHUNK_S(hdr);
115     uint8_t* next_hdr = chunk_ptr + sz;
116
117     // make sure the ptr we are 'bout to free makes sense
118     //   the size trick is stolen from glibc's malloc/malloc.c:4437 ;P
119     
120     assert_msg(((uintptr_t)ptr < (uintptr_t)(-sz)) && !((uintptr_t)ptr & 0x3),
121                "free(): invalid pointer");
122     
123     assert_msg(sz > WSIZE,
124                "free(): invalid size");
125
126     SW(chunk_ptr, hdr & ~M_ALLOCATED);
127     SW(FPTR(chunk_ptr, sz), hdr & ~M_ALLOCATED);
128     SW(next_hdr, LW(next_hdr) | M_PREV_FREE);
129     
130     coalesce(chunk_ptr);
131
132     mutex_unlock(&kheap.lock);
133 }
134
135
136 void*
137 lx_malloc_internal(heap_context_t* heap, size_t size)
138 {
139     // Simplest first fit approach.
140
141     if (!size) {
142         return NULL;
143     }
144
145     uint8_t* ptr = heap->start;
146     // round to largest 4B aligned value
147     //  and space for header
148     size = ROUNDUP(size + WSIZE, BOUNDARY);
149     while (ptr < (uint8_t*)heap->brk) {
150         uint32_t header = *((uint32_t*)ptr);
151         size_t chunk_size = CHUNK_S(header);
152         if (!chunk_size && CHUNK_A(header)) {
153             break;
154         }
155         if (chunk_size >= size && !CHUNK_A(header)) {
156             // found!
157             place_chunk(ptr, size);
158             return BPTR(ptr);
159         }
160         ptr += chunk_size;
161     }
162
163     // if heap is full (seems to be!), then allocate more space (if it's
164     // okay...)
165     if ((ptr = lx_grow_heap(heap, size))) {
166         place_chunk(ptr, size);
167         return BPTR(ptr);
168     }
169
170     // Well, we are officially OOM!
171     return NULL;
172 }
173
174 void
175 place_chunk(uint8_t* ptr, size_t size)
176 {
177     uint32_t header = *((uint32_t*)ptr);
178     size_t chunk_size = CHUNK_S(header);
179     *((uint32_t*)ptr) = PACK(size, CHUNK_PF(header) | M_ALLOCATED);
180     uint8_t* n_hdrptr = (uint8_t*)(ptr + size);
181     uint32_t diff = chunk_size - size;
182
183     if (!diff) {
184         // if the current free block is fully occupied
185         uint32_t n_hdr = LW(n_hdrptr);
186         // notify the next block about our avaliability
187         SW(n_hdrptr, n_hdr & ~0x2);
188     } else {
189         // if there is remaining free space left
190         uint32_t remainder_hdr = PACK(diff, M_NOT_ALLOCATED | M_PREV_ALLOCATED);
191         SW(n_hdrptr, remainder_hdr);
192         SW(FPTR(n_hdrptr, diff), remainder_hdr);
193
194         /*
195             | xxxx |      |         |
196
197                         |
198                         v
199                         
200             | xxxx |                |
201         */
202         coalesce(n_hdrptr);
203     }
204 }
205
206 void*
207 coalesce(uint8_t* chunk_ptr)
208 {
209     uint32_t hdr = LW(chunk_ptr);
210     uint32_t pf = CHUNK_PF(hdr);
211     uint32_t sz = CHUNK_S(hdr);
212
213     uint32_t n_hdr = LW(chunk_ptr + sz);
214
215     if (CHUNK_A(n_hdr) && pf) {
216         // case 1: prev is free
217         uint32_t prev_ftr = LW(chunk_ptr - WSIZE);
218         size_t prev_chunk_sz = CHUNK_S(prev_ftr);
219         uint32_t new_hdr = PACK(prev_chunk_sz + sz, CHUNK_PF(prev_ftr));
220         SW(chunk_ptr - prev_chunk_sz, new_hdr);
221         SW(FPTR(chunk_ptr, sz), new_hdr);
222         chunk_ptr -= prev_chunk_sz;
223     } else if (!CHUNK_A(n_hdr) && !pf) {
224         // case 2: next is free
225         size_t next_chunk_sz = CHUNK_S(n_hdr);
226         uint32_t new_hdr = PACK(next_chunk_sz + sz, pf);
227         SW(chunk_ptr, new_hdr);
228         SW(FPTR(chunk_ptr, sz + next_chunk_sz), new_hdr);
229     } else if (!CHUNK_A(n_hdr) && pf) {
230         // case 3: both free
231         uint32_t prev_ftr = LW(chunk_ptr - WSIZE);
232         size_t next_chunk_sz = CHUNK_S(n_hdr);
233         size_t prev_chunk_sz = CHUNK_S(prev_ftr);
234         uint32_t new_hdr =
235           PACK(next_chunk_sz + prev_chunk_sz + sz, CHUNK_PF(prev_ftr));
236         SW(chunk_ptr - prev_chunk_sz, new_hdr);
237         SW(FPTR(chunk_ptr, sz + next_chunk_sz), new_hdr);
238         chunk_ptr -= prev_chunk_sz;
239     }
240
241     // (fall through) case 4: prev and next are not free
242     return chunk_ptr;
243 }
244
245
246 void*
247 lx_grow_heap(heap_context_t* heap, size_t sz)
248 {
249     void* start;
250
251     // The "+ WSIZE" capture the overhead for epilogue marker
252     if (!(start = lxsbrk(heap, sz + WSIZE))) {
253         return NULL;
254     }
255     sz = ROUNDUP(sz, BOUNDARY);
256
257     // minus the overhead for epilogue, keep the invariant.
258     heap->brk -= WSIZE;
259
260     uint32_t old_marker = *((uint32_t*)start);
261     uint32_t free_hdr = PACK(sz, CHUNK_PF(old_marker));
262     SW(start, free_hdr);
263     SW(FPTR(start, sz), free_hdr);
264     SW(NEXT_CHK(start), PACK(0, M_ALLOCATED | M_PREV_FREE));
265
266     return coalesce(start);
267 }