shape_runtime/engine/
mod.rs

1//! Shape Engine - Unified execution interface
2//!
3//! This module provides a single, unified interface for executing all Shape code,
4//! replacing the multiple specialized executors with one powerful engine.
5
6// Submodules
7mod builder;
8mod execution;
9mod query_extraction;
10mod stdlib;
11mod types;
12
13// Re-export public types
14pub use crate::query_result::QueryType;
15pub use builder::ShapeEngineBuilder;
16use shape_value::ValueWord;
17pub use types::{
18    EngineBootstrapState, ExecutionMetrics, ExecutionResult, ExecutionType, Message, MessageLevel,
19};
20
21use crate::Runtime;
22use crate::data::DataFrame;
23use shape_ast::error::{Result, ShapeError};
24
25#[cfg(feature = "jit")]
26use std::collections::HashMap;
27
28use crate::hashing::HashDigest;
29use crate::snapshot::{ContextSnapshot, ExecutionSnapshot, SemanticSnapshot, SnapshotStore};
30use serde::Serialize;
31use shape_ast::Program;
32use shape_wire::WireValue;
33
34/// Trait for evaluating individual expressions and statement blocks.
35///
36/// This is used by StreamExecutor, WindowExecutor, and JoinExecutor
37/// to evaluate expressions without needing full program compilation.
38/// shape-vm implements this for BytecodeExecutor.
39pub trait ExpressionEvaluator: Send + Sync {
40    /// Evaluate a slice of statements and return the result.
41    fn eval_statements(
42        &self,
43        stmts: &[shape_ast::Statement],
44        ctx: &mut crate::context::ExecutionContext,
45    ) -> Result<ValueWord>;
46
47    /// Evaluate a single expression and return the result.
48    fn eval_expr(
49        &self,
50        expr: &shape_ast::Expr,
51        ctx: &mut crate::context::ExecutionContext,
52    ) -> Result<ValueWord>;
53}
54
55/// Result from ProgramExecutor::execute_program
56pub struct ProgramExecutorResult {
57    pub wire_value: WireValue,
58    pub type_info: Option<shape_wire::metadata::TypeInfo>,
59    pub execution_type: ExecutionType,
60    pub content_json: Option<serde_json::Value>,
61    pub content_html: Option<String>,
62    pub content_terminal: Option<String>,
63}
64
65/// Trait for executing Shape programs
66pub trait ProgramExecutor {
67    fn execute_program(
68        &mut self,
69        engine: &mut ShapeEngine,
70        program: &Program,
71    ) -> Result<ProgramExecutorResult>;
72}
73
74/// The unified Shape execution engine
75pub struct ShapeEngine {
76    /// The runtime environment
77    pub runtime: Runtime,
78    /// Default data for expressions/assignments
79    pub default_data: DataFrame,
80    /// JIT compilation cache (source hash -> compiled program)
81    #[cfg(feature = "jit")]
82    pub(crate) jit_cache: HashMap<u64, ()>,
83    /// Current source text for error messages (set before execution)
84    pub(crate) current_source: Option<String>,
85    /// Optional snapshot store for resumability
86    pub(crate) snapshot_store: Option<SnapshotStore>,
87    /// Last snapshot ID created
88    pub(crate) last_snapshot: Option<HashDigest>,
89    /// Script path for snapshot metadata
90    pub(crate) script_path: Option<String>,
91    /// Exported symbol names (persisted across REPL commands)
92    pub(crate) exported_symbols: std::collections::HashSet<String>,
93}
94
95impl ShapeEngine {
96    /// Create a new Shape engine
97    pub fn new() -> Result<Self> {
98        let mut runtime = Runtime::new_without_stdlib();
99        runtime.enable_persistent_context_without_data();
100
101        Ok(Self {
102            runtime,
103            default_data: DataFrame::default(),
104            #[cfg(feature = "jit")]
105            jit_cache: HashMap::new(),
106            current_source: None,
107            snapshot_store: None,
108            last_snapshot: None,
109            script_path: None,
110            exported_symbols: std::collections::HashSet::new(),
111        })
112    }
113
114    /// Create engine with data
115    pub fn with_data(data: DataFrame) -> Result<Self> {
116        let mut runtime = Runtime::new_without_stdlib();
117        runtime.enable_persistent_context(&data);
118        Ok(Self {
119            runtime,
120            default_data: data,
121            #[cfg(feature = "jit")]
122            jit_cache: HashMap::new(),
123            current_source: None,
124            snapshot_store: None,
125            last_snapshot: None,
126            script_path: None,
127            exported_symbols: std::collections::HashSet::new(),
128        })
129    }
130
131    /// Create engine with async data provider (Phase 6)
132    ///
133    /// This constructor sets up the engine with an async data provider.
134    /// Call `execute_async()` instead of `execute()` to use async prefetching.
135    pub fn with_async_provider(provider: crate::data::SharedAsyncProvider) -> Result<Self> {
136        let runtime_handle = tokio::runtime::Handle::try_current()
137            .map_err(|_| ShapeError::RuntimeError {
138                message: "No tokio runtime available. Ensure with_async_provider is called within a tokio context.".to_string(),
139                location: None,
140            })?;
141        let mut runtime = Runtime::new_without_stdlib();
142
143        // Create ExecutionContext with async provider
144        let ctx = crate::context::ExecutionContext::with_async_provider(provider, runtime_handle);
145        runtime.set_persistent_context(ctx);
146
147        Ok(Self {
148            runtime,
149            default_data: DataFrame::default(),
150            #[cfg(feature = "jit")]
151            jit_cache: HashMap::new(),
152            current_source: None,
153            snapshot_store: None,
154            last_snapshot: None,
155            script_path: None,
156            exported_symbols: std::collections::HashSet::new(),
157        })
158    }
159
160    /// Initialize REPL mode
161    ///
162    /// Call this once after creating the engine and loading stdlib,
163    /// but before executing any REPL commands. This configures output adapters
164    /// for REPL-friendly display.
165    pub fn init_repl(&mut self) {
166        // Set REPL output adapter to preserve PrintResult spans
167        if let Some(ctx) = self.runtime.persistent_context_mut() {
168            ctx.set_output_adapter(Box::new(crate::output_adapter::ReplAdapter));
169        }
170    }
171
172    /// Capture semantic/runtime state after stdlib bootstrap.
173    ///
174    /// Call this on an engine that has already loaded stdlib.
175    pub fn capture_bootstrap_state(&self) -> Result<EngineBootstrapState> {
176        let context =
177            self.runtime
178                .persistent_context()
179                .cloned()
180                .ok_or_else(|| ShapeError::RuntimeError {
181                    message: "No persistent context available for bootstrap capture".to_string(),
182                    location: None,
183                })?;
184        Ok(EngineBootstrapState {
185            semantic: SemanticSnapshot {
186                exported_symbols: self.exported_symbols.clone(),
187            },
188            context,
189        })
190    }
191
192    /// Apply a previously captured stdlib bootstrap state.
193    pub fn apply_bootstrap_state(&mut self, state: &EngineBootstrapState) {
194        self.exported_symbols = state.semantic.exported_symbols.clone();
195        self.runtime.set_persistent_context(state.context.clone());
196    }
197
198    /// Set the script path for snapshot metadata.
199    pub fn set_script_path(&mut self, path: impl Into<String>) {
200        self.script_path = Some(path.into());
201    }
202
203    /// Get the current script path, if set.
204    pub fn script_path(&self) -> Option<&str> {
205        self.script_path.as_deref()
206    }
207
208    /// Enable snapshotting with a content-addressed store.
209    pub fn enable_snapshot_store(&mut self, store: SnapshotStore) {
210        self.snapshot_store = Some(store);
211    }
212
213    /// Get last snapshot ID, if any.
214    pub fn last_snapshot(&self) -> Option<&HashDigest> {
215        self.last_snapshot.as_ref()
216    }
217
218    /// Access the snapshot store (if configured).
219    pub fn snapshot_store(&self) -> Option<&SnapshotStore> {
220        self.snapshot_store.as_ref()
221    }
222
223    /// Store a serializable blob in the snapshot store and return its hash.
224    pub fn store_snapshot_blob<T: Serialize>(&self, value: &T) -> Result<HashDigest> {
225        let store = self
226            .snapshot_store
227            .as_ref()
228            .ok_or_else(|| ShapeError::RuntimeError {
229                message: "Snapshot store not configured".to_string(),
230                location: None,
231            })?;
232        Ok(store.put_struct(value)?)
233    }
234
235    /// Create a snapshot of semantic/runtime state, with optional VM/bytecode hashes supplied by the executor.
236    pub fn snapshot_with_hashes(
237        &mut self,
238        vm_hash: Option<HashDigest>,
239        bytecode_hash: Option<HashDigest>,
240    ) -> Result<HashDigest> {
241        let store = self
242            .snapshot_store
243            .as_ref()
244            .ok_or_else(|| ShapeError::RuntimeError {
245                message: "Snapshot store not configured".to_string(),
246                location: None,
247            })?;
248
249        let semantic = SemanticSnapshot {
250            exported_symbols: self.exported_symbols.clone(),
251        };
252        let semantic_hash = store.put_struct(&semantic)?;
253
254        let context = if let Some(ctx) = self.runtime.persistent_context() {
255            ctx.snapshot(store)?
256        } else {
257            return Err(ShapeError::RuntimeError {
258                message: "No persistent context for snapshot".to_string(),
259                location: None,
260            });
261        };
262        let context_hash = store.put_struct(&context)?;
263
264        let snapshot = ExecutionSnapshot {
265            version: crate::snapshot::SNAPSHOT_VERSION,
266            created_at_ms: chrono::Utc::now().timestamp_millis(),
267            semantic_hash,
268            context_hash,
269            vm_hash,
270            bytecode_hash,
271            script_path: self.script_path.clone(),
272        };
273
274        let snapshot_hash = store.put_snapshot(&snapshot)?;
275        self.last_snapshot = Some(snapshot_hash.clone());
276        Ok(snapshot_hash)
277    }
278
279    /// Load a snapshot and return its components (semantic/context + optional vm/bytecode hashes).
280    pub fn load_snapshot(
281        &self,
282        snapshot_id: &HashDigest,
283    ) -> Result<(
284        SemanticSnapshot,
285        ContextSnapshot,
286        Option<HashDigest>,
287        Option<HashDigest>,
288    )> {
289        let store = self
290            .snapshot_store
291            .as_ref()
292            .ok_or_else(|| ShapeError::RuntimeError {
293                message: "Snapshot store not configured".to_string(),
294                location: None,
295            })?;
296        let snapshot = store.get_snapshot(snapshot_id)?;
297        let semantic: SemanticSnapshot =
298            store
299                .get_struct(&snapshot.semantic_hash)
300                .map_err(|e| ShapeError::RuntimeError {
301                    message: format!("failed to deserialize SemanticSnapshot: {e}"),
302                    location: None,
303                })?;
304        let context: ContextSnapshot =
305            store
306                .get_struct(&snapshot.context_hash)
307                .map_err(|e| ShapeError::RuntimeError {
308                    message: format!("failed to deserialize ContextSnapshot: {e}"),
309                    location: None,
310                })?;
311        Ok((semantic, context, snapshot.vm_hash, snapshot.bytecode_hash))
312    }
313
314    /// Apply a semantic/context snapshot to the current engine.
315    pub fn apply_snapshot(
316        &mut self,
317        semantic: SemanticSnapshot,
318        context: ContextSnapshot,
319    ) -> Result<()> {
320        self.exported_symbols = semantic.exported_symbols;
321        if let Some(ctx) = self.runtime.persistent_context_mut() {
322            let store = self
323                .snapshot_store
324                .as_ref()
325                .ok_or_else(|| ShapeError::RuntimeError {
326                    message: "Snapshot store not configured".to_string(),
327                    location: None,
328                })?;
329            ctx.restore_from_snapshot(context, store)?;
330            Ok(())
331        } else {
332            Err(ShapeError::RuntimeError {
333                message: "No persistent context for snapshot".to_string(),
334                location: None,
335            })
336        }
337    }
338
339    /// Register extension module namespaces with the runtime.
340    /// Must be called before execute() so the module loader recognizes modules like `duckdb`.
341    pub fn register_extension_modules(
342        &mut self,
343        modules: &[crate::extensions::ParsedModuleSchema],
344    ) {
345        self.runtime.register_extension_module_artifacts(modules);
346    }
347
348    /// Set the current source text for error messages
349    ///
350    /// Call this before execute() to enable source-contextualized error messages.
351    /// The source is used during bytecode compilation to populate debug info.
352    pub fn set_source(&mut self, source: &str) {
353        self.current_source = Some(source.to_string());
354    }
355
356    /// Get the current source text (if set)
357    pub fn current_source(&self) -> Option<&str> {
358        self.current_source.as_deref()
359    }
360
361    /// Register a data provider (Phase 8)
362    ///
363    /// Registers a named provider for runtime data access.
364    ///
365    /// # Example
366    ///
367    /// ```ignore
368    /// let adapter = Arc::new(DataFrameAdapter::new(...));
369    /// engine.register_provider("data", adapter);
370    /// ```
371    pub fn register_provider(&mut self, name: &str, provider: crate::data::SharedAsyncProvider) {
372        if let Some(ctx) = self.runtime.persistent_context_mut() {
373            ctx.register_provider(name, provider);
374        }
375    }
376
377    /// Set default data provider (Phase 8)
378    ///
379    /// Sets which provider to use for runtime data access when no provider is specified.
380    pub fn set_default_provider(&mut self, name: &str) -> Result<()> {
381        if let Some(ctx) = self.runtime.persistent_context_mut() {
382            ctx.set_default_provider(name)
383        } else {
384            Err(ShapeError::RuntimeError {
385                message: "No execution context available".to_string(),
386                location: None,
387            })
388        }
389    }
390
391    /// Register a type mapping (Phase 8)
392    ///
393    /// Registers a type mapping that defines the expected DataFrame structure
394    /// for a given type name. Type mappings enable validation and JIT optimization.
395    ///
396    /// # Example
397    ///
398    /// ```ignore
399    /// use shape_core::runtime::type_mapping::TypeMapping;
400    ///
401    /// // Register the Candle type (from stdlib)
402    /// let candle_mapping = TypeMapping::new("Candle".to_string())
403    ///     .add_field("timestamp", "timestamp")
404    ///     .add_field("open", "open")
405    ///     .add_field("high", "high")
406    ///     .add_field("low", "low")
407    ///     .add_field("close", "close")
408    ///     .add_field("volume", "volume")
409    ///     .add_required("timestamp")
410    ///     .add_required("open")
411    ///     .add_required("high")
412    ///     .add_required("low")
413    ///     .add_required("close");
414    ///
415    /// engine.register_type_mapping("Candle", candle_mapping);
416    /// ```
417    pub fn register_type_mapping(
418        &mut self,
419        type_name: &str,
420        mapping: crate::type_mapping::TypeMapping,
421    ) {
422        if let Some(ctx) = self.runtime.persistent_context_mut() {
423            ctx.register_type_mapping(type_name, mapping);
424        }
425    }
426
427    /// Get the current runtime state (for REPL)
428    pub fn get_runtime(&self) -> &Runtime {
429        &self.runtime
430    }
431
432    /// Get mutable runtime (for REPL state updates)
433    pub fn get_runtime_mut(&mut self) -> &mut Runtime {
434        &mut self.runtime
435    }
436
437    /// Get the format hint for a variable (if any)
438    ///
439    /// Returns the format hint specified in the variable's type annotation.
440    /// Example: `let rate: Number @ Percent = 0.05` → Some("Percent")
441    pub fn get_variable_format_hint(&self, name: &str) -> Option<String> {
442        self.runtime
443            .persistent_context()
444            .and_then(|ctx| ctx.get_variable_format_hint(name))
445    }
446
447    // ========================================================================
448    // Format Execution (Shape Runtime Formats)
449    // ========================================================================
450
451    /// Format a value using Shape runtime format evaluation
452    ///
453    /// This uses the format definitions from stdlib (e.g., stdlib/core/formats.shape)
454    /// instead of Rust fallback formatters.
455    ///
456    /// # Arguments
457    ///
458    /// * `value` - The value to format (as f64 for numbers)
459    /// * `type_name` - The Shape type name ("Number", "String", etc.)
460    /// * `format_name` - Optional format name (e.g., "Percent", "Currency"). Uses default if None.
461    /// * `params` - Format parameters as JSON (e.g., {"decimals": 1})
462    ///
463    /// # Returns
464    ///
465    /// Formatted string on success
466    ///
467    /// # Example
468    ///
469    /// ```ignore
470    /// let formatted = engine.format_value_string(
471    ///     0.1234,
472    ///     "Number",
473    ///     Some("Percent"),
474    ///     &HashMap::new()
475    /// )?;
476    /// assert_eq!(formatted, "12.34%");
477    /// ```
478    pub fn format_value_string(
479        &mut self,
480        value: f64,
481        type_name: &str,
482        format_name: Option<&str>,
483        params: &std::collections::HashMap<String, serde_json::Value>,
484    ) -> Result<String> {
485        use std::sync::Arc;
486
487        // Resolve type aliases and merge meta parameter overrides
488        let (resolved_type_name, merged_params) =
489            self.resolve_type_alias_for_formatting(type_name, params)?;
490
491        // Convert merged JSON params to runtime ValueWord values
492        let param_values: std::collections::HashMap<String, ValueWord> = merged_params
493            .iter()
494            .map(|(k, v)| {
495                let runtime_val = match v {
496                    serde_json::Value::Number(n) => ValueWord::from_f64(n.as_f64().unwrap_or(0.0)),
497                    serde_json::Value::String(s) => ValueWord::from_string(Arc::new(s.clone())),
498                    serde_json::Value::Bool(b) => ValueWord::from_bool(*b),
499                    _ => ValueWord::none(),
500                };
501                (k.clone(), runtime_val)
502            })
503            .collect();
504
505        // Convert value to runtime ValueWord
506        let runtime_value = ValueWord::from_f64(value);
507
508        // Call format with resolved type name and merged parameters
509        self.runtime.format_value(
510            runtime_value,
511            resolved_type_name.as_str(),
512            format_name,
513            param_values,
514        )
515    }
516
517    /// Resolve type alias to base type and merge meta parameter overrides
518    ///
519    /// If type_name is an alias (e.g., "Percent4"), resolves to base type ("Percent")
520    /// and merges stored parameter overrides with passed params.
521    fn resolve_type_alias_for_formatting(
522        &self,
523        type_name: &str,
524        params: &std::collections::HashMap<String, serde_json::Value>,
525    ) -> Result<(String, std::collections::HashMap<String, serde_json::Value>)> {
526        // Check if type_name is a type alias through the runtime context
527        let resolved = self
528            .runtime
529            .persistent_context()
530            .map(|ctx| ctx.resolve_type_for_format(type_name));
531
532        if let Some((base_type, Some(overrides))) = resolved {
533            if base_type != type_name {
534                let mut merged = std::collections::HashMap::new();
535
536                // First, add stored overrides from the alias (convert ValueWord to JSON)
537                for (key, val) in overrides {
538                    let json_val = if let Some(n) = val.as_f64() {
539                        serde_json::json!(n)
540                    } else if val.is_bool() {
541                        serde_json::json!(val.as_bool())
542                    } else {
543                        // Skip non-primitive override values
544                        continue;
545                    };
546                    merged.insert(key, json_val);
547                }
548
549                // Then, overlay with passed params (these take precedence)
550                for (key, val) in params {
551                    merged.insert(key.clone(), val.clone());
552                }
553
554                return Ok((base_type, merged));
555            }
556        }
557
558        // Not an alias, use as-is
559        Ok((type_name.to_string(), params.clone()))
560    }
561
562    // ========================================================================
563    // Extension Management
564    // ========================================================================
565
566    /// Load a data source extension from a shared library
567    ///
568    /// # Arguments
569    ///
570    /// * `path` - Path to the extension shared library (.so, .dll, .dylib)
571    /// * `config` - Configuration value for the extension
572    ///
573    /// # Returns
574    ///
575    /// Information about the loaded extension
576    ///
577    /// # Safety
578    ///
579    /// Loading extensions executes arbitrary code. Only load from trusted sources.
580    ///
581    /// # Example
582    ///
583    /// ```ignore
584    /// let info = engine.load_extension(Path::new("./libshape_ext_csv.so"), &json!({}))?;
585    /// println!("Loaded: {} v{}", info.name, info.version);
586    /// ```
587    pub fn load_extension(
588        &mut self,
589        path: &std::path::Path,
590        config: &serde_json::Value,
591    ) -> Result<crate::extensions::LoadedExtension> {
592        if let Some(ctx) = self.runtime.persistent_context_mut() {
593            ctx.load_extension(path, config)
594        } else {
595            Err(ShapeError::RuntimeError {
596                message: "No execution context available for extension loading".to_string(),
597                location: None,
598            })
599        }
600    }
601
602    /// Unload an extension by name
603    ///
604    /// # Arguments
605    ///
606    /// * `name` - Extension name to unload
607    ///
608    /// # Returns
609    ///
610    /// true if plugin was unloaded, false if not found
611    pub fn unload_extension(&mut self, name: &str) -> bool {
612        if let Some(ctx) = self.runtime.persistent_context_mut() {
613            ctx.unload_extension(name)
614        } else {
615            false
616        }
617    }
618
619    /// List all loaded extension names
620    pub fn list_extensions(&self) -> Vec<String> {
621        if let Some(ctx) = self.runtime.persistent_context() {
622            ctx.list_extensions()
623        } else {
624            Vec::new()
625        }
626    }
627
628    /// Get query schema for an extension (for LSP autocomplete)
629    ///
630    /// # Arguments
631    ///
632    /// * `name` - Extension name
633    ///
634    /// # Returns
635    ///
636    /// The query schema if extension exists
637    pub fn get_extension_query_schema(
638        &self,
639        name: &str,
640    ) -> Option<crate::extensions::ParsedQuerySchema> {
641        if let Some(ctx) = self.runtime.persistent_context() {
642            ctx.get_extension_query_schema(name)
643        } else {
644            None
645        }
646    }
647
648    /// Get output schema for an extension (for LSP autocomplete)
649    ///
650    /// # Arguments
651    ///
652    /// * `name` - Extension name
653    ///
654    /// # Returns
655    ///
656    /// The output schema if extension exists
657    pub fn get_extension_output_schema(
658        &self,
659        name: &str,
660    ) -> Option<crate::extensions::ParsedOutputSchema> {
661        if let Some(ctx) = self.runtime.persistent_context() {
662            ctx.get_extension_output_schema(name)
663        } else {
664            None
665        }
666    }
667
668    /// Get an extension data source by name
669    pub fn get_extension(
670        &self,
671        name: &str,
672    ) -> Option<std::sync::Arc<crate::extensions::ExtensionDataSource>> {
673        if let Some(ctx) = self.runtime.persistent_context() {
674            ctx.get_extension(name)
675        } else {
676            None
677        }
678    }
679
680    /// Get extension module schema by module namespace.
681    pub fn get_extension_module_schema(
682        &self,
683        module_name: &str,
684    ) -> Option<crate::extensions::ParsedModuleSchema> {
685        if let Some(ctx) = self.runtime.persistent_context() {
686            ctx.get_extension_module_schema(module_name)
687        } else {
688            None
689        }
690    }
691
692    /// Build VM extension modules from loaded extension module capabilities.
693    pub fn module_exports_from_extensions(&self) -> Vec<crate::module_exports::ModuleExports> {
694        if let Some(ctx) = self.runtime.persistent_context() {
695            ctx.module_exports_from_extensions()
696        } else {
697            Vec::new()
698        }
699    }
700
701    /// Return all loaded language runtimes, keyed by language identifier.
702    pub fn language_runtimes(
703        &self,
704    ) -> std::collections::HashMap<String, std::sync::Arc<crate::plugins::language_runtime::PluginLanguageRuntime>>
705    {
706        if let Some(ctx) = self.runtime.persistent_context() {
707            ctx.language_runtimes()
708        } else {
709            std::collections::HashMap::new()
710        }
711    }
712
713    /// Invoke one loaded module export via module namespace.
714    pub fn invoke_extension_module_nb(
715        &self,
716        module_name: &str,
717        function: &str,
718        args: &[shape_value::ValueWord],
719    ) -> Result<shape_value::ValueWord> {
720        if let Some(ctx) = self.runtime.persistent_context() {
721            ctx.invoke_extension_module_nb(module_name, function, args)
722        } else {
723            Err(shape_ast::error::ShapeError::RuntimeError {
724                message: "No runtime context available".to_string(),
725                location: None,
726            })
727        }
728    }
729
730    /// Invoke one loaded module export via module namespace.
731    pub fn invoke_extension_module_wire(
732        &self,
733        module_name: &str,
734        function: &str,
735        args: &[shape_wire::WireValue],
736    ) -> Result<shape_wire::WireValue> {
737        if let Some(ctx) = self.runtime.persistent_context() {
738            ctx.invoke_extension_module_wire(module_name, function, args)
739        } else {
740            Err(shape_ast::error::ShapeError::RuntimeError {
741                message: "No runtime context available".to_string(),
742                location: None,
743            })
744        }
745    }
746
747    // ========================================================================
748    // Progress Tracking
749    // ========================================================================
750
751    /// Enable progress tracking and return the registry for subscriptions
752    ///
753    /// Call this before executing code that may report progress.
754    /// The returned registry can be used to subscribe to progress events.
755    ///
756    /// # Example
757    ///
758    /// ```ignore
759    /// let registry = engine.enable_progress_tracking();
760    /// let mut receiver = registry.subscribe();
761    ///
762    /// // In a separate task
763    /// while let Ok(event) = receiver.recv().await {
764    ///     println!("Progress: {:?}", event);
765    /// }
766    /// ```
767    pub fn enable_progress_tracking(
768        &mut self,
769    ) -> std::sync::Arc<crate::progress::ProgressRegistry> {
770        // ProgressRegistry::new() already returns Arc<Self>
771        let registry = crate::progress::ProgressRegistry::new();
772        if let Some(ctx) = self.runtime.persistent_context_mut() {
773            ctx.set_progress_registry(registry.clone());
774        }
775        registry
776    }
777
778    /// Get the current progress registry if enabled
779    pub fn progress_registry(&self) -> Option<std::sync::Arc<crate::progress::ProgressRegistry>> {
780        self.runtime
781            .persistent_context()
782            .and_then(|ctx| ctx.progress_registry())
783            .cloned()
784    }
785
786    /// Check if there are pending progress events
787    pub fn has_pending_progress(&self) -> bool {
788        if let Some(registry) = self.progress_registry() {
789            !registry.is_empty()
790        } else {
791            false
792        }
793    }
794
795    /// Poll for progress events (non-blocking)
796    ///
797    /// Returns the next progress event if available, or None if queue is empty.
798    pub fn poll_progress(&self) -> Option<crate::progress::ProgressEvent> {
799        self.progress_registry()
800            .and_then(|registry| registry.try_recv())
801    }
802}
803
804impl Default for ShapeEngine {
805    fn default() -> Self {
806        Self::new().expect("Failed to create default Shape engine")
807    }
808}
809
810#[cfg(test)]
811mod tests {
812    use super::*;
813    use crate::extensions::{ParsedModuleArtifact, ParsedModuleSchema};
814
815    #[test]
816    fn test_register_extension_modules_registers_module_loader_artifacts() {
817        let mut engine = ShapeEngine::new().expect("engine should create");
818
819        engine.register_extension_modules(&[ParsedModuleSchema {
820            module_name: "duckdb".to_string(),
821            functions: Vec::new(),
822            artifacts: vec![ParsedModuleArtifact {
823                module_path: "duckdb".to_string(),
824                source: Some("pub fn connect(uri) { uri }".to_string()),
825                compiled: None,
826            }],
827        }]);
828
829        let mut loader = engine.runtime.configured_module_loader();
830        let module = loader
831            .load_module("duckdb")
832            .expect("registered extension module artifact should load");
833        assert!(
834            module.exports.contains_key("connect"),
835            "expected connect export"
836        );
837    }
838}
shape_runtime/engine/mod.rs

shape_runtime/engine/
mod.rs
