abstract frameBuffer_ into a FrameBuffer class

Author: NicolasHug

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -133,10 +133,6 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
   TORCH_CHECK(g_cuda_beta, "BetaCudaDeviceInterface was not registered!");
   TORCH_CHECK(
       device_.type() == torch::kCUDA, "Unsupported device: ", device_.str());
-
-  // TODONVDEC P1: init size should probably be min_num_decode_surfaces from
-  // video format
-  frameBuffer_.resize(4);
 }
 
 BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
@@ -344,7 +340,7 @@ int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* pPicParams) {
   dispInfo.repeat_first_field = 0;
   dispInfo.timestamp = guessedPts;
 
-  FrameBufferSlot* slot = findEmptySlot();
+  FrameBuffer::Slot* slot = frameBuffer_.findEmptySlot();
   slot->dispInfo = dispInfo;
   slot->guessedPts = guessedPts;
   slot->occupied = true;
@@ -358,7 +354,7 @@ int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* pPicParams) {
 int BetaCudaDeviceInterface::receiveFrame(
     UniqueAVFrame& avFrame,
     int64_t desiredPts) {
-  FrameBufferSlot* slot = findFrameWithExactPts(desiredPts);
+  FrameBuffer::Slot* slot = frameBuffer_.findFrameWithExactPts(desiredPts);
   if (slot == nullptr) {
     // No frame found, instruct caller to try again later after sending more
     // packets.
@@ -532,9 +528,8 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
       preAllocatedOutputTensor);
 }
 
-// TODONVDEC P0: Don't let buffer grow indefinitely.
-BetaCudaDeviceInterface::FrameBufferSlot*
-BetaCudaDeviceInterface::findEmptySlot() {
+BetaCudaDeviceInterface::FrameBuffer::Slot*
+BetaCudaDeviceInterface::FrameBuffer::findEmptySlot() {
   for (auto& slot : frameBuffer_) {
     if (!slot.occupied) {
       return &slot;
@@ -544,8 +539,9 @@ BetaCudaDeviceInterface::findEmptySlot() {
   return &frameBuffer_.back();
 }
 
-BetaCudaDeviceInterface::FrameBufferSlot*
-BetaCudaDeviceInterface::findFrameWithExactPts(int64_t desiredPts) {
+BetaCudaDeviceInterface::FrameBuffer::Slot*
+BetaCudaDeviceInterface::FrameBuffer::findFrameWithExactPts(
+    int64_t desiredPts) {
   for (auto& slot : frameBuffer_) {
     if (slot.occupied && slot.guessedPts == desiredPts) {
       return &slot;

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -64,6 +64,40 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   int frameReadyForDecoding(CUVIDPICPARAMS* pPicParams);
 
  private:
+  class FrameBuffer {
+   public:
+    struct Slot {
+      CUVIDPARSERDISPINFO dispInfo;
+      int64_t guessedPts;
+      bool occupied = false;
+
+      Slot() : guessedPts(-1), occupied(false) {
+        std::memset(&dispInfo, 0, sizeof(dispInfo));
+      }
+    };
+
+    // TODONVDEC P1: init size should probably be min_num_decode_surfaces from
+    // video format
+    FrameBuffer() : frameBuffer_(4) {}
+
+    ~FrameBuffer() = default;
+
+    Slot* findEmptySlot();
+    Slot* findFrameWithExactPts(int64_t desiredPts);
+
+    // Iterator support for range-based for loops
+    auto begin() {
+      return frameBuffer_.begin();
+    }
+
+    auto end() {
+      return frameBuffer_.end();
+    }
+
+   private:
+    std::vector<Slot> frameBuffer_;
+  };
+
   UniqueAVFrame convertCudaFrameToAVFrame(
       CUdeviceptr framePtr,
       unsigned int pitch,
@@ -73,19 +107,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
   UniqueCUvideodecoder decoder_;
   CUVIDEOFORMAT videoFormat_ = {};
 
-  struct FrameBufferSlot {
-    CUVIDPARSERDISPINFO dispInfo;
-    int64_t guessedPts;
-    bool occupied = false;
-
-    FrameBufferSlot() : guessedPts(-1), occupied(false) {
-      std::memset(&dispInfo, 0, sizeof(dispInfo));
-    }
-  };
-
-  std::vector<FrameBufferSlot> frameBuffer_;
-  FrameBufferSlot* findEmptySlot();
-  FrameBufferSlot* findFrameWithExactPts(int64_t desiredPts);
+  FrameBuffer frameBuffer_;
 
   std::queue<int64_t> packetsPtsQueue;