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