use std::str::FromStr;
use std::collections::BTreeSet;
use powdr_ast::parsed::{*, asm::*, types::*};
use powdr_number::{BigInt, BigUint};
use crate::{ParserContext, unescape_string};
use powdr_parser_util::Error;
use std::sync::Arc;

use lalrpop_util::ParseError;

grammar(ctx: &ParserContext);

extern {
    type Error = Error;
}

match {
    r"\s*" => { },
    r"//[^\n\r]*[\n\r]*" => { }, // Skip `// comments`
    r"/\*[^*]*\*+(?:[^/*][^*]*\*+)*/" => { },  // Skip `/* comments */`
    _,
}

pub PILFile: PILFile = {
    <PilStatement*> => PILFile(<>)
};

pub ASMModule: ASMModule = {
    (<ModuleStatement>)* => ASMModule { statements: <> }
};

ModuleStatement: ModuleStatement = {
    <MachineDefinition> => ModuleStatement::SymbolDefinition(<>),
    <PilStatementAtModuleLevel> => ModuleStatement::PilStatement(<>),
    <Import> => ModuleStatement::SymbolDefinition(<>),
    <ModuleDefinition> => ModuleStatement::SymbolDefinition(<>),
}

ModuleDefinition: SymbolDefinition = {
    "mod" <name:Identifier> ";" => SymbolDefinition { name: name.clone(), value: Module::External(name).into() }, 
    "mod" <name:Identifier> "{" <module:ASMModule> "}" => SymbolDefinition { name, value: Module::Local(module).into() }
}

Import: SymbolDefinition = {
    "use" <path:SymbolPath> <name:( "as" <Identifier> )?> ";" =>
        SymbolDefinition {
            name: name.unwrap_or(path.name().clone().try_into().unwrap()),
            value: Import {path}.into()
        }
}

pub SymbolPath: SymbolPath = {
    <abs:"::"?> <parts:( <Part> "::" )*> <end:Part>  => {
        SymbolPath::from_parts([
            abs.map(|_| vec![Part::Named(String::new())]).unwrap_or_default(),
            parts,
            vec![end],
        ].concat())
    },
}

Part: Part = {
    "super" => Part::Super,
    <name:Identifier> => Part::Named(name),
}

// Same as SymbolPath plus we allow "::<...>" at the end.
GenericSymbolPath: (SymbolPath, Option<Vec<Type<Expression>>>) = {
    // If we "inline" SymbolPath here, we get an ambiguity error.
    <abs:"::"?> <parts:( <Part> "::" )*> <end:Part> <types:("::" "<" <TypeTermList<Expression>> ">")?>  => (
        SymbolPath::from_parts([
            abs.map(|_| vec![Part::Named(String::new())]).unwrap_or_default(),
            parts,
            vec![end],
        ].concat()),
        types
    ),
}

/// Same as SymbolPath except that we do not allow 'int' and 'fe' to be parsed as identifiers.
TypeSymbolPath: SymbolPath = {
    <abs:"::"?> <parts:( <TypeSymbolPathPart> "::" )*> <end:TypeSymbolPathPart>  => {
        SymbolPath::from_parts([
            abs.map(|_| vec![Part::Named(String::new())]).unwrap_or_default(),
            parts,
            vec![end],
        ].concat())
    },
}

TypeSymbolPathPart: Part = {
    "super" => Part::Super,
    <name:NonSpecialIdentifier> => Part::Named(name),
}

// ---------------------------- PIL part -----------------------------

pub PilStatementAtModuleLevel = {
    LetStatement,
    <start:@L> <decl:EnumDeclaration> <end:@R> => PilStatement::EnumDeclaration(ctx.source_ref(start, end), decl),
    <start:@L> <decl:StructDeclaration> <end:@R> => PilStatement::StructDeclaration(ctx.source_ref(start, end), decl),
    <start:@L> <impl_:TraitImplementation> <end:@R> => PilStatement::TraitImplementation(ctx.source_ref(start, end), impl_),
    <start:@L> <decl:TraitDeclaration> <end:@R> => PilStatement::TraitDeclaration(ctx.source_ref(start, end), decl),
}

pub PilStatement = {
    Include,
    Namespace,
    LetStatement,
    PolynomialDefinition,
    PublicDeclaration,
    PolynomialConstantDefinition,
    PolynomialCommitDeclaration,
    <start:@L> <decl:EnumDeclaration> <end:@R> => PilStatement::EnumDeclaration(ctx.source_ref(start, end), decl),
    <start:@L> <decl:StructDeclaration> <end:@R> => PilStatement::StructDeclaration(ctx.source_ref(start, end), decl),
    <start:@L> <impl_:TraitImplementation> <end:@R> => PilStatement::TraitImplementation(ctx.source_ref(start, end), impl_),
    <start:@L> <decl:TraitDeclaration> <end:@R> => PilStatement::TraitDeclaration(ctx.source_ref(start, end), decl),
    ExpressionStatement,
};

Include: PilStatement = {
    <start:@L> "include" <file:StringLiteral> <end:@R> ";" => PilStatement::Include(ctx.source_ref(start, end), file)
};

NamespaceDegree: NamespaceDegree = {
    <min:Expression> ".." <max:Expression> => NamespaceDegree {min, max},
    <degree:Expression> => NamespaceDegree { min: degree.clone(), max: degree }
}

Namespace: PilStatement = {
    <start:@L> "namespace" <name:(SymbolPath)?> <pol_degree:("(" <NamespaceDegree> ")")?> <end:@R> ";"
        => PilStatement::Namespace(ctx.source_ref(start, end), name.unwrap_or_default(), pol_degree)
}

LetStatement: PilStatement = {
    <start:@L> "let" <name:GenericTypedName> <expr:( "=" <Expression> )?> <end:@R> ";" =>
        PilStatement::LetStatement(ctx.source_ref(start, end), name.0, name.1, expr)
}

PolynomialDefinition: PilStatement = {
    <start:@L> PolCol <name:PolynomialName> "=" <expr:Expression> <end:@R> ";" => PilStatement::PolynomialDefinition(ctx.source_ref(start, end), name, expr)
}

PublicDeclaration: PilStatement = {
    <start:@L> "public" <id:Identifier> "="
        <poly:NamespacedPolynomialReference>
        <expr1:("[" <Expression> "]")?>
        "(" <expr2:Expression> ")" <end:@R> ";" => PilStatement::PublicDeclaration(ctx.source_ref(start, end), id, poly, expr1, expr2)
}

PolynomialConstantDefinition: PilStatement = {
    <start:@L> PolCol ConstantFixed <id:Identifier> <def:FunctionDefinition> <end:@R> ";"
        => PilStatement::PolynomialConstantDefinition(ctx.source_ref(start, end), id, def)
}

FunctionDefinition: FunctionDefinition = {
    <start:@L> "(" <params:ParameterList> ")" <body:BlockExpression> <end:@R>
        => FunctionDefinition::Expression(Expression::LambdaExpression(ctx.source_ref(start, end), LambdaExpression{
            kind: FunctionKind::Pure,
            params,
            body,
            param_types: vec![],
        })),
    <start:@L> "=" <array:ArrayLiteralExpression> <end:@R> => FunctionDefinition::Array(array),
}

ParameterList: Vec<Pattern> = {
    PatternList
}

ArrayLiteralExpression: ArrayExpression = {
    <ArrayLiteralExpression> "+" <ArrayLiteralTerm> => ArrayExpression::concat(<>),
    ArrayLiteralTerm,
}

ArrayLiteralTerm: ArrayExpression = {
    "[" <ExpressionList> "]" => ArrayExpression::value(<>),
    "[" <ExpressionList> "]" "*" => ArrayExpression::repeated_value(<>),
}

PolynomialCommitDeclaration: PilStatement = {
    <start:@L> PolCol CommitWitness <stage:Stage?> <list:PolynomialNameList> <end:@R> ";" => PilStatement::PolynomialCommitDeclaration(ctx.source_ref(start, end), stage, list, None),
    <start:@L> PolCol CommitWitness <stage:Stage?> <name:PolynomialName> "(" <params:ParameterList> ")" "query" <body:BoxedExpression> <end:@R> ";"
     => PilStatement::PolynomialCommitDeclaration(
        ctx.source_ref(start, end),
        stage,
        vec![name],
        Some(FunctionDefinition::Expression(Expression::LambdaExpression(ctx.source_ref(start, end), LambdaExpression{
            kind: FunctionKind::Query,
            params,
            body,
            param_types: vec![],
        })))
    )
}

Stage: u32 = {
    "stage" "(" <Number> ")" => u32::try_from(<>).unwrap()
}

PolynomialNameList: Vec<PolynomialName> = {
    <mut list:( <PolynomialName> "," )*> <end:PolynomialName> ","? => { list.push(end); list }
}

PolynomialName: PolynomialName = {
    <name:Identifier> <array_size:( "[" <Expression> "]" )?> => PolynomialName{<>}
}

ExpressionStatement: PilStatement = {
    <ExpressionStatementWithoutSemicolon> ";"
}

ExpressionStatementWithoutSemicolon: PilStatement = {
    <start:@L> <expr:Expression> <end:@R> => PilStatement::Expression(ctx.source_ref(start, end), expr)
}

PolCol = {
    "pol", "col"
} 

CommitWitness = {
    "commit", "witness"
} 

ConstantFixed = {
    "constant", "fixed"
}

GenericTypedName: (String, Option<TypeScheme<Expression>>) = {
    <name:Identifier> => (name, None),
    <vars:("<" <TypeVarBounds> ">")?> <name:Identifier> <ty:(":" <Type<Expression>>)> =>
        (name, Some(TypeScheme{ vars: vars.unwrap_or_default(), ty }))
}


// ---------------------------- ASM part -----------------------------

MachineDefinition: SymbolDefinition = {
    "machine" <name:Identifier> <params:MachineParams> <properties:("with" <MachineProperties>)?> "{" <statements:(MachineStatement)*> "}" => SymbolDefinition { name, value: Machine { params, properties: properties.unwrap_or_default(), statements}.into() },
}

MachineProperties: MachineProperties = {
    <start:@L> <mut list:( <MachineProperty> "," )*> <last:MachineProperty> ","? <end:@R> =>? {
        list.push(last);
        Ok(MachineProperties::try_from_prop_list(ctx.source_ref(start, end), list)?)
    }
}

MachineProperty: (String, Expression) = {
    <name:Identifier> ":" <value:RestrictedExpression> => (name, value),
}

MachineParams: MachineParams = {
    => MachineParams::default(),
    "(" ")" => MachineParams::default(),
    <start:@L> "(" <mut list:( <Param> "," )*> <last:Param> ","? ")" <end:@R> =>? {
        list.push(last);
        Ok(MachineParams::try_from_params(ctx.source_ref(start, end), list)?)
    }
}

MachineStatement: MachineStatement = {
    Submachine,
    RegisterDeclaration,
    InstructionDeclaration,
    LinkDeclarationStatement,
    PilStatementInMachine,
    FunctionDeclaration,
    OperationDeclaration,
}

PilStatementInMachine: MachineStatement = {
    <start:@L> <stmt:PilStatement> <end:@R> => MachineStatement::Pil(ctx.source_ref(start, end), stmt)
}

Submachine: MachineStatement = {
    <start:@L> <path:SymbolPath> <id:Identifier> <args:("(" <ExpressionList> ")")?> <end:@R> ";" => MachineStatement::Submachine(ctx.source_ref(start, end), path, id, args.unwrap_or_default()),
}

pub RegisterDeclaration: MachineStatement = {
    // TODO default update
    <start:@L> "reg" <id:Identifier> <flag:( "[" <RegisterFlag> "]" )?> <end:@R> ";" => MachineStatement::RegisterDeclaration(ctx.source_ref(start, end), id, flag)

}

RegisterFlag: RegisterFlag = {
    "@pc" => RegisterFlag::IsPC,
    "<=" => RegisterFlag::IsAssignment,
    "@r" => RegisterFlag::IsReadOnly,
}

pub InstructionDeclaration: MachineStatement = {
    <start:@L> "instr" <id:Identifier> <instr:Instruction> <end:@R> => MachineStatement::InstructionDeclaration(ctx.source_ref(start, end), id, instr)
}

pub Instruction: Instruction = {
    <params:Params> <links:(<LinkDeclaration>)*> <body:InstructionBody> => Instruction { params, body, links },
    <params:Params> <links:(<LinkDeclaration>)+> ";" => Instruction { params, body: InstructionBody(vec![]), links },
}

pub LinkDeclarationStatement: MachineStatement = {
    <start:@L> <link:LinkDeclaration> <end:@R> ";" => MachineStatement::LinkDeclaration(ctx.source_ref(start, end), link),
}

pub LinkDeclaration: LinkDeclaration = {
    "link" <flag:LinkFlag> "=>" <link:RestrictedExpression> =>? Ok(LinkDeclaration { flag, link: link.try_into()?, is_permutation: false, }),
    "link" <flag:LinkFlag> "~>" <link:RestrictedExpression> =>? Ok(LinkDeclaration { flag, link: link.try_into()?, is_permutation: true, }),
}

pub LinkFlag: Expression = {
    "if" <Expression>,
    => 1.into(),
}

pub InstructionBody: InstructionBody = {
    "{}" => InstructionBody(vec![]),
    "{" <InstructionBodyElements> "}" => InstructionBody(<>),
}

InstructionBodyElements: Vec<PilStatement> = {
    <mut list:( <InstructionBodyElement> "," )*> <end:InstructionBodyElement> ","?  => { list.push(end); list },
    => vec![]
}

InstructionBodyElement: PilStatement = {
    ExpressionStatementWithoutSemicolon,
}

Params: Params<Param> = {
    <_input: ParamList> "->" <output: ParamList> => Params::new(_input, output),
    // we can ommit the arrow if there are no outputs
    <_input: ParamList> => Params::new(_input, vec![])
}

ParamList: Vec<Param> = {
    => vec![],
    <mut list:( <Param> "," )*> <end:Param> ","? => { list.push(end); list }
}

Param: Param = {
    <start:@L> <name: Identifier> <index:("[" <Number> "]")?> <ty:(":" <SymbolPath>)?> <end:@R> => Param{source: ctx.source_ref(start, end), name, index, ty},
}

FunctionDeclaration: MachineStatement = {
    <start:@L> "function" <id:Identifier> <params:Params> "{" <stmt:(<FunctionStatement>)*> "}" <end:@R> => MachineStatement::FunctionDeclaration(ctx.source_ref(start, end), id, params, stmt)
}

OperationDeclaration: MachineStatement = {
    <start:@L> "operation" <id:Identifier> <op:OperationId> <params:Params> <end:@R> ";" => MachineStatement::OperationDeclaration(ctx.source_ref(start, end), id, op, params)
}

OperationId: OperationId = {
    "<" <id:Number> ">" => OperationId { id: Some(id.into()) },
    => OperationId { id: None }
}

pub FunctionStatement: FunctionStatement = {
    AssignmentStatement,
    LabelStatement,
    DebugDirectiveStatement,
    ReturnStatement,
    InstructionStatement,
}

AssignmentStatement: FunctionStatement = {
    <start:@L> <ids:IdentifierList> <op:AssignOperator> <expr:BoxedExpression> <end:@R> ";" => FunctionStatement::Assignment(ctx.source_ref(start, end), ids, op, expr)
}

IdentifierList: Vec<String> = {
    <mut list:( <Identifier> "," )*> <end:Identifier> ","? => { list.push(end); list },
    => vec![]
}

AssignOperator: Option<Vec<AssignmentRegister>> = {
    "<==" => None,
    "<=" <AssignmentRegisterList> "=" => Some(<>)
}

AssignmentRegisterList: Vec<AssignmentRegister> = {
    <mut list:( <AssignmentRegister> "," )*> <end:AssignmentRegister> ","? => { list.push(end); list },
    => vec![]
}

AssignmentRegister: AssignmentRegister = {
    <Identifier> => AssignmentRegister::Register(<>),
    "_" => AssignmentRegister::Wildcard,
}

ReturnStatement: FunctionStatement = {
    <start:@L> "return" <list:ExpressionList> <end:@R> ";" => FunctionStatement::Return(ctx.source_ref(start, end), list)
}

InstructionStatement: FunctionStatement = {
    <start:@L> <id:Identifier> <list:ExpressionList> <end:@R> ";" => FunctionStatement::Instruction(ctx.source_ref(start, end), id, list)
}

DebugDirectiveStatement: FunctionStatement = {
    <start:@L> ".debug" "file" <n:UnsignedInteger> <d:StringLiteral> <f:StringLiteral> <end:@R> ";"
        => FunctionStatement::DebugDirective(ctx.source_ref(start, end), DebugDirective::File(n.try_into().unwrap(), d, f)),
    <start:@L> ".debug" "loc" <f:UnsignedInteger> <line:UnsignedInteger> <col:UnsignedInteger> <end:@R> ";"
        => FunctionStatement::DebugDirective(ctx.source_ref(start, end), DebugDirective::Loc(f.try_into().unwrap(), line.try_into().unwrap(), col.try_into().unwrap())),
    <start:@L> ".debug" "insn" <insn:StringLiteral> <end:@R> ";"
        => FunctionStatement::DebugDirective(ctx.source_ref(start, end), DebugDirective::OriginalInstruction(insn)),
}

LabelStatement: FunctionStatement = {
    <start:@L> <id:Identifier> ":" <end:@R> => FunctionStatement::Label(ctx.source_ref(start, end), id)
}

// ---------------------------- Expressions -----------------------------
// Expressions are parameterized with StructOption because in some cases, we want to disallow the parsing of struct expressions.
// This allow us to use expression in two diferrent ways:
//  Expression<TermWithStruct> allows structs (renamed as Expression)
//  Expression<TermWithoutStruct> does not allow structs (renamed as RestrictedExpression).
//  Same for BoxedExpression.

ExpressionList: Vec<Expression> = {
    => vec![],
    <mut list:( <Expression> "," )*> <end:Expression> ","?  => { list.push(end); list }
}

pub Expression = MacroExpression<TermWithStruct>;
pub RestrictedExpression = MacroExpression<TermWithoutStruct>;

pub BoxedExpression = MacroBoxedExpression<TermWithStruct>;
pub RestrictedBoxedExpression = MacroBoxedExpression<TermWithoutStruct>;

MacroExpression<StructOption>: Expression = {
    MacroBoxedExpression<StructOption> => *<>,
}

MacroBoxedExpression<StructOption>: Box<Expression> = {
    LookupExpression<StructOption>,
}

LookupExpression<StructOption>: Box<Expression> = {
    <start:@L> <left:SelectedExpressions<StructOption>> "in" <right:SelectedExpressions<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation {left, op: BinaryOperator::In, right}, start, end),
    <start:@L> <left:SelectedExpressions<StructOption>> "is" <right:SelectedExpressions<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation {left, op: BinaryOperator::Is, right}, start, end),
    <start:@L> <left:SelectedExpressions<StructOption>> "connect" <right:SelectedExpressions<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation {left, op: BinaryOperator::Connect, right}, start, end),
    SelectedExpressions<StructOption>
}

SelectedExpressions<StructOption>: Box<Expression> = {
     <start:@L> <left:LambdaExpression<StructOption>> "$" <right:LambdaExpression<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation {left, op: BinaryOperator::Select, right}, start, end),
    LambdaExpression<StructOption>
}

LambdaExpression<StructOption>: Box<Expression> = {
    <start:@L> <kind:FunctionKind> "||" <body:LambdaExpression<StructOption>> <end:@R>
        => ctx.to_expr_with_source_ref(LambdaExpression{kind, params: vec![], body, param_types: vec![]}, start, end),
    <start:@L> <kind:FunctionKind> "|" <params:ParameterList> "|" <body:LambdaExpression<StructOption>> <end:@R>
        => ctx.to_expr_with_source_ref(LambdaExpression{kind, params, body, param_types: vec![]}, start, end),
    LogicalOr<StructOption>
}

FunctionKind: FunctionKind = {
    => FunctionKind::Pure,
    "query" => FunctionKind::Query,
    "constr" => FunctionKind::Constr,
}

LogicalOr<StructOption>: Box<Expression> = {
    <start:@L> <left:LogicalOr<StructOption>> "||" <right:LogicalAnd<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op: BinaryOperator::LogicalOr, right }, start, end),
    LogicalAnd<StructOption>,
}

LogicalAnd<StructOption>: Box<Expression> = {
    <start:@L> <left:LogicalAnd<StructOption>> "&&" <right:Comparison<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op: BinaryOperator::LogicalAnd, right }, start, end),
    Comparison<StructOption>,
}

Comparison<StructOption>: Box<Expression> = {
    <start:@L> <left:BinaryOr<StructOption>> <op:ComparisonOp> <right:BinaryOr<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op: op, right }, start, end),
    BinaryOr<StructOption>
}

ComparisonOp: BinaryOperator = {
    "<" => BinaryOperator::Less,
    "<=" => BinaryOperator::LessEqual,
    "==" => BinaryOperator::Equal,
    "=" => BinaryOperator::Identity,
    "!=" => BinaryOperator::NotEqual,
    ">=" => BinaryOperator::GreaterEqual,
    ">" => BinaryOperator::Greater,
}

BinaryOr<StructOption>: Box<Expression> = {
    <start:@L> <left:BinaryOr<StructOption>> <op:BinaryOrOp> <right:BinaryXor<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    BinaryXor<StructOption>,
}

BinaryOrOp: BinaryOperator = {
    "|" => BinaryOperator::BinaryOr,
}

BinaryXor<StructOption>: Box<Expression> = {
    <start:@L> <left:BinaryXor<StructOption>> <op:BinaryXorOp> <right:BinaryAnd<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    BinaryAnd<StructOption>,
}

BinaryXorOp: BinaryOperator = {
    "^" => BinaryOperator::BinaryXor,
}

BinaryAnd<StructOption>: Box<Expression> = {
    <start:@L> <left:BinaryAnd<StructOption>> <op:BinaryAndOp> <right:BitShift<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    BitShift<StructOption>,
}

BinaryAndOp: BinaryOperator = {
    "&" => BinaryOperator::BinaryAnd,
}

BitShift<StructOption>: Box<Expression> = {
    <start:@L> <left:BitShift<StructOption>> <op:BitShiftOp> <right:Sum<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    Sum<StructOption>,
}

BitShiftOp: BinaryOperator = {
    "<<" => BinaryOperator::ShiftLeft,
    ">>" => BinaryOperator::ShiftRight,
}

Sum<StructOption>: Box<Expression> = {
    <start:@L> <left:Sum<StructOption>> <op:SumOp> <right:Product<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    Product<StructOption>,
}

SumOp: BinaryOperator = {
    "+" => BinaryOperator::Add,
    "-" => BinaryOperator::Sub,
}

Product<StructOption>: Box<Expression> = {
    <start:@L> <left:Product<StructOption>> <op:ProductOp> <right:Power<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    Power<StructOption>,
}

ProductOp: BinaryOperator = {
    "*" => BinaryOperator::Mul,
    "/" => BinaryOperator::Div,
    "%" => BinaryOperator::Mod,
}

Power<StructOption>: Box<Expression> = {
    <start:@L> <left:Term<StructOption>> <op:PowOp> <right:Power<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(BinaryOperation { left, op, right }, start, end),
    Unary<StructOption>,
}

PowOp: BinaryOperator = {
    "**" => BinaryOperator::Pow,
}

Unary<StructOption>: Box<Expression> = {
    <start:@L> <op:PrefixUnaryOp> <expr:PostfixUnary<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(UnaryOperation{op, expr}, start, end),
    PostfixUnary<StructOption>,
}

PrefixUnaryOp: UnaryOperator = {
    "-" => UnaryOperator::Minus,
    "!" => UnaryOperator::LogicalNot,
}

PostfixUnary<StructOption>: Box<Expression> = {
    <start:@L> <t:Term<StructOption>> <op:PostfixUnaryOp> <end:@R> => ctx.to_expr_with_source_ref(UnaryOperation{op, expr: t}, start, end),
    Term<StructOption>,
}

PostfixUnaryOp: UnaryOperator = {
    "'" => UnaryOperator::Next,
}

Term<StructOption>: Box<Expression> = {
    StructOption,
}

TermWithStruct: Box<Expression> = {
    TermCommon<TermWithStruct>,
    StructExpression,
}

TermWithoutStruct: Box<Expression> = {
    TermCommon<TermWithoutStruct>,
}


TermCommon<StructOption>: Box<Expression> = {
    <start:@L> <i:IndexAccess<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(i, start, end),
    <start:@L> <f:FunctionCall<StructOption>> <end:@R> => ctx.to_expr_with_source_ref(f, start, end),
    <start:@L> <i:ConstantIdentifier> <end:@R> => ctx.to_expr_with_source_ref(NamespacedPolynomialReference::from_identifier(i), start, end),
    <start:@L> <g:GenericReference> <end:@R> => ctx.to_expr_with_source_ref(g, start, end),
    <start:@L> <value:Number> <end:@R> => ctx.to_expr_with_source_ref(Number {value: value.into(), type_: None}, start, end),
    <start:@L> <number:TypedNumber> <end:@R> => { let (value, ty) = number; ctx.to_expr_with_source_ref(Number {value, type_: Some(ty)}, start, end) },
    <start:@L> <s:StringLiteral> <end:@R> => Box::new(Expression::String(ctx.source_ref(start, end), s)),
    MatchExpression,
    IfExpression,
    BlockExpression,
    <start:@L> "[" <items:ExpressionList> "]" <end:@R> => ctx.to_expr_with_source_ref(ArrayLiteral{items}, start, end),
    <start:@L> "(" <head:Expression> "," <tail:ExpressionList> ")" <end:@R> => { let mut list = vec![head]; list.extend(tail); Box::new(Expression::Tuple(ctx.source_ref(start, end), list)) },
    <start:@L> "(" ")" <end:@R> => Box::new(Expression::Tuple(ctx.source_ref(start, end), vec![])),
    "(" <BoxedExpression> ")",
    <start:@L> "${" <b:BoxedExpression> "}" <end:@R> => Box::new(Expression::FreeInput(ctx.source_ref(start, end), b))
}

IndexAccess<StructOption>: IndexAccess = {
    <array:Term<StructOption>> "[" <index:BoxedExpression> "]" => IndexAccess{<>},
}

FunctionCall<StructOption>: FunctionCall = {
    <function:Term<StructOption>> "(" <arguments:ExpressionList> ")" => FunctionCall {<>},
}

NamespacedPolynomialReference: NamespacedPolynomialReference = {
    <SymbolPath> => <>.into(),
    <namespace:Identifier> "." <name:Identifier> => SymbolPath::from_parts([namespace, name].into_iter().map(Part::Named)).into(),
}

GenericReference: NamespacedPolynomialReference = {
    <namespace:Identifier> "." <name:Identifier> => SymbolPath::from_parts([namespace, name].into_iter().map(Part::Named)).into(),
    <path:GenericSymbolPath> => NamespacedPolynomialReference{path: path.0, type_args: path.1},
}

MatchExpression: Box<Expression> = {
    <start:@L> "match" <scrutinee:RestrictedBoxedExpression> "{" <arms:MatchArms> "}" <end:@R> => ctx.to_expr_with_source_ref(MatchExpression{scrutinee, arms}, start, end),
}

MatchArms: Vec<MatchArm> = {
    => vec![],
    <mut list:( <MatchArm> "," )*> <end:MatchArm> ","?  => { list.push(end); list }
}

MatchArm: MatchArm = {
    <pattern:Pattern> "=>" <value: Expression> => MatchArm{pattern, value},
}

IfExpression: Box<Expression> = {
    <start:@L> "if" <condition:RestrictedBoxedExpression>
        <body:BlockExpression>
        "else"
        <else_body:BlockExpression>
        <end:@R> => ctx.to_expr_with_source_ref(IfExpression{condition, body, else_body}, start, end),
}

StructExpression: Box<Expression> = {	
    <start:@L> <name:NamespacedPolynomialReference> "{" <fields:NamedExpressions> "}" <end:@R> => ctx.to_expr_with_source_ref(StructExpression{name, fields}, start, end),
}

BlockExpression: Box<Expression> = {
    <start:@L> "{}" <end:@R> => ctx.to_expr_with_source_ref(BlockExpression{statements: vec![], expr: None}, start, end),
    <start:@L> "{" <statements:StatementInsideBlock*> <expr:BoxedExpression?> "}" <end:@R> => ctx.to_expr_with_source_ref(BlockExpression{statements, expr}, start, end),
}

StatementInsideBlock: StatementInsideBlock = {
    "let" <pattern:Pattern> <ty:(":" <Type<ArrayLengthNumber>>)?> <value:("=" <Expression>)?> ";" => StatementInsideBlock::LetStatement(LetStatementInsideBlock{<>}),
    <Expression> ";" => StatementInsideBlock::Expression(<>),
}

NamedExpressions: Vec<NamedExpression<Box<Expression>>> = {
    => vec![],
    <mut list:( <NamedExpression> "," )*> <end:NamedExpression> ","?  => { list.push(end); list }
}

NamedExpression: NamedExpression<Box<Expression>> = {
    <name:Identifier> ":" <body:BoxedExpression> => NamedExpression { name, body }
}

NamedTypes: Vec<NamedType<Expression>> = {
    => vec![],
    <mut list:( <NamedType> "," )*> <end:NamedType> ","?  => { list.push(end); list }
}

NamedType: NamedType<Expression> = {
    <name:Identifier> ":" <ty:TypeTerm<Expression>> => NamedType { name, ty }
}

// ---------------------------- Pattern -----------------------------

Pattern: Pattern = {
    <start:@L> "_" <end:@R> => Pattern::CatchAll(ctx.source_ref(start, end)),
    <start:@L> <n:Number> <end:@R> => Pattern::Number(ctx.source_ref(start, end), n.into()),
    <start:@L> "-" <n:Number> <end:@R> => Pattern::Number(ctx.source_ref(start, end), -BigInt::from(n)),
    <start:@L> <l:StringLiteral> <end:@R> => Pattern::String(ctx.source_ref(start, end), l),
    TuplePattern,
    ArrayPattern,
    // We never parse Pattern::Variable directly. Instead, we
    // parse all identifiers (and symbol paths) as Pattern::Enum.
    // At symbol resolution stage, single identifiers that do not
    // exist in the current scope are then turned into Pattern::Variable
    // to allow `match x { (y, None) => ... }` to bind the first component
    // to `y` and match the enum variant `None` at the same time.
    EnumPattern,
}

PatternIncludingEllipsis: Pattern = {
    Pattern => <>,
    <start:@L> ".." <end:@R> => Pattern::Ellipsis(ctx.source_ref(start, end)),
}

TuplePattern: Pattern = {
    <start:@L> "(" ")" <end:@R> => Pattern::Tuple(ctx.source_ref(start, end), vec![]),
    <start:@L> "(" <mut items:(<Pattern> ",")+> <last:Pattern> ")" <end:@R> => Pattern::Tuple(ctx.source_ref(start, end), {items.push(last); items})
}

ArrayPattern: Pattern = {
    <start:@L> "[" "]" <end:@R> => Pattern::Array(ctx.source_ref(start, end), vec![]),
    <start:@L> "[" <mut items:(<PatternIncludingEllipsis> ",")*> <last:PatternIncludingEllipsis> "]" <end:@R> => Pattern::Array(ctx.source_ref(start, end), {items.push(last); items})
}

EnumPattern: Pattern = {
    <start:@L> <n:SymbolPath> <items:( "(" <PatternList> ")" )?> <end:@R> => Pattern::Enum(ctx.source_ref(start, end), n, items)
}

PatternList: Vec<Pattern> = {
    <mut list:( <Pattern> "," )*> <end:Pattern> ","? => { list.push(end); list },
    => vec![]
}

// ---------------------------- Trait/Impl Declaration -----------------------------

TraitDeclaration: TraitDeclaration<Expression> = {
    "trait" <name:Identifier> <type_vars:("<" <TraitVars> ">")> "{" <functions:TraitFunctions> "}" => TraitDeclaration { name, type_vars, functions }
}

TraitVars: Vec<String> = {
    <mut list:( <TypeVar> "," )*> <end:TypeVar> ","?  => { list.push(end); list }
}

TraitFunctions: Vec<NamedType<Expression>> = {
    => vec![],
    <mut list:( <TraitFunction> "," )*> <end:TraitFunction> ","?  => { list.push(end); list }
}

TraitFunction: NamedType<Expression> = {
    <name:Identifier> ":" <params:TypeTermList<Expression>> "->" <value:TypeTermBox<Expression>> => NamedType { name, ty: Type::Function(FunctionType{params, value}) }
}

TraitImplementation: TraitImplementation<Expression> = {
    <start:@L> "impl" <type_scheme: GenericTraitName> "{" <functions:NamedArcExpressions> "}" <end:@L> => TraitImplementation { name: type_scheme.0, source_ref: ctx.source_ref(start, end), type_scheme: type_scheme.1, functions }
}


NamedArcExpressions: Vec<NamedExpression<Arc<Expression>>> = {
    => vec![],
    <mut list:( <NamedArcExpression> "," )*> <end:NamedArcExpression> ","?  => { list.push(end); list }
}

NamedArcExpression: NamedExpression<Arc<Expression>> = {
    <name:Identifier> ":" <body:Expression> => NamedExpression { name, body: body.into() }
}

GenericTraitName: (SymbolPath, TypeScheme) = {
    <vars:("<" <TypeVarBounds> ">")?> <name:SymbolPath> <items:("<" <TypeTermList<ArrayLengthNumber>> ">")> =>
        (name, TypeScheme{ vars: vars.unwrap_or_default(), ty: Type::Tuple(TupleType{items}) })
}


// ---------------------------- Type Declarations -----------------------------

EnumDeclaration: EnumDeclaration<Expression> = {
    "enum" <name:Identifier> <type_vars:("<" <TypeVarBounds> ">")?> "{" <variants:EnumVariants> "}"
        => EnumDeclaration{name, type_vars: type_vars.unwrap_or_default(), variants}
}

EnumVariants: Vec<EnumVariant<Expression>> = {
    => vec![],
    <mut list:( <EnumVariant> "," )*> <end:EnumVariant> ","?  => { list.push(end); list }
}

EnumVariant: EnumVariant<Expression> = {
    <name:Identifier> <fields:("(" <TypeTermList<Expression>> ")")?> => EnumVariant{<>}
}

StructDeclaration: StructDeclaration<Expression> = {
    "struct" <name:Identifier> <type_vars:("<" <TypeVarBounds> ">")?> "{" <fields:NamedTypes> "}"
        => StructDeclaration{name, type_vars: type_vars.unwrap_or_default(), fields}
}

// ---------------------------- Type Names -----------------------------

pub TypeExpr = Type<Expression>;
pub TypeNumber = Type<Number>;

Type<ArrayLength>: Type<ArrayLength> = {
    <params:TypeTermList<ArrayLength>> "->" <value:TypeTermBox<ArrayLength>> => Type::Function(FunctionType{<>}),
    TypeTerm<ArrayLength>
}

TypeTermList<ArrayLength>: Vec<Type<ArrayLength>> = {
    => vec![],
    <mut list:( <TypeTerm<ArrayLength>> "," )*> <end:TypeTerm<ArrayLength>> ","? => { list.push(end); list }
}

TypeTermBox<ArrayLength>: Box<Type<ArrayLength>> = {
    TypeTerm<ArrayLength> => Box::new(<>)
}

TypeTerm<ArrayLength>: Type<ArrayLength> = {
    // The parser parses all identifiers as NamedTypes, some are translated
    // to TypeVars later.
    TypeSymbolPath ("<" <TypeTermList<ArrayLength>> ">")? => Type::NamedType(<>),
    "!" => Type::Bottom,
    "bool" => Type::Bool,
    "int" => Type::Int,
    "fe" => Type::Fe,
    "string" => Type::String,
    "col" => Type::Col,
    "inter" => Type::Inter,
    "expr" => Type::Expr,
    <base:TypeTerm<ArrayLength>> "[" <length:ArrayLength?> "]" => Type::Array(ArrayType{base: Box::new(base), length}),
    "(" <mut items:( <TypeTerm<ArrayLength>> "," )+> <end:TypeTerm<ArrayLength>> ")" => { items.push(end); Type::Tuple(TupleType{items}) },
    "(" ")" => Type::Tuple(TupleType{items: vec![]}),
    "(" <Type<ArrayLength>> ")",
}

ArrayLengthExpr: Expression = {
    Expression
}

ArrayLengthNumber: u64 = {
    Number => u64::try_from(<>).unwrap()
}

TypeVar: String = {
    UppercaseIdentifier => <>,
}


// bounds should be SymbolPaths in the future.
pub TypeVarBounds: TypeBounds = {
    => Default::default(),
    <list:( <TypeVarWithBounds> "," )*> <end:TypeVarWithBounds> ","? => TypeBounds::new(list.into_iter().chain(std::iter::once(end)))
}

TypeVarWithBounds: (String, BTreeSet<String>) = {
    <TypeVar> <TypeBoundsList> => (<>)
}

TypeBoundsList: BTreeSet<String> = {
    => Default::default(),
    ":" <list:( <Identifier> "+" )*> <end:Identifier>  => list.into_iter().chain(std::iter::once(end)).collect(),
}



// ---------------------------- Terminals -----------------------------



StringLiteral: String = {
    r#""[^\\"\n\r]*(\\[tnfbrx'"\\0-9][^\\"\n\r]*)*""# => unescape_string(<>)
}

Identifier: String = {
    r"[a-z_][a-zA-Z$_0-9@]*" => <>.to_string(),
    UppercaseIdentifier,
    SpecialIdentifier => <>.to_string(),
}

NonSpecialIdentifier: String = {
    r"[a-z_][a-zA-Z$_0-9@]*" => <>.to_string(),
    UppercaseIdentifier => <>.to_string(),
}

UppercaseIdentifier: String = {
    r"[A-Z][a-zA-Z$_0-9@]*" => <>.to_string(),
}

// These identifier are special in the way that the lexer would treat them as keywords,
// but these keywords can never occur in a place where an identifier can be expected,
// so we allow them as identifiers as well.
SpecialIdentifier: &'input str = {
    "file",
    "loc",
    "insn",
    "int",
    "inter",
    "fe",
    "expr",
    "bool",
}

ConstantIdentifier: String = {
    r"%[a-zA-Z_][a-zA-Z$_0-9@]*" => <>.to_string(),
}

Number: BigUint = {
    r"[0-9][0-9_]*" => BigUint::from_str(&<>.replace('_', "")).unwrap().into(),
    <start:@L> <v:r"0x[0-9A-Fa-f][0-9A-Fa-f_]*"> <end:@R> =>?
        if v.ends_with("_fe") {
            // It would be too easy to confuse e.g. 0x12_fe with the field element 0x12.
            Err(ParseError::User{ error: ctx.source_ref(start, end).with_error("Hex literals cannot end in '_fe'.".to_string()) })
        } else {
            Ok(BigUint::from_str_radix(&v[2..].replace('_', ""), 16).unwrap().into())
        },
}

TypedNumber: (BigUint, Type) = {
    r"[0-9][0-9_]*_fe" => (BigUint::from_str(&<>.strip_suffix("_fe").unwrap().replace('_', "")).unwrap().into(), Type::Fe),
    // No hex with _fe suffix for now, because 'fe' is valid hex.
    r"[0-9][0-9_]*_int" => (BigUint::from_str(&<>.strip_suffix("_int").unwrap().replace('_', "")).unwrap().into(), Type::Int),
    r"0x[0-9A-Fa-f][0-9A-Fa-f_]*_int" => (BigUint::from_str_radix(&<>[2..].strip_suffix("_int").unwrap().replace('_', ""), 16).unwrap().into(), Type::Int),
    r"[0-9][0-9_]*_expr" => (BigUint::from_str(&<>.strip_suffix("_expr").unwrap().replace('_', "")).unwrap().into(), Type::Expr),
    r"0x[0-9A-Fa-f][0-9A-Fa-f_]*_expr" => (BigUint::from_str_radix(&<>[2..].strip_suffix("_expr").unwrap().replace('_', ""), 16).unwrap().into(), Type::Expr),
}

UnsignedInteger: BigUint = {
    r"[0-9][0-9_]*" => BigUint::from_str(&<>.replace('_', "")).unwrap(),
    r"0x[0-9A-Fa-f][0-9A-Fa-f_]*" => BigUint::from_str_radix(&<>[2..].replace('_', ""), 16).unwrap(),
}
