scirs2_core/array_protocol/
jit_impl.rs

1// Copyright (c) 2025, `SciRS2` Team
2//
3// Licensed under the Apache License, Version 2.0
4// (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
5//
6
7//! Just-In-Time (JIT) compilation support for array operations.
8//!
9//! This module provides functionality for JIT-compiling operations on arrays,
10//! allowing for faster execution of custom operations.
11
12use std::any::{Any, TypeId};
13use std::collections::HashMap;
14use std::fmt::Debug;
15use std::marker::PhantomData;
16use std::sync::{Arc, LazyLock, RwLock};
17
18use crate::array_protocol::{
19    ArrayFunction, ArrayProtocol, JITArray, JITFunction, JITFunctionFactory,
20};
21use crate::error::{CoreError, CoreResult, ErrorContext};
22
23/// JIT compilation backends
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum JITBackend {
26    /// LLVM backend
27    LLVM,
28
29    /// Cranelift backend
30    Cranelift,
31
32    /// WebAssembly backend
33    WASM,
34
35    /// Custom backend
36    Custom(TypeId),
37}
38
39impl Default for JITBackend {
40    fn default() -> Self {
41        Self::LLVM
42    }
43}
44
45/// Configuration for JIT compilation
46#[derive(Debug, Clone)]
47pub struct JITConfig {
48    /// The JIT backend to use
49    pub backend: JITBackend,
50
51    /// Whether to optimize the generated code
52    pub optimize: bool,
53
54    /// Optimization level (0-3)
55    pub opt_level: usize,
56
57    /// Whether to cache compiled functions
58    pub use_cache: bool,
59
60    /// Additional backend-specific options
61    pub backend_options: HashMap<String, String>,
62}
63
64impl Default for JITConfig {
65    fn default() -> Self {
66        Self {
67            backend: JITBackend::default(),
68            optimize: true,
69            opt_level: 2,
70            use_cache: true,
71            backend_options: HashMap::new(),
72        }
73    }
74}
75
76/// Type alias for the complex function type
77pub type JITFunctionType = dyn Fn(&[Box<dyn Any>]) -> CoreResult<Box<dyn Any>> + Send + Sync;
78
79/// A compiled JIT function
80pub struct JITFunctionImpl {
81    /// The source code of the function
82    source: String,
83
84    /// The compiled function. `Arc` (rather than `Box`) so `clone_box` can
85    /// produce a function that is byte-for-byte identical in behavior by
86    /// sharing the same underlying closure, instead of trying to
87    /// reconstruct an equivalent one from scratch.
88    function: Arc<JITFunctionType>,
89
90    /// Information about the compilation
91    compile_info: HashMap<String, String>,
92}
93
94impl Debug for JITFunctionImpl {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("JITFunctionImpl")
97            .field("source", &self.source)
98            .field("compile_info", &self.compile_info)
99            .finish_non_exhaustive()
100    }
101}
102
103impl JITFunctionImpl {
104    /// Create a new JIT function.
105    #[must_use]
106    pub fn new(
107        source: String,
108        function: Box<JITFunctionType>,
109        compile_info: HashMap<String, String>,
110    ) -> Self {
111        Self {
112            source,
113            function: Arc::from(function),
114            compile_info,
115        }
116    }
117}
118
119impl JITFunction for JITFunctionImpl {
120    fn evaluate(&self, args: &[Box<dyn Any>]) -> CoreResult<Box<dyn Any>> {
121        (self.function)(args)
122    }
123
124    fn source(&self) -> String {
125        self.source.clone()
126    }
127
128    fn compile_info(&self) -> HashMap<String, String> {
129        self.compile_info.clone()
130    }
131
132    fn clone_box(&self) -> Box<dyn JITFunction> {
133        Box::new(Self {
134            source: self.source.clone(),
135            function: Arc::clone(&self.function),
136            compile_info: self.compile_info.clone(),
137        })
138    }
139}
140
141// ============================================================================
142// A minimal expression parser + interpreter.
143//
144// Neither an LLVM nor a Cranelift toolchain is wired into this crate (Pure
145// Rust policy, and pulling either in as a hard dependency of `scirs2-core`
146// would be a large, separate undertaking), so `LLVMFunctionFactory` and
147// `CraneliftFunctionFactory` below both fall back to interpreting the
148// requested expression directly rather than JIT-compiling native code for
149// it. This is a real (if unoptimized) implementation of "evaluate this
150// expression for the given arguments" — not a stand-in that ignores its
151// input — for a small but useful arithmetic grammar:
152//
153//   expr   := term (('+' | '-') term)*
154//   term   := unary (('*' | '/') unary)*
155//   unary  := '-' unary | power
156//   power  := primary ('^' unary)?          (right-associative)
157//   primary:= number | ident ('(' expr (',' expr)* ')')? | '(' expr ')'
158//
159// with `sin/cos/tan/exp/ln/sqrt/abs` (1 argument) and `min/max` (2
160// arguments) as built-in function calls.
161//
162// `JITFunction::evaluate`'s `args: &[Box<dyn Any>]` are `f64` values bound
163// positionally to the expression's free variables in order of first
164// occurrence (left to right) — e.g. for `"x + y"`, `args[0]` is `x` and
165// `args[1]` is `y`.
166// ============================================================================
167
168/// An expression AST node.
169#[derive(Debug, Clone)]
170enum JITExpr {
171    Number(f64),
172    Var(String),
173    Neg(Box<JITExpr>),
174    Add(Box<JITExpr>, Box<JITExpr>),
175    Sub(Box<JITExpr>, Box<JITExpr>),
176    Mul(Box<JITExpr>, Box<JITExpr>),
177    Div(Box<JITExpr>, Box<JITExpr>),
178    Pow(Box<JITExpr>, Box<JITExpr>),
179    Call(String, Vec<JITExpr>),
180}
181
182#[derive(Debug, Clone, PartialEq)]
183enum JITToken {
184    Number(f64),
185    Ident(String),
186    Plus,
187    Minus,
188    Star,
189    Slash,
190    Caret,
191    LParen,
192    RParen,
193    Comma,
194}
195
196fn jit_tokenize(input: &str) -> Result<Vec<JITToken>, String> {
197    let chars: Vec<char> = input.chars().collect();
198    let mut tokens = Vec::new();
199    let mut i = 0;
200    while i < chars.len() {
201        let c = chars[i];
202        if c.is_whitespace() {
203            i += 1;
204        } else if c.is_ascii_digit()
205            || (c == '.' && chars.get(i + 1).is_some_and(char::is_ascii_digit))
206        {
207            let start = i;
208            while i < chars.len() && (chars[i].is_ascii_digit() || chars[i] == '.') {
209                i += 1;
210            }
211            // Support a trailing exponent, e.g. "1e-3".
212            if i < chars.len() && (chars[i] == 'e' || chars[i] == 'E') {
213                let mut j = i + 1;
214                if j < chars.len() && (chars[j] == '+' || chars[j] == '-') {
215                    j += 1;
216                }
217                if j < chars.len() && chars[j].is_ascii_digit() {
218                    i = j;
219                    while i < chars.len() && chars[i].is_ascii_digit() {
220                        i += 1;
221                    }
222                }
223            }
224            let text: String = chars[start..i].iter().collect();
225            let value: f64 = text
226                .parse()
227                .map_err(|_| format!("invalid number literal: '{text}'"))?;
228            tokens.push(JITToken::Number(value));
229        } else if c.is_alphabetic() || c == '_' {
230            let start = i;
231            while i < chars.len() && (chars[i].is_alphanumeric() || chars[i] == '_') {
232                i += 1;
233            }
234            tokens.push(JITToken::Ident(chars[start..i].iter().collect()));
235        } else {
236            let token = match c {
237                '+' => JITToken::Plus,
238                '-' => JITToken::Minus,
239                '*' => JITToken::Star,
240                '/' => JITToken::Slash,
241                '^' => JITToken::Caret,
242                '(' => JITToken::LParen,
243                ')' => JITToken::RParen,
244                ',' => JITToken::Comma,
245                other => return Err(format!("unexpected character: '{other}'")),
246            };
247            tokens.push(token);
248            i += 1;
249        }
250    }
251    Ok(tokens)
252}
253
254struct JITParser {
255    tokens: Vec<JITToken>,
256    pos: usize,
257}
258
259impl JITParser {
260    fn peek(&self) -> Option<&JITToken> {
261        self.tokens.get(self.pos)
262    }
263
264    fn advance(&mut self) -> Option<JITToken> {
265        let tok = self.tokens.get(self.pos).cloned();
266        if tok.is_some() {
267            self.pos += 1;
268        }
269        tok
270    }
271
272    fn expect(&mut self, expected: &JITToken) -> Result<(), String> {
273        match self.advance() {
274            Some(ref tok) if tok == expected => Ok(()),
275            Some(other) => Err(format!("expected {expected:?}, found {other:?}")),
276            None => Err(format!("expected {expected:?}, found end of expression")),
277        }
278    }
279
280    fn parse_expr(&mut self) -> Result<JITExpr, String> {
281        let mut node = self.parse_term()?;
282        loop {
283            match self.peek() {
284                Some(JITToken::Plus) => {
285                    self.advance();
286                    node = JITExpr::Add(Box::new(node), Box::new(self.parse_term()?));
287                }
288                Some(JITToken::Minus) => {
289                    self.advance();
290                    node = JITExpr::Sub(Box::new(node), Box::new(self.parse_term()?));
291                }
292                _ => break,
293            }
294        }
295        Ok(node)
296    }
297
298    fn parse_term(&mut self) -> Result<JITExpr, String> {
299        let mut node = self.parse_unary()?;
300        loop {
301            match self.peek() {
302                Some(JITToken::Star) => {
303                    self.advance();
304                    node = JITExpr::Mul(Box::new(node), Box::new(self.parse_unary()?));
305                }
306                Some(JITToken::Slash) => {
307                    self.advance();
308                    node = JITExpr::Div(Box::new(node), Box::new(self.parse_unary()?));
309                }
310                _ => break,
311            }
312        }
313        Ok(node)
314    }
315
316    fn parse_unary(&mut self) -> Result<JITExpr, String> {
317        if matches!(self.peek(), Some(JITToken::Minus)) {
318            self.advance();
319            Ok(JITExpr::Neg(Box::new(self.parse_unary()?)))
320        } else {
321            self.parse_power()
322        }
323    }
324
325    fn parse_power(&mut self) -> Result<JITExpr, String> {
326        let base = self.parse_primary()?;
327        if matches!(self.peek(), Some(JITToken::Caret)) {
328            self.advance();
329            let exponent = self.parse_unary()?;
330            Ok(JITExpr::Pow(Box::new(base), Box::new(exponent)))
331        } else {
332            Ok(base)
333        }
334    }
335
336    fn parse_primary(&mut self) -> Result<JITExpr, String> {
337        match self.advance() {
338            Some(JITToken::Number(n)) => Ok(JITExpr::Number(n)),
339            Some(JITToken::Ident(name)) => {
340                if matches!(self.peek(), Some(JITToken::LParen)) {
341                    self.advance();
342                    let mut args = Vec::new();
343                    if !matches!(self.peek(), Some(JITToken::RParen)) {
344                        args.push(self.parse_expr()?);
345                        while matches!(self.peek(), Some(JITToken::Comma)) {
346                            self.advance();
347                            args.push(self.parse_expr()?);
348                        }
349                    }
350                    self.expect(&JITToken::RParen)?;
351                    Ok(JITExpr::Call(name, args))
352                } else {
353                    Ok(JITExpr::Var(name))
354                }
355            }
356            Some(JITToken::LParen) => {
357                let inner = self.parse_expr()?;
358                self.expect(&JITToken::RParen)?;
359                Ok(inner)
360            }
361            Some(other) => Err(format!("unexpected token: {other:?}")),
362            None => Err("unexpected end of expression".to_string()),
363        }
364    }
365}
366
367/// Parses `expression` into an AST, returning an error string describing
368/// the syntax problem on failure (never panics on malformed input).
369fn jit_parse_expression(expression: &str) -> Result<JITExpr, String> {
370    let tokens = jit_tokenize(expression)?;
371    let mut parser = JITParser { tokens, pos: 0 };
372    let expr = parser.parse_expr()?;
373    if parser.pos != parser.tokens.len() {
374        return Err(format!(
375            "unexpected trailing token: {:?}",
376            parser.tokens[parser.pos]
377        ));
378    }
379    Ok(expr)
380}
381
382/// Collects the expression's free variables in order of first (left-to-right)
383/// occurrence — this fixes the positional order `evaluate`'s `args` are
384/// bound in.
385fn jit_collect_vars(
386    expr: &JITExpr,
387    order: &mut Vec<String>,
388    seen: &mut std::collections::HashSet<String>,
389) {
390    match expr {
391        JITExpr::Number(_) => {}
392        JITExpr::Var(name) => {
393            if seen.insert(name.clone()) {
394                order.push(name.clone());
395            }
396        }
397        JITExpr::Neg(inner) => jit_collect_vars(inner, order, seen),
398        JITExpr::Add(a, b)
399        | JITExpr::Sub(a, b)
400        | JITExpr::Mul(a, b)
401        | JITExpr::Div(a, b)
402        | JITExpr::Pow(a, b) => {
403            jit_collect_vars(a, order, seen);
404            jit_collect_vars(b, order, seen);
405        }
406        JITExpr::Call(_, args) => {
407            for a in args {
408                jit_collect_vars(a, order, seen);
409            }
410        }
411    }
412}
413
414/// Evaluates `expr` given a binding of variable name to value.
415fn jit_eval(expr: &JITExpr, vars: &HashMap<String, f64>) -> Result<f64, String> {
416    match expr {
417        JITExpr::Number(n) => Ok(*n),
418        JITExpr::Var(name) => vars
419            .get(name)
420            .copied()
421            .ok_or_else(|| format!("undefined variable: '{name}'")),
422        JITExpr::Neg(inner) => Ok(-jit_eval(inner, vars)?),
423        JITExpr::Add(a, b) => Ok(jit_eval(a, vars)? + jit_eval(b, vars)?),
424        JITExpr::Sub(a, b) => Ok(jit_eval(a, vars)? - jit_eval(b, vars)?),
425        JITExpr::Mul(a, b) => Ok(jit_eval(a, vars)? * jit_eval(b, vars)?),
426        JITExpr::Div(a, b) => Ok(jit_eval(a, vars)? / jit_eval(b, vars)?),
427        JITExpr::Pow(a, b) => Ok(jit_eval(a, vars)?.powf(jit_eval(b, vars)?)),
428        JITExpr::Call(name, argexprs) => {
429            let values = argexprs
430                .iter()
431                .map(|e| jit_eval(e, vars))
432                .collect::<Result<Vec<f64>, String>>()?;
433            match (name.as_str(), values.as_slice()) {
434                ("sin", [x]) => Ok(x.sin()),
435                ("cos", [x]) => Ok(x.cos()),
436                ("tan", [x]) => Ok(x.tan()),
437                ("exp", [x]) => Ok(x.exp()),
438                ("ln", [x]) => Ok(x.ln()),
439                ("sqrt", [x]) => Ok(x.sqrt()),
440                ("abs", [x]) => Ok(x.abs()),
441                ("min", [a, b]) => Ok(a.min(*b)),
442                ("max", [a, b]) => Ok(a.max(*b)),
443                (other, vals) => Err(format!(
444                    "unknown function or wrong number of arguments: {other}({n} args)",
445                    n = vals.len()
446                )),
447            }
448        }
449    }
450}
451
452/// Parses `expression` and builds a real (interpreted) [`JITFunctionType`]
453/// closure for it, used by both [`LLVMFunctionFactory`] and
454/// [`CraneliftFunctionFactory`] as their fallback "JIT" backend (see the
455/// module-level comment above). Returns a descriptive [`CoreError::JITError`]
456/// for anything the small grammar doesn't cover, rather than silently
457/// returning a placeholder value.
458fn jit_make_interpreted_function(expression: &str) -> CoreResult<Box<JITFunctionType>> {
459    let expr = jit_parse_expression(expression).map_err(|e| {
460        CoreError::JITError(ErrorContext::new(format!(
461            "failed to parse expression '{expression}': {e}"
462        )))
463    })?;
464
465    let mut order = Vec::new();
466    let mut seen = std::collections::HashSet::new();
467    jit_collect_vars(&expr, &mut order, &mut seen);
468
469    let function: Box<JITFunctionType> = Box::new(move |args: &[Box<dyn Any>]| {
470        if args.len() < order.len() {
471            return Err(CoreError::JITError(ErrorContext::new(format!(
472                "expression requires {need} argument(s) ({vars}) but only {got} were provided",
473                need = order.len(),
474                vars = order.join(", "),
475                got = args.len()
476            ))));
477        }
478        let mut vars = HashMap::with_capacity(order.len());
479        for (name, arg) in order.iter().zip(args.iter()) {
480            let value = arg.downcast_ref::<f64>().copied().ok_or_else(|| {
481                CoreError::JITError(ErrorContext::new(format!(
482                    "argument for variable '{name}' must be an f64"
483                )))
484            })?;
485            vars.insert(name.clone(), value);
486        }
487        let result = jit_eval(&expr, &vars).map_err(|e| {
488            CoreError::JITError(ErrorContext::new(format!("evaluation error: {e}")))
489        })?;
490        Ok(Box::new(result) as Box<dyn Any>)
491    });
492
493    Ok(function)
494}
495
496/// A factory for creating JIT functions using the LLVM backend
497pub struct LLVMFunctionFactory {
498    /// Configuration for JIT compilation
499    config: JITConfig,
500
501    /// Cache of compiled functions
502    cache: HashMap<String, Arc<dyn JITFunction>>,
503}
504
505impl LLVMFunctionFactory {
506    /// Create a new LLVM function factory.
507    pub fn new(config: JITConfig) -> Self {
508        Self {
509            config,
510            cache: HashMap::new(),
511        }
512    }
513
514    /// Compile a function for the given expression.
515    ///
516    /// No LLVM toolchain is linked into this crate (Pure Rust policy), so
517    /// this interprets the expression directly instead — a real
518    /// implementation of "evaluate this expression" for the grammar
519    /// documented on [`jit_make_interpreted_function`], not a placeholder.
520    fn compile(&self, expression: &str, array_typeid: TypeId) -> CoreResult<Arc<dyn JITFunction>> {
521        let mut compile_info = HashMap::new();
522        compile_info.insert("backend".to_string(), "LLVM".to_string());
523        compile_info.insert("opt_level".to_string(), self.config.opt_level.to_string());
524        compile_info.insert("array_type".to_string(), format!("{array_typeid:?}"));
525
526        let source = expression.to_string();
527        let function = jit_make_interpreted_function(expression)?;
528
529        let jit_function = JITFunctionImpl::new(source, function, compile_info);
530
531        Ok(Arc::new(jit_function))
532    }
533}
534
535impl JITFunctionFactory for LLVMFunctionFactory {
536    fn create_jit_function(
537        &self,
538        expression: &str,
539        array_typeid: TypeId,
540    ) -> CoreResult<Box<dyn JITFunction>> {
541        // Check if the function is already in the cache
542        if self.config.use_cache {
543            let cache_key = format!("{expression}-{array_typeid:?}");
544            if let Some(cached_fn) = self.cache.get(&cache_key) {
545                return Ok(cached_fn.as_ref().clone_box());
546            }
547        }
548
549        // Compile the function
550        let jit_function = self.compile(expression, array_typeid)?;
551
552        if self.config.use_cache {
553            // Add the function to the cache
554            let cache_key = format!("{expression}-{array_typeid:?}");
555            // In a real implementation, we'd need to handle this in a thread-safe way
556            // For now, we'll just clone the function
557            let mut cache = self.cache.clone();
558            cache.insert(cache_key, jit_function.clone());
559        }
560
561        // Clone the function and return it
562        Ok(jit_function.as_ref().clone_box())
563    }
564
565    fn supports_array_type(&self, _array_typeid: TypeId) -> bool {
566        // For simplicity, we'll say this factory supports all array types
567        true
568    }
569}
570
571/// A factory for creating JIT functions using the Cranelift backend
572pub struct CraneliftFunctionFactory {
573    /// Configuration for JIT compilation
574    config: JITConfig,
575
576    /// Cache of compiled functions
577    cache: HashMap<String, Arc<dyn JITFunction>>,
578}
579
580impl CraneliftFunctionFactory {
581    /// Create a new Cranelift function factory.
582    pub fn new(config: JITConfig) -> Self {
583        Self {
584            config,
585            cache: HashMap::new(),
586        }
587    }
588
589    /// Compile a function for the given expression.
590    ///
591    /// No Cranelift toolchain is linked into this crate (Pure Rust policy),
592    /// so this interprets the expression directly instead — a real
593    /// implementation of "evaluate this expression" for the grammar
594    /// documented on [`jit_make_interpreted_function`], not a placeholder.
595    fn compile(&self, expression: &str, array_typeid: TypeId) -> CoreResult<Arc<dyn JITFunction>> {
596        let mut compile_info = HashMap::new();
597        compile_info.insert("backend".to_string(), "Cranelift".to_string());
598        compile_info.insert("opt_level".to_string(), self.config.opt_level.to_string());
599        compile_info.insert("array_type".to_string(), format!("{array_typeid:?}"));
600
601        let source = expression.to_string();
602        let function = jit_make_interpreted_function(expression)?;
603
604        let jit_function = JITFunctionImpl::new(source, function, compile_info);
605
606        Ok(Arc::new(jit_function))
607    }
608}
609
610impl JITFunctionFactory for CraneliftFunctionFactory {
611    fn create_jit_function(
612        &self,
613        expression: &str,
614        array_typeid: TypeId,
615    ) -> CoreResult<Box<dyn JITFunction>> {
616        // Check if the function is already in the cache
617        if self.config.use_cache {
618            let cache_key = format!("{expression}-{array_typeid:?}");
619            if let Some(cached_fn) = self.cache.get(&cache_key) {
620                return Ok(cached_fn.as_ref().clone_box());
621            }
622        }
623
624        // Compile the function
625        let jit_function = self.compile(expression, array_typeid)?;
626
627        if self.config.use_cache {
628            // Add the function to the cache
629            let cache_key = format!("{expression}-{array_typeid:?}");
630            // In a real implementation, we'd need to handle this in a thread-safe way
631            // For now, we'll just clone the function
632            let mut cache = self.cache.clone();
633            cache.insert(cache_key, jit_function.clone());
634        }
635
636        // Clone the function and return it
637        Ok(jit_function.as_ref().clone_box())
638    }
639
640    fn supports_array_type(&self, _array_typeid: TypeId) -> bool {
641        // For simplicity, we'll say this factory supports all array types
642        true
643    }
644}
645
646/// A JIT manager that selects the appropriate factory for a given array type
647pub struct JITManager {
648    /// The available JIT function factories
649    factories: Vec<Box<dyn JITFunctionFactory>>,
650
651    /// Default configuration for JIT compilation
652    defaultconfig: JITConfig,
653}
654
655impl JITManager {
656    /// Create a new JIT manager.
657    pub fn new(defaultconfig: JITConfig) -> Self {
658        Self {
659            factories: Vec::new(),
660            defaultconfig,
661        }
662    }
663
664    /// Register a JIT function factory.
665    pub fn register_factory(&mut self, factory: Box<dyn JITFunctionFactory>) {
666        self.factories.push(factory);
667    }
668
669    /// Get a JIT function factory that supports the given array type.
670    pub fn get_factory_for_array_type(
671        &self,
672        array_typeid: TypeId,
673    ) -> Option<&dyn JITFunctionFactory> {
674        for factory in &self.factories {
675            if factory.supports_array_type(array_typeid) {
676                return Some(&**factory);
677            }
678        }
679        None
680    }
681
682    /// Compile a JIT function for the given expression and array type.
683    pub fn compile(
684        &self,
685        expression: &str,
686        array_typeid: TypeId,
687    ) -> CoreResult<Box<dyn JITFunction>> {
688        // Find a factory that supports the array type
689        if let Some(factory) = self.get_factory_for_array_type(array_typeid) {
690            factory.create_jit_function(expression, array_typeid)
691        } else {
692            Err(CoreError::JITError(ErrorContext::new(format!(
693                "No JIT factory supports array type: {array_typeid:?}"
694            ))))
695        }
696    }
697
698    /// Initialize the JIT manager with default factories.
699    pub fn initialize(&mut self) {
700        // Create and register the default factories
701        let llvm_config = JITConfig {
702            backend: JITBackend::LLVM,
703            ..self.defaultconfig.clone()
704        };
705        let llvm_factory = Box::new(LLVMFunctionFactory::new(llvm_config));
706
707        let cranelift_config = JITConfig {
708            backend: JITBackend::Cranelift,
709            ..self.defaultconfig.clone()
710        };
711        let cranelift_factory = Box::new(CraneliftFunctionFactory::new(cranelift_config));
712
713        self.register_factory(llvm_factory);
714        self.register_factory(cranelift_factory);
715    }
716
717    /// Get the global JIT manager instance.
718    #[must_use]
719    pub fn global() -> &'static RwLock<Self> {
720        static INSTANCE: LazyLock<RwLock<JITManager>> = LazyLock::new(|| {
721            RwLock::new(JITManager {
722                factories: Vec::new(),
723                defaultconfig: JITConfig {
724                    backend: JITBackend::LLVM,
725                    optimize: true,
726                    opt_level: 2,
727                    use_cache: true,
728                    backend_options: HashMap::new(),
729                },
730            })
731        });
732        &INSTANCE
733    }
734}
735
736/// An array that supports JIT compilation
737pub struct JITEnabledArray<T, A> {
738    /// The underlying array
739    inner: A,
740
741    /// Phantom data for the element type
742    phantom: PhantomData<T>,
743}
744
745impl<T, A> JITEnabledArray<T, A> {
746    /// Create a new JIT-enabled array.
747    pub fn new(inner: A) -> Self {
748        Self {
749            inner,
750            phantom: PhantomData,
751        }
752    }
753
754    /// Get a reference to the inner array.
755    pub const fn inner(&self) -> &A {
756        &self.inner
757    }
758}
759
760impl<T, A: Clone> Clone for JITEnabledArray<T, A> {
761    fn clone(&self) -> Self {
762        Self {
763            inner: self.inner.clone(),
764            phantom: PhantomData::<T>,
765        }
766    }
767}
768
769impl<T, A> JITArray for JITEnabledArray<T, A>
770where
771    T: Send + Sync + 'static,
772    A: ArrayProtocol + Clone + Send + Sync + 'static,
773{
774    fn compile(&self, expression: &str) -> CoreResult<Box<dyn JITFunction>> {
775        // Get the JIT manager
776        let jit_manager = JITManager::global();
777        let jit_manager = jit_manager.read().expect("Operation failed");
778
779        // Compile the function
780        (*jit_manager).compile(expression, TypeId::of::<A>())
781    }
782
783    fn supports_jit(&self) -> bool {
784        // Check if there's a factory that supports this array type
785        let jit_manager = JITManager::global();
786        let jit_manager = jit_manager.read().expect("Operation failed");
787
788        jit_manager
789            .get_factory_for_array_type(TypeId::of::<A>())
790            .is_some()
791    }
792
793    fn jit_info(&self) -> HashMap<String, String> {
794        let mut info = HashMap::new();
795
796        // Check if JIT is supported
797        let supported = self.supports_jit();
798        info.insert("supports_jit".to_string(), supported.to_string());
799
800        if supported {
801            // Get the JIT manager
802            let jit_manager = JITManager::global();
803            let jit_manager = jit_manager.read().expect("Operation failed");
804
805            // Get the factory
806            if jit_manager
807                .get_factory_for_array_type(TypeId::of::<A>())
808                .is_some()
809            {
810                // Get the factory's info
811                info.insert("factory".to_string(), "JIT factory available".to_string());
812            }
813        }
814
815        info
816    }
817}
818
819impl<T, A> ArrayProtocol for JITEnabledArray<T, A>
820where
821    T: Send + Sync + 'static,
822    A: ArrayProtocol + Clone + Send + Sync + 'static,
823{
824    fn array_function(
825        &self,
826        func: &ArrayFunction,
827        types: &[TypeId],
828        args: &[Box<dyn Any>],
829        kwargs: &HashMap<String, Box<dyn Any>>,
830    ) -> Result<Box<dyn Any>, crate::array_protocol::NotImplemented> {
831        // For now, just delegate to the inner array
832        self.inner.array_function(func, types, args, kwargs)
833    }
834
835    fn as_any(&self) -> &dyn Any {
836        self
837    }
838
839    fn shape(&self) -> &[usize] {
840        self.inner.shape()
841    }
842
843    fn dtype(&self) -> TypeId {
844        self.inner.dtype()
845    }
846
847    fn box_clone(&self) -> Box<dyn ArrayProtocol> {
848        // Clone the inner array directly
849        let inner_clone = self.inner.clone();
850        Box::new(Self {
851            inner: inner_clone,
852            phantom: PhantomData::<T>,
853        })
854    }
855}
856
857#[cfg(test)]
858mod tests {
859    use super::*;
860    use crate::array_protocol::NdarrayWrapper;
861    use ::ndarray::Array2;
862
863    #[test]
864    fn test_jit_function_creation() {
865        // Create a JIT function factory
866        let config = JITConfig {
867            backend: JITBackend::LLVM,
868            ..Default::default()
869        };
870        let factory = LLVMFunctionFactory::new(config);
871
872        // Create a simple expression
873        let expression = "x + y";
874
875        // Compile the function
876        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
877        let jit_function = factory
878            .create_jit_function(expression, array_typeid)
879            .expect("Operation failed");
880
881        // Check the function's properties
882        assert_eq!(jit_function.source(), expression);
883        let compile_info = jit_function.compile_info();
884        assert_eq!(
885            compile_info.get("backend").expect("Operation failed"),
886            "LLVM"
887        );
888    }
889
890    #[test]
891    fn test_jit_manager() {
892        // Initialize the JIT manager
893        let mut jit_manager = JITManager::new(JITConfig::default());
894        jit_manager.initialize();
895
896        // Check that the factories were registered
897        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
898        assert!(jit_manager
899            .get_factory_for_array_type(array_typeid)
900            .is_some());
901
902        // Compile a function
903        let expression = "x + y";
904        let jit_function = jit_manager
905            .compile(expression, array_typeid)
906            .expect("Operation failed");
907
908        // Check the function's properties
909        assert_eq!(jit_function.source(), expression);
910    }
911
912    #[test]
913    fn test_jit_enabled_array() {
914        // Create an ndarray
915        let array = Array2::<f64>::ones((10, 5));
916        let wrapped = NdarrayWrapper::new(array);
917
918        // Create a JIT-enabled array
919        let jit_array: JITEnabledArray<f64, _> = JITEnabledArray::new(wrapped);
920
921        // Initialize the JIT manager
922        {
923            let mut jit_manager = JITManager::global().write().expect("Operation failed");
924            jit_manager.initialize();
925        }
926
927        // Check if JIT is supported
928        assert!(jit_array.supports_jit());
929
930        // Compile a function
931        let expression = "x + y";
932        let jit_function = jit_array.compile(expression).expect("Operation failed");
933
934        // Check the function's properties
935        assert_eq!(jit_function.source(), expression);
936    }
937
938    /// Regression test: both backends used to always return the constant
939    /// 42.0 for any expression/arguments. Uses non-constant, non-42 data
940    /// specifically so a lingering hardcoded-42.0 fallback would fail.
941    #[test]
942    fn test_llvm_and_cranelift_backends_evaluate_for_real() {
943        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
944
945        for factory in [
946            Box::new(LLVMFunctionFactory::new(JITConfig::default())) as Box<dyn JITFunctionFactory>,
947            Box::new(CraneliftFunctionFactory::new(JITConfig::default()))
948                as Box<dyn JITFunctionFactory>,
949        ] {
950            let jit_function = factory
951                .create_jit_function("x * y + 1", array_typeid)
952                .expect("compile should succeed");
953
954            let args: Vec<Box<dyn Any>> = vec![Box::new(3.0_f64), Box::new(4.0_f64)];
955            let result = jit_function
956                .evaluate(&args)
957                .expect("evaluate should succeed");
958            let value = *result.downcast_ref::<f64>().expect("result should be f64");
959            assert_eq!(value, 3.0 * 4.0 + 1.0);
960
961            // Different (non-constant) arguments must give a different
962            // result — this is what a hardcoded-42.0 stub could never do.
963            let args2: Vec<Box<dyn Any>> = vec![Box::new(10.0_f64), Box::new(-2.0_f64)];
964            let result2 = jit_function
965                .evaluate(&args2)
966                .expect("evaluate should succeed");
967            let value2 = *result2.downcast_ref::<f64>().expect("result should be f64");
968            assert_eq!(value2, 10.0 * -2.0 + 1.0);
969            assert_ne!(value, value2);
970        }
971    }
972
973    #[test]
974    fn test_jit_expression_grammar_operators_and_functions() {
975        let factory = LLVMFunctionFactory::new(JITConfig::default());
976        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
977
978        let cases: &[(&str, &[f64], f64)] = &[
979            ("2 + 3 * 4", &[], 14.0),
980            ("(2 + 3) * 4", &[], 20.0),
981            ("2 ^ 10", &[], 1024.0),
982            ("-x + 1", &[5.0], -4.0),
983            ("sqrt(x)", &[16.0], 4.0),
984            ("max(x, y)", &[3.0, 7.0], 7.0),
985            ("min(x, y)", &[3.0, 7.0], 3.0),
986            ("abs(x)", &[-9.5], 9.5),
987        ];
988
989        for (expr, args, expected) in cases {
990            let jit_function = factory
991                .create_jit_function(expr, array_typeid)
992                .unwrap_or_else(|e| panic!("compiling '{expr}' should succeed: {e}"));
993            let boxed_args: Vec<Box<dyn Any>> =
994                args.iter().map(|&v| Box::new(v) as Box<dyn Any>).collect();
995            let result = jit_function
996                .evaluate(&boxed_args)
997                .unwrap_or_else(|e| panic!("evaluating '{expr}' should succeed: {e}"));
998            let value = *result
999                .downcast_ref::<f64>()
1000                .unwrap_or_else(|| panic!("'{expr}' result should be f64"));
1001            assert!(
1002                (value - expected).abs() < 1e-9,
1003                "'{expr}' evaluated to {value}, expected {expected}"
1004            );
1005        }
1006    }
1007
1008    #[test]
1009    fn test_jit_invalid_expression_is_honest_error() {
1010        let factory = LLVMFunctionFactory::new(JITConfig::default());
1011        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
1012
1013        let result = factory.create_jit_function("x + * y", array_typeid);
1014        assert!(
1015            result.is_err(),
1016            "a syntactically invalid expression must be rejected, not silently compiled"
1017        );
1018    }
1019
1020    #[test]
1021    fn test_jit_clone_box_preserves_behavior() {
1022        let factory = LLVMFunctionFactory::new(JITConfig::default());
1023        let array_typeid = TypeId::of::<NdarrayWrapper<f64, crate::ndarray::Ix2>>();
1024        let jit_function = factory
1025            .create_jit_function("x * x", array_typeid)
1026            .expect("compile should succeed");
1027
1028        let cloned = jit_function.clone_box();
1029        assert_eq!(cloned.source(), jit_function.source());
1030
1031        let args: Vec<Box<dyn Any>> = vec![Box::new(6.0_f64)];
1032        let original_result = *jit_function
1033            .evaluate(&args)
1034            .expect("evaluate should succeed")
1035            .downcast_ref::<f64>()
1036            .expect("result should be f64");
1037        let cloned_result = *cloned
1038            .evaluate(&args)
1039            .expect("evaluate should succeed")
1040            .downcast_ref::<f64>()
1041            .expect("result should be f64");
1042
1043        assert_eq!(original_result, 36.0);
1044        assert_eq!(cloned_result, 36.0);
1045    }
1046}
scirs2_core/array_protocol/jit_impl.rs

scirs2_core/array_protocol/
jit_impl.rs
