luaur_code_gen/records/
shared_code_allocator.rs

1use crate::records::code_allocator::CodeAllocator;
2use crate::records::native_module::NativeModule;
3use crate::records::native_module_ref::NativeModuleRef;
4use crate::type_aliases::module_id::ModuleId;
5use crate::type_aliases::native_proto_exec_data_ptr::NativeProtoExecDataPtr;
6use alloc::boxed::Box;
7use alloc::collections::BTreeMap;
8use alloc::vec::Vec;
9use core::sync::atomic::{AtomicUsize, Ordering};
10
11#[derive(Debug)]
12pub struct SharedCodeAllocator {
13    pub(crate) identified_modules: BTreeMap<ModuleId, Box<NativeModule>>,
14    pub(crate) anonymous_module_count: AtomicUsize,
15    pub(crate) code_allocator: *mut CodeAllocator,
16}
17
18impl SharedCodeAllocator {
19    pub fn erase_native_module_if_unreferenced(&mut self, native_module: &NativeModule) {
20        self.shared_code_allocator_erase_native_module_if_unreferenced(native_module);
21    }
22
23    pub fn get_or_insert_native_module(
24        &mut self,
25        module_id: &ModuleId,
26        native_protos: Vec<NativeProtoExecDataPtr>,
27        data: *const u8,
28        data_size: usize,
29        code: *const u8,
30        code_size: usize,
31    ) -> (NativeModuleRef, bool) {
32        use luaur_common::FFlag;
33
34        let existing_module = self.try_get_native_module_with_lock_held(module_id);
35        if !existing_module.native_module.is_null() {
36            return (existing_module, false);
37        }
38
39        unsafe {
40            if FFlag::LuauCodegenFreeBlocks.get() {
41                let result = (*self.code_allocator).allocate(data, data_size, code, code_size);
42
43                if result.start.is_null() {
44                    return (NativeModuleRef::default(), false);
45                }
46
47                let native_module = Box::new(
48                    NativeModule::native_module_shared_code_allocator_optional_module_id_code_allocation_data_vector_native_proto_exec_data_ptr(
49                        self as *mut SharedCodeAllocator,
50                        &Some(*module_id),
51                        result,
52                        native_protos,
53                    ),
54                );
55                let native_module_ptr: *const NativeModule = &*native_module;
56                self.identified_modules.insert(*module_id, native_module);
57
58                (
59                    NativeModuleRef::native_module_ref_native_module(native_module_ptr),
60                    true,
61                )
62            } else {
63                let mut native_data = core::ptr::null_mut();
64                let mut native_data_size = 0;
65                let mut code_start = core::ptr::null_mut();
66
67                if !(*self.code_allocator).allocate_deprecated(
68                    data,
69                    data_size,
70                    code,
71                    code_size,
72                    &mut native_data,
73                    &mut native_data_size,
74                    &mut code_start,
75                ) {
76                    return (NativeModuleRef::default(), false);
77                }
78
79                let native_module = Box::new(
80                    NativeModule::native_module_shared_code_allocator_optional_module_id_u8_vector_native_proto_exec_data_ptr(
81                        self as *mut SharedCodeAllocator,
82                        &Some(*module_id),
83                        code_start,
84                        native_protos,
85                    ),
86                );
87                let native_module_ptr: *const NativeModule = &*native_module;
88                self.identified_modules.insert(*module_id, native_module);
89
90                (
91                    NativeModuleRef::native_module_ref_native_module(native_module_ptr),
92                    true,
93                )
94            }
95        }
96    }
97
98    pub fn insert_anonymous_native_module(
99        &mut self,
100        native_protos: Vec<NativeProtoExecDataPtr>,
101        data: *const u8,
102        data_size: usize,
103        code: *const u8,
104        code_size: usize,
105    ) -> NativeModuleRef {
106        self.shared_code_allocator_insert_anonymous_native_module(
107            native_protos,
108            data,
109            data_size,
110            code,
111            code_size,
112        )
113    }
114
115    pub fn operator_assign_copy(
116        &mut self,
117        _other: &SharedCodeAllocator,
118    ) -> &mut SharedCodeAllocator {
119        unreachable!("Deleted operator=");
120    }
121
122    pub fn operator_assign_move(
123        &mut self,
124        _other: &mut SharedCodeAllocator,
125    ) -> &mut SharedCodeAllocator {
126        unreachable!("Deleted operator=");
127    }
128
129    pub fn shared_code_allocator_move(&mut self) {
130        unreachable!("Deleted constructor");
131    }
132
133    pub fn shared_code_allocator_code_allocator(&mut self, code_allocator: *mut CodeAllocator) {
134        self.identified_modules.clear();
135        self.anonymous_module_count = AtomicUsize::new(0);
136        self.code_allocator = code_allocator;
137    }
138
139    pub fn shared_code_allocator_copy(&mut self) {
140        unreachable!("Deleted constructor");
141    }
142
143    pub fn try_get_native_module(&self, module_id: &ModuleId) -> NativeModuleRef {
144        self.try_get_native_module_with_lock_held(module_id)
145    }
146
147    pub fn try_get_native_module_with_lock_held(&self, module_id: &ModuleId) -> NativeModuleRef {
148        match self.identified_modules.get(module_id) {
149            Some(native_module) => unsafe {
150                NativeModuleRef::native_module_ref_native_module(&**native_module)
151            },
152            None => NativeModuleRef::default(),
153        }
154    }
155}
156
157impl Default for SharedCodeAllocator {
158    fn default() -> Self {
159        Self {
160            identified_modules: BTreeMap::new(),
161            anonymous_module_count: AtomicUsize::new(0),
162            code_allocator: core::ptr::null_mut(),
163        }
164    }
165}
166
167impl Drop for SharedCodeAllocator {
168    fn drop(&mut self) {
169        crate::macros::codegen_assert::CODEGEN_ASSERT!(self.identified_modules.is_empty());
170        crate::macros::codegen_assert::CODEGEN_ASSERT!(
171            self.anonymous_module_count.load(Ordering::Relaxed) == 0
172        );
173    }
174}
luaur_code_gen/records/shared_code_allocator.rs

luaur_code_gen/records/
shared_code_allocator.rs
