atoxide-parser 0.1.3

//! Statement parsing for the Ato language using chumsky 0.12 with error recovery.

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

use super::expr::*;

/// 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)
}

/// Parse a pragma statement.
fn pragma_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Pragma)
        .then_ignore(tok(TokenKind::Newline).or_not())
        .map_with(|t: Token, e| {
            let span = e.span();
            Statement::Pragma(PragmaStmt {
                content: t.text.clone(),
                span: to_span(span),
            })
        })
}

/// Parse an import statement.
fn import_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Import)
        .ignore_then(type_reference())
        .then(choice((
            // Deprecated form: import X from "path"
            tok(TokenKind::From)
                .ignore_then(string_literal())
                .map(ImportKind::Deprecated),
            // Standard form with additional imports
            tok(TokenKind::Comma)
                .ignore_then(type_reference())
                .repeated()
                .collect::<Vec<_>>()
                .map(ImportKind::Standard),
        )))
        .map_with(|(first, kind), e| {
            let span = e.span();
            match kind {
                ImportKind::Deprecated(path) => Statement::DepImport(DepImportStmt {
                    type_ref: first,
                    from_path: path,
                    span: to_span(span),
                }),
                ImportKind::Standard(rest) => {
                    let mut imports = vec![first];
                    imports.extend(rest);
                    Statement::Import(ImportStmt {
                        from_path: None,
                        imports,
                        span: to_span(span),
                    })
                }
            }
        })
}

enum ImportKind {
    Deprecated(StringLiteral),
    Standard(Vec<TypeRef>),
}

/// Parse a from...import statement.
fn from_import_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::From)
        .ignore_then(string_literal())
        .then_ignore(tok(TokenKind::Import))
        .then(
            type_reference()
                .separated_by(tok(TokenKind::Comma))
                .at_least(1)
                .collect::<Vec<_>>(),
        )
        .map_with(|(path, imports), e| {
            let span = e.span();
            Statement::Import(ImportStmt {
                from_path: Some(path),
                imports,
                span: to_span(span),
            })
        })
}

/// Parse a pin name.
fn pin_name<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, PinName, Err<'a>> + Clone {
    choice((
        identifier().map(PinName::Identifier),
        number_literal().map(PinName::Number),
        string_literal().map(PinName::String),
    ))
}

/// Parse a pin declaration.
fn pin_declaration<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Pin)
        .ignore_then(pin_name())
        .map_with(|name, e| {
            let span = e.span();
            Statement::PinDeclaration(PinDeclaration {
                name,
                span: to_span(span),
            })
        })
}

/// Parse a connectable element.
fn connectable<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Connectable, Err<'a>> + Clone {
    choice((
        tok(TokenKind::Signal)
            .ignore_then(identifier())
            .map_with(|name, e| {
                let span = e.span();
                Connectable::SignalDef(SignalDef {
                    name,
                    span: to_span(span),
                })
            }),
        tok(TokenKind::Pin)
            .ignore_then(pin_name())
            .map_with(|name, e| {
                let span = e.span();
                Connectable::PinDef(PinDeclaration {
                    name,
                    span: to_span(span),
                })
            }),
        field_reference().map(Connectable::FieldRef),
    ))
}

/// Parse a signal definition (possibly followed by connection).
fn signal_def<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Signal)
        .ignore_then(identifier())
        .then(
            choice((
                // signal ~ connectable
                tok(TokenKind::Wire)
                    .ignore_then(connectable())
                    .map(ConnectionSuffix::Simple),
                // signal ~> connectable ~> ...
                tok(TokenKind::Sperm)
                    .ignore_then(connectable())
                    .then(
                        tok(TokenKind::Sperm)
                            .ignore_then(connectable())
                            .repeated()
                            .collect::<Vec<_>>(),
                    )
                    .map(|(first, rest)| {
                        let mut elements = vec![first];
                        elements.extend(rest);
                        ConnectionSuffix::Forward(elements)
                    }),
                // signal <~ connectable <~ ...
                tok(TokenKind::LSperm)
                    .ignore_then(connectable())
                    .then(
                        tok(TokenKind::LSperm)
                            .ignore_then(connectable())
                            .repeated()
                            .collect::<Vec<_>>(),
                    )
                    .map(|(first, rest)| {
                        let mut elements = vec![first];
                        elements.extend(rest);
                        ConnectionSuffix::Backward(elements)
                    }),
            ))
            .or_not(),
        )
        .map_with(|(name, suffix), e| {
            let span = e.span();
            let signal = SignalDef {
                name: name.clone(),
                span: to_span(span),
            };

            match suffix {
                None => Statement::SignalDef(signal),
                Some(ConnectionSuffix::Simple(right)) => Statement::Connection(Connection {
                    left: Connectable::SignalDef(signal),
                    right,
                    span: to_span(span),
                }),
                Some(ConnectionSuffix::Forward(rest)) => {
                    let mut elements = vec![Connectable::SignalDef(signal)];
                    elements.extend(rest);
                    Statement::DirectedConnection(DirectedConnection {
                        direction: ConnectionDirection::Forward,
                        elements,
                        span: to_span(span),
                    })
                }
                Some(ConnectionSuffix::Backward(rest)) => {
                    let mut elements = vec![Connectable::SignalDef(signal)];
                    elements.extend(rest);
                    Statement::DirectedConnection(DirectedConnection {
                        direction: ConnectionDirection::Backward,
                        elements,
                        span: to_span(span),
                    })
                }
            }
        })
}

enum ConnectionSuffix {
    Simple(Connectable),
    Forward(Vec<Connectable>),
    Backward(Vec<Connectable>),
}

/// Parse a pass statement.
fn pass_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Pass).map_with(|_, e| {
        let span = e.span();
        Statement::Pass(PassStmt {
            span: to_span(span),
        })
    })
}

/// Parse a string statement (docstring).
fn string_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    string_literal().map(|s| {
        Statement::StringStmt(StringStmt {
            span: s.span,
            value: s,
        })
    })
}

/// Parse an assert statement.
fn assert_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Assert)
        .ignore_then(comparison())
        .map_with(|comp, e| {
            let span = e.span();
            Statement::Assert(AssertStmt {
                comparison: comp,
                span: to_span(span),
            })
        })
}

/// Parse a template argument.
fn template_arg<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, TemplateArg, Err<'a>> + Clone {
    identifier()
        .then_ignore(tok(TokenKind::Assign))
        .then(literal())
        .map_with(|(name, value), e| {
            let span = e.span();
            TemplateArg {
                name,
                value,
                span: to_span(span),
            }
        })
}

/// Parse a template (<arg=value, ...>).
fn template<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Template, Err<'a>> + Clone {
    template_arg()
        .separated_by(tok(TokenKind::Comma))
        .allow_trailing()
        .collect::<Vec<_>>()
        .delimited_by(tok(TokenKind::LessThan), tok(TokenKind::GreaterThan))
        .map_with(|args, e| {
            let span = e.span();
            Template {
                args,
                span: to_span(span),
            }
        })
}

/// Parse a trait statement.
fn trait_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::Trait)
        .ignore_then(type_reference())
        .then(
            tok(TokenKind::DoubleColon)
                .ignore_then(identifier())
                .or_not(),
        )
        .then(template().or_not())
        .map_with(|((type_ref, constructor), template), e| {
            let span = e.span();
            Statement::Trait(TraitStmt {
                target: None,
                type_ref,
                constructor,
                template,
                span: to_span(span),
            })
        })
}

/// Parse a new expression.
fn new_expr<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, NewExpr, Err<'a>> + Clone {
    tok(TokenKind::New)
        .ignore_then(type_reference())
        .then(
            number_literal()
                .delimited_by(tok(TokenKind::OpenBracket), tok(TokenKind::CloseBracket))
                .or_not(),
        )
        .then(template().or_not())
        .map_with(|((type_ref, count), template), e| {
            let span = e.span();
            NewExpr {
                type_ref,
                count,
                template,
                span: to_span(span),
            }
        })
}

/// Parse an assignable value.
fn assignable<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Assignable, Err<'a>> + Clone {
    choice((
        string_literal().map(Assignable::String),
        new_expr().map(Assignable::New),
        bool_literal().map(Assignable::Boolean),
        // arithmetic_expression() handles physical literals as atoms, so it can
        // also parse expressions like `300 * y` or `voltage * 1.5`. We extract
        // standalone physical literals into Assignable::Physical for downstream
        // lowering that distinguishes Range/Bilateral from simple quantities.
        arithmetic_expression().map(|expr| {
            if let Expression::Literal(Literal::Physical(phys)) = expr {
                Assignable::Physical(phys)
            } else {
                Assignable::Arithmetic(expr)
            }
        }),
    ))
}

/// Parse a statement starting with a name (assignment, connection, etc.).
fn name_starting_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    field_reference()
        .then(choice((
            // Declaration with optional assignment: field: type (= value)?
            tok(TokenKind::Colon)
                .ignore_then(identifier())
                .then(tok(TokenKind::Assign).ignore_then(assignable()).or_not())
                .map(|(type_info, value)| NameSuffix::Declaration { type_info, value }),
            // Assignment: field = value
            tok(TokenKind::Assign)
                .ignore_then(assignable())
                .map(NameSuffix::Assign),
            // Cumulative assignment: field += value or field -= value
            choice((
                tok(TokenKind::AddAssign).to(CumOperator::Add),
                tok(TokenKind::SubAssign).to(CumOperator::Sub),
            ))
            .then(arithmetic_expression())
            .map(|(op, value)| NameSuffix::CumAssign { op, value }),
            // Set assignment: field |= value or field &= value
            choice((
                tok(TokenKind::OrAssign).to(SetOperator::Or),
                tok(TokenKind::AndAssign).to(SetOperator::And),
            ))
            .then(arithmetic_expression())
            .map(|(op, value)| NameSuffix::SetAssign { op, value }),
            // Connection: field ~ connectable
            tok(TokenKind::Wire)
                .ignore_then(connectable())
                .map(NameSuffix::Connect),
            // Directed connection forward: field ~> connectable ~> ...
            tok(TokenKind::Sperm)
                .ignore_then(connectable())
                .then(
                    tok(TokenKind::Sperm)
                        .ignore_then(connectable())
                        .repeated()
                        .collect::<Vec<_>>(),
                )
                .map(|(first, rest)| NameSuffix::ForwardConnect { first, rest }),
            // Directed connection backward: field <~ connectable <~ ...
            tok(TokenKind::LSperm)
                .ignore_then(connectable())
                .then(
                    tok(TokenKind::LSperm)
                        .ignore_then(connectable())
                        .repeated()
                        .collect::<Vec<_>>(),
                )
                .map(|(first, rest)| NameSuffix::BackwardConnect { first, rest }),
            // Retype: field -> Type
            tok(TokenKind::Arrow)
                .ignore_then(type_reference())
                .map(NameSuffix::Retype),
        )))
        .map_with(|(field_ref, suffix), e| {
            let span = e.span();
            match suffix {
                NameSuffix::Declaration { type_info, value } => {
                    let decl = Declaration {
                        field: field_ref.clone(),
                        type_info,
                        span: to_span(span),
                    };
                    match value {
                        Some(v) => Statement::Assignment(Assignment {
                            target: AssignTarget::Declaration(decl),
                            value: v,
                            span: to_span(span),
                        }),
                        None => Statement::Declaration(decl),
                    }
                }
                NameSuffix::Assign(value) => Statement::Assignment(Assignment {
                    target: AssignTarget::FieldRef(field_ref),
                    value,
                    span: to_span(span),
                }),
                NameSuffix::CumAssign { op, value } => Statement::CumAssignment(CumAssignment {
                    target: AssignTarget::FieldRef(field_ref),
                    operator: op,
                    value,
                    span: to_span(span),
                }),
                NameSuffix::SetAssign { op, value } => Statement::SetAssignment(SetAssignment {
                    target: AssignTarget::FieldRef(field_ref),
                    operator: op,
                    value,
                    span: to_span(span),
                }),
                NameSuffix::Connect(right) => Statement::Connection(Connection {
                    left: Connectable::FieldRef(field_ref),
                    right,
                    span: to_span(span),
                }),
                NameSuffix::ForwardConnect { first, rest } => {
                    let mut elements = vec![Connectable::FieldRef(field_ref), first];
                    elements.extend(rest);
                    Statement::DirectedConnection(DirectedConnection {
                        direction: ConnectionDirection::Forward,
                        elements,
                        span: to_span(span),
                    })
                }
                NameSuffix::BackwardConnect { first, rest } => {
                    let mut elements = vec![Connectable::FieldRef(field_ref), first];
                    elements.extend(rest);
                    Statement::DirectedConnection(DirectedConnection {
                        direction: ConnectionDirection::Backward,
                        elements,
                        span: to_span(span),
                    })
                }
                NameSuffix::Retype(new_type) => Statement::Retype(Retype {
                    field: field_ref,
                    new_type,
                    span: to_span(span),
                }),
            }
        })
}

enum NameSuffix {
    Declaration {
        type_info: Identifier,
        value: Option<Assignable>,
    },
    Assign(Assignable),
    CumAssign {
        op: CumOperator,
        value: Expression,
    },
    SetAssign {
        op: SetOperator,
        value: Expression,
    },
    Connect(Connectable),
    ForwardConnect {
        first: Connectable,
        rest: Vec<Connectable>,
    },
    BackwardConnect {
        first: Connectable,
        rest: Vec<Connectable>,
    },
    Retype(TypeRef),
}

/// Parse a slice ([start:stop:step]).
fn slice<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Slice, Err<'a>> + Clone {
    choice((
        // [::step]
        tok(TokenKind::DoubleColon)
            .ignore_then(number_literal().or_not())
            .map(|step| (None, None, step)),
        // [start:stop:step] or [start:stop] or [:stop] etc
        number_literal()
            .or_not()
            .then(
                tok(TokenKind::Colon)
                    .ignore_then(number_literal().or_not())
                    .then(
                        tok(TokenKind::Colon)
                            .ignore_then(number_literal().or_not())
                            .or_not(),
                    )
                    .or_not(),
            )
            .map(|(start, rest)| match rest {
                Some((stop, step)) => (start, stop, step.flatten()),
                None => (start, None, None),
            }),
    ))
    .delimited_by(tok(TokenKind::OpenBracket), tok(TokenKind::CloseBracket))
    .map_with(|(start, stop, step), e| {
        let span = e.span();
        Slice {
            start,
            stop,
            step,
            span: to_span(span),
        }
    })
}

/// Parse an iterable for a for loop.
fn iterable<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Iterable, Err<'a>> + Clone {
    choice((
        // List literal: [field1, field2, ...]
        field_reference()
            .separated_by(tok(TokenKind::Comma))
            .allow_trailing()
            .collect::<Vec<_>>()
            .delimited_by(tok(TokenKind::OpenBracket), tok(TokenKind::CloseBracket))
            .map(Iterable::List),
        // Field reference with optional slice
        field_reference()
            .then(slice().or_not())
            .map(|(field, slice)| Iterable::FieldRef { field, slice }),
    ))
}

/// Parse a simple statement.
fn simple_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    choice((
        import_stmt(),
        from_import_stmt(),
        pin_declaration(),
        signal_def(),
        assert_stmt(),
        pass_stmt(),
        trait_stmt(),
        string_stmt(),
        name_starting_stmt(),
    ))
}

/// Parse semicolon-separated statements on one line.
fn simple_stmts<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Vec<Statement>, Err<'a>> + Clone {
    simple_stmt()
        .separated_by(tok(TokenKind::Semicolon))
        .at_least(1)
        .collect::<Vec<_>>()
        .then_ignore(tok(TokenKind::Semicolon).or_not())
        .then_ignore(tok(TokenKind::Newline).or_not())
}

/// Parse a block (indented or single-line).
fn block<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Vec<Statement>, Err<'a>> + Clone {
    recursive(|block| {
        choice((
            // Single-line block
            simple_stmts(),
            // Multi-line block
            tok(TokenKind::Newline)
                .ignore_then(tok(TokenKind::Indent))
                .ignore_then(
                    choice((
                        pragma_stmt().map(|s| vec![s]),
                        block_def(block.clone()).map(|s| vec![s]),
                        for_stmt(block.clone()).map(|s| vec![s]),
                        simple_stmts(),
                    ))
                    .then_ignore(tok(TokenKind::Newline).repeated().collect::<Vec<_>>())
                    .repeated()
                    .at_least(1)
                    .collect::<Vec<Vec<Statement>>>()
                    .map(|vecs| vecs.into_iter().flatten().collect::<Vec<_>>()),
                )
                .then_ignore(tok(TokenKind::Dedent)),
        ))
    })
}

/// Parse a block definition (module, component, interface).
fn block_def<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>(
    block: impl Parser<'a, I, Vec<Statement>, Err<'a>> + Clone,
) -> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    let block_kind = choice((
        tok(TokenKind::Module).to(BlockKind::Module),
        tok(TokenKind::Component).to(BlockKind::Component),
        tok(TokenKind::Interface).to(BlockKind::Interface),
    ));

    block_kind
        .then(identifier())
        .then(tok(TokenKind::From).ignore_then(type_reference()).or_not())
        .then_ignore(tok(TokenKind::Colon))
        .then(block)
        .map_with(|(((kind, name), super_type), body), e| {
            let span = e.span();
            Statement::BlockDef(BlockDef {
                kind,
                name,
                super_type,
                body,
                span: to_span(span),
            })
        })
}

/// Parse a for statement.
fn for_stmt<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>(
    block: impl Parser<'a, I, Vec<Statement>, Err<'a>> + Clone,
) -> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    tok(TokenKind::For)
        .ignore_then(identifier())
        .then_ignore(tok(TokenKind::In))
        .then(iterable())
        .then_ignore(tok(TokenKind::Colon))
        .then(block)
        .map_with(|((variable, iterable), body), e| {
            let span = e.span();
            Statement::For(ForStmt {
                variable,
                iterable,
                body,
                span: to_span(span),
            })
        })
}

/// Parse a top-level statement with error recovery.
fn statement<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, Statement, Err<'a>> + Clone {
    recursive(|_| {
        let blk = block();
        choice((
            pragma_stmt(),
            block_def(blk.clone()),
            for_stmt(blk),
            simple_stmt(),
        ))
    })
}

/// Parse an entire file.
///
/// This parser parses top-level statements and collects parse errors.
/// For error recovery, use parse_recovery() which will continue after errors.
pub fn file_parser<'a, I: ValueInput<'a, Token = Token, Span = SimpleSpan>>()
-> impl Parser<'a, I, File, Err<'a>> {
    tok(TokenKind::Newline)
        .repeated()
        .collect::<Vec<_>>()
        .ignore_then(
            statement()
                .then_ignore(tok(TokenKind::Newline).repeated().collect::<Vec<_>>())
                .repeated()
                .collect::<Vec<_>>(),
        )
        .then_ignore(tok(TokenKind::Eof).or_not())
        .then_ignore(end())
        .map_with(|statements, e| {
            let span = e.span();
            File {
                statements,
                span: to_span(span),
            }
        })
}

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

    fn parse_file(source: &str) -> (Option<File>, usize) {
        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 (ast, errors) = file_parser().parse(input).into_output_errors();
        (ast, errors.len())
    }

    #[test]
    fn test_parse_empty() {
        let (ast, error_count) = parse_file("");
        assert_eq!(error_count, 0);
        assert!(ast.is_some());
        assert!(ast.unwrap().statements.is_empty());
    }

    #[test]
    fn test_parse_pragma() {
        let (ast, error_count) = parse_file("#pragma experiment(\"FOR_LOOP\")\n");
        assert_eq!(error_count, 0);
        assert!(ast.is_some());
    }

    #[test]
    fn test_parse_module() {
        let (ast, error_count) = parse_file("module M:\n    pass\n");
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        assert_eq!(file.statements.len(), 1);
    }

    #[test]
    fn test_parse_nested_block() {
        let (ast, error_count) = parse_file("module A:\n    module B:\n        pass\n");
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(outer) = &file.statements[0] {
            assert_eq!(outer.name.name, "A");
            if let Statement::BlockDef(inner) = &outer.body[0] {
                assert_eq!(inner.name.name, "B");
            } else {
                panic!("Expected inner block");
            }
        } else {
            panic!("Expected outer block");
        }
    }

    #[test]
    fn test_keyword_in_as_pin_name() {
        let source = "module M:\n    pin in\n";
        let (ast, error_count) = parse_file(source);
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(block) = &file.statements[0] {
            if let Statement::PinDeclaration(pin) = &block.body[0] {
                if let PinName::Identifier(id) = &pin.name {
                    assert_eq!(id.name, "in");
                } else {
                    panic!("Expected identifier pin name");
                }
            } else {
                panic!("Expected pin declaration, got {:?}", block.body[0]);
            }
        } else {
            panic!("Expected block def");
        }
    }

    #[test]
    fn test_keyword_in_as_signal_name() {
        let source = "module M:\n    signal in\n";
        let (ast, error_count) = parse_file(source);
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(block) = &file.statements[0] {
            if let Statement::SignalDef(sig) = &block.body[0] {
                assert_eq!(sig.name.name, "in");
            } else {
                panic!("Expected signal def, got {:?}", block.body[0]);
            }
        } else {
            panic!("Expected block def");
        }
    }

    #[test]
    fn test_keyword_in_field_reference() {
        let source = "module M:\n    x.in ~ y.to\n";
        let (ast, error_count) = parse_file(source);
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(block) = &file.statements[0] {
            if let Statement::Connection(conn) = &block.body[0] {
                if let Connectable::FieldRef(fr) = &conn.left {
                    assert_eq!(fr.parts.len(), 2);
                    assert_eq!(fr.parts[0].name.name, "x");
                    assert_eq!(fr.parts[1].name.name, "in");
                } else {
                    panic!("Expected field ref");
                }
                if let Connectable::FieldRef(fr) = &conn.right {
                    assert_eq!(fr.parts.len(), 2);
                    assert_eq!(fr.parts[0].name.name, "y");
                    assert_eq!(fr.parts[1].name.name, "to");
                } else {
                    panic!("Expected field ref");
                }
            } else {
                panic!("Expected connection, got {:?}", block.body[0]);
            }
        } else {
            panic!("Expected block def");
        }
    }

    #[test]
    fn test_for_loop_still_works_with_keyword_identifiers() {
        let source =
            "module M:\n    container = new M[3]\n    for item in container:\n        pass\n";
        let (ast, error_count) = parse_file(source);
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(block) = &file.statements[0] {
            // Second statement should be the for loop
            let for_found = block.body.iter().any(|s| matches!(s, Statement::For(_)));
            assert!(for_found, "Should have a for loop. Body: {:?}", block.body);
        } else {
            panic!("Expected block def");
        }
    }

    #[test]
    fn test_multiple_keywords_as_pin_names() {
        let source = "module M:\n    pin to\n    pin from\n    pin is\n    pin within\n";
        let (ast, error_count) = parse_file(source);
        assert_eq!(error_count, 0);
        let file = ast.unwrap();
        if let Statement::BlockDef(block) = &file.statements[0] {
            assert_eq!(block.body.len(), 4);
            let names: Vec<String> = block
                .body
                .iter()
                .filter_map(|s| {
                    if let Statement::PinDeclaration(p) = s
                        && let PinName::Identifier(id) = &p.name
                    {
                        return Some(id.name.clone());
                    }
                    None
                })
                .collect();
            assert_eq!(names, vec!["to", "from", "is", "within"]);
        } else {
            panic!("Expected block def");
        }
    }
}