1use crate::indexing::plan::{build_index_plan, IndexPlan};
2use crate::indexing::selectors::{
3 build_slice_selectors, index_scalar_from_value, materialize_index_value, SliceSelector,
4};
5use runmat_builtins::{CellArray, ComplexTensor, StringArray, Tensor, Value};
6use runmat_runtime::RuntimeError;
7
8pub fn build_numeric_subsref_cell(numeric: &[Value]) -> Result<Value, RuntimeError> {
9 let cell = CellArray::new(numeric.to_vec(), 1, numeric.len())
10 .map_err(|e| format!("subsref build error: {e}"))?;
11 Ok(Value::Cell(cell))
12}
13
14pub async fn object_subsref_paren(base: Value, numeric: &[Value]) -> Result<Value, RuntimeError> {
15 let cell = build_numeric_subsref_cell(numeric)?;
16 match base {
17 Value::Object(obj) => {
18 let args = vec![
19 Value::Object(obj),
20 Value::String("subsref".to_string()),
21 Value::String("()".to_string()),
22 cell,
23 ];
24 runmat_runtime::call_builtin_async("call_method", &args).await
25 }
26 Value::HandleObject(handle) => {
27 let args = vec![
28 Value::HandleObject(handle),
29 Value::String("subsref".to_string()),
30 Value::String("()".to_string()),
31 cell,
32 ];
33 runmat_runtime::call_builtin_async("call_method", &args).await
34 }
35 other => Err(format!("slice subsref requires object/handle, got {other:?}").into()),
36 }
37}
38
39pub async fn read_tensor_slice_1d(
40 tensor: &Tensor,
41 colon_mask: u32,
42 end_mask: u32,
43 numeric: &[Value],
44) -> Result<Value, RuntimeError> {
45 let total = tensor.data.len();
46 let mut idxs: Vec<usize> = Vec::new();
47 let mut idx_shape: Option<Vec<usize>> = None;
48 let is_colon = (colon_mask & 1u32) != 0;
49 let is_end = (end_mask & 1u32) != 0;
50 if is_colon {
51 idxs = (1..=total).collect();
52 } else if is_end {
53 idxs = vec![total];
54 } else if let Some(v) = numeric.first() {
55 let materialized = materialize_index_value(v).await?;
56 if let Some(i) = index_scalar_from_value(&materialized).await? {
57 if i < 1 {
58 return Err(crate::interpreter::errors::mex(
59 "IndexOutOfBounds",
60 "Index out of bounds",
61 ));
62 }
63 idxs = vec![i as usize];
64 } else {
65 match &materialized {
66 Value::Tensor(idx_t) => {
67 idx_shape = Some(idx_t.shape.clone());
68 for &val in &idx_t.data {
69 let i = val as isize;
70 if i < 1 || (i as usize) > total {
71 return Err(crate::interpreter::errors::mex(
72 "IndexOutOfBounds",
73 "Index out of bounds",
74 ));
75 }
76 idxs.push(i as usize);
77 }
78 }
79 Value::Bool(b) => {
80 if *b {
81 idxs = vec![1];
82 }
83 }
84 Value::LogicalArray(la) => {
85 if la.data.len() != total {
86 return Err(crate::interpreter::errors::mex(
87 "IndexShape",
88 "Logical mask length mismatch for linear indexing",
89 ));
90 }
91 for (i, &val) in la.data.iter().enumerate() {
92 if val != 0 {
93 idxs.push(i + 1);
94 }
95 }
96 }
97 _ => {
98 return Err(crate::interpreter::errors::mex(
99 "UnsupportedIndexType",
100 "Unsupported index type",
101 ))
102 }
103 }
104 }
105 } else {
106 return Err(crate::interpreter::errors::mex(
107 "MissingNumericIndex",
108 "missing numeric index",
109 ));
110 }
111 if idxs.iter().any(|&i| i == 0 || i > total) {
112 return Err(crate::interpreter::errors::mex(
113 "IndexOutOfBounds",
114 "Index out of bounds",
115 ));
116 }
117 if idxs.len() == 1 {
118 Ok(Value::Num(tensor.data[idxs[0] - 1]))
119 } else if idxs.is_empty() {
120 let shape = idx_shape.unwrap_or_else(|| vec![0, 1]);
121 let tens = Tensor::new(vec![], shape).map_err(|e| format!("Slice error: {e}"))?;
122 Ok(Value::Tensor(tens))
123 } else {
124 let mut out = Vec::with_capacity(idxs.len());
125 for &i in &idxs {
126 out.push(tensor.data[i - 1]);
127 }
128 let shape = idx_shape.unwrap_or_else(|| vec![idxs.len(), 1]);
129 let tens = Tensor::new(out, shape).map_err(|e| format!("Slice error: {e}"))?;
130 Ok(Value::Tensor(tens))
131 }
132}
133
134pub fn try_tensor_slice_2d_fast_path(
135 tensor: &Tensor,
136 dims: usize,
137 selectors: &[SliceSelector],
138) -> Result<Option<Value>, RuntimeError> {
139 if dims != 2 {
140 return Ok(None);
141 }
142 let rows = tensor.shape.first().copied().unwrap_or(1);
143 let cols = tensor.shape.get(1).copied().unwrap_or(1);
144 match (&selectors[0], &selectors[1]) {
145 (SliceSelector::Colon, SliceSelector::Scalar(j)) => {
146 let j0 = *j - 1;
147 if j0 >= cols {
148 return Err(crate::interpreter::errors::mex(
149 "IndexOutOfBounds",
150 "Index out of bounds",
151 ));
152 }
153 let start = j0 * rows;
154 let out = tensor.data[start..start + rows].to_vec();
155 if out.len() == 1 {
156 Ok(Some(Value::Num(out[0])))
157 } else {
158 let tens =
159 Tensor::new(out, vec![rows, 1]).map_err(|e| format!("Slice error: {e}"))?;
160 Ok(Some(Value::Tensor(tens)))
161 }
162 }
163 (SliceSelector::Scalar(i), SliceSelector::Colon) => {
164 let i0 = *i - 1;
165 if i0 >= rows {
166 return Err(crate::interpreter::errors::mex(
167 "IndexOutOfBounds",
168 "Index out of bounds",
169 ));
170 }
171 let mut out: Vec<f64> = Vec::with_capacity(cols);
172 for c in 0..cols {
173 out.push(tensor.data[i0 + c * rows]);
174 }
175 if out.len() == 1 {
176 Ok(Some(Value::Num(out[0])))
177 } else {
178 let tens =
179 Tensor::new(out, vec![1, cols]).map_err(|e| format!("Slice error: {e}"))?;
180 Ok(Some(Value::Tensor(tens)))
181 }
182 }
183 (SliceSelector::Colon, SliceSelector::Indices(js)) => {
184 if js.is_empty() {
185 let tens = Tensor::new(Vec::new(), vec![rows, 0])
186 .map_err(|e| format!("Slice error: {e}"))?;
187 Ok(Some(Value::Tensor(tens)))
188 } else {
189 let mut out: Vec<f64> = Vec::with_capacity(rows * js.len());
190 for &j in js {
191 let j0 = j - 1;
192 if j0 >= cols {
193 return Err(crate::interpreter::errors::mex(
194 "IndexOutOfBounds",
195 "Index out of bounds",
196 ));
197 }
198 let start = j0 * rows;
199 out.extend_from_slice(&tensor.data[start..start + rows]);
200 }
201 let tens = Tensor::new(out, vec![rows, js.len()])
202 .map_err(|e| format!("Slice error: {e}"))?;
203 Ok(Some(Value::Tensor(tens)))
204 }
205 }
206 (SliceSelector::Indices(is), SliceSelector::Colon) => {
207 if is.is_empty() {
208 let tens = Tensor::new(Vec::new(), vec![0, cols])
209 .map_err(|e| format!("Slice error: {e}"))?;
210 Ok(Some(Value::Tensor(tens)))
211 } else {
212 let mut out: Vec<f64> = Vec::with_capacity(is.len() * cols);
213 for c in 0..cols {
214 for &i in is {
215 let i0 = i - 1;
216 if i0 >= rows {
217 return Err(crate::interpreter::errors::mex(
218 "IndexOutOfBounds",
219 "Index out of bounds",
220 ));
221 }
222 out.push(tensor.data[i0 + c * rows]);
223 }
224 }
225 let tens = Tensor::new(out, vec![is.len(), cols])
226 .map_err(|e| format!("Slice error: {e}"))?;
227 Ok(Some(Value::Tensor(tens)))
228 }
229 }
230 _ => Ok(None),
231 }
232}
233
234pub async fn read_tensor_slice_nd(
235 tensor: &Tensor,
236 dims: usize,
237 colon_mask: u32,
238 end_mask: u32,
239 numeric: &[Value],
240) -> Result<Value, RuntimeError> {
241 let selectors =
242 build_slice_selectors(dims, colon_mask, end_mask, numeric, &tensor.shape).await?;
243 if let Some(value) = try_tensor_slice_2d_fast_path(tensor, dims, &selectors)? {
244 return Ok(value);
245 }
246 let plan = build_index_plan(&selectors, dims, &tensor.shape)?;
247 if plan.indices.is_empty() {
248 let out_tensor =
249 Tensor::new(Vec::new(), plan.output_shape).map_err(|e| format!("Slice error: {e}"))?;
250 return Ok(Value::Tensor(out_tensor));
251 }
252 let mut out_data: Vec<f64> = Vec::with_capacity(plan.indices.len());
253 for &lin in &plan.indices {
254 out_data.push(tensor.data[lin as usize]);
255 }
256 if out_data.len() == 1 {
257 Ok(Value::Num(out_data[0]))
258 } else {
259 let out_tensor =
260 Tensor::new(out_data, plan.output_shape).map_err(|e| format!("Slice error: {e}"))?;
261 Ok(Value::Tensor(out_tensor))
262 }
263}
264
265pub fn read_tensor_slice_from_plan(
266 tensor: &Tensor,
267 plan: &IndexPlan,
268) -> Result<Value, RuntimeError> {
269 if plan.indices.is_empty() {
270 let out_tensor = Tensor::new(Vec::new(), plan.output_shape.clone())
271 .map_err(|e| format!("Slice error: {e}"))?;
272 return Ok(Value::Tensor(out_tensor));
273 }
274 let mut out_data: Vec<f64> = Vec::with_capacity(plan.indices.len());
275 for &lin in &plan.indices {
276 out_data.push(tensor.data[lin as usize]);
277 }
278 if out_data.len() == 1 {
279 Ok(Value::Num(out_data[0]))
280 } else {
281 let out_tensor = Tensor::new(out_data, plan.output_shape.clone())
282 .map_err(|e| format!("Slice error: {e}"))?;
283 Ok(Value::Tensor(out_tensor))
284 }
285}
286
287pub async fn read_complex_slice(
288 tensor: &ComplexTensor,
289 dims: usize,
290 colon_mask: u32,
291 end_mask: u32,
292 numeric: &[Value],
293) -> Result<Value, RuntimeError> {
294 let selectors =
295 build_slice_selectors(dims, colon_mask, end_mask, numeric, &tensor.shape).await?;
296 let plan = build_index_plan(&selectors, dims, &tensor.shape)?;
297 read_complex_slice_from_plan(tensor, &plan)
298}
299
300pub fn read_complex_slice_from_plan(
301 tensor: &ComplexTensor,
302 plan: &IndexPlan,
303) -> Result<Value, RuntimeError> {
304 if plan.indices.is_empty() {
305 let empty = ComplexTensor::new(Vec::new(), plan.output_shape.clone())
306 .map_err(|e| format!("Slice error: {e}"))?;
307 return Ok(Value::ComplexTensor(empty));
308 }
309 if plan.indices.len() == 1 {
310 let lin = plan.indices[0] as usize;
311 let (re, im) = tensor.data.get(lin).copied().ok_or_else(|| {
312 crate::interpreter::errors::mex(
313 "IndexOutOfBounds",
314 "Slice error: complex index out of bounds",
315 )
316 })?;
317 return Ok(Value::Complex(re, im));
318 }
319 let mut out = Vec::with_capacity(plan.indices.len());
320 for &lin in &plan.indices {
321 let idx = lin as usize;
322 let value = tensor.data.get(idx).copied().ok_or_else(|| {
323 crate::interpreter::errors::mex(
324 "IndexOutOfBounds",
325 "Slice error: complex index out of bounds",
326 )
327 })?;
328 out.push(value);
329 }
330 let out_ct = ComplexTensor::new(out, plan.output_shape.clone())
331 .map_err(|e| format!("Slice error: {e}"))?;
332 Ok(Value::ComplexTensor(out_ct))
333}
334
335pub async fn read_gpu_slice(
336 handle: &runmat_accelerate_api::GpuTensorHandle,
337 dims: usize,
338 colon_mask: u32,
339 end_mask: u32,
340 numeric: &[Value],
341) -> Result<Value, RuntimeError> {
342 let base_shape = handle.shape.clone();
343 let selectors = build_slice_selectors(dims, colon_mask, end_mask, numeric, &base_shape).await?;
344 let plan = build_index_plan(&selectors, dims, &base_shape)?;
345 read_gpu_slice_from_plan(handle, &plan)
346}
347
348pub fn read_gpu_slice_from_plan(
349 handle: &runmat_accelerate_api::GpuTensorHandle,
350 plan: &IndexPlan,
351) -> Result<Value, RuntimeError> {
352 let provider = runmat_accelerate_api::provider().ok_or_else(|| {
353 crate::interpreter::errors::mex(
354 "AccelerationProviderUnavailable",
355 "No acceleration provider registered",
356 )
357 })?;
358 if plan.indices.is_empty() {
359 let zeros = provider
360 .zeros(&plan.output_shape)
361 .map_err(|e| format!("slice: {e}"))?;
362 Ok(Value::GpuTensor(zeros))
363 } else {
364 let result = provider
365 .gather_linear(handle, &plan.indices, &plan.output_shape)
366 .map_err(|e| format!("slice: {e}"))?;
367 Ok(Value::GpuTensor(result))
368 }
369}
370
371pub async fn read_string_slice(
372 sa: &StringArray,
373 dims: usize,
374 colon_mask: u32,
375 end_mask: u32,
376 numeric: &[Value],
377) -> Result<Value, RuntimeError> {
378 let rank = sa.shape.len();
379 if dims == 1 {
380 let total = sa.data.len();
381 let mut idxs: Vec<usize> = Vec::new();
382 let is_colon = (colon_mask & 1u32) != 0;
383 let is_end = (end_mask & 1u32) != 0;
384 if is_colon {
385 idxs = (1..=total).collect();
386 } else if is_end {
387 idxs = vec![total];
388 } else if let Some(v) = numeric.first() {
389 let materialized = materialize_index_value(v).await?;
390 if let Some(i) = index_scalar_from_value(&materialized).await? {
391 if i < 1 {
392 return Err(crate::interpreter::errors::mex(
393 "IndexOutOfBounds",
394 "Index out of bounds",
395 ));
396 }
397 idxs = vec![i as usize];
398 } else {
399 match &materialized {
400 Value::Tensor(idx_t) => {
401 let len = idx_t.shape.iter().product::<usize>();
402 if len == total {
403 for (i, &val) in idx_t.data.iter().enumerate() {
404 if val != 0.0 {
405 idxs.push(i + 1);
406 }
407 }
408 } else {
409 for &val in &idx_t.data {
410 let i = val as isize;
411 if i < 1 {
412 return Err(crate::interpreter::errors::mex(
413 "IndexOutOfBounds",
414 "Index out of bounds",
415 ));
416 }
417 idxs.push(i as usize);
418 }
419 }
420 }
421 _ => {
422 return Err(crate::interpreter::errors::mex(
423 "UnsupportedIndexType",
424 "Unsupported index type",
425 ))
426 }
427 }
428 }
429 } else {
430 return Err(crate::interpreter::errors::mex(
431 "MissingNumericIndex",
432 "missing numeric index",
433 ));
434 }
435 if idxs.iter().any(|&i| i == 0 || i > total) {
436 return Err(crate::interpreter::errors::mex(
437 "IndexOutOfBounds",
438 "Index out of bounds",
439 ));
440 }
441 if idxs.len() == 1 {
442 Ok(Value::String(sa.data[idxs[0] - 1].clone()))
443 } else {
444 let mut out: Vec<String> = Vec::with_capacity(idxs.len());
445 for &i in &idxs {
446 out.push(sa.data[i - 1].clone());
447 }
448 let out_sa = StringArray::new(out, vec![idxs.len(), 1])
449 .map_err(|e| format!("Slice error: {e}"))?;
450 Ok(Value::StringArray(out_sa))
451 }
452 } else {
453 let mut selectors: Vec<SliceSelector> = Vec::with_capacity(dims);
454 let mut num_iter = 0usize;
455 for d in 0..dims {
456 let is_colon = (colon_mask & (1u32 << d)) != 0;
457 let is_end = (end_mask & (1u32 << d)) != 0;
458 if is_colon {
459 selectors.push(SliceSelector::Colon);
460 } else if is_end {
461 let dim_len = *sa.shape.get(d).unwrap_or(&1);
462 selectors.push(SliceSelector::Scalar(dim_len));
463 } else {
464 let v = numeric.get(num_iter).ok_or_else(|| {
465 crate::interpreter::errors::mex("MissingNumericIndex", "missing numeric index")
466 })?;
467 num_iter += 1;
468 let materialized = materialize_index_value(v).await?;
469 if let Some(idx) = index_scalar_from_value(&materialized).await? {
470 if idx < 1 {
471 return Err(crate::interpreter::errors::mex(
472 "IndexOutOfBounds",
473 "Index out of bounds",
474 ));
475 }
476 selectors.push(SliceSelector::Scalar(idx as usize));
477 } else {
478 match &materialized {
479 Value::Tensor(idx_t) => {
480 let dim_len = *sa.shape.get(d).unwrap_or(&1);
481 let len = idx_t.shape.iter().product::<usize>();
482 let is_binary_mask =
483 len == dim_len && idx_t.data.iter().all(|&x| x == 0.0 || x == 1.0);
484 if is_binary_mask {
485 let mut v = Vec::new();
486 for (i, &val) in idx_t.data.iter().enumerate() {
487 if val != 0.0 {
488 v.push(i + 1);
489 }
490 }
491 selectors.push(SliceSelector::Indices(v));
492 } else {
493 let mut v = Vec::with_capacity(len);
494 for &val in &idx_t.data {
495 let idx = val as isize;
496 if idx < 1 {
497 return Err(crate::interpreter::errors::mex(
498 "IndexOutOfBounds",
499 "Index out of bounds",
500 ));
501 }
502 v.push(idx as usize);
503 }
504 selectors.push(SliceSelector::Indices(v));
505 }
506 }
507 _ => {
508 return Err(crate::interpreter::errors::mex(
509 "UnsupportedIndexType",
510 "Unsupported index type",
511 ))
512 }
513 }
514 }
515 }
516 }
517
518 let mut out_dims: Vec<usize> = Vec::new();
519 let mut per_dim_indices: Vec<Vec<usize>> = Vec::with_capacity(dims);
520 for (d, sel) in selectors.iter().enumerate().take(dims) {
521 let dim_len = *sa.shape.get(d).unwrap_or(&1);
522 let idxs = match sel {
523 SliceSelector::Colon => (1..=dim_len).collect::<Vec<usize>>(),
524 SliceSelector::Scalar(i) => vec![*i],
525 SliceSelector::Indices(v) => v.clone(),
526 SliceSelector::LinearIndices { values, .. } => values.clone(),
527 };
528 if idxs.iter().any(|&i| i == 0 || i > dim_len) {
529 return Err(crate::interpreter::errors::mex(
530 "IndexOutOfBounds",
531 "Index out of bounds",
532 ));
533 }
534 if idxs.len() > 1 {
535 out_dims.push(idxs.len());
536 } else {
537 out_dims.push(1);
538 }
539 per_dim_indices.push(idxs);
540 }
541 if dims == 2 {
542 match (
543 &per_dim_indices[0].as_slice(),
544 &per_dim_indices[1].as_slice(),
545 ) {
546 (i_list, j_list) if i_list.len() > 1 && j_list.len() == 1 => {
547 out_dims = vec![i_list.len(), 1];
548 }
549 (i_list, j_list) if i_list.len() == 1 && j_list.len() > 1 => {
550 out_dims = vec![1, j_list.len()];
551 }
552 _ => {}
553 }
554 }
555 let full_shape: Vec<usize> = if rank < dims {
556 let mut s = sa.shape.clone();
557 s.resize(dims, 1);
558 s
559 } else {
560 sa.shape.clone()
561 };
562 let mut strides: Vec<usize> = vec![0; dims];
563 let mut acc = 1usize;
564 for (d, stride) in strides.iter_mut().enumerate().take(dims) {
565 *stride = acc;
566 acc *= full_shape[d];
567 }
568 let total_out: usize = out_dims.iter().product();
569 if total_out == 0 {
570 return Ok(Value::StringArray(
571 StringArray::new(Vec::new(), out_dims).map_err(|e| format!("Slice error: {e}"))?,
572 ));
573 }
574 let mut out_data: Vec<String> = Vec::with_capacity(total_out);
575 let mut idx = vec![0usize; dims];
576 loop {
577 let current: Vec<usize> = (0..dims).map(|d| per_dim_indices[d][idx[d]]).collect();
578 let mut lin = 0usize;
579 for d in 0..dims {
580 let i0 = current[d] - 1;
581 lin += i0 * strides[d];
582 }
583 out_data.push(sa.data[lin].clone());
584 let mut d = 0usize;
585 while d < dims {
586 idx[d] += 1;
587 if idx[d] < per_dim_indices[d].len() {
588 break;
589 }
590 idx[d] = 0;
591 d += 1;
592 }
593 if d == dims {
594 break;
595 }
596 }
597 if out_data.len() == 1 {
598 Ok(Value::String(out_data[0].clone()))
599 } else {
600 let out_sa =
601 StringArray::new(out_data, out_dims).map_err(|e| format!("Slice error: {e}"))?;
602 Ok(Value::StringArray(out_sa))
603 }
604 }
605}
606
607pub fn gather_string_slice(sa: &StringArray, plan: &IndexPlan) -> Result<Value, RuntimeError> {
608 if plan.indices.is_empty() {
609 let empty = StringArray::new(Vec::new(), plan.output_shape.clone())
610 .map_err(|e| format!("Slice error: {e}"))?;
611 return Ok(Value::StringArray(empty));
612 }
613 if plan.indices.len() == 1 {
614 let lin = plan.indices[0] as usize;
615 let value = sa.data.get(lin).cloned().ok_or_else(|| {
616 crate::interpreter::errors::mex(
617 "IndexOutOfBounds",
618 "Slice error: string index out of bounds",
619 )
620 })?;
621 return Ok(Value::String(value));
622 }
623 let mut out = Vec::with_capacity(plan.indices.len());
624 for &lin in &plan.indices {
625 let idx = lin as usize;
626 let value = sa.data.get(idx).cloned().ok_or_else(|| {
627 crate::interpreter::errors::mex(
628 "IndexOutOfBounds",
629 "Slice error: string index out of bounds",
630 )
631 })?;
632 out.push(value);
633 }
634 let out_sa = StringArray::new(out, plan.output_shape.clone())
635 .map_err(|e| format!("Slice error: {e}"))?;
636 Ok(Value::StringArray(out_sa))
637}