From f62c8658e435cbdb854d8871ba92fc0769120c0c Mon Sep 17 00:00:00 2001
From: Arun Raghavan <arun@asymptotic.io>
Date: Tue, 10 Sep 2024 12:50:24 -0400
Subject: [PATCH] gstreamer: Add API to take an event and buffers in a pad
 probe

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer-rs/-/merge_requests/1448>
---
 gstreamer/src/pad.rs | 54 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/gstreamer/src/pad.rs b/gstreamer/src/pad.rs
index 490cff89f..0419a14a9 100644
--- a/gstreamer/src/pad.rs
+++ b/gstreamer/src/pad.rs
@@ -113,6 +113,47 @@ impl<'a> PadProbeInfo<'a> {
             _ => None,
         }
     }
+
+    // rustdoc-stripper-ignore-next
+    /// Takes over the buffer in the probe info. As the data is no longer valid for the caller, the
+    /// probe will be considered dropped after this point.
+    pub fn take_buffer(&mut self) -> Option<Buffer> {
+        if matches!(self.data, Some(PadProbeData::Buffer(..))) {
+            match self.data.take() {
+                Some(PadProbeData::Buffer(b)) => Some(b),
+                _ => unreachable!(),
+            }
+        } else {
+            None
+        }
+    }
+
+    // rustdoc-stripper-ignore-next
+    /// Takes over the buffer in the probe info. As the data is no longer valid for the caller, the
+    /// probe will be considered dropped after this point.
+    pub fn take_buffer_list(&mut self) -> Option<BufferList> {
+        if matches!(self.data, Some(PadProbeData::BufferList(..))) {
+            match self.data.take() {
+                Some(PadProbeData::BufferList(b)) => Some(b),
+                _ => unreachable!(),
+            }
+        } else {
+            None
+        }
+    }
+    // rustdoc-stripper-ignore-next
+    /// Takes over the event in the probe info. As the data is no longer valid for the caller, the
+    /// probe will be considered dropped after this point.
+    pub fn take_event(&mut self) -> Option<Event> {
+        if matches!(self.data, Some(PadProbeData::Event(..))) {
+            match self.data.take() {
+                Some(PadProbeData::Event(e)) => Some(e),
+                _ => unreachable!(),
+            }
+        } else {
+            None
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -1007,6 +1048,19 @@ unsafe fn update_probe_info(
             }
             other => panic!("Bad data for {data_type:?} pad probe returning Handled: {other:?}"),
         }
+    } else if ret == PadProbeReturn::Drop {
+        // We may have consumed the object via PadProbeInfo::take_*() functions
+        match probe_info.data {
+            None if data_type == Some(Buffer::static_type())
+                || data_type == Some(BufferList::static_type())
+                || data_type == Some(Event::static_type()) =>
+            {
+                (*info).data = ptr::null_mut();
+            }
+            _ => {
+                // Nothing to do, it's going to be dropped
+            }
+        }
     } else {
         match probe_info.data {
             Some(PadProbeData::Buffer(buffer)) => {