From 726cc895d80ae074277042283b8f4b53915affce Mon Sep 17 00:00:00 2001
From: He Junyan <junyan.he@intel.com>
Date: Thu, 10 Nov 2022 20:36:05 +0800
Subject: [PATCH] va: encoder: Add copy_output_data() helper function

Part-of: <https://gitlab.freedesktop.org/gstreamer/gstreamer/-/merge_requests/3015>
---
 .../gst-plugins-bad/sys/va/gstvabaseenc.c     | 47 +++++++++++++++++++
 .../gst-plugins-bad/sys/va/gstvabaseenc.h     |  5 +-
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.c b/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.c
index 9a9da9ef8a..8188a8e951 100644
--- a/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.c
+++ b/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.c
@@ -330,6 +330,53 @@ error:
   return NULL;
 }
 
+/* Return 0 means error and -1 means not enough data. */
+gint
+gst_va_base_enc_copy_output_data (GstVaBaseEnc * base,
+    GstVaEncodePicture * picture, guint8 * data, gint size)
+{
+  guint coded_size;
+  VASurfaceID surface;
+  VACodedBufferSegment *seg, *seg_list;
+  gint ret_sz = 0;
+
+  /* Wait for encoding to finish */
+  surface = gst_va_encode_picture_get_raw_surface (picture);
+  if (!va_sync_surface (base->display, surface))
+    goto out;
+
+  seg_list = NULL;
+  if (!va_map_buffer (base->display, GST_MAP_READ, picture->coded_buffer,
+          (gpointer *) & seg_list))
+    goto out;
+
+  if (!seg_list) {
+    va_unmap_buffer (base->display, picture->coded_buffer);
+    GST_WARNING_OBJECT (base, "coded buffer has no segment list");
+    goto out;
+  }
+
+  coded_size = 0;
+  for (seg = seg_list; seg; seg = seg->next)
+    coded_size += seg->size;
+
+  if (coded_size > size) {
+    GST_DEBUG_OBJECT (base, "Not enough space for coded data");
+    ret_sz = -1;
+    goto out;
+  }
+
+  for (seg = seg_list; seg; seg = seg->next) {
+    memcpy (data + ret_sz, seg->buf, seg->size);
+    ret_sz += seg->size;
+  }
+
+  va_unmap_buffer (base->display, picture->coded_buffer);
+
+out:
+  return ret_sz;
+}
+
 static GstAllocator *
 _allocator_from_caps (GstVaBaseEnc * base, GstCaps * caps)
 {
diff --git a/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.h b/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.h
index 68092bad0f..bb6491d095 100644
--- a/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.h
+++ b/subprojects/gst-plugins-bad/sys/va/gstvabaseenc.h
@@ -139,7 +139,10 @@ void                  gst_va_base_enc_add_codec_tag       (GstVaBaseEnc * base,
 void                  gst_va_base_enc_reset_state         (GstVaBaseEnc * base);
 GstBuffer *           gst_va_base_enc_create_output_buffer (GstVaBaseEnc * base,
                                                             GstVaEncodePicture * picture);
-
+gint                  gst_va_base_enc_copy_output_data    (GstVaBaseEnc * base,
+                                                           GstVaEncodePicture * picture,
+                                                           guint8 * data,
+                                                           gint size);
 void                  gst_va_base_enc_update_property_uint (GstVaBaseEnc * base,
                                                             guint32 * old_val,
                                                             guint32 new_val,