wowlab_engine_domain/rotation/
compiler.rs

1// #t(file: rust_pub_api_docs) thin entry shell around the generic lowerer
2// #t(file: rust_panic_in_result_fn, rust_unwrap_in_lib) inkwell builder methods return Result; unwrap is safe for well-formed IR
3
4//! LLVM JIT shell: builds the function prototype and drives the generic lowerer.
5
6use std::sync::Once;
7
8use inkwell::{
9    AddressSpace, OptimizationLevel,
10    context::Context,
11    execution_engine::ExecutionEngine,
12    targets::{InitializationConfig, Target},
13};
14use ouroboros::self_referencing;
15
16use super::{
17    buffer::DenseBuffer,
18    context::ContextSchema,
19    decision_trace::DecisionTraceTarget,
20    error::{Error, Result},
21    jit_backend::{Intrinsics, JitBackend, JitEntry, JitEnvironment},
22    jit_common::{RotationFn, SyncFnPtr},
23    prepared::PreparedRotation,
24    reference,
25    result::EvalResult,
26};
27
28const ROTATION_FN_NAME: &str = "rotation";
29
30const ROTATION_MODULE_NAME: &str = "rotation";
31
32#[self_referencing]
33struct LlvmBundle {
34    context: Context,
35    #[borrows(context)]
36    #[covariant]
37    execution_engine: ExecutionEngine<'this>,
38}
39
40pub struct CompiledRotation {
41    func_ptr: SyncFnPtr,
42    prepared: PreparedRotation,
43    trace: Option<DecisionTraceTarget>,
44    // Field-drop order matters: bundle drops last so JIT code outlives in-flight calls.
45    _bundle: LlvmBundle,
46}
47
48impl CompiledRotation {
49    pub(super) fn compile(prepared: PreparedRotation) -> Result<Self> {
50        compile_inner(prepared)
51    }
52
53    pub(super) fn evaluate(&mut self, buffer: &mut DenseBuffer, now_secs: f64) -> EvalResult {
54        if let Some(sink) = &self.trace {
55            return reference::evaluate_trace(&self.prepared, buffer, now_secs, sink.clone());
56        }
57
58        // SAFETY: `_bundle` owns the ExecutionEngine that keeps `func_ptr` valid.
59
60        let packed = unsafe { (self.func_ptr.0)(buffer.as_mut_ptr(), now_secs) };
61        let (kind, spell_id, payload) = wowlab_buffer_contract::decode_eval_result(packed);
62        let empower_rank = if matches!(
63            kind,
64            super::result::KIND_CAST | super::result::KIND_USE_ITEM
65        ) {
66            u8::try_from(payload).unwrap_or(0)
67        } else {
68            0
69        };
70
71        EvalResult {
72            kind,
73            empower_rank,
74            spell_id,
75            wait_time: if matches!(kind, super::result::KIND_WAIT | super::result::KIND_POOL) {
76                f32::from_bits(payload)
77            } else {
78                0.0
79            },
80        }
81    }
82
83    pub(super) fn schema(&self) -> &ContextSchema {
84        self.prepared.schema()
85    }
86
87    pub(super) const fn prepared(&self) -> &PreparedRotation {
88        &self.prepared
89    }
90
91    pub(super) fn set_trace(&mut self, sink: DecisionTraceTarget) {
92        self.trace = Some(sink);
93    }
94
95    pub(super) const fn has_trace(&self) -> bool {
96        self.trace.is_some()
97    }
98}
99
100impl std::fmt::Debug for CompiledRotation {
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.debug_struct("CompiledRotation").finish_non_exhaustive()
103    }
104}
105
106// SAFETY: bundle is owned by one CompiledRotation; evaluate takes &mut self, so the
107// non-atomic Rc in ExecutionEngine cannot be accessed concurrently.
108unsafe impl Send for CompiledRotation {}
109
110static INIT_NATIVE: Once = Once::new();
111
112fn ensure_native_target() {
113    INIT_NATIVE.call_once(|| {
114        Target::initialize_native(&InitializationConfig::default())
115            .expect("failed to initialize native LLVM target");
116        ExecutionEngine::link_in_mc_jit();
117    });
118}
119
120fn compile_inner(prepared: PreparedRotation) -> Result<CompiledRotation> {
121    ensure_native_target();
122
123    let mut captured_fn: Option<RotationFn> = None;
124    let bundle = LlvmBundle::try_new::<Error>(Context::create(), |context| {
125        let module = context.create_module(ROTATION_MODULE_NAME);
126        let builder = context.create_builder();
127
128        let i64_type = context.i64_type();
129        let f64_type = context.f64_type();
130        let ptr_type = context.ptr_type(AddressSpace::default());
131
132        let fn_type = i64_type.fn_type(&[ptr_type.into(), f64_type.into()], false);
133        let function = module.add_function(ROTATION_FN_NAME, fn_type, None);
134
135        let intrinsics = Intrinsics::declare(context, &module);
136
137        let entry_block = context.append_basic_block(function, "entry");
138
139        builder.position_at_end(entry_block);
140
141        let buf_param = function.get_nth_param(0).unwrap().into_pointer_value();
142        let now_param = function.get_nth_param(1).unwrap().into_float_value();
143
144        let environment = JitEnvironment {
145            context,
146            builder: &builder,
147        };
148        let entry = JitEntry {
149            function,
150            buffer: buf_param,
151            now: now_param,
152        };
153        let mut backend = JitBackend::new(&environment, &entry, intrinsics);
154
155        prepared.lower(&mut backend);
156
157        if let Err(msg) = module.verify() {
158            let ir = module.print_to_string().to_string();
159
160            return Err(Error::compilation(format!(
161                "LLVM module verification failed: {}\nIR:\n{}",
162                msg.to_string(),
163                ir,
164            )));
165        }
166
167        let ee = module
168            .create_jit_execution_engine(OptimizationLevel::Aggressive)
169            .map_err(|error| {
170                Error::compilation(format!("failed to create LLVM JIT engine: {error}"))
171            })?;
172
173        // SAFETY: signature matches RotationFn; EE moves into the bundle below.
174        let jit_fn = unsafe { ee.get_function::<RotationFn>(ROTATION_FN_NAME).unwrap() };
175        // SAFETY: same as above; bundle keeps the code alive past this scope.
176
177        captured_fn = Some(unsafe { jit_fn.as_raw() });
178
179        Ok(ee)
180    })?;
181
182    let raw_fn =
183        captured_fn.ok_or_else(|| Error::compilation("JIT function pointer not captured"))?;
184
185    Ok(CompiledRotation {
186        func_ptr: SyncFnPtr(raw_fn),
187        prepared,
188        trace: None,
189        _bundle: bundle,
190    })
191}
wowlab_engine_domain/rotation/compiler.rs

wowlab_engine_domain/rotation/
compiler.rs
