1use std::collections::{BTreeSet, HashSet};
4use std::ffi::c_void;
5use std::fmt;
6use std::path::{Path, PathBuf};
7use std::ptr::NonNull;
8use std::sync::{Arc, Mutex};
9
10use latexsnipper_foundation::Result;
11use latexsnipper_runtime::{
12 DeviceKind, RunRequest, RunResponse, RuntimeArtifacts, RuntimeCapabilities, RuntimeDevice,
13 RuntimeFactory, RuntimeKind, RuntimeOptions, RuntimeProbe, RuntimeSession, SessionMetadata,
14 TensorMap,
15};
16use latexsnipper_tensor::{Tensor, TensorData};
17use libloading::Library;
18
19use crate::abi::*;
20use crate::descriptor::RuntimePluginDescriptor;
21use crate::error::{plugin_runtime_error, PluginRuntimeResult};
22
23const MAX_STRING_BYTES: usize = 1024 * 1024;
24const MAX_TENSORS: usize = 4096;
25const MAX_DEVICES: usize = 256;
26
27struct LoadedPlugin {
28 _library: Library,
29 table: LatexSnipperRuntimePluginV1,
30 runtime_id: String,
31 plugin_version: String,
32 library_path: PathBuf,
33}
34
35unsafe impl Send for LoadedPlugin {}
39unsafe impl Sync for LoadedPlugin {}
40
41impl LoadedPlugin {
42 unsafe fn load(
43 library_path: &Path,
44 descriptor: &RuntimePluginDescriptor,
45 ) -> PluginRuntimeResult<Arc<Self>> {
46 let library = unsafe { Library::new(library_path) }.map_err(|error| {
49 plugin_runtime_error(format!(
50 "load trusted runtime plugin '{}': {error}",
51 library_path.display()
52 ))
53 })?;
54 let entry = unsafe { library.get::<RuntimePluginEntryV1>(LATEXSNIPPER_RUNTIME_PLUGIN_ENTRY_V1) }
57 .map_err(|error| {
58 plugin_runtime_error(format!(
59 "trusted library '{}' does not export latexsnipper_runtime_plugin_entry_v1: {error}",
60 library_path.display()
61 ))
62 })?;
63 let table_pointer = unsafe { entry() };
65 if table_pointer.is_null() {
66 return Err(plugin_runtime_error(
67 "runtime plugin entry returned a null function table",
68 ));
69 }
70 let struct_size = unsafe { std::ptr::addr_of!((*table_pointer).struct_size).read() };
73 if struct_size < std::mem::size_of::<LatexSnipperRuntimePluginV1>() {
74 return Err(plugin_runtime_error(format!(
75 "runtime plugin function table is too small: expected at least {}, got {}",
76 std::mem::size_of::<LatexSnipperRuntimePluginV1>(),
77 struct_size
78 )));
79 }
80 let table = unsafe { table_pointer.read() };
82 if table.abi_version != LATEXSNIPPER_RUNTIME_PLUGIN_ABI_V1 {
83 return Err(plugin_runtime_error(format!(
84 "runtime plugin ABI mismatch: host supports v1, plugin reports v{}",
85 table.abi_version
86 )));
87 }
88 require_functions(&table)?;
89 let runtime_id = copy_utf8(table.runtime_id, "runtime id", 128)?;
90 let plugin_version = copy_utf8(table.plugin_version, "plugin version", 128)?;
91 if runtime_id != descriptor.runtime_id || plugin_version != descriptor.plugin_version {
92 return Err(plugin_runtime_error(format!(
93 "runtime plugin identity differs from descriptor: descriptor={}/{}, library={}/{}",
94 descriptor.runtime_id, descriptor.plugin_version, runtime_id, plugin_version
95 )));
96 }
97 Ok(Arc::new(Self {
98 _library: library,
99 table,
100 runtime_id,
101 plugin_version,
102 library_path: library_path.to_path_buf(),
103 }))
104 }
105
106 fn error_message(&self, operation: &str) -> String {
107 let Some(last_error) = self.table.last_error else {
108 return format!("{operation} failed without a plugin error message");
109 };
110 let view = unsafe { last_error() };
112 match copy_utf8(view, "last error", 16 * 1024) {
113 Ok(message) if !message.is_empty() => message,
114 _ => format!("{operation} failed with an invalid plugin error message"),
115 }
116 }
117}
118
119fn require_functions(table: &LatexSnipperRuntimePluginV1) -> PluginRuntimeResult<()> {
120 let missing = [
121 (table.probe.is_none(), "probe"),
122 (table.create_session.is_none(), "create_session"),
123 (table.destroy_session.is_none(), "destroy_session"),
124 (table.run.is_none(), "run"),
125 (table.free_output.is_none(), "free_output"),
126 (table.last_error.is_none(), "last_error"),
127 ]
128 .into_iter()
129 .filter_map(|(missing, name)| missing.then_some(name))
130 .collect::<Vec<_>>();
131 if missing.is_empty() {
132 Ok(())
133 } else {
134 Err(plugin_runtime_error(format!(
135 "runtime plugin function table omits required functions: {}",
136 missing.join(", ")
137 )))
138 }
139}
140
141#[derive(Clone)]
142pub struct RuntimePluginFactory {
143 plugin: Arc<LoadedPlugin>,
144}
145
146impl fmt::Debug for RuntimePluginFactory {
147 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
148 formatter
149 .debug_struct("RuntimePluginFactory")
150 .field("runtime_id", &self.plugin.runtime_id)
151 .field("plugin_version", &self.plugin.plugin_version)
152 .field("library_path", &self.plugin.library_path)
153 .finish()
154 }
155}
156
157impl RuntimePluginFactory {
158 pub(crate) unsafe fn load_trusted(
159 library_path: &Path,
160 descriptor: &RuntimePluginDescriptor,
161 ) -> PluginRuntimeResult<Self> {
162 let plugin = unsafe { LoadedPlugin::load(library_path, descriptor) }?;
164 Ok(Self { plugin })
165 }
166
167 pub fn runtime_id(&self) -> &str {
168 &self.plugin.runtime_id
169 }
170
171 pub fn plugin_version(&self) -> &str {
172 &self.plugin.plugin_version
173 }
174}
175
176impl RuntimeFactory for RuntimePluginFactory {
177 fn kind(&self) -> RuntimeKind {
178 RuntimeKind::Custom(self.plugin.runtime_id.clone())
179 }
180
181 fn probe(&self) -> RuntimeProbe {
182 match probe_plugin(&self.plugin) {
183 Ok(probe) => probe,
184 Err(error) => RuntimeProbe::unavailable(error.to_string()),
185 }
186 }
187
188 fn create_session(
189 &self,
190 artifacts: &RuntimeArtifacts,
191 options: &RuntimeOptions,
192 ) -> Result<Box<dyn RuntimeSession>> {
193 let expected = RuntimeKind::Custom(self.plugin.runtime_id.clone());
194 if artifacts.runtime != expected {
195 return Err(plugin_runtime_error(format!(
196 "custom runtime '{}' received '{}' artifacts",
197 self.plugin.runtime_id, artifacts.runtime
198 )));
199 }
200 if !artifacts.buffers.is_empty() {
201 return Err(plugin_runtime_error(
202 "runtime plugin ABI v1 accepts installed file artifacts, not in-memory model buffers",
203 ));
204 }
205
206 let owned_artifacts = artifacts
207 .files
208 .iter()
209 .map(|(role, path)| {
210 let path = path.to_str().ok_or_else(|| {
211 plugin_runtime_error(format!(
212 "runtime plugin artifact path is not UTF-8: {}",
213 path.display()
214 ))
215 })?;
216 Ok((role.as_bytes().to_vec(), path.as_bytes().to_vec()))
217 })
218 .collect::<PluginRuntimeResult<Vec<_>>>()?;
219 let artifact_views = owned_artifacts
220 .iter()
221 .map(|(role, path)| LatexSnipperArtifactV1 {
222 role: LatexSnipperBytesV1::from_slice(role),
223 path: LatexSnipperBytesV1::from_slice(path),
224 })
225 .collect::<Vec<_>>();
226 let artifact_options = serde_json::to_vec(&artifacts.options)
227 .map_err(|error| plugin_runtime_error(error.to_string()))?;
228 let runtime_options =
229 serde_json::to_vec(options).map_err(|error| plugin_runtime_error(error.to_string()))?;
230 let request = LatexSnipperSessionCreateRequestV1 {
231 artifacts: artifact_views.as_ptr(),
232 artifact_count: artifact_views.len(),
233 artifact_options_json: LatexSnipperBytesV1::from_slice(&artifact_options),
234 runtime_options_json: LatexSnipperBytesV1::from_slice(&runtime_options),
235 };
236 let mut output = LatexSnipperSessionV1::empty();
237 let create = self
238 .plugin
239 .table
240 .create_session
241 .expect("required function was validated while loading plugin");
242 let status = unsafe { create(&request, &mut output) };
244 let handle = NonNull::new(output.handle);
245 if status != LS_RUNTIME_OK {
246 let message = self.plugin.error_message("create session");
247 if let Some(handle) = handle {
248 destroy_handle(&self.plugin, handle);
249 }
250 return Err(plugin_runtime_error(message));
251 }
252 let handle = handle.ok_or_else(|| {
253 plugin_runtime_error("runtime plugin create_session succeeded with a null handle")
254 })?;
255 let metadata_bytes =
256 match copy_bytes(output.metadata_json, "session metadata", MAX_STRING_BYTES) {
257 Ok(bytes) => bytes,
258 Err(error) => {
259 destroy_handle(&self.plugin, handle);
260 return Err(error);
261 }
262 };
263 let metadata: SessionMetadata = match serde_json::from_slice(&metadata_bytes) {
264 Ok(metadata) => metadata,
265 Err(error) => {
266 destroy_handle(&self.plugin, handle);
267 return Err(plugin_runtime_error(format!(
268 "runtime plugin returned invalid session metadata JSON: {error}"
269 )));
270 }
271 };
272 if metadata.runtime != expected {
273 destroy_handle(&self.plugin, handle);
274 return Err(plugin_runtime_error(format!(
275 "runtime plugin session metadata reports '{}', expected '{}'",
276 metadata.runtime, expected
277 )));
278 }
279 validate_metadata(&metadata).inspect_err(|_| destroy_handle(&self.plugin, handle))?;
280 Ok(Box::new(RuntimePluginSession {
281 metadata,
282 plugin: Arc::clone(&self.plugin),
283 handle,
284 execution: Mutex::new(()),
285 }))
286 }
287}
288
289fn probe_plugin(plugin: &LoadedPlugin) -> PluginRuntimeResult<RuntimeProbe> {
290 let probe = plugin
291 .table
292 .probe
293 .expect("required function was validated while loading plugin");
294 let mut native = LatexSnipperRuntimeProbeV1::empty();
295 let status = unsafe { probe(&mut native) };
297 if status != LS_RUNTIME_OK {
298 return Err(plugin_runtime_error(plugin.error_message("probe")));
299 }
300 if native.device_count > MAX_DEVICES || (native.device_count != 0 && native.devices.is_null()) {
301 return Err(plugin_runtime_error(
302 "runtime plugin probe returned an invalid device list",
303 ));
304 }
305 let native_devices = if native.device_count == 0 {
306 &[]
307 } else {
308 unsafe { std::slice::from_raw_parts(native.devices, native.device_count) }
310 };
311 let devices = native_devices
312 .iter()
313 .map(|device| {
314 Ok(RuntimeDevice {
315 name: copy_utf8(device.name, "device name", 1024)?,
316 kind: device_kind(device.kind)?,
317 memory_bytes: (device.has_memory_bytes != 0).then_some(device.memory_bytes),
318 })
319 })
320 .collect::<PluginRuntimeResult<Vec<_>>>()?;
321 let capabilities = if native.capabilities_json.len == 0 {
322 RuntimeCapabilities::default()
323 } else {
324 let bytes = copy_bytes(
325 native.capabilities_json,
326 "probe capabilities",
327 MAX_STRING_BYTES,
328 )?;
329 serde_json::from_slice(&bytes).map_err(|error| {
330 plugin_runtime_error(format!(
331 "runtime plugin returned invalid capabilities JSON: {error}"
332 ))
333 })?
334 };
335 let version = optional_utf8(native.version, "runtime version", 1024)?;
336 let reason = optional_utf8(native.reason_unavailable, "unavailable reason", 16 * 1024)?;
337 let available = native.available != 0;
338 Ok(RuntimeProbe {
339 available,
340 version,
341 devices,
342 reason_unavailable: if available {
343 None
344 } else {
345 Some(reason.unwrap_or_else(|| "plugin runtime reported unavailable".to_owned()))
346 },
347 capabilities,
348 })
349}
350
351struct RuntimePluginSession {
352 metadata: SessionMetadata,
353 plugin: Arc<LoadedPlugin>,
354 handle: NonNull<c_void>,
355 execution: Mutex<()>,
356}
357
358unsafe impl Send for RuntimePluginSession {}
361unsafe impl Sync for RuntimePluginSession {}
362
363impl RuntimeSession for RuntimePluginSession {
364 fn metadata(&self) -> &SessionMetadata {
365 &self.metadata
366 }
367
368 fn run(&self, request: RunRequest) -> Result<RunResponse> {
369 let _execution = self
370 .execution
371 .lock()
372 .map_err(|_| plugin_runtime_error("runtime plugin session lock was poisoned"))?;
373 validate_run_request(&self.metadata, &request)?;
374 let prepared = self
375 .metadata
376 .inputs
377 .iter()
378 .map(|spec| {
379 let tensor = request
380 .inputs
381 .get(&spec.name)
382 .expect("input names were validated before preparation");
383 PreparedTensor::new(&spec.name, tensor)
384 })
385 .collect::<PluginRuntimeResult<Vec<_>>>()?;
386 let input_views = prepared
387 .iter()
388 .map(PreparedTensor::view)
389 .collect::<Vec<_>>();
390 let method = request.method.as_deref().map(str::as_bytes);
391 let requested_views = request
392 .requested_outputs
393 .as_ref()
394 .map(|names| {
395 names
396 .iter()
397 .map(|name| LatexSnipperBytesV1::from_slice(name.as_bytes()))
398 .collect::<Vec<_>>()
399 })
400 .unwrap_or_default();
401 let native_request = LatexSnipperRunRequestV1 {
402 has_method: u8::from(method.is_some()),
403 method: method
404 .map(LatexSnipperBytesV1::from_slice)
405 .unwrap_or_else(LatexSnipperBytesV1::empty),
406 inputs: input_views.as_ptr(),
407 input_count: input_views.len(),
408 has_requested_outputs: u8::from(request.requested_outputs.is_some()),
409 requested_outputs: requested_views.as_ptr(),
410 requested_output_count: requested_views.len(),
411 };
412 let mut native_output = LatexSnipperOwnedTensorListV1::empty();
413 let run = self
414 .plugin
415 .table
416 .run
417 .expect("required function was validated while loading plugin");
418 let status = unsafe { run(self.handle.as_ptr(), &native_request, &mut native_output) };
420 let output_guard = PluginOutputGuard::new(&self.plugin, self.handle, native_output);
421 if status != LS_RUNTIME_OK {
422 return Err(plugin_runtime_error(self.plugin.error_message("run")));
423 }
424 let views = output_guard.views()?;
425 let declared = self
426 .metadata
427 .outputs
428 .iter()
429 .map(|spec| spec.name.as_str())
430 .collect::<BTreeSet<_>>();
431 let requested = request
432 .requested_outputs
433 .as_ref()
434 .map(|names| names.iter().map(String::as_str).collect::<BTreeSet<_>>())
435 .unwrap_or_else(|| declared.clone());
436 let mut outputs = TensorMap::new();
437 for view in views {
438 let tensor = copy_native_tensor(view)?;
439 let name = tensor.name().to_owned();
440 let Some(spec) = self.metadata.outputs.iter().find(|spec| spec.name == name) else {
441 return Err(plugin_runtime_error(format!(
442 "runtime plugin returned undeclared output '{name}'"
443 )));
444 };
445 if tensor.dtype().as_str() != spec.dtype
446 || tensor.shape().len() != spec.shape.len()
447 || spec
448 .shape
449 .iter()
450 .zip(tensor.shape())
451 .any(|(expected, actual)| {
452 expected.is_some_and(|value| i64::try_from(*actual).ok() != Some(value))
453 })
454 {
455 return Err(plugin_runtime_error(format!(
456 "runtime plugin output '{name}' does not match dtype/shape metadata"
457 )));
458 }
459 if requested.contains(name.as_str()) && outputs.insert(name.clone(), tensor).is_some() {
460 return Err(plugin_runtime_error(format!(
461 "runtime plugin returned duplicate output '{name}'"
462 )));
463 }
464 }
465 if outputs.keys().map(String::as_str).collect::<BTreeSet<_>>() != requested {
466 return Err(plugin_runtime_error(
467 "runtime plugin did not return every requested output",
468 ));
469 }
470 Ok(RunResponse { outputs })
471 }
472}
473
474impl Drop for RuntimePluginSession {
475 fn drop(&mut self) {
476 destroy_handle(&self.plugin, self.handle);
477 }
478}
479
480fn destroy_handle(plugin: &LoadedPlugin, handle: NonNull<c_void>) {
481 let destroy = plugin
482 .table
483 .destroy_session
484 .expect("required function was validated while loading plugin");
485 unsafe { destroy(handle.as_ptr()) };
487}
488
489struct PluginOutputGuard<'plugin> {
490 plugin: &'plugin LoadedPlugin,
491 session: NonNull<c_void>,
492 output: LatexSnipperOwnedTensorListV1,
493 active: bool,
494}
495
496impl<'plugin> PluginOutputGuard<'plugin> {
497 fn new(
498 plugin: &'plugin LoadedPlugin,
499 session: NonNull<c_void>,
500 output: LatexSnipperOwnedTensorListV1,
501 ) -> Self {
502 Self {
503 plugin,
504 session,
505 active: output.owns_allocation(),
506 output,
507 }
508 }
509
510 fn views(&self) -> PluginRuntimeResult<&[LatexSnipperTensorViewV1]> {
511 if self.output.tensor_count > MAX_TENSORS
512 || (self.output.tensor_count != 0 && self.output.tensors.is_null())
513 {
514 return Err(plugin_runtime_error(
515 "runtime plugin returned an invalid output tensor list",
516 ));
517 }
518 if self.output.tensor_count == 0 {
519 Ok(&[])
520 } else {
521 Ok(
524 unsafe {
525 std::slice::from_raw_parts(self.output.tensors, self.output.tensor_count)
526 },
527 )
528 }
529 }
530}
531
532impl Drop for PluginOutputGuard<'_> {
533 fn drop(&mut self) {
534 if !self.active {
535 return;
536 }
537 self.active = false;
538 let free = self
539 .plugin
540 .table
541 .free_output
542 .expect("required function was validated while loading plugin");
543 unsafe { free(self.session.as_ptr(), &mut self.output) };
546 self.output = LatexSnipperOwnedTensorListV1::empty();
547 }
548}
549
550struct PreparedTensor {
551 name: Vec<u8>,
552 shape: Vec<i64>,
553 dtype: i32,
554 bytes: Vec<u8>,
555}
556
557impl PreparedTensor {
558 fn new(name: &str, tensor: &Tensor) -> PluginRuntimeResult<Self> {
559 let count = checked_element_count(tensor.shape())?;
560 if tensor_data_len(tensor.data()) != count {
561 return Err(plugin_runtime_error(format!(
562 "input tensor '{name}' shape/data length mismatch"
563 )));
564 }
565 let shape = tensor
566 .shape()
567 .iter()
568 .map(|dimension| {
569 i64::try_from(*dimension)
570 .map_err(|_| plugin_runtime_error("tensor dimension exceeds i64::MAX"))
571 })
572 .collect::<PluginRuntimeResult<Vec<_>>>()?;
573 let (dtype, bytes) = encode_tensor_data(tensor.data());
574 Ok(Self {
575 name: name.as_bytes().to_vec(),
576 shape,
577 dtype,
578 bytes,
579 })
580 }
581
582 fn view(&self) -> LatexSnipperTensorViewV1 {
583 LatexSnipperTensorViewV1 {
584 name: LatexSnipperBytesV1::from_slice(&self.name),
585 dtype: self.dtype,
586 shape: self.shape.as_ptr(),
587 rank: self.shape.len(),
588 data: if self.bytes.is_empty() {
589 std::ptr::null()
590 } else {
591 self.bytes.as_ptr().cast::<c_void>()
592 },
593 byte_len: self.bytes.len(),
594 }
595 }
596}
597
598fn copy_native_tensor(view: &LatexSnipperTensorViewV1) -> PluginRuntimeResult<Tensor> {
599 let name = copy_utf8(view.name, "output tensor name", 4096)?;
600 let shape = copy_shape(view.shape, view.rank, &name)?;
601 let count = checked_element_count(&shape)?;
602 let width = dtype_width(view.dtype)?;
603 let expected = count
604 .checked_mul(width)
605 .ok_or_else(|| plugin_runtime_error("output tensor byte length overflow"))?;
606 if view.byte_len != expected || (view.byte_len != 0 && view.data.is_null()) {
607 return Err(plugin_runtime_error(format!(
608 "output tensor '{name}' requires {expected} bytes, plugin returned {}",
609 view.byte_len
610 )));
611 }
612 let bytes = if view.byte_len == 0 {
613 &[]
614 } else {
615 unsafe { std::slice::from_raw_parts(view.data.cast::<u8>(), view.byte_len) }
617 };
618 match view.dtype {
619 LS_DTYPE_FLOAT32 => Ok(Tensor::float32(
620 &name,
621 shape,
622 decode::<4, f32>(bytes, f32::from_ne_bytes),
623 )),
624 LS_DTYPE_FLOAT16 => Ok(Tensor::float16_bits(
625 &name,
626 shape,
627 decode::<2, u16>(bytes, u16::from_ne_bytes),
628 )),
629 LS_DTYPE_INT64 => Ok(Tensor::int64(
630 &name,
631 shape,
632 decode::<8, i64>(bytes, i64::from_ne_bytes),
633 )),
634 LS_DTYPE_INT32 => Ok(Tensor::int32(
635 &name,
636 shape,
637 decode::<4, i32>(bytes, i32::from_ne_bytes),
638 )),
639 LS_DTYPE_UINT8 => Ok(Tensor::u8(&name, shape, bytes.to_vec())),
640 LS_DTYPE_BOOL => Ok(Tensor::boolean(
641 &name,
642 shape,
643 bytes.iter().map(|byte| *byte != 0).collect(),
644 )),
645 _ => unreachable!("dtype width validation accepted only known codes"),
646 }
647}
648
649fn validate_metadata(metadata: &SessionMetadata) -> PluginRuntimeResult<()> {
650 for (direction, specs) in [("input", &metadata.inputs), ("output", &metadata.outputs)] {
651 let mut names = HashSet::new();
652 for spec in specs {
653 if spec.name.is_empty() || !names.insert(spec.name.as_str()) {
654 return Err(plugin_runtime_error(format!(
655 "runtime plugin metadata contains an empty or duplicate {direction} name"
656 )));
657 }
658 if dtype_code(&spec.dtype).is_none() {
659 return Err(plugin_runtime_error(format!(
660 "runtime plugin metadata uses unsupported dtype '{}'",
661 spec.dtype
662 )));
663 }
664 }
665 }
666 let mut methods = HashSet::new();
667 if metadata
668 .methods
669 .iter()
670 .any(|method| method.is_empty() || !methods.insert(method.as_str()))
671 {
672 return Err(plugin_runtime_error(
673 "runtime plugin metadata contains an empty or duplicate method",
674 ));
675 }
676 Ok(())
677}
678
679fn validate_run_request(
680 metadata: &SessionMetadata,
681 request: &RunRequest,
682) -> PluginRuntimeResult<()> {
683 if let Some(method) = &request.method {
684 if !metadata.methods.iter().any(|candidate| candidate == method) {
685 return Err(plugin_runtime_error(format!(
686 "runtime plugin session does not declare method '{method}'"
687 )));
688 }
689 }
690 let expected = metadata
691 .inputs
692 .iter()
693 .map(|spec| spec.name.as_str())
694 .collect::<BTreeSet<_>>();
695 let provided = request
696 .inputs
697 .keys()
698 .map(String::as_str)
699 .collect::<BTreeSet<_>>();
700 if expected != provided {
701 return Err(plugin_runtime_error(format!(
702 "runtime plugin input names differ; missing={:?}, unexpected={:?}",
703 expected.difference(&provided).collect::<Vec<_>>(),
704 provided.difference(&expected).collect::<Vec<_>>()
705 )));
706 }
707 for spec in &metadata.inputs {
708 let tensor = &request.inputs[&spec.name];
709 if tensor.dtype().as_str() != spec.dtype
710 || tensor.shape().len() != spec.shape.len()
711 || spec
712 .shape
713 .iter()
714 .zip(tensor.shape())
715 .any(|(expected, actual)| {
716 expected.is_some_and(|value| i64::try_from(*actual).ok() != Some(value))
717 })
718 {
719 return Err(plugin_runtime_error(format!(
720 "runtime plugin input '{}' does not match dtype/shape metadata",
721 spec.name
722 )));
723 }
724 }
725 if let Some(requested) = &request.requested_outputs {
726 let declared = metadata
727 .outputs
728 .iter()
729 .map(|spec| spec.name.as_str())
730 .collect::<HashSet<_>>();
731 let mut seen = HashSet::new();
732 if requested
733 .iter()
734 .any(|name| !declared.contains(name.as_str()) || !seen.insert(name.as_str()))
735 {
736 return Err(plugin_runtime_error(
737 "runtime plugin request contains an undeclared or duplicate output",
738 ));
739 }
740 }
741 Ok(())
742}
743
744fn copy_shape(pointer: *const i64, rank: usize, name: &str) -> PluginRuntimeResult<Vec<usize>> {
745 if rank > 64 || (rank != 0 && pointer.is_null()) {
746 return Err(plugin_runtime_error(format!(
747 "output tensor '{name}' has an invalid shape view"
748 )));
749 }
750 if rank == 0 {
751 return Ok(Vec::new());
752 }
753 unsafe { std::slice::from_raw_parts(pointer, rank) }
755 .iter()
756 .map(|dimension| {
757 usize::try_from(*dimension).map_err(|_| {
758 plugin_runtime_error(format!(
759 "output tensor '{name}' has invalid dimension {dimension}"
760 ))
761 })
762 })
763 .collect()
764}
765
766fn checked_element_count(shape: &[usize]) -> PluginRuntimeResult<usize> {
767 shape.iter().try_fold(1usize, |count, dimension| {
768 count
769 .checked_mul(*dimension)
770 .ok_or_else(|| plugin_runtime_error("tensor element count overflow"))
771 })
772}
773
774fn tensor_data_len(data: &TensorData) -> usize {
775 match data {
776 TensorData::Float32(values) => values.len(),
777 TensorData::Float16(values) => values.len(),
778 TensorData::Int64(values) => values.len(),
779 TensorData::Int32(values) => values.len(),
780 TensorData::UInt8(values) => values.len(),
781 TensorData::Bool(values) => values.len(),
782 }
783}
784
785fn encode_tensor_data(data: &TensorData) -> (i32, Vec<u8>) {
786 match data {
787 TensorData::Float32(values) => (LS_DTYPE_FLOAT32, encode(values, f32::to_ne_bytes)),
788 TensorData::Float16(values) => (LS_DTYPE_FLOAT16, encode(values, u16::to_ne_bytes)),
789 TensorData::Int64(values) => (LS_DTYPE_INT64, encode(values, i64::to_ne_bytes)),
790 TensorData::Int32(values) => (LS_DTYPE_INT32, encode(values, i32::to_ne_bytes)),
791 TensorData::UInt8(values) => (LS_DTYPE_UINT8, values.clone()),
792 TensorData::Bool(values) => (
793 LS_DTYPE_BOOL,
794 values.iter().copied().map(u8::from).collect(),
795 ),
796 }
797}
798
799fn encode<T, const WIDTH: usize>(values: &[T], to_bytes: impl Fn(T) -> [u8; WIDTH]) -> Vec<u8>
800where
801 T: Copy,
802{
803 values.iter().copied().flat_map(to_bytes).collect()
804}
805
806fn decode<const WIDTH: usize, T>(bytes: &[u8], from_bytes: impl Fn([u8; WIDTH]) -> T) -> Vec<T> {
807 bytes
808 .chunks_exact(WIDTH)
809 .map(|chunk| {
810 let mut value = [0; WIDTH];
811 value.copy_from_slice(chunk);
812 from_bytes(value)
813 })
814 .collect()
815}
816
817fn dtype_width(dtype: i32) -> PluginRuntimeResult<usize> {
818 match dtype {
819 LS_DTYPE_FLOAT32 | LS_DTYPE_INT32 => Ok(4),
820 LS_DTYPE_FLOAT16 => Ok(2),
821 LS_DTYPE_INT64 => Ok(8),
822 LS_DTYPE_UINT8 | LS_DTYPE_BOOL => Ok(1),
823 other => Err(plugin_runtime_error(format!(
824 "runtime plugin returned unknown dtype code {other}"
825 ))),
826 }
827}
828
829fn dtype_code(dtype: &str) -> Option<i32> {
830 match dtype {
831 "f32" => Some(LS_DTYPE_FLOAT32),
832 "f16" => Some(LS_DTYPE_FLOAT16),
833 "i64" => Some(LS_DTYPE_INT64),
834 "i32" => Some(LS_DTYPE_INT32),
835 "u8" => Some(LS_DTYPE_UINT8),
836 "bool" => Some(LS_DTYPE_BOOL),
837 _ => None,
838 }
839}
840
841fn device_kind(kind: i32) -> PluginRuntimeResult<DeviceKind> {
842 match kind {
843 LS_DEVICE_AUTO => Ok(DeviceKind::Auto),
844 LS_DEVICE_CPU => Ok(DeviceKind::Cpu),
845 LS_DEVICE_GPU => Ok(DeviceKind::Gpu),
846 LS_DEVICE_NPU => Ok(DeviceKind::Npu),
847 other => Err(plugin_runtime_error(format!(
848 "runtime plugin returned unknown device kind {other}"
849 ))),
850 }
851}
852
853fn optional_utf8(
854 view: LatexSnipperBytesV1,
855 label: &str,
856 maximum: usize,
857) -> PluginRuntimeResult<Option<String>> {
858 if view.len == 0 {
859 Ok(None)
860 } else {
861 copy_utf8(view, label, maximum).map(Some)
862 }
863}
864
865fn copy_utf8(
866 view: LatexSnipperBytesV1,
867 label: &str,
868 maximum: usize,
869) -> PluginRuntimeResult<String> {
870 let bytes = copy_bytes(view, label, maximum)?;
871 String::from_utf8(bytes)
872 .map_err(|_| plugin_runtime_error(format!("runtime plugin {label} is not valid UTF-8")))
873}
874
875fn copy_bytes(
876 view: LatexSnipperBytesV1,
877 label: &str,
878 maximum: usize,
879) -> PluginRuntimeResult<Vec<u8>> {
880 if view.len > maximum || (view.len != 0 && view.data.is_null()) {
881 return Err(plugin_runtime_error(format!(
882 "runtime plugin {label} byte view is invalid or exceeds {maximum} bytes"
883 )));
884 }
885 if view.len == 0 {
886 Ok(Vec::new())
887 } else {
888 Ok(unsafe { std::slice::from_raw_parts(view.data, view.len) }.to_vec())
891 }
892}