factorial/
factorial.rs

1
2use asmkit::core::jit_allocator::{JitAllocator, JitAllocatorOptions};
3
4fn main() {
5    {
6        use asmkit::core::buffer::CodeBuffer;
7        use asmkit::x86::*;
8        use formatter::pretty_disassembler;
9        let mut buf = CodeBuffer::new();
10        let mut asm = Assembler::new(&mut buf);
11
12        let label = asm.get_label();
13        let fac = asm.get_label();
14
15        asm.bind_label(fac);
16        asm.mov64ri(RAX, imm(1));
17        asm.test64rr(RDI, RDI);
18        asm.jnz(label);
19        asm.ret();
20
21        {
22            asm.bind_label(label);
23            asm.pushr(RBX);
24            asm.mov64rr(RBX, RDI);
25            asm.lea64rm(RDI, ptr64(RDI, -1));
26            asm.call(fac);
27            asm.mov64rr(RDX, RAX);
28            asm.mov64rr(RAX, RBX);
29            asm.imul64rr(RAX, RDX);
30            asm.popr(RBX);
31            asm.ret();
32        }
33
34        let result = buf.finish();
35
36        let mut jit = JitAllocator::new(JitAllocatorOptions::default());
37
38        let mut span = jit
39            .alloc(result.data().len())
40            .expect("failed to allocate code");
41        unsafe {
42            jit.write(&mut span, |span| {
43                span.rw()
44                    .copy_from_nonoverlapping(result.data().as_ptr(), result.data().len());
45            })
46            .unwrap();
47            let mut out = String::new();
48            pretty_disassembler(&mut out, 64, result.data(), span.rx() as _).unwrap();
49            println!("{}", out);
50            #[cfg(target_arch = "x86_64")]
51            {
52                let f: extern "C" fn(u64) -> u64 = std::mem::transmute(span.rx());
53
54                println!("X86 factorial(5) = {:?}", f(5));
55            }
56        }
57    }
58
59    {
60        use asmkit::core::buffer::CodeBuffer;
61        use asmkit::riscv::*;
62        use formatter::pretty_disassembler;
63        let mut buf = CodeBuffer::new();
64        let mut asm = Assembler::new(&mut buf);
65
66        let label = asm.get_label();
67        let fac = asm.get_label();
68        asm.bind_label(fac);
69        asm.bnez(A0, label);
70        asm.addi(A0, ZERO, imm(1));
71        asm.ret();
72        {
73            asm.bind_label(label);
74            asm.addi(SP, SP, imm(-16));
75            asm.sd(SP, RA, imm(8));
76            asm.sd(SP, S0, imm(0));
77            asm.mv(S0, A0);
78            asm.addi(A0, A0, imm(-1));
79
80            asm.call(fac);
81            asm.mul(A0, S0, A0);
82            asm.ld(RA, SP, imm(8));
83            asm.ld(S0, SP, imm(0));
84            asm.addi(SP, SP, imm(16));
85            asm.ret();
86        }
87
88        let result = buf.finish();
89
90        let mut jit = JitAllocator::new(JitAllocatorOptions::default());
91
92        let mut span = jit
93            .alloc(result.data().len())
94            .expect("failed to allocate code");
95        unsafe {
96            jit.write(&mut span, |span| {
97                span.rw()
98                    .copy_from_nonoverlapping(result.data().as_ptr(), result.data().len());
99            })
100            .unwrap(); 
101            
102
103            let mut out = String::new();
104            pretty_disassembler(&mut out, 64, result.data(), span.rx() as _).unwrap();
105            println!("{}", out);
106            #[cfg(target_arch = "riscv64")]
107            {
108                let f: extern "C" fn(u64) -> u64 = std::mem::transmute(span.rx());
109
110                println!("RV64 factorial(5) = {:?}", f(5));
111            }
112        }
113    }
114}