55
66#include < string.h>
77
8+ LUAU_FASTFLAGVARIABLE (LuauCodegenFreeBlocks)
9+
810#if defined(_WIN32)
911
1012#ifndef WIN32_LEAN_AND_MEAN
@@ -61,6 +63,15 @@ static void freePagesImpl(uint8_t* mem, size_t size)
6163 return VirtualProtect (mem, size, PAGE_EXECUTE_READ , &oldProtect) != 0 ;
6264}
6365
66+ [[nodiscard]] static bool makePagesNotExecutable (uint8_t * mem, size_t size)
67+ {
68+ CODEGEN_ASSERT ((uintptr_t (mem) & (kPageSize - 1 )) == 0 );
69+ CODEGEN_ASSERT (size == alignToPageSize (size));
70+
71+ DWORD oldProtect;
72+ return VirtualProtect (mem, size, PAGE_READWRITE , &oldProtect) != 0 ;
73+ }
74+
6475static void flushInstructionCache (uint8_t * mem, size_t size)
6576{
6677#if WINAPI_FAMILY_PARTITION(WINAPI_PARTITION_APP | WINAPI_PARTITION_SYSTEM)
@@ -98,6 +109,14 @@ static void freePagesImpl(uint8_t* mem, size_t size)
98109 return mprotect (mem, size, PROT_READ | PROT_EXEC ) == 0 ;
99110}
100111
112+ [[nodiscard]] static bool makePagesNotExecutable (uint8_t * mem, size_t size)
113+ {
114+ CODEGEN_ASSERT ((uintptr_t (mem) & (kPageSize - 1 )) == 0 );
115+ CODEGEN_ASSERT (size == alignToPageSize (size));
116+
117+ return mprotect (mem, size, PROT_READ | PROT_WRITE ) == 0 ;
118+ }
119+
101120static void flushInstructionCache (uint8_t * mem, size_t size)
102121{
103122#ifdef __APPLE__
@@ -136,11 +155,14 @@ CodeAllocator::~CodeAllocator()
136155 destroyBlockUnwindInfo (context, unwindInfo);
137156 }
138157
158+ if (FFlag::LuauCodegenFreeBlocks)
159+ CODEGEN_ASSERT (liveAllocations == 0 );
160+
139161 for (uint8_t * block : blocks)
140162 freePages (block, blockSize);
141163}
142164
143- bool CodeAllocator::allocate (
165+ bool CodeAllocator::allocate_DEPRECATED (
144166 const uint8_t * data,
145167 size_t dataSize,
146168 const uint8_t * code,
@@ -150,6 +172,8 @@ bool CodeAllocator::allocate(
150172 uint8_t *& resultCodeStart
151173)
152174{
175+ CODEGEN_ASSERT (!FFlag::LuauCodegenFreeBlocks);
176+
153177 // 'Round up' to preserve code alignment
154178 size_t alignedDataSize = (dataSize + (kCodeAlignment - 1 )) & ~(kCodeAlignment - 1 );
155179
@@ -209,6 +233,92 @@ bool CodeAllocator::allocate(
209233 return true ;
210234}
211235
236+ CodeAllocationData CodeAllocator::allocate (const uint8_t * data, size_t dataSize, const uint8_t * code, size_t codeSize)
237+ {
238+ CODEGEN_ASSERT (FFlag::LuauCodegenFreeBlocks);
239+
240+ // 'Round up' to preserve code alignment
241+ size_t alignedDataSize = (dataSize + (kCodeAlignment - 1 )) & ~(kCodeAlignment - 1 );
242+
243+ size_t totalSize = alignedDataSize + codeSize;
244+
245+ // Function has to fit into a single block with unwinding information
246+ if (totalSize > blockSize - kMaxReservedDataSize )
247+ return {};
248+
249+ size_t startOffset = 0 ;
250+
251+ // We might need a new block
252+ if (totalSize > size_t (blockEnd - blockPos))
253+ {
254+ if (!allocateNewBlock (startOffset))
255+ return {};
256+
257+ CODEGEN_ASSERT (totalSize <= size_t (blockEnd - blockPos));
258+ }
259+
260+ CODEGEN_ASSERT ((uintptr_t (blockPos) & (kPageSize - 1 )) == 0 ); // Allocation starts on page boundary
261+
262+ size_t dataOffset = startOffset + alignedDataSize - dataSize;
263+ size_t codeOffset = startOffset + alignedDataSize;
264+
265+ if (dataSize != 0 )
266+ memcpy (blockPos + dataOffset, data, dataSize);
267+ if (codeSize != 0 )
268+ memcpy (blockPos + codeOffset, code, codeSize);
269+
270+ size_t pageAlignedSize = alignToPageSize (startOffset + totalSize);
271+
272+ if (!makePagesExecutable (blockPos, pageAlignedSize))
273+ return {};
274+
275+ liveAllocations++;
276+
277+ flushInstructionCache (blockPos + codeOffset, codeSize);
278+
279+ CodeAllocationData result;
280+
281+ result.start = blockPos + startOffset;
282+ result.size = totalSize;
283+ result.codeStart = blockPos + codeOffset;
284+
285+ result.allocationStart = blockPos;
286+ result.allocationSize = pageAlignedSize;
287+
288+ // Ensure that future allocations from the block start from a page boundary.
289+ // This is important since we use W^X, and writing to the previous page would require briefly removing
290+ // executable bit from it, which may result in access violations if that code is being executed concurrently.
291+ if (pageAlignedSize <= size_t (blockEnd - blockPos))
292+ {
293+ blockPos += pageAlignedSize;
294+ CODEGEN_ASSERT ((uintptr_t (blockPos) & (kPageSize - 1 )) == 0 );
295+ CODEGEN_ASSERT (blockPos <= blockEnd);
296+ }
297+ else
298+ {
299+ // Future allocations will need to allocate fresh blocks
300+ blockPos = blockEnd;
301+ }
302+
303+ return result;
304+ }
305+
306+ void CodeAllocator::deallocate (CodeAllocationData codeAllocationData)
307+ {
308+ CODEGEN_ASSERT (FFlag::LuauCodegenFreeBlocks);
309+
310+ if (codeAllocationData.allocationStart == nullptr )
311+ return ;
312+
313+ [[maybe_unused]] bool result = makePagesNotExecutable (codeAllocationData.allocationStart , codeAllocationData.allocationSize );
314+ CODEGEN_ASSERT (result);
315+
316+ CODEGEN_ASSERT (liveAllocations != 0 );
317+ liveAllocations--;
318+
319+ // TODO: new allocations should be able to reuse the freed pages (but note that first block page contains unwind data)
320+ }
321+
212322bool CodeAllocator::allocateNewBlock (size_t & unwindInfoSize)
213323{
214324 // Stop allocating once we reach a global limit
0 commit comments