#include <stdalign.h>
#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <util.h>

#define	NBBY	8

struct xsave {
	union {
		char bytes[512];
	} 		xs_xsave;
	uint64_t	xs_xstate_bv;
	uint64_t	xs_xcomp_bv;
	uint8_t		xs_rsvd0[8];
	uint8_t		xs_rsvd1[40];
} __aligned(64);

__CTASSERT(sizeof(struct xsave) == 576);
__CTASSERT(alignof(struct xsave) == 64);

#define	XSAVE_X87	__BIT(0)
#define	XSAVE_SSE	__BIT(1)
#define	XSAVE_YMM_Hi128	__BIT(2)
#define	XSAVE_Opmask	__BIT(5)
#define	XSAVE_ZMM_Hi256	__BIT(6)
#define	XSAVE_Hi16_ZMM	__BIT(7)

#define	XSAVE_FMT	"\177\020"					      \
	"b\000"		"x87\0"						      \
	"b\001"		"SSE\0"						      \
	"b\002"		"YMM_Hi128\0"					      \
	"b\005"		"Opmask\0"					      \
	"b\006"		"ZMM_Hi256\0"					      \
	"b\007"		"Hi16_ZMM\0"					      \
	"b\077"		"COMP\0"					      \
	""

static void
xsave(struct xsave *xs, uint64_t bits)
{

	asm("xsave %[xs]"
	    : [xs]"=m"(*xs)
	    : "d"(__SHIFTOUT(bits, __BITS(63,32))),
	      "a"(__SHIFTOUT(bits, __BITS(31,0))));
}

static uint64_t
rdxcr(uint32_t n)
{
	uint64_t lo, hi;

	asm("xgetbv"
	    : "=a"(lo), "=d"(hi)
	    : "c"(n));
	return __SHIFTIN(hi, __BITS(63,32)) | __SHIFTIN(lo, __BITS(31,0));
}

static void
showbits(const char *title, uint64_t bits)
{
	struct xsave xs;
	char xcr0buf[1024], xstatebuf[1024], xcompbuf[1024];

	memset(&xs, 0, sizeof(xs));
	xsave(&xs, bits);
	snprintb(xcr0buf, sizeof(xcr0buf), XSAVE_FMT, rdxcr(0));
	snprintb(xstatebuf, sizeof(xstatebuf), XSAVE_FMT, xs.xs_xstate_bv);
	snprintb(xcompbuf, sizeof(xcompbuf), XSAVE_FMT, xs.xs_xcomp_bv);
	printf("%s: xcr0=%s xstate_bv=%s xcomp_bv=%s\n", title,
	    xcr0buf, xstatebuf, xcompbuf);
	fflush(stdout);
}

struct cpuid {
	uint32_t eax, ebx, ecx, edx;
};

static struct cpuid
cpuid(uint32_t eax, uint32_t ecx)
{
	struct cpuid c;

	asm("cpuid"
	    : "=a"(c.eax),
	      "=b"(c.ebx),
	      "=c"(c.ecx),
	      "=d"(c.edx)
	    : "a"(eax), "c"(ecx));
	return c;
}

int
main(void)
{
	uint8_t buf[32] __aligned(32) = {123};
	uint64_t bits =
	    XSAVE_X87 |
	    XSAVE_SSE |
	    XSAVE_YMM_Hi128 |
	    XSAVE_ZMM_Hi256 |
	    XSAVE_Hi16_ZMM;

	printf("CPUID[EAX=0x01].ECX=0x%x\n", cpuid(1, 0).ecx);
	fflush(stdout);
	showbits("0", bits);

	asm volatile("fldpi");
	showbits("1", bits);

	asm volatile("movdqa %[buf],%%xmm15" :: [buf]"m"(buf));
	showbits("2", bits);

	asm volatile("vmovdqa %[buf],%%ymm15" :: [buf]"m"(buf));
	showbits("3", bits);

	if (rdxcr(0) & XSAVE_ZMM_Hi256) {
		asm volatile("vmovdqa32 %[buf],%%zmm15" :: [buf]"m"(buf));
		showbits("4", bits);
	} else {
		printf("4: skip (no ZMM_Hi256)\n");
	}

	asm volatile("vzeroupper");
	showbits("5", bits);

	if (rdxcr(0) & XSAVE_Hi16_ZMM) {
		asm volatile("vmovdqa32 %[buf],%%zmm16" :: [buf]"m"(buf));
		showbits("6", bits);
	} else {
		printf("6: skip (no Hi16_ZMM)\n");
	}

	asm volatile("vzeroupper");
	showbits("7", bits);

	asm volatile("vzeroall");
	showbits("8", bits);

	fflush(stdout);
	return ferror(stdout);
}