#include <tests/doctest.h>

#include <iostream>

#include <sstream>

#include <libasr/codegen/x86_assembler.h>

using LCompilers::X86Reg;

// Print any vector like iterable to a string

template <class T>

inline std::ostream &print_vec(std::ostream &out, T &d)

{

out << "[";

for (auto p = d.begin(); p != d.end(); p++) {

if (p != d.begin())

out << ", ";

out << *p;

}

out << "]";

return out;

}

// Specialization for `const std::vector<uint8_t>`

template <>

inline std::ostream &print_vec(std::ostream &out, const std::vector<uint8_t> &d)

{

out << "[";

for (auto p = d.begin(); p != d.end(); p++) {

if (p != d.begin())

out << ", ";

out << LCompilers::i2s(*p);

}

out << "]";

return out;

}

namespace doctest {

// Convert std::vector<T> to string for doctest

template<typename T> struct StringMaker<std::vector<T>> {

static String convert(const std::vector<T> &value) {

std::ostringstream oss;

print_vec(oss, value);

return oss.str().c_str();

}

};

}

// Strips the first character (\n)

std::string S(const std::string &s) {

return s.substr(1, s.size());

}

TEST_CASE("Store and get instructions") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

a.asm_pop_r32(X86Reg::eax);

a.asm_jz_imm8(13);

LCompilers::Vec<uint8_t> &code = a.get_machine_code();

CHECK(code.size() == 3);

#ifdef LFORTRAN_ASM_PRINT

std::string asm_code = a.get_asm();

std::string ref = S(R"""(

BITS 32

org 0x08048000

pop eax

jz 0x0d

)""");

CHECK(asm_code == ref);

#endif

}

TEST_CASE("modrm_sib_disp") {

Allocator al(1024);

LCompilers::Vec<uint8_t> code;

LCompilers::X86Reg base;

LCompilers::X86Reg index;

std::vector<uint8_t> ref;

base = LCompilers::X86Reg::ecx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 0, false);

ref = {0xd9};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ecx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 0, true);

ref = {0x19};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::esi;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 0, true);

ref = {0x1e};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 0, true);

ref = {0x5d, 0x00};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 3, true);

ref = {0x5d, 0x03};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, -3, true);

ref = {0x5d, 0xfd};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 127, true);

ref = {0x5d, 0x7f};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 128, true);

ref = {0x9d, 0x80, 0x00, 0x00, 0x00};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, -128, true);

ref = {0x5d, 0x80};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, -129, true);

ref = {0x9d, 0x7f, 0xff, 0xff, 0xff};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::esp;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 0, true);

ref = {0x1c, 0x24};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::eax;

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 1, 0, true);

ref = {0x1c, 0x18};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 1, 0, true);

ref = {0x1c, 0x03};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 2, 0, true);

ref = {0x1c, 0x43};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 4, 0, true);

ref = {0x1c, 0x83};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 8, 0, true);

ref = {0x1c, 0xc3};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 8, 127, true);

ref = {0x5c, 0xc3, 0x7f};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 8, 128, true);

ref = {0x9c, 0xc3, 0x80, 0x00, 0x00, 0x00};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

index = LCompilers::X86Reg::eax;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, &index, 2, 1, true);

ref = {0x5c, 0x43, 0x01};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ebx, &base, nullptr, 1, 1, true);

ref = {0x5b, 0x01};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, &base, &index, 4, 1, true);

ref = {0x4c, 0x9d, 0x01};

CHECK( code.as_vector() == ref );

base = LCompilers::X86Reg::ebp;

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, &base, &index, 4, 128, true);

ref = {0x8c, 0x9d, 0x80, 0x00, 0x00, 0x00};

CHECK( code.as_vector() == ref );

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, nullptr, &index, 4, 1, true);

ref = {0x0c, 0x9d, 0x01, 0x00, 0x00, 0x00};

CHECK( code.as_vector() == ref );

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

CHECK_THROWS_AS(LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, nullptr, &index, 3, 1, true),

LCompilers::AssemblerError);

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

CHECK_THROWS_AS(LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, nullptr, nullptr, 1, 0, false),

LCompilers::AssemblerError);

index = LCompilers::X86Reg::ebx;

code.reserve(al, 32);

CHECK_THROWS_AS(LCompilers::modrm_sib_disp(code, al,

LCompilers::X86Reg::ecx, nullptr, nullptr, 1, 0, true),

LCompilers::AssemblerError);

}

TEST_CASE("Memory operand") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

X86Reg base = X86Reg::ebx;

X86Reg index = X86Reg::ecx;

a.asm_inc_m32(&base, nullptr, 1, 0);

a.asm_inc_m32(&base, nullptr, 1, 3);

a.asm_inc_m32(&base, nullptr, 1, -2);

a.asm_inc_m32(&base, &index, 1, 2);

a.asm_inc_m32(&base, &index, 2, 2);

a.asm_inc_m32(&base, &index, 4, 2);

a.asm_inc_m32(&base, &index, 8, 2);

a.asm_inc_m32(&base, &index, 1, -2);

a.asm_inc_m32(&base, &index, 1, 1024);

a.asm_inc_m32(&base, &index, 2, 0);

a.asm_inc_m32(nullptr, &index, 1, 2);

a.asm_inc_m32(nullptr, &index, 2, 2);

a.asm_inc_m32(nullptr, &index, 4, 2);

a.asm_inc_m32(nullptr, &index, 8, 2);

a.asm_inc_m32(nullptr, &index, 1, -2);

a.asm_inc_m32(nullptr, &index, 1, 1024);

a.asm_inc_m32(nullptr, &index, 2, 0);

#ifdef LFORTRAN_ASM_PRINT

std::string asm_code = a.get_asm();

std::string ref = S(R"""(

BITS 32

org 0x08048000

inc [ebx]

inc [ebx+3]

inc [ebx-2]

inc [ebx+ecx+2]

inc [ebx+2*ecx+2]

inc [ebx+4*ecx+2]

inc [ebx+8*ecx+2]

inc [ebx+ecx-2]

inc [ebx+ecx+1024]

inc [ebx+2*ecx]

inc [ecx+2]

inc [2*ecx+2]

inc [4*ecx+2]

inc [8*ecx+2]

inc [ecx-2]

inc [ecx+1024]

inc [2*ecx]

)""");

CHECK(asm_code == ref);

#endif

}

TEST_CASE("elf32 binary") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

LCompilers::emit_elf32_header(a);

std::string msg = "Hello World!\n";

a.add_label("msg");

a.asm_db_imm8(msg.c_str(), msg.size());

a.add_label("_start");

// ssize_t write(int fd, const void *buf, size_t count);

a.asm_mov_r32_imm32(LCompilers::X86Reg::eax, 4); // sys_write

a.asm_mov_r32_imm32(LCompilers::X86Reg::ebx, 1); // fd (stdout)

a.asm_mov_r32_imm32(LCompilers::X86Reg::ecx, a.get_defined_symbol("msg").value); // buf

a.asm_mov_r32_imm32(LCompilers::X86Reg::edx, msg.size()); // count

a.asm_int_imm8(0x80);

a.asm_call_label("exit");

a.add_label("exit");

// void exit(int status);

a.asm_mov_r32_imm32(LCompilers::X86Reg::eax, 1); // sys_exit

a.asm_mov_r32_imm32(LCompilers::X86Reg::ebx, 0); // exit code

a.asm_int_imm8(0x80); // syscall

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("write32");

// Note: both the machine code as well as the assembly below is checked

// for consistency using

//

// nasm -f bin write32.asm

//

// and the nasm generated `write32` binary exactly agrees with our binary.

//

// If any of the two tests below are modified, they have to be checked that

// they are still consistent with each other and that they work.

#ifdef LFORTRAN_ASM_PRINT

a.save_asm("write32.asm");

std::string asm_code = a.get_asm();

std::string ref = S(R"""(

BITS 32

org 0x08048000

ehdr:

db 0x7f

db 0x45

db 0x4c

db 0x46

db 0x01

db 0x01

db 0x01

db 0x00

db 0x00

db 0x00

db 0x00

db 0x00

db 0x00

db 0x00

db 0x00

db 0x00

dw 0x0002

dw 0x0003

dd 0x00000001

dd _start

dd e_phoff

dd 0x00000000

dd 0x00000000

dw ehdrsize

dw phdrsize

dw 0x0001

dw 0x0000

dw 0x0000

dw 0x0000

phdr:

dd 0x00000001

dd 0x00000000

dd 0x08048000

dd 0x08048000

dd filesize

dd filesize

dd 0x00000005

dd 0x00001000

phdr_end:

ehdrsize equ phdr - ehdr

phdrsize equ phdr_end - phdr

e_phoff equ phdr - ehdr

msg:

db 0x48

db 0x65

db 0x6c

db 0x6c

db 0x6f

db 0x20

db 0x57

db 0x6f

db 0x72

db 0x6c

db 0x64

db 0x21

db 0x0a

_start:

mov eax, 0x00000004

mov ebx, 0x00000001

mov ecx, 0x08048054

mov edx, 0x0000000d

int 0x80

call exit

exit:

mov eax, 0x00000001

mov ebx, 0x00000000

int 0x80

footer:

filesize equ footer - ehdr

)""");

CHECK(asm_code == ref);

#endif

std::vector<uint8_t> cref;

cref = {0x7f, 0x45, 0x4c, 0x46, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00,

0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x03, 0x00, 0x01, 0x00,

0x00, 0x00, 0x61, 0x80, 0x04, 0x08, 0x34, 0x00, 0x00, 0x00, 0x00,

0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x34, 0x00, 0x20, 0x00,

0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00,

0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x04, 0x08, 0x00, 0x80,

0x04, 0x08, 0x88, 0x00, 0x00, 0x00, 0x88, 0x00, 0x00, 0x00, 0x05,

0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x48, 0x65, 0x6c, 0x6c,

0x6f, 0x20, 0x57, 0x6f, 0x72, 0x6c, 0x64, 0x21, 0x0a, 0xb8, 0x04,

0x00, 0x00, 0x00, 0xbb, 0x01, 0x00, 0x00, 0x00, 0xb9, 0x54, 0x80,

0x04, 0x08, 0xba, 0x0d, 0x00, 0x00, 0x00, 0xcd, 0x80, 0xe8, 0x00,

0x00, 0x00, 0x00, 0xb8, 0x01, 0x00, 0x00, 0x00, 0xbb, 0x00, 0x00,

0x00, 0x00, 0xcd, 0x80};

CHECK( a.get_machine_code().as_vector() == cref );

}

TEST_CASE("print") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

std::string msg = "Hello World!\n";

LCompilers::emit_elf32_header(a);

a.add_label("_start");

LCompilers::emit_print(a, "msg", msg.size());

a.asm_call_label("exit");

LCompilers::emit_exit(a, "exit", 0);

LCompilers::emit_data_string(a, "msg", msg);

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("print32");

}

TEST_CASE("cmp") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

std::string msg1 = "Branch 1\n";

std::string msg2 = "Branch 2\n";

LCompilers::emit_elf32_header(a);

a.add_label("_start");

// if (3 >= 5) then

a.asm_mov_r32_imm32(LCompilers::X86Reg::eax, 3);

a.asm_cmp_r32_imm8(LCompilers::X86Reg::eax, 5);

a.asm_jge_label(".then");

a.asm_jmp_label(".else");

a.add_label(".then");

LCompilers::emit_print(a, "msg1", msg1.size());

a.asm_jmp_label(".endif");

a.add_label(".else");

LCompilers::emit_print(a, "msg2", msg2.size());

a.add_label(".endif");

a.asm_call_label("exit");

LCompilers::emit_exit(a, "exit", 0);

LCompilers::emit_data_string(a, "msg1", msg1);

LCompilers::emit_data_string(a, "msg2", msg2);

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("cmp32");

}

TEST_CASE("subroutines") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

std::string msg1 = "Subroutine 1, calling 2\n";

std::string msg1b = "Subroutine 1, done\n";

std::string msg2 = "Subroutine 2, modifying a message\n";

LCompilers::emit_elf32_header(a, 7);

a.add_label("sub1");

LCompilers::emit_print(a, "msg1", msg1.size());

a.asm_call_label("sub2");

LCompilers::emit_print(a, "msg1b", msg1b.size());

a.asm_ret();

a.add_label("sub2");

LCompilers::emit_print(a, "msg2", msg2.size());

a.asm_mov_r32_imm32(LCompilers::X86Reg::eax, 0x42414443);

a.asm_mov_r32_label(LCompilers::X86Reg::ebx, "msg1b");

LCompilers::X86Reg base = LCompilers::X86Reg::ebx;

a.asm_mov_m32_r32(&base, nullptr, 1, 5, LCompilers::X86Reg::eax);

a.asm_ret();

LCompilers::emit_exit(a, "exit", 0);

a.add_label("_start");

a.asm_call_label("sub1");

a.asm_call_label("exit");

LCompilers::emit_data_string(a, "msg1", msg1);

LCompilers::emit_data_string(a, "msg1b", msg1b);

LCompilers::emit_data_string(a, "msg2", msg2);

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("subroutines32");

}

TEST_CASE("subroutine args") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

std::string msg1 = "Subroutine 1\n";

std::string msg2 = "Sum equal to 9\n";

std::string msg3 = "Sum not equal to 9\n";

LCompilers::emit_elf32_header(a, 7);

a.add_label("sub1");

// Initialize stack

a.asm_push_r32(X86Reg::ebp);

a.asm_mov_r32_r32(X86Reg::ebp, X86Reg::esp);

a.asm_sub_r32_imm8(X86Reg::esp, 4); // one local variable

LCompilers::emit_print(a, "msg1", msg1.size());

X86Reg base = X86Reg::ebp;

// mov eax, [ebp+8] // first argument

a.asm_mov_r32_m32(X86Reg::eax, &base, nullptr, 1, 8);

// mov ecx, [ebp+12] // second argument

a.asm_mov_r32_m32(X86Reg::ecx, &base, nullptr, 1, 12);

// mov [ebp-4], eax // move eax to a local variable

a.asm_mov_m32_r32(&base, nullptr, 1, -4, X86Reg::eax);

// add [ebp-4], ecx // add ecx

a.asm_add_m32_r32(&base, nullptr, 1, -4, X86Reg::ecx);

// mov eax, [ebp-4] // move the sum into the return value (eax)

a.asm_mov_r32_m32(X86Reg::eax, &base, nullptr, 1, -4);

// Restore stack

a.asm_mov_r32_r32(X86Reg::esp, X86Reg::ebp);

a.asm_pop_r32(X86Reg::ebp);

a.asm_ret();

LCompilers::emit_exit(a, "exit", 0);

a.add_label("_start");

// Push arguments to stack (last argument first)

a.asm_push_imm8(5); // second argument

a.asm_push_imm8(4); // first argument

a.asm_call_label("sub1");

// Remove arguments from stack: 2 arguments, 4 bytes each

a.asm_add_r32_imm8(LCompilers::X86Reg::esp, 2*4);

a.asm_cmp_r32_imm8(LCompilers::X86Reg::eax, 9);

a.asm_je_label(".then");

a.asm_jmp_label(".else");

a.add_label(".then");

LCompilers::emit_print(a, "msg2", msg2.size());

a.asm_jmp_label(".endif");

a.add_label(".else");

LCompilers::emit_print(a, "msg3", msg3.size());

a.add_label(".endif");

a.asm_call_label("exit");

LCompilers::emit_data_string(a, "msg1", msg1);

LCompilers::emit_data_string(a, "msg2", msg2);

LCompilers::emit_data_string(a, "msg3", msg3);

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("subroutines_args32");

}

TEST_CASE("print integer") {

Allocator al(1024);

LCompilers::X86Assembler a(al, false);

LCompilers::emit_elf32_header(a);

emit_data_string(a, "string_neg", "-"); // - symbol for printing negative ints/floats

LCompilers::emit_print_int(a, "print_int");

LCompilers::emit_exit(a, "exit", 0);

a.add_label("_start");

a.asm_push_imm32(1234);

a.asm_call_label("print_int");

a.asm_add_r32_imm8(LCompilers::X86Reg::esp, 4);

a.asm_push_imm32(5678);

a.asm_call_label("print_int");

a.asm_add_r32_imm8(LCompilers::X86Reg::esp, 4);

a.asm_call_label("exit");

LCompilers::emit_elf32_footer(a);

a.verify();

a.save_binary("print_integer");

}