1#include <asm.h>
2#include <enforce.h>
3#include <err.h>
4#include <math/align.h>
5#include <mem/anon_vma.h>
6#include <mem/fixed_size_alloc.h>
7#include <mem/mm.h>
8#include <mem/page.h>
9#include <mem/vmm.h>
10#include <smp/perdomain.h>
11#include <structures/rbt.h>
12#include <types/refcount.h>
13
14#define MM_USER_MIN 0x0000000000010000UL
15#define MM_USER_MAX 0x0000800000000000UL
16
17FIXED_SIZE_RANGE_PERDOMAIN_DEFINE(mm, .obj_size = sizeof(struct mm),
18 .obj_align = cc_alignof(struct mm));
19
20/* VMA tree is rbit keyed [vma_range_start, vma_range_end - 1], carrying custom
21 * agumentation to track max_gap to make gap search O(log n)
22 *
23 * Non overlapping VMAs make 'left subtree max end' exactly the
24 * predecessor's end, so no neighbor ptr needed here */
25static size_t vma_range_node_min_low(struct rbit_node *n) {
26 return n ? rbit_entry(n, struct vma_range, mm_node)->min_low : 0;
27}
28
29static size_t vma_range_node_gap(struct rbit_node *n) {
30 return n ? rbit_entry(n, struct vma_range, mm_node)->max_gap : 0;
31}
32
33static bool vma_range_tree_augment(struct rbit_node *n) {
34 struct vma_range *v = rbit_entry(n, struct vma_range, mm_node);
35 size_t old_max = n->max, old_min = v->min_low, old_gap = v->max_gap;
36
37 /* node->max: subtree max of interval.high (= max end - 1) */
38 size_t mx = n->interval.high;
39 if (rbit_node_max(n: n->left) > mx)
40 mx = rbit_node_max(n: n->left);
41
42 if (rbit_node_max(n: n->right) > mx)
43 mx = rbit_node_max(n: n->right);
44
45 size_t mn = n->left ? vma_range_node_min_low(n: n->left) : n->interval.low;
46
47 /* max_gap: biggest free run between consecutive VMAs in the subtree */
48 size_t gap = 0;
49 if (n->left) {
50 if (vma_range_node_gap(n: n->left) > gap)
51 gap = vma_range_node_gap(n: n->left);
52
53 size_t pred_end = rbit_node_max(n: n->left) + 1; /* predecessor's end */
54 if (n->interval.low > pred_end && n->interval.low - pred_end > gap)
55 gap = n->interval.low - pred_end;
56 }
57 if (n->right) {
58 if (vma_range_node_gap(n: n->right) > gap)
59 gap = vma_range_node_gap(n: n->right);
60
61 size_t end = n->interval.high + 1; /* this VMA's end */
62 size_t succ_low = vma_range_node_min_low(n: n->right);
63 if (succ_low > end && succ_low - end > gap)
64 gap = succ_low - end;
65 }
66
67 n->max = mx;
68 v->min_low = mn;
69 v->max_gap = gap;
70 return mx != old_max || mn != old_min || gap != old_gap;
71}
72
73enum err mm_pgtable_init(struct mm *mm) {
74 mm->pml4 = vmm_make_user_pml4();
75 return mm->pml4 ? ERR_OK : ERR_NO_MEM;
76}
77
78void mm_pgtable_free(struct mm *mm) {
79 vmm_unmap_all_user_pages(vmm_phys_to_pml4(mm->pml4));
80}
81
82void mm_activate(struct mm *mm) {
83 cr3_write(cr3: mm->pml4);
84}
85
86enum err mm_map_page(struct mm *mm, vaddr_t va, paddr_t pa, uint64_t pflags) {
87 vmm_map_page_user(vmm_phys_to_pml4(mm->pml4), va, pa, .page_flags = pflags,
88 .vmm_flags = VMM_FLAG_USER);
89 return ERR_OK;
90}
91
92struct mm *mm_alloc(void) {
93 struct mm *mm = FSR_PERDOMAIN_ALLOC(mm);
94 if (!mm)
95 return NULL;
96
97 rbit_init(rbit: &mm->vma_range_tree);
98 mm->vma_range_tree.augment =
99 vma_range_tree_augment; /* O(log n) gap search */
100 /* Anyone can be under this rwlock */
101 rwlock_init(&mm->lock, THREAD_PRIO_CLASS_URGENT);
102 mm->vas = NULL;
103 mm->mmap_cursor = MM_USER_MIN;
104
105 refcount_init(rc: &mm->users, val: 1);
106 refcount_init(rc: &mm->refcount, val: 1);
107
108 if (mm_pgtable_init(mm) != ERR_OK) {
109 FSR_PERDOMAIN_FREE(mm, mm);
110 return NULL;
111 }
112 return mm;
113}
114
115void mm_free(struct mm *mm) {
116 enforce(rbit_empty(&mm->vma_range_tree));
117 mm_pgtable_free(mm);
118 FSR_PERDOMAIN_FREE(mm, mm);
119}
120
121void mm_exit(struct mm *mm) {
122 struct rbit_node *n = rbit_first(tree: &mm->vma_range_tree);
123 while (n) {
124 struct rbit_node *next = rbit_next(node: n);
125 struct vma_range *vma_range = rbit_entry(n, struct vma_range, mm_node);
126
127 mm_vma_range_remove(mm, vma_range);
128 vma_range_unlink_anon_vmas(vma_range); /* drops AVCs + anon_vma refs */
129 /* TODO: drop this VMA's PTEs + folio rmap once page
130 * tables can be torn down. For now the mappings die with the pml4 */
131 vma_range_free(vma_range);
132
133 n = next;
134 }
135}
136
137struct vma_range *mm_vma_range_find(struct mm *mm, vaddr_t addr) {
138 struct rbit_node *node = mm->vma_range_tree.root;
139 struct vma_range *result = NULL;
140
141 while (node) {
142 struct vma_range *v = rbit_entry(node, struct vma_range, mm_node);
143 if (vma_range_end(vma_range: v) > addr) {
144 result = v; /* candidate; look left for an earlier one */
145 node = node->left;
146 } else {
147 node = node->right;
148 }
149 }
150 return result;
151}
152
153struct vma_range *mm_vma_range_find_intersection(struct mm *mm, vaddr_t s,
154 vaddr_t e) {
155 /* Inclusive interval [s, e-1] matches the half-open [s, e) reservation */
156 struct interval iv = {.low = s, .high = e - 1};
157 struct rbit_node *n = rbit_overlap_search(root: mm->vma_range_tree.root, iv);
158 return n ? rbit_entry(n, struct vma_range, mm_node) : NULL;
159}
160
161void mm_vma_range_insert(struct mm *mm, struct vma_range *vma_range) {
162 /* vma_range_init() already set the interval; rbit_insert() does the rest */
163 enforce(!mm_vma_range_find_intersection(mm, vma_range_start(vma_range),
164 vma_range_end(vma_range)));
165 rbit_insert(tree: &mm->vma_range_tree, new_node: &vma_range->mm_node);
166}
167
168void mm_vma_range_remove(struct mm *mm, struct vma_range *vma_range) {
169 rbit_delete(tree: &mm->vma_range_tree, z: &vma_range->mm_node);
170}
171
172struct gap_ctx {
173 size_t len;
174 size_t align;
175 vaddr_t high;
176};
177
178static vaddr_t gap_fits(vaddr_t floor, vaddr_t ceil, const struct gap_ctx *c) {
179 if (ceil <= floor)
180 return 0;
181 vaddr_t a = ALIGN_UP_WRAPPING(floor, c->align);
182 if (a < floor) /* ALIGN_UP wrapped */
183 return 0;
184 if (a + c->len < a) /* len overflow */
185 return 0;
186 if (a + c->len > ceil || a + c->len > c->high)
187 return 0;
188 return a;
189}
190
191static vaddr_t gap_search(struct rbit_node *n, vaddr_t *floor,
192 const struct gap_ctx *c) {
193 if (!n)
194 return 0;
195 struct vma_range *v = rbit_entry(n, struct vma_range, mm_node);
196
197 if (n->left) {
198 if (vma_range_node_gap(n: n->left) >= c->len ||
199 vma_range_node_min_low(n: n->left) > *floor) {
200 vaddr_t r = gap_search(n: n->left, floor, c);
201 if (r)
202 return r;
203 } else { /* prune: jump floor to the subtree's max end */
204 vaddr_t e = rbit_node_max(n: n->left) + 1;
205 if (e > *floor)
206 *floor = e;
207 }
208 }
209
210 /* gap right before this VMA */
211 vaddr_t r = gap_fits(floor: *floor, ceil: vma_range_start(vma_range: v), c);
212 if (r)
213 return r;
214 if (vma_range_start(vma_range: v) >=
215 c->high) /* this VMA and all to its right are out */
216 return 0;
217 if (vma_range_end(vma_range: v) > *floor)
218 *floor = vma_range_end(vma_range: v);
219
220 if (n->right) {
221 if (vma_range_node_gap(n: n->right) >= c->len ||
222 vma_range_node_min_low(n: n->right) > *floor) {
223 vaddr_t rr = gap_search(n: n->right, floor, c);
224 if (rr)
225 return rr;
226 } else {
227 vaddr_t e = rbit_node_max(n: n->right) + 1;
228 if (e > *floor)
229 *floor = e;
230 }
231 }
232 return 0;
233}
234
235vaddr_t mm_vma_range_find_gap(struct mm *mm, size_t len, size_t align,
236 vaddr_t low, vaddr_t high) {
237 enforce_precond(IS_POW2(align));
238 if (len == 0 || high <= low || high - low < len)
239 return 0;
240
241 struct gap_ctx c = {.len = len, .align = align, .high = high};
242 vaddr_t floor = low;
243
244 vaddr_t r = gap_search(n: mm->vma_range_tree.root, floor: &floor, c: &c);
245 if (r)
246 return r;
247
248 /* trailing gap [floor, high) */
249 return gap_fits(floor, ceil: high, c: &c);
250}
251
252vaddr_t mm_map(struct mm *mm, vaddr_t hint, size_t len,
253 enum vma_range_protection prot, enum mm_map_flags flags) {
254 len = PAGE_ALIGN_UP_WRAPPING(len);
255 if (len == 0)
256 return 0;
257
258 rw_write_lock(&mm->lock);
259
260 vaddr_t addr;
261 if (flags & MM_MAP_FIXED) {
262 addr = PAGE_ALIGN_DOWN(hint);
263 if (mm_vma_range_find_intersection(mm, s: addr, e: addr + len)) {
264 /* TODO: clobber overlap via mm_unmap()
265 * once it can tear down PTEs */
266
267 /* FIXED map on busy range fails */
268 rw_unlock(&mm->lock);
269 return 0;
270 }
271 } else {
272 vaddr_t low = hint ? PAGE_ALIGN_UP_WRAPPING(hint) : mm->mmap_cursor;
273 if (hint && low < hint) {
274 rw_unlock(&mm->lock);
275 return 0;
276 }
277 addr = mm_vma_range_find_gap(mm, len, PAGE_SIZE, low, MM_USER_MAX);
278 if (!addr)
279 addr = mm_vma_range_find_gap(mm, len, PAGE_SIZE, MM_USER_MIN,
280 MM_USER_MAX);
281
282 if (!addr) {
283 rw_unlock(&mm->lock);
284 return 0;
285 }
286 if (!hint) /* advance the cursor past what we just handed out */
287 mm->mmap_cursor = addr + len;
288 }
289
290 /* anon reservation only, pages fault in lazily, so no PTEs */
291 struct vma_range *vma_range = vma_range_alloc(mm, start: addr, end: addr + len, prot);
292 if (!vma_range) {
293 rw_unlock(&mm->lock);
294 return 0;
295 }
296 mm_vma_range_insert(mm, vma_range);
297
298 rw_unlock(&mm->lock);
299 return addr;
300}
301