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

static const reg_pool_t gp64 =
{
    15, 8,
    {
	rax, rbx, rcx, rdx, rdi, rsi, rbp,
	r8, r9, r10, r11, r12, r13, r14, r15
    }
};

static void
test_regalloc (assembler_t *as)
{
    typedef void (* func_type_t) (void);
    stack_man_t stack_man;
    reg_alloc_t reg_alloc;
    reg_context_t context, context2, context3;
    fragment_t *fragment, *frag2;
    uint8_t *code;
    int i;

    stack_manager_init (&stack_man);
    reg_alloc_init (&reg_alloc, &gp64, &stack_man, &context3, 0);

    fragment = fragment_new (as);

    /* Create a context where all registers are allocated. That
     * means any child of context2 will have to restore them.
     */
    reg_context_init (&context2, &context3,
		      15,
		      rax, rbx, rcx, rdx, rsi, rdi, rbp,
		      r8, r9, r10, r11, r12, r13, r14, r15);

    reg_context_init (&context, &context2, 0);

    for (i = 0; i < 15; ++i)
    {
	op_t reg = reg_context_alloc (&context, fragment);

	BEGIN_ASM (fragment)
	    I_mov, reg, IMM (7),
	END_ASM ();
    }

    reg_context_fini (&context, fragment, 0);
    reg_context_fini (&context2, fragment, 0);
    reg_context_fini (&context3, fragment, 0);

    frag2 = fragment_new (as);

    BEGIN_ASM (frag2)
	I_sub, rsp, IMM (stack_man.total_size),
	0,
    END_ASM ()

    BEGIN_ASM (fragment)
	I_add,  rsp, IMM (stack_man.total_size),
	I_ret,
	0,
    END_ASM ()

    code = assembler_link (as, frag2, fragment, NULL);

    ((func_type_t)code)();
}


static bool_t
avx_supported (assembler_t *as)
{
    uint8_t *code;
    typedef int (* func) (void);
    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 ();

    code = assembler_link (as, fragment, NULL);

    return ((func)code)();
}

static void
test_shifts (assembler_t *as)
{
    fragment_t *frag = fragment_new (as);
    uint8_t *c;
    uint32_t m;

    BEGIN_ASM (frag)
        I_add,          PTR (rax), al,
        I_movabs,       rbx, IMM64 (&m),
        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),
    END_ASM ();

    c = assembler_link (as, frag, NULL);
}

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

    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 ();

    c = assembler_link (as, frag, NULL);
}

int
main ()
{
    assembler_t *assembler = assembler_new ("pixman");
    fragment_t *frag = fragment_new (assembler);
    uint8_t sandbox[1024];
    uint8_t *code;

    assembler_set_verbose (assembler, TRUE);

    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 ();

    code = assembler_link (assembler, frag, NULL);

    frag = fragment_new (assembler);
    
    test_regalloc (assembler);

    test_shifts (assembler);
    test_float (assembler);
    
    printf ("AVX supported: %s\n",
	    avx_supported (assembler)? "yes" : "no");

    BEGIN_ASM (frag)
        I_lzcnt,        r12, r12,
        I_popcnt,       r12, r12,
        I_jmp,          LABEL ("skip xop"),
        I_add,          DWORD_PTR + BASE(eax,0x1234), IMM (17),
        I_add,          BASE(eax,0x1234), edx,
        I_vpcmov,       xmm0, xmm1, PTR (ebx), xmm3,
        I_vpcmov,       xmm0, xmm1, xmm2, PTR (ebx),
        I_vpcomltb,     xmm0, xmm1, PTR (ebx),
        I_vpcomfalseq,  xmm0, xmm13, PTR (r9),
        I_vpcomtrueub,  xmm1, xmm14, xmm2,
        I_bextr,        r12, r12, IMM (0xfedeabe),
        DEFINE_LABEL ("skip xop"),
    END_ASM ();
    
    BEGIN_ASM (frag)
        I_mfence,
        I_lfence,
	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 (&sandbox),
	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, rbx,
	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, esp, 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 (&printf),
	I_add,		eax,   eax,
	I_dq,		IMM64 (printf),
	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),
	
	DEFINE_VALUE64 ("blah",		0xff00ff00ff00ff00ULL),
	DEFINE_VALUE64 ("hello",	"Hello World\n"),
	DEFINE_VALUE64 ("hello2",	"Hello again, World\n"),
	DEFINE_VALUE64 ("printf",	printf),
	DEFINE_VALUE64 ("retval",	"return value: %d\n"),
    END_ASM ();

    code = assembler_link (assembler, frag, NULL);

    if (!code)
    {
	printf ("error\n");
    }
    else
    {
	typedef void (* func) (void);

	((func)(void *)code)();
    }

    return 0;
}