Skip to content

Instantly share code, notes, and snippets.

@ClarkeRemy

ClarkeRemy/unsafe_arena_allocation.rs

Last active January 25, 2026 19:15
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save ClarkeRemy/feaa42afda54e31d532d289d6963044a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save ClarkeRemy/feaa42afda54e31d532d289d6963044a to your computer and use it in GitHub Desktop.
Download ZIP
Learning about arena allocation.
Raw
unsafe_arena_allocation.rs
extern crate alloc;
struct Arena {
start : *mut u8,
pos : *mut u8,
end : *mut u8,
layout : core::alloc::Layout,
}
fn arena_new(layout : core::alloc::Layout) -> Arena {
let start = unsafe { alloc::alloc::alloc(layout) };
Arena { start, pos: start, end: unsafe { start.byte_add(layout.size()) }, layout }
}
fn arena_free(arena : Arena) {
unsafe { alloc::alloc::dealloc(arena.start, arena.layout ) };
}
fn arena_push<T:Copy>(arena : &mut Arena, t : T) -> *mut T {
let out = arena_alloc::<T>(arena);
unsafe { out.write(t) };
out
}
fn arena_alloc<T:Copy>(arena : &mut Arena) -> *mut T {
let layout = core::alloc::Layout::new::<T>();
arena_alloc_layout(arena, layout) as *mut T
}
fn arena_push_slice<T:Copy+Sized>(arena : &mut Arena, slice : &[T]) -> *mut [T] {
let out = arena_alloc_slice::<T>(arena, slice.len());
unsafe { (out as *mut T).copy_from_nonoverlapping(slice as *const [T] as *const T, slice.len()) };
out
}
fn arena_alloc_slice<T:Copy>(arena : &mut Arena, len : usize) -> *mut [T] {
let layout = core::alloc::Layout::array::<T>(len).unwrap();
unsafe { core::slice::from_raw_parts_mut({
arena_alloc_layout(arena, layout) as *mut T
}, len) }
}
fn arena_alloc_layout(arena : &mut Arena, layout : core::alloc::Layout) -> *mut u8 {
let out = unsafe { arena.pos.add(arena.pos.align_offset(layout.align())) };
arena.pos = unsafe { out.add(layout.size()) };
out
}
struct ArenaMark(*mut u8);
fn arena_mark(arena : &Arena) -> ArenaMark {
ArenaMark(arena.pos)
}
fn arena_set_to_mark(arena : &mut Arena, mark : ArenaMark) {
core::assert!( arena.start <= mark.0 && mark.0 <= arena.pos );
arena.pos = mark.0;
}
fn arena_reset(arena : &mut Arena) {
arena.pos = arena.start
}
type Tag = u8;
mod expr {
use super::Tag;
#[repr(u8)]
enum Expr {
Var,
Const,
Add,
Mul,
_Variants,
}
pub(crate) const VAR : Tag = Expr::Var as u8;
pub(crate) const CONST : Tag = Expr::Const as u8;
pub(crate) const ADD : Tag = Expr::Add as u8;
pub(crate) const MUL : Tag = Expr::Mul as u8;
pub(crate) const _VARIANTS : Tag = Expr::_Variants as u8;
}
type Handle = *const u8;
type RangeHandle = [*const u8;2];
#[derive(Clone, Copy)]
#[repr(C)]
struct Tagged<const LEN : usize>{ handles : [Handle; LEN], tags : [Tag; LEN] }
fn tagged_nth<const LEN : usize>( Tagged { handles, tags } : Tagged<LEN>, n : usize) -> Tagged<1> { Tagged { handles: [handles[n]] , tags: [tags[n]] } }
#[derive(Clone, Copy)]
#[repr(C)]
struct TaggedRanges<const LEN : usize>{ handles : [RangeHandle; LEN], tags : [Tag; LEN] }
type Var = RangeHandle;
type Int = i64;
type Const = Int;
type Binop = Tagged<2>;
type Add = Binop;
type Mul = Binop;
#[derive(Clone, Copy)]
struct Expression(Tagged<1>);
fn deref_handle<T:Copy>(handle : Handle) -> T {
core::debug_assert!((handle as *const T).is_aligned());
unsafe { *(handle as *const T) }
}
fn deref_var ( handle : Handle ) -> Var { deref_handle(handle) }
fn deref_const( handle : Handle ) -> Const { deref_handle(handle) }
fn deref_binop( handle : Handle ) -> Binop { deref_handle(handle) }
fn mk_string(arena : &mut Arena, string : &[u8]) -> RangeHandle {
let s = arena_push_slice(arena, string);
unsafe {[
(s as *const [u8]) as *const u8,
((s as *const [u8]) as *const u8).add(s.len())
]}
}
fn mk_var(arena : &mut Arena, string : &[u8]) -> Tagged<1> {
let s = mk_string(arena, string);
let handle = arena_push(arena, s);
Tagged { handles: [handle as *const u8], tags: [expr::VAR] }
}
fn mk_const(arena : &mut Arena, c : Const) -> Tagged<1> {
let s = arena_push(arena, c);
Tagged { handles: [s as *const u8], tags: [expr::CONST] }
}
fn mk_binop(arena : &mut Arena, tag : Tag, left : Tagged<1>, right : Tagged<1>) -> Tagged<1> {
let lr = arena_push(arena, Tagged{handles:[left.handles[0], right.handles[0]], tags : [left.tags[0], right.tags[0]]});
Tagged { handles: [lr as *const u8], tags: [tag] }
}
fn range_handle_bytes<'a>([start, end] : RangeHandle)->&'a [u8] {
unsafe { std::slice::from_raw_parts(start, end as usize - start as usize) }
}
fn write_expression(arena : &mut Arena, expression : Expression, buffer : &mut[u8] ) -> usize {
let mut buf_cursor = 0;
let mark = arena_mark(arena);
let alloc = arena_alloc_slice::<Expression>(arena, 300);
let end = alloc as *mut Expression;
let mut cur = end;
unsafe { cur.write(expression) };
loop {
let Expression(Tagged { handles: [handle], tags : [tag] }) = unsafe { *cur };
match tag {
expr::VAR => {
const START : &[u8] = b"Var \"";
let range = range_handle_bytes(deref_var(handle));
let total_len = range.len() + START.len() + 1;
core::assert!(buffer.len()-buf_cursor >= total_len);
let fmt = [START, range, b"\""];
let memo_cursor = buf_cursor;
for each in fmt {
let next_cursor = buf_cursor + each.len();
buffer[buf_cursor..next_cursor].copy_from_slice(each);
buf_cursor = next_cursor
}
core::assert_eq!( buf_cursor - memo_cursor, total_len);
}
expr::CONST => {
let const_ = deref_const(handle);
const START : &[u8] = b"Const ";
let mut div = const_.abs();
let mut const_buffer = [b'?'; (1 + 3 * Const::BITS/u8::BITS) as usize];
let mut cursor = const_buffer.len();
'decimals : loop {
cursor -= 1;
const_buffer[cursor] = b'0' + (div % 10) as u8;
div = div/10;
if div == 0 { break 'decimals ; }
}
if const_.is_negative() {
cursor -= 1;
const_buffer[cursor] = b'-';
};
let total_len = const_buffer.len() - cursor + START.len();
assert!(buffer.len() - buf_cursor >= total_len,
"buffer.len() {}\
\nbuf_cursor {}\
\nbuffer.len() - buf_cursor {}\
\ntotal_len {}\
\n",
buffer.len(), buf_cursor, buffer.len()-buf_cursor,total_len
);
let fmt = [START, &const_buffer[cursor..]];
let memo_cursor = buf_cursor;
for each in fmt {
let next_cursor = buf_cursor + each.len();
buffer[buf_cursor..next_cursor].copy_from_slice(each);
buf_cursor = next_cursor
}
core::assert_eq!( buf_cursor - memo_cursor, total_len);
}
expr::ADD | expr::MUL => {
let Tagged { handles, tags } = deref_binop(handle);
let start = match tag {
expr::ADD => b"Add (",
expr::MUL => b"Mul (",
_ => unreachable!()
};
core::assert!(buffer.len() - buf_cursor >= start.len(),
"buffer.len() {}\
\nbuf_cursor {}\
\nbuffer.len() - buf_cursor {}\
\n", buffer.len(), buf_cursor, buffer.len() - buf_cursor
);
let next_cursor = buf_cursor + start.len();
buffer[buf_cursor..next_cursor].copy_from_slice(start);
buf_cursor = next_cursor;
unsafe {
const PAR : &[u8] = &[1, b')'];
const COMMA : &[u8] = &[2, b',', b' '];
let push_slice =
[ Tagged { handles: [PAR as *const _ as *const u8], tags: [expr::_VARIANTS] }
, Tagged { handles: [handles[1]], tags: [tags[1]] }
, Tagged { handles: [COMMA as *const _ as *const u8], tags: [expr::_VARIANTS] }
, Tagged { handles: [handles[0]], tags: [tags[0]] }
].map(Expression);
cur.copy_from_nonoverlapping(push_slice.as_ptr(), push_slice.len());
cur = cur.add(push_slice.len());
}
}
expr::_VARIANTS => {
let len = unsafe { *handle } as usize;
core::assert!(buffer.len() - buf_cursor >= len);
let slice = unsafe { core::slice::from_raw_parts(handle.add(1), len) };
let next_cursor = buf_cursor + len;
buffer[buf_cursor..next_cursor].copy_from_slice(slice);
buf_cursor = next_cursor;
}
_ => unreachable!()
};
if cur == end { break; }
cur = unsafe { cur.sub(1) };
}
arena_set_to_mark(arena, mark);
buf_cursor
}
fn simplify1(arena : &mut Arena, expr @ Expression(Tagged { handles : [handle], tags : [tag] }) : Expression) -> Expression {
match tag {
expr::ADD => {
let binop @ Tagged { handles, tags } = deref_binop(handle);
match tags {
[expr::CONST, expr::CONST] => Expression(mk_const(arena, (|[l,r]:[Const;2]|l+r)(handles.map(deref_const)))),
[expr::CONST, _ ] => match deref_const(handles[0]) { 0 => Expression(tagged_nth(binop, 1)), _ => expr, }
[_ , expr::CONST] => match deref_const(handles[1]) { 0 => Expression(tagged_nth(binop, 0)), _ => expr, }
_ => expr
}
}
expr::MUL => {
let binop @ Tagged { handles , tags } = deref_binop(handle);
match tags {
[expr::CONST, expr::CONST] => {Expression(mk_const(arena, (|[l,r]:[Const;2]|l*r)(handles.map(deref_const))))}
[expr::CONST, _ ] => match deref_const(handles[0]) { 0 => Expression(tagged_nth(binop, 0)), 1 => Expression(tagged_nth(binop, 1)), _ => expr }
[_ , expr::CONST] => match deref_const(handles[1]) { 0 => Expression(tagged_nth(binop, 1)), 1 => Expression(tagged_nth(binop, 0)), _ => expr }
_ => expr
}
}
expr::CONST | expr::VAR => expr,
_ => unreachable!(),
}
}
fn rec_simplify(arena : &mut Arena, expr @ Expression(Tagged { handles : [handle], tags : [tag] }): Expression) -> Expression {
match tag {
expr::ADD | expr::MUL => {
let binop = deref_binop(handle);
let mut e0 = Expression(tagged_nth(binop, 0));
let mut e1 = Expression(tagged_nth(binop, 1));
// top down optimizasion
e0 = simplify1(arena,e0);
e1 = simplify1(arena,e1);
// recursive optimization
e0 = rec_simplify(arena, e0);
e1 = rec_simplify(arena, e1);
// bottom up optimization
let binop_reconstruct = Expression(mk_binop(arena, tag, e0.0, e1.0));
simplify1(arena, binop_reconstruct)
}
_ => simplify1(arena, expr),
}
}
#[test]
fn display(){
let mut arena_ = arena_new(core::alloc::Layout::array::<usize>(2usize.pow(16)).unwrap());
let arena = &mut arena_;
let c_0 = mk_const(arena, -15);
let c_1 = mk_const(arena, 0);
let add = mk_binop(arena, expr::ADD, c_0, c_1);
let v0 = mk_var( arena, b"x");
let mul = mk_binop(arena, expr::MUL, add, v0);
let c_2 = mk_const(arena, 1);
let mul1 = mk_binop(arena, expr::MUL, mul, c_2);
let e = Expression(mul1);
let e = rec_simplify(arena, e);
let mut s = String::with_capacity(5000);
let v = unsafe { s.as_mut_vec() };
v.extend_from_slice(&[0;2000]);
let len = write_expression(arena, e, v);
unsafe { v.set_len(len) };
println!("{:?}", s);
arena_free(arena_);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment