compliance-scanner-agent/compliance-agent/src/pipeline/plc/parser.rs

//! Recursive-descent parser for the security-relevant subset of Structured Text.
//!
//! Tolerant by design: it parses the POUs, variable sections, and statement
//! bodies it understands, and skips (with statement/POU-level recovery) anything
//! it does not, so a single odd construct never sinks the whole file.

use super::ast::*;
use super::lexer::{Keyword as K, Tok, Token};

pub struct Parser {
    toks: Vec<Token>,
    pos: usize,
}

impl Parser {
    pub fn new(toks: Vec<Token>) -> Self {
        Self { toks, pos: 0 }
    }

    // ── token helpers ──────────────────────────────────────────────
    fn peek(&self) -> &Tok {
        &self.toks[self.pos.min(self.toks.len() - 1)].kind
    }
    fn line(&self) -> u32 {
        self.toks[self.pos.min(self.toks.len() - 1)].line
    }
    fn at_end(&self) -> bool {
        matches!(self.peek(), Tok::Eof)
    }
    fn advance(&mut self) -> Tok {
        let t = self.toks[self.pos.min(self.toks.len() - 1)].kind.clone();
        if self.pos < self.toks.len() - 1 {
            self.pos += 1;
        }
        t
    }
    fn eat(&mut self, t: &Tok) -> bool {
        if self.peek() == t {
            self.advance();
            true
        } else {
            false
        }
    }
    fn eat_kw(&mut self, k: K) -> bool {
        if matches!(self.peek(), Tok::Kw(x) if *x == k) {
            self.advance();
            true
        } else {
            false
        }
    }
    fn at_kw(&self, k: K) -> bool {
        matches!(self.peek(), Tok::Kw(x) if *x == k)
    }
    fn ident(&mut self) -> Option<String> {
        if let Tok::Ident(s) = self.peek() {
            let s = s.clone();
            self.advance();
            Some(s)
        } else {
            None
        }
    }

    // ── top level ──────────────────────────────────────────────────
    /// Parse every POU in the token stream.
    pub fn parse_units(&mut self) -> Vec<Pou> {
        let mut pous = Vec::new();
        while !self.at_end() {
            match self.peek() {
                Tok::Kw(K::Program) => {
                    self.advance();
                    if let Some(p) = self.parse_pou(PouKind::Program, K::EndProgram) {
                        pous.push(p);
                    }
                }
                Tok::Kw(K::Function) => {
                    self.advance();
                    if let Some(p) = self.parse_pou(PouKind::Function, K::EndFunction) {
                        pous.push(p);
                    }
                }
                Tok::Kw(K::FunctionBlock) => {
                    self.advance();
                    if let Some(p) = self.parse_pou(PouKind::FunctionBlock, K::EndFunctionBlock) {
                        pous.push(p);
                    }
                }
                // Skip TYPE...END_TYPE and anything else at top level.
                _ => {
                    self.advance();
                }
            }
        }
        pous
    }

    fn parse_pou(&mut self, kind: PouKind, end: K) -> Option<Pou> {
        let line = self.line();
        let name = self.ident().unwrap_or_else(|| "<anonymous>".to_string());
        // Optional `: return_type` for functions.
        if self.eat(&Tok::Colon) {
            let _ = self.advance(); // return type token
        }

        let mut vars = Vec::new();
        // Variable sections precede the body.
        while let Some(section) = self.var_section_kw() {
            self.advance();
            let _ = self.eat_kw(K::Constant); // CONSTANT is informational for our rules
            self.parse_var_decls(section, &mut vars);
        }

        // Body statements until END_<kind>.
        let mut body = Vec::new();
        while !self.at_end() && !self.at_kw(end) {
            if let Some(s) = self.parse_stmt() {
                body.push(s);
            }
        }
        self.eat_kw(end);

        Some(Pou {
            name,
            kind,
            vars,
            body,
            line,
        })
    }

    fn var_section_kw(&self) -> Option<VarSection> {
        match self.peek() {
            Tok::Kw(K::Var) => Some(VarSection::Var),
            Tok::Kw(K::VarInput) => Some(VarSection::Input),
            Tok::Kw(K::VarOutput) => Some(VarSection::Output),
            Tok::Kw(K::VarInOut) => Some(VarSection::InOut),
            Tok::Kw(K::VarGlobal) => Some(VarSection::Global),
            Tok::Kw(K::VarTemp) => Some(VarSection::Temp),
            Tok::Kw(K::VarExternal) => Some(VarSection::External),
            _ => None,
        }
    }

    fn parse_var_decls(&mut self, section: VarSection, out: &mut Vec<VarDecl>) {
        while !self.at_end() && !self.at_kw(K::EndVar) {
            let line = self.line();
            // names: a, b, c
            let mut names = Vec::new();
            match self.ident() {
                Some(n) => names.push(n),
                None => {
                    // Not a declaration we understand — skip to next ; or END_VAR.
                    self.sync_decl();
                    continue;
                }
            }
            while self.eat(&Tok::Comma) {
                if let Some(n) = self.ident() {
                    names.push(n);
                }
            }
            if !self.eat(&Tok::Colon) {
                self.sync_decl();
                continue;
            }
            let (type_name, array_bounds) = self.parse_type();
            let init = if self.eat(&Tok::Assign) {
                Some(self.parse_expr())
            } else {
                None
            };
            self.eat(&Tok::Semi);
            for n in names {
                out.push(VarDecl {
                    name: n,
                    section,
                    type_name: type_name.clone(),
                    array_bounds,
                    init: init.clone(),
                    line,
                });
            }
        }
        self.eat_kw(K::EndVar);
    }

    /// Parse a (possibly ARRAY) type, returning its rendered name and literal
    /// bounds when present.
    fn parse_type(&mut self) -> (String, Option<(i64, i64)>) {
        if self.eat_kw(K::Array) {
            let mut bounds = None;
            if self.eat(&Tok::LBrack) {
                let lo = self.int_lit();
                self.eat(&Tok::DotDot);
                let hi = self.int_lit();
                if let (Some(lo), Some(hi)) = (lo, hi) {
                    bounds = Some((lo, hi));
                }
                // Skip any further dimensions / tokens to the closing bracket.
                while !self.at_end() && !self.eat(&Tok::RBrack) {
                    self.advance();
                }
            }
            self.eat_kw(K::Of);
            let elem = self.type_ident();
            (format!("ARRAY OF {elem}"), bounds)
        } else {
            (self.type_ident(), None)
        }
    }

    fn type_ident(&mut self) -> String {
        // Types can be qualified idents; keep it simple: one token, plus any
        // string-length suffix like STRING[80].
        let base = match self.advance() {
            Tok::Ident(s) => s,
            Tok::Kw(_) => "TYPE".to_string(),
            other => format!("{other:?}"),
        };
        if self.eat(&Tok::LBrack) {
            while !self.at_end() && !self.eat(&Tok::RBrack) {
                self.advance();
            }
        }
        base
    }

    fn int_lit(&mut self) -> Option<i64> {
        match self.peek() {
            Tok::Int(n) => {
                let n = *n;
                self.advance();
                Some(n)
            }
            Tok::Minus => {
                self.advance();
                if let Tok::Int(n) = self.peek() {
                    let n = -*n;
                    self.advance();
                    Some(n)
                } else {
                    None
                }
            }
            _ => None,
        }
    }

    fn sync_decl(&mut self) {
        while !self.at_end() && !self.eat(&Tok::Semi) && !self.at_kw(K::EndVar) {
            self.advance();
        }
    }
    fn sync_stmt(&mut self) {
        while !self.at_end() && !self.eat(&Tok::Semi) {
            // Stop at block terminators so recovery doesn't swallow structure.
            if matches!(
                self.peek(),
                Tok::Kw(
                    K::EndIf
                        | K::EndFor
                        | K::EndWhile
                        | K::EndCase
                        | K::EndRepeat
                        | K::EndProgram
                        | K::EndFunction
                        | K::EndFunctionBlock
                        | K::Else
                        | K::Elsif
                )
            ) {
                return;
            }
            self.advance();
        }
    }

    // ── statements ─────────────────────────────────────────────────
    fn parse_stmt(&mut self) -> Option<Stmt> {
        let line = self.line();
        match self.peek().clone() {
            Tok::Semi => {
                self.advance();
                None
            }
            Tok::Kw(K::If) => self.parse_if(),
            Tok::Kw(K::Case) => self.parse_case(),
            Tok::Kw(K::For) => self.parse_for(),
            Tok::Kw(K::While) => self.parse_while(),
            Tok::Kw(K::Repeat) => self.parse_repeat(),
            Tok::Kw(K::Return) => {
                self.advance();
                self.eat(&Tok::Semi);
                Some(Stmt::Return { line })
            }
            Tok::Kw(K::Exit) => {
                self.advance();
                self.eat(&Tok::Semi);
                Some(Stmt::Exit { line })
            }
            Tok::Kw(K::Jmp) => {
                self.advance();
                let label = self.ident().unwrap_or_default();
                self.eat(&Tok::Semi);
                Some(Stmt::Jump { label, line })
            }
            Tok::Ident(name) => {
                // Could be `label:`, `call(...)`, or an assignment.
                // Lookahead: ident ':' (not ':=') → label.
                if matches!(
                    self.toks.get(self.pos + 1).map(|t| &t.kind),
                    Some(Tok::Colon)
                ) && !matches!(self.toks.get(self.pos + 2).map(|t| &t.kind), Some(Tok::Eq))
                {
                    self.advance(); // ident
                    self.advance(); // ':'
                    return Some(Stmt::Label { name, line });
                }
                let lhs = self.parse_expr();
                if self.eat(&Tok::Assign) {
                    let value = self.parse_expr();
                    self.eat(&Tok::Semi);
                    Some(Stmt::Assign {
                        target: lhs,
                        value,
                        line,
                    })
                } else if let Expr::Call { callee, args, .. } = lhs {
                    self.eat(&Tok::Semi);
                    Some(Stmt::Call { callee, args, line })
                } else {
                    // Bare expression / FB invocation without args recognized —
                    // skip to the terminator.
                    self.sync_stmt();
                    None
                }
            }
            _ => {
                self.sync_stmt();
                None
            }
        }
    }

    fn parse_block_until(&mut self, terms: &[K]) -> Vec<Stmt> {
        let mut body = Vec::new();
        while !self.at_end() && !terms.iter().any(|k| self.at_kw(*k)) {
            if let Some(s) = self.parse_stmt() {
                body.push(s);
            }
        }
        body
    }

    fn parse_if(&mut self) -> Option<Stmt> {
        let line = self.line();
        self.eat_kw(K::If);
        let mut branches = Vec::new();
        let cond = self.parse_expr();
        self.eat_kw(K::Then);
        let body = self.parse_block_until(&[K::Elsif, K::Else, K::EndIf]);
        branches.push((cond, body));
        while self.eat_kw(K::Elsif) {
            let c = self.parse_expr();
            self.eat_kw(K::Then);
            let b = self.parse_block_until(&[K::Elsif, K::Else, K::EndIf]);
            branches.push((c, b));
        }
        let else_body = if self.eat_kw(K::Else) {
            Some(self.parse_block_until(&[K::EndIf]))
        } else {
            None
        };
        self.eat_kw(K::EndIf);
        self.eat(&Tok::Semi);
        Some(Stmt::If {
            branches,
            else_body,
            line,
        })
    }

    fn parse_case(&mut self) -> Option<Stmt> {
        let line = self.line();
        self.eat_kw(K::Case);
        let selector = self.parse_expr();
        self.eat_kw(K::Of);
        let mut arms = Vec::new();
        let mut else_body = None;
        while !self.at_end() && !self.at_kw(K::EndCase) {
            if self.eat_kw(K::Else) {
                else_body = Some(self.parse_block_until(&[K::EndCase]));
                break;
            }
            // labels: expr {, expr} :
            let mut labels = vec![self.parse_expr()];
            while self.eat(&Tok::Comma) {
                labels.push(self.parse_expr());
            }
            self.eat(&Tok::Colon);
            let body = self.parse_block_until(&[K::EndCase, K::Else]);
            arms.push((labels, body));
        }
        self.eat_kw(K::EndCase);
        self.eat(&Tok::Semi);
        Some(Stmt::Case {
            selector,
            arms,
            else_body,
            line,
        })
    }

    fn parse_for(&mut self) -> Option<Stmt> {
        let line = self.line();
        self.eat_kw(K::For);
        let var = self.ident().unwrap_or_default();
        self.eat(&Tok::Assign);
        let from = self.parse_expr();
        self.eat_kw(K::To);
        let to = self.parse_expr();
        let by = if self.eat_kw(K::By) {
            Some(self.parse_expr())
        } else {
            None
        };
        self.eat_kw(K::Do);
        let body = self.parse_block_until(&[K::EndFor]);
        self.eat_kw(K::EndFor);
        self.eat(&Tok::Semi);
        Some(Stmt::For {
            var,
            from,
            to,
            by,
            body,
            line,
        })
    }

    fn parse_while(&mut self) -> Option<Stmt> {
        let line = self.line();
        self.eat_kw(K::While);
        let cond = self.parse_expr();
        self.eat_kw(K::Do);
        let body = self.parse_block_until(&[K::EndWhile]);
        self.eat_kw(K::EndWhile);
        self.eat(&Tok::Semi);
        Some(Stmt::While { cond, body, line })
    }

    fn parse_repeat(&mut self) -> Option<Stmt> {
        let line = self.line();
        self.eat_kw(K::Repeat);
        let body = self.parse_block_until(&[K::Until, K::EndRepeat]);
        self.eat_kw(K::Until);
        let until = self.parse_expr();
        self.eat_kw(K::EndRepeat);
        self.eat(&Tok::Semi);
        Some(Stmt::Repeat { body, until, line })
    }

    // ── expressions (precedence climbing) ──────────────────────────
    pub fn parse_expr(&mut self) -> Expr {
        self.parse_or()
    }

    fn parse_or(&mut self) -> Expr {
        let mut lhs = self.parse_and();
        loop {
            let op = match self.peek() {
                Tok::Kw(K::Or) => BinOp::Or,
                Tok::Kw(K::Xor) => BinOp::Xor,
                _ => break,
            };
            let line = self.line();
            self.advance();
            let rhs = self.parse_and();
            lhs = Expr::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                line,
            };
        }
        lhs
    }

    fn parse_and(&mut self) -> Expr {
        let mut lhs = self.parse_cmp();
        while matches!(self.peek(), Tok::Kw(K::And) | Tok::Amp) {
            let op = BinOp::And;
            let line = self.line();
            self.advance();
            let rhs = self.parse_cmp();
            lhs = Expr::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                line,
            };
        }
        lhs
    }

    fn parse_cmp(&mut self) -> Expr {
        let mut lhs = self.parse_add();
        loop {
            let op = match self.peek() {
                Tok::Eq => BinOp::Eq,
                Tok::Ne => BinOp::Ne,
                Tok::Lt => BinOp::Lt,
                Tok::Le => BinOp::Le,
                Tok::Gt => BinOp::Gt,
                Tok::Ge => BinOp::Ge,
                _ => break,
            };
            let line = self.line();
            self.advance();
            let rhs = self.parse_add();
            lhs = Expr::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                line,
            };
        }
        lhs
    }

    fn parse_add(&mut self) -> Expr {
        let mut lhs = self.parse_mul();
        loop {
            let op = match self.peek() {
                Tok::Plus => BinOp::Add,
                Tok::Minus => BinOp::Sub,
                _ => break,
            };
            let line = self.line();
            self.advance();
            let rhs = self.parse_mul();
            lhs = Expr::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                line,
            };
        }
        lhs
    }

    fn parse_mul(&mut self) -> Expr {
        let mut lhs = self.parse_unary();
        loop {
            let op = match self.peek() {
                Tok::Star => BinOp::Mul,
                Tok::Slash => BinOp::Div,
                Tok::Kw(K::Mod) => BinOp::Mod,
                Tok::Power => BinOp::Pow,
                _ => break,
            };
            let line = self.line();
            self.advance();
            let rhs = self.parse_unary();
            lhs = Expr::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                line,
            };
        }
        lhs
    }

    fn parse_unary(&mut self) -> Expr {
        let line = self.line();
        match self.peek() {
            Tok::Kw(K::Not) => {
                self.advance();
                Expr::Unary {
                    op: UnOp::Not,
                    expr: Box::new(self.parse_unary()),
                    line,
                }
            }
            Tok::Minus => {
                self.advance();
                Expr::Unary {
                    op: UnOp::Neg,
                    expr: Box::new(self.parse_unary()),
                    line,
                }
            }
            _ => self.parse_postfix(),
        }
    }

    fn parse_postfix(&mut self) -> Expr {
        let mut e = self.parse_primary();
        loop {
            let line = self.line();
            match self.peek() {
                Tok::LBrack => {
                    self.advance();
                    let index = self.parse_expr();
                    self.eat(&Tok::RBrack);
                    e = Expr::Index {
                        base: Box::new(e),
                        index: Box::new(index),
                        line,
                    };
                }
                Tok::Dot => {
                    self.advance();
                    let field = self.ident().unwrap_or_default();
                    e = Expr::Member {
                        base: Box::new(e),
                        field,
                        line,
                    };
                }
                _ => break,
            }
        }
        e
    }

    fn parse_primary(&mut self) -> Expr {
        let line = self.line();
        match self.advance() {
            Tok::Int(n) => Expr::Int(n, line),
            Tok::Real(r) => Expr::Real(r, line),
            Tok::Bool(b) => Expr::Bool(b, line),
            Tok::Str(s) => Expr::Str(s, line),
            Tok::Time(t) => Expr::Time(t, line),
            Tok::LParen => {
                let e = self.parse_expr();
                self.eat(&Tok::RParen);
                e
            }
            Tok::Ident(name) => {
                if self.eat(&Tok::LParen) {
                    let args = self.parse_call_args();
                    Expr::Call {
                        callee: name,
                        args,
                        line,
                    }
                } else {
                    Expr::Ident(name, line)
                }
            }
            // Unrecognized start of expression — yield a placeholder identifier.
            _ => Expr::Ident(String::new(), line),
        }
    }

    fn parse_call_args(&mut self) -> Vec<CallArg> {
        let mut args = Vec::new();
        if self.eat(&Tok::RParen) {
            return args;
        }
        loop {
            // Named arg: ident := expr  (peek two tokens).
            if let Tok::Ident(name) = self.peek().clone() {
                if matches!(
                    self.toks.get(self.pos + 1).map(|t| &t.kind),
                    Some(Tok::Assign)
                ) {
                    self.advance(); // ident
                    self.advance(); // :=
                    let value = self.parse_expr();
                    args.push(CallArg {
                        name: Some(name),
                        value,
                    });
                    if self.eat(&Tok::Comma) {
                        continue;
                    }
                    break;
                }
            }
            let value = self.parse_expr();
            args.push(CallArg { name: None, value });
            if self.eat(&Tok::Comma) {
                continue;
            }
            break;
        }
        self.eat(&Tok::RParen);
        args
    }
}

/// Parse ST source into its POUs.
pub fn parse(src: &str) -> Vec<Pou> {
    let toks = super::lexer::lex(src);
    Parser::new(toks).parse_units()
}

#[cfg(test)]
mod tests {
    use super::*;

    const SAMPLE: &str = r#"
PROGRAM Main
VAR
    idx : INT;
    pw : STRING := 'admin123';
    buf : ARRAY[0..9] OF INT;
    ok : BOOL := FALSE;
END_VAR
    // a comment
    IF idx > 0 THEN
        buf[idx] := idx * 2;
    ELSE
        JMP done;
    END_IF;
    Comm(IP := '10.0.0.1', PORT := 502);
done:
    ok := TRUE;
END_PROGRAM
"#;

    #[test]
    fn parses_program_vars_and_body() {
        let pous = parse(SAMPLE);
        assert_eq!(pous.len(), 1, "one POU");
        let p = &pous[0];
        assert_eq!(p.name, "Main");
        assert_eq!(p.kind, PouKind::Program);
        // vars: idx, pw, buf, ok
        assert_eq!(p.vars.len(), 4);
        let pw = p.vars.iter().find(|v| v.name == "pw").expect("pw");
        assert!(matches!(&pw.init, Some(Expr::Str(s, _)) if s == "admin123"));
        let buf = p.vars.iter().find(|v| v.name == "buf").expect("buf");
        assert_eq!(buf.array_bounds, Some((0, 9)));
        // body has an IF, a Call, a Label, and an Assign
        assert!(p.body.iter().any(|s| matches!(s, Stmt::If { .. })));
        assert!(p
            .body
            .iter()
            .any(|s| matches!(s, Stmt::Call { callee, .. } if callee == "Comm")));
        assert!(p
            .body
            .iter()
            .any(|s| matches!(s, Stmt::Label { name, .. } if name == "done")));
    }

    #[test]
    fn jmp_inside_if_is_captured() {
        let pous = parse(SAMPLE);
        let p = &pous[0];
        // find the IF, check its else branch has a JMP
        let has_jmp = p.body.iter().any(|s| match s {
            Stmt::If { else_body, .. } => else_body
                .as_ref()
                .map(|b| b.iter().any(|s| matches!(s, Stmt::Jump { .. })))
                .unwrap_or(false),
            _ => false,
        });
        assert!(has_jmp, "JMP should be parsed inside the ELSE branch");
    }
}