Skip to main content

jay/
array.rs

1//! Dense multidimensional array: a shape, a flat buffer and the [`Layout`]
2//! that says how one indexes the other.
3//!
4//! Buffers are either owned or borrowed from foreign memory (Arrow, the
5//! Python buffer protocol) through [`Buf`], which is what makes the data
6//! boundary zero-copy.
7
8use 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
17/// Anything that keeps a foreign buffer's memory alive: the importing side
18/// stores its release guards here and the buffer outlives nothing else.
19pub type Owner = Arc<dyn Any + Send + Sync>;
20
21/// How many joined buffers have been joined — that is, how many times a set
22/// of columns that crossed the boundary without a copy has since had to be
23/// copied into one block.
24///
25/// It exists so that a test can assert a copy did not happen: the number is
26/// process-wide and only ever grows.
27static JOINS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
28
29/// How many times a column-major array has had its rows materialised —
30/// [`Array::to_row_major`] doing real work. Process-wide, only ever grows,
31/// and here so that a test can say which verbs need the rows and which do
32/// not.
33static LAYOUTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
34
35/// The count of joins made since the process started. See `JOINS`.
36pub fn joins_made() -> u64 {
37    JOINS.load(Ordering::Relaxed)
38}
39
40/// The count of row-major materialisations since the process started. See
41/// `LAYOUTS`.
42pub fn layouts_made() -> u64 {
43    LAYOUTS.load(Ordering::Relaxed)
44}
45
46/// A flat element buffer, owned or borrowed.
47///
48/// A borrowed (foreign) buffer points into memory owned by someone else — an
49/// Arrow C data interface import, a Python buffer — and holds an `Owner`
50/// handle that keeps that memory alive for at least as long as the buffer.
51/// An owned buffer is refcounted, and [`Buf::slice`] of one is a window over
52/// the same allocation rather than a copy. However the buffer was made,
53/// cloning is a refcount bump and mutation copies first if the memory is
54/// shared, foreign or a window ([`Buf::to_mut`]), so a `Buf` behaves as a
55/// private value however cheaply it was cloned.
56pub struct Buf<T> {
57    repr: Repr<T>,
58}
59
60enum Repr<T> {
61    /// The whole of a refcounted `Vec`.
62    Owned(Arc<Vec<T>>),
63    /// The window `[off, off + len)` of a refcounted `Vec`, which is what
64    /// taking a cell or a section out of an owned array gives: a view over
65    /// the same allocation, never a copy. Writing to one copies first, as
66    /// writing to a shared whole does.
67    Slice { buf: Arc<Vec<T>>, off: usize, len: usize },
68    Foreign { ptr: *const T, len: usize, owner: Owner },
69    /// Several buffers end to end, joined only if someone asks for the flat
70    /// slice. This is how a table of columns arrives: each column keeps
71    /// borrowing its own memory, and a reader that wants the columns takes
72    /// them ([`Buf::parts`]) rather than the join. The join, once made, is
73    /// kept and shared with every clone, so no buffer is ever built twice.
74    ///
75    /// `join` is how to make it. A plain element type takes the parallel
76    /// copy, which is what keeps the join from costing more than the weave
77    /// it replaced; a heap-backed one takes the sequential clone.
78    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
86// SAFETY: no variant hands out aliased mutable access; a join of parts is
87// made once behind a `OnceLock` and never written again. A foreign buffer
88// is read-only for its whole life and its `owner` keeps the memory alive; an
89// owned buffer shares its `Vec` through an `Arc` and only ever mutates it
90// through `Arc::make_mut`, which copies unless this buffer is the sole
91// holder; a window over part of one becomes a `Vec` of its own before any
92// write, so it never mutates the allocation it shares. So `Buf` is exactly
93// as shareable as the `&[T]` it derefs to —
94// which, because an owned buffer is an `Arc<Vec<T>>` that may be dropped or
95// read from any thread holding a clone, needs `T: Send + Sync` on both.
96unsafe impl<T: Send + Sync> Send for Buf<T> {}
97// SAFETY: as above; `&Buf<T>` only ever hands out `&[T]`.
98unsafe 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    /// Borrow `len` elements at `ptr`, keeping `owner` alive alongside them.
110    ///
111    /// # Safety
112    ///
113    /// `ptr` must be aligned for `T` and point to `len` initialised elements
114    /// that stay valid, and are not mutated by anyone, for as long as `owner`
115    /// is alive. `len == 0` accepts a dangling `ptr`.
116    pub unsafe fn foreign(ptr: *const T, len: usize, owner: Owner) -> Buf<T> {
117        Buf { repr: Repr::Foreign { ptr, len, owner } }
118    }
119
120    /// True while the buffer still borrows foreign memory. A join of
121    /// buffers borrows while any of its parts does and the join has not
122    /// been made.
123    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    /// Elements the buffer holds, without joining a set of parts.
134    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    /// True once a joined buffer has had its join made: the copy the
146    /// boundary avoided has since been paid for.
147    pub fn is_joined(&self) -> bool {
148        matches!(&self.repr, Repr::Cols { flat, .. } if flat.get().is_some())
149    }
150
151    /// The parts of a buffer that was made by joining several, in order —
152    /// None for every other buffer. A reader that can work part by part
153    /// (a column at a time) takes this and never makes the join.
154    pub fn parts(&self) -> Option<&[Buf<T>]> {
155        match &self.repr {
156            Repr::Cols { parts, .. } => Some(parts),
157            _ => None,
158        }
159    }
160
161    /// The handle keeping this buffer's memory alive, for a borrowed
162    /// buffer.
163    ///
164    /// What the handle holds is the importing side's business, and a reader
165    /// that recognises one of its own can act on it: a device upload leaves
166    /// the device allocation in here, which is how an array carries its
167    /// location without becoming a different kind of array.
168    pub fn owner(&self) -> Option<&Owner> {
169        match &self.repr {
170            Repr::Foreign { owner, .. } => Some(owner),
171            _ => None,
172        }
173    }
174}
175
176/// Join the parts one element at a time. Any element type at all, and the
177/// only choice for the heap-backed ones.
178fn 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
186/// Join the parts on the thread pool: each chunk of the result copies from
187/// whichever parts cover it. A fresh block of this size costs more to fault
188/// in than to fill, and that cost only comes down by spreading the writes.
189fn 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    /// One buffer holding `parts` end to end. Nothing is copied here: the
212    /// parts are joined when — and only when — a caller asks for the flat
213    /// slice, and then one element at a time.
214    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                    // SAFETY: the `foreign` contract guarantees `len`
232                    // initialised, aligned, immutable elements at `ptr`, kept
233                    // alive by the owner this buffer holds.
234                    unsafe { std::slice::from_raw_parts(*ptr, *len) }
235                }
236            }
237            // The join a set of parts was put off making. It is made once
238            // and kept, so a buffer asked for its flat form twice pays for
239            // it once.
240            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    /// The buffer as a uniquely owned `Vec`, copying once if it is foreign or
248    /// shared with another holder. Subsequent calls on the same buffer are
249    /// free until it is cloned again.
250    pub fn to_mut(&mut self) -> &mut Vec<T> {
251        // A window over part of a `Vec` becomes a `Vec` of its own first:
252        // what the caller writes — a change of length included — must not
253        // reach the other windows over the same allocation.
254        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    /// The contents as a `Vec`, moving it out when this buffer is the sole
264    /// holder of a whole one and copying otherwise.
265    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    /// Elements `[start, end)`, as a view: no element is copied, whatever
282    /// the buffer is. A foreign slice keeps borrowing and shares the same
283    /// owner; an owned one is a window over the same refcounted `Vec`, so
284    /// it holds that whole allocation alive for as long as it lives.
285    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                // SAFETY: `start <= len` keeps the offset inside the same
303                // allocation; the new buffer holds a clone of the owner.
304                unsafe { Buf::foreign(ptr.add(start), end - start, owner.clone()) }
305            }
306            // A range inside one part is that part's own slice, so taking a
307            // column out of a joined table copies nothing. Anything else
308            // crosses a seam and has to read the join.
309            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    /// The join, made if it was not made yet. Only a joined buffer has one.
331    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    /// [`Buf::join`] for a plain element type: the join, if it is ever
342    /// made, is made on the thread pool.
343    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
356/// Cloning never copies elements: every shape of buffer is a refcount bump,
357/// and the copy happens later, in [`Buf::to_mut`], only if someone writes
358/// while the memory is still shared.
359impl<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                // SAFETY: same pointer, same owner, same guarantees.
368                unsafe { Buf::foreign(*ptr, *len, owner.clone()) }
369            }
370            // The parts are refcount bumps, and the join is shared with
371            // every other holder: made at most once however many clones ask
372            // for it.
373            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    /// Arbitrary-precision integers. Like boxes, these are heap-backed
429    /// pointers rather than machine words: never foreign, never fused,
430    /// never vectorised.
431    Ext(Buf<Ext>),
432    /// Exact ratios, each in lowest terms. Heap-backed, as `Ext` is.
433    Rat(Buf<Rat>),
434    F64(Buf<f64>),
435    /// Complex numbers, interleaved `[re, im]` — the layout numpy, C and a
436    /// pair of Arrow float columns all share.
437    Complex(Buf<Cx>),
438    Char(Buf<char>),
439    /// Symbols: every element is an index into the process-wide symbol
440    /// table (see [`crate::symbol`]), so the buffer is as flat and as
441    /// cheap to copy as one of integers and the names live once each.
442    Symbol(Buf<crate::symbol::Id>),
443    /// Boxes: every element is a whole array. Foreign memory never holds
444    /// these, so a boxed buffer is always owned and cloning it is a
445    /// refcount bump like any other.
446    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    /// True while the payload still borrows foreign memory.
483    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    /// The handle keeping this payload's memory alive, for a borrowed one.
498    /// See [`Buf::owner`].
499    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    /// The fill element used by overtaking and framing. The boxed fill is
542    /// J's `a:`, a box holding an empty numeric list.
543    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    /// Append element `i` of `src`, which must hold the same type. Nothing
558    /// happens if the types disagree; every caller has checked already.
559    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    /// The columns end to end, as the flat buffer of a [`Layout::ColMajor`]
591    /// array of shape `[rows, columns.len()]`.
592    ///
593    /// Nothing is copied: each column keeps borrowing whatever memory it
594    /// arrived in, and the buffer joins them only if some reader asks for
595    /// the flat slice. This is the table boundary that does no work.
596    ///
597    /// None on the same disagreements [`Data::interleave`] refuses.
598    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                    // A column longer than the table contributes its first
609                    // `rows` elements, as the weave takes them.
610                    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    /// The `cols` runs of `rows` elements this buffer holds, each as a
629    /// buffer of its own: the columns of a [`Layout::ColMajor`] array.
630    /// Slicing a joined buffer at a seam copies nothing, so a table that
631    /// arrived as columns is read back as the columns it arrived as.
632    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    /// Weave column-major buffers into one row-major block of shape
637    /// `[rows, columns.len()]`.
638    ///
639    /// This is the table boundary: a DataFrame arrives as one buffer per
640    /// column and libjay works rows-leading, so the elements have to be
641    /// woven once. The weave reads every column in order and writes its
642    /// result straight through, split across threads at the sizes that pay.
643    ///
644    /// None when the columns disagree on element type, when one is shorter
645    /// than `rows`, or when there are no columns at all — the importing
646    /// side has already reported that.
647    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        /// One row of the output takes one element from each column, so a
655        /// chunk of the output is a run of whole rows plus, at either end,
656        /// the part of a row the neighbouring chunk does not hold.
657        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                // The tail of a row that began in the chunk before this one.
663                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                // The head of a row the next chunk finishes.
682                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        /// The same weave for the heap-backed types, which are neither
692        /// `Copy` nor worth a thread: Arrow carries none of them, so this
693        /// only ever runs on data libjay built itself.
694        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    /// Widen to `to`. Returns None for unsupported conversions.
728    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/// How an array's shape indexes its flat buffer.
772///
773/// The shape is always the LOGICAL one — rows leading, the contract every
774/// frontend and every diagnostic reads by — and the layout says only where
775/// element `(i0 … ik)` sits in the buffer. Rank 0 and rank 1 have one
776/// possible answer and are always [`Layout::RowMajor`].
777#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
778pub enum Layout {
779    /// The last axis varies fastest: offset `((i0 * s1) + i1) * s2 + …`.
780    /// Everything in the runtime reads this unless it has asked.
781    #[default]
782    RowMajor,
783    /// The FIRST axis varies fastest, which for a matrix means each column
784    /// is contiguous — the layout a table of Arrow columns already has, and
785    /// the layout `|:` produces by flipping this flag instead of moving
786    /// 160 MB.
787    ColMajor,
788}
789
790/// An array: a logical shape, a flat buffer, and the layout joining them.
791///
792/// `data` is the buffer as it lies. A reader that indexes it must either
793/// honour [`Array::layout`] or take [`Array::to_row_major`] first; the
794/// runtime's rule is that a value reaching a verb has already been made
795/// row-major unless that verb asked for the other one.
796///
797/// A SPARSE array is the one exception to "the buffer holds every element":
798/// `shape` is still the logical shape, but `data` holds only the stored
799/// cells and [`crate::sparse::Sparse`] says where they sit. Only `$.`, the
800/// display and `":` read that form; every other reader takes
801/// [`Array::densified`] first.
802#[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
811/// Two arrays are equal when they hold the same elements at the same
812/// indices, whatever buffer order — or storage kind — each of them keeps.
813impl 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    /// A sparse array: the logical `shape`, the stored cells, and the
836    /// description of where they sit. `data` holds `entries` cells and not
837    /// one element per position, so this is the only constructor that does
838    /// not tie the buffer's length to the shape.
839    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    /// True while the array holds only its stored cells.
844    pub fn is_sparse(&self) -> bool {
845        self.sparse.is_some()
846    }
847
848    /// The item an array with no items would have held — APL's prototype.
849    ///
850    /// A simple array's type says what its fills look like, so nothing has
851    /// to be remembered; a nested one does, since an empty buffer of boxes
852    /// no longer says whether its items were pairs of numbers or of
853    /// characters. `0⍴⊂2 3⍴9` is such an array, and `↑` of it answers the
854    /// 2 by 3 table of zeros this holds. Only the operations that make an
855    /// empty out of a nested array set it, and only APL reads it.
856    pub fn proto(&self) -> Option<&Array> {
857        self.proto.as_deref()
858    }
859
860    /// The same array, remembering what its items looked like.
861    pub fn with_proto(mut self, proto: Array) -> Array {
862        self.proto = Some(std::sync::Arc::new(proto));
863        self
864    }
865
866    /// How this array is stored sparsely, or None for a dense one.
867    pub fn sparse_parts(&self) -> Option<&crate::sparse::Sparse> {
868        self.sparse.as_deref()
869    }
870
871    /// This array with every position materialised. A dense array is a
872    /// refcount bump; a sparse one is expanded here and nowhere else.
873    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    /// An array whose buffer holds its first axis fastest — the columns of
881    /// a matrix, end to end. Rank 0 and 1 have only one layout and take it.
882    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    /// The same buffer read the other way round. The caller is asserting
889    /// that the buffer really is in `layout` order for this shape.
890    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    /// How this array's buffer is ordered.
896    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    /// The flat buffer, for a reader that indexes it row-major. Debug builds
905    /// refuse a buffer that is not in that order, which is what keeps a
906    /// column-major table from being read as if it were rows.
907    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    /// This array with its elements in row-major order, materialising them
913    /// once if they are not. Already row-major: a refcount bump.
914    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    /// The buffer's elements in row-major order for this array's shape.
923    fn transposed_data(&self) -> Data {
924        let rows = self.shape[0];
925        let rest: usize = self.shape[1..].iter().product();
926        // A matrix is the weave the table boundary used to do eagerly: the
927        // columns are already contiguous, and it runs on the pool.
928        if self.rank() == 2
929            && let Some(d) = Data::interleave(&self.data.columns(rows, rest), rows)
930        {
931            return d;
932        }
933        // Higher rank: the first axis varies fastest, so reading the source
934        // at the transposed offset gives the row-major order.
935        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    /// `y` as a scalar box (J `<`).
988    pub fn boxed(value: Array) -> Array {
989        Array::new(vec![], Data::Box(vec![value].into()))
990    }
991
992    /// The element that fills a boxed array: J's `a:`, a box holding an
993    /// empty numeric list.
994    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    /// Total number of elements.
1007    pub fn count(&self) -> usize {
1008        self.shape.iter().product()
1009    }
1010
1011    /// Number of items (major cells): leading axis length, 1 for a scalar.
1012    pub fn items(&self) -> usize {
1013        self.shape.first().copied().unwrap_or(1)
1014    }
1015
1016    /// Elements per item.
1017    pub fn item_size(&self) -> usize {
1018        self.shape.iter().skip(1).product()
1019    }
1020
1021    /// Widen the elements. A cast reads and writes the buffer as it lies,
1022    /// so the layout comes through untouched.
1023    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    /// Split into cells: the trailing `cell_rank` axes form the cell shape,
1037    /// the leading axes form the frame.
1038    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    /// One cell without materialising all of them.
1052    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    /// Item `i` (major cell along the leading axis).
1060    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    /// The boxed elements, if the array holds boxes.
1073    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    /// The extended integers, if the array holds them.
1088    pub fn as_ext_slice(&self) -> Option<&[Ext]> {
1089        match &self.data {
1090            Data::Ext(v) => Some(v),
1091            _ => None,
1092        }
1093    }
1094
1095    /// The rationals, if the array holds them.
1096    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    /// Numeric contents widened to complex. None for character or boxed data.
1111    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    /// Numeric contents widened to f64. None for character data.
1124    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            // A complex value whose imaginary part is zero IS a real one
1132            // wherever a real is wanted: J answers `1 <. j. 0` with 0 and
1133            // `i. 3j0` with `0 1 2`, while `3!:0 j. 0` still reports the
1134            // complex type, so the demotion is at the use and not at the
1135            // making.
1136            Data::Complex(v) => {
1137                v.iter().map(|z| (z[1] == 0.0).then_some(z[0])).collect()
1138            }
1139            // An EMPTY array carries no value of the wrong type, so it is
1140            // acceptable numeric data whatever type it was written at:
1141            // `#. ''` is 0 in J and `¯3⊥''` is 0 in GNU APL. An empty BOX
1142            // is not: J refuses `2 #. 0$<1` where it answers `2 #. ''`.
1143            Data::Char(_) | Data::Symbol(_) if self.count() == 0 => Some(Vec::new()),
1144            Data::Char(_) | Data::Symbol(_) | Data::Box(_) => None,
1145        }
1146    }
1147
1148    /// Numeric contents as i64 if exactly representable.
1149    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            // An exact value converts only when it really is a machine
1154            // integer; anything else is a refusal, not a rounding.
1155            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            // The same two readings [`Array::to_f64_vec`] gives: a complex
1170            // with no imaginary part is the real it displays as, and an
1171            // empty of a non-numeric type holds no value to refuse.
1172            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    /// Numeric contents as i64 where a COUNT, a LENGTH or an INDEX is
1185    /// wanted, admitting a float that is merely near a whole number.
1186    ///
1187    /// Both references round such a float to the whole number beside it
1188    /// rather than refusing it — `⍳2-1E¯14` is `1 2` and `(2-1e_14) {. 1 2 3`
1189    /// is `1 2` — and neither admission is the comparison tolerance:
1190    /// `⎕CT←0` and `9!:19 (0)` leave both exactly where they are. The two
1191    /// admissions differ in shape. [`NearInt::J`] is relative, a value
1192    /// within `2^-44` of a whole number's magnitude; [`NearInt::Apl`] is
1193    /// absolute, `1e-10`, whatever the magnitude; [`NearInt::Tolerant`] is
1194    /// relative and follows the comparison tolerance in force. Everything a
1195    /// full integer apart is still a refusal in all three.
1196    pub fn to_i64_vec_near(&self, near: NearInt) -> Option<Vec<i64>> {
1197        let Data::F64(v) = &self.data else {
1198            // Every other type is exact or is refused outright; only a
1199            // float can be near a whole number without being one.
1200            return self.to_i64_vec();
1201        };
1202        v.iter().map(|&x| near.round(x)).collect()
1203    }
1204}
1205
1206/// The near-integer admission a count, a length or an index position uses.
1207///
1208/// In J and in GNU APL it is a language constant: neither lets a program
1209/// move it, and neither is the comparison tolerance the same program can
1210/// set. Dyalog's is the exception — relative and scaled by `⎕CT` — so the
1211/// rule is a dialect setting there, and [`NearInt::Tolerant`] carries the
1212/// tolerance in force with it.
1213#[derive(Clone, Copy, Debug, PartialEq)]
1214pub enum NearInt {
1215    /// J: `|x - n| ≤ 2^-44 × max(|x|, |n|)`, so the window grows with the
1216    /// magnitude and closes completely at zero.
1217    J,
1218    /// APL: `|x - n| < 1e-10` at every magnitude.
1219    Apl,
1220    /// Dyalog: the dialect's own tolerant equality against the whole
1221    /// number, so the window grows with the magnitude and `⎕CT` moves it.
1222    Tolerant(crate::verb::Tol),
1223}
1224
1225impl NearInt {
1226    /// J's relative admission, which is also the value J's comparison
1227    /// tolerance starts at — the two are separate settings that happen to
1228    /// share a number.
1229    pub const J_RELATIVE: f64 = 1.0 / 17_592_186_044_416.0;
1230    /// APL's absolute admission.
1231    pub const APL_ABSOLUTE: f64 = 1e-10;
1232
1233    /// The rule for a language, at that language's shipped dialect.
1234    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    /// The whole number `x` stands for, or None when it stands for none.
1242    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    /// Owns a vector and records its own drop, so a test can assert that a
1262    /// foreign buffer kept it alive.
1263    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        // SAFETY: the guard owns the vector, is moved into the buffer's owner
1279        // slot, and nothing mutates it afterwards.
1280        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        // SAFETY: zero length, so the dangling pointer is never dereferenced.
1297        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        // Sole holder again: further writes are in place.
1318        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        // The original memory is untouched and released with the owner.
1380        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}