JitAllocator

asmkit::core::jit_allocator

Struct JitAllocator 

Source 
pub struct JitAllocator { /* private fields */ }
Expand description

A simple implementation of memory manager that uses virtual_memory. functions to manage virtual memory for JIT compiled code.

Implementation notes:

  • Granularity of allocated blocks is different than granularity for a typical C malloc. In addition, the allocator can use several memory pools having a different granularity to minimize the maintenance overhead. Multiple pools feature requires use_multiple_pools flag to be set.

  • The allocator doesn’t store any information in executable memory, instead, the implementation uses two bit-vectors to manage allocated memory of each allocator-block. The first bit-vector called ‘used’ is used to track used memory (where each bit represents memory size defined by granularity) and the second bit vector called ‘stop’ is used as a sentinel to mark where the allocated area ends.

  • Internally, the allocator also uses RB tree to keep track of all blocks across all pools. Each inserted block is added to the tree so it can be matched fast during release() and shrink().

Implementations§

Source§

impl JitAllocator

Source

pub fn new(params: JitAllocatorOptions) -> Box<Self>

Creates a new JitAllocator instance.

Examples found in repository?
examples/reloc.rs (line 29)
10fn main() {
11    let mut buf = CodeBuffer::new();
12    let mut asm = Assembler::new(&mut buf);
13
14    let str_constant = asm.add_constant("Hello, World!\0");
15    let puts_sym = asm
16        .buffer
17        .add_symbol(ExternalName::Symbol("puts".into()), RelocDistance::Far);
18
19    asm.lea64rm(RDI, ptr64_label(str_constant, 0));
20    asm.callm(ptr64_sym(puts_sym, 0));
21    asm.ret();
22
23    let result = buf.finish();
24
25    for reloc in result.relocs() {
26        println!("{:?}", reloc);
27    }
28
29    let mut jit = JitAllocator::new(Default::default());
30
31    // allocate memory for GOT table and for code itself
32    let mut span = jit
33        .alloc(result.data().len() + result.relocs().len() * 8)
34        .unwrap();
35
36    let mut got_addr_rx = std::ptr::null();
37
38    unsafe {
39        jit.write(&mut span, |span| {
40            span.rw()
41                .copy_from_nonoverlapping(result.data().as_ptr(), result.data().len());
42            got_addr_rx = span.rx().add(result.data().len());
43            span.rw()
44                .add(result.data().len())
45                .cast::<usize>()
46                .write(puts as *const u8 as usize);
47            // we only link to one symbol in GOT table, don't bother with anything else...
48            perform_relocations(
49                span.rw(),
50                span.rx(),
51                &result.relocs(),
52                |_| unreachable!(),
53                |_| got_addr_rx,
54                |_| unreachable!(),
55            );
56        })
57        .unwrap();
58
59        let mut out = String::new();
60        pretty_disassembler(
61            &mut out,
62            64,
63            std::slice::from_raw_parts(span.rx(), result.data().len()),
64            span.rx() as _,
65        )
66        .unwrap();
67
68        println!("{}", out);
69        #[cfg(target_arch = "x86_64")]
70        {
71            let f: extern "C" fn() = std::mem::transmute(span.rx());
72
73            f();
74        }
75    }
76}
More examples
Hide additional examples
examples/factorial.rs (line 36)
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}
Source

pub unsafe fn reset(&mut self, reset_policy: ResetPolicy)

Resets current allocator by emptying all pools and blocks.

Frees all memory is ResetPolicy::Hard is specified or immediate_release in JitAllocatorOptions is specific.

Source

pub fn alloc(&mut self, size: usize) -> Result<Span, AsmError>

Allocates size bytes in the executable memory region. Returns two pointers. One points to Read-Execute mapping and another to Read-Write mapping. All code writes must go to the Read-Write mapping.

Examples found in repository?
examples/reloc.rs (line 33)
10fn main() {
11    let mut buf = CodeBuffer::new();
12    let mut asm = Assembler::new(&mut buf);
13
14    let str_constant = asm.add_constant("Hello, World!\0");
15    let puts_sym = asm
16        .buffer
17        .add_symbol(ExternalName::Symbol("puts".into()), RelocDistance::Far);
18
19    asm.lea64rm(RDI, ptr64_label(str_constant, 0));
20    asm.callm(ptr64_sym(puts_sym, 0));
21    asm.ret();
22
23    let result = buf.finish();
24
25    for reloc in result.relocs() {
26        println!("{:?}", reloc);
27    }
28
29    let mut jit = JitAllocator::new(Default::default());
30
31    // allocate memory for GOT table and for code itself
32    let mut span = jit
33        .alloc(result.data().len() + result.relocs().len() * 8)
34        .unwrap();
35
36    let mut got_addr_rx = std::ptr::null();
37
38    unsafe {
39        jit.write(&mut span, |span| {
40            span.rw()
41                .copy_from_nonoverlapping(result.data().as_ptr(), result.data().len());
42            got_addr_rx = span.rx().add(result.data().len());
43            span.rw()
44                .add(result.data().len())
45                .cast::<usize>()
46                .write(puts as *const u8 as usize);
47            // we only link to one symbol in GOT table, don't bother with anything else...
48            perform_relocations(
49                span.rw(),
50                span.rx(),
51                &result.relocs(),
52                |_| unreachable!(),
53                |_| got_addr_rx,
54                |_| unreachable!(),
55            );
56        })
57        .unwrap();
58
59        let mut out = String::new();
60        pretty_disassembler(
61            &mut out,
62            64,
63            std::slice::from_raw_parts(span.rx(), result.data().len()),
64            span.rx() as _,
65        )
66        .unwrap();
67
68        println!("{}", out);
69        #[cfg(target_arch = "x86_64")]
70        {
71            let f: extern "C" fn() = std::mem::transmute(span.rx());
72
73            f();
74        }
75    }
76}
More examples
Hide additional examples
examples/factorial.rs (line 39)
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}
Source

pub unsafe fn release(&mut self, rx_ptr: *const u8) -> Result<(), AsmError>

Releases the memory allocated by alloc.

§SAFETY
  • rx_ptr must have been returned from alloc
  • rx_ptr must have been allocaetd from this allocator
  • rx_ptr must not have been passed to release before
  • rx_ptr must point to read-execute part of memory returned from alloc.
Source

pub unsafe fn shrink( &mut self, rx_ptr: *const u8, new_size: usize, ) -> Result<(), AsmError>

Shrinks the memory allocated by alloc.

§SAFETY

rx_ptr must be a pointer returned by alloc.

Source

pub fn query(&self, rx_ptr: *const u8) -> Result<Span, AsmError>

Takes a pointer into the JIT memory and tries to query RX, RW mappings and size of the allocation.

Source

pub unsafe fn write( &mut self, span: &mut Span, write_func: impl FnMut(&mut Span), ) -> Result<(), AsmError>

Examples found in repository?
examples/reloc.rs (lines 39-56)
10fn main() {
11    let mut buf = CodeBuffer::new();
12    let mut asm = Assembler::new(&mut buf);
13
14    let str_constant = asm.add_constant("Hello, World!\0");
15    let puts_sym = asm
16        .buffer
17        .add_symbol(ExternalName::Symbol("puts".into()), RelocDistance::Far);
18
19    asm.lea64rm(RDI, ptr64_label(str_constant, 0));
20    asm.callm(ptr64_sym(puts_sym, 0));
21    asm.ret();
22
23    let result = buf.finish();
24
25    for reloc in result.relocs() {
26        println!("{:?}", reloc);
27    }
28
29    let mut jit = JitAllocator::new(Default::default());
30
31    // allocate memory for GOT table and for code itself
32    let mut span = jit
33        .alloc(result.data().len() + result.relocs().len() * 8)
34        .unwrap();
35
36    let mut got_addr_rx = std::ptr::null();
37
38    unsafe {
39        jit.write(&mut span, |span| {
40            span.rw()
41                .copy_from_nonoverlapping(result.data().as_ptr(), result.data().len());
42            got_addr_rx = span.rx().add(result.data().len());
43            span.rw()
44                .add(result.data().len())
45                .cast::<usize>()
46                .write(puts as *const u8 as usize);
47            // we only link to one symbol in GOT table, don't bother with anything else...
48            perform_relocations(
49                span.rw(),
50                span.rx(),
51                &result.relocs(),
52                |_| unreachable!(),
53                |_| got_addr_rx,
54                |_| unreachable!(),
55            );
56        })
57        .unwrap();
58
59        let mut out = String::new();
60        pretty_disassembler(
61            &mut out,
62            64,
63            std::slice::from_raw_parts(span.rx(), result.data().len()),
64            span.rx() as _,
65        )
66        .unwrap();
67
68        println!("{}", out);
69        #[cfg(target_arch = "x86_64")]
70        {
71            let f: extern "C" fn() = std::mem::transmute(span.rx());
72
73            f();
74        }
75    }
76}
More examples
Hide additional examples
examples/factorial.rs (lines 42-45)
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}
Source

pub unsafe fn copy_from_slice( &mut self, span: &mut Span, offset: usize, slice: &[u8], ) -> Result<(), AsmError>

Trait Implementations§

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.