Add ability to replace sections in RDCFile

qrenderdoc/Code/CaptureContext.cpp

@@ -538,12 +538,29 @@ bool CaptureContext::SaveCaptureTo(const QString &captureFile)
   {
     if(QFileInfo(GetCaptureFilename()).exists())
     {
-      // QFile::copy won't overwrite, so remove the destination first (the save dialog already
-      // prompted for overwrite)
-      QFile::remove(captureFile);
-      success = QFile::copy(GetCaptureFilename(), captureFile);
+      if(GetCaptureFilename() == captureFile)
+      {
+        success = true;
+      }
+      else
+      {
+        ICaptureFile *capFile = Replay().GetCaptureFile();
 
-      error = tr("Couldn't save to %1").arg(captureFile);
+        if(capFile)
+        {
+          // this will overwrite
+          success = capFile->CopyFileTo(captureFile.toUtf8().data());
+        }
+        else
+        {
+          // QFile::copy won't overwrite, so remove the destination first (the save dialog already
+          // prompted for overwrite)
+          QFile::remove(captureFile);
+          success = QFile::copy(GetCaptureFilename(), captureFile);
+        }
+
+        error = tr("Couldn't save to %1").arg(captureFile);
+      }
     }
     else
     {
@@ -568,9 +585,17 @@ bool CaptureContext::SaveCaptureTo(const QString &captureFile)
     return false;
   }
 
+  // if it was a temporary capture, remove the old instnace
+  if(m_CaptureTemporary)
+    QFile::remove(m_CaptureFile);
+
+  // Update the filename, and mark that it's local and not temporary now.
   m_CaptureFile = captureFile;
+  m_CaptureLocal = true;
   m_CaptureTemporary = false;
 
+  Replay().ReopenCaptureFile(captureFile);
+
   return true;
 }
 

qrenderdoc/Code/ReplayManager.cpp

@@ -355,6 +355,13 @@ void ReplayManager::PingRemote()
   }
 }
 
+void ReplayManager::ReopenCaptureFile(const QString &path)
+{
+  if(!m_CaptureFile)
+    m_CaptureFile = RENDERDOC_OpenCaptureFile();
+  m_CaptureFile->OpenFile(path.toUtf8().data(), "rdc");
+}
+
 uint32_t ReplayManager::ExecuteAndInject(const QString &exe, const QString &workingDir,
                                          const QString &cmdLine,
                                          const QList<EnvironmentModification> &env,

qrenderdoc/Code/ReplayManager.h

@@ -83,6 +83,11 @@ public:
       return m_Remote;
     return NULL;
   }
+
+  // may return NULL if the capture file is not open locally. Consider using ICaptureAccess above to
+  // work whether local or remote.
+  ICaptureFile *GetCaptureFile() { return m_CaptureFile; }
+  void ReopenCaptureFile(const QString &path);
   const RemoteHost *CurrentRemote() { return m_RemoteHost; }
   uint32_t ExecuteAndInject(const QString &exe, const QString &workingDir, const QString &cmdLine,
                             const QList<EnvironmentModification> &env, const QString &capturefile,

renderdoc/api/replay/renderdoc_replay.h

@@ -1310,6 +1310,20 @@ For the :param:`filetype` parameter, see :meth:`OpenFile`.
 )");
   virtual ReplayStatus OpenBuffer(const bytebuf &buffer, const char *filetype) = 0;
 
+  DOCUMENT(R"(When a capture file is opened, an exclusive lock is held on the file on disk. This
+makes it impossible to copy the file to another location at the user's request. Calling this
+function will copy the file on disk to a new location but otherwise won't affect the capture handle.
+The new file will be locked, the old file will be unlocked - to allow deleting if necessary.
+
+It is invalid to call this function if :meth:`OpenFile` has not previously been called to open the
+file.
+
+:param str filename: The filename to copy to.
+:return: ``True`` if the operation succeeded.
+:rtype: bool
+)");
+  virtual bool CopyFileTo(const char *filename) = 0;
+
   DOCUMENT(R"(Converts the currently loaded file to a given format and saves it to disk.
 
 This allows converting a native RDC to another representation, or vice-versa converting another

renderdoc/replay/capture_file.cpp

@@ -112,6 +112,7 @@ public:
 
   ReplayStatus OpenFile(const char *filename, const char *filetype);
   ReplayStatus OpenBuffer(const bytebuf &buffer, const char *filetype);
+  bool CopyFileTo(const char *filename);
 
   void Shutdown() { delete this; }
   ReplaySupport LocalReplaySupport() { return m_Support; }
@@ -253,6 +254,14 @@ ReplayStatus CaptureFile::OpenBuffer(const bytebuf &buffer, const char *filetype
   return Init();
 }
 
+bool CaptureFile::CopyFileTo(const char *filename)
+{
+  if(m_RDC)
+    return m_RDC->CopyFileTo(filename);
+
+  return false;
+}
+
 ReplayStatus CaptureFile::Init()
 {
   if(!m_RDC)

renderdoc/serialise/rdcfile.cpp

@@ -245,7 +245,7 @@ void RDCFile::Open(const char *path)
   RDCCOMPILE_ASSERT(offsetof(BinarySectionHeader, name) == sizeof(uint32_t) * 10,
                     "BinarySectionHeader size has changed or contains padding");
 
-  m_File = FileIO::fopen(path, "rb");
+  m_File = FileIO::fopen(path, "r+b");
   m_Filename = path;
 
   if(!m_File)
@@ -423,6 +423,8 @@ void RDCFile::Init(StreamReader &reader)
     BinarySectionHeader sectionHeader = {0};
     byte *reading = (byte *)&sectionHeader;
 
+    uint64_t headerOffset = reader.GetOffset();
+
     reader.Read(*reading);
     reading++;
 
@@ -520,7 +522,8 @@ void RDCFile::Init(StreamReader &reader)
       props.uncompressedSize = length;
 
       SectionLocation loc;
-      loc.offs = reader.GetOffset();
+      loc.headerOffset = headerOffset;
+      loc.dataOffset = reader.GetOffset();
       loc.diskLength = length;
 
       reader.SkipBytes(loc.diskLength);
@@ -564,7 +567,8 @@ void RDCFile::Init(StreamReader &reader)
         RETURNCORRUPT("Error reading binary section header");
 
       SectionLocation loc;
-      loc.offs = reader.GetOffset();
+      loc.headerOffset = headerOffset;
+      loc.dataOffset = reader.GetOffset();
       loc.diskLength = sectionHeader.sectionCompressedLength;
 
       m_Sections.push_back(props);
@@ -587,6 +591,29 @@ void RDCFile::Init(StreamReader &reader)
   }
 }
 
+bool RDCFile::CopyFileTo(const char *filename)
+{
+  if(!m_File)
+    return false;
+
+  // remember our position and close the file
+  uint64_t prevPos = FileIO::ftell64(m_File);
+  FileIO::fclose(m_File);
+
+  // try to move to the new location
+  bool success = FileIO::Copy(m_Filename.c_str(), filename, true);
+
+  // if it succeeded, update our filename
+  if(success)
+    m_Filename = filename;
+
+  // re-open the file (either the new one, or the old one if it failed) and re-seek
+  m_File = FileIO::fopen(m_Filename.c_str(), "r+b");
+  FileIO::fseek64(m_File, prevPos, SEEK_SET);
+
+  return success;
+}
+
 void RDCFile::SetData(RDCDriver driver, const char *driverName, uint64_t machineIdent,
                       const RDCThumb *thumb)
 {
@@ -657,6 +684,11 @@ void RDCFile::Create(const char *filename)
 
 int RDCFile::SectionIndex(SectionType type) const
 {
+  // Unknown is not a real type, any arbitrary sections with names will be listed as unknown, so
+  // don't return a false-positive index. This allows us to skip some special cases outside
+  if(type == SectionType::Unknown)
+    return -1;
+
   for(size_t i = 0; i < m_Sections.size(); i++)
     if(m_Sections[i].type == type)
       return int(i);
@@ -689,7 +721,7 @@ StreamReader *RDCFile::ReadSection(int index) const
 
   const SectionProperties &props = m_Sections[index];
   SectionLocation offsetSize = m_SectionLocations[index];
-  FileIO::fseek64(m_File, offsetSize.offs, SEEK_SET);
+  FileIO::fseek64(m_File, offsetSize.dataOffset, SEEK_SET);
 
   StreamReader *fileReader = new StreamReader(m_File, offsetSize.diskLength, Ownership::Nothing);
 
@@ -717,29 +749,7 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
   if(m_Error != ContainerError::NoError)
     return new StreamWriter(StreamWriter::InvalidStream);
 
-  // only handle the case of writing a section that doesn't exist yet.
-  // For handling a section that does exist, it depends on the section type:
-  // - For frame capture, then we just write to a new file since we want it
-  //   to be first. Once the writing is done, copy across any other sections
-  //   after it.
-  // - For non-frame capture, we remove the existing section and move up any
-  //   sections that were after it. Then just return a new writer that appends
-
-  if(props.type != SectionType::Unknown)
-  {
-    if(SectionIndex(props.type) >= 0)
-    {
-      RDCERR("Replacing sections is currently not supported.");
-      return new StreamWriter(StreamWriter::InvalidStream);
-    }
-    RDCASSERT((size_t)props.type < (size_t)SectionType::Count);
-  }
-
-  if(SectionIndex(props.name.c_str()) >= 0)
-  {
-    RDCERR("Replacing sections is currently not supported.");
-    return new StreamWriter(StreamWriter::InvalidStream);
-  }
+  RDCASSERT((size_t)props.type < (size_t)SectionType::Count);
 
   if(m_File == NULL)
   {
@@ -764,7 +774,7 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
     return new StreamWriter(StreamWriter::InvalidStream);
   }
 
-  if(m_CurrentWritingProps.type != SectionType::Count)
+  if(!m_CurrentWritingProps.name.empty())
   {
     RDCERR("Only one section can be written at once.");
     return new StreamWriter(StreamWriter::InvalidStream);
@@ -774,10 +784,216 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
   SectionType type = props.type;
 
   // normalise names for known sections
-  if(props.type != SectionType::Unknown)
+  if(type != SectionType::Unknown && type < SectionType::Count)
     name = SectionTypeNames[(size_t)type];
 
-  FileIO::fseek64(m_File, 0, SEEK_END);
+  if(name.empty())
+  {
+    RDCERR("Sections must have a name, either auto-populated from a known type or specified.");
+    return new StreamWriter(StreamWriter::InvalidStream);
+  }
+
+  // For handling a section that does exist, it depends on the section type:
+  // - For frame capture, then we just write to a new file since we want it
+  //   to be first. Once the writing is done, copy across any other sections
+  //   after it.
+  // - For non-frame capture, we remove the existing section and move up any
+  //   sections that were after it. Then just return a new writer that appends
+
+  // we store this callback here so that we can execute it after any post-section-writing header
+  // fixups. We need to be able to fixup any pre-existing sections that got shifted around.
+  StreamCloseCallback modifySectionCallback;
+
+  if(SectionIndex(type) >= 0 || SectionIndex(name.c_str()) >= 0)
+  {
+    if(type == SectionType::FrameCapture || name == SectionTypeNames[(int)SectionType::FrameCapture])
+    {
+      // simple case - if there are no other sections then we can just overwrite the existing frame
+      // capture.
+      if(NumSections() == 1)
+      {
+        // seek to the start of where the section is.
+        FileIO::fseek64(m_File, m_SectionLocations[0].headerOffset, SEEK_SET);
+
+        uint64_t oldLength = m_SectionLocations[0].diskLength;
+
+        // after writing, we need to be sure to fixup the size (in case we wrote less data).
+        modifySectionCallback = [this, oldLength]() {
+          if(oldLength > m_SectionLocations[0].diskLength)
+          {
+            FileIO::ftruncateat(
+                m_File, m_SectionLocations[0].dataOffset + m_SectionLocations[0].diskLength);
+          }
+        };
+      }
+      else
+      {
+        FILE *origFile = m_File;
+
+        // save the sections
+        std::vector<SectionProperties> origSections = m_Sections;
+        std::vector<SectionLocation> origSectionLocations = m_SectionLocations;
+
+        SectionLocation oldCaptureLocation = m_SectionLocations[0];
+
+        // remove section 0, the frame capture, since it will be fixed up separately
+        origSections.erase(origSections.begin());
+        origSectionLocations.erase(origSectionLocations.begin());
+
+        std::string tempFilename = FileIO::GetTempFolderFilename() + "capture_rewrite.rdc";
+
+        // create the file, this will overwrite m_File with the new file and file header using the
+        // existing loaded metadata
+        Create(tempFilename.c_str());
+
+        // after we've written the frame capture, we need to copy over the other sections into the
+        // temporary file and finally move the temporary file over the top of the existing file.
+        modifySectionCallback = [this, origFile, origSections, origSectionLocations, tempFilename]() {
+          // seek to write after the frame capture
+          FileIO::fseek64(
+              m_File, m_SectionLocations[0].dataOffset + m_SectionLocations[0].diskLength, SEEK_SET);
+
+          // write the old sections
+          for(size_t i = 0; i < origSections.size(); i++)
+          {
+            SectionLocation loc = origSectionLocations[i];
+
+            FileIO::fseek64(origFile, loc.headerOffset, SEEK_SET);
+
+            uint64_t newHeaderOffset = FileIO::ftell64(m_File);
+
+            // update the offsets to where they are in the new file
+            if(newHeaderOffset > loc.headerOffset)
+            {
+              uint64_t delta = newHeaderOffset - loc.headerOffset;
+
+              loc.headerOffset += delta;
+              loc.dataOffset += delta;
+            }
+            else if(newHeaderOffset < loc.headerOffset)
+            {
+              uint64_t delta = loc.headerOffset - newHeaderOffset;
+
+              loc.headerOffset -= delta;
+              loc.dataOffset -= delta;
+            }
+
+            uint64_t headerLen = loc.dataOffset - loc.headerOffset;
+
+            // copy header and data together
+            StreamWriter writer(m_File, Ownership::Nothing);
+            StreamReader reader(origFile, headerLen + loc.diskLength, Ownership::Nothing);
+
+            m_Sections.push_back(origSections[i]);
+            m_SectionLocations.push_back(loc);
+
+            StreamTransfer(&writer, &reader, NULL);
+          }
+
+          // close the file writing to the temp location
+          FileIO::fclose(m_File);
+
+          // move the temp file over the original
+          FileIO::Move(tempFilename.c_str(), m_Filename.c_str(), true);
+
+          // re-open the file after it's been overwritten.
+          m_File = FileIO::fopen(m_Filename.c_str(), "r+b");
+        };
+
+        // fall through - we'll write to m_File immediately after the file header
+      }
+
+      // the new section data for the framecapture will be pushed on after writing. Any others will
+      // be re-added in the fixup step above
+      m_Sections.clear();
+      m_SectionLocations.clear();
+    }
+    else
+    {
+      // we're writing some section after the frame capture. We'll do this in-place by reading the
+      // other sections out to memory (assuming that they are mostly small, and even if they are
+      // somewhat large, it's still much better to leave the frame capture (which should dominate
+      // file size) on disk where it is.
+      int index = SectionIndex(type);
+
+      if(index < 0)
+        index = SectionIndex(name.c_str());
+
+      RDCASSERT(index >= 0);
+
+      std::vector<bytebuf> origSectionData;
+      std::vector<uint64_t> origHeaderSizes;
+
+      uint64_t overwriteLocation = m_SectionLocations[index].headerOffset;
+      uint64_t oldLength = m_SectionLocations[index].diskLength;
+
+      // erase the target section. The others will be moved up to match
+      m_Sections.erase(m_Sections.begin() + index);
+      m_SectionLocations.erase(m_SectionLocations.begin() + index);
+
+      origSectionData.reserve(NumSections() - index);
+      origHeaderSizes.reserve(NumSections() - index);
+
+      // go through all subsequent sections after this one in the file, read them into memory.
+      // this could be optimised since we're going to write them back out below, we could do this
+      // just with an in-memory window large enough.
+      for(int i = index; i < NumSections(); i++)
+      {
+        const SectionLocation &loc = m_SectionLocations[i];
+
+        FileIO::fseek64(m_File, loc.headerOffset, SEEK_SET);
+
+        uint64_t headerLen = loc.dataOffset - loc.headerOffset;
+
+        // read header and data together
+        StreamReader reader(m_File, headerLen + loc.diskLength, Ownership::Nothing);
+
+        origHeaderSizes.push_back(headerLen);
+        origSectionData.push_back(bytebuf());
+
+        bytebuf &data = origSectionData.back();
+        data.resize((size_t)reader.GetSize());
+        reader.Read(data.data(), data.size());
+      }
+
+      // we write the sections now over where the old section used to be, so the newly written
+      // section is last in the file. This means if the same section is updated over and over, it
+      // doesn't require moving any sections once it's already at the end.
+
+      // seek to write to where the removed section started
+      FileIO::fseek64(m_File, overwriteLocation, SEEK_SET);
+
+      // write the old sections
+      for(size_t i = 0; i < origSectionData.size(); i++)
+      {
+        // update the offsets to where they are in the new file
+        m_SectionLocations[index + i].headerOffset = FileIO::ftell64(m_File);
+        m_SectionLocations[index + i].dataOffset =
+            m_SectionLocations[index + i].headerOffset + origHeaderSizes[i];
+
+        // write the data
+        StreamWriter writer(m_File, Ownership::Nothing);
+        writer.Write(origSectionData[i].data(), origSectionData[i].size());
+      }
+
+      // after writing, we need to be sure to fixup the size (in case we wrote less data).
+      modifySectionCallback = [this, oldLength]() {
+        if(oldLength > m_SectionLocations.back().diskLength)
+        {
+          FileIO::ftruncateat(
+              m_File, m_SectionLocations.back().dataOffset + m_SectionLocations.back().diskLength);
+        }
+      };
+
+      // fall through - we now write to m_File with the new section wherever we left off after the
+      // moved sections.
+    }
+  }
+  else
+  {
+    // we're adding a new section - seek to the end of the file to append it
+    FileIO::fseek64(m_File, 0, SEEK_END);
+  }
 
   uint64_t headerOffset = FileIO::ftell64(m_File);
 
@@ -830,11 +1046,14 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
         new StreamWriter(new ZSTDCompressor(fileWriter, Ownership::Stream), Ownership::Stream);
   }
 
+  uint64_t dataOffset = FileIO::ftell64(m_File);
+
   m_CurrentWritingProps = props;
   m_CurrentWritingProps.name = name;
 
   // register a destroy callback to tidy up the section at the end
-  fileWriter->AddCloseCallback([this, type, name, headerOffset, fileWriter, compWriter]() {
+  fileWriter->AddCloseCallback([this, type, name, headerOffset, dataOffset, fileWriter, compWriter]() {
+    FileIO::fflush(m_File);
 
     // the offset of the file writer is how many bytes were written to disk - the compressed length.
     uint64_t compressedLength = fileWriter->GetOffset();
@@ -852,6 +1071,11 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
     m_CurrentWritingProps.uncompressedSize = uncompressedLength;
 
     m_Sections.push_back(m_CurrentWritingProps);
+    SectionLocation loc;
+    loc.headerOffset = headerOffset;
+    loc.dataOffset = dataOffset;
+    loc.diskLength = compressedLength;
+    m_SectionLocations.push_back(loc);
 
     m_CurrentWritingProps = SectionProperties();
 
@@ -869,6 +1093,9 @@ StreamWriter *RDCFile::WriteSection(const SectionProperties &props)
     }
   });
 
+  if(modifySectionCallback)
+    fileWriter->AddCloseCallback(modifySectionCallback);
+
   // if we're compressing return that writer, otherwise return the file writer directly
   return compWriter ? compWriter : fileWriter;
 }

renderdoc/serialise/rdcfile.h

@@ -93,6 +93,8 @@ public:
   void Open(const char *filename);
   void Open(const std::vector<byte> &buffer);
 
+  bool CopyFileTo(const char *filename);
+
   // Sets the parameters of an RDCFile in memory.
   void SetData(RDCDriver driver, const char *driverName, uint64_t machineIdent,
                const RDCThumb *thumb);
@@ -136,7 +138,8 @@ private:
 
   struct SectionLocation
   {
-    uint64_t offs;
+    uint64_t headerOffset;
+    uint64_t dataOffset;
     uint64_t diskLength;
   };