shape_runtime/intrinsics/
mod.rs

1//! High-performance intrinsic functions for Shape
2//!
3//! Intrinsics are Rust-implemented functions that provide performance-critical
4//! operations while keeping domain logic in Shape stdlib.
5//!
6//! These functions are prefixed with `__intrinsic_` and should not be called
7//! directly by users - they are wrapped by Shape stdlib functions.
8
9use crate::context::ExecutionContext;
10use parking_lot::RwLock;
11use shape_ast::error::{Result, ShapeError};
12use shape_value::ValueWord;
13use std::collections::HashMap;
14use std::sync::Arc;
15
16pub mod array;
17pub mod array_transforms;
18pub mod convolution;
19pub mod distributions;
20pub mod fft;
21pub mod math;
22pub mod matrix;
23pub mod matrix_kernels;
24pub mod random;
25pub mod recurrence;
26pub mod rolling;
27pub mod scan;
28pub mod statistical;
29pub mod stochastic;
30pub mod vector;
31
32/// Function signature for intrinsics
33/// Takes evaluated arguments and execution context, returns a ValueWord value
34pub type IntrinsicFn = fn(&[ValueWord], &mut ExecutionContext) -> Result<ValueWord>;
35
36/// Global intrinsics registry
37///
38/// This registry holds all registered intrinsic functions and provides
39/// fast dispatch. It's thread-safe and can be shared across contexts.
40#[derive(Clone)]
41pub struct IntrinsicsRegistry {
42    functions: Arc<RwLock<HashMap<String, IntrinsicFn>>>,
43}
44
45impl std::fmt::Debug for IntrinsicsRegistry {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        f.debug_struct("IntrinsicsRegistry")
48            .field("num_intrinsics", &self.functions.read().len())
49            .finish()
50    }
51}
52
53impl IntrinsicsRegistry {
54    /// Create new registry and register all intrinsics
55    pub fn new() -> Self {
56        let mut functions = HashMap::new();
57
58        // Register all intrinsics by category
59        Self::register_math_intrinsics(&mut functions);
60        Self::register_random_intrinsics(&mut functions);
61        Self::register_distributions_intrinsics(&mut functions);
62        Self::register_stochastic_intrinsics(&mut functions);
63        Self::register_rolling_intrinsics(&mut functions);
64        Self::register_series_intrinsics(&mut functions);
65        Self::register_array_intrinsics(&mut functions);
66        Self::register_statistical_intrinsics(&mut functions);
67        Self::register_vector_intrinsics(&mut functions);
68        Self::register_matrix_intrinsics(&mut functions);
69        Self::register_recurrence_intrinsics(&mut functions);
70        Self::register_convolution_intrinsics(&mut functions);
71        Self::register_scan_intrinsics(&mut functions);
72        Self::register_fft_intrinsics(&mut functions);
73
74        Self {
75            functions: Arc::new(RwLock::new(functions)),
76        }
77    }
78
79    /// Register a single intrinsic
80    pub fn register(&self, name: &str, func: IntrinsicFn) {
81        let full_name = if name.starts_with("__intrinsic_") {
82            name.to_string()
83        } else {
84            format!("__intrinsic_{}", name)
85        };
86
87        self.functions.write().insert(full_name, func);
88    }
89
90    /// Call an intrinsic function
91    pub fn call(
92        &self,
93        name: &str,
94        args: &[ValueWord],
95        ctx: &mut ExecutionContext,
96    ) -> Result<ValueWord> {
97        let functions = self.functions.read();
98
99        let func = functions
100            .get(name)
101            .ok_or_else(|| ShapeError::RuntimeError {
102                message: format!(
103                    "Unknown intrinsic: {}. Available intrinsics: {:?}",
104                    name,
105                    functions.keys().take(5).collect::<Vec<_>>()
106                ),
107                location: None,
108            })?;
109
110        func(args, ctx)
111    }
112
113    /// Check if a function name is an intrinsic
114    pub fn is_intrinsic(&self, name: &str) -> bool {
115        self.functions.read().contains_key(name)
116    }
117
118    /// Get list of all registered intrinsics
119    pub fn list_intrinsics(&self) -> Vec<String> {
120        self.functions.read().keys().cloned().collect()
121    }
122
123    /// Register all math intrinsics
124    fn register_math_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
125        functions.insert("__intrinsic_sum".to_string(), math::intrinsic_sum);
126        functions.insert("__intrinsic_mean".to_string(), math::intrinsic_mean);
127        functions.insert("__intrinsic_min".to_string(), math::intrinsic_min);
128        functions.insert("__intrinsic_max".to_string(), math::intrinsic_max);
129        functions.insert("__intrinsic_std".to_string(), math::intrinsic_std);
130        functions.insert("__intrinsic_variance".to_string(), math::intrinsic_variance);
131        // Trigonometric intrinsics
132        functions.insert("__intrinsic_sin".to_string(), math::intrinsic_sin);
133        functions.insert("__intrinsic_cos".to_string(), math::intrinsic_cos);
134        functions.insert("__intrinsic_tan".to_string(), math::intrinsic_tan);
135        functions.insert("__intrinsic_asin".to_string(), math::intrinsic_asin);
136        functions.insert("__intrinsic_acos".to_string(), math::intrinsic_acos);
137        functions.insert("__intrinsic_atan".to_string(), math::intrinsic_atan);
138        functions.insert("__intrinsic_atan2".to_string(), math::intrinsic_atan2);
139        functions.insert("__intrinsic_sinh".to_string(), math::intrinsic_sinh);
140        functions.insert("__intrinsic_cosh".to_string(), math::intrinsic_cosh);
141        functions.insert("__intrinsic_tanh".to_string(), math::intrinsic_tanh);
142        // Character code intrinsics
143        functions.insert(
144            "__intrinsic_char_code".to_string(),
145            math::intrinsic_char_code,
146        );
147        functions.insert(
148            "__intrinsic_from_char_code".to_string(),
149            math::intrinsic_from_char_code,
150        );
151    }
152
153    /// Register all rolling window intrinsics
154    fn register_rolling_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
155        functions.insert(
156            "__intrinsic_rolling_sum".to_string(),
157            rolling::intrinsic_rolling_sum,
158        );
159        functions.insert(
160            "__intrinsic_rolling_mean".to_string(),
161            rolling::intrinsic_rolling_mean,
162        );
163        functions.insert(
164            "__intrinsic_rolling_std".to_string(),
165            rolling::intrinsic_rolling_std,
166        );
167        functions.insert(
168            "__intrinsic_rolling_min".to_string(),
169            rolling::intrinsic_rolling_min,
170        );
171        functions.insert(
172            "__intrinsic_rolling_max".to_string(),
173            rolling::intrinsic_rolling_max,
174        );
175        functions.insert("__intrinsic_ema".to_string(), rolling::intrinsic_ema);
176    }
177
178    /// Register all column transformation intrinsics
179    fn register_series_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
180        functions.insert(
181            "__intrinsic_shift".to_string(),
182            array_transforms::intrinsic_shift,
183        );
184        functions.insert(
185            "__intrinsic_diff".to_string(),
186            array_transforms::intrinsic_diff,
187        );
188        functions.insert(
189            "__intrinsic_pct_change".to_string(),
190            array_transforms::intrinsic_pct_change,
191        );
192        functions.insert(
193            "__intrinsic_fillna".to_string(),
194            array_transforms::intrinsic_fillna,
195        );
196        functions.insert(
197            "__intrinsic_cumsum".to_string(),
198            array_transforms::intrinsic_cumsum,
199        );
200        functions.insert(
201            "__intrinsic_cumprod".to_string(),
202            array_transforms::intrinsic_cumprod,
203        );
204        functions.insert(
205            "__intrinsic_clip".to_string(),
206            array_transforms::intrinsic_clip,
207        );
208        functions.insert(
209            "__intrinsic_series".to_string(),
210            array_transforms::intrinsic_column_select,
211        );
212    }
213
214    /// Register array operation intrinsics
215    /// Note: map/filter/reduce are now handled directly by the VM via call_value_immediate_nb.
216    fn register_array_intrinsics(_functions: &mut HashMap<String, IntrinsicFn>) {
217        // Previously registered intrinsic_map, intrinsic_filter, intrinsic_reduce here.
218        // These are now handled by the VM executor directly (array_transform.rs, array_aggregation.rs).
219    }
220
221    /// Register statistical intrinsics
222    fn register_statistical_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
223        functions.insert(
224            "__intrinsic_correlation".to_string(),
225            statistical::intrinsic_correlation,
226        );
227        functions.insert(
228            "__intrinsic_covariance".to_string(),
229            statistical::intrinsic_covariance,
230        );
231        functions.insert(
232            "__intrinsic_percentile".to_string(),
233            statistical::intrinsic_percentile,
234        );
235        functions.insert(
236            "__intrinsic_median".to_string(),
237            statistical::intrinsic_median,
238        );
239    }
240
241    /// Register vector intrinsics
242    fn register_vector_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
243        functions.insert("__intrinsic_vec_abs".to_string(), vector::intrinsic_vec_abs);
244        functions.insert(
245            "__intrinsic_vec_sqrt".to_string(),
246            vector::intrinsic_vec_sqrt,
247        );
248        functions.insert("__intrinsic_vec_ln".to_string(), vector::intrinsic_vec_ln);
249        functions.insert("__intrinsic_vec_exp".to_string(), vector::intrinsic_vec_exp);
250        functions.insert("__intrinsic_vec_add".to_string(), vector::intrinsic_vec_add);
251        functions.insert("__intrinsic_vec_sub".to_string(), vector::intrinsic_vec_sub);
252        functions.insert("__intrinsic_vec_mul".to_string(), vector::intrinsic_vec_mul);
253        functions.insert("__intrinsic_vec_div".to_string(), vector::intrinsic_vec_div);
254        functions.insert("__intrinsic_vec_max".to_string(), vector::intrinsic_vec_max);
255        functions.insert("__intrinsic_vec_min".to_string(), vector::intrinsic_vec_min);
256        functions.insert(
257            "__intrinsic_vec_select".to_string(),
258            vector::intrinsic_vec_select,
259        );
260    }
261
262    /// Register recurrence intrinsics
263    fn register_recurrence_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
264        functions.insert(
265            "__intrinsic_linear_recurrence".to_string(),
266            recurrence::intrinsic_linear_recurrence,
267        );
268    }
269
270    /// Register matrix intrinsics
271    fn register_matrix_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
272        functions.insert(
273            "__intrinsic_matmul_vec".to_string(),
274            matrix::intrinsic_matmul_vec,
275        );
276        functions.insert(
277            "__intrinsic_matmul_mat".to_string(),
278            matrix::intrinsic_matmul_mat,
279        );
280    }
281
282    /// Register random number generation intrinsics
283    fn register_random_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
284        functions.insert("__intrinsic_random".to_string(), random::intrinsic_random);
285        functions.insert(
286            "__intrinsic_random_int".to_string(),
287            random::intrinsic_random_int,
288        );
289        functions.insert(
290            "__intrinsic_random_seed".to_string(),
291            random::intrinsic_random_seed,
292        );
293        functions.insert(
294            "__intrinsic_random_normal".to_string(),
295            random::intrinsic_random_normal,
296        );
297        functions.insert(
298            "__intrinsic_random_array".to_string(),
299            random::intrinsic_random_array,
300        );
301    }
302
303    /// Register statistical distribution intrinsics
304    fn register_distributions_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
305        functions.insert(
306            "__intrinsic_dist_uniform".to_string(),
307            distributions::intrinsic_dist_uniform,
308        );
309        functions.insert(
310            "__intrinsic_dist_lognormal".to_string(),
311            distributions::intrinsic_dist_lognormal,
312        );
313        functions.insert(
314            "__intrinsic_dist_exponential".to_string(),
315            distributions::intrinsic_dist_exponential,
316        );
317        functions.insert(
318            "__intrinsic_dist_poisson".to_string(),
319            distributions::intrinsic_dist_poisson,
320        );
321        functions.insert(
322            "__intrinsic_dist_sample_n".to_string(),
323            distributions::intrinsic_dist_sample_n,
324        );
325    }
326
327    /// Register stochastic process intrinsics
328    fn register_stochastic_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
329        functions.insert(
330            "__intrinsic_brownian_motion".to_string(),
331            stochastic::intrinsic_brownian_motion,
332        );
333        functions.insert("__intrinsic_gbm".to_string(), stochastic::intrinsic_gbm);
334        functions.insert(
335            "__intrinsic_ou_process".to_string(),
336            stochastic::intrinsic_ou_process,
337        );
338        functions.insert(
339            "__intrinsic_random_walk".to_string(),
340            stochastic::intrinsic_random_walk,
341        );
342    }
343
344    /// Register convolution intrinsics
345    fn register_convolution_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
346        functions.insert(
347            "__intrinsic_stencil".to_string(),
348            convolution::intrinsic_stencil,
349        );
350    }
351
352    /// Register scan intrinsics
353    fn register_scan_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
354        functions.insert("__intrinsic_scan".to_string(), scan::intrinsic_scan);
355    }
356
357    /// Register FFT (Fast Fourier Transform) intrinsics
358    fn register_fft_intrinsics(functions: &mut HashMap<String, IntrinsicFn>) {
359        functions.insert("__intrinsic_fft".to_string(), fft::intrinsic_fft);
360        functions.insert("__intrinsic_ifft".to_string(), fft::intrinsic_ifft);
361        functions.insert("__intrinsic_psd".to_string(), fft::intrinsic_psd);
362        functions.insert(
363            "__intrinsic_dominant_frequency".to_string(),
364            fft::intrinsic_dominant_frequency,
365        );
366        functions.insert("__intrinsic_bandpass".to_string(), fft::intrinsic_bandpass);
367        functions.insert(
368            "__intrinsic_harmonics".to_string(),
369            fft::intrinsic_harmonics,
370        );
371    }
372}
373
374impl Default for IntrinsicsRegistry {
375    fn default() -> Self {
376        Self::new()
377    }
378}
379
380// ============================================================================
381// Common arg extraction helpers (DRY across all intrinsic modules)
382//
383// These are `pub` so that shape-vm can reuse them when delegating to runtime
384// intrinsics without duplicating extraction/conversion logic.
385// ============================================================================
386
387/// Extract a f64 from a ValueWord argument, coercing int to float.
388pub fn extract_f64(nb: &ValueWord, label: &str) -> Result<f64> {
389    nb.as_number_coerce()
390        .ok_or_else(|| ShapeError::RuntimeError {
391            message: format!("{} must be a number", label),
392            location: None,
393        })
394}
395
396/// Extract a usize from a ValueWord argument (for window sizes, counts, etc.).
397pub fn extract_usize(nb: &ValueWord, label: &str) -> Result<usize> {
398    let n = nb
399        .as_number_coerce()
400        .ok_or_else(|| ShapeError::RuntimeError {
401            message: format!("{} must be a number", label),
402            location: None,
403        })?;
404    Ok(n as usize)
405}
406
407/// Extract a Vec<f64> from a ValueWord array argument.
408///
409/// Supports typed arrays (IntArray, FloatArray) with zero-copy fast paths.
410pub fn extract_f64_array(nb: &ValueWord, label: &str) -> Result<Vec<f64>> {
411    let view = nb.as_any_array().ok_or_else(|| ShapeError::RuntimeError {
412        message: format!("{} must be an array", label),
413        location: None,
414    })?;
415    if let Some(slice) = view.as_f64_slice() {
416        return Ok(slice.to_vec());
417    }
418    if let Some(slice) = view.as_i64_slice() {
419        return Ok(slice.iter().map(|&v| v as f64).collect());
420    }
421    let arr = view.to_generic();
422    arr.iter()
423        .map(|v| {
424            v.as_number_coerce()
425                .ok_or_else(|| ShapeError::RuntimeError {
426                    message: format!("{} must contain only numeric values", label),
427                    location: None,
428                })
429        })
430        .collect()
431}
432
433/// Extract a string reference from a ValueWord argument.
434pub fn extract_str<'a>(nb: &'a ValueWord, label: &str) -> Result<&'a str> {
435    nb.as_str().ok_or_else(|| ShapeError::RuntimeError {
436        message: format!("{} must be a string", label),
437        location: None,
438    })
439}
440
441/// Build a ValueWord array from a Vec<f64>.
442pub fn f64_vec_to_nb_array(data: Vec<f64>) -> ValueWord {
443    ValueWord::from_array(std::sync::Arc::new(
444        data.into_iter().map(ValueWord::from_f64).collect(),
445    ))
446}
447
448/// Build a ValueWord IntArray from a Vec<i64>.
449///
450/// Returns a typed IntArray (preserves integer type fidelity) rather than
451/// a generic array of boxed ValueWords.
452pub fn i64_vec_to_nb_int_array(data: Vec<i64>) -> ValueWord {
453    ValueWord::from_int_array(std::sync::Arc::new(data.into()))
454}
455
456/// Try to get an i64 slice directly from a ValueWord's IntArray heap value.
457///
458/// Zero-copy: returns a reference into the Arc<TypedBuffer<i64>>.
459/// Returns `None` for all non-IntArray values (caller should fall back to f64 path).
460pub fn try_extract_i64_slice(nb: &ValueWord) -> Option<&[i64]> {
461    nb.as_int_array().map(|buf| buf.as_slice())
462}
463
464/// Build a ValueWord IntArray with validity bitmap from Vec<Option<i64>>.
465///
466/// `None` entries become null (validity bit = 0), `Some(v)` entries become valid.
467/// Used by rolling window i64 paths where positions before the window is full
468/// have no value.
469pub fn option_i64_vec_to_nb(data: Vec<Option<i64>>) -> ValueWord {
470    use shape_value::typed_buffer::TypedBuffer;
471    let mut buf = TypedBuffer::<i64>::with_capacity(data.len());
472    for item in data {
473        match item {
474            Some(v) => buf.push(v),
475            None => buf.push_null(),
476        }
477    }
478    ValueWord::from_int_array(std::sync::Arc::new(buf))
479}
shape_runtime/intrinsics/mod.rs

shape_runtime/intrinsics/
mod.rs
