CoopAllTheThings/host/src/capture/shared_texture.cpp

#include "capture/shared_texture.hpp"

#include "capture/dxgi_format.hpp"
#include "coop/protocol.hpp"
#include "coop/shared_memory.hpp"

namespace coop
{

bool SharedTextureSource::init(ID3D11Device* device)
{
	if (device == nullptr || FAILED(device->QueryInterface(IID_PPV_ARGS(&device_))))
	{
		return false;
	}
	device_->GetImmediateContext(&ctx_);
	return ctx_ != nullptr;
}

void SharedTextureSource::reset()
{
	srv_.Reset();
	private_.Reset();
	mutex_.Reset();
	shared_.Reset();
	pid_ = 0;
	width_ = height_ = format_ = 0;
	last_generation_ = 0;
	frames_copied_ = 0;
	frames_missed_ = 0;
}

bool SharedTextureSource::reopen(unsigned long pid, const VideoShareView& share)
{
	srv_.Reset();
	private_.Reset();
	mutex_.Reset();
	shared_.Reset();
	width_ = height_ = format_ = 0;
	pid_ = pid;

	if (pid == 0 || share.width == 0 || share.height == 0)
	{
		return false; // the hook hasn't shared a backbuffer yet
	}

	const std::wstring name = video_share_name(pid);
	if (FAILED(device_->OpenSharedResourceByName(name.c_str(),
	                                             DXGI_SHARED_RESOURCE_READ | DXGI_SHARED_RESOURCE_WRITE,
	                                             IID_PPV_ARGS(&shared_))) ||
	    shared_ == nullptr)
	{
		return false;
	}
	if (FAILED(shared_.As(&mutex_)) || mutex_ == nullptr)
	{
		shared_.Reset();
		return false;
	}

	// Private copy we sample from, so we only hold the keyed mutex during the copy.
	// Use the plain-UNORM sibling of the (possibly sRGB) backbuffer format so the
	// SRV passes the bytes through without an sRGB->linear decode that would darken
	// the mirror (CopyResource is allowed within the shared typeless group).
	D3D11_TEXTURE2D_DESC desc{};
	desc.Width = share.width;
	desc.Height = share.height;
	desc.MipLevels = 1;
	desc.ArraySize = 1;
	desc.Format = srgb_to_unorm(static_cast<DXGI_FORMAT>(share.format));
	desc.SampleDesc.Count = 1;
	desc.Usage = D3D11_USAGE_DEFAULT;
	desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	if (FAILED(device_->CreateTexture2D(&desc, nullptr, &private_)) || private_ == nullptr)
	{
		mutex_.Reset();
		shared_.Reset();
		return false;
	}
	if (FAILED(device_->CreateShaderResourceView(private_.Get(), nullptr, &srv_)))
	{
		srv_.Reset();
		private_.Reset();
		mutex_.Reset();
		shared_.Reset();
		return false;
	}

	width_ = share.width;
	height_ = share.height;
	format_ = share.format;
	return true;
}

bool SharedTextureSource::read_pixel(std::uint32_t x, std::uint32_t y, std::uint8_t out[4])
{
	if (private_ == nullptr || ctx_ == nullptr || device_ == nullptr)
	{
		return false;
	}
	D3D11_TEXTURE2D_DESC desc{};
	private_->GetDesc(&desc);
	if (x >= desc.Width || y >= desc.Height)
	{
		return false;
	}
	D3D11_TEXTURE2D_DESC staging_desc = desc;
	staging_desc.Usage = D3D11_USAGE_STAGING;
	staging_desc.BindFlags = 0;
	staging_desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	staging_desc.MiscFlags = 0;
	Microsoft::WRL::ComPtr<ID3D11Texture2D> staging;
	if (FAILED(device_->CreateTexture2D(&staging_desc, nullptr, &staging)))
	{
		return false;
	}
	ctx_->CopyResource(staging.Get(), private_.Get());
	D3D11_MAPPED_SUBRESOURCE map{};
	if (FAILED(ctx_->Map(staging.Get(), 0, D3D11_MAP_READ, 0, &map)))
	{
		return false;
	}
	const auto* px = static_cast<const std::uint8_t*>(map.pData) + static_cast<std::size_t>(y) * map.RowPitch +
	                 static_cast<std::size_t>(x) * 4; // R8G8B8A8_UNORM
	out[0] = px[0];
	out[1] = px[1];
	out[2] = px[2];
	out[3] = px[3];
	ctx_->Unmap(staging.Get(), 0);
	return true;
}

bool SharedTextureSource::update(const VideoShareView& share, unsigned long pid)
{
	if (device_ == nullptr || pid == 0)
	{
		reset();
		return false;
	}

	// (Re)open whenever the target or the published backbuffer geometry changes.
	if (pid != pid_ || share.width != width_ || share.height != height_ || share.format != format_)
	{
		if (!reopen(pid, share))
		{
			return srv_ != nullptr; // couldn't open yet; keep any prior frame
		}
		last_generation_ = 0; // force a copy of the current frame
	}

	if (shared_ == nullptr || mutex_ == nullptr)
	{
		return srv_ != nullptr;
	}
	if (share.generation == last_generation_)
	{
		return srv_ != nullptr; // no new frame; keep showing the last copy
	}

	// Bounded wait so a stalled producer can't hang the host's render thread.
	if (mutex_->AcquireSync(kVideoMutexKey, 8) == S_OK)
	{
		ctx_->CopyResource(private_.Get(), shared_.Get());
		mutex_->ReleaseSync(kVideoMutexKey);
		// Generations between the last copy and this one were published but never shown
		// (we only ever copy the newest). last_generation_ == 0 is the first copy after a
		// (re)open, where the gap to a large generation is meaningless, so skip it.
		if (last_generation_ != 0 && share.generation > last_generation_ + 1)
		{
			frames_missed_ += share.generation - last_generation_ - 1;
		}
		last_generation_ = share.generation;
		++frames_copied_;
	}
	return srv_ != nullptr;
}

} // namespace coop