holey-bytes/depell/wasm-hbc/src/lib.rs

#![feature(alloc_error_handler)]
#![no_std]

use {
    alloc::{borrow::ToOwned, string::String, vec::Vec},
    core::{
        alloc::{GlobalAlloc, Layout},
        cell::UnsafeCell,
    },
    hblang::{codegen::Codegen, parser::FileId},
};

extern crate alloc;

const ARENA_SIZE: usize = 4 * 128 * 1024;
const MAX_PANIC_SIZE: usize = 1024 * 4;
const MAX_INPUT_SIZE: usize = 32 * 4 * 1024;

#[cfg(target_arch = "wasm32")]
#[panic_handler]
pub fn handle_panic(_info: &core::panic::PanicInfo) -> ! {
    unsafe {
        use core::fmt::Write;
        let mut f = Write(&mut PANIC_MESSAGE[..]);
        _ = writeln!(f, "{}", info);
        PANIC_MESSAGE_LEN = 1024 - f.0.len();
    }

    core::arch::wasm32::unreachable();
}

#[no_mangle]
static mut PANIC_MESSAGE: [u8; MAX_PANIC_SIZE] = [0; MAX_PANIC_SIZE];
#[no_mangle]
static mut PANIC_MESSAGE_LEN: usize = 0;

#[global_allocator]
static ALLOCATOR: ArenaAllocator = ArenaAllocator::new();

#[cfg(target_arch = "wasm32")]
#[alloc_error_handler]
fn alloc_error(_: core::alloc::Layout) -> ! {
    core::arch::wasm32::unreachable()
}

#[repr(C, align(32))]
struct ArenaAllocator {
    arena: UnsafeCell<[u8; ARENA_SIZE]>,
    head: UnsafeCell<*mut u8>,
}

impl ArenaAllocator {
    const fn new() -> Self {
        ArenaAllocator {
            arena: UnsafeCell::new([0; ARENA_SIZE]),
            head: UnsafeCell::new(core::ptr::null_mut()),
        }
    }

    unsafe fn reset(&self) {
        (*self.head.get()) = self.arena.get().cast::<u8>().add(ARENA_SIZE);
    }
}

unsafe impl Sync for ArenaAllocator {}

unsafe impl GlobalAlloc for ArenaAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        let align = layout.align();

        let until = self.arena.get() as *mut u8;

        let new_head = (*self.head.get()).sub(size);
        let aligned_head = (new_head as usize & !(1 << (align - 1))) as *mut u8;

        if until > aligned_head {
            return core::ptr::null_mut();
        }

        *self.head.get() = aligned_head;
        aligned_head
    }

    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
        /* lol */
    }
}

#[no_mangle]
static MAX_INPUT: usize = MAX_INPUT_SIZE;
#[no_mangle]
static mut INPUT_LEN: usize = 0;
#[no_mangle]
static mut INPUT: [u8; MAX_INPUT_SIZE] = [0; MAX_INPUT_SIZE];

#[no_mangle]
unsafe fn compile_and_run() {
    ALLOCATOR.reset();

    struct File<'a> {
        path: &'a str,
        code: &'a mut str,
    }

    let files = {
        let mut input_bytes =
            core::slice::from_raw_parts_mut(core::ptr::addr_of_mut!(INPUT).cast::<u8>(), INPUT_LEN);

        let mut files = Vec::with_capacity(32);
        while let Some((&mut path_len, rest)) = input_bytes.split_first_chunk_mut() {
            let (path, rest) = rest.split_at_mut(u16::from_le_bytes(path_len) as usize);
            let (&mut code_len, rest) = rest.split_first_chunk_mut().unwrap();
            let (code, rest) = rest.split_at_mut(u16::from_le_bytes(code_len) as usize);
            files.push(File {
                path: core::str::from_utf8_unchecked(path),
                code: core::str::from_utf8_unchecked_mut(code),
            });
            input_bytes = rest;
        }
        files.sort_unstable_by_key(|f| f.path);
        files
    };

    let files = {
        let mut ctx = hblang::parser::ParserCtx::default();
        let paths = files.iter().map(|f| f.path).collect::<Vec<_>>();
        let mut loader = |path: &str, _: &str| Ok(paths.binary_search(&path).unwrap() as FileId);
        files
            .into_iter()
            .map(|f| {
                hblang::parser::Ast::new(
                    f.path,
                    // since 'free' does nothing this is fine
                    String::from_raw_parts(f.code.as_mut_ptr(), f.code.len(), f.code.len()),
                    &mut ctx,
                    &mut loader,
                )
            })
            .collect::<Vec<_>>()
    };

    let code = {
        let mut c = Codegen::default();
        c.files = files;
        c.generate();
        let mut buf = Vec::with_capacity(1024 * 8);
        c.assemble(&mut buf);
        buf
    };
}