Skip to main content

vyre_foundation/ir_inner/model/program/
buffer_decl.rs

1use std::ops::Range;
2use std::sync::Arc;
3
4use crate::ir_inner::model::spec_types::{BufferAccess, DataType};
5
6use super::{MemoryHints, MemoryKind};
7
8/// Linear-type discipline for a buffer binding.
9///
10/// Vyre's IR is moving from an unrestricted-by-default world toward
11/// a substructural type system: a buffer can be marked `Linear`
12/// (must be used exactly once on each path through the Program),
13/// `Affine` (used at most once  -  drops are fine), `Relevant`
14/// (used at least once), or `Unrestricted` (the historical default).
15/// The type-checker pass (P-1.0-V2.2) verifies these assertions
16/// before lowering; backends that hit a violation reject the
17/// program at validation time instead of producing wrong code.
18///
19/// `Unrestricted` is the safe default when authoring a `BufferDecl`
20/// for back-compat  -  every existing program continues to type-check.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
22#[non_exhaustive]
23pub enum LinearType {
24    /// Use exactly once on every path. Forbids both drop-without-use
25    /// and double-use.
26    Linear,
27    /// Use at most once on every path. Allows drop-without-use,
28    /// forbids double-use.
29    Affine,
30    /// Use at least once on every path. Forbids drop-without-use,
31    /// allows double-use.
32    Relevant,
33    /// No discipline applied. Default for back-compat with the
34    /// pre-V2.x IR.
35    #[default]
36    Unrestricted,
37}
38
39impl LinearType {
40    /// Whether this discipline forbids dropping a buffer without
41    /// using it (`Linear` or `Relevant`).
42    #[must_use]
43    #[inline]
44    pub const fn forbids_drop(self) -> bool {
45        matches!(self, Self::Linear | Self::Relevant)
46    }
47
48    /// Whether this discipline forbids using a buffer more than once
49    /// (`Linear` or `Affine`).
50    #[must_use]
51    #[inline]
52    pub const fn forbids_reuse(self) -> bool {
53        matches!(self, Self::Linear | Self::Affine)
54    }
55}
56
57/// Refinement predicate over a buffer's element count (P-1.0-V3.1).
58///
59/// Represents a small grammar of constraints a `BufferDecl` author
60/// can attach. The validator (P-1.0-V3.2) checks each predicate
61/// against the program's static count and the optimizer (P-1.0-V3.3)
62/// uses verified predicates to prove loop-bound and alignment
63/// invariants for vectorization.
64///
65/// `None` (the default) is "unconstrained"; existing programs keep
66/// their current behavior.
67#[derive(Debug, Clone, PartialEq, Eq, Hash)]
68#[non_exhaustive]
69pub enum ShapePredicate {
70    /// `count >= n`. Holds when the runtime element count is at
71    /// least `n`. Used to prove non-empty workgroup buffers and
72    /// minimum vectorization tile sizes.
73    AtLeast(u32),
74    /// `count <= n`. Holds when the count never exceeds `n`. Used
75    /// to bound dispatch sizes and prevent oversized allocations.
76    AtMost(u32),
77    /// `count == n`. The strongest constraint; the count is fixed.
78    Exactly(u32),
79    /// `count % n == 0`. Used for alignment proofs (e.g. SIMD lanes).
80    MultipleOf(u32),
81    /// `count % modulus == remainder`. Invalid modular forms evaluate
82    /// false, so static validation catches impossible declarations.
83    ModEquals {
84        /// Divisor used by the modular equality.
85        modulus: u32,
86        /// Required remainder. Must be less than `modulus` to match.
87        remainder: u32,
88    },
89    /// `min <= count * scale + offset <= max`, evaluated with wide
90    /// arithmetic for frontend-derived affine constraints.
91    AffineRange {
92        /// Multiplicative coefficient applied to `count`.
93        scale: i64,
94        /// Constant term added after scaling.
95        offset: i64,
96        /// Inclusive lower bound for the affine expression.
97        min: i64,
98        /// Inclusive upper bound for the affine expression.
99        max: i64,
100    },
101    /// Conjunction of two predicates (`p1 && p2`). Both must hold.
102    And(Box<ShapePredicate>, Box<ShapePredicate>),
103    /// Disjunction of two predicates (`p1 || p2`). Either may hold.
104    Or(Box<ShapePredicate>, Box<ShapePredicate>),
105    /// Negation of a predicate.
106    Not(Box<ShapePredicate>),
107}
108
109impl ShapePredicate {
110    /// Evaluate the predicate against a concrete `count`. Returns
111    /// `true` when the predicate holds. P-1.0-V3.2 uses this from
112    /// the `validate()` pass; P-1.0-V3.3 calls it from optimizer
113    /// passes that need a yes/no proof.
114    #[must_use]
115    pub fn holds(&self, count: u32) -> bool {
116        match self {
117            Self::AtLeast(n) => count >= *n,
118            Self::AtMost(n) => count <= *n,
119            Self::Exactly(n) => count == *n,
120            Self::MultipleOf(n) => *n != 0 && count % *n == 0,
121            Self::ModEquals { modulus, remainder } => {
122                *modulus != 0 && *remainder < *modulus && count % *modulus == *remainder
123            }
124            Self::AffineRange {
125                scale,
126                offset,
127                min,
128                max,
129            } => {
130                let value = i128::from(count) * i128::from(*scale) + i128::from(*offset);
131                value >= i128::from(*min) && value <= i128::from(*max)
132            }
133            Self::And(a, b) => a.holds(count) && b.holds(count),
134            Self::Or(a, b) => a.holds(count) || b.holds(count),
135            Self::Not(inner) => !inner.holds(count),
136        }
137    }
138
139    /// Evaluate the predicate against a concrete count.
140    #[must_use]
141    pub fn evaluate(&self, count: u32) -> bool {
142        self.holds(count)
143    }
144
145    /// Whether this predicate proves that the count cannot be zero.
146    #[must_use]
147    pub fn proves_non_empty(&self) -> bool {
148        match self {
149            Self::AtLeast(n) | Self::Exactly(n) => *n > 0,
150            Self::ModEquals { modulus, remainder } => {
151                *modulus != 0 && *remainder < *modulus && *remainder > 0
152            }
153            Self::AffineRange {
154                offset, min, max, ..
155            } => {
156                let zero_value = i128::from(*offset);
157                zero_value < i128::from(*min) || zero_value > i128::from(*max)
158            }
159            Self::And(left, right) => left.proves_non_empty() || right.proves_non_empty(),
160            Self::Or(left, right) => left.proves_non_empty() && right.proves_non_empty(),
161            _ => false,
162        }
163    }
164
165    /// Human-readable form for error messages.
166    #[must_use]
167    pub fn describe(&self) -> String {
168        match self {
169            Self::AtLeast(n) => format!("count >= {n}"),
170            Self::AtMost(n) => format!("count <= {n}"),
171            Self::Exactly(n) => format!("count == {n}"),
172            Self::MultipleOf(n) => format!("count % {n} == 0"),
173            Self::ModEquals { modulus, remainder } => format!("count % {modulus} == {remainder}"),
174            Self::AffineRange {
175                scale,
176                offset,
177                min,
178                max,
179            } => {
180                format!("{min} <= count * {scale} + {offset} <= {max}")
181            }
182            Self::And(a, b) => format!("({}) && ({})", a.describe(), b.describe()),
183            Self::Or(a, b) => format!("({}) || ({})", a.describe(), b.describe()),
184            Self::Not(inner) => format!("!({})", inner.describe()),
185        }
186    }
187}
188
189/// A named buffer binding in a program.
190///
191/// # Examples
192///
193/// ```
194/// use vyre::ir::{BufferDecl, BufferAccess, DataType};
195///
196/// let buf = BufferDecl::read("input", 0, DataType::U32);
197/// assert_eq!(buf.name(), "input");
198/// assert_eq!(buf.binding(), 0);
199/// ```
200#[derive(Debug, Clone, PartialEq, Eq, Hash)]
201pub struct BufferDecl {
202    /// Human-readable name. Referenced by `Expr::Load`, `Node::Store`, etc.
203    pub name: Arc<str>,
204    /// Binding slot: `@binding(N)`. All buffers are in `@group(0)`.
205    /// Ignored for `BufferAccess::Workgroup`.
206    pub binding: u32,
207    /// Access mode.
208    pub access: BufferAccess,
209    /// Memory tier.
210    pub kind: MemoryKind,
211    /// Element data type.
212    pub element: DataType,
213    /// Number of elements.
214    ///
215    /// For `Workgroup` memory this is the static array length.
216    /// For storage and uniform buffers this is `0` (runtime-sized).
217    pub count: u32,
218    /// Whether this buffer is the scalar expression output for composition inlining.
219    pub is_output: bool,
220    /// Whether the end-to-end pipeline reads this buffer after Program execution.
221    ///
222    /// Passes must treat this as an externally-visible sink even when the IR
223    /// itself does not read the buffer again.
224    pub pipeline_live_out: bool,
225    /// Optional byte range to read back from this output buffer.
226    ///
227    /// `None` preserves the historical behavior and reads back the full
228    /// declared output buffer.
229    pub output_byte_range: Option<Range<usize>>,
230    /// Non-binding backend optimization hints.
231    pub hints: MemoryHints,
232    /// When true, admits `DataType::Bytes` load/store despite V013.
233    ///
234    /// Bytes-producing or bytes-extraction ops (decode.base64,
235    /// `compression.lz4_decompress`, `match.dfa_scan` position emission, etc.)
236    /// opt into V013 relaxation per-buffer. Default false keeps scalar
237    /// arithmetic protected from accidental bytes-blob reinterpretation.
238    pub bytes_extraction: bool,
239    /// Linear-type discipline for this buffer (P-1.0-V2.1).
240    ///
241    /// Defaults to `LinearType::Unrestricted` so existing programs
242    /// continue to type-check. Authors opt in by calling
243    /// [`BufferDecl::with_linear_type`]. The type-checker pass
244    /// (`crate::validate::linear_type`) walks the IR and
245    /// rejects programs that violate the declared discipline; backends
246    /// that hit a violation surface it as a validation error before
247    /// lowering.
248    pub linear_type: LinearType,
249    /// Optional shape-refinement predicate (P-1.0-V3.1).
250    ///
251    /// `None` is the default (no shape constraint, identical to the
252    /// pre-V3.x IR). Authors opt in via
253    /// [`BufferDecl::with_shape_predicate`]. The validator
254    /// ([`crate::validate::shape_predicate::check_shape_predicates`])
255    /// evaluates each predicate against the program's static `count`
256    /// at `validate()` time and rejects programs whose static shape
257    /// contradicts the declaration.
258    pub shape_predicate: Option<ShapePredicate>,
259}
260
261impl BufferDecl {
262    /// Create a storage buffer declaration.
263    ///
264    /// # Examples
265    ///
266    /// ```
267    /// use vyre::ir::{BufferDecl, BufferAccess, DataType};
268    /// let _ = BufferDecl::storage("a", 0, BufferAccess::ReadOnly, DataType::U32);
269    /// ```
270    #[must_use]
271    #[inline]
272    pub fn storage(name: &str, binding: u32, access: BufferAccess, element: DataType) -> Self {
273        let kind = match &access {
274            BufferAccess::ReadOnly => MemoryKind::Readonly,
275            BufferAccess::Uniform => MemoryKind::Uniform,
276            BufferAccess::Workgroup => MemoryKind::Shared,
277            _ => MemoryKind::Global,
278        };
279        Self {
280            name: Arc::from(name),
281            binding,
282            access,
283            kind,
284            element,
285            count: 0,
286            is_output: false,
287            pipeline_live_out: false,
288            output_byte_range: None,
289            hints: MemoryHints::default(),
290            bytes_extraction: false,
291            linear_type: LinearType::default(),
292            shape_predicate: None,
293        }
294    }
295
296    /// Shorthand for a read-only storage buffer.
297    ///
298    /// # Examples
299    ///
300    /// ```
301    /// use vyre::ir::{BufferDecl, DataType};
302    /// let _ = BufferDecl::read("a", 0, DataType::U32);
303    /// ```
304    #[must_use]
305    #[inline]
306    pub fn read(name: &str, binding: u32, element: DataType) -> Self {
307        Self::storage(name, binding, BufferAccess::ReadOnly, element)
308    }
309
310    /// Shorthand for a read-write storage buffer.
311    ///
312    /// # Examples
313    ///
314    /// ```
315    /// use vyre::ir::{BufferDecl, DataType};
316    /// let _ = BufferDecl::read_write("a", 0, DataType::U32);
317    /// ```
318    #[must_use]
319    #[inline]
320    pub fn read_write(name: &str, binding: u32, element: DataType) -> Self {
321        Self::storage(name, binding, BufferAccess::ReadWrite, element)
322    }
323
324    /// Shorthand for the read-write result buffer used by call inlining.
325    ///
326    /// # Examples
327    ///
328    /// ```
329    /// use vyre::ir::{BufferDecl, DataType};
330    /// let _ = BufferDecl::output("a", 0, DataType::U32);
331    /// ```
332    #[must_use]
333    #[inline]
334    pub fn output(name: &str, binding: u32, element: DataType) -> Self {
335        Self {
336            is_output: true,
337            pipeline_live_out: true,
338            ..Self::read_write(name, binding, element)
339        }
340    }
341
342    /// Mark whether a caller/backend observes this buffer after Program execution.
343    #[must_use]
344    #[inline]
345    pub fn with_pipeline_live_out(mut self, flag: bool) -> Self {
346        self.pipeline_live_out = flag;
347        self
348    }
349
350    /// Attach an output byte range for backends that can read back a slice.
351    #[must_use]
352    #[inline]
353    pub fn with_output_byte_range(mut self, range: Range<usize>) -> Self {
354        self.output_byte_range = Some(range);
355        self
356    }
357
358    /// Set the static element count for storage-style buffers.
359    ///
360    /// Set the element count. A count of `0` retains the IR's
361    /// runtime-sized-buffer representation; validators reject zero-sized
362    /// workgroup allocations before dispatch.
363    #[must_use]
364    #[inline]
365    pub fn with_count(mut self, count: u32) -> Self {
366        self.count = count;
367        self
368    }
369
370    /// Shorthand for a uniform buffer.
371    ///
372    /// # Examples
373    ///
374    /// ```
375    /// use vyre::ir::{BufferDecl, DataType};
376    /// let _ = BufferDecl::uniform("a", 0, DataType::U32);
377    /// ```
378    #[must_use]
379    #[inline]
380    pub fn uniform(name: &str, binding: u32, element: DataType) -> Self {
381        Self::storage(name, binding, BufferAccess::Uniform, element)
382    }
383
384    /// Shorthand for a workgroup-local shared array.
385    ///
386    /// `count` is the static number of elements visible to all invocations
387    /// in the same workgroup.
388    ///
389    /// # Examples
390    ///
391    /// ```
392    /// use vyre::ir::{BufferAccess, BufferDecl, DataType, MemoryKind};
393    ///
394    /// let scratch = BufferDecl::workgroup("scratch", 64, DataType::U32);
395    ///
396    /// assert_eq!(scratch.name(), "scratch");
397    /// assert_eq!(scratch.access(), BufferAccess::Workgroup);
398    /// assert_eq!(scratch.kind(), MemoryKind::Shared);
399    /// assert_eq!(scratch.count(), 64);
400    /// ```
401    #[must_use]
402    #[inline]
403    pub fn workgroup(name: &str, count: u32, element: DataType) -> Self {
404        Self {
405            name: Arc::from(name),
406            binding: 0,
407            access: BufferAccess::Workgroup,
408            kind: MemoryKind::Shared,
409            element,
410            count,
411            is_output: false,
412            pipeline_live_out: false,
413            output_byte_range: None,
414            hints: MemoryHints::default(),
415            bytes_extraction: false,
416            linear_type: LinearType::default(),
417            shape_predicate: None,
418        }
419    }
420
421    /// Mark this buffer as a bytes-extraction context so V013 admits Bytes load/store.
422    #[must_use]
423    #[inline]
424    pub fn with_bytes_extraction(mut self, flag: bool) -> Self {
425        self.bytes_extraction = flag;
426        self
427    }
428
429    /// Set the linear-type discipline (P-1.0-V2.1).
430    ///
431    /// Defaults to [`LinearType::Unrestricted`] from the constructor;
432    /// the type-checker pass enforces stricter disciplines when set.
433    #[must_use]
434    #[inline]
435    pub fn with_linear_type(mut self, linear_type: LinearType) -> Self {
436        self.linear_type = linear_type;
437        self
438    }
439
440    /// Set the shape-refinement predicate (P-1.0-V3.1).
441    ///
442    /// Defaults to `None` (unconstrained); the validator
443    /// ([`crate::validate::shape_predicate::check_shape_predicates`])
444    /// rejects programs whose static `count` violates the predicate.
445    #[must_use]
446    #[inline]
447    pub fn with_shape_predicate(mut self, predicate: ShapePredicate) -> Self {
448        self.shape_predicate = Some(predicate);
449        self
450    }
451
452    /// Override the memory tier.
453    #[must_use]
454    #[inline]
455    pub fn with_kind(mut self, kind: MemoryKind) -> Self {
456        self.kind = kind;
457        self
458    }
459
460    /// Override memory optimization hints.
461    #[must_use]
462    #[inline]
463    pub fn with_hints(mut self, hints: MemoryHints) -> Self {
464        self.hints = hints;
465        self
466    }
467
468    /// Buffer name.
469    #[must_use]
470    #[inline]
471    pub fn name(&self) -> &str {
472        &self.name
473    }
474
475    /// Binding slot.
476    #[must_use]
477    #[inline]
478    pub fn binding(&self) -> u32 {
479        self.binding
480    }
481
482    /// Buffer access mode.
483    #[must_use]
484    #[inline]
485    pub fn access(&self) -> BufferAccess {
486        self.access.clone()
487    }
488
489    /// Memory tier.
490    #[must_use]
491    #[inline]
492    pub fn kind(&self) -> MemoryKind {
493        self.kind
494    }
495
496    /// Non-binding memory hints.
497    #[must_use]
498    #[inline]
499    pub fn hints(&self) -> MemoryHints {
500        self.hints
501    }
502
503    /// Element data type.
504    #[must_use]
505    #[inline]
506    pub fn element(&self) -> DataType {
507        self.element.clone()
508    }
509
510    /// Static element count for workgroup buffers.
511    #[must_use]
512    #[inline]
513    pub fn count(&self) -> u32 {
514        self.count
515    }
516
517    /// Whether [`Self::count`] is a static array length that reaches generated
518    /// backend code, rather than a runtime-sized binding length.
519    ///
520    /// This mirrors, arm for arm, the `MemoryClass::Shared` and
521    /// `MemoryClass::Scratch` cases of `vyre_lower::lower::memory_class`, which
522    /// is the single Program-to-descriptor boundary every emitter reads. Those
523    /// two classes are the ones whose `element_count` becomes a fixed-length
524    /// array in emitted code (`.shared` byte length in PTX,
525    /// `array<T, N>` in WGSL); every other class emits a runtime-sized array
526    /// and ignores the count.
527    ///
528    /// `Persistent` is excluded because it is rejected before classification.
529    ///
530    /// Compiled-pipeline cache identity uses this to decide whether `count`
531    /// belongs in the cache key: including it for a runtime-sized storage
532    /// buffer would recompile the shader on every buffer resize, and omitting
533    /// it for a workgroup buffer would let two different shared-memory array
534    /// lengths share one cache entry.
535    #[must_use]
536    #[inline]
537    pub fn has_static_element_count(&self) -> bool {
538        match (self.kind, &self.access) {
539            (MemoryKind::Persistent, _) => false,
540            (MemoryKind::Shared, _) | (_, BufferAccess::Workgroup) => true,
541            (MemoryKind::Local, _) => true,
542            _ => false,
543        }
544    }
545
546    /// Static packed byte length for fixed-size buffers.
547    ///
548    /// Returns `Ok(None)` for runtime-sized buffer declarations (`count == 0`)
549    /// and for fixed-count buffers whose element type is runtime-sized. Sub-byte
550    /// element types use their packed bit width, so three `I4` elements occupy
551    /// two bytes rather than three conservative one-byte lanes.
552    ///
553    /// # Errors
554    ///
555    /// Returns an actionable diagnostic when the packed byte count overflows.
556    pub fn static_byte_len(&self) -> Result<Option<usize>, String> {
557        let count = usize::try_from(self.count).map_err(|error| {
558            format!(
559                "buffer `{}` static element count {} cannot fit usize ({error}). Fix: split the buffer or reduce its element count.",
560                self.name, self.count
561            )
562        })?;
563        if count == 0 {
564            return Ok(None);
565        }
566        self.element.packed_size_bytes(count).map_err(|error| {
567            format!(
568                "buffer `{}` static byte length could not be computed: {error}. Fix: use a fixed-width element type or split the buffer.",
569                self.name
570            )
571        })
572    }
573
574    /// Return true when this buffer is the unique inlining result buffer.
575    #[must_use]
576    #[inline]
577    pub fn is_output(&self) -> bool {
578        self.is_output
579    }
580
581    /// Return true when the buffer must survive IR-local deadness analysis.
582    #[must_use]
583    #[inline]
584    pub fn is_pipeline_live_out(&self) -> bool {
585        self.pipeline_live_out
586    }
587
588    /// True when a backend ALLOCATES this buffer's storage (an output it writes) rather
589    /// than reading it from the dispatch inputs: an `is_output` buffer, any `WriteOnly`
590    /// buffer, or a `pipeline_live_out` `ReadWrite` intermediate.
591    ///
592    /// This is the SINGLE cross-backend definition of "backend-allocated output": the
593    /// reference interpreter, the CpuRef backend, and the device drivers MUST all agree
594    /// on which buffers they allocate-and-write vs. read from inputs, so each calls this
595    /// method instead of re-deriving the predicate. Drift here would make the interpreter
596    /// and a backend disagree on a program's outputs (a silent readback bug).
597    #[must_use]
598    #[inline]
599    pub fn is_backend_allocated_output(&self) -> bool {
600        self.is_output()
601            || matches!(self.access, BufferAccess::WriteOnly)
602            || (self.is_pipeline_live_out() && matches!(self.access, BufferAccess::ReadWrite))
603    }
604
605    /// Refuse this buffer when it is backend-allocated and has no static size.
606    ///
607    /// A buffer selected by [`Self::is_backend_allocated_output`] never receives
608    /// host bytes, so `count == 0` leaves an executor nothing to size its
609    /// allocation or its readback from. Every execution path calls this rather
610    /// than re-deriving the condition, so the reference interpreter refuses
611    /// exactly what the device backends refuse. An oracle that accepts what its
612    /// targets reject certifies programs that cannot run.
613    ///
614    /// A writable buffer that is NOT backend-allocated (a plain `ReadWrite`) is
615    /// deliberately accepted: it consumes one host input slot, so its element
616    /// count is inferable from the bytes the caller supplies and is resolved per
617    /// dispatch.
618    ///
619    /// # Examples
620    ///
621    /// ```
622    /// use vyre::ir::{BufferDecl, DataType};
623    ///
624    /// // No count and backend-allocated: refused, and the message names the remedy.
625    /// let error = BufferDecl::output("out", 0, DataType::U32)
626    ///     .require_static_readback_size()
627    ///     .expect_err("a countless output has no readback size");
628    /// assert!(error.contains(".with_count(n)"));
629    ///
630    /// // A count makes it well-formed.
631    /// assert!(BufferDecl::output("out", 0, DataType::U32)
632    ///     .with_count(4)
633    ///     .require_static_readback_size()
634    ///     .is_ok());
635    ///
636    /// // A plain read_write takes its size from the caller's bytes, so it is fine.
637    /// assert!(BufferDecl::read_write("rw", 1, DataType::U32)
638    ///     .require_static_readback_size()
639    ///     .is_ok());
640    /// ```
641    ///
642    /// # Errors
643    ///
644    /// Returns the operator-facing message when this buffer is backend-allocated
645    /// and its readback size cannot be determined. A `count` of zero means
646    /// "runtime-sized" rather than "zero elements", since `count` defaults to
647    /// zero and so cannot represent a declared empty buffer. An explicit
648    /// `output_byte_range` states the readback size directly, so it satisfies
649    /// this check on its own: that is how a legitimately EMPTY output declares
650    /// itself, with `.with_output_byte_range(0..0)`. Without that escape an
651    /// empty-input program is indistinguishable from a mis-declared one, and the
652    /// only way to pass is to inflate the count to a nonzero value the buffer
653    /// does not have.
654    #[inline]
655    pub fn require_static_readback_size(&self) -> Result<(), String> {
656        if self.is_backend_allocated_output() && self.count == 0 && self.output_byte_range.is_none()
657        {
658            return Err(format!(
659                "backend-allocated output buffer `{}` has no static element count and no output byte range, so its readback size is unknown. Fix: declare it with .with_count(n), or with .with_output_byte_range(0..0) if it is genuinely empty.",
660                self.name()
661            ));
662        }
663        Ok(())
664    }
665
666    /// Byte range the consumer needs from this output buffer, if declared.
667    #[must_use]
668    #[inline]
669    pub fn output_byte_range(&self) -> Option<Range<usize>> {
670        self.output_byte_range.clone()
671    }
672
673    /// Linear-type discipline (P-1.0-V2.1).
674    #[must_use]
675    #[inline]
676    pub fn linear_type(&self) -> LinearType {
677        self.linear_type
678    }
679
680    /// Shape-refinement predicate (P-1.0-V3.1).
681    #[must_use]
682    #[inline]
683    pub fn shape_predicate(&self) -> Option<&ShapePredicate> {
684        self.shape_predicate.as_ref()
685    }
686}
687
688#[cfg(test)]
689mod linear_type_tests {
690    use super::*;
691
692    #[test]
693    fn default_is_unrestricted() {
694        let buf = BufferDecl::read("a", 0, DataType::U32);
695        assert_eq!(buf.linear_type(), LinearType::Unrestricted);
696        assert!(!LinearType::Unrestricted.forbids_drop());
697        assert!(!LinearType::Unrestricted.forbids_reuse());
698    }
699
700    #[test]
701    fn linear_forbids_both() {
702        assert!(LinearType::Linear.forbids_drop());
703        assert!(LinearType::Linear.forbids_reuse());
704    }
705
706    #[test]
707    fn affine_forbids_only_reuse() {
708        assert!(!LinearType::Affine.forbids_drop());
709        assert!(LinearType::Affine.forbids_reuse());
710    }
711
712    #[test]
713    fn relevant_forbids_only_drop() {
714        assert!(LinearType::Relevant.forbids_drop());
715        assert!(!LinearType::Relevant.forbids_reuse());
716    }
717
718    #[test]
719    fn with_linear_type_is_round_trip() {
720        for lt in [
721            LinearType::Linear,
722            LinearType::Affine,
723            LinearType::Relevant,
724            LinearType::Unrestricted,
725        ] {
726            let buf = BufferDecl::read("a", 0, DataType::U32).with_linear_type(lt);
727            assert_eq!(buf.linear_type(), lt);
728        }
729    }
730
731    #[test]
732    fn workgroup_constructor_defaults_to_unrestricted() {
733        let buf = BufferDecl::workgroup("scratch", 64, DataType::U32);
734        assert_eq!(buf.linear_type(), LinearType::Unrestricted);
735    }
736
737    #[test]
738    fn static_byte_len_uses_packed_subbyte_width() {
739        let buf = BufferDecl::read("packed_i4", 0, DataType::I4).with_count(3);
740        assert_eq!(
741            buf.static_byte_len()
742                .expect("Fix: packed I4 byte length must compute"),
743            Some(2)
744        );
745    }
746
747    #[test]
748    fn static_byte_len_marks_runtime_sized_buffers_dynamic() {
749        let zero_count = BufferDecl::read("dynamic_count", 0, DataType::U32);
750        assert_eq!(
751            zero_count
752                .static_byte_len()
753                .expect("Fix: zero-count buffer must be representable"),
754            None
755        );
756
757        let dynamic_element = BufferDecl::read("tensor", 0, DataType::Tensor).with_count(4);
758        assert_eq!(
759            dynamic_element
760                .static_byte_len()
761                .expect("Fix: runtime-sized element must be representable"),
762            None
763        );
764    }
765}
766
767#[cfg(test)]
768mod shape_predicate_tests {
769    use super::*;
770
771    #[test]
772    fn at_least_holds_when_count_meets_minimum() {
773        let p = ShapePredicate::AtLeast(64);
774        assert!(p.holds(64));
775        assert!(p.holds(128));
776        assert!(!p.holds(32));
777    }
778
779    #[test]
780    fn at_most_holds_when_count_within_bound() {
781        let p = ShapePredicate::AtMost(64);
782        assert!(p.holds(0));
783        assert!(p.holds(64));
784        assert!(!p.holds(65));
785    }
786
787    #[test]
788    fn exactly_holds_only_for_match() {
789        let p = ShapePredicate::Exactly(7);
790        assert!(p.holds(7));
791        assert!(!p.holds(6));
792        assert!(!p.holds(8));
793    }
794
795    #[test]
796    fn multiple_of_holds_for_aligned_count() {
797        let p = ShapePredicate::MultipleOf(64);
798        assert!(p.holds(0));
799        assert!(p.holds(64));
800        assert!(p.holds(128));
801        assert!(!p.holds(63));
802        assert!(!p.holds(65));
803    }
804
805    #[test]
806    fn multiple_of_zero_never_holds() {
807        let p = ShapePredicate::MultipleOf(0);
808        assert!(!p.holds(0));
809        assert!(!p.holds(64));
810    }
811
812    #[test]
813    fn and_combines_two_predicates() {
814        // count >= 64 && count % 32 == 0
815        let p = ShapePredicate::And(
816            Box::new(ShapePredicate::AtLeast(64)),
817            Box::new(ShapePredicate::MultipleOf(32)),
818        );
819        assert!(p.holds(64));
820        assert!(p.holds(96));
821        assert!(!p.holds(32)); // satisfies MultipleOf but not AtLeast
822        assert!(!p.holds(80)); // satisfies AtLeast but not MultipleOf
823    }
824
825    #[test]
826    fn or_accepts_either_predicate() {
827        let p = ShapePredicate::Or(
828            Box::new(ShapePredicate::Exactly(8)),
829            Box::new(ShapePredicate::Exactly(16)),
830        );
831        assert!(p.holds(8));
832        assert!(p.holds(16));
833        assert!(!p.holds(12));
834    }
835
836    #[test]
837    fn not_inverts_predicate() {
838        let p = ShapePredicate::Not(Box::new(ShapePredicate::AtMost(64)));
839        assert!(!p.holds(64));
840        assert!(p.holds(65));
841    }
842
843    #[test]
844    fn mod_equals_requires_valid_modular_form() {
845        assert!(ShapePredicate::ModEquals {
846            modulus: 16,
847            remainder: 4,
848        }
849        .holds(20));
850        assert!(!ShapePredicate::ModEquals {
851            modulus: 16,
852            remainder: 4,
853        }
854        .holds(21));
855        assert!(!ShapePredicate::ModEquals {
856            modulus: 0,
857            remainder: 0,
858        }
859        .holds(0));
860        assert!(!ShapePredicate::ModEquals {
861            modulus: 4,
862            remainder: 4,
863        }
864        .holds(4));
865    }
866
867    #[test]
868    fn affine_range_uses_wide_arithmetic() {
869        let p = ShapePredicate::AffineRange {
870            scale: 4,
871            offset: -8,
872            min: 24,
873            max: 40,
874        };
875        assert!(!p.holds(7));
876        assert!(p.holds(8));
877        assert!(p.holds(12));
878        assert!(!p.holds(13));
879        assert!(!ShapePredicate::AffineRange {
880            scale: i64::MAX,
881            offset: i64::MAX,
882            min: i64::MIN,
883            max: i64::MAX,
884        }
885        .holds(u32::MAX));
886    }
887
888    #[test]
889    fn buffer_decl_default_shape_predicate_is_none() {
890        let buf = BufferDecl::read("a", 0, DataType::U32);
891        assert_eq!(buf.shape_predicate(), None);
892    }
893
894    #[test]
895    fn with_shape_predicate_round_trip() {
896        let buf = BufferDecl::read("a", 0, DataType::U32)
897            .with_shape_predicate(ShapePredicate::MultipleOf(32));
898        assert_eq!(buf.shape_predicate(), Some(&ShapePredicate::MultipleOf(32)));
899    }
900
901    #[test]
902    fn describe_renders_human_readable() {
903        assert_eq!(
904            ShapePredicate::And(
905                Box::new(ShapePredicate::AtLeast(64)),
906                Box::new(ShapePredicate::MultipleOf(32)),
907            )
908            .describe(),
909            "(count >= 64) && (count % 32 == 0)"
910        );
911    }
912}