furiosa_mapping/
ext.rs

1//! Impl traits exposing Mapping/FMapping/Index/Division methods via the `furiosa-mapping-impl` shared library.
2//!
3//! This trait is to invoke the methods of the types as if they are defined by `impl for T`, i.e. the method invocation
4//! notation `x.foo()`, not a function call `foo(x)`.  Note that the implementation of `impl for T` should be in the
5//! same crate with `T` due to the orphan rule, but we would like to provide/hide the method definitions as a separate
6//! library.
7
8use abi_stable::std_types::{ROption, RResult, RSlice, RVec, Tuple2};
9use furiosa_mapping_types::{
10    Atom, Division, DivisionError, DivisionMode, DivisionSide, FMapping, Factor, Ident, Mapping, PaddingKind,
11    RSortedMap, Relaxed, Span, Strict, Term, TermBounds, TermPosition,
12};
13
14use crate::{Index, IndexValueError};
15
16#[expect(improper_ctypes, reason = "all types are #[repr(C)] + StableAbi")]
17unsafe extern "C-unwind" {
18    fn factor_padding(size: usize, kind: PaddingKind) -> Factor;
19    fn factor_idents(slf: &Factor) -> RVec<Ident>;
20    fn atom_size(slf: &Atom) -> usize;
21    fn atom_idents(slf: &Atom) -> RVec<Ident>;
22    fn atom_contains_ident(slf: &Atom, ident: Ident) -> bool;
23    fn atom_find_symbol_size(slf: &Atom, ident: Ident) -> ROption<usize>;
24    fn term_depth(slf: &Term) -> usize;
25    fn term_idents(slf: &Term) -> RVec<Ident>;
26    fn mapping_size(slf: &Mapping) -> usize;
27    fn mapping_pair(slf: Mapping, other: Mapping) -> Mapping;
28    fn mapping_pairs(ms: RVec<Mapping>) -> Mapping;
29    fn mapping_divide_span(slf: &Mapping, divisor: &Mapping, span: usize) -> RResult<Division<Span>, DivisionError>;
30    fn mapping_divide_relaxed(slf: &Mapping, divisor: &Mapping) -> Division<Relaxed>;
31    fn mapping_divide_strict(slf: &Mapping, divisor: &Mapping) -> Division<Strict>;
32    fn mapping_factorize(slf: &Mapping) -> FMapping;
33    fn fmapping_from_axes(axes: RSlice<Term>) -> FMapping;
34    fn fmapping_pop(slf: &mut FMapping) -> ROption<Factor>;
35    fn fmapping_into_inner(slf: FMapping) -> RVec<Factor>;
36    fn fmapping_is_padding(slf: &FMapping) -> bool;
37    fn fmapping_has_terms(slf: &FMapping) -> bool;
38    fn fmapping_terms_with_stride(slf: &FMapping) -> RVec<TermPosition>;
39    fn fmapping_factors(slf: &FMapping) -> RSlice<'_, Factor>;
40    fn fmapping_contains_ident(slf: &FMapping, ident: Ident) -> bool;
41    fn fmapping_idents(slf: &FMapping) -> RVec<Ident>;
42    fn fmapping_find_symbol_size(slf: &FMapping, ident: Ident) -> ROption<usize>;
43    fn fmapping_eval(slf: &FMapping, position: usize) -> Index;
44    fn fmapping_is_ident_isolated(slf: &FMapping, ident: Ident) -> bool;
45    fn fmapping_into_factor(slf: FMapping) -> Factor;
46    fn fmapping_mul(slf: FMapping, inner: FMapping) -> FMapping;
47    fn fmapping_stride(slf: FMapping, stride: usize) -> FMapping;
48    fn fmapping_modulo(slf: FMapping, modulo: usize) -> FMapping;
49    fn fmapping_is_resize_of(slf: &FMapping, original: &FMapping) -> bool;
50    fn fmapping_padding(slf: FMapping, padding: usize, kind: PaddingKind) -> FMapping;
51    fn fmapping_to_mapping(slf: &FMapping) -> Mapping;
52    fn fmapping_divide_span(slf: FMapping, divisor: FMapping, span: usize) -> RResult<Division<Span>, DivisionError>;
53    fn fmapping_divide_relaxed(slf: FMapping, divisor: FMapping) -> Division<Relaxed>;
54    fn fmapping_divide_strict(slf: FMapping, divisor: FMapping) -> Division<Strict>;
55    fn fmapping_normalize(slf: FMapping) -> FMapping;
56    fn fmapping_remove_padding(slf: FMapping) -> FMapping;
57    fn fmapping_split_at(slf: &FMapping, target: usize) -> Tuple2<FMapping, FMapping>;
58    fn fmapping_pad(slf: FMapping, target: usize) -> FMapping;
59    fn division_exact_relaxed(slf: Division<Relaxed>) -> RResult<Division<Relaxed>, DivisionError>;
60    fn division_exact_strict(slf: Division<Strict>) -> RResult<Division<Strict>, DivisionError>;
61    fn division_exact_span(slf: Division<Span>) -> RResult<Division<Span>, DivisionError>;
62    fn division_strict_remainder(slf: &Division<Strict>, side: DivisionSide) -> FMapping;
63    fn division_strict_bounds(slf: &Division<Strict>) -> RVec<TermBounds>;
64    fn index_new() -> Index;
65    fn index_add(slf: &mut Index, other: Index);
66    fn index_mark_invalid(slf: &mut Index);
67    fn index_add_term(slf: &mut Index, term: Term, value: usize);
68    fn index_add_mapping(slf: &mut Index, mapping: FMapping, value: usize);
69    fn index_ident_value(slf: &Index, ident: Ident) -> RResult<usize, IndexValueError>;
70    fn index_finalize(slf: Index) -> RResult<RSortedMap<Term, usize>, ()>;
71    fn index_gen_indexes(slf: &Index, mapping: FMapping) -> RVec<Index>;
72}
73
74// ── MappingExt ────────────────────────────────────────────────────────────────
75
76/// Methods for [`Mapping`].
77pub trait MappingExt: Sized {
78    /// Returns the size of the mapping expression.
79    fn size(&self) -> usize;
80    /// Pairs two mapping expressions and simplifies.
81    fn pair(self, other: Mapping) -> Mapping;
82    /// Pairs multiple mapping expressions and simplifies.
83    fn pairs(ms: RVec<Mapping>) -> Mapping;
84    /// Divides the mapping expression by a single-term divisor with a span.
85    fn divide_span(&self, divisor: &Mapping, span: usize) -> RResult<Division<Span>, DivisionError>;
86    /// Divides the mapping expression without exposing analysis.
87    fn divide_relaxed(&self, divisor: &Mapping) -> Division<Relaxed>;
88    /// Divides the mapping expression.
89    fn divide_strict(&self, divisor: &Mapping) -> Division<Strict>;
90    /// Factorizes the mapping expression into factor representation.
91    fn factorize(&self) -> FMapping;
92}
93
94impl MappingExt for Mapping {
95    fn size(&self) -> usize {
96        unsafe { mapping_size(self) }
97    }
98    fn pair(self, other: Mapping) -> Mapping {
99        unsafe { mapping_pair(self, other) }
100    }
101    fn pairs(ms: RVec<Mapping>) -> Mapping {
102        unsafe { mapping_pairs(ms) }
103    }
104    fn divide_span(&self, divisor: &Mapping, span: usize) -> RResult<Division<Span>, DivisionError> {
105        unsafe { mapping_divide_span(self, divisor, span) }
106    }
107    fn divide_relaxed(&self, divisor: &Mapping) -> Division<Relaxed> {
108        unsafe { mapping_divide_relaxed(self, divisor) }
109    }
110    fn divide_strict(&self, divisor: &Mapping) -> Division<Strict> {
111        unsafe { mapping_divide_strict(self, divisor) }
112    }
113    fn factorize(&self) -> FMapping {
114        unsafe { mapping_factorize(self) }
115    }
116}
117
118// ── FMappingExt ───────────────────────────────────────────────────────────────
119
120/// Methods for [`FMapping`].
121pub trait FMappingExt: Sized {
122    /// Creates a new empty factor mapping.
123    fn new() -> Self;
124    /// Creates a factor mapping from a slice of terms (axes).
125    fn from_axes(axes: RSlice<Term>) -> Self;
126    /// Pops the outermost factor of the mapping expression.
127    fn pop(&mut self) -> ROption<Factor>;
128    /// Converts the factor mapping into a vector of factors.
129    fn into_inner(self) -> RVec<Factor>;
130    /// Checks if the factor mapping is only padding.
131    fn is_padding(&self) -> bool;
132    /// Returns true if any factor is a Term (not Padding).
133    fn has_terms(&self) -> bool;
134    /// Extracts each Term with its effective stride in this FMapping.
135    fn terms_with_stride(&self) -> RVec<TermPosition>;
136    /// Returns a reference to the factors (innermost first).
137    fn factors(&self) -> RSlice<'_, Factor>;
138    /// Returns true if any term in this FMapping contains the given ident.
139    fn contains_ident(&self, ident: Ident) -> bool;
140    /// Returns the unique idents referenced in this FMapping.
141    fn idents(&self) -> RVec<Ident>;
142    /// Returns the original declared size of the given ident's Symbol.
143    fn find_symbol_size(&self, ident: Ident) -> ROption<usize>;
144    /// Evaluates this FMapping at the given position.
145    fn eval(&self, position: usize) -> Index;
146    /// Returns true if the given ident is isolated.
147    fn is_ident_isolated(&self, ident: Ident) -> bool;
148    /// Converts the factor mapping into a single factor.
149    fn into_factor(self) -> Factor;
150    /// Multiplies two factor mappings (composes outer × inner).
151    fn mul(self, inner: FMapping) -> FMapping;
152    /// Applies stride to the factor mapping.
153    fn stride(self, stride: usize) -> FMapping;
154    /// Applies modulo to the factor mapping.
155    fn modulo(self, modulo: usize) -> FMapping;
156    /// Returns true if `self` is a resize of `original`.
157    fn is_resize_of(&self, original: &FMapping) -> bool;
158    /// Converts the factor mapping into a mapping expression.
159    fn to_mapping(&self) -> Mapping;
160    /// Divides the factor mapping by a single-term divisor with a span.
161    fn divide_span(self, divisor: FMapping, span: usize) -> RResult<Division<Span>, DivisionError>;
162    /// Divides the factor mapping with read-accessible (Top) matched-hole padding.
163    fn divide_relaxed(self, divisor: FMapping) -> Division<Relaxed>;
164    /// Divides the factor mapping with analysis-capable (Bottom) matched-hole padding.
165    fn divide_strict(self, divisor: FMapping) -> Division<Strict>;
166    /// Normalizes the factor mapping via round-trip through `Mapping`.
167    fn normalize(self) -> FMapping;
168    /// Removes all right padding.
169    fn remove_padding(self) -> FMapping;
170    /// Splits the mapping at `target` from the innermost side.
171    fn split_at(&self, target: usize) -> Tuple2<FMapping, FMapping>;
172    /// Removes existing padding and pads to `target` size.
173    fn pad(self, target: usize) -> FMapping;
174    /// Applies padding to the factor mapping.
175    fn padding(self, padding: usize, kind: PaddingKind) -> FMapping;
176}
177
178impl FMappingExt for FMapping {
179    fn new() -> Self {
180        FMapping::new()
181    }
182    fn from_axes(axes: RSlice<Term>) -> Self {
183        unsafe { fmapping_from_axes(axes) }
184    }
185    fn pop(&mut self) -> ROption<Factor> {
186        unsafe { fmapping_pop(self) }
187    }
188    fn into_inner(self) -> RVec<Factor> {
189        unsafe { fmapping_into_inner(self) }
190    }
191    fn is_padding(&self) -> bool {
192        unsafe { fmapping_is_padding(self) }
193    }
194    fn has_terms(&self) -> bool {
195        unsafe { fmapping_has_terms(self) }
196    }
197    fn terms_with_stride(&self) -> RVec<TermPosition> {
198        unsafe { fmapping_terms_with_stride(self) }
199    }
200    fn factors(&self) -> RSlice<'_, Factor> {
201        unsafe { fmapping_factors(self) }
202    }
203    fn contains_ident(&self, ident: Ident) -> bool {
204        unsafe { fmapping_contains_ident(self, ident) }
205    }
206    fn idents(&self) -> RVec<Ident> {
207        unsafe { fmapping_idents(self) }
208    }
209    fn find_symbol_size(&self, ident: Ident) -> ROption<usize> {
210        unsafe { fmapping_find_symbol_size(self, ident) }
211    }
212    fn eval(&self, position: usize) -> Index {
213        unsafe { fmapping_eval(self, position) }
214    }
215    fn is_ident_isolated(&self, ident: Ident) -> bool {
216        unsafe { fmapping_is_ident_isolated(self, ident) }
217    }
218    fn into_factor(self) -> Factor {
219        unsafe { fmapping_into_factor(self) }
220    }
221    fn mul(self, inner: FMapping) -> FMapping {
222        unsafe { fmapping_mul(self, inner) }
223    }
224    fn stride(self, stride: usize) -> FMapping {
225        unsafe { fmapping_stride(self, stride) }
226    }
227    fn modulo(self, modulo: usize) -> FMapping {
228        unsafe { fmapping_modulo(self, modulo) }
229    }
230    fn is_resize_of(&self, original: &FMapping) -> bool {
231        unsafe { fmapping_is_resize_of(self, original) }
232    }
233    fn to_mapping(&self) -> Mapping {
234        unsafe { fmapping_to_mapping(self) }
235    }
236    fn divide_span(self, divisor: FMapping, span: usize) -> RResult<Division<Span>, DivisionError> {
237        unsafe { fmapping_divide_span(self, divisor, span) }
238    }
239    fn divide_relaxed(self, divisor: FMapping) -> Division<Relaxed> {
240        unsafe { fmapping_divide_relaxed(self, divisor) }
241    }
242    fn divide_strict(self, divisor: FMapping) -> Division<Strict> {
243        unsafe { fmapping_divide_strict(self, divisor) }
244    }
245    fn normalize(self) -> FMapping {
246        unsafe { fmapping_normalize(self) }
247    }
248    fn remove_padding(self) -> FMapping {
249        unsafe { fmapping_remove_padding(self) }
250    }
251    fn split_at(&self, target: usize) -> Tuple2<FMapping, FMapping> {
252        unsafe { fmapping_split_at(self, target) }
253    }
254    fn pad(self, target: usize) -> FMapping {
255        unsafe { fmapping_pad(self, target) }
256    }
257    fn padding(self, padding: usize, kind: PaddingKind) -> FMapping {
258        unsafe { fmapping_padding(self, padding, kind) }
259    }
260}
261
262// ── AtomExt ───────────────────────────────────────────────────────────────────
263
264/// Methods for [`Atom`].
265pub trait AtomExt {
266    /// Returns the size of the atomic mapping expression.
267    fn size(&self) -> usize;
268    /// Returns the idents contained in this atom.
269    fn idents(&self) -> RVec<Ident>;
270    /// Returns true if this atom contains the given ident (recursively).
271    fn contains_ident(&self, ident: Ident) -> bool;
272    /// Returns the original declared size of the given ident's Symbol.
273    fn find_symbol_size(&self, ident: Ident) -> ROption<usize>;
274}
275
276impl AtomExt for Atom {
277    fn size(&self) -> usize {
278        unsafe { atom_size(self) }
279    }
280    fn idents(&self) -> RVec<Ident> {
281        unsafe { atom_idents(self) }
282    }
283    fn contains_ident(&self, ident: Ident) -> bool {
284        unsafe { atom_contains_ident(self, ident) }
285    }
286    fn find_symbol_size(&self, ident: Ident) -> ROption<usize> {
287        unsafe { atom_find_symbol_size(self, ident) }
288    }
289}
290
291// ── TermExt ───────────────────────────────────────────────────────────────────
292
293/// Methods for [`Term`].
294pub trait TermExt {
295    /// Returns the depth of the term.
296    fn depth(&self) -> usize;
297    /// Returns the idents contained in this term.
298    fn idents(&self) -> RVec<Ident>;
299}
300
301impl TermExt for Term {
302    fn depth(&self) -> usize {
303        unsafe { term_depth(self) }
304    }
305    fn idents(&self) -> RVec<Ident> {
306        unsafe { term_idents(self) }
307    }
308}
309
310// ── FactorExt ─────────────────────────────────────────────────────────────────
311
312/// Methods for [`Factor`].
313pub trait FactorExt {
314    /// Creates a padding factor.
315    fn padding(size: usize, kind: PaddingKind) -> Factor;
316    /// Returns the idents contained in this factor.
317    fn idents(&self) -> RVec<Ident>;
318}
319
320impl FactorExt for Factor {
321    fn padding(size: usize, kind: PaddingKind) -> Factor {
322        unsafe { factor_padding(size, kind) }
323    }
324    fn idents(&self) -> RVec<Ident> {
325        unsafe { factor_idents(self) }
326    }
327}
328
329// ── IndexExt ──────────────────────────────────────────────────────────────────
330
331/// Methods for [`Index`].
332pub trait IndexExt: Sized {
333    /// Creates a new empty index.
334    fn new() -> Self;
335    /// Adds another index to this index.
336    fn add(&mut self, other: Index);
337    /// Marks this index as invalid.
338    fn mark_invalid(&mut self);
339    /// Adds a term to this index.
340    fn add_term(&mut self, term: Term, value: usize);
341    /// Adds a mapping to this index.
342    fn add_mapping<I: crate::M>(&mut self, value: usize);
343    /// Returns the contribution of a specific ident to this Index.
344    fn ident_value(&self, ident: Ident) -> RResult<usize, IndexValueError>;
345    /// Finalizes the index.
346    fn finalize(self) -> RResult<RSortedMap<Term, usize>, ()>;
347    /// Generates all possible indexes based on the given mapping.
348    fn gen_indexes(&self, mapping: FMapping) -> RVec<Index>;
349}
350
351impl IndexExt for Index {
352    fn new() -> Self {
353        unsafe { index_new() }
354    }
355    fn add(&mut self, other: Index) {
356        unsafe { index_add(self, other) }
357    }
358    fn mark_invalid(&mut self) {
359        unsafe { index_mark_invalid(self) }
360    }
361    fn add_term(&mut self, term: Term, value: usize) {
362        unsafe { index_add_term(self, term, value) }
363    }
364    fn add_mapping<I: crate::M>(&mut self, value: usize) {
365        unsafe { index_add_mapping(self, mapping_factorize(&I::to_value()), value) }
366    }
367    fn ident_value(&self, ident: Ident) -> RResult<usize, IndexValueError> {
368        unsafe { index_ident_value(self, ident) }
369    }
370    fn finalize(self) -> RResult<RSortedMap<Term, usize>, ()> {
371        unsafe { index_finalize(self) }
372    }
373    fn gen_indexes(&self, mapping: FMapping) -> RVec<Index> {
374        unsafe { index_gen_indexes(self, mapping) }
375    }
376}
377
378// ── DivisionExt ───────────────────────────────────────────────────────────────
379
380/// Methods for [`Division<M>`].
381pub trait DivisionExt<Md: DivisionMode> {
382    /// Returns `Ok(self)` if all divisor terms were matched.
383    fn exact(self) -> RResult<Division<Md>, DivisionError>;
384}
385
386impl DivisionExt<Relaxed> for Division<Relaxed> {
387    fn exact(self) -> RResult<Division<Relaxed>, DivisionError> {
388        unsafe { division_exact_relaxed(self) }
389    }
390}
391
392impl DivisionExt<Strict> for Division<Strict> {
393    fn exact(self) -> RResult<Division<Strict>, DivisionError> {
394        unsafe { division_exact_strict(self) }
395    }
396}
397
398impl DivisionExt<Span> for Division<Span> {
399    fn exact(self) -> RResult<Division<Span>, DivisionError> {
400        unsafe { division_exact_span(self) }
401    }
402}
403
404/// Methods specific to [`Division<Strict>`].
405pub trait StrictDivisionExt {
406    /// Returns the original mapping with matched terms removed.
407    fn remainder(&self, side: DivisionSide) -> FMapping;
408    /// Returns per-term padding bounds reconstructed from the strict residue.
409    fn bounds(&self) -> RVec<TermBounds>;
410}
411
412impl StrictDivisionExt for Division<Strict> {
413    fn remainder(&self, side: DivisionSide) -> FMapping {
414        unsafe { division_strict_remainder(self, side) }
415    }
416
417    /// Returns per-term padding bounds reconstructed from the strict residue.
418    ///
419    /// This reports the compact block bounds used by [`remainder`](Self::remainder).  For example, `m![A #
420    /// 16].divide_strict(m![A]).bounds()` reports dividend bounds `4..16` and divisor bounds `4..4` when `A = 4`.
421    ///
422    /// ```rust
423    /// use furiosa_mapping::*;
424    /// use furiosa_opt_macro::m;
425    ///
426    /// axes![A = 4];
427    /// let division = <m![A # 16]>::to_value()
428    ///     .factorize()
429    ///     .divide_strict(<m![A]>::to_value().factorize());
430    /// let [row]: [TermBounds; 1] = division.bounds().into_vec().try_into().unwrap();
431    ///
432    /// assert_eq!(row.dividend, BlockBounds { min: 4, max: 16 });
433    /// assert_eq!(row.divisor, BlockBounds { min: 4, max: 4 });
434    /// ```
435    fn bounds(&self) -> RVec<TermBounds> {
436        unsafe { division_strict_bounds(self) }
437    }
438}