use shape_ast::ast::TypeAnnotation;
use shape_runtime::module_exports::ModuleExports;
use shape_runtime::type_schema::typed_object_from_pairs;
use shape_runtime::type_system::BuiltinTypes;
use shape_runtime::typed_module_exports::{
ConcreteReturn, ConcreteType, TypedReturn, register_typed_function,
};
use shape_value::heap_value::HeapValue;
use shape_value::KindedSlot;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
pub(crate) fn build_type_reflection_snapshot(
compiler: &super::BytecodeCompiler,
enclosing_type_params: &[String],
) -> TypeReflectionSnapshot {
let mut snapshot = TypeReflectionSnapshot::default();
for (name, (field_names, _span)) in &compiler.struct_types {
let field_types = compiler
.struct_generic_info
.get(name)
.map(|info| info.runtime_field_types.clone())
.unwrap_or_default();
let ordered: Vec<(String, TypeAnnotation)> = field_names
.iter()
.filter_map(|fname| {
field_types
.get(fname)
.cloned()
.map(|ann| (fname.clone(), ann))
})
.collect();
snapshot.struct_defs.insert(name.clone(), ordered);
}
for (alias_name, _target) in &compiler.type_aliases {
snapshot.alias_defs.insert(
alias_name.clone(),
TypeAnnotation::Basic(_target.clone()),
);
}
for type_name in compiler
.type_tracker
.schema_registry()
.type_names()
.map(|s| s.to_string())
.collect::<Vec<_>>()
{
if let Some(schema) = compiler.type_tracker.schema_registry().get(&type_name) {
if let Some(enum_info) = schema.get_enum_info() {
let variants: Vec<String> = enum_info
.variants
.iter()
.map(|v| v.name.clone())
.collect();
snapshot.enum_defs.insert(type_name.clone(), variants);
}
}
}
for tp in enclosing_type_params {
snapshot.known_type_params.insert(tp.clone());
}
snapshot
}
#[derive(Debug, Clone, Default)]
pub(crate) struct TypeReflectionSnapshot {
pub(crate) struct_defs: HashMap<String, Vec<(String, TypeAnnotation)>>,
pub(crate) enum_defs: HashMap<String, Vec<String>>,
pub(crate) alias_defs: HashMap<String, TypeAnnotation>,
pub(crate) known_type_params: HashSet<String>,
}
#[derive(Debug, Clone)]
pub(crate) enum ComptimeDirective {
Extend(shape_ast::ast::ExtendStatement),
RemoveTarget,
SetParamType {
param_name: String,
type_annotation: shape_ast::ast::TypeAnnotation,
},
SetParamValue {
param_name: String,
value: KindedSlot,
},
SetReturnType {
type_annotation: shape_ast::ast::TypeAnnotation,
},
ReplaceBody {
body: Vec<shape_ast::ast::Statement>,
},
ReplaceModule {
items: Vec<shape_ast::ast::Item>,
},
}
thread_local! {
static COMPTIME_DIRECTIVES: RefCell<Vec<ComptimeDirective>> = const { RefCell::new(Vec::new()) };
}
pub(crate) fn clear_comptime_directives() {
COMPTIME_DIRECTIVES.with(|directives| {
let mut directives = directives.borrow_mut();
directives.clear();
});
}
pub(crate) fn take_comptime_directives() -> Vec<ComptimeDirective> {
COMPTIME_DIRECTIVES.with(|directives| {
let mut directives = directives.borrow_mut();
std::mem::take(&mut *directives)
})
}
fn push_comptime_directive(directive: ComptimeDirective) -> Result<(), String> {
COMPTIME_DIRECTIVES.with(|directives| {
let mut directives = directives.borrow_mut();
directives.push(directive);
});
Ok(())
}
fn parse_type_annotation_payload(payload: &str) -> Result<shape_ast::ast::TypeAnnotation, String> {
if let Ok(parsed) = serde_json::from_str::<shape_ast::ast::TypeAnnotation>(payload) {
return Ok(parsed);
}
let snippet = format!("fn __type_probe(value: {}) {{ value }}", payload);
let program = shape_ast::parse_program(&snippet)
.map_err(|e| format!("invalid type payload '{}': {}", payload, e))?;
let maybe_ann = program.items.into_iter().find_map(|item| match item {
shape_ast::ast::Item::Function(func, _) => {
func.params.first().and_then(|p| p.type_annotation.clone())
}
_ => None,
});
maybe_ann.ok_or_else(|| format!("could not parse type payload '{}'", payload))
}
fn parse_function_body_payload(payload: &str) -> Result<Vec<shape_ast::ast::Statement>, String> {
if let Ok(parsed) = serde_json::from_str::<Vec<shape_ast::ast::Statement>>(payload) {
return Ok(parsed);
}
let snippet = format!("fn __body_probe() {{ {} }}", payload);
let program = shape_ast::parse_program(&snippet)
.map_err(|e| format!("invalid replacement body payload: {}", e))?;
let maybe_body = program.items.into_iter().find_map(|item| match item {
shape_ast::ast::Item::Function(func, _) => Some(func.body),
_ => None,
});
maybe_body.ok_or_else(|| "could not parse replacement function body payload".to_string())
}
fn parse_module_items_payload(payload: &str) -> Result<Vec<shape_ast::ast::Item>, String> {
if let Ok(parsed) = serde_json::from_str::<Vec<shape_ast::ast::Item>>(payload) {
return Ok(parsed);
}
let snippet = format!("mod __module_probe__ {{ {} }}", payload);
let program = shape_ast::parse_program(&snippet)
.map_err(|e| format!("invalid replacement module payload: {}", e))?;
let maybe_items = program.items.into_iter().find_map(|item| match item {
shape_ast::ast::Item::Module(module, _) if module.name == "__module_probe__" => {
Some(module.items)
}
_ => None,
});
maybe_items.ok_or_else(|| "could not parse replacement module payload".to_string())
}
fn nb_str(s: &str) -> KindedSlot {
KindedSlot::from_string_arc(Arc::new(s.to_string()))
}
pub(crate) fn create_comptime_builtins_module(
trait_impl_keys: HashSet<String>,
type_snapshot: TypeReflectionSnapshot,
) -> ModuleExports {
let mut module = ModuleExports::new("__comptime__");
register_typed_function(
&mut module,
"implements",
"Check if a type implements a trait at compile time",
vec![],
ConcreteType::Bool,
move |nb_args, _ctx| {
let type_name = match nb_args.first().and_then(|nb| nb.as_str()) {
Some(s) => s.to_string(),
None => return Ok(TypedReturn::Concrete(ConcreteReturn::Bool(false))),
};
let trait_name = match nb_args.get(1).and_then(|nb| nb.as_str()) {
Some(s) => s.to_string(),
None => return Ok(TypedReturn::Concrete(ConcreteReturn::Bool(false))),
};
let has_impl = |ty: &str| {
let legacy = format!("{}::{}", trait_name, ty);
let canonical_prefix = format!("{}::{}::", trait_name, ty);
trait_impl_keys.contains(&legacy)
|| trait_impl_keys
.iter()
.any(|key| key.starts_with(&canonical_prefix))
};
if has_impl(&type_name) {
return Ok(TypedReturn::Concrete(ConcreteReturn::Bool(true)));
}
if BuiltinTypes::is_integer_type_name(&type_name) {
for widen_to in &["number", "float", "f64"] {
if has_impl(widen_to) {
return Ok(TypedReturn::Concrete(ConcreteReturn::Bool(true)));
}
}
}
Ok(TypedReturn::Concrete(ConcreteReturn::Bool(false)))
},
);
register_typed_function(
&mut module,
"warning",
"Emit a compile-time warning",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
if let Some(msg) = nb_args.first().and_then(|nb| nb.as_str()) {
eprintln!("[comptime warning] {}", msg);
}
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"error",
"Emit a compile-time error and abort comptime execution",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let msg = match nb_args.first() {
Some(nb) => match nb.as_str() {
Some(s) => s.to_string(),
None => format!("<{:?}>", nb.kind()),
},
None => "comptime error".to_string(),
};
Err(format!("[comptime error] {}", msg))
},
);
let _build_config_schema = shape_runtime::type_schema::register_predeclared_any_schema(&[
"debug".to_string(),
"target_arch".to_string(),
"target_os".to_string(),
"version".to_string(),
]);
register_typed_function(
&mut module,
"build_config",
"Return build-time configuration",
vec![],
ConcreteType::Object,
|_args, _ctx| {
let kinded = typed_object_from_pairs(&[
("debug", KindedSlot::from_bool(cfg!(debug_assertions))),
("version", nb_str(env!("CARGO_PKG_VERSION"))),
("target_os", nb_str(std::env::consts::OS)),
("target_arch", nb_str(std::env::consts::ARCH)),
]);
let bits = kinded.slot().raw();
let ptr = bits as *const shape_value::heap_value::TypedObjectStorage;
unsafe { shape_value::v2::refcount::v2_retain(&(*ptr).header); }
drop(kinded);
Ok(TypedReturn::Concrete(ConcreteReturn::OpaqueTypedObject(
Arc::new(HeapValue::TypedObject(
shape_value::heap_value::TypedObjectPtr::new(ptr),
)),
)))
},
);
let _type_info_schema = shape_runtime::type_schema::register_predeclared_any_schema(&[
"kind".to_string(),
"name".to_string(),
]);
let snapshot_for_type_info = type_snapshot;
register_typed_function(
&mut module,
"type_info",
"Return the TypeInfo reflection record for the named type",
vec![],
ConcreteType::OpaqueTypedObject("TypeInfo".to_string()),
move |nb_args, _ctx| {
let raw_name = nb_args
.first()
.and_then(|nb| nb.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| {
"__type_info_marshal_pending__".to_string()
});
let type_info_hv = build_type_info_heap_value(&raw_name, &snapshot_for_type_info)?;
Ok(TypedReturn::Concrete(ConcreteReturn::OpaqueTypedObject(
Arc::new(type_info_hv),
)))
},
);
register_typed_function(
&mut module,
"__emit_extend",
"Internal: emit extend directive payload",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let json = nb_args
.first()
.and_then(|nb| nb.as_str())
.ok_or_else(|| "__emit_extend expects a JSON string payload".to_string())?;
let extend: shape_ast::ast::ExtendStatement =
serde_json::from_str(json).map_err(|e| format!("invalid extend payload: {}", e))?;
push_comptime_directive(ComptimeDirective::Extend(extend))?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_remove",
"Internal: remove the current annotation target",
vec![],
ConcreteType::Unit,
|_nb_args, _ctx| {
push_comptime_directive(ComptimeDirective::RemoveTarget)?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_set_param_type",
"Internal: set a parameter type by name",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let param_name = nb_args
.first()
.and_then(|nb| nb.as_str())
.ok_or_else(|| {
"__emit_set_param_type expects param name as first string arg".to_string()
})?
.to_string();
let payload = nb_args.get(1).and_then(|nb| nb.as_str()).ok_or_else(|| {
"__emit_set_param_type expects type annotation as second string arg".to_string()
})?;
let type_annotation = parse_type_annotation_payload(payload)?;
push_comptime_directive(ComptimeDirective::SetParamType {
param_name,
type_annotation,
})?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_set_param_value",
"Internal: set a parameter default value by name",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let param_name = nb_args
.first()
.and_then(|nb| nb.as_str())
.ok_or_else(|| {
"__emit_set_param_value expects param name as first string arg".to_string()
})?
.to_string();
let value = nb_args.get(1).cloned().ok_or_else(|| {
"__emit_set_param_value expects a value as second arg".to_string()
})?;
push_comptime_directive(ComptimeDirective::SetParamValue { param_name, value })?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_set_return_type",
"Internal: set the function return type",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let payload = nb_args.first().and_then(|nb| nb.as_str()).ok_or_else(|| {
"__emit_set_return_type expects a type annotation string".to_string()
})?;
let type_annotation = parse_type_annotation_payload(payload)?;
push_comptime_directive(ComptimeDirective::SetReturnType { type_annotation })?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_replace_body",
"Internal: replace function body from AST payload",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let payload = nb_args.first().and_then(|nb| nb.as_str()).ok_or_else(|| {
"__emit_replace_body expects a function body source string".to_string()
})?;
let body = parse_function_body_payload(payload)?;
push_comptime_directive(ComptimeDirective::ReplaceBody { body })?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
register_typed_function(
&mut module,
"__emit_replace_module",
"Internal: replace module items from source payload",
vec![],
ConcreteType::Unit,
|nb_args, _ctx| {
let payload = nb_args.first().and_then(|nb| nb.as_str()).ok_or_else(|| {
"__emit_replace_module expects a module body source string".to_string()
})?;
let items = parse_module_items_payload(payload)?;
push_comptime_directive(ComptimeDirective::ReplaceModule { items })?;
Ok(TypedReturn::Concrete(ConcreteReturn::Unit))
},
);
module
}
#[derive(Debug, Clone, Copy)]
enum TypeKindLabel {
Int,
Float,
Bool,
String,
Decimal,
BigInt,
Array,
HashMap,
Option,
Result,
TypedObject,
TraitObject,
TypeVar,
Function,
Tuple,
Unit,
Unknown,
}
impl TypeKindLabel {
fn as_str(self) -> &'static str {
match self {
TypeKindLabel::Int => "Int",
TypeKindLabel::Float => "Float",
TypeKindLabel::Bool => "Bool",
TypeKindLabel::String => "String",
TypeKindLabel::Decimal => "Decimal",
TypeKindLabel::BigInt => "BigInt",
TypeKindLabel::Array => "Array",
TypeKindLabel::HashMap => "HashMap",
TypeKindLabel::Option => "Option",
TypeKindLabel::Result => "Result",
TypeKindLabel::TypedObject => "TypedObject",
TypeKindLabel::TraitObject => "TraitObject",
TypeKindLabel::TypeVar => "TypeVar",
TypeKindLabel::Function => "Function",
TypeKindLabel::Tuple => "Tuple",
TypeKindLabel::Unit => "Unit",
TypeKindLabel::Unknown => "Unknown",
}
}
}
fn classify_bare_type_name(
name: &str,
snapshot: &TypeReflectionSnapshot,
) -> TypeKindLabel {
if snapshot.known_type_params.contains(name) {
return TypeKindLabel::TypeVar;
}
match name {
"int" | "i64" | "i32" | "i16" | "i8" | "u64" | "u32" | "u16" | "u8" => {
TypeKindLabel::Int
}
"number" | "f64" | "f32" | "float" => TypeKindLabel::Float,
"bool" => TypeKindLabel::Bool,
"string" | "str" => TypeKindLabel::String,
"decimal" => TypeKindLabel::Decimal,
"bigint" => TypeKindLabel::BigInt,
"()" | "unit" | "void" => TypeKindLabel::Unit,
_ => {
if snapshot.struct_defs.contains_key(name)
|| snapshot.alias_defs.contains_key(name)
|| snapshot.enum_defs.contains_key(name)
{
TypeKindLabel::TypedObject
} else {
TypeKindLabel::Unknown
}
}
}
}
fn build_type_info_heap_value(
type_name: &str,
snapshot: &TypeReflectionSnapshot,
) -> Result<HeapValue, String> {
let label = classify_bare_type_name(type_name, snapshot);
let kinded = typed_object_from_pairs(&[
("kind", nb_str(label.as_str())),
("name", nb_str(type_name)),
]);
let bits = kinded.slot().raw();
let ptr = bits as *const shape_value::heap_value::TypedObjectStorage;
unsafe {
shape_value::v2::refcount::v2_retain(&(*ptr).header);
}
drop(kinded);
Ok(HeapValue::TypedObject(
shape_value::heap_value::TypedObjectPtr::new(ptr),
))
}
#[cfg(all(test, feature = "deep-tests"))]
mod tests {
use super::*;
use shape_runtime::type_schema::TypeSchemaRegistry;
fn test_ctx() -> shape_runtime::module_exports::ModuleContext<'static> {
let registry = Box::leak(Box::new(TypeSchemaRegistry::new()));
shape_runtime::module_exports::ModuleContext {
schemas: registry,
invoke_callable: None,
raw_invoker: None,
function_hashes: None,
vm_state: None,
granted_permissions: None,
scope_constraints: None,
set_pending_resume: None,
set_pending_frame_resume: None,
}
}
#[test]
fn test_comptime_builtins_module_created() {
let module = create_comptime_builtins_module(Default::default(), Default::default());
assert_eq!(module.name, "__comptime__");
}
#[test]
fn test_comptime_warning_builtin() {
let ctx = test_ctx();
let module = create_comptime_builtins_module(Default::default(), Default::default());
let args = vec![nb_str("test warning")];
let result = module
.invoke_export("warning", &args, &ctx)
.expect("warning function should exist");
assert!(result.is_ok());
assert!(result.unwrap().is_unit());
}
#[test]
fn test_comptime_error_builtin() {
let ctx = test_ctx();
let module = create_comptime_builtins_module(Default::default(), Default::default());
let args = vec![nb_str("something failed")];
let result = module
.invoke_export("error", &args, &ctx)
.expect("error function should exist");
assert!(result.is_err());
let err = result.unwrap_err();
assert!(format!("{}", err).contains("something failed"));
}
#[test]
fn test_comptime_implements_returns_false_when_not_registered() {
let ctx = test_ctx();
let module = create_comptime_builtins_module(Default::default(), Default::default());
let args = vec![nb_str("Currency"), nb_str("Display")];
let result = module
.invoke_export("implements", &args, &ctx)
.expect("implements function should exist");
assert!(result.is_ok());
assert_eq!(result.unwrap().as_bool(), Some(false));
}
#[test]
fn test_comptime_implements_returns_true_when_registered() {
let ctx = test_ctx();
let mut impls = HashSet::new();
impls.insert("Serializable::number".to_string());
impls.insert("Display::Currency".to_string());
let module = create_comptime_builtins_module(impls, Default::default());
let result = module
.invoke_export(
"implements",
&[nb_str("number"), nb_str("Serializable")],
&ctx,
)
.expect("implements function should exist");
assert_eq!(result.unwrap().as_bool(), Some(true));
let result = module
.invoke_export(
"implements",
&[nb_str("Currency"), nb_str("Display")],
&ctx,
)
.expect("implements function should exist");
assert_eq!(result.unwrap().as_bool(), Some(true));
let result = module
.invoke_export(
"implements",
&[nb_str("string"), nb_str("Serializable")],
&ctx,
)
.expect("implements function should exist");
assert_eq!(result.unwrap().as_bool(), Some(false));
}
#[test]
fn test_comptime_implements_numeric_widening() {
let ctx = test_ctx();
let mut impls = HashSet::new();
impls.insert("Serializable::number".to_string());
let module = create_comptime_builtins_module(impls, Default::default());
let result = module
.invoke_export(
"implements",
&[nb_str("int"), nb_str("Serializable")],
&ctx,
)
.expect("implements function should exist");
assert_eq!(result.unwrap().as_bool(), Some(true));
let result = module
.invoke_export(
"implements",
&[nb_str("i64"), nb_str("Serializable")],
&ctx,
)
.expect("implements function should exist");
assert_eq!(result.unwrap().as_bool(), Some(true));
}
#[test]
fn test_comptime_build_config_builtin() {
let ctx = test_ctx();
let module = create_comptime_builtins_module(Default::default(), Default::default());
let result = module
.invoke_export("build_config", &[], &ctx)
.expect("build_config function should exist");
assert!(result.is_ok());
assert_eq!(result.unwrap().clone().type_name(), "object");
}
}