#include <stdio.h>
#include "simplex86.h"
#include "simple-reg.h"
#include "stack-man.h"
#include "crc32.h"

static uint8_t *
avx_supported (assembler_t *as)
{
    fragment_t *fragment = fragment_new (as);

    BEGIN_ASM (fragment)
	I_push,   rcx,
	I_push,   rdx,
        I_mov,    eax, IMM (1),
	I_cpuid,
	I_and,    ecx, IMM (0x18000000),
	I_cmp,    ecx, IMM (0x18000000),
	I_jne,    LABEL ("not_supported"),
	/* processor supports AVX instructions and XGETBV is enabled by OS */
	I_xor,    ecx, ecx,  /* specify 0 for XFEATURE_ENABLED_MASK register */
	I_xgetbv,	     /* result in EDX:EAX */
	I_and,    al, IMM (0x06),
	I_cmp,    al, IMM (0x06),
	I_jne,    LABEL ("not_supported"),
	I_mov,    eax, IMM (1),
	I_jmp,    LABEL ("done"),
	DEFINE_LABEL ("not_supported"),
	I_xor,	  rax, rax,
	DEFINE_LABEL ("done"),
	I_pop,    rdx,
	I_pop,    rcx,
	I_ret,
    END_ASM ();

    return assembler_link (as, fragment, NULL);
}

static uint8_t *
test_shifts (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);

    BEGIN_ASM (frag)
	I_vgatherdd,	ymm0, INDEX(r9, 23417, ymm0, 8), ymm2,
	I_vgatherdd,	ymm0, INDEX(eax, 23423, ymm9, 2), ymm12,
	I_vgatherdd,	ymm0, INDEX(rax, 23423, ymm9, 2), ymm12,
	I_vgatherdps,	ymm0, INDEX(r9, 23417, ymm0, 8), ymm2,
	I_vgatherdps,	ymm0, INDEX(eax, 23423, ymm9, 2), ymm12,
	I_vgatherqpd,	ymm0, INDEX(rax, 23423, ymm9, 2), ymm12,
        I_add,          PTR (rax), al,
        I_movabs,       rbx, IMM64 (0xaaaabbbbeeeeffff),
        I_stmxcsr,      PTR (rbx),
        I_vstmxcsr,     PTR (rbx),
        I_ldmxcsr,      PTR (rbx),
        I_vldmxcsr,     PTR (rbx),
        I_lfence,
        I_rdtsc,
        I_clflush,      PTR (r13),
        I_and,          esi, IMM (0x7fff00ff),
        I_shrd,         ax, cx, IMM (77),
        I_shrd,         INDEX (rax, 17, rax, 8), r9w, cl,
        I_shld,         INDEX (r9d, 17, r9d, 8), r9d, IMM (77),
        I_shld,         r13, rcx, IMM (77),
        I_shrd,         INDEX (rcx, 0, rax, 4), eax, cl,
        I_shrd,         INDEX (ecx, 17, ecx, 2), rcx, cl,

        I_add,          DWORD_PTR + INDEX (r14, 17, r14, 2), IMM (32),

        I_psllq,        mm0, mm7,
        I_psllq,        xmm13, IMM (7),
        I_vpsllq,       xmm13, xmm13, IMM(7),
        I_psrad,        mm7, mm0,
        I_vpsrad,       xmm13, xmm13, IMM (2),
        I_psrlw,        xmm7, PTR (ebx),
        I_psrlw,        xmm7, IMM (2),
        I_vpsrlq,       xmm13, xmm13, BASE (ebx, 17),
        I_psrlq,        mm0, IMM (64),

        /* AVX2 */
        I_vpsllq,       ymm1, ymm9, PTR (ebx),
        I_vpsllq,       ymm13, ymm12, IMM(7),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_movdq (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    uint8_t *result;

    BEGIN_ASM (frag)
	I_movd, mm7, DWORD_PTR + PTR (r9),
	I_movd, DWORD_PTR + INDEX (r12, 0, rax, 8), mm7,
	I_movd, xmm12, DWORD_PTR + PTR (r10),
	I_movd, DWORD_PTR + INDEX (r11, 0, rax, 8), xmm12,

	I_vmovd, xmm12, DWORD_PTR + PTR (r10),
	I_vmovd, DWORD_PTR + INDEX (r11, 0, rax, 8), xmm12,

	/* These variants should use the RM64 variant
	 * due to r9 being a 64 bit register
	 */
	I_movq, mm7, PTR (r9),
	I_movq, PTR(r9), mm7,
	I_movq, xmm7, PTR (r10),
	I_movq, PTR (r10), xmm7,

	/* These will use the SSEM variant due to eax being
	 * a 32 bit register
	 */
	I_movq, PTR (eax), xmm7,
	I_movq, xmm7, PTR (eax),
	I_movq, PTR (eax), mm7,
	I_movq, mm7, PTR (eax),

	I_vmovq, PTR (eax), xmm7,
	I_vmovq, xmm7, PTR (eax),
    END_ASM ();

    result = assembler_link (as, frag, NULL);

    return result;
}

static uint8_t *
test_float (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);

    BEGIN_ASM (frag)
        I_addps,        xmm0, xmm7,
        I_addpd,        xmm0, PTR (rbx),
        I_vaddps,       xmm0, xmm1, xmm2,
        I_vaddpd,       xmm0, xmm0, PTR (rbx),
        I_vaddpd,       ymm0, ymm0, PTR (rbx),
        I_vaddss,       xmm0, xmm7, xmm3,
        I_vaddsd,       xmm9, xmm3, PTR (rbx),
        I_vblendps,     ymm9, ymm0, PTR (r9), IMM (7),
        I_vrcpps,       xmm0, xmm0,
        I_vrcpps,       ymm15, PTR (rax),
        I_vroundps,     ymm0, ymm0, PTR (eax), IMM (2),
        I_vshufpd,      ymm1, ymm13, ymm13, IMM (12),
        I_cmpps,        xmm2, PTR (eax), IMM (-1),
    END_ASM ();

    BEGIN_ASM (frag)
        I_vmaskmovps,   xmm0, xmm7, PTR (ebx),
        I_vmaskmovps,   ymm0, ymm7, PTR (ebx),
        I_vmaskmovps,   PTR (ebx), xmm0, xmm9,
        I_vmaskmovps,   PTR (ebx), ymm0, ymm12,
        I_vmaskmovpd,   xmm0, xmm7, PTR (ebx),
        I_vmaskmovpd,   ymm0, ymm7, PTR (ebx),
        I_vmaskmovpd,   PTR (ebx), xmm0, xmm9,
        I_vmaskmovpd,   PTR (ebx), ymm0, ymm12,
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_convert (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);

    BEGIN_ASM (frag)
        I_cbw,
        I_cwde,
        I_cdqe,
        I_bswap, eax,
        I_bswap, r13,
        I_movbe, ax, PTR (eax),
        I_movbe, eax, PTR (eax),
        I_movbe, rax, PTR (eax),
        I_movbe, PTR (eax), ax,
        I_movbe, PTR (eax), eax,
        I_movbe, PTR (rax), rax,
        I_bsf, ax, PTR (eax),
        I_bsr, eax, eax,
        I_bsf, r9d, BASE (r9d, 17),
        I_bsr, rcx, rcx,
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_crc32 (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    
    BEGIN_ASM (frag)
        I_crc32,        eax, BYTE_PTR + PTR (ebx),
        I_crc32,        eax, WORD_PTR + PTR (eax),
        I_crc32,        eax, eax,
        I_crc32,        r9, QWORD_PTR + PTR (r9),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_pshuflw (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    
    BEGIN_ASM (frag)
	I_pshuflw,	xmm2, xmm2, UIMM (0xff),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_reg_regm (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);

    BEGIN_ASM (frag)
	I_mov, rax, PTR (rsp),
	I_mov, rax, PTR (r12),
	I_mov, rax, PTR (rbp),
	I_mov, rax, PTR (r13),
	I_mov, rax, BASE (rsp, 10),
	I_mov, rax, BASE (r12, 10),
	I_mov, rax, BASE (rbp, 10),
	I_mov, rax, BASE (r13, 10),
	I_mov, rax, INDEX (rsp, 10, r12, 4),
	I_mov, rax, INDEX (r12, 10, r12, 4),
	I_mov, rax, INDEX (rbp, 10, r12, 4),
	I_mov, rax, INDEX (r13, 10, r12, 4),
	I_mov, rax, INDEX (rsp, 10, r12, 1),
	I_mov, rax, INDEX (r12, 10, r12, 1),
	I_mov, rax, INDEX (rbp, 10, r12, 1),
	I_mov, rax, INDEX (r13, 10, r12, 1),
	I_mov, rax, INDEX (rsp, 0, r12, 1),
	I_mov, rax, INDEX (r12, 0, r12, 1),
	I_mov, rax, INDEX (rbp, 0, rbp, 1),
	I_mov, rax, INDEX (r13, 0, r12, 1),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_movd (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    
    BEGIN_ASM (frag)
	I_movd, DWORD_PTR + PTR (rdx), xmm0,
	I_movd, xmm5, DWORD_PTR + PTR (r12),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

static uint8_t *
test_riprel (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    
    BEGIN_ASM (frag)
        I_and,  al, IMM (0),
        I_xor,  al, al,
        I_neg, WORD_PTR + RIP_REL ("llll"),
        I_neg, WORD_PTR + PTR (eax),
        I_movzx, eax, BYTE_PTR + INDEX (eax, 17, ebx, 2),
        I_movsx, rax, WORD_PTR + BASE (esp, 3),
        DEFINE_LABEL ("llll"),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}
static uint8_t *
test_misc (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    
    BEGIN_ASM (frag)
        I_mfence,
        I_lfence,
        I_add,          DWORD_PTR + BASE(eax,0x1234), IMM (17),
        I_add,          DWORD_PTR + BASE(eax,0x1234), edx,
	I_jmp,		LABEL("far_away"),
	DEFINE_LABEL ("begin"),
	I_lock, I_xchg, rax, PTR (rsp),
	I_lock, I_xchg, PTR (rsp), rax,
        I_lea,          eax, PTR (rip),
        I_lea,          ax, PTR (rip),
        I_lea,          ax, PTR (esp),
        I_lea,          ax, PTR (rsp),
        I_lea,          rbx, INDEX (rax, 17, rax, 8),
        I_add,          edx, ebx,
        I_add,          ebx, edx,
        I_add,          rax, IMM (17),
        I_add,          eax, IMM (17),
        I_mul,          rbx,
        I_mul,          al,
        I_mul,          ebp,
        I_mul,          r9l,

        I_imul2,        r9, r9,
        
        I_imul3,        ax, r9w, IMM (17),
        I_imul3,        r9d, RIP_REL ("printf"), IMM (0x7befe57),
        
        I_movabs,       rbx, IMM64 (0xaaaaeeeedddd0000),
	I_movq,		xmm0, BASE (rbx, 127),

        I_mov,          al, IMM (0),
	I_mov,		r13, RIP_REL ("printf"),
	I_mov,		rdi, RIP_REL ("hello"),
	I_call,		r13,
        I_sub,          ax, IMM (234),
	I_jmp,		LABEL ("aligned"),
	I_align,	IMM(16),
	DEFINE_LABEL ("aligned"),
        I_movzx,        esi, ax,
	I_align,	IMM (11),
        I_movsx,        rax, ax,
	I_align,	IMM (5),
        I_movsx,        rax, DWORD_PTR + INDEX (rbx, 1234, rax, 8),
	I_align,	IMM (5),
        I_movsx,        rax, al,
	I_align,	IMM (7),
        I_movzx,        r9d, cl,
	I_align,	IMM (5),
        I_movzx,        r9, cl,
	I_align,	IMM (5),

        I_jmp,          LABEL ("skipavx"),
        I_vzeroupper,
        I_vzeroall,
        I_vcvtph2ps,    ymm15, xmm12,
        I_vcvtph2ps,    ymm2, RIP_REL ("hello2"),
        I_vcvtph2ps,    ymm15, xmm12,

        I_vcvtps2ph,    RIP_REL ("hello2"), xmm12, IMM (7),
        I_vcvtps2ph,    xmm0, xmm12, IMM (7),
        I_vcvtps2ph,    RIP_REL ("hello2"), ymm12, IMM (0x77),
        I_vcvtps2ph,    xmm9, ymm12, IMM (0x77),

        I_pshufhw,      xmm15, xmm12, IMM (13),
	I_align,	IMM (4),
        I_vpshufhw,     xmm0, PTR (rsp), IMM (3),
        I_vpmovzxbq,    xmm0, INDEX (r14, 0x8000, rax, 8),
        I_pmovzxbq,     xmm0, INDEX (r14, 0x8000, rax, 8),
        DEFINE_LABEL ("skipavx"),
        
        /* The ABI spec claims that %al is used as an "upper bound" of
         * the number of vector registers used in a vararg function.
         * However, if we don't set %al to 0 here, printf crashes.
         */
        I_mov,          al, IMM (0),
        I_mov,          r9l, IMM (0),
        I_mov,          ah, IMM (8),
	I_xor,		bh, r9l,
	I_xor,		r9l, al,
	I_mov,		r13, RIP_REL ("printf"),
	I_mov,		rdi, RIP_REL ("hello2"),
        I_call,         r13,

        I_mov,          eax, IMM(32),
        I_movq,         mm7, mm2,
        I_movq,         mm3, RIP_REL ("printf"),
        I_movq,         INDEX (rbx, 17, rax, 4), mm2,
        I_movq,         xmm9, RIP_REL ("printf"),
        I_movq,         INDEX (rbx, 17, rax, 4), xmm2,
        I_movq,         xmm9, xmm7,
        I_movq2dq,      xmm9, mm7,
        
	I_mov,		rsi, rax,
	I_mov,		rdi, RIP_REL ("retval"),
	I_xor,		rax, rax,
	I_call,		r13,

        I_xor,          al, al,
	I_mov,		rdi, RIP_REL ("hello2"),
        I_call,         RIP_REL("printf"),

	I_add,		rax, rbx,
	I_jnz,		LABEL ("end"),
	I_add,		ebx,	   BASE (ebx, 127),
	I_or,		ebx,	   IMM (1234),
	I_xor,		ebx,	   BASE (ebx, 127),
	I_cpuid,
	I_push,		esi,
	I_mov,		WORD_PTR + PTR (ebx), IMM (7),
	I_mov,		PTR (ebx), eax,
	I_mov,		eax,       PTR (ebx),
	I_mov,		rax,       PTR (ebx),
	I_mov,		r9,        PTR (r9),
	I_push,		PTR (r9),
	I_mov,		esp,       ebp,
	I_mov,		ebp,       rsp,
	I_jnz,		LABEL ("begin"),
	I_call,		r10,
	I_jmp,		r10,
	I_jmp,		LABEL ("begin"),
	I_nop,

	DEFINE_LABEL ("end"),

	I_sal,		r10, IMM (17),
	I_sal,		r10d, IMM (17),
	I_rcl,		r10,   cl,
	I_sal,		eax,   IMM (1),
	I_add,		rbx,   IMM (3),
	I_pabsb,	xmm9,  BASE (rbx, (-19 + 16)),
	I_pabsb,	mm7,   BASE (rbx, (-19 + 16)),
	I_vpabsb,	xmm9,  BASE (rbx, (-19 + 16)),
	I_xor,		rax, rax,
	I_vpacksswb,	xmm12, xmm10, INDEX (rbx, 17, r10, 8),
	I_vpacksswb,	xmm12, xmm10, INDEX (rbx, 17, rax, 8),
	I_vpackusdw,	xmm12, xmm10, INDEX (rbx, 17, rax, 8),
	I_vpackuswb,	xmm12, xmm10, INDEX (rbx, 17, rax, 8),
	I_vpaddw,	xmm12, xmm10, xmm7,
	I_mov,		r10, rdx,
	I_palignr,	mm7,   PTR(r10), IMM(17),
    END_ASM ();
    
    BEGIN_ASM (frag)
	I_vpalignr,	xmm7,  xmm8, PTR (r10), IMM (17),
    END_ASM ();

    BEGIN_ASM (frag)
	I_vpandn,	xmm7,  xmm8, xmm9,
	I_vpsubb,	xmm7,  xmm8, RIP_REL ("blah"),
	I_vpaddb,	xmm3,  xmm3, xmm7,
	I_mov,		eax,   ADDRESS32 (12345),
	I_mov,		rax,   ADDRESS32 (12345),
	I_mov,		rax,   BASE (rsp, 17),
	I_mov,		rax,   PTR (r13),
	I_mov,		rax,   INDEX (r13, 0, rbp, 1),
	I_mov,		eax,   ADDRESS32 (0xabcdef),
    END_ASM ();

    BEGIN_ASM (frag)
	I_dq, IMM64 (234),
    END_ASM ();

    BEGIN_ASM (frag)
	I_mov,		rax,   BASE (rip, 0x12345),
	I_movabs,	r10,   IMM64 (0xfedeabefedeabeULL),
	I_movabs,	rax,   IMM64 (0xfedeabefedeabeULL),
	I_movabs,	rax,   IMM64 (0x324897982374e),
	I_add,		eax,   eax,
	I_dq,		IMM64 (0x287364876),
	I_dw,		IMM (12345),
	DEFINE_LABEL ("far_away"),
	I_jmp,		LABEL ("begin"),

	I_align,	IMM(4),
	I_dd,		IMM (0x90909090),
	I_db,		IMM ((int8_t)0xc3),
        I_sub,          al, IMM (7),
        I_sbb,          ax, IMM (1),
        I_sub,          eax, IMM (0),
        I_xor,          rax, IMM (234),
	
	DEFINE_VALUE64 ("blah",		0xff00ff00ff00ff00ULL),
	DEFINE_VALUE64 ("hello",	0xabeabeabeabeabeaULL),
	DEFINE_VALUE64 ("hello2",	0x1870987324987987ULL),
	DEFINE_VALUE64 ("printf",	0x1239874691823764ULL),
	DEFINE_VALUE64 ("retval",	0x8726348726348762ULL),
    END_ASM ();

    return assembler_link (as, frag, NULL);
}

typedef uint8_t *(* generator_t) (assembler_t *as);

static bool_t
run_test (assembler_t *as, 
	  const char *name, generator_t generator, uint32_t crc32)
{
    uint8_t *code = generator (as);
    size_t size = assembler_get_last_size (as);
    uint32_t c;

    if ((c = compute_crc32 (0xafafafaf, code, size)) != crc32)
    {
	printf ("Test %s failed. Got crc32: %x; should be %x\n",
		name, c, crc32);
	return FALSE;
    }

    return TRUE;
}

int
main ()
{
    assembler_t *as;
    bool_t success = TRUE;

    putenv ("PIXMAN_JIT_SANITY_CHECK=1");
    
    as = assembler_new ("pixman");

#if 0
    assembler_set_verbose (as, 1);
#endif

/* FIXME: We may produce different binaries on on machines that don't
 * support multibyte nop. We need a way to force the assembler to 
 * pretend that the machine has or hasn't some specific features.
 *
 * Also we will produce different output on 32 vs 64 bit, at least
 * due to differences in default address size.
 */
    success &= run_test (as, "test_crc32", test_crc32, 0x3ce8ed36 );
    success &= run_test (as, "avx_supported", avx_supported, 0xeeaec30c);
    success &= run_test (as, "test_shifts", test_shifts, 0x4a641094);
    success &= run_test (as, "test_float", test_float, 0xc193030 );
    success &= run_test (as, "test_convert", test_convert, 0x9d52c83b);
    success &= run_test (as, "test_misc", test_misc, 0x1ea2e7a);
    success &= run_test (as, "test_riprel", test_riprel, 0xf5497642);
    success &= run_test (as, "test_movdq", test_movdq, 0xedd4f79d);
    success &= run_test (as, "test_pshuflw", test_pshuflw, 0x7a8c7f59);
    success &= run_test (as, "test_movd", test_movd, 0xcb69ef60);
    success &= run_test (as, "test_reg_regm", test_reg_regm, 0x61dd330d);

    if (success)
	printf ("Test suite PASSED\n");
    else
	printf ("Test suite FAILED\n");

    return !success;
}