hazelnut/agent/CollabVMFbcAgent/capture_nvfbc.cpp

#include "capture.hpp"

// clang-format off
#include "Utils.hpp"

#include "NvFBCLibrary.h"
#include <NvFBC/nvFBCToSys.h>
// clang-format on

namespace hazelnut {

	class FramebufferCapture_NVFBC final : public IFramebufferCapture {
		NVFBCRESULT nvfbcStatus;
		NvFBCLibrary nvfbcLib;

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
		NvFBCFrameGrabInfo grabInfo;
		NvFBCToSys* nvfbc;

		u32 width;
		u32 height;

		u32* pRawFramebuffer;
		unique_buffer<u32> convertedFramebuffer;

		u8* pDiffMap;
		u32 diffmapWidth;
		u32 diffmapHeight;

	   public:
		virtual ~FramebufferCapture_NVFBC() { Shutdown(); }

		void Shutdown() {
			NvfbcDestroyInstance();
			nvfbcLib.close();
		}

		bool Initialize() override {
			if(!nvfbcLib.load()) {
				return false;
			}

			if(!NvfbcInitSession())
				return false;

			return true;
		}

		void NvfbcDestroyInstance() {
			if(nvfbc) {
				nvfbc->NvFBCToSysRelease();
				nvfbc = nullptr;

				// reset state while we're here I guess
				pRawFramebuffer = nullptr;
				pDiffMap = nullptr;

				width = 0;
				height = 0;
				diffmapWidth = 0;
				diffmapHeight = 0;
			}
		}

		bool NvfbcCreateInstance() {
			// Destroy an existing NVFBC session to avoid memory leak
			if(nvfbc) {
				NvfbcDestroyInstance();
			}

			DWORD maxDisplayWidth = -1;
			DWORD maxDisplayHeight = -1;
			nvfbc = (NvFBCToSys*)nvfbcLib.create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);

			if(!nvfbc) {
				return false;
			}

			return true;
		}

		bool NvfbcInitSession() {
			if(!NvfbcCreateInstance())
				return false;

			// set up a session
			NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams = { 0 };
			fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
			fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
			fbcSysSetupParams.bWithHWCursor = true;
			fbcSysSetupParams.ppBuffer = (void**)&pRawFramebuffer;

			// enable a 32x32 difference map
			fbcSysSetupParams.bDiffMap = TRUE;
			fbcSysSetupParams.ppDiffMap = (void**)&pDiffMap;
			fbcSysSetupParams.eDiffMapBlockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;

			auto status = nvfbc->NvFBCToSysSetUp(&fbcSysSetupParams);
			if(status != NVFBC_SUCCESS)
				return false;

			memset(&fbcSysGrabParams, 0, sizeof(NVFBC_TOSYS_GRAB_FRAME_PARAMS));

			// set up grab params
			fbcSysGrabParams.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
			fbcSysGrabParams.dwFlags = NVFBC_TOSYS_WAIT_WITH_TIMEOUT;
			fbcSysGrabParams.dwTargetWidth = 0;
			fbcSysGrabParams.dwTargetHeight = 0;
			fbcSysGrabParams.dwStartX = 0;
			fbcSysGrabParams.dwStartY = 0;
			fbcSysGrabParams.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL;
			fbcSysGrabParams.pNvFBCFrameGrabInfo = &grabInfo;
			fbcSysGrabParams.dwWaitTime = 16;
			return true;
		}

		DisplayCaptureResult CaptureFrame() override {
			auto nvStatus = nvfbc->NvFBCToSysGrabFrame(&this->fbcSysGrabParams);

			switch(nvStatus) {
				case NVFBC_SUCCESS: break;

				// Need to recreate the session. If it fails then we fail too.
				case NVFBC_ERROR_INVALIDATED_SESSION: {
					if(!NvfbcInitSession())
						return DisplayCaptureResult::Fail;

					// Recurse. This looks naughty, but will allow us to directly retry
					// the capture. (Plus if this causes issues whatever has happened is probably beyond saving)
					return CaptureFrame();
				} break;

				default: return DisplayCaptureResult::Fail;
			}


			if(width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
				width = grabInfo.dwWidth;
				height = grabInfo.dwHeight;

				convertedFramebuffer.resize(width * height);

				return DisplayCaptureResult::OkButResized;
			}

			// Splat to the converted framebuffer. It doesn't have padding on it.
			auto* pBufferData = (u8*)convertedFramebuffer.data();
			auto* pSrcData = (u8*)&pRawFramebuffer[0];

			for(u32 y = 0; y < grabInfo.dwHeight; ++y) {
				usize srcStart = (y * grabInfo.dwBufferWidth) * 4;
				usize dstStart = (y * width) * 4;

				// Convert to BGRA
				for(u32 x = 0; x < grabInfo.dwWidth * 4; x += 4) {
					pBufferData[(dstStart + x) + 0] = pSrcData[(srcStart + x) + 2]; // B
					pBufferData[(dstStart + x) + 1] = pSrcData[(srcStart + x) + 1]; // G
					pBufferData[(dstStart + x) + 2] = pSrcData[(srcStart + x) + 0]; // R
					pBufferData[(dstStart + x) + 3] = 0xff;							// A
				}

				// memcpy(&pBufferData[y * width], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], grabInfo.dwWidth * 4);
			}

			// sleep
			// nvfbc->NvFBCToSysGPUBasedCPUSleep(16000);

			return DisplayCaptureResult::Ok;
		}

		FramebufferInformation GetFramebufferInformation() override { return FramebufferInformation { convertedFramebuffer.data(), width, height }; }

		DiffInformation GetDiffInformation() override {
			diffmapWidth = (u32)ceil((f32)width / 32);
			diffmapHeight = (u32)ceil((f32)height / 32);
			return DiffInformation { pDiffMap, diffmapWidth, diffmapHeight };
		}
	};

	IFramebufferCapture* CreateFramebufferCapture_NVFBC() {
		return new FramebufferCapture_NVFBC();
	}
} // namespace hazelnut