Skip to main content

air_parser/ast/
types.rs

1use super::*;
2
3/// The types of values which can be represented in an AirScript program
4#[derive(Hash, Debug, Copy, Clone, PartialEq, Eq)]
5pub enum Type {
6    /// A field element
7    Felt,
8    /// A vector of N integers
9    Vector(usize),
10    /// A matrix of N rows and M columns
11    Matrix(usize, usize),
12}
13impl Type {
14    /// Returns true if this type is an aggregate
15    #[inline]
16    pub fn is_aggregate(&self) -> bool {
17        match self {
18            Self::Felt => false,
19            Self::Vector(_) | Self::Matrix(_, _) => true,
20        }
21    }
22
23    /// Returns true if this type is a scalar
24    #[inline]
25    pub fn is_scalar(&self) -> bool {
26        matches!(self, Self::Felt)
27    }
28
29    /// Returns true if this type is a valid iterable in a comprehension
30    #[inline]
31    pub fn is_iterable(&self) -> bool {
32        self.is_vector()
33    }
34
35    /// Returns true if this type is a vector
36    #[inline]
37    pub fn is_vector(&self) -> bool {
38        matches!(self, Self::Vector(_))
39    }
40
41    /// Return a new [Type] representing the type of the value produced by the given [AccessType]
42    pub fn access(&self, access_type: AccessType) -> Result<Self, InvalidAccessError> {
43        match *self {
44            ty if access_type == AccessType::Default => Ok(ty),
45            Self::Felt => Err(InvalidAccessError::IndexIntoScalar),
46            Self::Vector(len) => match access_type {
47                AccessType::Slice(range) => {
48                    let slice_range = range.to_slice_range();
49                    if slice_range.end > len {
50                        Err(InvalidAccessError::IndexOutOfBounds)
51                    } else {
52                        Ok(Self::Vector(slice_range.len()))
53                    }
54                }
55                AccessType::Index(idx) if idx >= len => Err(InvalidAccessError::IndexOutOfBounds),
56                AccessType::Index(_) => Ok(Self::Felt),
57                AccessType::Matrix(_, _) => Err(InvalidAccessError::IndexIntoScalar),
58                _ => unreachable!(),
59            },
60            Self::Matrix(rows, cols) => match access_type {
61                AccessType::Slice(range) => {
62                    let slice_range = range.to_slice_range();
63                    if slice_range.end > rows {
64                        Err(InvalidAccessError::IndexOutOfBounds)
65                    } else {
66                        Ok(Self::Matrix(slice_range.len(), cols))
67                    }
68                }
69                AccessType::Index(idx) if idx >= rows => Err(InvalidAccessError::IndexOutOfBounds),
70                AccessType::Index(_) => Ok(Self::Vector(cols)),
71                AccessType::Matrix(row, col) if row >= rows || col >= cols => {
72                    Err(InvalidAccessError::IndexOutOfBounds)
73                }
74                AccessType::Matrix(_, _) => Ok(Self::Felt),
75                _ => unreachable!(),
76            },
77        }
78    }
79}
80impl fmt::Display for Type {
81    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
82        match self {
83            Self::Felt => f.write_str("felt"),
84            Self::Vector(n) => write!(f, "felt[{n}]"),
85            Self::Matrix(rows, cols) => write!(f, "felt[{rows}, {cols}]"),
86        }
87    }
88}
89
90/// Represents the type signature of a function
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub enum FunctionType {
93    /// An evaluator function, which has no results, and has
94    /// a complex type signature due to the nature of trace bindings
95    Evaluator(Vec<TraceSegment>),
96    /// A standard function with one or more inputs, and a result
97    Function(Vec<Type>, Type),
98}
99impl FunctionType {
100    pub fn result(&self) -> Option<Type> {
101        match self {
102            Self::Evaluator(_) => None,
103            Self::Function(_, result) => Some(*result),
104        }
105    }
106}