1 files changed, 151 insertions, 29 deletions
diff --git a/source/Plugins/LanguageRuntime/RenderScript/RenderScriptRuntime/RenderScriptRuntime.cpp b/source/Plugins/LanguageRuntime/RenderScript/RenderScriptRuntime/RenderScriptRuntime.cpp
index 5d82ded..8e5d31b 100644
--- a/source/Plugins/LanguageRuntime/RenderScript/RenderScriptRuntime/RenderScriptRuntime.cpp
+++ b/source/Plugins/LanguageRuntime/RenderScript/RenderScriptRuntime/RenderScriptRuntime.cpp
@@ -235,16 +235,28 @@ struct RenderScriptRuntime::AllocationDetails
         }
     };
 
-    // Header for reading and writing allocation contents
-    // to a binary file.
+    // The FileHeader struct specifies the header we use for writing allocations to a binary file.
+    // Our format begins with the ASCII characters "RSAD", identifying the file as an allocation dump.
+    // Member variables dims and hdr_size are then written consecutively, immediately followed by an instance of
+    // the ElementHeader struct. Because Elements can contain subelements, there may be more than one instance
+    // of the ElementHeader struct. With this first instance being the root element, and the other instances being
+    // the root's descendants. To identify which instances are an ElementHeader's children, each struct
+    // is immediately followed by a sequence of consecutive offsets to the start of its child structs.
+    // These offsets are 4 bytes in size, and the 0 offset signifies no more children.
     struct FileHeader
     {
         uint8_t ident[4];      // ASCII 'RSAD' identifying the file
-        uint16_t hdr_size;     // Header size in bytes, for backwards compatability
-        uint16_t type;         // DataType enum
-        uint32_t kind;         // DataKind enum
         uint32_t dims[3];      // Dimensions
-        uint32_t element_size; // Size of a single element, including padding
+        uint16_t hdr_size;     // Header size in bytes, including all element headers
+    };
+
+    struct ElementHeader
+    {
+        uint16_t type;          // DataType enum
+        uint32_t kind;          // DataKind enum
+        uint32_t element_size;  // Size of a single element, including padding
+        uint16_t vector_size;   // Vector width
+        uint32_t array_size;    // Number of elements in array
     };
 
     // Monotonically increasing from 1
@@ -286,7 +298,6 @@ struct RenderScriptRuntime::AllocationDetails
     }
 };
 
-
 const ConstString &
 RenderScriptRuntime::Element::GetFallbackStructName()
 {
@@ -2084,37 +2095,62 @@ RenderScriptRuntime::LoadAllocation(Stream &strm, const uint32_t alloc_id, const
 
     // Cast start of buffer to FileHeader and use pointer to read metadata
     void* file_buffer = data_sp->GetBytes();
-    const AllocationDetails::FileHeader* head = static_cast<AllocationDetails::FileHeader*>(file_buffer);
+    if (file_buffer == NULL || data_sp->GetByteSize() <
+        (sizeof(AllocationDetails::FileHeader) + sizeof(AllocationDetails::ElementHeader)))
+    {
+        strm.Printf("Error: File %s does not contain enough data for header", filename);
+        strm.EOL();
+        return false;
+    }
+    const AllocationDetails::FileHeader* file_header = static_cast<AllocationDetails::FileHeader*>(file_buffer);
 
-    // Advance buffer past header
-    file_buffer = static_cast<uint8_t*>(file_buffer) + head->hdr_size;
+    // Check file starts with ascii characters "RSAD"
+    if (file_header->ident[0] != 'R' || file_header->ident[1] != 'S' || file_header->ident[2] != 'A'
+        || file_header->ident[3] != 'D')
+    {
+        strm.Printf("Error: File doesn't contain identifier for an RS allocation dump. Are you sure this is the correct file?");
+        strm.EOL();
+        return false;
+    }
+
+    // Look at the type of the root element in the header
+    AllocationDetails::ElementHeader root_element_header;
+    memcpy(&root_element_header, static_cast<uint8_t*>(file_buffer) + sizeof(AllocationDetails::FileHeader),
+           sizeof(AllocationDetails::ElementHeader));
 
     if (log)
         log->Printf("RenderScriptRuntime::LoadAllocation - header type %u, element size %u",
-                    head->type, head->element_size);
+                    root_element_header.type, root_element_header.element_size);
 
     // Check if the target allocation and file both have the same number of bytes for an Element
-    if (*alloc->element.datum_size.get() != head->element_size)
+    if (*alloc->element.datum_size.get() != root_element_header.element_size)
     {
         strm.Printf("Warning: Mismatched Element sizes - file %u bytes, allocation %u bytes",
-                    head->element_size, *alloc->element.datum_size.get());
+                    root_element_header.element_size, *alloc->element.datum_size.get());
         strm.EOL();
     }
 
-    // Check if the target allocation and file both have the same integral type
-    const unsigned int type = static_cast<unsigned int>(*alloc->element.type.get());
-    if (type != head->type)
+    // Check if the target allocation and file both have the same type
+    const unsigned int alloc_type = static_cast<unsigned int>(*alloc->element.type.get());
+    const unsigned int file_type = root_element_header.type;
+
+    if (file_type > Element::RS_TYPE_FONT)
+    {
+        strm.Printf("Warning: File has unknown allocation type");
+        strm.EOL();
+    }
+    else if (alloc_type != file_type)
     {
         // Enum value isn't monotonous, so doesn't always index RsDataTypeToString array
-        unsigned int printable_target_type_index = type;
-        unsigned int printable_head_type_index = head->type;
-        if (type >= Element::RS_TYPE_ELEMENT && type <= Element::RS_TYPE_FONT)
+        unsigned int printable_target_type_index = alloc_type;
+        unsigned int printable_head_type_index = file_type;
+        if (alloc_type >= Element::RS_TYPE_ELEMENT && alloc_type <= Element::RS_TYPE_FONT)
             printable_target_type_index = static_cast<Element::DataType>(
-                                         (type - Element::RS_TYPE_ELEMENT) + Element::RS_TYPE_MATRIX_2X2 + 1);
+                                         (alloc_type - Element::RS_TYPE_ELEMENT) + Element::RS_TYPE_MATRIX_2X2 + 1);
 
-        if (head->type >= Element::RS_TYPE_ELEMENT && head->type <= Element::RS_TYPE_FONT)
+        if (file_type >= Element::RS_TYPE_ELEMENT && file_type <= Element::RS_TYPE_FONT)
             printable_head_type_index = static_cast<Element::DataType>(
-                                        (head->type - Element::RS_TYPE_ELEMENT) + Element::RS_TYPE_MATRIX_2X2 + 1);
+                                        (file_type - Element::RS_TYPE_ELEMENT) + Element::RS_TYPE_MATRIX_2X2 + 1);
 
         const char* file_type_cstr = AllocationDetails::RsDataTypeToString[printable_head_type_index][0];
         const char* target_type_cstr = AllocationDetails::RsDataTypeToString[printable_target_type_index][0];
@@ -2124,8 +2160,11 @@ RenderScriptRuntime::LoadAllocation(Stream &strm, const uint32_t alloc_id, const
         strm.EOL();
     }
 
+    // Advance buffer past header
+    file_buffer = static_cast<uint8_t*>(file_buffer) + file_header->hdr_size;
+
     // Calculate size of allocation data in file
-    size_t length = data_sp->GetByteSize() - head->hdr_size;
+    size_t length = data_sp->GetByteSize() - file_header->hdr_size;
 
     // Check if the target allocation and file both have the same total data size.
     const unsigned int alloc_size = *alloc->size.get();
@@ -2154,6 +2193,62 @@ RenderScriptRuntime::LoadAllocation(Stream &strm, const uint32_t alloc_id, const
     return true;
 }
 
+// Function takes as parameters a byte buffer, which will eventually be written to file as the element header,
+// an offset into that buffer, and an Element that will be saved into the buffer at the parametrised offset.
+// Return value is the new offset after writing the element into the buffer.
+// Elements are saved to the file as the ElementHeader struct followed by offsets to the structs of all the element's children.
+size_t
+RenderScriptRuntime::PopulateElementHeaders(const std::shared_ptr<uint8_t> header_buffer, size_t offset, const Element& elem)
+{
+    // File struct for an element header with all the relevant details copied from elem.
+    // We assume members are valid already.
+    AllocationDetails::ElementHeader elem_header;
+    elem_header.type = *elem.type.get();
+    elem_header.kind = *elem.type_kind.get();
+    elem_header.element_size = *elem.datum_size.get();
+    elem_header.vector_size = *elem.type_vec_size.get();
+    elem_header.array_size = elem.array_size.isValid() ? *elem.array_size.get() : 0;
+    const size_t elem_header_size = sizeof(AllocationDetails::ElementHeader);
+
+    // Copy struct into buffer and advance offset
+    // We assume that header_buffer has been checked for NULL before this method is called
+    memcpy(header_buffer.get() + offset, &elem_header, elem_header_size);
+    offset += elem_header_size;
+
+    // Starting offset of child ElementHeader struct
+    size_t child_offset = offset + ((elem.children.size() + 1) * sizeof(uint32_t));
+    for (const RenderScriptRuntime::Element& child : elem.children)
+    {
+        // Recursively populate the buffer with the element header structs of children.
+        // Then save the offsets where they were set after the parent element header.
+        memcpy(header_buffer.get() + offset, &child_offset, sizeof(uint32_t));
+        offset += sizeof(uint32_t);
+
+        child_offset = PopulateElementHeaders(header_buffer, child_offset, child);
+    }
+
+    // Zero indicates no more children
+    memset(header_buffer.get() + offset, 0, sizeof(uint32_t));
+
+    return child_offset;
+}
+
+// Given an Element object this function returns the total size needed in the file header to store the element's details.
+// Taking into account the size of the element header struct, plus the offsets to all the element's children.
+// Function is recursive so that the size of all ancestors is taken into account.
+size_t
+RenderScriptRuntime::CalculateElementHeaderSize(const Element& elem)
+{
+    size_t size = (elem.children.size() + 1) * sizeof(uint32_t); // Offsets to children plus zero terminator
+    size += sizeof(AllocationDetails::ElementHeader); // Size of header struct with type details
+
+    // Calculate recursively for all descendants
+    for (const Element& child : elem.children)
+        size += CalculateElementHeaderSize(child);
+
+    return size;
+}
+
 // Function copies allocation contents into a binary file.
 // This file can then be loaded later into a different allocation.
 // There is a header, FileHeader, before the allocation data containing meta-data.
@@ -2209,17 +2304,44 @@ RenderScriptRuntime::SaveAllocation(Stream &strm, const uint32_t alloc_id, const
     // Create the file header
     AllocationDetails::FileHeader head;
     head.ident[0] = 'R'; head.ident[1] = 'S'; head.ident[2] = 'A'; head.ident[3] = 'D';
-    head.hdr_size = static_cast<uint16_t>(sizeof(AllocationDetails::FileHeader));
-    head.type = static_cast<uint16_t>(*alloc->element.type.get());
-    head.kind = static_cast<uint32_t>(*alloc->element.type_kind.get());
     head.dims[0] = static_cast<uint32_t>(alloc->dimension.get()->dim_1);
     head.dims[1] = static_cast<uint32_t>(alloc->dimension.get()->dim_2);
     head.dims[2] = static_cast<uint32_t>(alloc->dimension.get()->dim_3);
-    head.element_size = static_cast<uint32_t>(*alloc->element.datum_size.get());
+
+    const size_t element_header_size = CalculateElementHeaderSize(alloc->element);
+    assert((sizeof(AllocationDetails::FileHeader) + element_header_size) < UINT16_MAX && "Element header too large");
+    head.hdr_size = static_cast<uint16_t>(sizeof(AllocationDetails::FileHeader) + element_header_size);
 
     // Write the file header
     size_t num_bytes = sizeof(AllocationDetails::FileHeader);
-    Error err = file.Write(static_cast<const void*>(&head), num_bytes);
+    if (log)
+        log->Printf("RenderScriptRuntime::SaveAllocation - Writing File Header, 0x%zX bytes", num_bytes);
+
+    Error err = file.Write(&head, num_bytes);
+    if (!err.Success())
+    {
+        strm.Printf("Error: '%s' when writing to file '%s'", err.AsCString(), filename);
+        strm.EOL();
+        return false;
+    }
+
+    // Create the headers describing the element type of the allocation.
+    std::shared_ptr<uint8_t> element_header_buffer(new uint8_t[element_header_size]);
+    if (element_header_buffer == nullptr)
+    {
+        strm.Printf("Internal Error: Couldn't allocate %zu bytes on the heap", element_header_size);
+        strm.EOL();
+        return false;
+    }
+
+    PopulateElementHeaders(element_header_buffer, 0, alloc->element);
+
+    // Write headers for allocation element type to file
+    num_bytes = element_header_size;
+    if (log)
+        log->Printf("RenderScriptRuntime::SaveAllocation - Writing Element Headers, 0x%zX bytes", num_bytes);
+
+    err = file.Write(element_header_buffer.get(), num_bytes);
     if (!err.Success())
     {
         strm.Printf("Error: '%s' when writing to file '%s'", err.AsCString(), filename);
@@ -2230,7 +2352,7 @@ RenderScriptRuntime::SaveAllocation(Stream &strm, const uint32_t alloc_id, const
     // Write allocation data to file
     num_bytes = static_cast<size_t>(*alloc->size.get());
     if (log)
-        log->Printf("RenderScriptRuntime::SaveAllocation - Writing 0x%" PRIx64 " bytes from %p", (uint64_t) num_bytes, (void*) buffer.get());
+        log->Printf("RenderScriptRuntime::SaveAllocation - Writing 0x%zX bytes", num_bytes);
 
     err = file.Write(buffer.get(), num_bytes);
     if (!err.Success())