1use std::any::Any;
9use std::ops::Deref;
10use std::sync::atomic::Ordering;
11use std::sync::{Arc, OnceLock};
12
13use crate::complex::Cx;
14use crate::dtype::DType;
15use crate::exact::{Ext, Rat};
16
17pub type Owner = Arc<dyn Any + Send + Sync>;
20
21static JOINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
28
29static LAYOUTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
34
35pub fn joins_made() -> u64 {
37 JOINS.load(Ordering::Relaxed)
38}
39
40pub fn layouts_made() -> u64 {
43 LAYOUTS.load(Ordering::Relaxed)
44}
45
46pub struct Buf<T> {
57 repr: Repr<T>,
58}
59
60enum Repr<T> {
61 Owned(Arc<Vec<T>>),
63 Slice { buf: Arc<Vec<T>>, off: usize, len: usize },
68 Foreign { ptr: *const T, len: usize, owner: Owner },
69 Cols {
79 parts: Vec<Buf<T>>,
80 len: usize,
81 flat: Arc<OnceLock<Arc<Vec<T>>>>,
82 join: fn(&[Buf<T>], usize) -> Vec<T>,
83 },
84}
85
86unsafe impl<T: Send + Sync> Send for Buf<T> {}
97unsafe impl<T: Send + Sync> Sync for Buf<T> {}
99
100impl<T> Buf<T> {
101 pub fn new() -> Buf<T> {
102 Buf { repr: Repr::Owned(Arc::new(Vec::new())) }
103 }
104
105 pub fn from_vec(v: Vec<T>) -> Buf<T> {
106 Buf { repr: Repr::Owned(Arc::new(v)) }
107 }
108
109 pub unsafe fn foreign(ptr: *const T, len: usize, owner: Owner) -> Buf<T> {
117 Buf { repr: Repr::Foreign { ptr, len, owner } }
118 }
119
120 pub fn is_foreign(&self) -> bool {
124 match &self.repr {
125 Repr::Foreign { .. } => true,
126 Repr::Cols { parts, flat, .. } => {
127 flat.get().is_none() && parts.iter().any(Buf::is_foreign)
128 }
129 _ => false,
130 }
131 }
132
133 pub fn len(&self) -> usize {
135 match &self.repr {
136 Repr::Owned(v) => v.len(),
137 Repr::Slice { len, .. } | Repr::Foreign { len, .. } | Repr::Cols { len, .. } => *len,
138 }
139 }
140
141 pub fn is_empty(&self) -> bool {
142 self.len() == 0
143 }
144
145 pub fn is_joined(&self) -> bool {
148 matches!(&self.repr, Repr::Cols { flat, .. } if flat.get().is_some())
149 }
150
151 pub fn parts(&self) -> Option<&[Buf<T>]> {
155 match &self.repr {
156 Repr::Cols { parts, .. } => Some(parts),
157 _ => None,
158 }
159 }
160
161 pub fn owner(&self) -> Option<&Owner> {
169 match &self.repr {
170 Repr::Foreign { owner, .. } => Some(owner),
171 _ => None,
172 }
173 }
174}
175
176fn join_sequential<T: Clone>(parts: &[Buf<T>], len: usize) -> Vec<T> {
179 let mut v = Vec::with_capacity(len);
180 for part in parts {
181 v.extend_from_slice(part.as_slice());
182 }
183 v
184}
185
186fn join_parallel<T: Copy + Default + Send + Sync>(parts: &[Buf<T>], len: usize) -> Vec<T> {
190 let slices: Vec<&[T]> = parts.iter().map(Buf::as_slice).collect();
191 let (out, ok) = crate::par::fill(len, |start, dst: &mut [T]| {
192 let mut at = 0;
193 let mut written = 0;
194 for s in &slices {
195 let (from, to) = (at, at + s.len());
196 at = to;
197 let lo = start.max(from);
198 let hi = (start + dst.len()).min(to);
199 if lo < hi {
200 dst[lo - start..hi - start].copy_from_slice(&s[lo - from..hi - from]);
201 written += hi - lo;
202 }
203 }
204 written == dst.len()
205 });
206 debug_assert!(ok, "the parts do not cover the join");
207 out
208}
209
210impl<T: Clone> Buf<T> {
211 pub fn join(parts: Vec<Buf<T>>) -> Buf<T> {
215 Buf::joined(parts, join_sequential)
216 }
217
218 fn joined(parts: Vec<Buf<T>>, join: fn(&[Buf<T>], usize) -> Vec<T>) -> Buf<T> {
219 let len = parts.iter().map(Buf::len).sum();
220 Buf { repr: Repr::Cols { parts, len, flat: Arc::new(OnceLock::new()), join } }
221 }
222
223 pub fn as_slice(&self) -> &[T] {
224 match &self.repr {
225 Repr::Owned(v) => v,
226 Repr::Slice { buf, off, len } => &buf[*off..*off + *len],
227 Repr::Foreign { ptr, len, .. } => {
228 if *len == 0 {
229 &[]
230 } else {
231 unsafe { std::slice::from_raw_parts(*ptr, *len) }
235 }
236 }
237 Repr::Cols { parts, len, flat, join } => flat.get_or_init(|| {
241 JOINS.fetch_add(1, Ordering::Relaxed);
242 Arc::new(join(parts, *len))
243 }),
244 }
245 }
246
247 pub fn to_mut(&mut self) -> &mut Vec<T> {
251 if !matches!(self.repr, Repr::Owned(_)) {
255 self.repr = Repr::Owned(Arc::new(self.as_slice().to_vec()));
256 }
257 match &mut self.repr {
258 Repr::Owned(v) => Arc::make_mut(v),
259 _ => unreachable!("just converted to a whole owned buffer"),
260 }
261 }
262
263 pub fn into_vec(self) -> Vec<T> {
266 match self.repr {
267 Repr::Owned(v) => Arc::try_unwrap(v).unwrap_or_else(|v| v.as_slice().to_vec()),
268 Repr::Slice { ref buf, off, len } => buf[off..off + len].to_vec(),
269 Repr::Foreign { .. } | Repr::Cols { .. } => self.as_slice().to_vec(),
270 }
271 }
272
273 pub fn push(&mut self, value: T) {
274 self.to_mut().push(value);
275 }
276
277 pub fn extend_from_slice(&mut self, other: &[T]) {
278 self.to_mut().extend_from_slice(other);
279 }
280
281 pub fn slice(&self, start: usize, end: usize) -> Buf<T> {
286 match &self.repr {
287 Repr::Owned(v) => {
288 assert!(start <= end && end <= v.len(), "slice out of range");
289 if start == 0 && end == v.len() {
290 return Buf { repr: Repr::Owned(Arc::clone(v)) };
291 }
292 Buf { repr: Repr::Slice { buf: Arc::clone(v), off: start, len: end - start } }
293 }
294 Repr::Slice { buf, off, len } => {
295 assert!(start <= end && end <= *len, "slice out of range");
296 let repr =
297 Repr::Slice { buf: Arc::clone(buf), off: off + start, len: end - start };
298 Buf { repr }
299 }
300 Repr::Foreign { ptr, len, owner } => {
301 assert!(start <= end && end <= *len, "slice out of range");
302 unsafe { Buf::foreign(ptr.add(start), end - start, owner.clone()) }
305 }
306 Repr::Cols { parts, len, flat, .. } => {
310 assert!(start <= end && end <= *len, "slice out of range");
311 if flat.get().is_none() {
312 let mut at = 0;
313 for part in parts {
314 let stop = at + part.len();
315 if start >= at && end <= stop {
316 return part.slice(start - at, end - at);
317 }
318 at = stop;
319 }
320 }
321 let whole = Arc::clone(self.flat_arc());
322 if start == 0 && end == whole.len() {
323 return Buf { repr: Repr::Owned(whole) };
324 }
325 Buf { repr: Repr::Slice { buf: whole, off: start, len: end - start } }
326 }
327 }
328 }
329
330 fn flat_arc(&self) -> &Arc<Vec<T>> {
332 self.as_slice();
333 match &self.repr {
334 Repr::Cols { flat, .. } => flat.get().expect("just initialised"),
335 _ => unreachable!("only a joined buffer is asked for its join"),
336 }
337 }
338}
339
340impl<T: Copy + Default + Send + Sync> Buf<T> {
341 pub fn join_fast(parts: Vec<Buf<T>>) -> Buf<T> {
344 Buf::joined(parts, join_parallel)
345 }
346}
347
348impl<T: Clone> Deref for Buf<T> {
349 type Target = [T];
350
351 fn deref(&self) -> &[T] {
352 self.as_slice()
353 }
354}
355
356impl<T: Clone> Clone for Buf<T> {
360 fn clone(&self) -> Buf<T> {
361 match &self.repr {
362 Repr::Owned(v) => Buf { repr: Repr::Owned(Arc::clone(v)) },
363 Repr::Slice { buf, off, len } => {
364 Buf { repr: Repr::Slice { buf: Arc::clone(buf), off: *off, len: *len } }
365 }
366 Repr::Foreign { ptr, len, owner } => {
367 unsafe { Buf::foreign(*ptr, *len, owner.clone()) }
369 }
370 Repr::Cols { parts, len, flat, join } => Buf {
374 repr: Repr::Cols {
375 parts: parts.clone(),
376 len: *len,
377 flat: Arc::clone(flat),
378 join: *join,
379 },
380 },
381 }
382 }
383}
384
385impl<T> Default for Buf<T> {
386 fn default() -> Buf<T> {
387 Buf::new()
388 }
389}
390
391impl<T: Clone + std::fmt::Debug> std::fmt::Debug for Buf<T> {
392 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
393 std::fmt::Debug::fmt(self.as_slice(), f)
394 }
395}
396
397impl<T: Clone + PartialEq> PartialEq for Buf<T> {
398 fn eq(&self, other: &Buf<T>) -> bool {
399 self.as_slice() == other.as_slice()
400 }
401}
402
403impl<T> From<Vec<T>> for Buf<T> {
404 fn from(v: Vec<T>) -> Buf<T> {
405 Buf::from_vec(v)
406 }
407}
408
409impl<'a, T: Clone> IntoIterator for &'a Buf<T> {
410 type Item = &'a T;
411 type IntoIter = std::slice::Iter<'a, T>;
412
413 fn into_iter(self) -> std::slice::Iter<'a, T> {
414 self.as_slice().iter()
415 }
416}
417
418impl<T> FromIterator<T> for Buf<T> {
419 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Buf<T> {
420 Buf::from_vec(Vec::from_iter(iter))
421 }
422}
423
424#[derive(Clone, Debug, PartialEq)]
425pub enum Data {
426 Bool(Buf<u8>),
427 I64(Buf<i64>),
428 Ext(Buf<Ext>),
432 Rat(Buf<Rat>),
434 F64(Buf<f64>),
435 Complex(Buf<Cx>),
438 Char(Buf<char>),
439 Symbol(Buf<crate::symbol::Id>),
443 Box(Buf<Array>),
447}
448
449impl Data {
450 pub fn dtype(&self) -> DType {
451 match self {
452 Data::Bool(_) => DType::Bool,
453 Data::I64(_) => DType::I64,
454 Data::Ext(_) => DType::Ext,
455 Data::Rat(_) => DType::Rat,
456 Data::F64(_) => DType::F64,
457 Data::Complex(_) => DType::Complex,
458 Data::Char(_) => DType::Char,
459 Data::Symbol(_) => DType::Symbol,
460 Data::Box(_) => DType::Box,
461 }
462 }
463
464 pub fn len(&self) -> usize {
465 match self {
466 Data::Bool(v) => v.len(),
467 Data::I64(v) => v.len(),
468 Data::Ext(v) => v.len(),
469 Data::Rat(v) => v.len(),
470 Data::F64(v) => v.len(),
471 Data::Complex(v) => v.len(),
472 Data::Char(v) => v.len(),
473 Data::Symbol(v) => v.len(),
474 Data::Box(v) => v.len(),
475 }
476 }
477
478 pub fn is_empty(&self) -> bool {
479 self.len() == 0
480 }
481
482 pub fn is_foreign(&self) -> bool {
484 match self {
485 Data::Bool(v) => v.is_foreign(),
486 Data::I64(v) => v.is_foreign(),
487 Data::Ext(v) => v.is_foreign(),
488 Data::Rat(v) => v.is_foreign(),
489 Data::F64(v) => v.is_foreign(),
490 Data::Complex(v) => v.is_foreign(),
491 Data::Char(v) => v.is_foreign(),
492 Data::Symbol(v) => v.is_foreign(),
493 Data::Box(v) => v.is_foreign(),
494 }
495 }
496
497 pub fn owner(&self) -> Option<&Owner> {
500 match self {
501 Data::Bool(v) => v.owner(),
502 Data::I64(v) => v.owner(),
503 Data::Ext(v) => v.owner(),
504 Data::Rat(v) => v.owner(),
505 Data::F64(v) => v.owner(),
506 Data::Complex(v) => v.owner(),
507 Data::Char(v) => v.owner(),
508 Data::Symbol(v) => v.owner(),
509 Data::Box(v) => v.owner(),
510 }
511 }
512
513 pub fn slice(&self, start: usize, end: usize) -> Data {
514 match self {
515 Data::Bool(v) => Data::Bool(v.slice(start, end)),
516 Data::I64(v) => Data::I64(v.slice(start, end)),
517 Data::Ext(v) => Data::Ext(v.slice(start, end)),
518 Data::Rat(v) => Data::Rat(v.slice(start, end)),
519 Data::F64(v) => Data::F64(v.slice(start, end)),
520 Data::Complex(v) => Data::Complex(v.slice(start, end)),
521 Data::Char(v) => Data::Char(v.slice(start, end)),
522 Data::Symbol(v) => Data::Symbol(v.slice(start, end)),
523 Data::Box(v) => Data::Box(v.slice(start, end)),
524 }
525 }
526
527 pub fn empty(dtype: DType) -> Data {
528 match dtype {
529 DType::Bool => Data::Bool(Buf::new()),
530 DType::I64 => Data::I64(Buf::new()),
531 DType::Ext => Data::Ext(Buf::new()),
532 DType::Rat => Data::Rat(Buf::new()),
533 DType::F64 => Data::F64(Buf::new()),
534 DType::Complex => Data::Complex(Buf::new()),
535 DType::Char => Data::Char(Buf::new()),
536 DType::Symbol => Data::Symbol(Buf::new()),
537 DType::Box => Data::Box(Buf::new()),
538 }
539 }
540
541 pub fn push_fill(&mut self) {
544 match self {
545 Data::Bool(v) => v.push(0),
546 Data::I64(v) => v.push(0),
547 Data::Ext(v) => v.push(Ext::default()),
548 Data::Rat(v) => v.push(Rat::zero()),
549 Data::F64(v) => v.push(0.0),
550 Data::Complex(v) => v.push(crate::complex::ZERO),
551 Data::Char(v) => v.push(' '),
552 Data::Symbol(v) => v.push(crate::symbol::EMPTY),
553 Data::Box(v) => v.push(Array::box_fill()),
554 }
555 }
556
557 pub fn push_from(&mut self, src: &Data, i: usize) {
560 match (self, src) {
561 (Data::Bool(a), Data::Bool(b)) => a.push(b[i]),
562 (Data::I64(a), Data::I64(b)) => a.push(b[i]),
563 (Data::Ext(a), Data::Ext(b)) => a.push(b[i].clone()),
564 (Data::Rat(a), Data::Rat(b)) => a.push(b[i].clone()),
565 (Data::F64(a), Data::F64(b)) => a.push(b[i]),
566 (Data::Complex(a), Data::Complex(b)) => a.push(b[i]),
567 (Data::Char(a), Data::Char(b)) => a.push(b[i]),
568 (Data::Symbol(a), Data::Symbol(b)) => a.push(b[i]),
569 (Data::Box(a), Data::Box(b)) => a.push(b[i].clone()),
570 _ => {}
571 }
572 }
573
574 pub fn extend_from(&mut self, other: &Data) -> bool {
575 match (self, other) {
576 (Data::Bool(a), Data::Bool(b)) => a.extend_from_slice(b),
577 (Data::I64(a), Data::I64(b)) => a.extend_from_slice(b),
578 (Data::Ext(a), Data::Ext(b)) => a.extend_from_slice(b),
579 (Data::Rat(a), Data::Rat(b)) => a.extend_from_slice(b),
580 (Data::F64(a), Data::F64(b)) => a.extend_from_slice(b),
581 (Data::Complex(a), Data::Complex(b)) => a.extend_from_slice(b),
582 (Data::Char(a), Data::Char(b)) => a.extend_from_slice(b),
583 (Data::Symbol(a), Data::Symbol(b)) => a.extend_from_slice(b),
584 (Data::Box(a), Data::Box(b)) => a.extend_from_slice(b),
585 _ => return false,
586 }
587 true
588 }
589
590 pub fn join(columns: &[Data], rows: usize) -> Option<Data> {
599 let first = columns.first()?;
600 if columns.iter().any(|c| c.dtype() != first.dtype() || c.len() < rows) {
601 return None;
602 }
603 macro_rules! by {
604 ($variant:ident, $join:expr) => {{
605 let mut parts = Vec::with_capacity(columns.len());
606 for c in columns {
607 let Data::$variant(v) = c else { return None };
608 parts.push(if v.len() == rows { v.clone() } else { v.slice(0, rows) });
611 }
612 Some(Data::$variant($join(parts)))
613 }};
614 }
615 match first.dtype() {
616 DType::Bool => by!(Bool, Buf::join_fast),
617 DType::I64 => by!(I64, Buf::join_fast),
618 DType::F64 => by!(F64, Buf::join_fast),
619 DType::Complex => by!(Complex, Buf::join_fast),
620 DType::Char => by!(Char, Buf::join),
621 DType::Symbol => by!(Symbol, Buf::join_fast),
622 DType::Ext => by!(Ext, Buf::join),
623 DType::Rat => by!(Rat, Buf::join),
624 DType::Box => by!(Box, Buf::join),
625 }
626 }
627
628 pub fn columns(&self, rows: usize, cols: usize) -> Vec<Data> {
633 (0..cols).map(|j| self.slice(j * rows, (j + 1) * rows)).collect()
634 }
635
636 pub fn interleave(columns: &[Data], rows: usize) -> Option<Data> {
648 let cols = columns.len();
649 let first = columns.first()?;
650 if columns.iter().any(|c| c.dtype() != first.dtype() || c.len() < rows) {
651 return None;
652 }
653
654 fn weave<T: Copy + Default + Send + Sync>(columns: &[&[T]], rows: usize) -> Vec<T> {
658 let cols = columns.len();
659 let (out, _) = crate::par::fill(rows * cols, |start, part: &mut [T]| {
660 let mut rest = &mut part[..];
661 let mut at = start;
662 let lead = ((cols - at % cols) % cols).min(rest.len());
664 if lead > 0 {
665 let (head, tail) = rest.split_at_mut(lead);
666 let r = at / cols;
667 for (k, slot) in head.iter_mut().enumerate() {
668 *slot = columns[at % cols + k][r];
669 }
670 at += lead;
671 rest = tail;
672 }
673 let whole = rest.len() / cols;
674 let (body, tail) = rest.split_at_mut(whole * cols);
675 let r0 = at / cols;
676 for (k, row) in body.chunks_exact_mut(cols).enumerate() {
677 for (slot, col) in row.iter_mut().zip(columns) {
678 *slot = col[r0 + k];
679 }
680 }
681 let r = r0 + whole;
683 for (c, slot) in tail.iter_mut().enumerate() {
684 *slot = columns[c][r];
685 }
686 true
687 });
688 out
689 }
690
691 fn weave_cloned<T: Clone>(columns: &[&[T]], rows: usize) -> Vec<T> {
695 let mut out = Vec::with_capacity(rows * columns.len());
696 for r in 0..rows {
697 for c in columns {
698 out.push(c[r].clone());
699 }
700 }
701 out
702 }
703
704 macro_rules! by {
705 ($variant:ident, $weave:ident) => {{
706 let mut s = Vec::with_capacity(cols);
707 for c in columns {
708 let Data::$variant(v) = c else { return None };
709 s.push(v.as_slice());
710 }
711 Some(Data::$variant($weave(&s, rows).into()))
712 }};
713 }
714 match first.dtype() {
715 DType::Bool => by!(Bool, weave),
716 DType::I64 => by!(I64, weave),
717 DType::F64 => by!(F64, weave),
718 DType::Complex => by!(Complex, weave),
719 DType::Char => by!(Char, weave),
720 DType::Symbol => by!(Symbol, weave),
721 DType::Ext => by!(Ext, weave_cloned),
722 DType::Rat => by!(Rat, weave_cloned),
723 DType::Box => by!(Box, weave_cloned),
724 }
725 }
726
727 pub fn cast(&self, to: DType) -> Option<Data> {
729 if self.dtype() == to {
730 return Some(self.clone());
731 }
732 match (self, to) {
733 (Data::Bool(v), DType::I64) => Some(Data::I64(v.iter().map(|&x| x as i64).collect())),
734 (Data::Bool(v), DType::F64) => Some(Data::F64(v.iter().map(|&x| x as f64).collect())),
735 (Data::I64(v), DType::F64) => Some(Data::F64(v.iter().map(|&x| x as f64).collect())),
736 (Data::Bool(v), DType::Ext) => Some(Data::Ext(v.iter().map(|&x| Ext::from(x)).collect())),
737 (Data::I64(v), DType::Ext) => Some(Data::Ext(v.iter().map(|&x| Ext::from(x)).collect())),
738 (Data::Bool(v), DType::Rat) => {
739 Some(Data::Rat(v.iter().map(|&x| Rat::from_int(Ext::from(x))).collect()))
740 }
741 (Data::I64(v), DType::Rat) => {
742 Some(Data::Rat(v.iter().map(|&x| Rat::from_int(Ext::from(x))).collect()))
743 }
744 (Data::Ext(v), DType::Rat) => {
745 Some(Data::Rat(v.iter().map(|x| Rat::from_int(x.clone())).collect()))
746 }
747 (Data::Ext(v), DType::F64) => {
748 Some(Data::F64(v.iter().map(crate::exact::ext_to_f64).collect()))
749 }
750 (Data::Rat(v), DType::F64) => Some(Data::F64(v.iter().map(Rat::to_f64).collect())),
751 (Data::Bool(v), DType::Complex) => {
752 Some(Data::Complex(v.iter().map(|&x| [x as f64, 0.0]).collect()))
753 }
754 (Data::I64(v), DType::Complex) => {
755 Some(Data::Complex(v.iter().map(|&x| [x as f64, 0.0]).collect()))
756 }
757 (Data::Ext(v), DType::Complex) => {
758 Some(Data::Complex(v.iter().map(|x| [crate::exact::ext_to_f64(x), 0.0]).collect()))
759 }
760 (Data::Rat(v), DType::Complex) => {
761 Some(Data::Complex(v.iter().map(|x| [x.to_f64(), 0.0]).collect()))
762 }
763 (Data::F64(v), DType::Complex) => {
764 Some(Data::Complex(v.iter().map(|&x| [x, 0.0]).collect()))
765 }
766 _ => None,
767 }
768 }
769}
770
771#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
778pub enum Layout {
779 #[default]
782 RowMajor,
783 ColMajor,
788}
789
790#[derive(Clone, Debug)]
803pub struct Array {
804 pub shape: Vec<usize>,
805 pub data: Data,
806 layout: Layout,
807 sparse: Option<crate::sparse::Handle>,
808 proto: Option<std::sync::Arc<Array>>,
809}
810
811impl PartialEq for Array {
814 fn eq(&self, other: &Array) -> bool {
815 if self.shape != other.shape {
816 return false;
817 }
818 if self.sparse.is_some() || other.sparse.is_some() {
819 let (a, b) = (self.densified(), other.densified());
820 return a.to_row_major().data == b.to_row_major().data;
821 }
822 if self.layout == other.layout {
823 return self.data == other.data;
824 }
825 self.to_row_major().data == other.to_row_major().data
826 }
827}
828
829impl Array {
830 pub fn new(shape: Vec<usize>, data: Data) -> Array {
831 debug_assert_eq!(shape.iter().product::<usize>(), data.len());
832 Array { shape, data, layout: Layout::RowMajor, sparse: None, proto: None }
833 }
834
835 pub fn sparse(shape: Vec<usize>, data: Data, sparse: crate::sparse::Sparse) -> Array {
840 Array { shape, data, layout: Layout::RowMajor, sparse: Some(std::sync::Arc::new(sparse)), proto: None }
841 }
842
843 pub fn is_sparse(&self) -> bool {
845 self.sparse.is_some()
846 }
847
848 pub fn proto(&self) -> Option<&Array> {
857 self.proto.as_deref()
858 }
859
860 pub fn with_proto(mut self, proto: Array) -> Array {
862 self.proto = Some(std::sync::Arc::new(proto));
863 self
864 }
865
866 pub fn sparse_parts(&self) -> Option<&crate::sparse::Sparse> {
868 self.sparse.as_deref()
869 }
870
871 pub fn densified(&self) -> Array {
874 match &self.sparse {
875 None => self.clone(),
876 Some(s) => crate::sparse::densify(self, s),
877 }
878 }
879
880 pub fn col_major(shape: Vec<usize>, data: Data) -> Array {
883 debug_assert_eq!(shape.iter().product::<usize>(), data.len());
884 let layout = if shape.len() < 2 { Layout::RowMajor } else { Layout::ColMajor };
885 Array { shape, data, layout, sparse: None, proto: None }
886 }
887
888 pub fn with_layout(mut self, layout: Layout) -> Array {
891 self.layout = if self.shape.len() < 2 { Layout::RowMajor } else { layout };
892 self
893 }
894
895 pub fn layout(&self) -> Layout {
897 self.layout
898 }
899
900 pub fn is_row_major(&self) -> bool {
901 self.layout == Layout::RowMajor
902 }
903
904 pub fn row_major_data(&self) -> &Data {
908 debug_assert!(self.is_row_major(), "a column-major buffer read as row-major");
909 &self.data
910 }
911
912 pub fn to_row_major(&self) -> Array {
915 if self.is_row_major() {
916 return self.clone();
917 }
918 LAYOUTS.fetch_add(1, Ordering::Relaxed);
919 Array::new(self.shape.clone(), self.transposed_data())
920 }
921
922 fn transposed_data(&self) -> Data {
924 let rows = self.shape[0];
925 let rest: usize = self.shape[1..].iter().product();
926 if self.rank() == 2
929 && let Some(d) = Data::interleave(&self.data.columns(rows, rest), rows)
930 {
931 return d;
932 }
933 let n = self.count();
936 let mut out = Data::empty(self.dtype());
937 let mut coord = vec![0usize; self.rank()];
938 for _ in 0..n {
939 let mut idx = 0;
940 let mut stride = 1;
941 for (k, &len) in self.shape.iter().enumerate() {
942 idx += coord[k] * stride;
943 stride *= len;
944 }
945 out.push_from(&self.data, idx);
946 let mut k = self.rank();
947 while k > 0 {
948 k -= 1;
949 coord[k] += 1;
950 if coord[k] < self.shape[k] {
951 break;
952 }
953 coord[k] = 0;
954 }
955 }
956 out
957 }
958
959 pub fn scalar_i64(v: i64) -> Array {
960 Array::new(vec![], Data::I64(vec![v].into()))
961 }
962
963 pub fn scalar_f64(v: f64) -> Array {
964 Array::new(vec![], Data::F64(vec![v].into()))
965 }
966
967 pub fn scalar_bool(v: bool) -> Array {
968 Array::new(vec![], Data::Bool(vec![v as u8].into()))
969 }
970
971 pub fn from_i64(values: Vec<i64>) -> Array {
972 Array::new(vec![values.len()], Data::I64(values.into()))
973 }
974
975 pub fn from_f64(values: Vec<f64>) -> Array {
976 Array::new(vec![values.len()], Data::F64(values.into()))
977 }
978
979 pub fn from_chars(values: Vec<char>) -> Array {
980 Array::new(vec![values.len()], Data::Char(values.into()))
981 }
982
983 pub fn empty(dtype: DType) -> Array {
984 Array::new(vec![0], Data::empty(dtype))
985 }
986
987 pub fn boxed(value: Array) -> Array {
989 Array::new(vec![], Data::Box(vec![value].into()))
990 }
991
992 pub fn box_fill() -> Array {
995 Array::empty(DType::I64)
996 }
997
998 pub fn dtype(&self) -> DType {
999 self.data.dtype()
1000 }
1001
1002 pub fn rank(&self) -> usize {
1003 self.shape.len()
1004 }
1005
1006 pub fn count(&self) -> usize {
1008 self.shape.iter().product()
1009 }
1010
1011 pub fn items(&self) -> usize {
1013 self.shape.first().copied().unwrap_or(1)
1014 }
1015
1016 pub fn item_size(&self) -> usize {
1018 self.shape.iter().skip(1).product()
1019 }
1020
1021 pub fn cast(&self, to: DType) -> Option<Array> {
1024 if self.is_sparse() {
1025 return self.densified().cast(to);
1026 }
1027 Some(Array {
1028 shape: self.shape.clone(),
1029 data: self.data.cast(to)?,
1030 layout: self.layout,
1031 sparse: None,
1032 proto: self.proto.clone(),
1033 })
1034 }
1035
1036 pub fn cells(&self, frame_rank: usize) -> Vec<Array> {
1039 debug_assert!(frame_rank <= self.rank());
1040 debug_assert!(self.is_row_major(), "cells of a column-major buffer");
1041 let cell_shape: Vec<usize> = self.shape[frame_rank..].to_vec();
1042 let cell_size: usize = cell_shape.iter().product();
1043 let n: usize = self.shape[..frame_rank].iter().product();
1044 (0..n)
1045 .map(|i| {
1046 Array::new(cell_shape.clone(), self.data.slice(i * cell_size, (i + 1) * cell_size))
1047 })
1048 .collect()
1049 }
1050
1051 pub fn cell_at(&self, frame_rank: usize, index: usize) -> Array {
1053 debug_assert!(self.is_row_major(), "a cell of a column-major buffer");
1054 let cell_shape: Vec<usize> = self.shape[frame_rank..].to_vec();
1055 let cell_size: usize = cell_shape.iter().product();
1056 Array::new(cell_shape, self.data.slice(index * cell_size, (index + 1) * cell_size))
1057 }
1058
1059 pub fn item(&self, i: usize) -> Array {
1061 debug_assert!(self.rank() >= 1);
1062 self.cell_at(1, i)
1063 }
1064
1065 pub fn as_i64_slice(&self) -> Option<&[i64]> {
1066 match &self.data {
1067 Data::I64(v) => Some(v),
1068 _ => None,
1069 }
1070 }
1071
1072 pub fn as_boxes(&self) -> Option<&[Array]> {
1074 match &self.data {
1075 Data::Box(v) => Some(v),
1076 _ => None,
1077 }
1078 }
1079
1080 pub fn as_f64_slice(&self) -> Option<&[f64]> {
1081 match &self.data {
1082 Data::F64(v) => Some(v),
1083 _ => None,
1084 }
1085 }
1086
1087 pub fn as_ext_slice(&self) -> Option<&[Ext]> {
1089 match &self.data {
1090 Data::Ext(v) => Some(v),
1091 _ => None,
1092 }
1093 }
1094
1095 pub fn as_rat_slice(&self) -> Option<&[Rat]> {
1097 match &self.data {
1098 Data::Rat(v) => Some(v),
1099 _ => None,
1100 }
1101 }
1102
1103 pub fn as_complex_slice(&self) -> Option<&[Cx]> {
1104 match &self.data {
1105 Data::Complex(v) => Some(v),
1106 _ => None,
1107 }
1108 }
1109
1110 pub fn to_complex_vec(&self) -> Option<Vec<Cx>> {
1112 match &self.data {
1113 Data::Bool(v) => Some(v.iter().map(|&x| [x as f64, 0.0]).collect()),
1114 Data::I64(v) => Some(v.iter().map(|&x| [x as f64, 0.0]).collect()),
1115 Data::Ext(v) => Some(v.iter().map(|x| [crate::exact::ext_to_f64(x), 0.0]).collect()),
1116 Data::Rat(v) => Some(v.iter().map(|x| [x.to_f64(), 0.0]).collect()),
1117 Data::F64(v) => Some(v.iter().map(|&x| [x, 0.0]).collect()),
1118 Data::Complex(v) => Some(v.to_vec()),
1119 Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
1120 }
1121 }
1122
1123 pub fn to_f64_vec(&self) -> Option<Vec<f64>> {
1125 match &self.data {
1126 Data::Bool(v) => Some(v.iter().map(|&x| x as f64).collect()),
1127 Data::I64(v) => Some(v.iter().map(|&x| x as f64).collect()),
1128 Data::Ext(v) => Some(v.iter().map(crate::exact::ext_to_f64).collect()),
1129 Data::Rat(v) => Some(v.iter().map(Rat::to_f64).collect()),
1130 Data::F64(v) => Some(v.to_vec()),
1131 Data::Complex(v) => {
1137 v.iter().map(|z| (z[1] == 0.0).then_some(z[0])).collect()
1138 }
1139 Data::Char(_) | Data::Symbol(_) if self.count() == 0 => Some(Vec::new()),
1144 Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
1145 }
1146 }
1147
1148 pub fn to_i64_vec(&self) -> Option<Vec<i64>> {
1150 match &self.data {
1151 Data::Bool(v) => Some(v.iter().map(|&x| x as i64).collect()),
1152 Data::I64(v) => Some(v.to_vec()),
1153 Data::Ext(v) => v.iter().map(crate::exact::ext_to_i64).collect(),
1156 Data::Rat(v) => {
1157 v.iter().map(|x| x.to_int().as_ref().and_then(crate::exact::ext_to_i64)).collect()
1158 }
1159 Data::F64(v) => {
1160 let mut out = Vec::with_capacity(v.len());
1161 for &x in v.iter() {
1162 if x.fract() != 0.0 || x.abs() >= i64::MAX as f64 {
1163 return None;
1164 }
1165 out.push(x as i64);
1166 }
1167 Some(out)
1168 }
1169 Data::Complex(v) => v
1173 .iter()
1174 .map(|z| {
1175 (z[1] == 0.0 && z[0].fract() == 0.0 && z[0].abs() < i64::MAX as f64)
1176 .then_some(z[0] as i64)
1177 })
1178 .collect(),
1179 Data::Char(_) | Data::Symbol(_) if self.count() == 0 => Some(Vec::new()),
1180 Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
1181 }
1182 }
1183
1184 pub fn to_i64_vec_near(&self, near: NearInt) -> Option<Vec<i64>> {
1197 let Data::F64(v) = &self.data else {
1198 return self.to_i64_vec();
1201 };
1202 v.iter().map(|&x| near.round(x)).collect()
1203 }
1204}
1205
1206#[derive(Clone, Copy, Debug, PartialEq)]
1214pub enum NearInt {
1215 J,
1218 Apl,
1220 Tolerant(crate::verb::Tol),
1223}
1224
1225impl NearInt {
1226 pub const J_RELATIVE: f64 = 1.0 / 17_592_186_044_416.0;
1230 pub const APL_ABSOLUTE: f64 = 1e-10;
1232
1233 pub fn of(lang: crate::Lang) -> NearInt {
1235 match lang {
1236 crate::Lang::J => NearInt::J,
1237 crate::Lang::Apl => NearInt::Apl,
1238 }
1239 }
1240
1241 pub fn round(self, x: f64) -> Option<i64> {
1243 if x.fract() == 0.0 {
1244 return (x.abs() < i64::MAX as f64).then_some(x as i64);
1245 }
1246 let n = x.round();
1247 let within = match self {
1248 NearInt::J => (x - n).abs() <= Self::J_RELATIVE * x.abs().max(n.abs()),
1249 NearInt::Apl => (x - n).abs() < Self::APL_ABSOLUTE,
1250 NearInt::Tolerant(tol) => tol.eq(x, n),
1251 };
1252 (within && n.abs() < i64::MAX as f64).then_some(n as i64)
1253 }
1254}
1255
1256#[cfg(test)]
1257mod tests {
1258 use super::*;
1259 use std::sync::atomic::{AtomicBool, Ordering};
1260
1261 struct Guard {
1264 values: Vec<i64>,
1265 dropped: Arc<AtomicBool>,
1266 }
1267
1268 impl Drop for Guard {
1269 fn drop(&mut self) {
1270 self.dropped.store(true, Ordering::SeqCst);
1271 }
1272 }
1273
1274 fn foreign_buf(values: Vec<i64>, dropped: Arc<AtomicBool>) -> Buf<i64> {
1275 let guard = Arc::new(Guard { values, dropped });
1276 let ptr = guard.values.as_ptr();
1277 let len = guard.values.len();
1278 unsafe { Buf::foreign(ptr, len, guard) }
1281 }
1282
1283 #[test]
1284 fn owned_buf_derefs_to_its_slice() {
1285 let b: Buf<i64> = vec![1, 2, 3].into();
1286 assert!(!b.is_foreign());
1287 assert_eq!(&b[..], &[1, 2, 3]);
1288 assert_eq!(b.len(), 3);
1289 assert_eq!(b.iter().sum::<i64>(), 6);
1290 }
1291
1292 #[test]
1293 fn empty_buf_is_a_valid_empty_slice() {
1294 let b: Buf<f64> = Buf::new();
1295 assert_eq!(&b[..], &[] as &[f64]);
1296 let f = unsafe { Buf::<f64>::foreign(std::ptr::null(), 0, Arc::new(())) };
1298 assert_eq!(&f[..], &[] as &[f64]);
1299 }
1300
1301 #[test]
1302 fn cloning_an_owned_buf_shares_the_same_memory() {
1303 let b: Buf<i64> = vec![1, 2, 3].into();
1304 let c = b.clone();
1305 assert_eq!(b.as_ptr(), c.as_ptr(), "owned clone copied the elements");
1306 assert_eq!(&c[..], &[1, 2, 3]);
1307 }
1308
1309 #[test]
1310 fn writing_to_a_shared_owned_buf_copies_first() {
1311 let b: Buf<i64> = vec![1, 2, 3].into();
1312 let mut c = b.clone();
1313 c.to_mut()[0] = 99;
1314 assert_eq!(&b[..], &[1, 2, 3], "the other holder saw the write");
1315 assert_eq!(&c[..], &[99, 2, 3]);
1316 assert_ne!(b.as_ptr(), c.as_ptr());
1317 let ptr = c.as_ptr();
1319 c.to_mut()[1] = 98;
1320 assert_eq!(c.as_ptr(), ptr, "unshared write copied");
1321 }
1322
1323 #[test]
1324 fn into_vec_moves_when_sole_holder_and_copies_when_shared() {
1325 let b: Buf<i64> = vec![1, 2, 3].into();
1326 let ptr = b.as_ptr();
1327 let v = b.into_vec();
1328 assert_eq!(v.as_ptr(), ptr, "sole holder copied instead of moving");
1329
1330 let b: Buf<i64> = vec![1, 2, 3].into();
1331 let c = b.clone();
1332 let v = b.into_vec();
1333 assert_eq!(v, vec![1, 2, 3]);
1334 assert_eq!(&c[..], &[1, 2, 3]);
1335 }
1336
1337 #[test]
1338 fn foreign_buf_reads_borrowed_memory_and_keeps_the_owner_alive() {
1339 let dropped = Arc::new(AtomicBool::new(false));
1340 let b = foreign_buf(vec![10, 20, 30], dropped.clone());
1341 assert!(b.is_foreign());
1342 assert_eq!(&b[..], &[10, 20, 30]);
1343 assert!(!dropped.load(Ordering::SeqCst), "owner dropped while borrowed");
1344 drop(b);
1345 assert!(dropped.load(Ordering::SeqCst), "owner leaked after the buffer died");
1346 }
1347
1348 #[test]
1349 fn cloning_a_foreign_buf_shares_the_same_memory() {
1350 let dropped = Arc::new(AtomicBool::new(false));
1351 let b = foreign_buf(vec![1, 2, 3], dropped.clone());
1352 let c = b.clone();
1353 assert!(c.is_foreign());
1354 assert_eq!(b.as_ptr(), c.as_ptr());
1355 drop(b);
1356 assert!(!dropped.load(Ordering::SeqCst), "owner dropped while a clone lives");
1357 assert_eq!(&c[..], &[1, 2, 3]);
1358 }
1359
1360 #[test]
1361 fn slicing_a_foreign_buf_keeps_borrowing() {
1362 let dropped = Arc::new(AtomicBool::new(false));
1363 let b = foreign_buf(vec![1, 2, 3, 4], dropped.clone());
1364 let s = b.slice(1, 3);
1365 assert!(s.is_foreign());
1366 assert_eq!(&s[..], &[2, 3]);
1367 drop(b);
1368 assert_eq!(&s[..], &[2, 3]);
1369 assert!(!dropped.load(Ordering::SeqCst));
1370 }
1371
1372 #[test]
1373 fn mutating_a_foreign_buf_copies_first() {
1374 let dropped = Arc::new(AtomicBool::new(false));
1375 let mut b = foreign_buf(vec![1, 2, 3], dropped.clone());
1376 b.push(4);
1377 assert!(!b.is_foreign());
1378 assert_eq!(&b[..], &[1, 2, 3, 4]);
1379 drop(b);
1381 assert!(dropped.load(Ordering::SeqCst));
1382 }
1383
1384 #[test]
1385 fn copy_on_write_leaves_other_holders_alone() {
1386 let dropped = Arc::new(AtomicBool::new(false));
1387 let b = foreign_buf(vec![1, 2, 3], dropped.clone());
1388 let mut c = b.clone();
1389 c.to_mut()[0] = 99;
1390 assert_eq!(&b[..], &[1, 2, 3]);
1391 assert_eq!(&c[..], &[99, 2, 3]);
1392 }
1393
1394 #[test]
1395 fn foreign_data_slices_without_copying() {
1396 let dropped = Arc::new(AtomicBool::new(false));
1397 let a = Array::new(vec![2, 2], Data::I64(foreign_buf(vec![1, 2, 3, 4], dropped)));
1398 assert!(a.data.is_foreign());
1399 let row = a.item(1);
1400 assert!(row.data.is_foreign());
1401 assert_eq!(row.as_i64_slice(), Some(&[3, 4][..]));
1402 }
1403
1404 #[test]
1405 fn foreign_data_extends_by_copying() {
1406 let dropped = Arc::new(AtomicBool::new(false));
1407 let mut d = Data::I64(foreign_buf(vec![1, 2], dropped));
1408 assert!(d.is_foreign());
1409 assert!(d.extend_from(&Data::I64(vec![3].into())));
1410 assert!(!d.is_foreign());
1411 assert_eq!(d, Data::I64(vec![1, 2, 3].into()));
1412 }
1413}