1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
5 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
6 IntValue, StructValue, Value,
7};
8use runmat_macros::runtime_builtin;
9use std::io::{self, Write};
10use std::net::TcpStream;
11
12use crate::builtins::common::spec::{
13 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
14 ReductionNaN, ResidencyPolicy, ShapeRequirements,
15};
16use crate::{build_runtime_error, gather_if_needed_async, BuiltinResult, RuntimeError};
17
18use super::accept::{client_handle, configure_stream, CLIENT_HANDLE_FIELD};
19
20const BUILTIN_NAME: &str = "write";
21
22const WRITE_OUTPUT_COUNT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
23 name: "count",
24 ty: BuiltinParamType::NumericScalar,
25 arity: BuiltinParamArity::Required,
26 default: None,
27 description: "Number of elements written to the socket.",
28}];
29const WRITE_INPUTS_CLIENT_DATA: [BuiltinParamDescriptor; 2] = [
30 BuiltinParamDescriptor {
31 name: "client",
32 ty: BuiltinParamType::Any,
33 arity: BuiltinParamArity::Required,
34 default: None,
35 description: "tcpclient handle struct.",
36 },
37 BuiltinParamDescriptor {
38 name: "data",
39 ty: BuiltinParamType::Any,
40 arity: BuiltinParamArity::Required,
41 default: None,
42 description: "Payload to send.",
43 },
44];
45const WRITE_INPUTS_CLIENT_DATA_DATATYPE: [BuiltinParamDescriptor; 3] = [
46 BuiltinParamDescriptor {
47 name: "client",
48 ty: BuiltinParamType::Any,
49 arity: BuiltinParamArity::Required,
50 default: None,
51 description: "tcpclient handle struct.",
52 },
53 BuiltinParamDescriptor {
54 name: "data",
55 ty: BuiltinParamType::Any,
56 arity: BuiltinParamArity::Required,
57 default: None,
58 description: "Payload to send.",
59 },
60 BuiltinParamDescriptor {
61 name: "datatype",
62 ty: BuiltinParamType::StringScalar,
63 arity: BuiltinParamArity::Optional,
64 default: Some("\"uint8\""),
65 description: "Data type label (for example \"uint8\", \"double\", \"char\", \"string\").",
66 },
67];
68const WRITE_SIGNATURES: [BuiltinSignatureDescriptor; 2] = [
69 BuiltinSignatureDescriptor {
70 label: "count = write(client, data)",
71 inputs: &WRITE_INPUTS_CLIENT_DATA,
72 outputs: &WRITE_OUTPUT_COUNT,
73 },
74 BuiltinSignatureDescriptor {
75 label: "count = write(client, data, datatype)",
76 inputs: &WRITE_INPUTS_CLIENT_DATA_DATATYPE,
77 outputs: &WRITE_OUTPUT_COUNT,
78 },
79];
80
81const WRITE_ERROR_INVALID_CLIENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
82 code: "RM.WRITE.INVALID_CLIENT",
83 identifier: Some("RunMat:write:InvalidTcpClient"),
84 when: "Client handle is missing, malformed, invalid, or disconnected.",
85 message: "write: invalid tcpclient handle",
86};
87const WRITE_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
88 code: "RM.WRITE.INVALID_INPUT",
89 identifier: Some("RunMat:write:InvalidInput"),
90 when: "Argument list shape is unsupported for write.",
91 message: "write: invalid argument list",
92};
93const WRITE_ERROR_INVALID_DATA: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
94 code: "RM.WRITE.INVALID_DATA",
95 identifier: Some("RunMat:write:InvalidData"),
96 when: "Payload cannot be converted to the requested datatype.",
97 message: "write: invalid data payload",
98};
99const WRITE_ERROR_INVALID_DATATYPE: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
100 code: "RM.WRITE.INVALID_DATATYPE",
101 identifier: Some("RunMat:write:InvalidDataType"),
102 when: "Datatype argument is not a supported scalar text label.",
103 message: "write: invalid datatype argument",
104};
105const WRITE_ERROR_NOT_CONNECTED: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
106 code: "RM.WRITE.NOT_CONNECTED",
107 identifier: Some("RunMat:write:NotConnected"),
108 when: "Client has no active socket connection.",
109 message: "write: tcpclient is disconnected",
110};
111const WRITE_ERROR_TIMEOUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
112 code: "RM.WRITE.TIMEOUT",
113 identifier: Some("RunMat:write:Timeout"),
114 when: "Socket write exceeds configured timeout.",
115 message: "write: timed out while sending data",
116};
117const WRITE_ERROR_CONNECTION_CLOSED: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
118 code: "RM.WRITE.CONNECTION_CLOSED",
119 identifier: Some("RunMat:write:ConnectionClosed"),
120 when: "Peer closes socket before payload is fully written.",
121 message: "write: connection closed before all data was sent",
122};
123const WRITE_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
124 code: "RM.WRITE.INTERNAL",
125 identifier: Some("RunMat:write:InternalError"),
126 when: "Internal socket/control-flow conversion fails.",
127 message: "write: internal socket error",
128};
129const WRITE_ERRORS: [BuiltinErrorDescriptor; 8] = [
130 WRITE_ERROR_INVALID_CLIENT,
131 WRITE_ERROR_INVALID_INPUT,
132 WRITE_ERROR_INVALID_DATA,
133 WRITE_ERROR_INVALID_DATATYPE,
134 WRITE_ERROR_NOT_CONNECTED,
135 WRITE_ERROR_TIMEOUT,
136 WRITE_ERROR_CONNECTION_CLOSED,
137 WRITE_ERROR_INTERNAL,
138];
139pub const WRITE_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
140 signatures: &WRITE_SIGNATURES,
141 output_mode: BuiltinOutputMode::Fixed,
142 completion_policy: BuiltinCompletionPolicy::Public,
143 errors: &WRITE_ERRORS,
144};
145
146#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::io::net::write")]
147pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
148 name: "write",
149 op_kind: GpuOpKind::Custom("network"),
150 supported_precisions: &[],
151 broadcast: BroadcastSemantics::None,
152 provider_hooks: &[],
153 constant_strategy: ConstantStrategy::InlineLiteral,
154 residency: ResidencyPolicy::GatherImmediately,
155 nan_mode: ReductionNaN::Include,
156 two_pass_threshold: None,
157 workgroup_size: None,
158 accepts_nan_mode: false,
159 notes: "Socket writes always execute on the host CPU; GPU providers are never consulted.",
160};
161
162fn write_error_with_message(
163 message: impl Into<String>,
164 error: &'static BuiltinErrorDescriptor,
165) -> RuntimeError {
166 let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
167 if let Some(identifier) = error.identifier {
168 builder = builder.with_identifier(identifier);
169 }
170 builder.build()
171}
172
173fn write_error(error: &'static BuiltinErrorDescriptor) -> RuntimeError {
174 write_error_with_message(error.message, error)
175}
176
177fn write_error_with_detail(
178 error: &'static BuiltinErrorDescriptor,
179 detail: impl AsRef<str>,
180) -> RuntimeError {
181 let detail = detail.as_ref();
182 let detail = detail.strip_prefix("write: ").unwrap_or(detail);
183 write_error_with_message(format!("{}: {}", error.message, detail), error)
184}
185
186fn write_flow(error: &'static BuiltinErrorDescriptor, message: impl AsRef<str>) -> RuntimeError {
187 write_error_with_detail(error, message)
188}
189
190fn map_write_flow(err: RuntimeError, error: &'static BuiltinErrorDescriptor) -> RuntimeError {
191 let mut builder = build_runtime_error(format!("{BUILTIN_NAME}: {}", err.message()))
192 .with_builtin(BUILTIN_NAME)
193 .with_source(err);
194 if let Some(identifier) = error.identifier {
195 builder = builder.with_identifier(identifier);
196 }
197 builder.build()
198}
199
200#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::io::net::write")]
201pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
202 name: "write",
203 shape: ShapeRequirements::Any,
204 constant_strategy: ConstantStrategy::InlineLiteral,
205 elementwise: None,
206 reduction: None,
207 emits_nan: false,
208 notes: "Networking builtin executed eagerly on the CPU.",
209};
210
211#[runtime_builtin(
212 name = "write",
213 category = "io/net",
214 summary = "Write numeric or text payloads to TCP client connections.",
215 keywords = "write,tcpclient,networking",
216 type_resolver(crate::builtins::io::type_resolvers::write_type),
217 descriptor(crate::builtins::io::net::write::WRITE_DESCRIPTOR),
218 builtin_path = "crate::builtins::io::net::write"
219)]
220async fn write_builtin(
221 client: Value,
222 data: Value,
223 rest: Vec<Value>,
224) -> crate::BuiltinResult<Value> {
225 let client = gather_if_needed_async(&client)
226 .await
227 .map_err(|flow| map_write_flow(flow, &WRITE_ERROR_INVALID_CLIENT))?;
228 let data = gather_if_needed_async(&data)
229 .await
230 .map_err(|flow| map_write_flow(flow, &WRITE_ERROR_INVALID_DATA))?;
231
232 let mut gathered_rest = Vec::with_capacity(rest.len());
233 for value in rest {
234 gathered_rest.push(
235 gather_if_needed_async(&value)
236 .await
237 .map_err(|flow| map_write_flow(flow, &WRITE_ERROR_INVALID_DATATYPE))?,
238 );
239 }
240 let datatype = parse_arguments(&gathered_rest)?;
241
242 let client_struct = match &client {
243 Value::Struct(st) => st,
244 _ => {
245 return Err(write_flow(
246 &WRITE_ERROR_INVALID_CLIENT,
247 "write: expected tcpclient struct as first argument",
248 ))
249 }
250 };
251
252 let client_id = extract_client_id(client_struct)?;
253 let handle = client_handle(client_id).ok_or_else(|| {
254 write_flow(
255 &WRITE_ERROR_INVALID_CLIENT,
256 "write: tcpclient handle is no longer valid",
257 )
258 })?;
259
260 let (mut stream, timeout, byte_order) = {
261 let guard = handle.lock().unwrap_or_else(|poison| poison.into_inner());
262 if !guard.connected {
263 return Err(write_error(&WRITE_ERROR_NOT_CONNECTED));
264 }
265 let timeout = guard.timeout;
266 let byte_order = parse_byte_order(&guard.byte_order);
267 let stream = guard.stream.try_clone().map_err(|err| {
268 write_flow(
269 &WRITE_ERROR_INTERNAL,
270 format!("write: clone failed ({err})"),
271 )
272 })?;
273 (stream, timeout, byte_order)
274 };
275
276 if let Err(err) = configure_stream(&stream, timeout) {
277 return Err(write_flow(
278 &WRITE_ERROR_INTERNAL,
279 format!("write: unable to configure socket timeout ({err})"),
280 ));
281 }
282
283 let payload = prepare_payload(&data, datatype, byte_order)?;
284 if payload.bytes.is_empty() {
285 return Ok(Value::Num(0.0));
286 }
287
288 match write_bytes(&mut stream, &payload.bytes) {
289 Ok(_) => Ok(Value::Num(payload.elements as f64)),
290 Err(WriteError::Timeout) => Err(write_error(&WRITE_ERROR_TIMEOUT)),
291 Err(WriteError::ConnectionClosed) => {
292 if let Ok(mut guard) = handle.lock() {
293 guard.connected = false;
294 }
295 Err(write_error(&WRITE_ERROR_CONNECTION_CLOSED))
296 }
297 Err(WriteError::Io(err)) => Err(write_flow(
298 &WRITE_ERROR_INTERNAL,
299 format!("write: socket error ({err})"),
300 )),
301 }
302}
303
304#[derive(Clone, Copy)]
305enum DataType {
306 UInt8,
307 Int8,
308 UInt16,
309 Int16,
310 UInt32,
311 Int32,
312 UInt64,
313 Int64,
314 Single,
315 Double,
316 Char,
317 String,
318}
319
320impl DataType {
321 fn default() -> Self {
322 DataType::UInt8
323 }
324
325 fn element_size(self) -> usize {
326 match self {
327 DataType::UInt8 | DataType::Int8 | DataType::Char | DataType::String => 1,
328 DataType::UInt16 | DataType::Int16 => 2,
329 DataType::UInt32 | DataType::Int32 | DataType::Single => 4,
330 DataType::UInt64 | DataType::Int64 | DataType::Double => 8,
331 }
332 }
333}
334
335#[derive(Clone, Copy)]
336enum ByteOrder {
337 Little,
338 Big,
339}
340
341struct Payload {
342 bytes: Vec<u8>,
343 elements: usize,
344}
345
346fn parse_arguments(args: &[Value]) -> BuiltinResult<DataType> {
347 match args.len() {
348 0 => Ok(DataType::default()),
349 1 => parse_datatype(&args[0]),
350 _ => Err(write_flow(
351 &WRITE_ERROR_INVALID_INPUT,
352 "write: expected at most one datatype argument",
353 )),
354 }
355}
356
357fn parse_datatype(value: &Value) -> BuiltinResult<DataType> {
358 let text = scalar_string(value)?;
359 let lowered = text.trim().to_ascii_lowercase();
360 if lowered.is_empty() {
361 return Err(write_flow(
362 &WRITE_ERROR_INVALID_DATATYPE,
363 "write: datatype must not be empty",
364 ));
365 }
366 let dtype = match lowered.as_str() {
367 "uint8" => DataType::UInt8,
368 "int8" => DataType::Int8,
369 "uint16" => DataType::UInt16,
370 "int16" => DataType::Int16,
371 "uint32" => DataType::UInt32,
372 "int32" => DataType::Int32,
373 "uint64" => DataType::UInt64,
374 "int64" => DataType::Int64,
375 "single" => DataType::Single,
376 "double" => DataType::Double,
377 "char" => DataType::Char,
378 "string" => DataType::String,
379 _ => {
380 return Err(write_flow(
381 &WRITE_ERROR_INVALID_DATATYPE,
382 format!("write: unsupported datatype '{text}'"),
383 ))
384 }
385 };
386 Ok(dtype)
387}
388
389fn prepare_payload(data: &Value, datatype: DataType, order: ByteOrder) -> BuiltinResult<Payload> {
390 match datatype {
391 DataType::Char => char_payload(data),
392 DataType::String => string_payload(data),
393 _ => numeric_payload(data, datatype, order),
394 }
395}
396
397fn numeric_payload(data: &Value, datatype: DataType, order: ByteOrder) -> BuiltinResult<Payload> {
398 let values = flatten_numeric(data)?;
399 let mut bytes = Vec::with_capacity(values.len() * datatype.element_size());
400 for value in values.iter().copied() {
401 match datatype {
402 DataType::UInt8 => bytes.push(cast_to_u8(value)),
403 DataType::Int8 => bytes.push(cast_to_i8(value) as u8),
404 DataType::UInt16 => extend_u16(&mut bytes, cast_to_u16(value), order),
405 DataType::Int16 => extend_i16(&mut bytes, cast_to_i16(value), order),
406 DataType::UInt32 => extend_u32(&mut bytes, cast_to_u32(value), order),
407 DataType::Int32 => extend_i32(&mut bytes, cast_to_i32(value), order),
408 DataType::UInt64 => extend_u64(&mut bytes, cast_to_u64(value), order),
409 DataType::Int64 => extend_i64(&mut bytes, cast_to_i64(value), order),
410 DataType::Single => extend_f32(&mut bytes, cast_to_f32(value), order),
411 DataType::Double => extend_f64(&mut bytes, value, order),
412 DataType::Char | DataType::String => unreachable!(),
413 }
414 }
415 Ok(Payload {
416 bytes,
417 elements: values.len(),
418 })
419}
420
421fn char_payload(data: &Value) -> BuiltinResult<Payload> {
422 let bytes = match data {
423 Value::CharArray(ca) => ca.data.iter().map(|&ch| (ch as u32 & 0xFF) as u8).collect(),
424 Value::String(text) => text.bytes().collect(),
425 Value::StringArray(sa) => {
426 if sa.data.len() != 1 {
427 return Err(write_flow(
428 &WRITE_ERROR_INVALID_DATA,
429 "write: string array input must be scalar when using 'char'",
430 ));
431 }
432 sa.data[0].as_bytes().to_vec()
433 }
434 Value::Tensor(t) => t.data.iter().map(|&v| cast_to_u8(v)).collect::<Vec<u8>>(),
435 Value::Num(n) => vec![cast_to_u8(*n)],
436 Value::Int(iv) => vec![cast_to_u8(iv.to_f64())],
437 Value::Bool(b) => vec![if *b { 1 } else { 0 }],
438 Value::LogicalArray(la) => la
439 .data
440 .iter()
441 .map(|&b| if b != 0 { 1 } else { 0 })
442 .collect(),
443 _ => {
444 return Err(write_flow(
445 &WRITE_ERROR_INVALID_DATA,
446 "write: unsupported input for 'char' datatype",
447 ))
448 }
449 };
450 Ok(Payload {
451 elements: bytes.len(),
452 bytes,
453 })
454}
455
456fn string_payload(data: &Value) -> BuiltinResult<Payload> {
457 match data {
458 Value::String(text) => Ok(Payload {
459 elements: 1,
460 bytes: text.as_bytes().to_vec(),
461 }),
462 Value::CharArray(ca) => {
463 let string: String = ca.data.iter().collect();
464 Ok(Payload {
465 elements: 1,
466 bytes: string.into_bytes(),
467 })
468 }
469 Value::StringArray(sa) => {
470 if sa.data.is_empty() {
471 return Ok(Payload {
472 elements: 0,
473 bytes: Vec::new(),
474 });
475 }
476 if sa.data.len() != 1 {
477 return Err(write_flow(
478 &WRITE_ERROR_INVALID_DATA,
479 "write: string array input must be scalar when using 'string'",
480 ));
481 }
482 Ok(Payload {
483 elements: 1,
484 bytes: sa.data[0].as_bytes().to_vec(),
485 })
486 }
487 _ => Err(write_flow(
488 &WRITE_ERROR_INVALID_DATA,
489 "write: expected text input when using 'string' datatype",
490 )),
491 }
492}
493
494fn flatten_numeric(value: &Value) -> BuiltinResult<Vec<f64>> {
495 match value {
496 Value::Tensor(t) => Ok(t.data.clone()),
497 Value::SparseTensor(s) => {
498 let total_elements = s.rows.checked_mul(s.cols).ok_or_else(|| {
499 write_error_with_message(
500 "write: sparse matrix dimensions overflow",
501 &WRITE_ERROR_INTERNAL,
502 )
503 })?;
504 if total_elements > 10_000_000 {
505 return Err(write_error_with_message(
506 format!("write: cannot densify sparse tensor {}x{} ({} elements exceeds safe threshold)", s.rows, s.cols, total_elements),
507 &WRITE_ERROR_INTERNAL,
508 ));
509 }
510 s.to_dense().map(|dense| dense.data).map_err(|err| {
511 write_error_with_message(format!("write: {err}"), &WRITE_ERROR_INTERNAL)
512 })
513 }
514 Value::Num(n) => Ok(vec![*n]),
515 Value::Int(iv) => Ok(vec![iv.to_f64()]),
516 Value::Bool(b) => Ok(vec![if *b { 1.0 } else { 0.0 }]),
517 Value::LogicalArray(la) => Ok(la
518 .data
519 .iter()
520 .map(|&b| if b != 0 { 1.0 } else { 0.0 })
521 .collect()),
522 Value::CharArray(ca) => Ok(ca
523 .data
524 .iter()
525 .map(|&ch| (ch as u32 & 0xFF) as f64)
526 .collect()),
527 Value::String(text) => Ok(text.chars().map(|ch| (ch as u32) as f64).collect()),
528 Value::StringArray(sa) => {
529 if sa.data.len() != 1 {
530 return Err(write_flow(
531 &WRITE_ERROR_INVALID_DATA,
532 "write: string array input must be scalar",
533 ));
534 }
535 Ok(sa.data[0].chars().map(|ch| (ch as u32) as f64).collect())
536 }
537 Value::Complex(_, _) | Value::ComplexTensor(_) => Err(write_flow(
538 &WRITE_ERROR_INVALID_DATA,
539 "write: complex data is not supported",
540 )),
541 Value::Symbolic(_) | Value::SymbolicArray(_) => Err(write_flow(
542 &WRITE_ERROR_INVALID_DATA,
543 "write: symbolic data is not supported",
544 )),
545 Value::Cell(_)
546 | Value::Struct(_)
547 | Value::Object(_)
548 | Value::HandleObject(_)
549 | Value::Listener(_)
550 | Value::FunctionHandle(_)
551 | Value::ExternalFunctionHandle(_)
552 | Value::MethodFunctionHandle(_)
553 | Value::BoundFunctionHandle { .. }
554 | Value::Closure(_)
555 | Value::ClassRef(_)
556 | Value::MException(_)
557 | Value::OutputList(_) => Err(write_flow(
558 &WRITE_ERROR_INVALID_DATA,
559 "write: unsupported input type",
560 )),
561 Value::GpuTensor(_) => Err(write_flow(
562 &WRITE_ERROR_INVALID_DATA,
563 "write: GPU tensor should have been gathered before encoding",
564 )),
565 }
566}
567
568fn cast_to_u8(value: f64) -> u8 {
569 let rounded = rounded_scalar(value);
570 if !rounded.is_finite() {
571 return if rounded.is_sign_negative() {
572 0
573 } else {
574 u8::MAX
575 };
576 }
577 if rounded < 0.0 {
578 0
579 } else if rounded > u8::MAX as f64 {
580 u8::MAX
581 } else {
582 rounded as u8
583 }
584}
585
586fn cast_to_i8(value: f64) -> i8 {
587 let rounded = rounded_scalar(value);
588 if !rounded.is_finite() {
589 return if rounded.is_sign_negative() {
590 i8::MIN
591 } else {
592 i8::MAX
593 };
594 }
595 if rounded < i8::MIN as f64 {
596 i8::MIN
597 } else if rounded > i8::MAX as f64 {
598 i8::MAX
599 } else {
600 rounded as i8
601 }
602}
603
604fn cast_to_u16(value: f64) -> u16 {
605 let rounded = rounded_scalar(value);
606 if !rounded.is_finite() {
607 return if rounded.is_sign_negative() {
608 0
609 } else {
610 u16::MAX
611 };
612 }
613 if rounded < 0.0 {
614 0
615 } else if rounded > u16::MAX as f64 {
616 u16::MAX
617 } else {
618 rounded as u16
619 }
620}
621
622fn cast_to_i16(value: f64) -> i16 {
623 let rounded = rounded_scalar(value);
624 if !rounded.is_finite() {
625 return if rounded.is_sign_negative() {
626 i16::MIN
627 } else {
628 i16::MAX
629 };
630 }
631 if rounded < i16::MIN as f64 {
632 i16::MIN
633 } else if rounded > i16::MAX as f64 {
634 i16::MAX
635 } else {
636 rounded as i16
637 }
638}
639
640fn cast_to_u32(value: f64) -> u32 {
641 let rounded = rounded_scalar(value);
642 if !rounded.is_finite() {
643 return if rounded.is_sign_negative() {
644 0
645 } else {
646 u32::MAX
647 };
648 }
649 if rounded < 0.0 {
650 0
651 } else if rounded > u32::MAX as f64 {
652 u32::MAX
653 } else {
654 rounded as u32
655 }
656}
657
658fn cast_to_i32(value: f64) -> i32 {
659 let rounded = rounded_scalar(value);
660 if !rounded.is_finite() {
661 return if rounded.is_sign_negative() {
662 i32::MIN
663 } else {
664 i32::MAX
665 };
666 }
667 if rounded < i32::MIN as f64 {
668 i32::MIN
669 } else if rounded > i32::MAX as f64 {
670 i32::MAX
671 } else {
672 rounded as i32
673 }
674}
675
676fn cast_to_u64(value: f64) -> u64 {
677 let rounded = rounded_scalar(value);
678 if !rounded.is_finite() {
679 return if rounded.is_sign_negative() {
680 0
681 } else {
682 u64::MAX
683 };
684 }
685 if rounded < 0.0 {
686 0
687 } else if rounded > u64::MAX as f64 {
688 u64::MAX
689 } else {
690 rounded as u64
691 }
692}
693
694fn cast_to_i64(value: f64) -> i64 {
695 let rounded = rounded_scalar(value);
696 if !rounded.is_finite() {
697 return if rounded.is_sign_negative() {
698 i64::MIN
699 } else {
700 i64::MAX
701 };
702 }
703 if rounded < i64::MIN as f64 {
704 i64::MIN
705 } else if rounded > i64::MAX as f64 {
706 i64::MAX
707 } else {
708 rounded as i64
709 }
710}
711
712fn cast_to_f32(value: f64) -> f32 {
713 value as f32
714}
715
716fn rounded_scalar(value: f64) -> f64 {
717 if value.is_nan() {
718 0.0
719 } else {
720 value.round()
721 }
722}
723
724fn extend_u16(buffer: &mut Vec<u8>, value: u16, order: ByteOrder) {
725 match order {
726 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
727 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
728 }
729}
730
731fn extend_i16(buffer: &mut Vec<u8>, value: i16, order: ByteOrder) {
732 match order {
733 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
734 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
735 }
736}
737
738fn extend_u32(buffer: &mut Vec<u8>, value: u32, order: ByteOrder) {
739 match order {
740 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
741 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
742 }
743}
744
745fn extend_i32(buffer: &mut Vec<u8>, value: i32, order: ByteOrder) {
746 match order {
747 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
748 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
749 }
750}
751
752fn extend_u64(buffer: &mut Vec<u8>, value: u64, order: ByteOrder) {
753 match order {
754 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
755 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
756 }
757}
758
759fn extend_i64(buffer: &mut Vec<u8>, value: i64, order: ByteOrder) {
760 match order {
761 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
762 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
763 }
764}
765
766fn extend_f32(buffer: &mut Vec<u8>, value: f32, order: ByteOrder) {
767 match order {
768 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
769 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
770 }
771}
772
773fn extend_f64(buffer: &mut Vec<u8>, value: f64, order: ByteOrder) {
774 match order {
775 ByteOrder::Little => buffer.extend_from_slice(&value.to_le_bytes()),
776 ByteOrder::Big => buffer.extend_from_slice(&value.to_be_bytes()),
777 }
778}
779
780fn parse_byte_order(text: &str) -> ByteOrder {
781 if text.eq_ignore_ascii_case("big-endian") || text.eq_ignore_ascii_case("big endian") {
782 ByteOrder::Big
783 } else {
784 ByteOrder::Little
785 }
786}
787
788fn scalar_string(value: &Value) -> BuiltinResult<String> {
789 match value {
790 Value::String(s) => Ok(s.clone()),
791 Value::CharArray(ca) if ca.rows == 1 => Ok(ca.data.iter().collect()),
792 Value::StringArray(sa) if sa.data.len() == 1 => Ok(sa.data[0].clone()),
793 _ => Err(write_flow(
794 &WRITE_ERROR_INVALID_DATATYPE,
795 "write: datatype argument must be a string scalar or character row vector",
796 )),
797 }
798}
799
800fn extract_client_id(struct_value: &StructValue) -> BuiltinResult<u64> {
801 let id_value = struct_value
802 .fields
803 .get(CLIENT_HANDLE_FIELD)
804 .ok_or_else(|| {
805 write_flow(
806 &WRITE_ERROR_INVALID_CLIENT,
807 "write: tcpclient struct is missing internal handle",
808 )
809 })?;
810 match id_value {
811 Value::Int(IntValue::U64(id)) => Ok(*id),
812 Value::Int(iv) => Ok(iv.to_i64() as u64),
813 _ => Err(write_flow(
814 &WRITE_ERROR_INVALID_CLIENT,
815 "write: tcpclient struct has invalid handle field",
816 )),
817 }
818}
819
820enum WriteError {
821 Timeout,
822 ConnectionClosed,
823 Io(io::Error),
824}
825
826fn write_bytes(stream: &mut TcpStream, bytes: &[u8]) -> Result<(), WriteError> {
827 let mut offset = 0usize;
828 while offset < bytes.len() {
829 match stream.write(&bytes[offset..]) {
830 Ok(0) => return Err(WriteError::ConnectionClosed),
831 Ok(n) => offset += n,
832 Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
833 Err(err) if is_timeout(&err) => return Err(WriteError::Timeout),
834 Err(err) if is_connection_closed_error(&err) => {
835 return Err(WriteError::ConnectionClosed)
836 }
837 Err(err) => return Err(WriteError::Io(err)),
838 }
839 }
840 Ok(())
841}
842
843fn is_timeout(err: &io::Error) -> bool {
844 matches!(
845 err.kind(),
846 io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock
847 )
848}
849
850fn is_connection_closed_error(err: &io::Error) -> bool {
851 matches!(
852 err.kind(),
853 io::ErrorKind::BrokenPipe
854 | io::ErrorKind::ConnectionReset
855 | io::ErrorKind::ConnectionAborted
856 | io::ErrorKind::NotConnected
857 | io::ErrorKind::UnexpectedEof
858 )
859}
860
861#[cfg(test)]
862pub(crate) mod tests {
863 use super::*;
864 use crate::builtins::io::net::accept::{
865 configure_stream, insert_client, remove_client_for_test,
866 };
867 use runmat_builtins::{CharArray, IntValue, StructValue, Tensor};
868 use std::io::Read;
869 use std::net::{TcpListener, TcpStream};
870 use std::sync::{Arc, Barrier};
871 use std::thread;
872
873 fn make_client(stream: TcpStream, timeout: f64, byte_order: &str) -> Value {
874 let peer_addr = stream.peer_addr().expect("peer addr");
875 configure_stream(&stream, timeout).expect("configure stream");
876 let client_id = insert_client(stream, 0, peer_addr, timeout, byte_order.to_string());
877 let mut st = StructValue::new();
878 st.fields.insert(
879 CLIENT_HANDLE_FIELD.to_string(),
880 Value::Int(IntValue::U64(client_id)),
881 );
882 Value::Struct(st)
883 }
884
885 fn client_id(client: &Value) -> u64 {
886 match client {
887 Value::Struct(st) => match st.fields.get(CLIENT_HANDLE_FIELD) {
888 Some(Value::Int(IntValue::U64(id))) => *id,
889 Some(Value::Int(iv)) => iv.to_i64() as u64,
890 other => panic!("unexpected id field {other:?}"),
891 },
892 other => panic!("expected struct, got {other:?}"),
893 }
894 }
895
896 fn assert_error_identifier(err: RuntimeError, expected: &str) {
897 assert_eq!(err.identifier(), Some(expected));
898 }
899
900 fn run_write(client: Value, data: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
901 futures::executor::block_on(write_builtin(client, data, rest))
902 }
903
904 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
905 #[test]
906 fn write_descriptor_signatures_cover_core_forms() {
907 let labels: Vec<&str> = WRITE_DESCRIPTOR
908 .signatures
909 .iter()
910 .map(|sig| sig.label)
911 .collect();
912 assert!(labels.contains(&"count = write(client, data)"));
913 assert!(labels.contains(&"count = write(client, data, datatype)"));
914 }
915
916 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
917 #[test]
918 fn write_default_uint8_sends_bytes() {
919 let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
920 let port = listener.local_addr().unwrap().port();
921 let handle = thread::spawn(move || {
922 let (mut stream, _) = listener.accept().expect("accept");
923 let mut received = Vec::new();
924 stream.read_to_end(&mut received).unwrap_or_default();
925 received
926 });
927
928 let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect");
929 let client = make_client(stream, 1.0, "little-endian");
930 let tensor = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![1, 4]).unwrap();
931 let result = run_write(client.clone(), Value::Tensor(tensor), Vec::new()).expect("write");
932 match result {
933 Value::Num(count) => assert_eq!(count, 4.0),
934 other => panic!("expected numeric result, got {other:?}"),
935 }
936 remove_client_for_test(client_id(&client));
937 let received = handle.join().expect("join");
938 assert_eq!(received, vec![1, 2, 3, 4]);
939 }
940
941 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
942 #[test]
943 fn write_double_big_endian_encodes_correctly() {
944 let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
945 let port = listener.local_addr().unwrap().port();
946 let handle = thread::spawn(move || {
947 let (mut stream, _) = listener.accept().expect("accept");
948 let mut buf = [0u8; 24];
949 stream.read_exact(&mut buf).expect("read");
950 buf
951 });
952
953 let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect");
954 let client = make_client(stream, 1.0, "big-endian");
955 let tensor = Tensor::new(vec![1.5, 2.5, 3.5], vec![1, 3]).unwrap();
956 let result = run_write(
957 client.clone(),
958 Value::Tensor(tensor),
959 vec![Value::from("double")],
960 )
961 .expect("write");
962 match result {
963 Value::Num(count) => assert_eq!(count, 3.0),
964 other => panic!("expected numeric count, got {other:?}"),
965 }
966 remove_client_for_test(client_id(&client));
967
968 let received = handle.join().expect("join");
969 let mut expected = Vec::new();
970 extend_f64(&mut expected, 1.5, ByteOrder::Big);
971 extend_f64(&mut expected, 2.5, ByteOrder::Big);
972 extend_f64(&mut expected, 3.5, ByteOrder::Big);
973 assert_eq!(received.to_vec(), expected);
974 }
975
976 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
977 #[test]
978 fn write_char_payload_encodes_ascii() {
979 let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
980 let port = listener.local_addr().unwrap().port();
981 let handle = thread::spawn(move || {
982 let (mut stream, _) = listener.accept().expect("accept");
983 let mut buf = Vec::new();
984 stream.read_to_end(&mut buf).unwrap_or_default();
985 buf
986 });
987
988 let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect");
989 let client = make_client(stream, 1.0, "little-endian");
990 let chars = CharArray::new("RunMat".chars().collect(), 1, 6).unwrap();
991 let result = run_write(
992 client.clone(),
993 Value::CharArray(chars),
994 vec![Value::from("char")],
995 )
996 .expect("write");
997 match result {
998 Value::Num(count) => assert_eq!(count, 6.0),
999 other => panic!("expected numeric count, got {other:?}"),
1000 }
1001 remove_client_for_test(client_id(&client));
1002 let received = handle.join().expect("join");
1003 assert_eq!(received, b"RunMat");
1004 }
1005
1006 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1007 #[test]
1008 fn write_errors_when_client_disconnected() {
1009 let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
1010 let port = listener.local_addr().unwrap().port();
1011 let barrier = Arc::new(Barrier::new(2));
1012 let thread_barrier = barrier.clone();
1013 let handle = thread::spawn(move || {
1014 let (stream, _) = listener.accept().expect("accept");
1015 thread_barrier.wait();
1016 drop(stream);
1017 });
1018
1019 let stream = TcpStream::connect(("127.0.0.1", port)).expect("connect");
1020 let client = make_client(stream, 1.0, "little-endian");
1021 let id = client_id(&client);
1022 if let Some(handle_ref) = client_handle(id) {
1023 if let Ok(mut guard) = handle_ref.lock() {
1024 guard.connected = false;
1025 }
1026 }
1027
1028 let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![1, 3]).unwrap();
1029 let err = run_write(client.clone(), Value::Tensor(tensor), Vec::new()).expect_err("write");
1030 assert_error_identifier(err, WRITE_ERROR_NOT_CONNECTED.identifier.unwrap());
1031
1032 remove_client_for_test(id);
1033 barrier.wait();
1034 handle.join().expect("join");
1035 }
1036}