Make subtitle references ref-counted to ensure the memory remains valid while its still used.

Fixes issue 299.

4 files changed, 72 insertions, 57 deletions

diff --git a/decoder/LAVVideo/subtitles/LAVSubtitleFrame.cpp b/decoder/LAVVideo/subtitles/LAVSubtitleFrame.cpp
index 6954cd23..630be4b8 100644
--- a/decoder/LAVVideo/subtitles/LAVSubtitleFrame.cpp
+++ b/decoder/LAVVideo/subtitles/LAVSubtitleFrame.cpp
@@ -34,8 +34,7 @@ CLAVSubtitleFrame::CLAVSubtitleFrame(void)
 CLAVSubtitleFrame::~CLAVSubtitleFrame(void)
 {
   for (int i = 0; i < m_NumBitmaps; i++) {
-    if (m_Bitmaps[i].freePixels)
-      SAFE_CO_FREE(m_Bitmaps[i].pixels);
+    m_Bitmaps[i]->Release();
   }
   SAFE_CO_FREE(m_Bitmaps);
 }
@@ -52,30 +51,22 @@ STDMETHODIMP CLAVSubtitleFrame::SetClipRect(RECT clipRect)
   return S_OK;
 }
 
-STDMETHODIMP CLAVSubtitleFrame::AddBitmap(const LAVSubRect &subRect)
+STDMETHODIMP CLAVSubtitleFrame::AddBitmap(CLAVSubRect *subRect)
 {
   // Allocate memory for the new block
   void *mem = CoTaskMemRealloc(m_Bitmaps, sizeof(*m_Bitmaps) * (m_NumBitmaps+1));
   if (!mem) {
     return E_OUTOFMEMORY;
   }
-  m_Bitmaps = (LAVSubRect *)mem;
+
+  m_Bitmaps = (CLAVSubRect **)mem;
   m_Bitmaps[m_NumBitmaps] = subRect;
   m_NumBitmaps++;
 
-  return S_OK;
-}
+  // Hold reference on the subtitle rect
+  subRect->AddRef();
 
-STDMETHODIMP CLAVSubtitleFrame::AddBitmap(ULONGLONG id, POINT position, SIZE size, LPVOID pixels, int pitch)
-{
-  LAVSubRect rect;
-  rect.id       = id;
-  rect.position = position;
-  rect.size     = size;
-  rect.pixels   = pixels;
-  rect.pitch    = pitch;
-
-  return AddBitmap(rect);
+  return S_OK;
 }
 
 STDMETHODIMP CLAVSubtitleFrame::GetOutputRect(RECT *outputRect)
@@ -109,11 +100,11 @@ STDMETHODIMP CLAVSubtitleFrame::GetBitmap(int index, ULONGLONG *id, POINT *posit
   CheckPointer(pixels,   E_POINTER);
   CheckPointer(pitch,    E_POINTER);
 
-  *id       = m_Bitmaps[index].id;
-  *position = m_Bitmaps[index].position;
-  *size     = m_Bitmaps[index].size;
-  *pixels   = m_Bitmaps[index].pixels;
-  *pitch    = m_Bitmaps[index].pitch;
+  *id       = m_Bitmaps[index]->id;
+  *position = m_Bitmaps[index]->position;
+  *size     = m_Bitmaps[index]->size;
+  *pixels   = m_Bitmaps[index]->pixels;
+  *pitch    = m_Bitmaps[index]->pitch;
 
   return S_OK;
 }
diff --git a/decoder/LAVVideo/subtitles/LAVSubtitleFrame.h b/decoder/LAVVideo/subtitles/LAVSubtitleFrame.h
index 5486aea5..8e12c6cc 100644
--- a/decoder/LAVVideo/subtitles/LAVSubtitleFrame.h
+++ b/decoder/LAVVideo/subtitles/LAVSubtitleFrame.h
@@ -19,7 +19,27 @@
 
 #pragma once
 
-typedef struct LAVSubRect {
+class CLAVSubRect : public IUnknown
+{
+public:
+  CLAVSubRect() : m_cRef(0), rtStart(AV_NOPTS_VALUE), rtStop(AV_NOPTS_VALUE), id(0), pixels(NULL), pixelsPal(NULL), pitch(0), forced(false), freePixels(false) { memset(&position, 0, sizeof(position)); memset(&size, 0, sizeof(size));}
+  ~CLAVSubRect() { SAFE_CO_FREE(pixels); SAFE_CO_FREE(pixelsPal); DbgLog((LOG_TRACE, 10, L"Rect at %I64d deleted", rtStart));}
+
+  // IUnknown
+  STDMETHODIMP QueryInterface(REFIID riid, void **ppvObject)
+  {
+    if (riid == IID_IUnknown) {
+      AddRef();
+      *ppvObject = (IUnknown *)this;
+    } else {
+      return E_NOINTERFACE;
+    }
+    return S_OK;
+  }
+  STDMETHODIMP_(ULONG) AddRef() { ULONG lRef = InterlockedIncrement( &m_cRef ); return max(ULONG(lRef),1ul); }
+  STDMETHODIMP_(ULONG) Release() { ULONG lRef = InterlockedDecrement( &m_cRef ); if (lRef == 0) { m_cRef++; delete this; return 0; } return max(ULONG(lRef),1ul); }
+
+public:
   REFERENCE_TIME rtStart;  ///< Start Time
   REFERENCE_TIME rtStop;   ///< Stop time
   ULONGLONG id;            ///< Unique Identifier (same ID = same subtitle)
@@ -31,7 +51,10 @@ typedef struct LAVSubRect {
   bool forced;             ///< Forced/Menu
 
   bool freePixels;         ///< If true, pixel data is free'ed upon destroy
-} LAVSubRect;
+
+private:
+  ULONG m_cRef;
+};
 
 class CLAVSubtitleFrame : public ISubRenderFrame, public CUnknown
 {
@@ -48,8 +71,7 @@ public:
 
   STDMETHODIMP SetOutputRect(RECT outputRect);
   STDMETHODIMP SetClipRect(RECT clipRect);
-  STDMETHODIMP AddBitmap(const LAVSubRect &subRect);
-  STDMETHODIMP AddBitmap(ULONGLONG id, POINT position, SIZE size, LPVOID pixels, int pitch);
+  STDMETHODIMP AddBitmap(CLAVSubRect *subRect);
 
   BOOL Empty() const { return m_NumBitmaps == 0; };
 
@@ -57,6 +79,6 @@ private:
   RECT m_outputRect;
   RECT m_clipRect;
 
-  LAVSubRect *m_Bitmaps;
+  CLAVSubRect **m_Bitmaps;
   int m_NumBitmaps;
 };
diff --git a/decoder/LAVVideo/subtitles/LAVSubtitleProvider.cpp b/decoder/LAVVideo/subtitles/LAVSubtitleProvider.cpp
index cea94c92..6b41fe3b 100644
--- a/decoder/LAVVideo/subtitles/LAVSubtitleProvider.cpp
+++ b/decoder/LAVVideo/subtitles/LAVSubtitleProvider.cpp
@@ -121,16 +121,14 @@ STDMETHODIMP CLAVSubtitleProvider::RequestFrame(REFERENCE_TIME start, REFERENCE_
   {
     CAutoLock lock(this);
     for (auto it = m_SubFrames.begin(); it != m_SubFrames.end(); it++) {
-      LAVSubRect *pRect = *it;
+      CLAVSubRect *pRect = *it;
       if ((pRect->rtStart == AV_NOPTS_VALUE) || ((pRect->rtStop == AV_NOPTS_VALUE || pRect->rtStop > mid) && pRect->rtStart <= mid)
         && (m_bComposit || pRect->forced)) {
 
-        LAVSubRect rect = *pRect;
-
         if (m_pHLI && PTS2RT(m_pHLI->StartPTM) <= mid && PTS2RT(m_pHLI->EndPTM) >= mid) {
-          ProcessDVDHLI(rect);
+          pRect = ProcessDVDHLI(pRect);
         }
-        subtitleFrame->AddBitmap(rect);
+        subtitleFrame->AddBitmap(pRect);
       }
     }
   }
@@ -211,9 +209,7 @@ void CLAVSubtitleProvider::ClearSubtitleRects()
 {
   CAutoLock lock(this);
   for (auto it = m_SubFrames.begin(); it != m_SubFrames.end(); it++) {
-    CoTaskMemFree((LPVOID)(*it)->pixels);
-    CoTaskMemFree((LPVOID)(*it)->pixelsPal);
-    delete *it;
+    (*it)->Release();
   }
   m_SubFrames.clear();
 }
@@ -226,9 +222,7 @@ void CLAVSubtitleProvider::TimeoutSubtitleRects(REFERENCE_TIME rt)
   while (it != m_SubFrames.end()) {
     if ((*it)->rtStop != AV_NOPTS_VALUE && (*it)->rtStop < timestamp) {
       DbgLog((LOG_TRACE, 10, L"Timed out subtitle at %I64d", (*it)->rtStart));
-      CoTaskMemFree((LPVOID)(*it)->pixels);
-      CoTaskMemFree((LPVOID)(*it)->pixelsPal);
-      delete *it;
+      (*it)->Release();
       it = m_SubFrames.erase(it);
     } else {
       it++;
@@ -421,7 +415,7 @@ void CLAVSubtitleProvider::ProcessSubtitleRect(AVSubtitleRect *rect, REFERENCE_T
   // Store the rect
   POINT position = { rect->x - hpad, rect->y - vpad };
   SIZE size = { width, height };
-  LAVSubRect *lavRect = new LAVSubRect();
+  CLAVSubRect *lavRect = new CLAVSubRect();
   if (!lavRect) return;
   lavRect->id       = m_SubPicId++;
   lavRect->pitch    = rgbStride;
@@ -466,9 +460,10 @@ void CLAVSubtitleProvider::ProcessSubtitleRect(AVSubtitleRect *rect, REFERENCE_T
   AddSubtitleRect(lavRect);
 }
 
-void CLAVSubtitleProvider::AddSubtitleRect(LAVSubRect *rect)
+void CLAVSubtitleProvider::AddSubtitleRect(CLAVSubRect *rect)
 {
   CAutoLock lock(this);
+  rect->AddRef();
   m_SubFrames.push_back(rect);
 }
 
@@ -539,30 +534,36 @@ STDMETHODIMP CLAVSubtitleProvider::SetDVDHLI(struct _AM_PROPERTY_SPHLI *pHLI)
   return S_OK;
 }
 
-void CLAVSubtitleProvider::ProcessDVDHLI(LAVSubRect &rect)
+CLAVSubRect* CLAVSubtitleProvider::ProcessDVDHLI(CLAVSubRect *rect)
 {
   DVDSubContext *ctx = (DVDSubContext *)m_pAVCtx->priv_data;
-  if (!m_pHLI || !rect.pixelsPal || !ctx->has_palette)
-    return;
+  if (!m_pHLI || !rect->pixelsPal || !ctx->has_palette)
+    return rect;
+
+  LPVOID newPixels = CoTaskMemAlloc(rect->pitch * rect->size.cy * 4);
+  if (!newPixels) return rect;
+
+  // copy pixels before modification
+  memcpy(newPixels, rect->pixels, rect->pitch * rect->size.cy * 4);
 
-  LPVOID newPixels = CoTaskMemAlloc(rect.pitch * rect.size.cy * 4);
-  if (!newPixels) return;
-  memcpy(newPixels, rect.pixels, rect.pitch * rect.size.cy * 4);
+  uint8_t *originalPalPixels = (uint8_t *)rect->pixelsPal;
 
-  rect.pixels = newPixels;
-  rect.freePixels = true;
+  // create new object
+  rect = new CLAVSubRect(*rect);
+  rect->pixels = newPixels;
+  rect->pixelsPal = NULL;
 
   // Need to assign a new Id since we're modifying it here..
-  rect.id = m_SubPicId++;
+  rect->id = m_SubPicId++;
 
   uint8_t *palette   = (uint8_t *)ctx->palette;
-  for (int y = 0; y < rect.size.cy; y++) {
-    if (y+rect.position.y < m_pHLI->StartY || y+rect.position.y > m_pHLI->StopY)
+  for (int y = 0; y < rect->size.cy; y++) {
+    if (y+rect->position.y < m_pHLI->StartY || y+rect->position.y > m_pHLI->StopY)
       continue;
-    uint8_t *pixelsPal = ((uint8_t *)rect.pixelsPal) + rect.pitch * y;
-    uint8_t *pixels = ((uint8_t *)rect.pixels) + rect.pitch * y * 4;
-    for (int x = 0; x < rect.size.cx; x++) {
-      if (x+rect.position.x < m_pHLI->StartX || x+rect.position.x > m_pHLI->StopX)
+    uint8_t *pixelsPal = originalPalPixels + rect->pitch * y;
+    uint8_t *pixels = ((uint8_t *)rect->pixels) + rect->pitch * y * 4;
+    for (int x = 0; x < rect->size.cx; x++) {
+      if (x+rect->position.x < m_pHLI->StartX || x+rect->position.x > m_pHLI->StopX)
         continue;
       uint8_t idx = pixelsPal[x];
       uint8_t alpha;
@@ -596,6 +597,7 @@ void CLAVSubtitleProvider::ProcessDVDHLI(LAVSubRect &rect)
       pixels[(x << 2) + 3] = a;
     }
   }
+  return rect;
 }
 
 STDMETHODIMP CLAVSubtitleProvider::SetDVDComposit(BOOL bComposit)
diff --git a/decoder/LAVVideo/subtitles/LAVSubtitleProvider.h b/decoder/LAVVideo/subtitles/LAVSubtitleProvider.h
index eeff6edf..d2ae3912 100644
--- a/decoder/LAVVideo/subtitles/LAVSubtitleProvider.h
+++ b/decoder/LAVVideo/subtitles/LAVSubtitleProvider.h
@@ -63,8 +63,8 @@ private:
 
   void ProcessSubtitleFrame(AVSubtitle *sub, REFERENCE_TIME rtStart, REFERENCE_TIME rtStop);
   void ProcessSubtitleRect(AVSubtitleRect *rect, REFERENCE_TIME rtStart, REFERENCE_TIME rtStop);
-  void AddSubtitleRect(LAVSubRect *rect);
-  void ProcessDVDHLI(LAVSubRect &rect);
+  void AddSubtitleRect(CLAVSubRect *rect);
+  CLAVSubRect* ProcessDVDHLI(CLAVSubRect *rect);
   void ClearSubtitleRects();
   void TimeoutSubtitleRects(REFERENCE_TIME rtStop);
 
@@ -82,7 +82,7 @@ private:
   ULONGLONG             m_SubPicId;
   BOOL                  m_bComposit;
 
-  std::list<LAVSubRect*> m_SubFrames;
+  std::list<CLAVSubRect *> m_SubFrames;
 
   struct _AM_PROPERTY_SPHLI *m_pHLI;
 };