air_parser/ast/
trace.rs

1use std::fmt;
2
3use miden_diagnostics::{SourceSpan, Spanned};
4
5use super::*;
6
7/// The id of a trace segment is its index in the trace_columns declaration
8pub type TraceSegmentId = usize;
9
10/// The index of a column in a particular trace segment
11pub type TraceColumnIndex = usize;
12
13#[derive(Clone, Spanned)]
14pub struct TraceSegment {
15    #[span]
16    pub span: SourceSpan,
17    /// The index of this segment in the trace_columns declaration
18    pub id: TraceSegmentId,
19    /// The name of this trace segment, e.g. `$main`
20    ///
21    /// NOTE: The name of a trace segment is always a special identifier (i.e. has the `$` prefix)
22    pub name: Identifier,
23    /// The number of columns in this trace segment
24    pub size: usize,
25    /// Bindings declared in this segment, without the segment-wide binding, e.g. `$main`
26    pub bindings: Vec<TraceBinding>,
27    /// A vector of `size` elements which tracks for every column whether a
28    /// constraint has been applied to that column, and on what boundaries.
29    pub boundary_constrained: Vec<Span<ColumnBoundaryFlags>>,
30}
31impl TraceSegment {
32    /// Constructs a new [TraceSegment] given a span, segment id, name, and a vector of (Identifier, size) pairs.
33    pub fn new(
34        span: SourceSpan,
35        id: TraceSegmentId,
36        name: Identifier,
37        raw_bindings: Vec<Span<(Identifier, usize)>>,
38    ) -> Self {
39        let mut bindings = Vec::with_capacity(raw_bindings.len());
40        let mut offset = 0;
41        for binding in raw_bindings.into_iter() {
42            let (name, size) = binding.item;
43            let ty = match size {
44                1 => Type::Felt,
45                n => Type::Vector(n),
46            };
47            bindings.push(TraceBinding::new(
48                binding.span(),
49                name,
50                id,
51                offset,
52                size,
53                ty,
54            ));
55            offset += size;
56        }
57
58        // The size of the segment is the sum of the sizes of all the bindings
59        let size = offset;
60        Self {
61            span,
62            id,
63            name,
64            size,
65            bindings,
66            boundary_constrained: vec![
67                Span::new(SourceSpan::UNKNOWN, ColumnBoundaryFlags::EMPTY);
68                size
69            ],
70        }
71    }
72
73    /// Returns true if `column` is constrained on `boundary`
74    pub fn is_boundary_constrained(&self, column: TraceColumnIndex, boundary: Boundary) -> bool {
75        self.boundary_constrained[column].is_constrained(boundary)
76    }
77
78    /// Marks `column` as constrained on `boundary`, and associates it with a span
79    /// that is responsible for the constraint.
80    ///
81    /// Returns `None` if the column was previously unconstrained on `boundary`,
82    /// otherwise it returns the span responsible for the previous constraint for
83    /// use in diagnostics
84    pub fn mark_constrained(
85        &mut self,
86        span: SourceSpan,
87        column: TraceColumnIndex,
88        boundary: Boundary,
89    ) -> Option<SourceSpan> {
90        let flags = &mut self.boundary_constrained[column];
91        if flags.is_constrained(boundary) {
92            Some(flags.span())
93        } else {
94            *flags = Span::new(span, flags.item | boundary);
95            None
96        }
97    }
98
99    #[inline]
100    pub fn is_empty(&self) -> bool {
101        self.size == 0
102    }
103}
104impl fmt::Debug for TraceSegment {
105    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
106        f.debug_struct("TraceSegment")
107            .field("id", &self.id)
108            .field("name", &self.name)
109            .field("size", &self.size)
110            .field("bindings", &self.bindings)
111            .field(
112                "boundary_constrained",
113                &FormatConstrainedFlags(&self.boundary_constrained),
114            )
115            .finish()
116    }
117}
118impl fmt::Display for TraceSegment {
119    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
120        if !self.name.is_generated() {
121            if let Some(name) = self.name.as_str().strip_prefix('$') {
122                write!(f, "{name}: ")?;
123            } else {
124                write!(f, "{}: ", self.name)?;
125            }
126        }
127        if self.bindings.is_empty() {
128            write!(f, "[{}]", self.size)
129        } else {
130            write!(f, "{}", DisplayList(self.bindings.as_slice()))
131        }
132    }
133}
134impl Eq for TraceSegment {}
135impl PartialEq for TraceSegment {
136    fn eq(&self, other: &Self) -> bool {
137        self.name == other.name && self.bindings == other.bindings && self.size == other.size
138    }
139}
140
141#[derive(Copy, Clone, PartialEq, Eq, Default)]
142pub struct ColumnBoundaryFlags(u8);
143impl ColumnBoundaryFlags {
144    /// An empty flag set that indicates the column is unconstrained
145    pub const EMPTY: Self = Self(0);
146    /// A flag set that indicates the column is constrained on the first boundary
147    pub const FIRST: Self = Self(0b001);
148    /// A flag set that indicates the column is constrained on the last boundary
149    pub const LAST: Self = Self(0b010);
150    /// A flag set that indicates the column is constrained on both boundaries
151    pub const BOTH: Self = Self(0b011);
152
153    /// Returns true if this column is constrained on `boundary`
154    pub fn is_constrained(&self, boundary: Boundary) -> bool {
155        *self & boundary
156    }
157}
158impl fmt::Debug for ColumnBoundaryFlags {
159    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
160        match self.0 {
161            0b000 => f.write_str("*"),
162            0b001 => f.write_str("F"),
163            0b010 => f.write_str("L"),
164            0b011 => f.write_str("B"),
165            _ => unreachable!(),
166        }
167    }
168}
169impl std::ops::BitOr<Boundary> for ColumnBoundaryFlags {
170    type Output = ColumnBoundaryFlags;
171
172    fn bitor(self, boundary: Boundary) -> Self {
173        Self(
174            self.0
175                | match boundary {
176                    Boundary::First => Self::FIRST.0,
177                    Boundary::Last => Self::LAST.0,
178                },
179        )
180    }
181}
182impl std::ops::BitAnd<Boundary> for ColumnBoundaryFlags {
183    type Output = bool;
184
185    fn bitand(self, boundary: Boundary) -> bool {
186        let bit = match boundary {
187            Boundary::First => Self::FIRST.0,
188            Boundary::Last => Self::LAST.0,
189        };
190        self.0 & bit == bit
191    }
192}
193
194/// Used to help format the boundary constraint flags
195struct FormatConstrainedFlags<'a>(&'a [Span<ColumnBoundaryFlags>]);
196impl fmt::Debug for FormatConstrainedFlags<'_> {
197    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
198        f.debug_list()
199            .entries(self.0.iter().map(|c| c.item))
200            .finish()
201    }
202}
203
204/// [TraceBinding] is used to represent one or more columns in the execution trace that are bound to
205/// a name. For single columns, the size is 1. For groups, the size is the number of columns in the
206/// group. The offset is the column index in the trace where the first column of the binding starts.
207#[derive(Copy, Clone, Spanned)]
208pub struct TraceBinding {
209    #[span]
210    pub span: SourceSpan,
211    /// The name of this binding, if applicable
212    pub name: Option<Identifier>,
213    /// The id of the segment to which this binding belongs
214    pub segment: TraceSegmentId,
215    /// The offset to the first column of the segment which is bound by this binding
216    pub offset: usize,
217    /// The number of columns which are bound
218    pub size: usize,
219    /// The effective type of this binding
220    pub ty: Type,
221}
222impl TraceBinding {
223    /// Creates a new trace binding.
224    pub const fn new(
225        span: SourceSpan,
226        name: Identifier,
227        segment: TraceSegmentId,
228        offset: usize,
229        size: usize,
230        ty: Type,
231    ) -> Self {
232        Self {
233            span,
234            name: Some(name),
235            segment,
236            offset,
237            size,
238            ty,
239        }
240    }
241
242    /// Returns a [Type] that describes what type of value this binding represents
243    #[inline]
244    pub fn ty(&self) -> Type {
245        self.ty
246    }
247
248    #[inline]
249    pub fn is_scalar(&self) -> bool {
250        self.ty.is_scalar()
251    }
252
253    /// Derive a new [TraceBinding] derived from the current one given an [AccessType]
254    pub fn access(&self, access_type: AccessType) -> Result<Self, InvalidAccessError> {
255        match access_type {
256            AccessType::Default => Ok(*self),
257            AccessType::Slice(_) if self.is_scalar() => Err(InvalidAccessError::SliceOfScalar),
258            AccessType::Slice(range) => {
259                let slice_range = range.to_slice_range();
260                if slice_range.end > self.size {
261                    Err(InvalidAccessError::IndexOutOfBounds)
262                } else {
263                    let offset = self.offset + slice_range.start;
264                    let size = slice_range.len();
265                    Ok(Self {
266                        offset,
267                        size,
268                        ty: Type::Vector(size),
269                        ..*self
270                    })
271                }
272            }
273            AccessType::Index(_) if self.is_scalar() => Err(InvalidAccessError::IndexIntoScalar),
274            AccessType::Index(idx) if idx >= self.size => Err(InvalidAccessError::IndexOutOfBounds),
275            AccessType::Index(idx) => {
276                let offset = self.offset + idx;
277                Ok(Self {
278                    offset,
279                    size: 1,
280                    ty: Type::Felt,
281                    ..*self
282                })
283            }
284            AccessType::Matrix(_, _) => Err(InvalidAccessError::IndexIntoScalar),
285        }
286    }
287}
288impl Eq for TraceBinding {}
289impl PartialEq for TraceBinding {
290    fn eq(&self, other: &Self) -> bool {
291        self.segment == other.segment
292            && self.name == other.name
293            && self.offset == other.offset
294            && self.size == other.size
295            && self.ty == other.ty
296    }
297}
298impl fmt::Debug for TraceBinding {
299    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
300        f.debug_struct("TraceBinding")
301            .field("name", &self.name)
302            .field("segment", &self.segment)
303            .field("offset", &self.offset)
304            .field("size", &self.size)
305            .field("ty", &self.ty)
306            .finish()
307    }
308}
309impl fmt::Display for TraceBinding {
310    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
311        if self.size == 1 {
312            write!(
313                f,
314                "{}",
315                self.name.as_ref().map(|n| n.as_str()).unwrap_or("?")
316            )
317        } else {
318            write!(
319                f,
320                "{}[{}]",
321                self.name.as_ref().map(|n| n.as_str()).unwrap_or("?"),
322                self.size
323            )
324        }
325    }
326}
air_parser/ast/trace.rs

air_parser/ast/
trace.rs
