Indicate how bytecode serializer allocates return buffer

This change allows for callers to ByteCodeSerializer to specify if the buffer containing the serialized bytecode (and parser cache state) should be allocated with ANew or CoTaskMemAlloc. The caller is then responsible for freeing the memory appropriately after use.

Author: Thomas Moore (CHAKRA)

lib/Jsrt/Jsrt.cpp

@@ -3819,7 +3819,7 @@ JsErrorCode JsSerializeScriptCore(const byte *script, size_t cb,
         // However, the PAL defines DWORD for us on linux as unsigned int so the cast is safe here.
         HRESULT hr = Js::ByteCodeSerializer::SerializeToBuffer(scriptContext,
             tempAllocator, static_cast<DWORD>(cSourceCodeLength), utf8Code,
-            functionBody, functionBody->GetHostSrcInfo(), false, &buffer,
+            functionBody, functionBody->GetHostSrcInfo(), &buffer,
             (DWORD*) bufferSize, dwFlags);
         END_TEMP_ALLOCATOR(tempAllocator, scriptContext);
 

lib/Parser/BGParseManager.cpp

@@ -541,10 +541,9 @@ void BGParseWorkItem::ParseUTF8Core(Js::ScriptContext* scriptContext)
             this->script,
             functionBody,
             functionBody->GetHostSrcInfo(),
-            true, // allocateBuffer
             &this->bufferReturn,
             &this->bufferReturnBytes,
-            GENERATE_BYTE_CODE_PARSER_STATE
+            GENERATE_BYTE_CODE_PARSER_STATE | GENERATE_BYTE_CODE_COTASKMEMALLOC
         );
         END_TEMP_ALLOCATOR(tempAllocator, scriptContext);
         Assert(this->parseHR == S_OK);

lib/Runtime/Base/ScriptContext.cpp

@@ -2229,7 +2229,7 @@ namespace Js
 #ifdef ENABLE_WININET_PROFILE_DATA_CACHE
         byte* serializeParserStateCacheBuffer = parserStateCacheBuffer;
         DWORD serializeParserStateCacheSize = parserStateCacheByteCount;
-        DWORD dwFlags = GENERATE_BYTE_CODE_PARSER_STATE;
+        DWORD dwFlags = GENERATE_BYTE_CODE_PARSER_STATE | GENERATE_BYTE_CODE_ALLOC_ANEW;
         DebugOnly(auto url = !srcInfo->sourceContextInfo->isHostDynamicDocument ? srcInfo->sourceContextInfo->url : this->GetUrl());
 
         // If we already have a parser state cache serialized into a buffer, we should skip creating it again
@@ -2242,7 +2242,7 @@ namespace Js
             BEGIN_TEMP_ALLOCATOR(tempAllocator, this, _u("ByteCodeSerializer"));
             hr = Js::ByteCodeSerializer::SerializeToBuffer(this,
                 tempAllocator, (DWORD)cbLength, pszSrc, func->GetFunctionBody(),
-                func->GetHostSrcInfo(), true, &serializeParserStateCacheBuffer,
+                func->GetHostSrcInfo(), &serializeParserStateCacheBuffer,
                 &serializeParserStateCacheSize, dwFlags);
             END_TEMP_ALLOCATOR(tempAllocator, this);
 
@@ -2410,7 +2410,7 @@ namespace Js
 
                 return Js::ByteCodeSerializer::SerializeToBuffer(this,
                     alloc, static_cast<DWORD>(cSourceCodeLength), utf8Code,
-                    functionBody, functionBody->GetHostSrcInfo(), false, buffer,
+                    functionBody, functionBody->GetHostSrcInfo(), buffer,
                     bufferSize, dwFlags);
             }
             else

lib/Runtime/ByteCode/ByteCodeSerializeFlags.h

@@ -4,6 +4,11 @@
 //-------------------------------------------------------------------------------------------------------
 #pragma once
 
-#define GENERATE_BYTE_CODE_BUFFER_LIBRARY 0x00000001
-#define GENERATE_BYTE_CODE_FOR_NATIVE 0x00000002
-#define GENERATE_BYTE_CODE_PARSER_STATE 0x00000004
+#define GENERATE_BYTE_CODE_BUFFER_LIBRARY   0x00000001
+#define GENERATE_BYTE_CODE_FOR_NATIVE       0x00000002
+#define GENERATE_BYTE_CODE_PARSER_STATE     0x00000004
+// When the return buffer needs to be allocated while generating bytecode, use one
+// of these flags to indicate how to allocate the memory. The absence of both flags
+// indicates that no allocation is needed.
+#define GENERATE_BYTE_CODE_COTASKMEMALLOC   0x00000008
+#define GENERATE_BYTE_CODE_ALLOC_ANEW       0x00000010

lib/Runtime/ByteCode/ByteCodeSerializer.cpp

@@ -420,6 +420,16 @@ class ByteCodeBufferBuilder
         return (dwFlags & GENERATE_BYTE_CODE_PARSER_STATE) != 0;
     }
 
+    bool ShouldAllocWithCoTaskMem() const
+    {
+        return (dwFlags & GENERATE_BYTE_CODE_COTASKMEMALLOC);
+    }
+
+    bool ShouldAllocWithANew() const
+    {
+        return (dwFlags & GENERATE_BYTE_CODE_ALLOC_ANEW);
+    }
+
 public:
 
     ByteCodeBufferBuilder(uint32 sourceSize, uint32 sourceCharLength, LPCUTF8 utf8Source, Utf8SourceInfo* sourceInfo, ScriptContext * scriptContext, ArenaAllocator * alloc, DWORD dwFlags, int builtInPropertyCount)
@@ -528,7 +538,7 @@ class ByteCodeBufferBuilder
         string16ToId = Anew(alloc, TString16ToId, alloc);
     }
 
-    HRESULT Create(bool allocateBuffer, byte ** buffer, DWORD * bufferBytes)
+    HRESULT Create(byte ** buffer, DWORD * bufferBytes)
     {
         BufferBuilderList all(_u("Final"));
 
@@ -574,15 +584,16 @@ class ByteCodeBufferBuilder
         totalSize.value = size;
 
         // Allocate the bytes
-        if (allocateBuffer)
+        if (ShouldAllocWithANew() || ShouldAllocWithCoTaskMem())
         {
             *bufferBytes = size;
-            if (GenerateParserStateCache())
+            if (ShouldAllocWithANew())
             {
                 *buffer = AnewArray(scriptContext->SourceCodeAllocator(), byte, *bufferBytes);
             }
             else
             {
+                Assert(ShouldAllocWithCoTaskMem());
                 *buffer = (byte*)CoTaskMemAlloc(*bufferBytes);
             }
 
@@ -4701,7 +4712,7 @@ FunctionInfo* ByteCodeCache::LookupFunctionInfo(ScriptContext * scriptContext, L
 }
 
 // Serialize function body
-HRESULT ByteCodeSerializer::SerializeToBuffer(ScriptContext * scriptContext, ArenaAllocator * alloc, DWORD sourceByteLength, LPCUTF8 utf8Source, FunctionBody * function, SRCINFO const* srcInfo, bool allocateBuffer, byte ** buffer, DWORD * bufferBytes, DWORD dwFlags)
+HRESULT ByteCodeSerializer::SerializeToBuffer(ScriptContext * scriptContext, ArenaAllocator * alloc, DWORD sourceByteLength, LPCUTF8 utf8Source, FunctionBody * function, SRCINFO const* srcInfo, byte ** buffer, DWORD * bufferBytes, DWORD dwFlags)
 {
     int builtInPropertyCount = (dwFlags & GENERATE_BYTE_CODE_BUFFER_LIBRARY) != 0 ?  PropertyIds::_countJSOnlyProperty : TotalNumberOfBuiltInProperties;
 
@@ -4715,7 +4726,8 @@ HRESULT ByteCodeSerializer::SerializeToBuffer(ScriptContext * scriptContext, Are
 
     ArenaAllocator* codeAllocator = nullptr;
     ByteCodeCache* cache = nullptr;
-    if (((dwFlags & GENERATE_BYTE_CODE_PARSER_STATE) != 0) && allocateBuffer)
+    DWORD shouldUseCodeAllocator = GENERATE_BYTE_CODE_PARSER_STATE | GENERATE_BYTE_CODE_ALLOC_ANEW;
+    if ((dwFlags & shouldUseCodeAllocator) == shouldUseCodeAllocator) // does this apply for cotaskmemalloc as well?
     {
         codeAllocator = scriptContext->SourceCodeAllocator();
         cache = Anew(codeAllocator, ByteCodeCache, scriptContext, builtInPropertyCount);
@@ -4728,7 +4740,7 @@ HRESULT ByteCodeSerializer::SerializeToBuffer(ScriptContext * scriptContext, Are
 
     if (SUCCEEDED(hr))
     {
-        hr = builder.Create(allocateBuffer, buffer, bufferBytes);
+        hr = builder.Create(buffer, bufferBytes);
     }
 
     if (SUCCEEDED(hr) && cache != nullptr)

lib/Runtime/ByteCode/ByteCodeSerializer.h

@@ -167,7 +167,7 @@ namespace Js
     struct ByteCodeSerializer
     {
         // Serialize a function body.
-        static HRESULT SerializeToBuffer(ScriptContext * scriptContext, ArenaAllocator * alloc, DWORD sourceByteLength, LPCUTF8 utf8Source, FunctionBody * function, SRCINFO const* srcInfo, bool allocateBuffer, byte ** buffer, DWORD * bufferBytes, DWORD dwFlags = 0);
+        static HRESULT SerializeToBuffer(ScriptContext * scriptContext, ArenaAllocator * alloc, DWORD sourceByteLength, LPCUTF8 utf8Source, FunctionBody * function, SRCINFO const* srcInfo, byte ** buffer, DWORD * bufferBytes, DWORD dwFlags = 0);
 
         // Deserialize a function body. The content of utf8Source must be the same as was originally passed to SerializeToBuffer
         static HRESULT DeserializeFromBuffer(ScriptContext * scriptContext, uint32 scriptFlags, LPCUTF8 utf8Source, SRCINFO const * srcInfo, byte * buffer, NativeModule *nativeModule, Field(FunctionBody*)* function, uint sourceIndex = Js::Constants::InvalidSourceIndex);