1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
5 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor, Value,
6};
7use runmat_macros::runtime_builtin;
8
9use crate::builtins::common::spec::{
10 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
11 ReductionNaN, ResidencyPolicy, ShapeRequirements,
12};
13use crate::console::{record_console_line, ConsoleStream};
14use crate::gather_if_needed_async;
15
16use super::disp::{empty_return_value, format_for_disp};
17
18#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::io::display")]
19pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
20 name: "display",
21 op_kind: GpuOpKind::Custom("sink"),
22 supported_precisions: &[],
23 broadcast: BroadcastSemantics::None,
24 provider_hooks: &[],
25 constant_strategy: ConstantStrategy::InlineLiteral,
26 residency: ResidencyPolicy::GatherImmediately,
27 nan_mode: ReductionNaN::Include,
28 two_pass_threshold: None,
29 workgroup_size: None,
30 accepts_nan_mode: false,
31 notes: "Always formats on the CPU; GPU tensors are gathered via the active provider before display.",
32};
33
34#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::io::display")]
35pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
36 name: "display",
37 shape: ShapeRequirements::Any,
38 constant_strategy: ConstantStrategy::InlineLiteral,
39 elementwise: None,
40 reduction: None,
41 emits_nan: false,
42 notes: "Side-effecting sink; excluded from fusion planning.",
43};
44
45const DISPLAY_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
46 name: "ans",
47 ty: BuiltinParamType::NumericArray,
48 arity: BuiltinParamArity::Required,
49 default: None,
50 description: "Empty matrix placeholder returned by sink invocation.",
51}];
52
53const DISPLAY_INPUTS_VALUE: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
54 name: "X",
55 ty: BuiltinParamType::Any,
56 arity: BuiltinParamArity::Required,
57 default: None,
58 description: "Value to display with a variable-name header.",
59}];
60
61const DISPLAY_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
62 label: "display(X)",
63 inputs: &DISPLAY_INPUTS_VALUE,
64 outputs: &DISPLAY_OUTPUT,
65}];
66
67const DISPLAY_ERROR_ARG_CONFIG: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
68 code: "RM.DISPLAY.ARG_CONFIG",
69 identifier: Some("RunMat:display:TooManyInputs"),
70 when: "Too many input arguments are passed to display.",
71 message: "display: too many input arguments",
72};
73
74const DISPLAY_ERROR_GATHER: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
75 code: "RM.DISPLAY.GATHER",
76 identifier: Some("RunMat:display:GatherFailed"),
77 when: "Input value cannot be gathered onto the host for rendering.",
78 message: "display: failed to gather value for display",
79};
80
81const DISPLAY_ERRORS: [BuiltinErrorDescriptor; 2] =
82 [DISPLAY_ERROR_ARG_CONFIG, DISPLAY_ERROR_GATHER];
83
84pub const DISPLAY_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
85 signatures: &DISPLAY_SIGNATURES,
86 output_mode: BuiltinOutputMode::Fixed,
87 completion_policy: BuiltinCompletionPolicy::Public,
88 errors: &DISPLAY_ERRORS,
89};
90
91fn display_error(error: &'static BuiltinErrorDescriptor) -> crate::RuntimeError {
92 display_error_with(error, error.message)
93}
94
95fn display_error_with(
96 error: &'static BuiltinErrorDescriptor,
97 message: impl Into<String>,
98) -> crate::RuntimeError {
99 let mut builder = crate::build_runtime_error(message).with_builtin("display");
100 if let Some(identifier) = error.identifier {
101 builder = builder.with_identifier(identifier);
102 }
103 builder.build()
104}
105
106#[runtime_builtin(
107 name = "display",
108 category = "io",
109 summary = "Display values with a variable-name header.",
110 keywords = "display,disp,print,object,gpu",
111 sink = true,
112 accel = "sink",
113 suppress_auto_output = true,
114 type_resolver(crate::builtins::io::type_resolvers::disp_type),
115 descriptor(crate::builtins::io::display::DISPLAY_DESCRIPTOR),
116 builtin_path = "crate::builtins::io::display"
117)]
118pub async fn display_builtin(value: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
119 if !rest.is_empty() {
120 return Err(display_error(&DISPLAY_ERROR_ARG_CONFIG));
121 }
122
123 if dispatch_custom_display_method(&value).await? {
124 return Ok(empty_return_value());
125 }
126
127 let host_value = gather_if_needed_async(&value)
128 .await
129 .map_err(|e| display_error_with(&DISPLAY_ERROR_GATHER, format!("display: {e}")))?;
130 let label = display_label();
131 let body = format_display_body(&label, &format_for_disp(&host_value));
132 record_console_line(ConsoleStream::Stdout, body);
133 Ok(empty_return_value())
134}
135
136async fn dispatch_custom_display_method(value: &Value) -> crate::BuiltinResult<bool> {
137 match value {
138 Value::Object(_) | Value::HandleObject(_) => {}
139 _ => return Ok(false),
140 }
141 let args = vec![value.clone()];
142 crate::dispatcher::try_call_registered_instance_method("display", &args, Some(0))
143 .await
144 .map(|result| result.is_some())
145}
146
147fn display_label() -> String {
148 crate::callsite::arg_text(0)
149 .map(|text| text.trim().to_string())
150 .filter(|text| is_simple_identifier(text))
151 .unwrap_or_else(|| "ans".to_string())
152}
153
154fn is_simple_identifier(text: &str) -> bool {
155 let mut chars = text.chars();
156 let Some(first) = chars.next() else {
157 return false;
158 };
159 (first == '_' || first.is_ascii_alphabetic())
160 && chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
161}
162
163fn format_display_body(label: &str, lines: &[String]) -> String {
164 match lines {
165 [] => format!("{label} ="),
166 [single] if !single.contains('\n') => format!("{label} = {single}"),
167 _ => format!("{label} =\n{}", lines.join("\n")),
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174 use crate::console::{reset_thread_buffer, take_thread_buffer};
175 use futures::executor::block_on;
176 use runmat_builtins::{Access, ClassDef, MethodDef, ObjectInstance, Tensor};
177 use runmat_hir::{SourceId, Span};
178 use std::collections::HashMap;
179 use std::sync::Arc;
180
181 fn span_of(source: &str, needle: &str) -> Span {
182 let start = source.find(needle).expect("needle present");
183 Span {
184 start,
185 end: start + needle.len(),
186 }
187 }
188
189 fn stdout_text() -> String {
190 take_thread_buffer()
191 .into_iter()
192 .filter(|entry| entry.stream == ConsoleStream::Stdout)
193 .map(|entry| entry.text)
194 .collect::<String>()
195 }
196
197 #[test]
198 fn display_descriptor_signatures_cover_core_forms() {
199 let labels: Vec<&str> = DISPLAY_DESCRIPTOR
200 .signatures
201 .iter()
202 .map(|sig| sig.label)
203 .collect();
204 assert!(labels.contains(&"display(X)"));
205 }
206
207 #[test]
208 fn display_uses_simple_callsite_variable_name() {
209 let source = "display(alpha);";
210 let _catalog_guard = crate::source_context::replace_source_catalog(vec![(
211 SourceId(81),
212 "/tmp/display_name.m".to_string(),
213 source.to_string(),
214 )]);
215 let _callsite_guard = crate::callsite::push_callsite(
216 Some(SourceId(81)),
217 Some(vec![span_of(source, "alpha")]),
218 );
219
220 reset_thread_buffer();
221 let result = block_on(display_builtin(Value::Num(42.0), Vec::new())).expect("display");
222
223 assert_eq!(result, Value::Tensor(Tensor::zeros(vec![0, 0])));
224 assert_eq!(stdout_text(), "alpha = 42\n");
225 }
226
227 #[test]
228 fn display_uses_ans_for_expression_callsite() {
229 let source = "display(alpha + 1);";
230 let _catalog_guard = crate::source_context::replace_source_catalog(vec![(
231 SourceId(82),
232 "/tmp/display_expr.m".to_string(),
233 source.to_string(),
234 )]);
235 let _callsite_guard = crate::callsite::push_callsite(
236 Some(SourceId(82)),
237 Some(vec![span_of(source, "alpha + 1")]),
238 );
239
240 reset_thread_buffer();
241 block_on(display_builtin(Value::Num(43.0), Vec::new())).expect("display");
242
243 assert_eq!(stdout_text(), "ans = 43\n");
244 }
245
246 #[test]
247 fn display_formats_multiline_values_under_header() {
248 reset_thread_buffer();
249 let tensor = Tensor::new(vec![1.0, 3.0, 2.0, 4.0], vec![2, 2]).expect("tensor");
250 block_on(display_builtin(Value::Tensor(tensor), Vec::new())).expect("display");
251
252 assert_eq!(stdout_text(), "ans =\n 1 2\n 3 4\n");
253 }
254
255 #[test]
256 fn display_dispatches_custom_object_display_method() {
257 let class_name = "DisplayHookObject".to_string();
258 runmat_builtins::register_class(ClassDef {
259 name: class_name.clone(),
260 parent: None,
261 properties: HashMap::new(),
262 methods: HashMap::from([(
263 "display".to_string(),
264 MethodDef {
265 name: "display".to_string(),
266 is_static: false,
267 is_abstract: false,
268 is_sealed: false,
269 access: Access::Public,
270 function_name: "DisplayHookObject.customDisplayImpl".to_string(),
271 implicit_class_argument: None,
272 },
273 )]),
274 });
275 let _resolver =
276 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|name| {
277 (name == "DisplayHookObject.customDisplayImpl").then_some(771)
278 })));
279 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
280 |function, args, requested_outputs| {
281 let args = args.to_vec();
282 Box::pin(async move {
283 assert_eq!(function, 771);
284 assert_eq!(requested_outputs, 0);
285 assert!(
286 matches!(args.first(), Some(Value::Object(object)) if object.class_name == "DisplayHookObject")
287 );
288 record_console_line(ConsoleStream::Stdout, "custom display");
289 Ok(empty_return_value())
290 })
291 },
292 )));
293
294 reset_thread_buffer();
295 block_on(display_builtin(
296 Value::Object(ObjectInstance::new(class_name)),
297 Vec::new(),
298 ))
299 .expect("display");
300
301 assert_eq!(stdout_text(), "custom display\n");
302 }
303
304 #[test]
305 fn display_dispatches_inherited_custom_object_display_method() {
306 let parent_class_name = "DisplayHookParent".to_string();
307 let child_class_name = "DisplayHookChild".to_string();
308 runmat_builtins::register_class(ClassDef {
309 name: parent_class_name.clone(),
310 parent: None,
311 properties: HashMap::new(),
312 methods: HashMap::from([(
313 "display".to_string(),
314 MethodDef {
315 name: "display".to_string(),
316 is_static: false,
317 is_abstract: false,
318 is_sealed: false,
319 access: Access::Public,
320 function_name: "DisplayHookParent.renderDisplay".to_string(),
321 implicit_class_argument: None,
322 },
323 )]),
324 });
325 runmat_builtins::register_class(ClassDef {
326 name: child_class_name.clone(),
327 parent: Some(parent_class_name),
328 properties: HashMap::new(),
329 methods: HashMap::new(),
330 });
331 let _resolver =
332 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|name| {
333 (name == "DisplayHookParent.renderDisplay").then_some(772)
334 })));
335 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
336 |function, args, requested_outputs| {
337 let args = args.to_vec();
338 Box::pin(async move {
339 assert_eq!(function, 772);
340 assert_eq!(requested_outputs, 0);
341 assert!(
342 matches!(args.first(), Some(Value::Object(object)) if object.class_name == "DisplayHookChild")
343 );
344 record_console_line(ConsoleStream::Stdout, "inherited display");
345 Ok(empty_return_value())
346 })
347 },
348 )));
349
350 reset_thread_buffer();
351 block_on(display_builtin(
352 Value::Object(ObjectInstance::new(child_class_name)),
353 Vec::new(),
354 ))
355 .expect("display");
356
357 assert_eq!(stdout_text(), "inherited display\n");
358 }
359
360 #[test]
361 fn display_rejects_extra_arguments() {
362 let err = block_on(display_builtin(Value::Num(1.0), vec![Value::Num(2.0)]))
363 .expect_err("extra argument should fail");
364 assert_eq!(err.identifier(), DISPLAY_ERROR_ARG_CONFIG.identifier);
365 }
366}