From d44d32071707aea4c79dd8e5a98bb0bc4bf0f47e Mon Sep 17 00:00:00 2001
From: Aryan Salmanpour <aryan.salmanpour@amd.com>
Date: Thu, 9 Nov 2023 08:47:58 -0500
Subject: [PATCH] Add support for rocDecGetDecodeStatus API (#46)

[ROCm/rocdecode commit: 6ab31ce40e4c734a65d3dbd620da5027f5286c14]
---
 projects/rocdecode/api/rocdecode.h            |  1 +
 .../rocdecode/src/rocdecode/roc_decoder.cpp   | 10 ++++----
 .../rocdecode/src/rocdecode/rocdecode_api.cpp |  3 +++
 .../rocdecode/vaapi/vaapi_videodecoder.cpp    | 24 ++++++++++++++++++-
 .../src/rocdecode/vaapi/vaapi_videodecoder.h  |  1 +
 5 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/projects/rocdecode/api/rocdecode.h b/projects/rocdecode/api/rocdecode.h
index 611060618f..51f204c0ac 100644
--- a/projects/rocdecode/api/rocdecode.h
+++ b/projects/rocdecode/api/rocdecode.h
@@ -138,6 +138,7 @@ typedef enum rocDecodeStatus_enum {
     // 3 to 7 enums are reserved for future use
     rocDecodeStatus_Error           = 8,   // Decode is completed with an error (error is not concealed)
     rocDecodeStatus_Error_Concealed = 9,   // Decode is completed with an error and error is concealed 
+    rocDecodeStatus_Displaying      = 10,  // Decode is completed, displaying in progress
 } rocDecDecodeStatus;
 
 /**************************************************************************************************************/
diff --git a/projects/rocdecode/src/rocdecode/roc_decoder.cpp b/projects/rocdecode/src/rocdecode/roc_decoder.cpp
index c37b4efddb..78ba95e220 100644
--- a/projects/rocdecode/src/rocdecode/roc_decoder.cpp
+++ b/projects/rocdecode/src/rocdecode/roc_decoder.cpp
@@ -64,10 +64,12 @@ rocDecStatus RocDecoder::decodeFrame(RocdecPicParams *pPicParams) {
 }
 
 rocDecStatus RocDecoder::getDecodeStatus(int nPicIdx, RocdecDecodeStatus* pDecodeStatus) {
-    // todo:: return appropriate decStatus
-    // init vaapi decoder to get the decoding status of the picture specified by nPicIndex 
-    // return status
-    return ROCDEC_NOT_IMPLEMENTED;
+    rocDecStatus rocdec_status = ROCDEC_SUCCESS;
+    rocdec_status = va_video_decoder_.GetDecodeStatus(nPicIdx, pDecodeStatus);
+    if (rocdec_status != ROCDEC_SUCCESS) {
+        ERR("ERROR: Failed to query the decode status!" + TOSTR(rocdec_status));
+    }
+    return rocdec_status;
 }
 
 rocDecStatus RocDecoder::reconfigureDecoder(RocdecReconfigureDecoderInfo *pDecReconfigParams) {
diff --git a/projects/rocdecode/src/rocdecode/rocdecode_api.cpp b/projects/rocdecode/src/rocdecode/rocdecode_api.cpp
index 73b9b1ce61..3e87275788 100644
--- a/projects/rocdecode/src/rocdecode/rocdecode_api.cpp
+++ b/projects/rocdecode/src/rocdecode/rocdecode_api.cpp
@@ -120,6 +120,9 @@ rocDecDecodeFrame(rocDecDecoderHandle hDecoder, RocdecPicParams *pPicParams) {
 /************************************************************************************************************/
 rocDecStatus ROCDECAPI 
 rocDecGetDecodeStatus(rocDecDecoderHandle hDecoder, int nPicIdx, RocdecDecodeStatus* pDecodeStatus) {
+    if (hDecoder == nullptr || pDecodeStatus == nullptr) {
+        return ROCDEC_INVALID_PARAMETER;
+    }
     auto handle = static_cast<DecHandle *> (hDecoder);
     rocDecStatus ret;
     try {
diff --git a/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.cpp b/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.cpp
index ada2a91e0b..7f1c56290b 100644
--- a/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.cpp
+++ b/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.cpp
@@ -151,4 +151,26 @@ rocDecStatus VaapiVideoDecoder::CreateContext() {
 rocDecStatus VaapiVideoDecoder::SubmitDecode(RocdecPicParams *pPicParams) {
     // Todo copy pic param, slice param, IQ matrix and slice data from RocdecPicParams to VAAPI struct buffers, then submit to VAAPI driver.
     return ROCDEC_SUCCESS;
-}
\ No newline at end of file
+}
+
+rocDecStatus VaapiVideoDecoder::GetDecodeStatus(int pic_idx, RocdecDecodeStatus *decode_status) {
+    VASurfaceStatus va_surface_status;
+    if (pic_idx >= va_surface_ids_.size() || decode_status == nullptr) {
+        return ROCDEC_INVALID_PARAMETER;
+    }
+    CHECK_VAAPI(vaQuerySurfaceStatus(va_display_, va_surface_ids_[pic_idx], &va_surface_status));
+    switch (va_surface_status) {
+        case VASurfaceRendering:
+            decode_status->decodeStatus = rocDecodeStatus_InProgress;
+            break;
+        case VASurfaceReady:
+            decode_status->decodeStatus = rocDecodeStatus_Success;
+            break;
+        case VASurfaceDisplaying:
+            decode_status->decodeStatus = rocDecodeStatus_Displaying;
+            break;
+        default:
+           decode_status->decodeStatus = rocDecodeStatus_Invalid;
+    }
+    return ROCDEC_SUCCESS;
+}
diff --git a/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.h b/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.h
index 1929aeb116..25f06bf9e5 100644
--- a/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.h
+++ b/projects/rocdecode/src/rocdecode/vaapi/vaapi_videodecoder.h
@@ -49,6 +49,7 @@ public:
     ~VaapiVideoDecoder();
     rocDecStatus InitializeDecoder(std::string gcn_arch_name);
     rocDecStatus SubmitDecode(RocdecPicParams *pPicParams);
+    rocDecStatus GetDecodeStatus(int pic_idx, RocdecDecodeStatus* decode_status);
 private:
     RocDecoderCreateInfo decoder_create_info_;
     int drm_fd_;