air-parser 0.4.0

use std::sync::Arc;
use std::collections::HashSet;

use miden_diagnostics::{CodeMap, DiagnosticsHandler, Severity, SourceSpan, Span, Spanned};

use crate::{
    ast::*,
    lexer::Token,
    parser::ParseError,
    sema::SemanticAnalysisError,
    symbols,
    Symbol
};

grammar(diagnostics: &DiagnosticsHandler, codemap: &Arc<CodeMap>, next_var: &mut usize);

// MACROS
// ================================================================================================

// Comma-delimited with at least one element
Comma<T>: Vec<T> = {
    <v:(<T> ",")*> <e:T> => {
        let mut v = v;
        v.push(e);
        v
    }
};

// AST NODE
// ================================================================================================

pub Source: Source = {
    <Program> => Source::Program(<>),
    <Module*> =>? Library::new(diagnostics, codemap.clone(), <>)
        .map_err(|err| ParseError::from(err).into())
        .map(Source::Library),
}

pub Program: Program = {
    <root:Root> <modules:Module*> =>? {
        let root_name = root.name;
        let mut modules = modules;
        modules.push(root);
        let library = match Library::new(diagnostics, codemap.clone(), modules) {
            Ok(lib) => lib,
            Err(err) => return Err(ParseError::from(err).into()),
        };
        Program::load(diagnostics, root_name, library)
            .map_err(|err| ParseError::from(err).into())
    }
}

pub AnyModule: Module = {
    Root,
    Module,
}

Root: Module = {
    <l:@L> "def" <name:Identifier> <decls:Declaration*> <r:@R> =>? {
        Module::from_declarations(diagnostics, ModuleType::Root, span!(l, r), name, decls)
            .map_err(|err| ParseError::Analysis(err).into())
    }
}

Module: Module = {
    <l:@L> "mod" <name:Identifier> <decls:Declaration*> <r:@R> =>? {
        Module::from_declarations(diagnostics, ModuleType::Library, span!(l, r), name, decls)
            .map_err(|err| ParseError::Analysis(err).into())
    }
}

Declaration: Declaration = {
    Import => Declaration::Import(<>),
    Constant => Declaration::Constant(<>),
    PeriodicColumns => Declaration::PeriodicColumns(<>),
    EvaluatorFunction => Declaration::EvaluatorFunction(<>),
    Function => Declaration::Function(<>),
    Buses => Declaration::Buses(<>),
    <l:@L> <trace:Trace> <r:@R> => Declaration::Trace(Span::new(span!(l, r), trace)),
    <PublicInputs> => Declaration::PublicInputs(<>),
    <BoundaryConstraints> => Declaration::BoundaryConstraints(<>),
    <IntegrityConstraints> => Declaration::IntegrityConstraints(<>),
}

Import: Span<Import> = {
    <l:@L> "use" <module:Identifier> "::" "*" ";" <r:@R> => Span::new(span!(l, r), Import::All { module: Identifier::new(span!(l, r), module.name()) }),
    <l:@L> "use" <module:Identifier> "::" <item:Identifier> ";" <r:@R> => {
        let mut items: HashSet<Identifier> = HashSet::default();
        items.insert(item);
        Span::new(span!(l, r), Import::Partial { module, items })
    }
}

// TRACE COLUMNS
// ================================================================================================

Trace: Vec<TraceSegment> = {
    <l:@L> "trace_columns" "{" <main: MainTraceBindings?> "}" <r:@R> =>?
        match main {
            Some(main) => Ok(vec![main]),
            None => {
                diagnostics.diagnostic(Severity::Error)
                    .with_message("declaration of main trace columns is required")
                    .with_primary_label(span!(l, r), "missing 'main' declaration in this section")
                    .emit();
                Err(ParseError::Failed.into())
            }
        }
}

MainSegmentId: Identifier = {
    <l:@L> "main" <r:@R> => Identifier::new(span!(l, r), symbols::Main),
}

MainTraceBindings: TraceSegment = {
    <l:@L> <name:MainSegmentId> ":" <bindings: Vector<TraceBinding>> "," <r:@R> =>
        TraceSegment::new(span!(l, r), 0, name, bindings),
}

TraceBinding: Span<(Identifier, usize)> = {
    <name: Identifier> => Span::new(name.span(), (name, 1)),
    <l:@L> <name: Identifier> <size: Size> <r:@R> => Span::new(span!(l, r), (name, size as usize)),
}

// CONSTANTS
// ================================================================================================

Constant: Constant = {
    <l:@L> "const" <name: Identifier> "=" <value: ConstExpr> ";" <r:@R>
        => Constant::new(span!(l, r), name, value),
}

ConstExpr: ConstantExpr = {
    <Num_u64> => ConstantExpr::Scalar(<>),
    <Vector<Num_u64>> => ConstantExpr::Vector(<>),
    <Matrix<Num_u64>> => ConstantExpr::Matrix(<>),
}

// PUBLIC INPUTS
// ================================================================================================

// At least one public input is required.
PublicInputs: Span<Vec<PublicInput>> = {
    <l:@L> "public_inputs" "{" <inputs:PublicInput+> "}" <r:@R> => Span::new(span!(l, r), inputs)
}

PublicInput: PublicInput = {
    <l:@L> <name: Identifier> ":" <size: Size> "," <r:@R>
        => PublicInput::new_vector(span!(l, r), name, size),
    <l:@L> <name: Identifier> ":" <table: TableSize> "," <r:@R>
        => PublicInput::new_table(span!(l, r), name, table),
}

// PERIODIC COLUMNS
// ================================================================================================

// Periodic columns are not required, and there is no limit to the number that can be provided.
PeriodicColumns: Span<Vec<PeriodicColumn>> = {
    <l:@L> "periodic_columns" "{" <columns:PeriodicColumn*> "}" <r:@R>
        => Span::new(span!(l, r), columns)
}

PeriodicColumn: PeriodicColumn = {
    <l:@L> <name: Identifier> ":" <values: Vector<Num_u64>> "," <r:@R>
        => PeriodicColumn::new(span!(l, r), name, values),
}


// BUSES
// ================================================================================================

// Buses are not required, but if the section is specified at least one Bus declaration is required.
Buses: Span<Vec<Bus>> = {
    <l:@L> "buses" "{" <bus:Bus+> "}" <r:@R>
        => Span::new(span!(l, r), bus)
}

Bus: Bus = {
    <l:@L> <bus_type: BusType> <name: Identifier> "," <r:@R>
        => Bus::new(span!(l, r), name, bus_type),
}

BusType: BusType = {
    "multiset" => BusType::Multiset,
    "logup" => BusType::Logup,
}

// EVALUATOR FUNCTIONS
// ================================================================================================

EvaluatorFunction: EvaluatorFunction = {
    <l:@L> "ev" <name: FunctionIdentifier> "(" <params: EvaluatorBindings> ")" "{" <body: StatementBlock> "}" <r:@R>
        => EvaluatorFunction::new(span!(l, r), name, params, body)
}

EvaluatorBindings: Vec<TraceSegment> = {
    <l:@L> <trace: Comma<EvaluatorSegmentBindings>> <r:@R> =>? {
        let mut segments = Vec::with_capacity(trace.len());

        for (segment, (span, bindings)) in trace.into_iter().enumerate() {
            // We generate a name for these segments to distinguish them from direct references
            // to the actual main columns. This is useful during the inlining phase
            let segment_name = Identifier::new(span, Symbol::intern(format!("%{}", *next_var)));
            *next_var += 1;
            segments.push(TraceSegment::new(span, segment, segment_name, bindings));
        }

        // the last segment of trace columns cannot be empty.
        if let Some(segment) = segments.last() {
            if segment.is_empty() {
                diagnostics.diagnostic(Severity::Error)
                  .with_message("invalid evaluator function definition")
                  .with_primary_label(segment.span(), "the last trace segment cannot be empty")
                  .emit();
                return Err(ParseError::Failed.into());
            }
        }
        Ok(segments)
    }
}

EvaluatorSegmentBindings: (SourceSpan, Vec<Span<(Identifier, usize)>>) = {
    <l:@L> <elems: Vector<TraceBinding>> <r:@R> => (span!(l, r), elems),
    <l:@L> "[" "]" <r:@R> => (span!(l, r), vec![]),
}

// FUNCTIONS
// ================================================================================================

Function: Function = {
    <l:@L> "fn" <name: FunctionIdentifier> "(" <params: FunctionBindings> ")" "->" <ty: FunctionBindingType> "{" <body: FunctionBody> "}" <r:@R>
        => Function::new(span!(l, r), name, params, ty, body)
}

FunctionBindings: Vec<(Identifier, Type)> = {
    <l:@L> <params: Comma<FunctionBinding>> <r:@R> => params,
}

FunctionBinding: (Identifier, Type) = {
    <name: Identifier> ":" <ty: FunctionBindingType> => (name, ty),
}

FunctionBindingType: Type = {
    <l:@L> "felt" <r:@R> => Type::Felt,
    <l:@L> "felt" <size: Size> <r:@R> => Type::Vector(size as usize),
    <l:@L> "felt" "[" <row_size: Num_u64> "," <col_size: Num_u64> "]" <r:@R> => Type::Matrix(row_size as usize, col_size as usize),
}

FunctionBody: Vec<Statement> = {
    <l:@L> <stmts: StatementBlock> <r:@R> =>? {
        if stmts.len() > 1 {
            diagnostics.diagnostic(Severity::Error)
            .with_message("invalid function definition")
            .with_primary_label(span!(l, r), "function should have 0 or more let statements followed by a return statement.")
            .emit();
            return Err(ParseError::Failed.into());
        }
        Ok(stmts)
    },
}

// BOUNDARY CONSTRAINTS
// ================================================================================================

BoundaryConstraints: Span<Vec<Statement>> = {
    <l:@L> "boundary_constraints" "{" <body: StatementBlock> "}" <r:@R> => Span::new(span!(l, r), body),
}

Boundary: Boundary = {
    "first" => Boundary::First,
    "last" => Boundary::Last
}

// INTEGRITY CONSTRAINTS
// ================================================================================================

IntegrityConstraints: Span<Vec<Statement>> = {
    <l:@L> "integrity_constraints" "{" <body: StatementBlock> "}" <r:@R>
        => Span::new(span!(l, r), body)
}

// STATEMENTS
// ================================================================================================

StatementBlock: Vec<Statement> = {
    <Let> => vec![Statement::Let(<>)],
    <stmts:ConstraintStatements> <last:Let> => {
        let mut stmts = stmts;
        stmts.push(Statement::Let(last));
        stmts
    },
    <ConstraintStatements>,
    <ReturnStatement> => vec![Statement::Expr(<>)],
}

Let: Let = {
    <l:@L> "let" <name: Identifier> "=" <value: Expr> ";" <r:@R> <body: StatementBlock>
        => Let::new(span!(l, r), name, value, body)
}

ConstraintStatements: Vec<Statement> = {
    <stmts:ConstraintStatement+> => {
        stmts.into_iter().flatten().collect::<Vec<_>>()
    }
}

ConstraintStatement: Vec<Statement> = {
    "enf" "match" "{" <MatchArm+> "}" ";" => <>,
    "enf" <ConstraintExpr> ";" => vec![<>],
    <BusConstraintExpr> ";" => vec![<>],
}

ReturnStatement: Expr = {
    <l:@L> "return" <expr: Expr> ";" <r:@R> => expr,
}

MatchArm: Statement = {
    <l:@L> "case" <selector:ScalarExpr> ":" <constraint:ScalarConstraintExpr> "," <r:@R> => {
        let generated_name = format!("%{}", *next_var);
        *next_var += 1;
        let generated_binding = Identifier::new(SourceSpan::UNKNOWN, Symbol::intern(generated_name));
        let context = vec![(generated_binding, Expr::Range(RangeExpr::from(0..1)))];
        Statement::EnforceAll(ListComprehension::new(span!(l, r), constraint, context, Some(selector)))
    }
}

// This grammar rules handles two types of constraints: simple and comprehension constraints.
// These constraints come in four syntax varieties:
//
// 1. `enf x' = x + y for x in xs when x`
// 2. `enf x' = x + y for x in xs`
// 3. `enf x' = x + y when x`
// 4. `enf x' = x + y`
//
// 1 and 2 are called "comprehension" constraints, i.e. it expresses a set of
// constraints for every element in a set of iterables, with an optional selector that causes the
// constraint to be vacuously true when the condition is false (or put another way, the constraint
// is only meaningful/enforced when the condition is true).
//
// 4 is called a "simple" constraint, i.e. it expresses a single constraint that is always enforced.
//
// 3 appears to be "simple" as well, except that it has a selector applied, which makes it somewhat
// of an odd duck. To keep the AST simple, we perform a simple transformation to constraints of this
// specific variation, such that they appear to be "comprehension" constraints. For a more visual
// example of what this transformation does, assume the following constraint of form 3 was written:
//
//    enf x' = x + y when x & y
//
// This will be transformed into form 1, like so:
//
//    enf x' = x + y for %0 in 0..1 when x & y
//
// Where `%0` is a generated variable (i.e. inexpressible in the language itself to avoid name conflicts).
//
ConstraintExpr: Statement = {
    <l:@L> <expr: ScalarConstraintExpr> <comprehension: ConstraintComprehension<ScalarExpr>?> <selector: WithSelector?> <r:@R> => {
        // If we parsed a comprehension, we've parsed either form 1 or 2
        if let Some(context) = comprehension {
            Statement::EnforceAll(ListComprehension::new(span!(l, r), expr, context, selector))
        } else {
            // If we didn't parse this as a comprehension, but a selector is present, the constraint is in form 3,
            // so transform it into form 1. Otherwise, if no selector is present, this is form 4, i.e. simple.
            if selector.is_some() {
                let generated_name = format!("%{}", *next_var);
                *next_var += 1;
                let generated_binding = Identifier::new(SourceSpan::UNKNOWN, Symbol::intern(generated_name));
                let context = vec![(generated_binding, Expr::Range(RangeExpr::from(0..1)))];
                Statement::EnforceAll(ListComprehension::new(span!(l, r), expr, context, selector))
            } else {
                Statement::Enforce(expr)
            }
        }
    }
}

// 1. `p.insert(a, b) when s`
// 2. `q.insert(a, b) with m`
BusConstraintExpr: Statement = {
    <l:@L> <expr: ScalarBusConstraintExpr> <selector: WithSelector> <r:@R> => {
        let generated_name = format!("%{}", *next_var);
        *next_var += 1;
        let generated_binding = Identifier::new(SourceSpan::UNKNOWN, Symbol::intern(generated_name));
        let context = vec![(generated_binding, Expr::Range(RangeExpr::from(0..1)))];
        Statement::BusEnforce(ListComprehension::new(span!(l, r), expr, context, Some(selector)))
    },
    <l:@L> <expr: ScalarBusConstraintExpr> <multiplicity: WithMultiplicity> <r:@R> => {
        let generated_name = format!("%{}", *next_var);
        *next_var += 1;
        let generated_binding = Identifier::new(SourceSpan::UNKNOWN, Symbol::intern(generated_name));
        let context = vec![(generated_binding, Expr::Range(RangeExpr::from(0..1)))];
        Statement::BusEnforce(ListComprehension::new(span!(l, r), expr, context, Some(multiplicity)))
    }
}

// 1. `p.first = null`
// 2. `q.first = ???` To define
BoundaryBusConstraintExpr: Statement = {
    <l:@L> <expr: ScalarBusConstraintExpr> <r:@R> => {
        let generated_name = format!("%{}", *next_var);
        *next_var += 1;
        let generated_binding = Identifier::new(SourceSpan::UNKNOWN, Symbol::intern(generated_name));
        let context = vec![(generated_binding, Expr::Range(RangeExpr::from(0..1)))];
        Statement::BusEnforce(ListComprehension::new(span!(l, r), expr, context, None))
    }
}

ScalarConstraintExpr: ScalarExpr = {
    FunctionCall,
    <l:@L> <lhs: ScalarExpr> "=" <rhs: ScalarExpr> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Eq, lhs, rhs)),
}

ScalarBusConstraintExpr: ScalarExpr = {
    <l:@L> <bus: Identifier> "." <bus_operator: BusOperator> "(" <args: Comma<Expr>> ")" <r:@R> => {
        ScalarExpr::BusOperation(BusOperation::new(span!(l, r), bus, bus_operator, args))
    }
}

BusOperator: BusOperator = {
    "insert" => BusOperator::Insert,
    "remove" => BusOperator::Remove
}

WithSelector: ScalarExpr = {
    "when" <ScalarExpr>,
}
WithMultiplicity: ScalarExpr = {
    "with" <ScalarExpr>,
}

Expr: Expr = {
    <ScalarExpr> =>? {
        match Expr::try_from(<>) {
            Ok(expr) => Ok(expr),
            Err(err) => Err(ParseError::from(SemanticAnalysisError::InvalidExpr(err)).into())
        }
    },
    <l:@L> <value:Vector<ScalarExpr>> <r:@R> =>? {
        let mut value = value;
        let mut elems = vec![];
        for v in value.drain(..) {
            match Expr::try_from(v) {
                Ok(expr) => {
                    elems.push(expr);
                }
                Err(err) => return Err(ParseError::from(SemanticAnalysisError::InvalidExpr(err)).into()),
            }
        }
        Ok(Expr::Vector(Span::new(span!(l, r), elems)))
    },
    <l:@L> <value:Matrix<ScalarExpr>> <r:@R> => Expr::Matrix(Span::new(span!(l, r), value)),
    "[" <ListComprehension<ScalarExpr>> "]" => Expr::ListComprehension(<>),
    <l:@L> "[" "]" <r:@R> => Expr::Vector(Span::new(span!(l, r), vec![])),
}

// --- SCALAR EXPRESSIONS WITH PRECEDENCE (LOWEST TO HIGHEST) ----------------------

ScalarExpr: ScalarExpr = {
    ScalarExprBase,
}

ScalarExprBase: ScalarExpr = {
    #[precedence(level="0")]
    SymbolAccess,
    <Int> => ScalarExpr::Const(<>),
    <l:@L> "null" <r:@R> => ScalarExpr::Null(Span::new(span!(l, r), ())),
    <l:@L> "unconstrained" <r:@R> => ScalarExpr::Unconstrained(Span::new(span!(l, r), ())),
    "(" <ScalarExpr> ")",

    #[precedence(level="1")]
    FunctionCall,

    #[precedence(level="2")]
    <l:@L> "!" <expr:ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Sub, ScalarExpr::Const(Span::new(span!(l, r), 1)), expr)),

    #[precedence(level="3")] #[assoc(side="left")]
    <l:@L> <lhs: ScalarExprBase> "^" <rhs: ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Exp, lhs, rhs)),

    #[precedence(level="4")] #[assoc(side="left")]
    <l:@L> <lhs: ScalarExprBase> "*" <rhs: ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Mul, lhs, rhs)),

    #[precedence(level="5")] #[assoc(side="left")]
    <l:@L> <lhs: ScalarExprBase> "+" <rhs: ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Add, lhs, rhs)),
    <l:@L> <lhs: ScalarExprBase> "-" <rhs: ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Sub, lhs, rhs)),

    #[precedence(level="6")] #[assoc(side="left")]
    <l:@L> <lhs:ScalarExprBase> "&" <rhs:ScalarExprBase> <r:@R>
        => ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Mul, lhs, rhs)),
    <l:@L> <lhs:ScalarExprBase> "|" <rhs:ScalarExprBase> <r:@R> => {
        let lhs2 = ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Add, lhs.clone(), rhs.clone()));
        let rhs2 = ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Mul, lhs, rhs));
        return ScalarExpr::Binary(BinaryExpr::new(span!(l, r), BinaryOp::Sub, lhs2, rhs2));
    }
}

FunctionCall: ScalarExpr = {
    <l:@L> <callee: FunctionIdentifier> "(" <args: Comma<Expr>> ")" <r:@R> => {
        ScalarExpr::Call(Call::new(span!(l, r), callee, args))
    }
}

SymbolAccess: ScalarExpr = {
    <l:@L> <symbol_access: SymbolAccessBaseSpanned> "." <boundary: Boundary> <r:@R>
        => ScalarExpr::BoundedSymbolAccess(BoundedSymbolAccess::new(span!(l, r), SymbolAccess::new(symbol_access.span(), symbol_access.item.0, symbol_access.item.1, 0), boundary)),
    <l:@L> <symbol_access: SymbolAccessBaseSpanned> "'" <r:@R>
        => ScalarExpr::SymbolAccess(SymbolAccess::new(span!(l, r), symbol_access.item.0, symbol_access.item.1, 1)),
    <symbol_access: SymbolAccessBaseSpanned>
        => ScalarExpr::SymbolAccess(SymbolAccess::new(symbol_access.span(), symbol_access.item.0, symbol_access.item.1, 0)),
}

SymbolAccessBaseSpanned: Span<(Identifier, AccessType)> = {
    <l:@L> <base:SymbolAccessBase> <r:@R> => Span::new(span!(l, r), base),
}

SymbolAccessBase: (Identifier, AccessType) = {
    <ident: Identifier> => (ident, AccessType::Default),
    <ident: Identifier> "[" <range: Range> "]" => (ident, AccessType::Slice(range)),
    <ident: Identifier> <idx: Index> =>  (ident, AccessType::Index(idx)),
    <ident: Identifier> <row: Index> <col: Index> => (ident, AccessType::Matrix(row, col)),
    // accessing an identifier used in a section declaration, like a named trace segment, e.g. $main
    <ident: DeclIdentifier> => (ident, AccessType::Default),
    <ident: DeclIdentifier> <idx: Index> => (ident, AccessType::Index(idx))

}

// COMPREHENSIONS
// ================================================================================================

ConstraintComprehension<T>: ComprehensionContext = {
    <l:@L> "for" <members: Members> "in" <iterables: Iterables> <r:@R> =>?
        if members.len() != iterables.len() {
            diagnostics.diagnostic(Severity::Error)
                .with_message("bindings and iterables lengths are mismatched")
                .with_primary_label(span!(l, r), "in this comprehension")
                .emit();
            Err(ParseError::Analysis(SemanticAnalysisError::Invalid).into())
        } else {
            Ok(members.into_iter().zip(iterables).collect::<Vec<_>>())
        }
}

ListComprehension<T>: ListComprehension = {
    <l:@L> <expr: T> "for" <members: Members> "in" <iterables: Iterables> <r:@R> =>?
        if members.len() != iterables.len() {
            diagnostics.diagnostic(Severity::Error)
                .with_message("bindings and iterables lengths are mismatched")
                .with_primary_label(span!(l, r), "in this comprehension")
                .emit();
            Err(ParseError::Analysis(SemanticAnalysisError::Invalid).into())
        } else {
            Ok(ListComprehension::new(span!(l, r), expr, members.into_iter().zip(iterables).collect::<Vec<_>>(), None))
        }
}

Members: Vec<Identifier> = {
    <member: Identifier> => vec![member],
    "(" <members: Comma<Identifier>> ")" => members
}

Iterables: Vec<Expr> = {
    <iterable: Iterable> => vec![iterable],
    "(" <iterables: Comma<Iterable>> ")" => iterables
}

Iterable: Expr = {
    <ident: Identifier> => Expr::SymbolAccess(SymbolAccess::new(ident.span(), ident, AccessType::Default, 0)),
    <range: Range> => Expr::Range(range),
    <l:@L> <ident: Identifier> "[" <range: Range> "]" <r:@R> => Expr::SymbolAccess(SymbolAccess::new(span!(l, r), ident, AccessType::Slice(range), 0)),
    <l:@L> <function_call: FunctionCall> <r:@R> => if let ScalarExpr::Call(call) = function_call {
        Expr::Call(call)
    } else {
        unreachable!()
    }
}

Range: RangeExpr = {
    <l:@L> <start: RangeBound> ".." <end: RangeBound> <r:@R> => RangeExpr { span: span!(l, r), start, end }
}

RangeBound: RangeBound = {
    <l:@L> <value:Num_u64> <r:@R> => RangeBound::Const(Span::new(span!(l, r), value as usize)),
    <l:@L> <name:Identifier> <r:@R> => RangeBound::SymbolAccess(ConstSymbolAccess::new(span!(l, r), name)),
}

// ATOMS
// ================================================================================================

Vector<T>: Vec<T> = {
    "[" <Comma<T>> "]" => <>,
}

Matrix<T>: Vec<Vec<T>> = {
    Vector<Vector<T>>,
}

Tuple<T>: Vec<T> = {
    "(" <v1:T> "," <v2:T> <v:("," <T>)*> ")" => {
        let mut v = v;
        v.insert(0, v2);
        v.insert(0, v1);
        v
    }
};

Size: u64 = {
    "[" <Num_u64> "]" => <>
}

Index: usize = {
    "[" <idx: Num_u64> "]" => idx as usize
}

TableSize: u64 = {
    "[" "[" <cols: Num_u64> "]" "]" => <>
}

DeclIdentifier: Identifier = {
    <l:@L> <name:decl_ident_ref> <r:@R> => Identifier::new(span!(l, r), name)
}

Identifier: Identifier = {
    <l:@L> <name:identifier> <r:@R> => Identifier::new(span!(l, r), name)
}

FunctionIdentifier: Identifier = {
    <l:@L> <name:function_identifier> <r:@R> => Identifier::new(span!(l, r), name)
}

Int: Span<u64> = {
   <l:@L> <i:Num_u64> <r:@R> => Span::new(span!(l, r), i),
}

Num_u64: u64 = {
    int,
}

// LEXER
// ================================================================================================

extern {
    type Error = ParseError;
    type Location = miden_diagnostics::SourceIndex;

    enum Token {
        identifier => Token::Ident(<Symbol>),
        decl_ident_ref => Token::DeclIdentRef(<Symbol>),
        function_identifier => Token::FunctionIdent(<Symbol>),
        int => Token::Num(<u64>),
        "def" => Token::Def,
        "mod" => Token::Mod,
        "use" => Token::Use,
        "let" => Token::Let,
        "for" => Token::For,
        "in" => Token::In,
        "const" => Token::Const,
        "trace_columns" => Token::TraceColumns,
        "main" => Token::Main,
        "public_inputs" => Token::PublicInputs,
        "periodic_columns" => Token::PeriodicColumns,
        "buses" => Token::Buses,
        "multiset" => Token::Multiset,
        "logup" => Token::Logup,
        "null" => Token::Null,
        "unconstrained" => Token::Unconstrained,
        "insert" => Token::Insert,
        "remove" => Token::Remove,
        "boundary_constraints" => Token::BoundaryConstraints,
        "first" => Token::First,
        "last" => Token::Last,
        "integrity_constraints" => Token::IntegrityConstraints,
        "ev" => Token::Ev,
        "fn" => Token::Fn,
        "enf" => Token::Enf,
        "return" => Token::Return,
        "match" => Token::Match,
        "case" => Token::Case,
        "when" => Token::When,
        "with" => Token::With,
        "felt" => Token::Felt,
        "'" => Token::Quote,
        "=" => Token::Equal,
        "+" => Token::Plus,
        "-" => Token::Minus,
        "*" => Token::Star,
        "^" => Token::Caret,
        "&" => Token::Ampersand,
        "|" => Token::Bar,
        "!" => Token::Bang,
        ":" => Token::Colon,
        "::" => Token::ColonColon,
        "," => Token::Comma,
        "[" => Token::LBracket,
        "]" => Token::RBracket,
        "(" => Token::LParen,
        ")" => Token::RParen,
        "{" => Token::LBrace,
        "}" => Token::RBrace,
        "." => Token::Dot,
        ".." => Token::DotDot,
        "->" => Token::Arrow,
        ";" => Token::SemiColon,
    }
}