Fix subobject raw bytes storage; don't point into RDAT for reflection (#1710)

* Fix subobject raw bytes storage; don't point into RDAT for reflection

DxilSubobject:
- combine string and raw bytes storage in m_BytesStorage
- use StringRef as key, since memcmp is used with explicit length

DxilRuntimeReflection:
- add m_BytesMap for local unique copy of raw bytes in RDAT
  to prevent pointing into RDAT blob for root sig
- clean up some insertion patterns

* Rename Get[SubobjectString/RawBytes] to Intern[String/RawBytes]

include/dxc/DXIL/DxilSubobject.h

@@ -108,8 +108,7 @@ private:
 class DxilSubobjects {
 public:
   typedef std::pair<std::unique_ptr<char[]>, size_t> StoredBytes;
-  typedef llvm::MapVector< llvm::StringRef, StoredBytes > StringStorage;
-  typedef llvm::MapVector< const void*, StoredBytes > RawBytesStorage;
+  typedef llvm::MapVector< llvm::StringRef, StoredBytes > BytesStorage;
   typedef llvm::MapVector< llvm::StringRef, std::unique_ptr<DxilSubobject> > SubobjectStorage;
   using Kind = DXIL::SubobjectKind;
 
@@ -121,9 +120,9 @@ public:
   DxilSubobjects &operator=(const DxilSubobjects &other) = delete;
 
   // Add/find string in owned subobject strings, returning canonical ptr
-  llvm::StringRef GetSubobjectString(llvm::StringRef value);
+  llvm::StringRef InternString(llvm::StringRef value);
   // Add/find raw bytes, returning canonical ptr
-  const void *GetRawBytes(const void *ptr, size_t size);
+  const void *InternRawBytes(const void *ptr, size_t size);
   DxilSubobject *FindSubobject(llvm::StringRef name);
   void RemoveSubobject(llvm::StringRef name);
   DxilSubobject &CloneSubobject(const DxilSubobject &Subobject, llvm::StringRef Name);
@@ -158,8 +157,7 @@ public:
 private:
   DxilSubobject &CreateSubobject(Kind kind, llvm::StringRef Name);
 
-  StringStorage m_StringStorage;
-  RawBytesStorage m_RawBytesStorage;
+  BytesStorage m_BytesStorage;
   SubobjectStorage m_Subobjects;
 };
 

include/dxc/DxilContainer/DxilRuntimeReflection.inl

@@ -263,6 +263,7 @@ namespace {
 class DxilRuntimeReflection_impl : public DxilRuntimeReflection {
 private:
   typedef std::unordered_map<const char *, std::unique_ptr<wchar_t[]>> StringMap;
+  typedef std::unordered_map<const void *, std::unique_ptr<char[]>> BytesMap;
   typedef std::vector<const wchar_t *> WStringList;
   typedef std::vector<DxilResourceDesc> ResourceList;
   typedef std::vector<DxilResourceDesc *> ResourceRefList;
@@ -271,17 +272,19 @@ private:
 
   DxilRuntimeData m_RuntimeData;
   StringMap m_StringMap;
+  BytesMap m_BytesMap;
   ResourceList m_Resources;
   FunctionList m_Functions;
   SubobjectList m_Subobjects;
   std::unordered_map<ResourceKey, DxilResourceDesc *> m_ResourceMap;
   std::unordered_map<DxilFunctionDesc *, ResourceRefList> m_FuncToResMap;
-  std::unordered_map<DxilFunctionDesc *, WStringList> m_FuncToStringMap;
-  std::unordered_map<DxilSubobjectDesc *, WStringList> m_SubobjectToStringMap;
+  std::unordered_map<DxilFunctionDesc *, WStringList> m_FuncToDependenciesMap;
+  std::unordered_map<DxilSubobjectDesc *, WStringList> m_SubobjectToExportsMap;
   bool m_initialized;
 
   const wchar_t *GetWideString(const char *ptr);
   void AddString(const char *ptr);
+  const void *GetBytes(const void *ptr, size_t size);
   void InitializeReflection();
   const DxilResourceDesc * const*GetResourcesForFunction(DxilFunctionDesc &function,
                              const FunctionReader &functionReader);
@@ -297,8 +300,8 @@ public:
   // TODO: Implement pipeline state validation with runtime data
   // TODO: Update BlobContainer.h to recognize 'RDAT' blob
   DxilRuntimeReflection_impl()
-      : m_RuntimeData(), m_StringMap(), m_Resources(), m_Functions(),
-        m_FuncToResMap(), m_FuncToStringMap(), m_SubobjectToStringMap(),
+      : m_RuntimeData(), m_StringMap(), m_BytesMap(), m_Resources(), m_Functions(),
+        m_FuncToResMap(), m_FuncToDependenciesMap(), m_SubobjectToExportsMap(),
         m_initialized(false) {}
   virtual ~DxilRuntimeReflection_impl() {}
   // This call will allocate memory for GetLibraryReflection call
@@ -326,7 +329,19 @@ const wchar_t *DxilRuntimeReflection_impl::GetWideString(const char *ptr) {
   return m_StringMap.at(ptr).get();
 }
 
+const void *DxilRuntimeReflection_impl::GetBytes(const void *ptr, size_t size) {
+  auto it = m_BytesMap.find(ptr);
+  if (it != m_BytesMap.end())
+    return it->second.get();
+
+  auto inserted = m_BytesMap.insert(std::make_pair(ptr, std::unique_ptr<char[]>(new char[size])));
+  void *newPtr = inserted.first->second.get();
+  memcpy(newPtr, ptr, size);
+  return newPtr;
+}
+
 bool DxilRuntimeReflection_impl::InitFromRDAT(const void *pRDAT, size_t size) {
+  assert(!m_initialized && "may only initialize once");
   m_initialized = m_RuntimeData.InitFromRDAT(pRDAT, size);
   if (m_initialized)
     InitializeReflection();
@@ -399,30 +414,28 @@ DxilRuntimeReflection_impl::AddResource(const ResourceReader &resourceReader) {
 
 const DxilResourceDesc * const*DxilRuntimeReflection_impl::GetResourcesForFunction(
     DxilFunctionDesc &function, const FunctionReader &functionReader) {
-  if (m_FuncToResMap.find(&function) == m_FuncToResMap.end())
-    m_FuncToResMap.insert(std::pair<DxilFunctionDesc *, ResourceRefList>(
-        &function, ResourceRefList()));
-  ResourceRefList &resourceList = m_FuncToResMap.at(&function);
-  if (resourceList.empty()) {
-    resourceList.reserve(functionReader.GetNumResources());
-    for (uint32_t i = 0; i < functionReader.GetNumResources(); ++i) {
-      const ResourceReader resourceReader = functionReader.GetResource(i);
-      ResourceKey key((uint32_t)resourceReader.GetResourceClass(),
-                      resourceReader.GetID());
-      auto it = m_ResourceMap.find(key);
-      assert(it != m_ResourceMap.end() && it->second && "Otherwise, resource was not in map, or was null");
-      resourceList.emplace_back(it->second);
-    }
+  if (!functionReader.GetNumResources())
+    return nullptr;
+  auto it = m_FuncToResMap.insert(std::make_pair(&function, ResourceRefList()));
+  assert(it.second && "otherwise, collision");
+  ResourceRefList &resourceList = it.first->second;
+  resourceList.reserve(functionReader.GetNumResources());
+  for (uint32_t i = 0; i < functionReader.GetNumResources(); ++i) {
+    const ResourceReader resourceReader = functionReader.GetResource(i);
+    ResourceKey key((uint32_t)resourceReader.GetResourceClass(),
+                    resourceReader.GetID());
+    auto it = m_ResourceMap.find(key);
+    assert(it != m_ResourceMap.end() && it->second && "Otherwise, resource was not in map, or was null");
+    resourceList.emplace_back(it->second);
   }
-  return resourceList.empty() ? nullptr : resourceList.data();
+  return resourceList.data();
 }
 
 const wchar_t **DxilRuntimeReflection_impl::GetDependenciesForFunction(
     DxilFunctionDesc &function, const FunctionReader &functionReader) {
-  if (m_FuncToStringMap.find(&function) == m_FuncToStringMap.end())
-    m_FuncToStringMap.insert(
-        std::pair<DxilFunctionDesc *, WStringList>(&function, WStringList()));
-  WStringList &wStringList = m_FuncToStringMap.at(&function);
+  auto it = m_FuncToDependenciesMap.insert(std::make_pair(&function, WStringList()));
+  assert(it.second && "otherwise, collision");
+  WStringList &wStringList = it.first->second;
   for (uint32_t i = 0; i < functionReader.GetNumDependencies(); ++i) {
     wStringList.emplace_back(GetWideString(functionReader.GetDependency(i)));
   }
@@ -455,10 +468,9 @@ DxilRuntimeReflection_impl::AddFunction(const FunctionReader &functionReader) {
 
 const wchar_t **DxilRuntimeReflection_impl::GetExportsForAssociation(
   DxilSubobjectDesc &subobject, const SubobjectReader &subobjectReader) {
-  if (m_SubobjectToStringMap.find(&subobject) == m_SubobjectToStringMap.end())
-    m_SubobjectToStringMap.insert(
-      std::pair<DxilSubobjectDesc *, WStringList>(&subobject, WStringList()));
-  WStringList &wStringList = m_SubobjectToStringMap.at(&subobject);
+  auto it = m_SubobjectToExportsMap.insert(std::make_pair(&subobject, WStringList()));
+  assert(it.second && "otherwise, collision");
+  WStringList &wStringList = it.first->second;
   for (uint32_t i = 0; i < subobjectReader.GetSubobjectToExportsAssociation_NumExports(); ++i) {
     wStringList.emplace_back(GetWideString(subobjectReader.GetSubobjectToExportsAssociation_Export(i)));
   }
@@ -481,6 +493,7 @@ DxilSubobjectDesc *DxilRuntimeReflection_impl::AddSubobject(const SubobjectReade
   case DXIL::SubobjectKind::LocalRootSignature:
     if (!subobjectReader.GetRootSignature(&subobject.RootSignature.pSerializedSignature, &subobject.RootSignature.SizeInBytes))
       return nullptr;
+    subobject.RootSignature.pSerializedSignature = GetBytes(subobject.RootSignature.pSerializedSignature, subobject.RootSignature.SizeInBytes);
     break;
   case DXIL::SubobjectKind::SubobjectToExportsAssociation:
     subobject.SubobjectToExportsAssociation.Subobject =

lib/DXIL/DxilSubobject.cpp

@@ -22,7 +22,7 @@ namespace hlsl {
 DxilSubobject::DxilSubobject(DxilSubobject &&other)
   : m_Owner(other.m_Owner),
     m_Kind(other.m_Kind),
-    m_Name(m_Owner.GetSubobjectString(other.m_Name)),
+    m_Name(m_Owner.InternString(other.m_Name)),
     m_Exports(std::move(other.m_Exports))
 {
   DXASSERT_NOMSG(DXIL::IsValidSubobjectKind(m_Kind));
@@ -32,7 +32,7 @@ DxilSubobject::DxilSubobject(DxilSubobject &&other)
 DxilSubobject::DxilSubobject(DxilSubobjects &owner, Kind kind, llvm::StringRef name)
   : m_Owner(owner),
     m_Kind(kind),
-    m_Name(m_Owner.GetSubobjectString(name)),
+    m_Name(m_Owner.InternString(name)),
     m_Exports()
 {
   DXASSERT_NOMSG(DXIL::IsValidSubobjectKind(m_Kind));
@@ -81,17 +81,17 @@ void DxilSubobject::CopyUnionedContents(const DxilSubobject &other) {
 
 void DxilSubobject::InternStrings() {
   // Transfer strings if necessary
-  m_Name = m_Owner.GetSubobjectString(m_Name).data();
+  m_Name = m_Owner.InternString(m_Name).data();
   switch (m_Kind) {
   case Kind::SubobjectToExportsAssociation:
-    SubobjectToExportsAssociation.Subobject = m_Owner.GetSubobjectString(SubobjectToExportsAssociation.Subobject).data();
+    SubobjectToExportsAssociation.Subobject = m_Owner.InternString(SubobjectToExportsAssociation.Subobject).data();
     for (auto &ptr : m_Exports)
-      ptr = m_Owner.GetSubobjectString(ptr).data();
+      ptr = m_Owner.InternString(ptr).data();
     break;
   case Kind::HitGroup:
-    HitGroup.AnyHit = m_Owner.GetSubobjectString(HitGroup.AnyHit).data();
-    HitGroup.ClosestHit = m_Owner.GetSubobjectString(HitGroup.ClosestHit).data();
-    HitGroup.Intersection = m_Owner.GetSubobjectString(HitGroup.Intersection).data();
+    HitGroup.AnyHit = m_Owner.InternString(HitGroup.AnyHit).data();
+    HitGroup.ClosestHit = m_Owner.InternString(HitGroup.ClosestHit).data();
+    HitGroup.Intersection = m_Owner.InternString(HitGroup.Intersection).data();
     break;
   default:
     break;
@@ -177,45 +177,40 @@ bool DxilSubobject::GetHitGroup(DXIL::HitGroupType &hitGroupType,
 
 
 DxilSubobjects::DxilSubobjects()
-  : m_StringStorage()
-  , m_RawBytesStorage()
+  : m_BytesStorage()
   , m_Subobjects()
 {}
 DxilSubobjects::DxilSubobjects(DxilSubobjects &&other)
-  : m_StringStorage(std::move(other.m_StringStorage))
-  , m_RawBytesStorage(std::move(other.m_RawBytesStorage))
+  : m_BytesStorage(std::move(other.m_BytesStorage))
   , m_Subobjects(std::move(other.m_Subobjects))
 {}
 DxilSubobjects::~DxilSubobjects() {}
 
 
-llvm::StringRef DxilSubobjects::GetSubobjectString(llvm::StringRef value) {
-  auto it = m_StringStorage.find(value);
-  if (it != m_StringStorage.end())
+llvm::StringRef DxilSubobjects::InternString(llvm::StringRef value) {
+  auto it = m_BytesStorage.find(value);
+  if (it != m_BytesStorage.end())
     return it->first;
 
   size_t size = value.size();
-  StoredBytes stored(std::unique_ptr<char[]>(new char[size + 1]), size + 1);
+  StoredBytes stored(std::make_pair(std::unique_ptr<char[]>(new char[size + 1]), size + 1));
   memcpy(stored.first.get(), value.data(), size);
   stored.first[size] = 0;
   llvm::StringRef key(stored.first.get(), size);
-  m_StringStorage[key] = std::move(stored);
+  m_BytesStorage[key] = std::move(stored);
   return key;
 }
 
-const void *DxilSubobjects::GetRawBytes(const void *ptr, size_t size) {
-  auto it = m_RawBytesStorage.find(ptr);
-  if (it != m_RawBytesStorage.end())
-    return ptr;
+const void *DxilSubobjects::InternRawBytes(const void *ptr, size_t size) {
+  auto it = m_BytesStorage.find(llvm::StringRef((const char *)ptr, size));
+  if (it != m_BytesStorage.end())
+    return it->first.data();
 
-  DXASSERT_NOMSG(size < UINT_MAX);
-  if (size >= UINT_MAX)
-    return nullptr;
-  StoredBytes stored(std::unique_ptr<char[]>(new char[size]), size);
+  StoredBytes stored(std::make_pair(std::unique_ptr<char[]>(new char[size]), size));
   memcpy(stored.first.get(), ptr, size);
-  ptr = stored.first.get();
-  m_RawBytesStorage[ptr] = std::move(stored);
-  return ptr;
+  llvm::StringRef key(stored.first.get(), size);
+  m_BytesStorage[key] = std::move(stored);
+  return key.data();
 }
 
 DxilSubobject *DxilSubobjects::FindSubobject(llvm::StringRef name) {
@@ -233,7 +228,7 @@ void DxilSubobjects::RemoveSubobject(llvm::StringRef name) {
 
 DxilSubobject &DxilSubobjects::CloneSubobject(
     const DxilSubobject &Subobject, llvm::StringRef Name) {
-  Name = GetSubobjectString(Name);
+  Name = InternString(Name);
   DXASSERT(FindSubobject(Name) == nullptr,
     "otherwise, name collision between subobjects");
   std::unique_ptr<DxilSubobject> ptr(new DxilSubobject(*this, Subobject, Name));
@@ -255,9 +250,9 @@ DxilSubobject &DxilSubobjects::CreateStateObjectConfig(
 DxilSubobject &DxilSubobjects::CreateRootSignature(
     llvm::StringRef Name, bool local, const void *Data, uint32_t Size, llvm::StringRef *pText /*= nullptr*/) {
   auto &obj = CreateSubobject(local ? Kind::LocalRootSignature : Kind::GlobalRootSignature, Name);
-  obj.RootSignature.Data = Data;
+  obj.RootSignature.Data = InternRawBytes(Data, Size);
   obj.RootSignature.Size = Size;
-  obj.RootSignature.Text = (pText ? GetSubobjectString(*pText).data() : nullptr);
+  obj.RootSignature.Text = (pText ? InternString(*pText).data() : nullptr);
   return obj;
 }
 
@@ -267,10 +262,10 @@ DxilSubobject &DxilSubobjects::CreateSubobjectToExportsAssociation(
     llvm::StringRef *Exports,
     uint32_t NumExports) {
   auto &obj = CreateSubobject(Kind::SubobjectToExportsAssociation, Name);
-  Subobject = GetSubobjectString(Subobject);
+  Subobject = InternString(Subobject);
   obj.SubobjectToExportsAssociation.Subobject = Subobject.data();
   for (unsigned i = 0; i < NumExports; i++) {
-    obj.m_Exports.emplace_back(GetSubobjectString(Exports[i]).data());
+    obj.m_Exports.emplace_back(InternString(Exports[i]).data());
   }
   return obj;
 }
@@ -299,9 +294,9 @@ DxilSubobject &DxilSubobjects::CreateHitGroup(llvm::StringRef Name,
                                               llvm::StringRef ClosestHit,
                                               llvm::StringRef Intersection) {
   auto &obj = CreateSubobject(Kind::HitGroup, Name);
-  AnyHit = GetSubobjectString(AnyHit);
-  ClosestHit = GetSubobjectString(ClosestHit);
-  Intersection = GetSubobjectString(Intersection);
+  AnyHit = InternString(AnyHit);
+  ClosestHit = InternString(ClosestHit);
+  Intersection = InternString(Intersection);
   obj.HitGroup.Type = hitGroupType;
   obj.HitGroup.AnyHit = AnyHit.data();
   obj.HitGroup.ClosestHit = ClosestHit.data();
@@ -310,7 +305,7 @@ DxilSubobject &DxilSubobjects::CreateHitGroup(llvm::StringRef Name,
 }
 
 DxilSubobject &DxilSubobjects::CreateSubobject(Kind kind, llvm::StringRef Name) {
-  Name = GetSubobjectString(Name);
+  Name = InternString(Name);
   IFTBOOLMSG(FindSubobject(Name) == nullptr, DXC_E_GENERAL_INTERNAL_ERROR, "Subobject name collision");
   IFTBOOLMSG(!Name.empty(), DXC_E_GENERAL_INTERNAL_ERROR, "Empty Subobject name");
   std::unique_ptr<DxilSubobject> ptr(new DxilSubobject(*this, kind, Name));