// SPDX-License-Identifier: GPL-2.0-only
#include <sys/prctl.h>
#include <unistd.h>
#include <errno.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <stdlib.h>

#include "kselftest_harness.h"
#include "v_helpers.h"

#define NEXT_PROGRAM "./vstate_exec_nolibc"

int test_and_compare_child(long provided, long expected, int inherit, int xtheadvector)
{
	int rc;

	rc = prctl(PR_RISCV_V_SET_CONTROL, provided);
	if (rc != 0) {
		printf("prctl with provided arg %lx failed with code %d\n",
		       provided, rc);
		return -1;
	}
	rc = launch_test(NEXT_PROGRAM, inherit, xtheadvector);
	if (rc != expected) {
		printf("Test failed, check %d != %ld\n", rc, expected);
		return -2;
	}
	return 0;
}

#define PR_RISCV_V_VSTATE_CTRL_CUR_SHIFT 0
#define PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT 2

TEST(get_control_no_v)
{
	long rc;

	if (is_vector_supported() || is_xtheadvector_supported())
		SKIP(return, "Test expects vector to be not supported");

	rc = prctl(PR_RISCV_V_GET_CONTROL);
	EXPECT_EQ(-1, rc)
	TH_LOG("GET_CONTROL should fail on kernel/hw without ZVE32X");
	EXPECT_EQ(EINVAL, errno)
	TH_LOG("GET_CONTROL should fail on kernel/hw without ZVE32X");
}

TEST(set_control_no_v)
{
	long rc;

	if (is_vector_supported() || is_xtheadvector_supported())
		SKIP(return, "Test expects vector to be not supported");

	rc = prctl(PR_RISCV_V_SET_CONTROL, PR_RISCV_V_VSTATE_CTRL_ON);
	EXPECT_EQ(-1, rc)
	TH_LOG("SET_CONTROL should fail on kernel/hw without ZVE32X");
	EXPECT_EQ(EINVAL, errno)
	TH_LOG("SET_CONTROL should fail on kernel/hw without ZVE32X");
}

TEST(vstate_on_current)
{
	long flag;
	long rc;

	if (!is_vector_supported() && !is_xtheadvector_supported())
		SKIP(return, "Vector not supported");

	flag = PR_RISCV_V_VSTATE_CTRL_ON;
	rc = prctl(PR_RISCV_V_SET_CONTROL, flag);
	EXPECT_EQ(0, rc) TH_LOG("Enabling V for current should always succeed");
}

TEST(vstate_off_eperm)
{
	long flag;
	long rc;

	if (!is_vector_supported() && !is_xtheadvector_supported())
		SKIP(return, "Vector not supported");

	flag = PR_RISCV_V_VSTATE_CTRL_OFF;
	rc = prctl(PR_RISCV_V_SET_CONTROL, flag);
	EXPECT_EQ(EPERM, errno)
	TH_LOG("Disabling V in current thread with V enabled must fail with EPERM(%d)", errno);
	EXPECT_EQ(-1, rc)
	TH_LOG("Disabling V in current thread with V enabled must fail with EPERM(%d)", errno);
}

TEST(vstate_on_no_nesting)
{
	long flag;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}

	/* Turn on next's vector explicitly and test */
	flag = PR_RISCV_V_VSTATE_CTRL_ON << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;

	EXPECT_EQ(0, test_and_compare_child(flag, PR_RISCV_V_VSTATE_CTRL_ON, 0, xtheadvector));
}

TEST(vstate_off_nesting)
{
	long flag;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}

	/* Turn off next's vector explicitly and test */
	flag = PR_RISCV_V_VSTATE_CTRL_OFF << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;

	EXPECT_EQ(0, test_and_compare_child(flag, PR_RISCV_V_VSTATE_CTRL_OFF, 1, xtheadvector));
}

TEST(vstate_on_inherit_no_nesting)
{
	long flag, expected;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}

	/* Turn on next's vector explicitly and test no inherit */
	flag = PR_RISCV_V_VSTATE_CTRL_ON << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;
	flag |= PR_RISCV_V_VSTATE_CTRL_INHERIT;
	expected = flag | PR_RISCV_V_VSTATE_CTRL_ON;

	EXPECT_EQ(0, test_and_compare_child(flag, expected, 0, xtheadvector));
}

TEST(vstate_on_inherit)
{
	long flag, expected;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}

	/* Turn on next's vector explicitly and test inherit */
	flag = PR_RISCV_V_VSTATE_CTRL_ON << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;
	flag |= PR_RISCV_V_VSTATE_CTRL_INHERIT;
	expected = flag | PR_RISCV_V_VSTATE_CTRL_ON;

	EXPECT_EQ(0, test_and_compare_child(flag, expected, 1, xtheadvector));
}

TEST(vstate_off_inherit_no_nesting)
{
	long flag, expected;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}
	/* Turn off next's vector explicitly and test no inherit */
	flag = PR_RISCV_V_VSTATE_CTRL_OFF << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;
	flag |= PR_RISCV_V_VSTATE_CTRL_INHERIT;
	expected = flag | PR_RISCV_V_VSTATE_CTRL_OFF;

	EXPECT_EQ(0, test_and_compare_child(flag, expected, 0, xtheadvector));
}

TEST(vstate_off_inherit)
{
	long flag, expected;
	int xtheadvector = 0;

	if (!is_vector_supported()) {
		if (is_xtheadvector_supported())
			xtheadvector = 1;
		else
			SKIP(return, "Vector not supported");
	}

	/* Turn off next's vector explicitly and test inherit */
	flag = PR_RISCV_V_VSTATE_CTRL_OFF << PR_RISCV_V_VSTATE_CTRL_NEXT_SHIFT;
	flag |= PR_RISCV_V_VSTATE_CTRL_INHERIT;
	expected = flag | PR_RISCV_V_VSTATE_CTRL_OFF;

	EXPECT_EQ(0, test_and_compare_child(flag, expected, 1, xtheadvector));
}

/* arguments should fail with EINVAL */
TEST(inval_set_control_1)
{
	int rc;

	if (!is_vector_supported() && !is_xtheadvector_supported())
		SKIP(return, "Vector not supported");

	rc = prctl(PR_RISCV_V_SET_CONTROL, 0xff0);
	EXPECT_EQ(-1, rc);
	EXPECT_EQ(EINVAL, errno);
}

/* arguments should fail with EINVAL */
TEST(inval_set_control_2)
{
	int rc;

	if (!is_vector_supported() && !is_xtheadvector_supported())
		SKIP(return, "Vector not supported");

	rc = prctl(PR_RISCV_V_SET_CONTROL, 0x3);
	EXPECT_EQ(-1, rc);
	EXPECT_EQ(EINVAL, errno);
}

/* arguments should fail with EINVAL */
TEST(inval_set_control_3)
{
	int rc;

	if (!is_vector_supported() && !is_xtheadvector_supported())
		SKIP(return, "Vector not supported");

	rc = prctl(PR_RISCV_V_SET_CONTROL, 0xc);
	EXPECT_EQ(-1, rc);
	EXPECT_EQ(EINVAL, errno);
}

TEST_HARNESS_MAIN