atoxide-parser 0.1.3

//! Expression parsing for the Ato language using chumsky 0.12.
//!
//! This module implements expression parsing with proper operator precedence
//! using manual precedence climbing (chumsky 0.9 doesn't have pratt()).

use crate::ast::*;
use atoxide_lexer::{Span, Token, TokenKind};
use chumsky::input::ValueInput;
use chumsky::prelude::*;

/// Chumsky error extra type alias for all parsers in this module.
type Err<'a> = extra::Err<Rich<'a, Token, SimpleSpan>>;

/// Convert a SimpleSpan to our Span type.
fn to_span(span: SimpleSpan) -> Span {
    Span::new(span.start, span.end, 1, 1)
}

/// Helper to create a token matcher.
pub fn tok<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>(
    kind: TokenKind,
) -> impl Parser<'a, I, Token, Err<'a>> + Clone {
    any().filter(move |t: &Token| t.kind == kind)
}

/// Parse an identifier.
/// Keywords can be used as identifiers in certain contexts (e.g., pin names,
/// signal names, field references like `module.in`, template arguments).
/// This is safe because in statement-level parsing, keyword-specific parsers
/// (block_def, for_stmt, import_stmt, etc.) are tried before name_starting_stmt,
/// so keywords are consumed as keywords when they start a statement.
pub fn identifier<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Identifier, Err<'a>> + Clone {
    any()
        .filter(|t: &Token| {
            matches!(
                t.kind,
                TokenKind::Name
                    // Active keywords that can also appear as identifiers
                    | TokenKind::In
                    | TokenKind::To
                    | TokenKind::From
                    | TokenKind::Is
                    | TokenKind::Within
                    | TokenKind::Pin
                    | TokenKind::Signal
                    | TokenKind::New
                    | TokenKind::Assert
                    | TokenKind::Pass
                    | TokenKind::Trait
                    | TokenKind::Component
                    | TokenKind::Module
                    | TokenKind::Interface
                    | TokenKind::Import
                    | TokenKind::For
                    // Reserved keywords (already allowed)
                    | TokenKind::Int
                    | TokenKind::Float
                    | TokenKind::StringKw
                    | TokenKind::Str
                    | TokenKind::Bytes
                    | TokenKind::Parameter
                    | TokenKind::Param
                    | TokenKind::Test
                    | TokenKind::Require
                    | TokenKind::Requires
                    | TokenKind::Check
                    | TokenKind::Report
                    | TokenKind::Ensure
            )
        })
        .map_with(|t: Token, e| {
            let span = e.span();
            Identifier {
                name: t.text.clone(),
                span: to_span(span),
            }
        })
}

/// Parse a number literal.
pub fn number_literal<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, NumberLiteral, Err<'a>> + Clone {
    tok(TokenKind::Number).map_with(|t: Token, e| {
        let span = e.span();
        NumberLiteral {
            value: t.text.clone(),
            span: to_span(span),
        }
    })
}

/// Parse a string literal.
pub fn string_literal<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, StringLiteral, Err<'a>> + Clone {
    tok(TokenKind::String).map_with(|t: Token, e| {
        let span = e.span();
        StringLiteral {
            value: t.text.clone(),
            span: to_span(span),
        }
    })
}

/// Parse a boolean literal.
pub fn bool_literal<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, BoolLiteral, Err<'a>> + Clone {
    choice((
        tok(TokenKind::True).map_with(|_, e| {
            let span = e.span();
            BoolLiteral {
                value: true,
                span: to_span(span),
            }
        }),
        tok(TokenKind::False).map_with(|_, e| {
            let span = e.span();
            BoolLiteral {
                value: false,
                span: to_span(span),
            }
        }),
    ))
}

/// Parse an array index [n].
fn array_index<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, NumberLiteral, Err<'a>> + Clone {
    tok(TokenKind::OpenBracket)
        .ignore_then(number_literal())
        .then_ignore(tok(TokenKind::CloseBracket))
}

/// Parse a field reference part (name with optional index).
fn field_ref_part<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, FieldRefPart, Err<'a>> + Clone {
    identifier()
        .then(array_index().or_not())
        .map_with(|(name, index), e| {
            let span = e.span();
            FieldRefPart {
                name,
                index,
                span: to_span(span),
            }
        })
}

/// Parse a field reference (a.b.c[0].d).
pub fn field_reference<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, FieldRef, Err<'a>> + Clone {
    field_ref_part()
        .separated_by(tok(TokenKind::Dot))
        .at_least(1)
        .collect::<Vec<_>>()
        .then(
            // Optional trailing pin reference (.1)
            tok(TokenKind::Dot).ignore_then(number_literal()).or_not(),
        )
        .map_with(|(parts, pin_ref), e| {
            let span = e.span();
            FieldRef {
                parts,
                pin_ref,
                span: to_span(span),
            }
        })
}

/// Parse a type reference (Name.Name.Name).
pub fn type_reference<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, TypeRef, Err<'a>> + Clone {
    identifier()
        .separated_by(tok(TokenKind::Dot))
        .at_least(1)
        .collect::<Vec<_>>()
        .map_with(|parts, e| {
            let span = e.span();
            TypeRef {
                parts,
                span: to_span(span),
            }
        })
}

/// Split a number token text into numeric and unit parts.
/// E.g., "10kohm" -> ("10", Some("kohm")), "42" -> ("42", None)
fn split_number_unit(text: &str) -> (&str, Option<&str>) {
    // Find where the unit suffix starts (first alphabetic character after digits/dots/exponent)
    let mut unit_start = text.len();
    let mut seen_e = false;
    let mut in_exponent = false;

    for (i, c) in text.char_indices() {
        if c == 'e' || c == 'E' {
            seen_e = true;
            in_exponent = true;
        } else if in_exponent && (c == '+' || c == '-' || c.is_ascii_digit()) {
            // Part of exponent (sign or digits)
        } else if c.is_alphabetic() || c == '_' {
            // This is the start of the unit suffix
            // Unless we're right after 'e'/'E' (which was an exponent marker)
            if !seen_e || (i > 0 && !text[..i].ends_with('e') && !text[..i].ends_with('E')) {
                unit_start = i;
                break;
            }
        } else if c.is_ascii_digit() || c == '.' {
            in_exponent = false;
        }
    }

    if unit_start == text.len() {
        (text, None)
    } else {
        (&text[..unit_start], Some(&text[unit_start..]))
    }
}

/// Parse a quantity (number with optional unit).
/// Handles both embedded units ("10kohm") and space-separated units ("300 A", "32 kHz").
pub fn quantity<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Quantity, Err<'a>> + Clone {
    // Optional sign
    let sign = choice((tok(TokenKind::Minus).to("-"), tok(TokenKind::Plus).to("")))
        .or_not()
        .map(|s| s.unwrap_or(""));

    sign.then(number_literal())
        .then(
            // Try to parse a separate unit token (Name) after the number
            any()
                .filter(|t: &Token| t.kind == TokenKind::Name)
                .map_with(|t: Token, e| Identifier {
                    name: t.text.clone(),
                    span: to_span(e.span()),
                })
                .or_not(),
        )
        .map_with(|((sign, num), separate_unit), e| {
            let span = e.span();
            let (num_part, embedded_unit) = split_number_unit(&num.value);

            let number_value = if sign == "-" {
                format!("-{}", num_part)
            } else {
                num_part.to_string()
            };

            // Prefer embedded unit (10kohm) over separate unit (10 kohm)
            // But use separate unit if no embedded unit exists
            let unit = embedded_unit
                .map(|u| Identifier {
                    name: u.to_string(),
                    span: num.span, // Approximate span
                })
                .or(separate_unit);

            Quantity {
                number: NumberLiteral {
                    value: number_value,
                    span: num.span,
                },
                unit,
                span: to_span(span),
            }
        })
}

/// Parse a tolerance value.
/// Handles cases like "5%", "5", "5mA" where the unit might be in the number token.
fn tolerance<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Tolerance, Err<'a>> + Clone {
    number_literal()
        .then(tok(TokenKind::Percent).or_not())
        .map_with(|(num, is_percent), e| {
            let span = e.span();
            let (num_part, unit_part) = split_number_unit(&num.value);

            let unit = unit_part.map(|u| Identifier {
                name: u.to_string(),
                span: num.span,
            });

            Tolerance {
                value: num_part.to_string(),
                is_percent: is_percent.is_some(),
                unit,
                span: to_span(span),
            }
        })
}

/// Parse a physical literal (quantity, range, or bilateral).
pub fn physical_literal<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, PhysicalLiteral, Err<'a>> + Clone {
    quantity()
        .then(
            choice((
                // Range: quantity TO quantity
                tok(TokenKind::To)
                    .ignore_then(quantity())
                    .map(PhysicalSuffix::Range),
                // Bilateral: quantity +/- tolerance
                tok(TokenKind::PlusOrMinus)
                    .ignore_then(tolerance())
                    .map(PhysicalSuffix::Bilateral),
            ))
            .or_not(),
        )
        .map_with(|(base, suffix), e| {
            let span = e.span();
            match suffix {
                Some(PhysicalSuffix::Range(to)) => PhysicalLiteral::Range(QuantityRange {
                    from: base,
                    to,
                    span: to_span(span),
                }),
                Some(PhysicalSuffix::Bilateral(tol)) => {
                    PhysicalLiteral::Bilateral(BilateralQuantity {
                        base,
                        tolerance: tol,
                        span: to_span(span),
                    })
                }
                None => PhysicalLiteral::Quantity(base),
            }
        })
}

/// Helper enum for physical literal parsing.
enum PhysicalSuffix {
    Range(Quantity),
    Bilateral(Tolerance),
}

/// Parse a literal.
pub fn literal<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Literal, Err<'a>> + Clone {
    choice((
        physical_literal().map(Literal::Physical),
        string_literal().map(Literal::String),
        bool_literal().map(Literal::Bool),
    ))
}

/// Parse a primary expression (atom).
/// Note: The grouped expression case needs to accept any arithmetic expression,
/// not just atoms. We achieve this by making arithmetic_expression recursive.
fn atom<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>(
    full_expr: impl Parser<'a, I, Expression, Err<'a>> + Clone,
) -> impl Parser<'a, I, Expression, Err<'a>> + Clone {
    choice((
        // Grouped expression - accepts full arithmetic expressions
        full_expr
            .delimited_by(tok(TokenKind::OpenParen), tok(TokenKind::CloseParen))
            .map(|e| Expression::Group(Box::new(e))),
        // Physical literal (must come before field_reference to handle numbers)
        physical_literal().map(|p| Expression::Literal(Literal::Physical(p))),
        // Field reference
        field_reference().map(Expression::FieldRef),
        // String literal
        string_literal().map(|s| Expression::Literal(Literal::String(s))),
        // Boolean literal
        bool_literal().map(|b| Expression::Literal(Literal::Bool(b))),
    ))
}

/// Parse an arithmetic expression with proper precedence.
/// This is the main entry point for expression parsing with full recursion support.
pub fn arithmetic_expression<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Expression, Err<'a>> + Clone {
    recursive(|full_expr| {
        // Power expression (** is right-associative)
        let power = recursive(|power| {
            atom(full_expr.clone())
                .then(tok(TokenKind::Power).ignore_then(power).or_not())
                .map_with(|(base, exp), e| {
                    let span = e.span();
                    match exp {
                        Some(e) => Expression::Binary(Box::new(BinaryExpr {
                            left: base,
                            operator: BinaryOp::Power,
                            right: e,
                            span: to_span(span),
                        })),
                        None => base,
                    }
                })
        });

        // Term (* and /)
        let term = power.clone().foldl(
            choice((
                tok(TokenKind::Star).to(BinaryOp::Mul),
                tok(TokenKind::Div).to(BinaryOp::Div),
            ))
            .then(power)
            .repeated(),
            |left, (op, right)| {
                let span = left.span().merge(&right.span());
                Expression::Binary(Box::new(BinaryExpr {
                    left,
                    operator: op,
                    right,
                    span,
                }))
            },
        );

        // Sum (+ and -)
        let sum = term.clone().foldl(
            choice((
                tok(TokenKind::Plus).to(BinaryOp::Add),
                tok(TokenKind::Minus).to(BinaryOp::Sub),
            ))
            .then(term)
            .repeated(),
            |left, (op, right)| {
                let span = left.span().merge(&right.span());
                Expression::Binary(Box::new(BinaryExpr {
                    left,
                    operator: op,
                    right,
                    span,
                }))
            },
        );

        // Bitwise operations (| and &, lowest precedence)
        sum.clone().foldl(
            choice((
                tok(TokenKind::OrOp).to(BinaryOp::BitOr),
                tok(TokenKind::AndOp).to(BinaryOp::BitAnd),
            ))
            .then(sum)
            .repeated(),
            |left, (op, right)| {
                let span = left.span().merge(&right.span());
                Expression::Binary(Box::new(BinaryExpr {
                    left,
                    operator: op,
                    right,
                    span,
                }))
            },
        )
    })
}

/// Parse a comparison operator.
fn compare_op<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, CompareOpKind, Err<'a>> + Clone {
    choice((
        tok(TokenKind::LessEq).to(CompareOpKind::LessEq),
        tok(TokenKind::GreaterEq).to(CompareOpKind::GreaterEq),
        tok(TokenKind::LessThan).to(CompareOpKind::LessThan),
        tok(TokenKind::GreaterThan).to(CompareOpKind::GreaterThan),
        tok(TokenKind::Within).to(CompareOpKind::Within),
        tok(TokenKind::Is).to(CompareOpKind::Is),
    ))
}

/// Parse a comparison expression.
pub fn comparison<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Comparison, Err<'a>> + Clone {
    arithmetic_expression()
        .then(
            compare_op()
                .then(arithmetic_expression())
                .map_with(|(kind, right), e| {
                    let span = e.span();
                    CompareOp {
                        kind,
                        right,
                        span: to_span(span),
                    }
                })
                .repeated()
                .at_least(1)
                .collect::<Vec<_>>(),
        )
        .map_with(|(left, operations), e| {
            let span = e.span();
            Comparison {
                left,
                operations,
                span: to_span(span),
            }
        })
}

#[cfg(test)]
mod tests {
    use super::*;
    use atoxide_lexer::Lexer;
    use chumsky::input::Input;

    fn parse_expr(source: &str) -> Expression {
        let lexer = Lexer::new(source);
        let tokens: Vec<Token> = lexer.collect();
        let len = source.len();

        let token_spans: Vec<(Token, SimpleSpan)> = tokens
            .into_iter()
            .map(|t| {
                let sp: SimpleSpan = (t.span.start..t.span.end).into();
                (t, sp)
            })
            .collect();
        let eoi: SimpleSpan = (len..len).into();
        let input = token_spans.as_slice().map(eoi, |(t, s)| (t, s));
        let (output, _errors) = arithmetic_expression()
            .then_ignore(tok(TokenKind::Newline).or_not())
            .then_ignore(tok(TokenKind::Eof).or_not())
            .then_ignore(end())
            .parse(input)
            .into_output_errors();
        output.unwrap()
    }

    #[test]
    fn test_parse_number() {
        let expr = parse_expr("42");
        assert!(matches!(expr, Expression::Literal(Literal::Physical(_))));
    }

    #[test]
    fn test_parse_quantity() {
        let expr = parse_expr("10kohm");
        if let Expression::Literal(Literal::Physical(PhysicalLiteral::Quantity(q))) = expr {
            assert_eq!(q.number.value, "10");
            assert!(q.unit.is_some());
            assert_eq!(q.unit.as_ref().unwrap().name, "kohm");
        } else {
            panic!("Expected quantity");
        }
    }

    #[test]
    fn test_parse_binary_add() {
        let expr = parse_expr("1 + 2");
        if let Expression::Binary(b) = expr {
            assert_eq!(b.operator, BinaryOp::Add);
        } else {
            panic!("Expected binary expression");
        }
    }

    #[test]
    fn test_parse_precedence() {
        // 1 + 2 * 3 should parse as 1 + (2 * 3)
        let expr = parse_expr("1 + 2 * 3");
        if let Expression::Binary(b) = expr {
            assert_eq!(b.operator, BinaryOp::Add);
            assert!(matches!(b.right, Expression::Binary(_)));
        } else {
            panic!("Expected binary expression");
        }
    }
}