wayland_vaapi_ffmpeg_hardware_player.c

/*
 * Copyright (c) 2021 Nathaniel R. Lewis <github@nrlewis.dev>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of mosquitto nor the names of its
 *    contributors may be used to endorse or promote products derived from
 *    this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#define EGL_EGLEXT_PROTOTYPES

#include <stdio.h>
#include <errno.h>
#include <assert.h>
#include <unistd.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/hwcontext.h>
#include <libavutil/hwcontext_vaapi.h>
#include <libavutil/pixdesc.h>

#include <libdrm/drm_fourcc.h>

#include <va/va.h>
#include <va/va_wayland.h>
#include <va/va_drmcommon.h>

#include <wayland-client.h>
#include <wayland-egl.h>

#include <EGL/egl.h>
#include <EGL/eglext.h>
#include <GLES2/gl2.h>
#include <GLES2/gl2ext.h>

// check out mpv "get_format_hwdec" and setup hardware frames

// display state
struct wl_display *waylandDisplay = NULL;
struct wl_registry *waylandRegistry = NULL;
struct wl_compositor *waylandCompositor = NULL;
struct wl_shell *waylandShell = NULL;
struct wl_surface *waylandSurface = NULL;
struct wl_shell_surface *waylandShellSurface = NULL;

// GL state
struct wl_egl_window *eglWindow = NULL;
EGLDisplay eglDisplay = EGL_NO_DISPLAY;
EGLSurface eglSurface;
EGLContext eglContext;

// VAAPI state
VADisplay vaDisplay = NULL;
AVBufferRef *avDeviceRef = NULL;
AVBufferRef *avHardwareFramesContextRef = NULL;

// video decoding state
AVFormatContext *formatContext = NULL;
AVCodecContext *decodeContext = NULL;
int streamId = -1;

// EGL extensions
PFNEGLCREATEIMAGEKHRPROC createImage;
PFNEGLDESTROYIMAGEKHRPROC destroyImage;
PFNGLEGLIMAGETARGETTEXTURE2DOESPROC imageTargetTexture2D;

// ---------------------- SHADERS ----------------------------------
const char *vertex_shader_source =
	"attribute vec2 position;                 \n"
	"attribute vec2 coordinate_in;            \n"
	"varying vec2 coordinate;                 \n"
	"                                         \n"
	"void main()                              \n"
	"{                                        \n"
	"    coordinate = coordinate_in;          \n"
	"    gl_Position = vec4(position, 0, 1);  \n"
	"}                                        \n";

const char *fragment_shader_source =
	"#extension GL_OES_EGL_image_external : require\n"
	"precision highp float;                        \n"
	"varying vec2 coordinate;                      \n"
	"uniform samplerExternalOES tex;               \n"
	"                                              \n"
	"void main()                                   \n"
	"{                                             \n"
	"    vec4 color = texture2D(tex, coordinate);  \n"
	"    gl_FragColor = vec4(color.rgb, 1);        \n"
	//"    gl_FragColor = vec4(coordinate, 0.0, 1);  \n"
	"}                                             \n";

// ---------------------- UTILITIES -------------------------------

GLuint compileShader(const char *source, GLenum type)
{
	GLuint shader = glCreateShader(type);
	glShaderSource(shader, 1, &source, NULL);
	glCompileShader(shader);

	int ret = 0;
	glGetShaderiv(shader, GL_COMPILE_STATUS, &ret);
	if (!ret)
	{
		glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &ret);
		if (ret > 1)
		{
			char* log = calloc(ret + 1, 1);
			glGetShaderInfoLog(shader, ret, NULL, log);
			fprintf(stderr, "failed to compile shader with message: %s\n", log);
			free(log);
		}
		else
		{
			fprintf(stderr, "failed to compile shader\n");
		}

		glDeleteShader(shader);
		return 0;
	}

	return shader;
}

GLuint linkProgram(GLuint vertex_shader, GLuint fragment_shader)
{
	GLuint program = glCreateProgram();

	glAttachShader(program, vertex_shader);
	glAttachShader(program, fragment_shader);

	int ret = 0;
	glLinkProgram(program);
	glGetProgramiv(program, GL_LINK_STATUS, &ret);
	if (!ret)
	{
		glGetProgramiv(program, GL_INFO_LOG_LENGTH, &ret);
		if (ret > 1)
		{
			char* log = calloc(ret + 1, 1);
			glGetProgramInfoLog(program, ret, NULL, log);
			fprintf(stderr, "failed to link program with message: %s\n", log);
			free(log);
		}
		else
		{
			fprintf(stderr, "failed to link program\n");
		}

		glDeleteProgram(program);
		return 0;
	}

	return program;
}

// ---------------------- API CALLBACKS ------------------------------

static void shell_surface_ping (void *data, struct wl_shell_surface *shell_surface, uint32_t serial)
{
	wl_shell_surface_pong (shell_surface, serial);
}

static void shell_surface_configure (void *data, struct wl_shell_surface *shell_surface, uint32_t edges, int32_t width, int32_t height)
{
	struct wl_egl_window *window = data;
	wl_egl_window_resize (window, width, height, 0, 0);
}

static void shell_surface_popup_done (void *data, struct wl_shell_surface *shell_surface)
{
}

static struct wl_shell_surface_listener shell_surface_listener = {&shell_surface_ping, &shell_surface_configure, &shell_surface_popup_done};

static void registry_add_object (void *data, struct wl_registry *registry, uint32_t name, const char *interface, uint32_t version)
{
	if (!strcmp(interface,"wl_compositor"))
	{
		waylandCompositor = wl_registry_bind (registry, name, &wl_compositor_interface, version);
	}

	else if (!strcmp(interface,"wl_shell"))
	{
		waylandShell = wl_registry_bind (registry, name, &wl_shell_interface, version);
	}
}

static void registry_remove_object (void *data, struct wl_registry *registry, uint32_t name)
{
}

static struct wl_registry_listener registry_listener = {&registry_add_object, &registry_remove_object};

static void vaapi_error_callback(const char *msg)
{
	fprintf(stderr, "[VAAPI ERROR]: %s\n", msg);
}

static void vaapi_info_callback(const char *msg)
{
	fprintf(stdout, "[VAAPI INFO]: %s\n", msg);
}

static enum AVPixelFormat avcodec_get_format(struct AVCodecContext *avctx, const enum AVPixelFormat *fmt)
{	
	avHardwareFramesContextRef = av_hwframe_ctx_alloc(avDeviceRef);
	if (!avHardwareFramesContextRef)
	{
		fprintf(stderr, "failed to allocate hw frames context.\n");
		return -1;
	}

	AVHWFramesContext *context = (AVHWFramesContext*) avHardwareFramesContextRef->data;
	context->format = AV_PIX_FMT_VAAPI_VLD;
	context->sw_format = AV_PIX_FMT_NV12;
	context->width = avctx->coded_width;
	context->height = avctx->coded_height;
	context->initial_pool_size = 8;

	int ret = av_hwframe_ctx_init(avHardwareFramesContextRef);
	if (ret < 0)
	{
		fprintf(stderr, "failed to initialize hwframe context.\n");
		av_buffer_unref(&avHardwareFramesContextRef);
		exit(EXIT_FAILURE);
	}

	avctx->hw_frames_ctx = av_buffer_ref(avHardwareFramesContextRef);

	return AV_PIX_FMT_VAAPI_VLD;
}

// ---------------------- INITIALIZATION FUNCTIONS ------------------------------

void print_fourcc(int fourcc)
{
	char c[4] = {
		(char) ((fourcc) & 0x000000ff),
		(char) ((fourcc >> 8) & 0x000000ff),
		(char) ((fourcc >> 16) & 0x000000ff),
		(char) ((fourcc >> 24) & 0x000000ff),
	};
	printf("got format: FOURCC('%c', '%c', '%c', '%c')\n", c[0], c[1], c[2], c[3]);
}	

int prepareEGL()
{
	waylandDisplay = wl_display_connect(NULL);
	if (!waylandDisplay)
	{
		fprintf(stderr, "failed to connect to wayland display.\n");
		return -1;
	}

	waylandRegistry = wl_display_get_registry(waylandDisplay);
	wl_registry_add_listener(waylandRegistry, &registry_listener, NULL);
	wl_display_roundtrip(waylandDisplay);

	eglDisplay = eglGetDisplay(waylandDisplay);
	if (eglDisplay == EGL_NO_DISPLAY)
	{
		fprintf(stderr, "failed to get EGL display from wayland.\n");
		return -1;
	}

	int version[2] = {0};
	if (eglInitialize(eglDisplay, &version[0], &version[1]) == EGL_FALSE)
	{
		fprintf(stderr, "failed to initialize EGL.\n");
		return -1;
	}

	printf("got EGL version: %d.%d\n", version[0], version[1]);
	return 0;
}

int prepareEGLWindow()
{
	if (eglBindAPI(EGL_OPENGL_ES_API) == EGL_FALSE)
	{
		fprintf(stderr, "failed to bind opengl es api.\n");
		return -1;
	}

	EGLint attributes[] = {
		EGL_SURFACE_TYPE, EGL_WINDOW_BIT,
		EGL_RED_SIZE, 8,
		EGL_GREEN_SIZE, 8,
		EGL_BLUE_SIZE, 8,
		EGL_ALPHA_SIZE, 8,
		EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT,
		EGL_NONE
	};

	EGLint context_attributes[] = {
		EGL_CONTEXT_CLIENT_VERSION, 2,
		EGL_NONE
	};

	EGLConfig config;
	EGLint num_config;
	eglChooseConfig (eglDisplay, attributes, &config, 1, &num_config);

	eglContext = eglCreateContext(eglDisplay, config, EGL_NO_CONTEXT, context_attributes);
	if (eglContext == EGL_NO_CONTEXT)
	{
		fprintf(stderr, "failed to create egl context.\n");
		return -1;
	}

	waylandSurface = wl_compositor_create_surface(waylandCompositor);
	waylandShellSurface = wl_shell_get_shell_surface(waylandShell, waylandSurface);
	eglWindow = wl_egl_window_create(waylandSurface, decodeContext->coded_width, decodeContext->coded_height);
	eglSurface = eglCreateWindowSurface(eglDisplay, config, eglWindow, NULL);

	wl_shell_surface_add_listener(waylandShellSurface, &shell_surface_listener, eglWindow);
	wl_shell_surface_set_toplevel(waylandShellSurface);

	if (eglMakeCurrent(eglDisplay, eglSurface, eglSurface, eglContext) == EGL_FALSE)
	{
		fprintf(stderr, "failed to make context current.\n");
		return -1;
	}

	return 0;
}

int prepareVAAPI()
{
	vaDisplay = vaGetDisplayWl(waylandDisplay);
	if (!vaDisplay)
	{
		fprintf(stderr, "failed to get vaapi display from wayland.\n");
		return -1;
	}

	int version[2] = {0};
	if (vaInitialize(vaDisplay, &version[0], &version[1]) != VA_STATUS_SUCCESS)
	{
		fprintf(stderr, "vaapi initialization encountered and error.\n");
		return -1;
	}

	printf("got VAAPI version: %d.%d\n", version[0], version[1]);

	int formatCount = vaMaxNumImageFormats(vaDisplay);
	VAImageFormat *formats = calloc (formatCount, sizeof(VAImageFormat));
	if (vaQueryImageFormats(vaDisplay, formats, &formatCount) != VA_STATUS_SUCCESS)
	{
		fprintf(stderr, "query vaapi image formats failed.\n");
		return -1;
	}

	// list formats
	for (int i = 0; i < formatCount; i++)
	{
		print_fourcc(formats[i].fourcc);
	}

	// Allocate ffmpeg hardware device context
	avDeviceRef = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_VAAPI);
	if (!avDeviceRef)
	{
		fprintf(stderr, "failed to allocate hardware device context.\n");
		return -1;
	}

	// Pass display into context
	AVHWDeviceContext *deviceContext = (AVHWDeviceContext*) avDeviceRef->data;
	AVVAAPIDeviceContext *vaapiDeviceContext = deviceContext->hwctx;
	vaapiDeviceContext->display = vaDisplay;

	if (av_hwdevice_ctx_init(avDeviceRef) < 0)
	{
		fprintf(stderr, "failed to initialize hardware device context.\n");
		av_buffer_unref(&avDeviceRef);
		return -1;
	}

	return 0;
}

int prepareVideoDecoder(const char *url)
{
	if (avformat_open_input(&formatContext, url, NULL, NULL))
	{
		fprintf(stderr, "failed to open video file!\n");
		return -1;
	}

	if (avformat_find_stream_info(formatContext, NULL) < 0)
	{
		fprintf(stderr, "failed to find stream info!\n");
		return -1;
	}

	// Get the first available video stream
	AVCodecParameters *codecParameters = NULL;
	for (int i = 0; i < formatContext->nb_streams; i++)
	{
		AVStream *stream = formatContext->streams[i];
		if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			streamId = i;
			codecParameters = stream->codecpar;
			break;
		}
	}

	if (!codecParameters)
	{
		fprintf(stderr, "failed to find a video stream!\n");
		return -1;
	}

	// Find a decoder for the video stream
	AVCodec *codec = avcodec_find_decoder(codecParameters->codec_id);
	if (!codec)
	{
		fprintf(stderr, "unsupported codec!\n");
		return -1;
	}
     
	decodeContext = avcodec_alloc_context3(codec);
	if (avcodec_parameters_to_context(decodeContext, codecParameters) < 0)
	{
		fprintf(stderr, "failed to copy context!\n");
		return EXIT_FAILURE;
	}

	decodeContext->get_format = avcodec_get_format;

	// Open decoder
	if (avcodec_open2(decodeContext, codec, NULL) < 0)
	{
		fprintf(stderr, "failed to open decoding context!\n");
		return EXIT_FAILURE;
	}

	return 0;
}

int main (int argc, char **argv)
{
	av_register_all();
	av_log_set_level(AV_LOG_VERBOSE);

	createImage = (PFNEGLCREATEIMAGEKHRPROC) eglGetProcAddress("eglCreateImageKHR");
	destroyImage = (PFNEGLDESTROYIMAGEKHRPROC) eglGetProcAddress("eglDestroyImageKHR");
	imageTargetTexture2D = (PFNGLEGLIMAGETARGETTEXTURE2DOESPROC) eglGetProcAddress("glEGLImageTargetTexture2DOES");
	if (!createImage || !destroyImage || !imageTargetTexture2D)
	{
		fprintf(stderr, "failed to load EGL extensions.\n");
		return EXIT_FAILURE;
	}

	if (argc < 2)
	{
		fprintf(stderr, "need to specify video file!\n");
		return EXIT_FAILURE;
	}

	if (prepareEGL() < 0)
	{
		return EXIT_FAILURE;
	}

	if (prepareVAAPI() < 0)
	{
		return EXIT_FAILURE;
	}

	if (prepareVideoDecoder(argv[1]) < 0)
	{
		return EXIT_FAILURE;
	}

	if (prepareEGLWindow() < 0)
	{
		return EXIT_FAILURE;
	}

	// Compile shaders
	GLuint vertexShader = compileShader(vertex_shader_source, GL_VERTEX_SHADER);
	GLuint fragmentShader = compileShader(fragment_shader_source, GL_FRAGMENT_SHADER);
	GLuint program = linkProgram(vertexShader, fragmentShader);

	// Prepare buffers
	GLfloat quad[] = {
		-1.0, -1.0, 0.0, 1.0,
		-1.0,  1.0, 0.0, 0.0,
		1.0, -1.0, 1.0, 1.0,
		1.0,  1.0, 1.0, 0.0,
	};
	
	GLuint quadVertexBuffer;
	glGenBuffers(1, &quadVertexBuffer);
	glBindBuffer(GL_ARRAY_BUFFER, quadVertexBuffer);
	glBufferData(GL_ARRAY_BUFFER, sizeof(float) * 16, quad, GL_STATIC_DRAW);
	glBindBuffer(GL_ARRAY_BUFFER, 0);
	
	GLuint texture;
	glGenTextures(1, &texture);
	glBindTexture(GL_TEXTURE_EXTERNAL_OES, texture);
	glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_EXTERNAL_OES, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glBindTexture(GL_TEXTURE_EXTERNAL_OES, 0);

	//eglSwapInterval(eglDisplay, 0);

	// Read video
	AVFrame *frame = av_frame_alloc();
	AVPacket packet;
	EGLImageKHR eglImage = EGL_NO_IMAGE_KHR;
	VADRMPRIMESurfaceDescriptor descriptor;

	int packetId = 0;
	int frameId = 0;

	printf("playing video of resolution: %dx%d\n", decodeContext->coded_width, decodeContext->coded_height);

	while (av_read_frame(formatContext, &packet) >= 0)
	{
		packetId++;
		if (packet.stream_index != streamId)
			continue;

		int ret = avcodec_send_packet(decodeContext, &packet);
		switch (ret)
		{
			case AVERROR(EAGAIN):
				fprintf(stderr, "avcodec_send_packet: EAGAIN\n");
				return EXIT_FAILURE;
				break;
			case AVERROR(EINVAL):
				fprintf(stderr, "avcodec_send_packet: EINVAL\n");
				return EXIT_FAILURE;
				break;
			case AVERROR(ENOMEM):
				fprintf(stderr, "avcodec_send_packet: ENOMEM\n");
				return EXIT_FAILURE;
				break;
			default:
				break;
		}

		ret = avcodec_receive_frame(decodeContext, frame);
		if (ret == AVERROR(EAGAIN))
			continue;

		// got a frame
		frameId++;
		//printf("> got frame #%d (format = %d) (packet idx %d)\n", frameId, frame->format, packetId);

		if (frame->data[3] == NULL || frame->hw_frames_ctx == NULL)
		{
			fprintf(stderr, "\tbad vaapu frame/no surface/not vaapi frame!\n");
			continue;
		}

		VAStatus status;

		// unmap old image
		if (eglImage != EGL_NO_IMAGE)
		{
			if (destroyImage(eglDisplay, eglImage) == EGL_FALSE)
			{
				fprintf(stderr, "\tfailed to destroy egl image!\n");
			}
		}

		// Convert surface to a prime buffer
		VASurfaceID surfaceId = (VASurfaceID) frame->data[3];
		status = vaExportSurfaceHandle(vaDisplay, surfaceId, VA_SURFACE_ATTRIB_MEM_TYPE_DRM_PRIME_2,
			VA_EXPORT_SURFACE_SEPARATE_LAYERS | VA_EXPORT_SURFACE_READ_ONLY, &descriptor);
		if (status != VA_STATUS_SUCCESS)
		{
			fprintf(stderr, "\tfailed to export surface handle!\n");
		}

		print_fourcc(descriptor.fourcc);

		// create an EGLImage from the descriptors
		EGLint attributes[] = {
			EGL_WIDTH, descriptor.width,
			EGL_HEIGHT, descriptor.height,
			EGL_LINUX_DRM_FOURCC_EXT, descriptor.fourcc,
			EGL_DMA_BUF_PLANE0_FD_EXT, descriptor.objects[descriptor.layers[0].object_index[0]].fd,
			EGL_DMA_BUF_PLANE0_OFFSET_EXT, descriptor.layers[0].offset[0],
			EGL_DMA_BUF_PLANE0_PITCH_EXT, descriptor.layers[0].pitch[0],
			EGL_DMA_BUF_PLANE1_FD_EXT, descriptor.objects[descriptor.layers[1].object_index[0]].fd,
			EGL_DMA_BUF_PLANE1_OFFSET_EXT, descriptor.layers[1].offset[0],
			EGL_DMA_BUF_PLANE1_PITCH_EXT, descriptor.layers[1].pitch[0],
			EGL_NONE,
		};
	
		eglImage = createImage(eglDisplay, EGL_NO_CONTEXT, EGL_LINUX_DMA_BUF_EXT, NULL, attributes);
		if (eglImage == EGL_NO_IMAGE_KHR)
		{
			fprintf(stderr, "failed to create khr image from prime buffer.\n");
			continue;
		}

		close(descriptor.objects[0].fd);
		close(descriptor.objects[1].fd);


		glClearColor(0, 0, 0, 1);
		glViewport(0, 0, decodeContext->coded_width, decodeContext->coded_height);
		glClear(GL_COLOR_BUFFER_BIT);

		glActiveTexture(GL_TEXTURE0);
		glBindTexture(GL_TEXTURE_EXTERNAL_OES, texture);
		imageTargetTexture2D(GL_TEXTURE_EXTERNAL_OES, eglImage);

		// draw fullscreen quad
		glUseProgram(program);
		glUniform1i(glGetUniformLocation(program, "tex"), 0);
		glBindBuffer(GL_ARRAY_BUFFER, quadVertexBuffer);
		glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 4*sizeof(float), 0);
		glEnableVertexAttribArray(0);
		glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 4*sizeof(float), (const void *)(sizeof(float) * 2));
		glEnableVertexAttribArray(1);
		glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
		glBindBuffer(GL_ARRAY_BUFFER, 0);
		glBindTexture(GL_TEXTURE_EXTERNAL_OES, 0);

		eglSwapBuffers(eglDisplay, eglSurface);
	}

	return EXIT_SUCCESS;
}