use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::Arc;
use inkwell::context::Context;
use inkwell::execution_engine::ExecutionEngine;
use inkwell::passes::PassBuilderOptions;
use inkwell::targets::{
CodeModel, InitializationConfig, RelocMode, Target, TargetMachine, TargetTriple,
};
use inkwell::OptimizationLevel;
use relon_eval_api::inplace_return::ArenaRegions;
use relon_eval_api::{ClosureData, Evaluator, RuntimeError, Scope, Thunk, Value};
use relon_parser::Node;
use crate::codegen::{
emit_fast_entry, emit_module_funcs, emit_module_funcs_closed_world,
emit_module_funcs_closed_world_wasm, emit_module_funcs_wasm, is_buffer_protocol_signature,
ConstPool, EntryShape, FastPathProfile, WorldMode, ENTRY_SYMBOL, ENTRY_SYMBOL_FAST,
};
use crate::error::LlvmError;
use crate::state::ArenaState;
use crate::str_helpers::RELON_LLVM_STR_CONTAINS_ARENA_SYMBOL;
use inkwell::module::Linkage;
use inkwell::targets::FileType;
use inkwell::values::FunctionValue;
use std::path::Path;
pub const MAX_LEGACY_ARITY: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodegenTarget {
Native,
Wasm32,
}
#[allow(dead_code)]
const WASM32_DATA_LAYOUT: &str = "e-m:e-p:32:32-p10:8:8-p20:8:8-i64:64-n32:64-S128-ni:1:10:20";
const WASM32_TRIPLE: &str = "wasm32-wasi";
type LegacyEntryFn4 = unsafe extern "C" fn(i64, i64, i64, i64) -> i64;
type LegacyEntryFn3 = unsafe extern "C" fn(i64, i64, i64) -> i64;
type LegacyEntryFn2 = unsafe extern "C" fn(i64, i64) -> i64;
type LegacyEntryFn1 = unsafe extern "C" fn(i64) -> i64;
type LegacyEntryFn0 = unsafe extern "C" fn() -> i64;
type BufferEntryFn = unsafe extern "C" fn(
*const ArenaState,
i32, i32, i32, i32, i64, ) -> i32;
type FastEntryFn0 = unsafe extern "C" fn() -> i64;
type FastEntryFn1 = unsafe extern "C" fn(i64) -> i64;
type FastEntryFn2 = unsafe extern "C" fn(i64, i64) -> i64;
type FastEntryFn3 = unsafe extern "C" fn(i64, i64, i64) -> i64;
type FastEntryFn4 = unsafe extern "C" fn(i64, i64, i64, i64) -> i64;
type FastEntryFn5 = unsafe extern "C" fn(i64, i64, i64, i64, i64) -> i64;
type FastEntryFn6 = unsafe extern "C" fn(i64, i64, i64, i64, i64, i64) -> i64;
type FastEntryFn7 = unsafe extern "C" fn(i64, i64, i64, i64, i64, i64, i64) -> i64;
type FastEntryFn8 = unsafe extern "C" fn(i64, i64, i64, i64, i64, i64, i64, i64) -> i64;
struct JitOwned {
_engine: ExecutionEngine<'static>,
entry_ptr: usize,
fast_entry_ptr: Option<usize>,
ir_dump: String,
_ctx: Box<Context>,
}
unsafe impl Send for JitOwned {}
unsafe impl Sync for JitOwned {}
struct BufferSchema {
main_schema: relon_eval_api::schema_canonical::Schema,
return_schema: relon_eval_api::schema_canonical::Schema,
main_layout: relon_eval_api::layout::OffsetTable,
return_layout: relon_eval_api::layout::OffsetTable,
}
pub struct LlvmAotEvaluator {
jit: JitOwned,
entry_shape: EntryShape,
entry_arity: usize,
param_names: Vec<String>,
buffer_schema: Option<BufferSchema>,
fast_path_arity: usize,
fast_path_auto_dispatch: bool,
const_data: Vec<u8>,
native_imports: Vec<relon_ir::ir::NativeImport>,
host_fns: Arc<crate::state::HostFnRegistry>,
caps_mask: i64,
step_budget: AtomicI64,
}
thread_local! {
static LLVM_ARENA_POOL: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
}
fn step_budget_to_i64(steps: Option<u64>) -> i64 {
match steps {
None => 0,
Some(0) => -1,
Some(n) => i64::try_from(n).unwrap_or(i64::MAX),
}
}
impl LlvmAotEvaluator {
pub fn from_ir_direct(
ir: relon_ir::ir::Module,
param_names: Vec<String>,
) -> Result<Self, LlvmError> {
Self::from_ir_inner(ir, param_names, None)
}
pub fn from_source(src: &str) -> Result<Self, LlvmError> {
Self::from_source_with_options_inner(src, None)
}
pub fn from_source_with_options(
src: &str,
options: &relon_analyzer::AnalyzeOptions,
) -> Result<Self, LlvmError> {
Self::from_source_with_options_inner(src, Some(options))
}
fn from_source_with_options_inner(
src: &str,
options: Option<&relon_analyzer::AnalyzeOptions>,
) -> Result<Self, LlvmError> {
let (ir, main_schema, return_schema) = Self::lower_source_with_options(src, options)?;
let main_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&main_schema)
.map_err(|e| LlvmError::Codegen(format!("main schema layout: {e}")))?;
let return_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&return_schema)
.map_err(|e| LlvmError::Codegen(format!("return schema layout: {e}")))?;
let param_names: Vec<String> = main_schema.fields.iter().map(|f| f.name.clone()).collect();
let schema = BufferSchema {
main_schema,
return_schema,
main_layout,
return_layout,
};
Self::from_ir_inner(ir, param_names, Some(schema))
}
fn lower_source_with_options(
src: &str,
options: Option<&relon_analyzer::AnalyzeOptions>,
) -> Result<
(
relon_ir::ir::Module,
relon_eval_api::schema_canonical::Schema,
relon_eval_api::schema_canonical::Schema,
),
LlvmError,
> {
let owned;
let options: &relon_analyzer::AnalyzeOptions = match options {
Some(o) => {
if o.strict_mode {
owned = relon_analyzer::AnalyzeOptions {
strict_mode: false,
..o.clone()
};
&owned
} else {
o
}
}
None => {
owned = relon_analyzer::AnalyzeOptions {
strict_mode: false,
..Default::default()
};
&owned
}
};
let lowered = relon_ir::frontend::compile(src, options).map_err(|e| match e {
relon_ir::FrontendError::Parse(msg) => LlvmError::Parse(msg),
relon_ir::FrontendError::Analyze(n) => LlvmError::Analyze(n),
relon_ir::FrontendError::Lowering(msg) => {
LlvmError::Codegen(format!("lower_workspace_single: {msg}"))
}
})?;
Ok((lowered.module, lowered.main_schema, lowered.return_schema))
}
pub fn from_source_closed_world(
src: &str,
options: &relon_analyzer::AnalyzeOptions,
host_shim_src: &str,
) -> Result<Self, LlvmError> {
let (ir, main_schema, return_schema) = Self::lower_source_with_options(src, Some(options))?;
let main_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&main_schema)
.map_err(|e| LlvmError::Codegen(format!("main schema layout: {e}")))?;
let return_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&return_schema)
.map_err(|e| LlvmError::Codegen(format!("return schema layout: {e}")))?;
let param_names: Vec<String> = main_schema.fields.iter().map(|f| f.name.clone()).collect();
let schema = BufferSchema {
main_schema,
return_schema,
main_layout,
return_layout,
};
Self::from_ir_inner_world(
ir,
param_names,
Some(schema),
WorldMode::ClosedWorld,
Some(host_shim_src),
)
}
fn from_ir_inner(
ir: relon_ir::ir::Module,
param_names: Vec<String>,
buffer_schema: Option<BufferSchema>,
) -> Result<Self, LlvmError> {
Self::from_ir_inner_world(ir, param_names, buffer_schema, WorldMode::OpenWorld, None)
}
fn from_ir_inner_world(
ir: relon_ir::ir::Module,
param_names: Vec<String>,
buffer_schema: Option<BufferSchema>,
world_mode: WorldMode,
host_shim_src: Option<&str>,
) -> Result<Self, LlvmError> {
let entry_idx = ir
.entry_func_index
.ok_or_else(|| LlvmError::Codegen("IR module has no entry function".into()))?;
let entry = &ir.funcs[entry_idx];
let buffer_shape = is_buffer_protocol_signature(&entry.params, entry.ret);
if !buffer_shape && entry.params.len() > MAX_LEGACY_ARITY {
return Err(LlvmError::UnsupportedSignature(format!(
"llvm-aot: {} params exceeds MAX_LEGACY_ARITY={MAX_LEGACY_ARITY}",
entry.params.len()
)));
}
let declared_arity = if buffer_shape {
buffer_schema
.as_ref()
.map(|s| s.main_schema.fields.len())
.unwrap_or(0)
} else {
entry.params.len()
};
if param_names.len() != declared_arity {
return Err(LlvmError::UnsupportedSignature(format!(
"llvm-aot: param_names len {} does not match declared arity {declared_arity}",
param_names.len()
)));
}
if buffer_shape && buffer_schema.is_none() {
return Err(LlvmError::UnsupportedSignature(
"llvm-aot: buffer-protocol IR requires schema metadata; use from_source".into(),
));
}
if !buffer_shape && buffer_schema.is_some() {
return Err(LlvmError::UnsupportedSignature(
"llvm-aot: schema metadata supplied for non-buffer entry".into(),
));
}
let ctx_box: Box<Context> = Box::new(Context::create());
let ctx_static: &'static Context = unsafe { &*(ctx_box.as_ref() as *const Context) };
let module = ctx_static.create_module("relon_llvm_aot");
let buffer_return_size = buffer_schema
.as_ref()
.map(|s| s.return_layout.root_size as u32)
.unwrap_or(0);
let const_pool = ConstPool::from_module(&ir)?;
let lambda_ir_idx_set: std::collections::HashSet<u32> =
ir.closure_table.iter().copied().collect();
let helpers: Vec<&relon_ir::ir::Func> = ir
.funcs
.iter()
.enumerate()
.filter(|(i, _)| *i != entry_idx && !lambda_ir_idx_set.contains(&(*i as u32)))
.map(|(_, f)| f)
.collect();
let helper_ir_indices: Vec<u32> = ir
.funcs
.iter()
.enumerate()
.filter(|(i, _)| *i != entry_idx && !lambda_ir_idx_set.contains(&(*i as u32)))
.map(|(i, _)| i as u32)
.collect();
let lambdas: Vec<&relon_ir::ir::Func> = ir
.closure_table
.iter()
.map(|&ir_idx| &ir.funcs[ir_idx as usize])
.collect();
let emit = match world_mode {
WorldMode::OpenWorld => emit_module_funcs,
WorldMode::ClosedWorld => emit_module_funcs_closed_world,
};
let (_llvm_fn, entry_shape, helper_table, closure_fn_table) = emit(
ctx_static,
&module,
entry,
buffer_return_size,
&const_pool,
&helpers,
Some(&helper_ir_indices),
&lambdas,
&ir.closure_table,
&ir.imports,
)?;
if matches!(world_mode, WorldMode::ClosedWorld) {
let shim = host_shim_src.ok_or_else(|| {
LlvmError::Codegen(
"from_ir_inner_world: ClosedWorld requires a host_shim_src".into(),
)
})?;
crate::cocompile::link_and_inline_host_shim(&module, shim, &ir.imports)?;
}
let fast_profile = buffer_schema
.as_ref()
.filter(|_| ir.closure_table.is_empty())
.and_then(|s| build_fast_path_profile(s).ok());
let fast_path_auto_dispatch = !body_may_raise_typed_trap(&entry.body);
let mut fast_emit_diagnostic: Option<String> = None;
if let Some(profile) = fast_profile.as_ref() {
match emit_fast_entry(
ctx_static,
&module,
entry,
profile,
&helper_table,
&closure_fn_table,
) {
Ok(_) => {}
Err(e) => {
fast_emit_diagnostic = Some(format!("{e}"));
if let Some(f) = module.get_function(ENTRY_SYMBOL_FAST) {
unsafe { f.delete() };
}
}
}
}
module
.verify()
.map_err(|e| LlvmError::Codegen(format!("LLVM verifier rejected module: {e}")))?;
stamp_host_target_attributes(&module);
let preopt_dump: Option<String> = std::env::var_os("RELON_LLVM_DUMP_PREOPT")
.map(|_| module.print_to_string().to_string());
run_default_o3_pipeline(&module)?;
let mut ir_dump = module.print_to_string().to_string();
if let Some(p) = preopt_dump {
ir_dump = format!("; --- PRE-OPT IR ---\n{p}\n; --- POST-OPT IR ---\n{ir_dump}");
}
if let Some(dir) = std::env::var_os("RELON_LLVM_DUMP_DIR") {
let dir = std::path::PathBuf::from(dir);
let _ = std::fs::create_dir_all(&dir);
let _ = std::fs::write(dir.join("module.post_o3.ll"), &ir_dump);
if let Ok(()) = Target::initialize_native(&InitializationConfig::default()) {
let triple_str = TargetMachine::get_default_triple();
if let Ok(target) = Target::from_triple(&triple_str) {
let cpu = TargetMachine::get_host_cpu_name();
let features = TargetMachine::get_host_cpu_features();
if let Ok(triple_utf8) = triple_str.as_str().to_str() {
let triple = TargetTriple::create(triple_utf8);
if let Some(machine) = target.create_target_machine(
&triple,
cpu.to_str().unwrap_or(""),
features.to_str().unwrap_or(""),
OptimizationLevel::Aggressive,
RelocMode::Default,
CodeModel::JITDefault,
) {
let _ = machine.write_to_file(
&module,
FileType::Assembly,
&dir.join("module.s"),
);
let _ = machine.write_to_file(
&module,
FileType::Object,
&dir.join("module.o"),
);
}
if let Some(machine) = target.create_target_machine(
&triple,
cpu.to_str().unwrap_or(""),
features.to_str().unwrap_or(""),
OptimizationLevel::Aggressive,
RelocMode::PIC,
CodeModel::Small,
) {
let _ = machine.write_to_file(
&module,
FileType::Assembly,
&dir.join("module.small_pic.s"),
);
}
if let Some(machine) = target.create_target_machine(
&triple,
cpu.to_str().unwrap_or(""),
features.to_str().unwrap_or(""),
OptimizationLevel::Aggressive,
RelocMode::Static,
CodeModel::Small,
) {
let _ = machine.write_to_file(
&module,
FileType::Assembly,
&dir.join("module.small_static.s"),
);
}
}
}
}
}
let uses_extern_shim = module
.get_function(crate::str_helpers::RELON_LLVM_STR_CONTAINS_ARENA_SYMBOL)
.is_some()
|| module
.get_function(crate::str_helpers::RELON_LLVM_F64_TO_STR_SYMBOL)
.is_some()
|| module
.get_function(crate::state::RELON_LLVM_CALL_NATIVE_SYMBOL)
.is_some();
let force_default_mcjit = std::env::var_os("RELON_LLVM_FORCE_DEFAULT_MCJIT").is_some();
let engine = if uses_extern_shim || force_default_mcjit {
module
.create_jit_execution_engine(OptimizationLevel::Aggressive)
.map_err(|e| LlvmError::Codegen(format!("create_jit_execution_engine: {e}")))?
} else {
let mm = crate::mcjit_mm::ContiguousCodeMemoryManager::new();
module
.create_mcjit_execution_engine_with_memory_manager(
mm,
OptimizationLevel::Aggressive,
inkwell::targets::CodeModel::Small,
false,
false,
)
.map_err(|e| {
LlvmError::Codegen(format!(
"create_mcjit_execution_engine_with_memory_manager (Small CodeModel): {e}"
))
})?
};
if let Some(shim_fn) =
module.get_function(crate::str_helpers::RELON_LLVM_STR_CONTAINS_ARENA_SYMBOL)
{
engine.add_global_mapping(
&shim_fn,
crate::str_helpers::relon_llvm_str_contains_arena_addr(),
);
}
if let Some(shim_fn) = module.get_function(crate::str_helpers::RELON_LLVM_F64_TO_STR_SYMBOL)
{
engine.add_global_mapping(&shim_fn, crate::str_helpers::relon_llvm_f64_to_str_addr());
}
if let Some(cn_fn) = module.get_function(crate::state::RELON_LLVM_CALL_NATIVE_SYMBOL) {
engine.add_global_mapping(&cn_fn, crate::state::relon_llvm_call_native_addr());
}
let entry_ptr = engine.get_function_address(ENTRY_SYMBOL).map_err(|e| {
LlvmError::Codegen(format!(
"ExecutionEngine could not resolve `{ENTRY_SYMBOL}`: {e}"
))
})?;
let (fast_entry_ptr, fast_path_arity) = match (&fast_profile, &fast_emit_diagnostic) {
(Some(profile), None) => match engine.get_function_address(ENTRY_SYMBOL_FAST) {
Ok(ptr) => (Some(ptr), profile.arg_offsets.len()),
Err(_) => (None, 0),
},
_ => (None, 0),
};
let ir_dump = match fast_emit_diagnostic {
Some(diag) => format!("; fast-emit diagnostic: {diag}\n{ir_dump}"),
None => ir_dump,
};
Ok(Self {
jit: JitOwned {
_engine: engine,
entry_ptr,
fast_entry_ptr,
ir_dump,
_ctx: ctx_box,
},
entry_shape,
entry_arity: entry.params.len(),
param_names,
buffer_schema,
fast_path_arity,
fast_path_auto_dispatch,
const_data: const_pool.bytes,
native_imports: ir.imports.clone(),
host_fns: Arc::new(crate::state::HostFnRegistry::new()),
caps_mask: 0,
step_budget: AtomicI64::new(0),
})
}
pub fn arity(&self) -> usize {
self.param_names.len()
}
pub fn param_names(&self) -> &[String] {
&self.param_names
}
pub fn native_imports(&self) -> &[relon_ir::ir::NativeImport] {
&self.native_imports
}
pub fn with_host_fns(
mut self,
host_fns: &std::collections::HashMap<String, Arc<dyn relon_eval_api::RelonFunction>>,
) -> Self {
let mut registry = crate::state::HostFnRegistry::new();
for (idx, imp) in self.native_imports.iter().enumerate() {
if let Some(func) = host_fns.get(&imp.name) {
registry.register(idx as u32, Arc::clone(func));
}
}
self.host_fns = Arc::new(registry);
self
}
pub fn with_granted_cap(mut self, bit: u32) -> Self {
if bit < 64 {
self.caps_mask |= 1i64 << bit;
}
self
}
pub fn with_caps(mut self, caps_mask: i64) -> Self {
self.caps_mask = caps_mask;
self
}
pub fn set_step_budget(&self, steps: Option<u64>) {
self.step_budget
.store(step_budget_to_i64(steps), Ordering::Relaxed);
}
pub fn with_step_budget(self, steps: Option<u64>) -> Self {
self.set_step_budget(steps);
self
}
pub fn run_main_legacy_i64(&self, args: &[i64]) -> Result<i64, RuntimeError> {
if self.entry_shape != EntryShape::LegacyI64 {
return Err(RuntimeError::Unsupported {
reason: "llvm-aot: run_main_legacy_i64 called on buffer-protocol entry".into(),
});
}
if args.len() != self.entry_arity {
return Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: #main expects {} arg(s), got {}",
self.entry_arity,
args.len()
),
});
}
let ptr = self.jit.entry_ptr;
unsafe {
match self.entry_arity {
0 => {
let f: LegacyEntryFn0 = std::mem::transmute(ptr);
Ok(f())
}
1 => {
let f: LegacyEntryFn1 = std::mem::transmute(ptr);
Ok(f(args[0]))
}
2 => {
let f: LegacyEntryFn2 = std::mem::transmute(ptr);
Ok(f(args[0], args[1]))
}
3 => {
let f: LegacyEntryFn3 = std::mem::transmute(ptr);
Ok(f(args[0], args[1], args[2]))
}
4 => {
let f: LegacyEntryFn4 = std::mem::transmute(ptr);
Ok(f(args[0], args[1], args[2], args[3]))
}
n => Err(RuntimeError::Unsupported {
reason: format!("llvm-aot: arity {n} > MAX_LEGACY_ARITY={MAX_LEGACY_ARITY}"),
}),
}
}
}
pub fn emit_ir_dump(&self) -> &str {
&self.jit.ir_dump
}
pub fn has_fast_path(&self) -> bool {
self.jit.fast_entry_ptr.is_some()
}
pub fn fast_path_arity(&self) -> Option<usize> {
self.jit.fast_entry_ptr.map(|_| self.fast_path_arity)
}
pub fn fast_entry_runtime_addr(&self) -> Option<usize> {
self.jit.fast_entry_ptr
}
pub fn entry_runtime_addr(&self) -> usize {
self.jit.entry_ptr
}
pub fn host_target_cpu() -> String {
TargetMachine::get_host_cpu_name()
.to_str()
.unwrap_or("")
.to_string()
}
pub fn run_main_legacy_i64_fast(&self, args: &[i64]) -> Result<i64, RuntimeError> {
let ptr = self
.jit
.fast_entry_ptr
.ok_or_else(|| RuntimeError::Unsupported {
reason:
"llvm-aot: fast entry not available; source not Int-only or fast-emit failed"
.into(),
})?;
let arity = self.fast_path_arity;
if args.len() != arity {
return Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot fast path: #main expects {arity} arg(s), got {}",
args.len()
),
});
}
unsafe {
let r = match arity {
0 => {
let f: FastEntryFn0 = std::mem::transmute(ptr);
f()
}
1 => {
let f: FastEntryFn1 = std::mem::transmute(ptr);
f(args[0])
}
2 => {
let f: FastEntryFn2 = std::mem::transmute(ptr);
f(args[0], args[1])
}
3 => {
let f: FastEntryFn3 = std::mem::transmute(ptr);
f(args[0], args[1], args[2])
}
4 => {
let f: FastEntryFn4 = std::mem::transmute(ptr);
f(args[0], args[1], args[2], args[3])
}
5 => {
let f: FastEntryFn5 = std::mem::transmute(ptr);
f(args[0], args[1], args[2], args[3], args[4])
}
6 => {
let f: FastEntryFn6 = std::mem::transmute(ptr);
f(args[0], args[1], args[2], args[3], args[4], args[5])
}
7 => {
let f: FastEntryFn7 = std::mem::transmute(ptr);
f(
args[0], args[1], args[2], args[3], args[4], args[5], args[6],
)
}
8 => {
let f: FastEntryFn8 = std::mem::transmute(ptr);
f(
args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7],
)
}
n => {
return Err(RuntimeError::Unsupported {
reason: format!("llvm-aot fast path: arity {n} > 8 dispatch cap"),
});
}
};
Ok(r)
}
}
fn try_run_main_fast(
&self,
args: &HashMap<String, Value>,
) -> Result<Option<Value>, RuntimeError> {
if self.jit.fast_entry_ptr.is_none() {
return Ok(None);
}
if !self.fast_path_auto_dispatch {
return Ok(None);
}
let arity = self.fast_path_arity;
if arity != self.param_names.len() {
return Ok(None);
}
let mut argv = [0i64; 8];
for (i, name) in self.param_names.iter().enumerate() {
match args.get(name) {
Some(Value::Int(v)) => argv[i] = *v,
_ => return Ok(None), }
}
let r = self.run_main_legacy_i64_fast(&argv[..arity])?;
if let Some(schema) = self.buffer_schema.as_ref() {
if is_single_value_wrapper(&schema.return_schema) {
Ok(Some(Value::Int(r)))
} else {
let field_name = schema.return_schema.fields[0].name.clone();
let mut map: HashMap<String, Value> = HashMap::with_capacity(1);
map.insert(field_name, Value::Int(r));
Ok(Some(Value::branded_dict(
map,
Some(schema.return_schema.name.clone()),
)))
}
} else {
Ok(Some(Value::Int(r)))
}
}
fn run_main_buffer(&self, args: HashMap<String, Value>) -> Result<Value, RuntimeError> {
let schema = self
.buffer_schema
.as_ref()
.ok_or_else(|| RuntimeError::Unsupported {
reason: "llvm-aot: run_main_buffer called without schema metadata".into(),
})?;
let mut builder = relon_eval_api::buffer::BufferBuilder::new(
&schema.main_layout,
&schema.main_schema.fields,
);
for field in &schema.main_schema.fields {
let value = args
.get(&field.name)
.ok_or_else(|| RuntimeError::Unsupported {
reason: format!("llvm-aot: missing #main arg `{}`", field.name),
})?;
write_value_into_builder(&mut builder, field, value)?;
}
let in_ptr_pre = relon_util::align_up(
u32::try_from(self.const_data.len()).map_err(|_| {
RuntimeError::IoError("llvm const-data section exceeds u32 range".into())
})?,
8,
);
let in_bytes = builder
.finish_arena_absolute(in_ptr_pre)
.map_err(buffer_to_runtime_error)?;
let in_len = in_bytes.len() as u32;
let out_root_size = schema.return_layout.root_size as u32;
let needs_pointer_indirect_return = return_needs_tail_region(&schema.return_schema);
let tail_cap: u32 = if needs_pointer_indirect_return {
65_536
} else {
0
};
let out_cap = relon_util::align_up(out_root_size.max(8) + tail_cap + 16, 8);
let const_data_len = u32::try_from(self.const_data.len()).map_err(|_| {
RuntimeError::IoError("llvm const-data section exceeds u32 range".into())
})?;
let in_ptr = relon_util::align_up(const_data_len, 8);
let out_ptr = relon_util::align_up(in_ptr + in_len, 8);
let scratch_base = relon_util::align_up(out_ptr + out_cap, 8);
let scratch_size: u32 = 1_048_576; let arena_size = (scratch_base + scratch_size) as usize;
LLVM_ARENA_POOL.with(|cell| match cell.try_borrow_mut() {
Ok(mut buf) => self.dispatch_with_arena(
schema,
&mut buf,
arena_size,
in_ptr,
in_len,
out_ptr,
out_cap,
scratch_base,
&in_bytes,
),
Err(_) => {
let mut fallback: Vec<u8> = Vec::new();
self.dispatch_with_arena(
schema,
&mut fallback,
arena_size,
in_ptr,
in_len,
out_ptr,
out_cap,
scratch_base,
&in_bytes,
)
}
})
}
#[allow(clippy::too_many_arguments)]
fn dispatch_with_arena(
&self,
schema: &BufferSchema,
arena: &mut Vec<u8>,
arena_size: usize,
in_ptr: u32,
in_len: u32,
out_ptr: u32,
out_cap: u32,
scratch_base: u32,
in_bytes: &[u8],
) -> Result<Value, RuntimeError> {
if arena.len() < arena_size {
arena.resize(arena_size, 0);
}
let observable_end = (out_ptr + out_cap) as usize;
debug_assert!(observable_end <= arena_size);
debug_assert!(self.const_data.len() <= in_ptr as usize);
arena[self.const_data.len()..observable_end].fill(0);
if !self.const_data.is_empty() {
arena[..self.const_data.len()].copy_from_slice(&self.const_data);
}
arena[in_ptr as usize..in_ptr as usize + in_bytes.len()].copy_from_slice(in_bytes);
let live_arena = &mut arena[..arena_size];
let state = ArenaState::new(live_arena, scratch_base);
state.set_step_budget(self.step_budget.load(Ordering::Relaxed));
unsafe {
state.install_host_fns(Arc::as_ptr(&self.host_fns));
}
let state_ptr: *const ArenaState = &state;
let bytes_written = {
let f: BufferEntryFn = unsafe { std::mem::transmute(self.jit.entry_ptr) };
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| unsafe {
f(
state_ptr,
in_ptr as i32,
in_len as i32,
out_ptr as i32,
out_cap as i32,
self.caps_mask,
)
}))
.map_err(|_| RuntimeError::Unsupported {
reason: "llvm-aot: JIT entry panicked (no trap-code recovery in Phase B)".into(),
})?
};
let trap_code = state.trap_code();
if trap_code != 0 {
return Err(crate::state::NativeTrap::runtime_error_from_code(trap_code));
}
self.decode_buffer_return(
schema,
arena,
ArenaRegions {
const_data_len: self.const_data.len(),
in_ptr,
in_len,
out_ptr,
out_cap,
scratch_base,
arena_size,
},
bytes_written,
)
}
fn decode_buffer_return(
&self,
schema: &BufferSchema,
arena: &[u8],
regions: ArenaRegions,
ret: i32,
) -> Result<Value, RuntimeError> {
if ret < 0 {
let root_abs = relon_eval_api::inplace_return::decode_inplace_sentinel(ret)?;
if !is_single_value_wrapper(&schema.return_schema) {
return Err(RuntimeError::IoError(
"llvm-aot in-place return on a non-single-value return schema".into(),
));
}
return relon_eval_api::inplace_return::decode_inplace_return(
"llvm-aot",
arena,
regions,
root_abs,
&schema.return_schema.fields[0],
&schema.return_layout,
&schema.return_schema.fields,
);
}
let bw = ret as usize;
let read_len = bw.max(schema.return_layout.root_size);
let out_ptr = regions.out_ptr as usize;
let read_end = out_ptr + read_len;
if read_end > regions.arena_size || read_end > arena.len() {
return Err(RuntimeError::IoError(
"llvm-aot arena too small for return decode".into(),
));
}
let arena = &arena[..regions.arena_size.min(arena.len())];
relon_eval_api::inplace_return::decode_object_return(
"llvm-aot",
arena,
out_ptr,
regions,
&schema.return_layout,
&schema.return_schema,
is_single_value_wrapper(&schema.return_schema),
)
}
pub fn wasm_buffer_plan(
&self,
args: &HashMap<String, Value>,
) -> Result<WasmBufferDispatch, RuntimeError> {
let schema = self
.buffer_schema
.as_ref()
.ok_or_else(|| RuntimeError::Unsupported {
reason: "llvm-aot: wasm_buffer_plan called without schema metadata".into(),
})?;
let mut builder = relon_eval_api::buffer::BufferBuilder::new(
&schema.main_layout,
&schema.main_schema.fields,
);
for field in &schema.main_schema.fields {
let value = args
.get(&field.name)
.ok_or_else(|| RuntimeError::Unsupported {
reason: format!("llvm-aot: missing #main arg `{}`", field.name),
})?;
write_value_into_builder(&mut builder, field, value)?;
}
let in_ptr_pre = relon_util::align_up(
u32::try_from(self.const_data.len()).map_err(|_| {
RuntimeError::IoError("llvm const-data section exceeds u32 range".into())
})?,
8,
);
let in_bytes = builder
.finish_arena_absolute(in_ptr_pre)
.map_err(buffer_to_runtime_error)?;
let in_len = in_bytes.len() as u32;
let out_root_size = schema.return_layout.root_size as u32;
let needs_pointer_indirect_return = return_needs_tail_region(&schema.return_schema);
let tail_cap: u32 = if needs_pointer_indirect_return {
65_536
} else {
0
};
let out_cap = relon_util::align_up(out_root_size.max(8) + tail_cap + 16, 8);
let const_data_len = u32::try_from(self.const_data.len()).map_err(|_| {
RuntimeError::IoError("llvm const-data section exceeds u32 range".into())
})?;
let in_ptr = relon_util::align_up(const_data_len, 8);
let out_ptr = relon_util::align_up(in_ptr + in_len, 8);
let scratch_base = relon_util::align_up(out_ptr + out_cap, 8);
let scratch_size: u32 = 1_048_576;
let arena_size = (scratch_base + scratch_size) as usize;
Ok(WasmBufferDispatch {
const_data: self.const_data.clone(),
in_bytes,
regions: ArenaRegions {
const_data_len: self.const_data.len(),
in_ptr,
in_len,
out_ptr,
out_cap,
scratch_base,
arena_size,
},
})
}
pub fn wasm_buffer_decode(
&self,
arena: &[u8],
regions: ArenaRegions,
ret: i32,
) -> Result<Value, RuntimeError> {
let schema = self
.buffer_schema
.as_ref()
.ok_or_else(|| RuntimeError::Unsupported {
reason: "llvm-aot: wasm_buffer_decode called without schema metadata".into(),
})?;
self.decode_buffer_return(schema, arena, regions, ret)
}
}
#[derive(Debug, Clone)]
pub struct WasmBufferDispatch {
pub const_data: Vec<u8>,
pub in_bytes: Vec<u8>,
pub regions: ArenaRegions,
}
impl Evaluator for LlvmAotEvaluator {
fn eval(&self, _node: &Node, _scope: &Arc<Scope>) -> Result<Value, RuntimeError> {
Err(RuntimeError::Unsupported {
reason: "llvm-aot: `eval` is not supported".into(),
})
}
fn eval_root(&self, _scope: &Arc<Scope>) -> Result<Value, RuntimeError> {
Err(RuntimeError::Unsupported {
reason: "llvm-aot: `eval_root` is not supported".into(),
})
}
fn run_main(&self, args: HashMap<String, Value>) -> Result<Value, RuntimeError> {
if let Some(v) = self.try_run_main_fast(&args)? {
return Ok(v);
}
match self.entry_shape {
EntryShape::Buffer => self.run_main_buffer(args),
EntryShape::LegacyI64 => {
let mut argv = [0i64; MAX_LEGACY_ARITY];
for (i, name) in self.param_names.iter().enumerate() {
let v = args.get(name).ok_or_else(|| RuntimeError::Unsupported {
reason: format!("llvm-aot: missing #main arg `{name}`"),
})?;
match v {
Value::Int(n) => argv[i] = *n,
other => {
return Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: legacy-i64 #main arg `{name}` is {} (Int only)",
other.type_name()
),
});
}
}
}
let r = self.run_main_legacy_i64(&argv[..self.entry_arity])?;
Ok(Value::Int(r))
}
}
}
fn force_thunk(&self, _thunk: &Arc<Thunk>) -> Result<Value, RuntimeError> {
Err(RuntimeError::Unsupported {
reason: "llvm-aot: `force_thunk` is not supported".into(),
})
}
fn invoke_closure(
&self,
_closure: &ClosureData,
_args: &[Value],
) -> Result<Value, RuntimeError> {
Err(RuntimeError::Unsupported {
reason: "llvm-aot: `invoke_closure` is not supported".into(),
})
}
}
fn buffer_to_runtime_error(e: relon_eval_api::buffer::BufferError) -> RuntimeError {
RuntimeError::IoError(format!("llvm-aot buffer: {e}"))
}
fn is_single_value_wrapper(schema: &relon_eval_api::schema_canonical::Schema) -> bool {
schema.name == relon_ir::MAIN_RETURN_SCHEMA_NAME
&& schema.fields.len() == 1
&& schema.fields[0].name == relon_ir::RETURN_VALUE_FIELD_NAME
}
fn is_single_int_field_record(schema: &relon_eval_api::schema_canonical::Schema) -> bool {
use relon_eval_api::schema_canonical::TypeRepr;
!schema.is_tuple && schema.fields.len() == 1 && matches!(schema.fields[0].ty, TypeRepr::Int)
}
fn write_value_into_builder(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
field: &relon_eval_api::schema_canonical::Field,
value: &Value,
) -> Result<(), RuntimeError> {
use relon_eval_api::schema_canonical::TypeRepr;
match (&field.ty, value) {
(TypeRepr::Int, Value::Int(v)) => marshal_int_in(builder, &field.name, *v),
(TypeRepr::Float, Value::Float(v)) => {
marshal_float_in(builder, &field.name, v.into_inner())
}
(TypeRepr::Float, Value::Int(v)) => marshal_float_in(builder, &field.name, *v as f64),
(TypeRepr::Bool, Value::Bool(v)) => marshal_bool_in(builder, &field.name, *v),
(TypeRepr::Unit, v) if v.is_option_none() => marshal_unit_in(builder, &field.name),
(TypeRepr::String, Value::String(s)) => marshal_string_in(builder, &field.name, s),
(TypeRepr::Schema { schema }, Value::Dict(dict)) if !schema.is_tuple => {
marshal_schema_in(builder, &field.name, schema, dict)
}
(TypeRepr::Schema { schema }, Value::Tuple(items)) if schema.is_tuple => {
marshal_tuple_in(builder, &field.name, schema, items.as_ref())
}
(TypeRepr::List { element }, Value::List(items)) => {
marshal_list_in(builder, &field.name, element, items)
}
(TypeRepr::Option { .. } | TypeRepr::Result { .. } | TypeRepr::Enum { .. }, _) => builder
.write_value(&field.name, &field.ty, value)
.map_err(buffer_to_runtime_error),
(ty, v) => Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: #main arg `{}` got {} but schema expects {ty:?}",
field.name,
v.type_name()
),
}),
}
}
fn marshal_int_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
v: i64,
) -> Result<(), RuntimeError> {
builder.write_int(name, v).map_err(buffer_to_runtime_error)
}
fn marshal_float_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
v: f64,
) -> Result<(), RuntimeError> {
builder
.write_float(name, v)
.map_err(buffer_to_runtime_error)
}
fn marshal_bool_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
v: bool,
) -> Result<(), RuntimeError> {
builder.write_bool(name, v).map_err(buffer_to_runtime_error)
}
fn marshal_unit_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
) -> Result<(), RuntimeError> {
builder.write_unit(name).map_err(buffer_to_runtime_error)
}
fn marshal_string_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
s: &str,
) -> Result<(), RuntimeError> {
builder
.write_string(name, s)
.map_err(buffer_to_runtime_error)
}
fn marshal_list_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
element: &relon_eval_api::schema_canonical::TypeRepr,
items: &[Value],
) -> Result<(), RuntimeError> {
use relon_eval_api::schema_canonical::TypeRepr;
let mismatch = |idx: usize, got: &Value, want: &str| RuntimeError::Unsupported {
reason: format!(
"llvm-aot: List<{want}> arg `{name}` element #{idx} got {} but expects {want}",
got.type_name()
),
};
match element {
TypeRepr::Int => {
let mut out = Vec::with_capacity(items.len());
for (i, it) in items.iter().enumerate() {
match it {
Value::Int(v) => out.push(*v),
other => return Err(mismatch(i, other, "Int")),
}
}
builder
.write_list_int(name, &out)
.map_err(buffer_to_runtime_error)
}
TypeRepr::Float => {
let mut out = Vec::with_capacity(items.len());
for (i, it) in items.iter().enumerate() {
match it {
Value::Float(v) => out.push(v.into_inner()),
Value::Int(v) => out.push(*v as f64),
other => return Err(mismatch(i, other, "Float")),
}
}
builder
.write_list_float(name, &out)
.map_err(buffer_to_runtime_error)
}
TypeRepr::Bool => {
let mut out = Vec::with_capacity(items.len());
for (i, it) in items.iter().enumerate() {
match it {
Value::Bool(v) => out.push(*v),
other => return Err(mismatch(i, other, "Bool")),
}
}
builder
.write_list_bool(name, &out)
.map_err(buffer_to_runtime_error)
}
TypeRepr::String => {
let mut out: Vec<&str> = Vec::with_capacity(items.len());
for (i, it) in items.iter().enumerate() {
match it {
Value::String(s) => out.push(s.as_str()),
other => return Err(mismatch(i, other, "String")),
}
}
builder
.write_list_string(name, &out)
.map_err(buffer_to_runtime_error)
}
TypeRepr::Schema { schema } => marshal_list_schema_in(builder, name, schema, items),
TypeRepr::List { element: inner } => marshal_list_list_in(builder, name, inner, items),
TypeRepr::Option { .. } | TypeRepr::Result { .. } | TypeRepr::Enum { .. } => {
let ty = TypeRepr::List {
element: Box::new(element.clone()),
};
builder
.write_value(name, &ty, &Value::List(Arc::new(items.to_vec())))
.map_err(buffer_to_runtime_error)
}
other => Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: List element type {other:?} for arg `{name}` is not yet materialised \
(List<Int/Float/Bool/String/Schema> + List<List<scalar>>)"
),
}),
}
}
fn marshal_list_schema_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
schema: &relon_eval_api::schema_canonical::Schema,
items: &[Value],
) -> Result<(), RuntimeError> {
let elem_layout = relon_eval_api::layout::SchemaLayout::offsets_for(schema).map_err(|e| {
RuntimeError::Unsupported {
reason: format!("llvm-aot: List<Schema> arg `{name}` element layout: {e}"),
}
})?;
let mut writer = builder
.list_record_writer(name, &elem_layout, schema)
.map_err(buffer_to_runtime_error)?;
for (i, it) in items.iter().enumerate() {
let mut child = writer.start_entry();
match it {
Value::Dict(dict) if !schema.is_tuple => {
write_schema_into_builder(&mut child, schema, dict, name)?;
}
Value::Tuple(tuple_items) if schema.is_tuple => {
write_tuple_into_builder(&mut child, schema, tuple_items.as_ref(), name)?;
}
other => {
return Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: List<Schema> arg `{name}` element #{i} got {} but expects {}",
other.type_name(),
if schema.is_tuple {
"a tuple"
} else {
"a branded record"
}
),
});
}
}
writer
.finish_entry(builder, child)
.map_err(buffer_to_runtime_error)?;
}
builder
.finish_list_record(writer)
.map_err(buffer_to_runtime_error)
}
fn marshal_list_list_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
inner: &relon_eval_api::schema_canonical::TypeRepr,
items: &[Value],
) -> Result<(), RuntimeError> {
use relon_eval_api::schema_canonical::TypeRepr;
match inner {
TypeRepr::Int | TypeRepr::Float | TypeRepr::Bool => {
relon_eval_api::buffer::write_nested_scalar_list(builder, name, inner, items)
.map_err(buffer_to_runtime_error)
}
_ => relon_eval_api::buffer::write_nested_pointer_array_list(builder, name, inner, items)
.map_err(buffer_to_runtime_error),
}
}
fn marshal_schema_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
schema: &relon_eval_api::schema_canonical::Schema,
dict: &relon_eval_api::ValueDict,
) -> Result<(), RuntimeError> {
let sub_layout = relon_eval_api::layout::SchemaLayout::offsets_for(schema).map_err(|e| {
RuntimeError::Unsupported {
reason: format!("llvm-aot: schema arg `{name}` layout: {e}"),
}
})?;
let mut child = builder
.sub_record(name, &sub_layout, &schema.fields)
.map_err(buffer_to_runtime_error)?;
write_schema_into_builder(&mut child, schema, dict, name)?;
builder
.finish_sub_record(name, child)
.map_err(buffer_to_runtime_error)
}
fn marshal_tuple_in(
builder: &mut relon_eval_api::buffer::BufferBuilder<'_>,
name: &str,
schema: &relon_eval_api::schema_canonical::Schema,
items: &[Value],
) -> Result<(), RuntimeError> {
let sub_layout = relon_eval_api::layout::SchemaLayout::offsets_for(schema).map_err(|e| {
RuntimeError::Unsupported {
reason: format!("llvm-aot: tuple arg `{name}` layout: {e}"),
}
})?;
let mut child = builder
.sub_record(name, &sub_layout, &schema.fields)
.map_err(buffer_to_runtime_error)?;
write_tuple_into_builder(&mut child, schema, items, name)?;
builder
.finish_sub_record(name, child)
.map_err(buffer_to_runtime_error)
}
fn write_schema_into_builder(
child: &mut relon_eval_api::buffer::BufferBuilder<'_>,
schema: &relon_eval_api::schema_canonical::Schema,
dict: &relon_eval_api::ValueDict,
parent_field: &str,
) -> Result<(), RuntimeError> {
for sub_field in &schema.fields {
let sub_value =
dict.map
.get(sub_field.name.as_str())
.ok_or_else(|| RuntimeError::Unsupported {
reason: format!(
"llvm-aot: schema arg `{parent_field}` is missing field `{}`",
sub_field.name
),
})?;
write_value_into_builder(child, sub_field, sub_value)?;
}
Ok(())
}
fn write_tuple_into_builder(
child: &mut relon_eval_api::buffer::BufferBuilder<'_>,
schema: &relon_eval_api::schema_canonical::Schema,
items: &[Value],
parent_field: &str,
) -> Result<(), RuntimeError> {
if items.len() != schema.fields.len() {
return Err(RuntimeError::Unsupported {
reason: format!(
"llvm-aot: tuple arg `{parent_field}` has arity {} but schema expects {}",
items.len(),
schema.fields.len()
),
});
}
for (sub_field, sub_value) in schema.fields.iter().zip(items.iter()) {
write_value_into_builder(child, sub_field, sub_value)?;
}
Ok(())
}
fn return_needs_tail_region(schema: &relon_eval_api::schema_canonical::Schema) -> bool {
use relon_eval_api::schema_canonical::TypeRepr;
schema.fields.iter().any(|f| {
matches!(
f.ty,
TypeRepr::String
| TypeRepr::List { .. }
| TypeRepr::Schema { .. }
| TypeRepr::Option { .. }
| TypeRepr::Result { .. }
| TypeRepr::Enum { .. }
)
})
}
fn fast_entry_emittable(entry: &relon_ir::ir::Func) -> bool {
!body_references_const_pool(&entry.body)
}
fn body_may_raise_typed_trap(body: &[relon_ir::ir::TaggedOp]) -> bool {
use relon_ir::ir::{IrType, Op};
for tagged in body {
let hit = match &tagged.op {
Op::Add(IrType::I64)
| Op::Sub(IrType::I64)
| Op::Mul(IrType::I64)
| Op::Div(IrType::I64)
| Op::Mod(IrType::I64)
| Op::Trap { .. }
| Op::CheckCap { .. }
| Op::CallNative { .. } => true,
Op::Block { body, .. } | Op::Loop { body, .. } => body_may_raise_typed_trap(body),
Op::If {
then_body,
else_body,
..
} => body_may_raise_typed_trap(then_body) || body_may_raise_typed_trap(else_body),
Op::Call { fn_index, .. } => {
let stdlib = relon_ir::stdlib::builtin_stdlib();
stdlib
.get(*fn_index as usize)
.map(|callee| body_may_raise_typed_trap(&callee.body_owned()))
.unwrap_or(true)
}
_ => false,
};
if hit {
return true;
}
}
false
}
fn body_references_const_pool(body: &[relon_ir::ir::TaggedOp]) -> bool {
use relon_ir::ir::Op;
for tagged in body {
let hit = match &tagged.op {
Op::ConstString { .. }
| Op::ConstListInt { .. }
| Op::ConstListFloat { .. }
| Op::ConstListBool { .. }
| Op::ConstListString { .. } => true,
Op::Block { body, .. } | Op::Loop { body, .. } => body_references_const_pool(body),
Op::If {
then_body,
else_body,
..
} => body_references_const_pool(then_body) || body_references_const_pool(else_body),
Op::Call { fn_index, .. } => {
let stdlib = relon_ir::stdlib::builtin_stdlib();
stdlib
.get(*fn_index as usize)
.map(|callee| body_references_const_pool(&callee.body_owned()))
.unwrap_or(false)
}
_ => false,
};
if hit {
return true;
}
}
false
}
fn compute_effectful_imports(ir: &relon_ir::ir::Module) -> Vec<bool> {
let mut effectful = vec![false; ir.imports.len()];
for func in &ir.funcs {
scan_body_effectful(&func.body, &mut effectful);
}
effectful
}
fn scan_body_effectful(body: &[relon_ir::ir::TaggedOp], effectful: &mut [bool]) {
use relon_ir::ir::Op;
let mut pending_check_caps: u32 = 0;
for tagged in body {
match &tagged.op {
Op::CheckCap { .. } => pending_check_caps += 1,
Op::CallNative { import_idx, .. } => {
if pending_check_caps > 0 {
if let Some(slot) = effectful.get_mut(*import_idx as usize) {
*slot = true;
}
}
pending_check_caps = 0;
}
Op::Block { body, .. } | Op::Loop { body, .. } => {
scan_body_effectful(body, effectful);
}
Op::If {
then_body,
else_body,
..
} => {
scan_body_effectful(then_body, effectful);
scan_body_effectful(else_body, effectful);
}
_ => {}
}
}
}
fn build_fast_path_profile(schema: &BufferSchema) -> Result<FastPathProfile, ()> {
use relon_eval_api::schema_canonical::TypeRepr;
for f in &schema.main_schema.fields {
if !matches!(f.ty, TypeRepr::Int) {
return Err(());
}
}
if !is_single_int_field_record(&schema.return_schema) {
return Err(());
}
let mut arg_offsets: Vec<u32> = Vec::with_capacity(schema.main_layout.fields.len());
for (i, f) in schema.main_schema.fields.iter().enumerate() {
let lo = schema.main_layout.fields.get(i).ok_or(())?;
if lo.name != f.name {
return Err(());
}
arg_offsets.push(lo.offset as u32);
}
if arg_offsets.len() > 8 {
return Err(());
}
let ret_offset = schema
.return_layout
.fields
.first()
.map(|f| f.offset as u32)
.ok_or(())?;
Ok(FastPathProfile {
arg_offsets,
ret_offset,
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmittedEntryShape {
FastInt,
Buffer,
}
#[derive(Debug, Clone)]
pub struct EmittedField {
pub name: String,
pub offset: u32,
pub ty: EmittedFieldType,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmittedFieldType {
Int,
Float,
Bool,
Unit,
String,
ListInt,
}
#[derive(Debug, Clone)]
pub struct EmitObjectInfo {
pub entry_symbol: String,
pub entry_arity: usize,
pub param_names: Vec<String>,
pub shape: EmittedEntryShape,
pub main_fields: Vec<EmittedField>,
pub return_fields: Vec<EmittedField>,
pub main_root_size: u32,
pub return_root_size: u32,
pub return_has_tail: bool,
pub const_data: Vec<u8>,
pub references_str_contains_shim: bool,
}
impl LlvmAotEvaluator {
pub fn emit_object(
src: &str,
entry_symbol: &str,
out_path: &Path,
) -> Result<EmitObjectInfo, LlvmError> {
let options = relon_analyzer::AnalyzeOptions {
strict_mode: false,
..Default::default()
};
Self::emit_object_with_options(
src,
entry_symbol,
out_path,
&options,
WorldMode::OpenWorld,
None,
)
}
pub fn emit_object_with_options(
src: &str,
entry_symbol: &str,
out_path: &Path,
options: &relon_analyzer::AnalyzeOptions,
world_mode: WorldMode,
host_shim_src: Option<&str>,
) -> Result<EmitObjectInfo, LlvmError> {
Self::emit_object_for_target(
src,
entry_symbol,
out_path,
options,
world_mode,
host_shim_src,
CodegenTarget::Native,
)
}
#[allow(clippy::too_many_arguments)]
pub fn emit_object_for_target(
src: &str,
entry_symbol: &str,
out_path: &Path,
options: &relon_analyzer::AnalyzeOptions,
world_mode: WorldMode,
host_shim_src: Option<&str>,
target: CodegenTarget,
) -> Result<EmitObjectInfo, LlvmError> {
let (ir, main_schema, return_schema) = Self::lower_source_with_options(src, Some(options))?;
let main_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&main_schema)
.map_err(|e| LlvmError::Codegen(format!("main schema layout: {e}")))?;
let return_layout = relon_eval_api::layout::SchemaLayout::offsets_for(&return_schema)
.map_err(|e| LlvmError::Codegen(format!("return schema layout: {e}")))?;
let param_names: Vec<String> = main_schema.fields.iter().map(|f| f.name.clone()).collect();
let schema = BufferSchema {
main_schema,
return_schema,
main_layout,
return_layout,
};
let descriptors_strict = matches!(target, CodegenTarget::Native);
let (main_fields, return_fields) = if descriptors_strict {
(
lower_field_descriptors(&schema.main_schema, &schema.main_layout)?,
lower_field_descriptors(&schema.return_schema, &schema.return_layout)?,
)
} else {
(
lower_field_descriptors(&schema.main_schema, &schema.main_layout)
.unwrap_or_default(),
lower_field_descriptors(&schema.return_schema, &schema.return_layout)
.unwrap_or_default(),
)
};
let entry_idx = ir
.entry_func_index
.ok_or_else(|| LlvmError::Codegen("IR module has no entry function".into()))?;
let entry = &ir.funcs[entry_idx];
if !crate::codegen::is_buffer_protocol_signature(&entry.params, entry.ret) {
return Err(LlvmError::UnsupportedSignature(
"relon-rs build: lowering produced a non-buffer entry shape".into(),
));
}
let fast_profile = match world_mode {
WorldMode::ClosedWorld => None,
WorldMode::OpenWorld if !ir.imports.is_empty() => None,
WorldMode::OpenWorld => build_fast_path_profile(&schema).ok(),
};
let ctx = Context::create();
let module = ctx.create_module("relon_rs_object");
let const_pool = ConstPool::from_module(&ir)?;
let fast_profile = match fast_profile {
Some(profile) if fast_entry_emittable(entry) && ir.closure_table.is_empty() => {
Some(profile)
}
_ => None,
};
let (shape, references_str_contains_shim) = match fast_profile {
Some(ref profile) => {
let helper_table: HashMap<u32, FunctionValue<'_>> = HashMap::new();
let closure_fn_table: Vec<FunctionValue<'_>> = Vec::new();
let llvm_fn = emit_fast_entry(
&ctx,
&module,
entry,
profile,
&helper_table,
&closure_fn_table,
)?;
llvm_fn.as_global_value().set_name(entry_symbol);
llvm_fn.set_linkage(Linkage::External);
(EmittedEntryShape::FastInt, false)
}
None => {
let buffer_return_size = schema.return_layout.root_size as u32;
let lambda_ir_idx_set: std::collections::HashSet<u32> =
ir.closure_table.iter().copied().collect();
let helpers: Vec<&relon_ir::ir::Func> = ir
.funcs
.iter()
.enumerate()
.filter(|(i, _)| *i != entry_idx && !lambda_ir_idx_set.contains(&(*i as u32)))
.map(|(_, f)| f)
.collect();
let helper_ir_indices: Vec<u32> = ir
.funcs
.iter()
.enumerate()
.filter(|(i, _)| *i != entry_idx && !lambda_ir_idx_set.contains(&(*i as u32)))
.map(|(i, _)| i as u32)
.collect();
let lambdas: Vec<&relon_ir::ir::Func> = ir
.closure_table
.iter()
.map(|&ir_idx| &ir.funcs[ir_idx as usize])
.collect();
let effectful_imports = compute_effectful_imports(&ir);
let llvm_fn = match (world_mode, target) {
(WorldMode::ClosedWorld, CodegenTarget::Wasm32) => {
emit_module_funcs_closed_world_wasm(
&ctx,
&module,
entry,
buffer_return_size,
&const_pool,
&helpers,
Some(&helper_ir_indices),
&lambdas,
&ir.closure_table,
&ir.imports,
&effectful_imports,
)?
.0
}
(world_mode, target) => {
let emit = match (world_mode, target) {
(WorldMode::OpenWorld, CodegenTarget::Wasm32) => emit_module_funcs_wasm,
(WorldMode::OpenWorld, CodegenTarget::Native) => emit_module_funcs,
(WorldMode::ClosedWorld, _) => emit_module_funcs_closed_world,
};
emit(
&ctx,
&module,
entry,
buffer_return_size,
&const_pool,
&helpers,
Some(&helper_ir_indices),
&lambdas,
&ir.closure_table,
&ir.imports,
)?
.0
}
};
llvm_fn.as_global_value().set_name(entry_symbol);
llvm_fn.set_linkage(Linkage::External);
if matches!(world_mode, WorldMode::ClosedWorld) {
let shim = host_shim_src.ok_or_else(|| {
LlvmError::Codegen(
"emit_object_with_options: ClosedWorld requires a host_shim_src \
(the #[no_mangle] extern \"C\" host crate to link + inline)"
.into(),
)
})?;
match target {
CodegenTarget::Wasm32 => {
crate::cocompile::link_and_inline_host_shim_wasm_pure_only(
&module,
shim,
&ir.imports,
&effectful_imports,
)?;
}
CodegenTarget::Native => {
crate::cocompile::link_and_inline_host_shim(
&module,
shim,
&ir.imports,
)?;
}
}
}
let needs_shim = module
.get_function(RELON_LLVM_STR_CONTAINS_ARENA_SYMBOL)
.is_some()
|| module
.get_function(crate::str_helpers::RELON_LLVM_F64_TO_STR_SYMBOL)
.is_some();
(EmittedEntryShape::Buffer, needs_shim)
}
};
module.verify().map_err(|e| {
LlvmError::Codegen(format!("LLVM verifier rejected object module: {e}"))
})?;
let (machine, target_triple) = create_object_target_machine(target)?;
module.set_triple(&TargetTriple::create(&target_triple));
module.set_data_layout(&machine.get_target_data().get_data_layout());
match target {
CodegenTarget::Native => {
stamp_host_target_attributes(&module);
run_default_o3_pipeline(&module)?;
}
CodegenTarget::Wasm32 => {
let opts = PassBuilderOptions::create();
module
.run_passes("default<O3>", &machine, opts)
.map_err(|e| LlvmError::Codegen(format!("wasm32 run_passes O3: {e}")))?;
}
}
if let Some(parent) = out_path.parent() {
if !parent.as_os_str().is_empty() {
std::fs::create_dir_all(parent)
.map_err(|e| LlvmError::Codegen(format!("create out dir `{parent:?}`: {e}")))?;
}
}
machine
.write_to_file(&module, FileType::Object, out_path)
.map_err(|e| LlvmError::Codegen(format!("write object `{out_path:?}`: {e}")))?;
let entry_arity = main_fields.len();
let main_root_size = schema.main_layout.root_size as u32;
let return_root_size = schema.return_layout.root_size as u32;
let return_has_tail = return_needs_tail_region(&schema.return_schema);
let const_data = match shape {
EmittedEntryShape::FastInt => Vec::new(),
EmittedEntryShape::Buffer => const_pool.bytes,
};
let (main_fields_out, return_fields_out, main_root_size_out, return_root_size_out) =
match shape {
EmittedEntryShape::FastInt => (Vec::new(), Vec::new(), 0, 0),
EmittedEntryShape::Buffer => {
(main_fields, return_fields, main_root_size, return_root_size)
}
};
Ok(EmitObjectInfo {
entry_symbol: entry_symbol.to_string(),
entry_arity,
param_names,
shape,
main_fields: main_fields_out,
return_fields: return_fields_out,
main_root_size: main_root_size_out,
return_root_size: return_root_size_out,
return_has_tail: matches!(shape, EmittedEntryShape::Buffer) && return_has_tail,
const_data,
references_str_contains_shim,
})
}
}
fn lower_field_descriptors(
schema: &relon_eval_api::schema_canonical::Schema,
layout: &relon_eval_api::layout::OffsetTable,
) -> Result<Vec<EmittedField>, LlvmError> {
let mut out = Vec::with_capacity(schema.fields.len());
for (i, f) in schema.fields.iter().enumerate() {
let lo = layout.fields.get(i).ok_or_else(|| {
LlvmError::Codegen(format!(
"lower_field_descriptors: layout missing slot for field `{}`",
f.name
))
})?;
if lo.name != f.name {
return Err(LlvmError::Codegen(format!(
"lower_field_descriptors: schema/layout name mismatch at slot {i}: schema=`{}`, layout=`{}`",
f.name, lo.name
)));
}
let ty = emitted_field_type_for(&f.ty).ok_or_else(|| {
LlvmError::UnsupportedSignature(format!(
"relon-rs build (Phase 2): field `{}` type {:?} not yet wired for marshalling",
f.name, f.ty
))
})?;
out.push(EmittedField {
name: f.name.clone(),
offset: lo.offset as u32,
ty,
});
}
Ok(out)
}
fn emitted_field_type_for(
ty: &relon_eval_api::schema_canonical::TypeRepr,
) -> Option<EmittedFieldType> {
use relon_eval_api::schema_canonical::TypeRepr;
match ty {
TypeRepr::Int => Some(EmittedFieldType::Int),
TypeRepr::Float => Some(EmittedFieldType::Float),
TypeRepr::Bool => Some(EmittedFieldType::Bool),
TypeRepr::Unit => Some(EmittedFieldType::Unit),
TypeRepr::String => Some(EmittedFieldType::String),
TypeRepr::List { element } if matches!(element.as_ref(), TypeRepr::Int) => {
Some(EmittedFieldType::ListInt)
}
_ => None,
}
}
fn stamp_host_target_attributes(module: &inkwell::module::Module<'_>) {
let cpu = TargetMachine::get_host_cpu_name();
let features = TargetMachine::get_host_cpu_features();
let cpu = cpu.to_str().unwrap_or("");
let features = features.to_str().unwrap_or("");
if cpu.is_empty() {
return;
}
let ctx = module.get_context();
let cpu_attr = ctx.create_string_attribute("target-cpu", cpu);
let features_attr = ctx.create_string_attribute("target-features", features);
let mut func = module.get_first_function();
while let Some(f) = func {
if f.count_basic_blocks() > 0 {
f.remove_string_attribute(inkwell::attributes::AttributeLoc::Function, "target-cpu");
f.remove_string_attribute(
inkwell::attributes::AttributeLoc::Function,
"target-features",
);
f.add_attribute(inkwell::attributes::AttributeLoc::Function, cpu_attr);
f.add_attribute(inkwell::attributes::AttributeLoc::Function, features_attr);
}
func = f.get_next_function();
}
}
fn run_default_o3_pipeline(module: &inkwell::module::Module<'_>) -> Result<(), LlvmError> {
Target::initialize_native(&InitializationConfig::default())
.map_err(|e| LlvmError::Codegen(format!("initialize_native: {e}")))?;
let triple_str = TargetMachine::get_default_triple();
let target = Target::from_triple(&triple_str)
.map_err(|e| LlvmError::Codegen(format!("target from_triple: {e}")))?;
let cpu = TargetMachine::get_host_cpu_name();
let features = TargetMachine::get_host_cpu_features();
let triple = TargetTriple::create(
triple_str
.as_str()
.to_str()
.map_err(|e| LlvmError::Codegen(format!("triple utf8: {e}")))?,
);
let machine = target
.create_target_machine(
&triple,
cpu.to_str().unwrap_or(""),
features.to_str().unwrap_or(""),
OptimizationLevel::Aggressive,
RelocMode::Default,
CodeModel::JITDefault,
)
.ok_or_else(|| LlvmError::Codegen("create_target_machine returned null".into()))?;
let opts = PassBuilderOptions::create();
module
.run_passes("default<O3>", &machine, opts)
.map_err(|e| LlvmError::Codegen(format!("run_passes O3: {e}")))?;
Ok(())
}
fn create_object_target_machine(
target: CodegenTarget,
) -> Result<(TargetMachine, String), LlvmError> {
match target {
CodegenTarget::Native => {
Target::initialize_native(&InitializationConfig::default())
.map_err(|e| LlvmError::Codegen(format!("initialize_native: {e}")))?;
let triple_str = TargetMachine::get_default_triple();
let t = Target::from_triple(&triple_str)
.map_err(|e| LlvmError::Codegen(format!("target from_triple: {e}")))?;
let cpu = TargetMachine::get_host_cpu_name();
let features = TargetMachine::get_host_cpu_features();
let triple = TargetTriple::create(
triple_str
.as_str()
.to_str()
.map_err(|e| LlvmError::Codegen(format!("triple utf8: {e}")))?,
);
let machine = t
.create_target_machine(
&triple,
cpu.to_str().unwrap_or(""),
features.to_str().unwrap_or(""),
OptimizationLevel::Aggressive,
RelocMode::PIC,
CodeModel::Default,
)
.ok_or_else(|| LlvmError::Codegen("create_target_machine returned null".into()))?;
let triple_owned = triple_str
.as_str()
.to_str()
.map_err(|e| LlvmError::Codegen(format!("triple utf8: {e}")))?
.to_string();
Ok((machine, triple_owned))
}
CodegenTarget::Wasm32 => {
Target::initialize_webassembly(&InitializationConfig::default());
let triple = TargetTriple::create(WASM32_TRIPLE);
let t = Target::from_triple(&triple)
.map_err(|e| LlvmError::Codegen(format!("wasm32 target from_triple: {e}")))?;
let machine = t
.create_target_machine(
&triple,
"",
"+bulk-memory",
OptimizationLevel::Aggressive,
RelocMode::Static,
CodeModel::Default,
)
.ok_or_else(|| {
LlvmError::Codegen("wasm32 create_target_machine returned null".into())
})?;
Ok((machine, WASM32_TRIPLE.to_string()))
}
}
}