wowlab_engine_domain/rotation/

compiler.rs

// #t(file: rust_pub_api_docs) thin entry shell around the generic lowerer
// #t(file: rust_panic_in_result_fn, rust_unwrap_in_lib) inkwell builder methods return Result; unwrap is safe for well-formed IR

//! LLVM JIT shell: builds the function prototype and drives the generic lowerer.

use std::sync::Once;

use inkwell::{
    AddressSpace, OptimizationLevel,
    context::Context,
    execution_engine::ExecutionEngine,
    targets::{InitializationConfig, Target},
};
use ouroboros::self_referencing;
use wowlab_types::sim::Rotation;

use super::{
    backend::RuntimeBackend,
    buffer::DenseBuffer,
    context::ContextSchema,
    error::{Error, Result},
    jit_backend::{Intrinsics, JitBackend, JitEntry, JitEnvironment},
    jit_common::{RotationFn, SyncFnPtr},
    lower,
    resolver::SpecResolver,
    result::EvalResult,
};

const ROTATION_FN_NAME: &str = "rotation";

const ROTATION_MODULE_NAME: &str = "rotation";

#[self_referencing]
struct LlvmBundle {
    context: Context,
    #[borrows(context)]
    #[covariant]
    execution_engine: ExecutionEngine<'this>,
}

pub struct CompiledRotation {
    func_ptr: SyncFnPtr,
    schema: ContextSchema,
    pub(crate) rotation: Rotation,
    pub(crate) resolver: SpecResolver,
    // Field-drop order matters: bundle drops last so JIT code outlives in-flight calls.
    _bundle: LlvmBundle,
}

impl CompiledRotation {
    /// Parses, validates, and JIT-compiles a rotation.
    ///
    /// # Errors
    /// Returns an error when parsing, lowering, or native compilation fails.
    pub fn compile_json(json: &str, resolver: &SpecResolver) -> Result<Self> {
        let rotation = super::ast::parse_and_validate(json)?;

        Self::compile_resolved(&rotation, resolver)
    }

    /// JIT-compiles an already parsed rotation.
    ///
    /// # Errors
    /// Returns an error when lowering or native compilation fails.
    pub fn compile(rotation: &Rotation, resolver: &SpecResolver) -> Result<Self> {
        Self::compile_resolved(rotation, resolver)
    }

    fn compile_resolved(rotation: &Rotation, resolver: &SpecResolver) -> Result<Self> {
        compile_inner(rotation, resolver)
    }
}

impl std::fmt::Debug for CompiledRotation {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CompiledRotation").finish_non_exhaustive()
    }
}

// SAFETY: bundle is owned by one CompiledRotation; evaluate takes &mut self, so the
// non-atomic Rc in ExecutionEngine cannot be accessed concurrently.
unsafe impl Send for CompiledRotation {}

impl RuntimeBackend for CompiledRotation {
    fn evaluate(&mut self, buffer: &mut DenseBuffer, now_secs: f64) -> EvalResult {
        // SAFETY: `_bundle` owns the ExecutionEngine that keeps `func_ptr` valid.
        let packed = unsafe { (self.func_ptr.0)(buffer.as_mut_ptr(), now_secs) };
        let (kind, spell_id, payload) = wowlab_buffer_contract::decode_eval_result(packed);
        let empower_rank = if matches!(
            kind,
            super::result::KIND_CAST | super::result::KIND_USE_ITEM
        ) {
            u8::try_from(payload).unwrap_or(0)
        } else {
            0
        };

        EvalResult {
            kind,
            empower_rank,
            spell_id,
            wait_time: if matches!(kind, super::result::KIND_WAIT | super::result::KIND_POOL) {
                f32::from_bits(payload)
            } else {
                0.0
            },
        }
    }

    fn schema(&self) -> &ContextSchema {
        &self.schema
    }
}

static INIT_NATIVE: Once = Once::new();

fn ensure_native_target() {
    INIT_NATIVE.call_once(|| {
        Target::initialize_native(&InitializationConfig::default())
            .expect("failed to initialize native LLVM target");
        ExecutionEngine::link_in_mc_jit();
    });
}

fn compile_inner(rotation: &Rotation, resolver: &SpecResolver) -> Result<CompiledRotation> {
    ensure_native_target();

    let (schema, table) = lower::prepare(rotation, resolver)?;

    let mut captured_fn: Option<RotationFn> = None;
    let bundle = LlvmBundle::try_new::<Error>(Context::create(), |context| {
        let module = context.create_module(ROTATION_MODULE_NAME);
        let builder = context.create_builder();

        let i64_type = context.i64_type();
        let f64_type = context.f64_type();
        let ptr_type = context.ptr_type(AddressSpace::default());

        let fn_type = i64_type.fn_type(&[ptr_type.into(), f64_type.into()], false);
        let function = module.add_function(ROTATION_FN_NAME, fn_type, None);

        let intrinsics = Intrinsics::declare(context, &module);

        let entry_block = context.append_basic_block(function, "entry");

        builder.position_at_end(entry_block);

        let buf_param = function.get_nth_param(0).unwrap().into_pointer_value();
        let now_param = function.get_nth_param(1).unwrap().into_float_value();

        let environment = JitEnvironment {
            context,
            builder: &builder,
        };
        let entry = JitEntry {
            function,
            buffer: buf_param,
            now: now_param,
        };
        let mut backend = JitBackend::new(&environment, &entry, intrinsics);

        lower::lower_rotation(&mut backend, rotation, &schema, resolver, &table);

        if let Err(msg) = module.verify() {
            let ir = module.print_to_string().to_string();

            return Err(Error::compilation(format!(
                "LLVM module verification failed: {}\nIR:\n{}",
                msg.to_string(),
                ir,
            )));
        }

        let ee = module
            .create_jit_execution_engine(OptimizationLevel::Aggressive)
            .map_err(|error| {
                Error::compilation(format!("failed to create LLVM JIT engine: {error}"))
            })?;

        // SAFETY: signature matches RotationFn; EE moves into the bundle below.
        let jit_fn = unsafe { ee.get_function::<RotationFn>(ROTATION_FN_NAME).unwrap() };
        // SAFETY: same as above; bundle keeps the code alive past this scope.

        captured_fn = Some(unsafe { jit_fn.as_raw() });

        Ok(ee)
    })?;

    let raw_fn =
        captured_fn.ok_or_else(|| Error::compilation("JIT function pointer not captured"))?;

    Ok(CompiledRotation {
        func_ptr: SyncFnPtr(raw_fn),
        schema,
        rotation: rotation.clone(),
        resolver: resolver.clone(),
        _bundle: bundle,
    })
}