spec-checker 0.9.4

//! Spec DSL parser. Hand-rolled recursive-descent.

use crate::ast::*;
use lex_syntax::token::{lex, Token, TokenKind};

#[derive(Debug, thiserror::Error)]
#[error("spec parse error at byte {pos}: {msg}")]
pub struct SpecParseError {
    pub pos: usize,
    pub msg: String,
}

pub fn parse_spec(src: &str) -> Result<Spec, SpecParseError> {
    let toks = lex(src).map_err(|e| SpecParseError {
        pos: e.span.start, msg: format!("lex: {}", e.snippet),
    })?;
    let mut p = Parser { toks, idx: 0 };
    let spec = p.parse_spec()?;
    p.skip_newlines();
    if !p.at_eof() {
        return Err(p.err("trailing input after spec"));
    }
    Ok(spec)
}

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

impl Parser {
    fn at_eof(&self) -> bool { self.idx >= self.toks.len() }

    fn peek(&self) -> Option<&TokenKind> { self.toks.get(self.idx).map(|t| &t.kind) }

    fn bump(&mut self) -> Option<Token> {
        let t = self.toks.get(self.idx).cloned();
        if t.is_some() { self.idx += 1; }
        t
    }

    fn pos(&self) -> usize {
        self.toks.get(self.idx).map(|t| t.span.start).unwrap_or(0)
    }

    fn err(&self, m: impl Into<String>) -> SpecParseError {
        SpecParseError { pos: self.pos(), msg: m.into() }
    }

    fn skip_newlines(&mut self) {
        while matches!(self.peek(), Some(TokenKind::Newline) | Some(TokenKind::Semi)) {
            self.idx += 1;
        }
    }

    fn eat(&mut self, k: &TokenKind) -> bool {
        self.skip_newlines();
        if let Some(cur) = self.peek() {
            if std::mem::discriminant(cur) == std::mem::discriminant(k) {
                self.bump();
                return true;
            }
        }
        false
    }

    fn expect_ident(&mut self, ctx: &str) -> Result<String, SpecParseError> {
        self.skip_newlines();
        match self.peek() {
            Some(TokenKind::Ident(_)) => match self.bump().unwrap().kind {
                TokenKind::Ident(n) => Ok(n), _ => unreachable!(),
            },
            other => Err(self.err(format!("expected ident {ctx}, got {other:?}"))),
        }
    }

    fn expect_keyword(&mut self, kw: &str) -> Result<(), SpecParseError> {
        // We don't have a `spec`/`forall`/`where` keyword in the lexer;
        // they come through as Ident.
        let id = self.expect_ident(&format!("(keyword `{kw}`)"))?;
        if id != kw {
            return Err(self.err(format!("expected `{kw}`, got `{id}`")));
        }
        Ok(())
    }

    fn parse_spec(&mut self) -> Result<Spec, SpecParseError> {
        // `spec <name> { forall ... : body }`
        self.expect_keyword("spec")?;
        let name = self.expect_ident("for spec name")?;
        if !self.eat(&TokenKind::LBrace) {
            return Err(self.err("expected `{` after spec name"));
        }
        self.expect_keyword("forall")?;

        let mut quantifiers = Vec::new();
        loop {
            quantifiers.push(self.parse_quantifier()?);
            self.skip_newlines();
            if self.eat(&TokenKind::Comma) {
                continue;
            }
            break;
        }

        // Optional `where <constraint>` applies to the *last* quantifier.
        if let Some(TokenKind::Ident(n)) = self.peek() {
            if n == "where" {
                self.bump();
                let c = self.parse_expr()?;
                if let Some(last) = quantifiers.last_mut() {
                    last.constraint = Some(c);
                }
            }
        }

        // Body separator: `:` or `=>`. We accept either.
        if !self.eat(&TokenKind::Colon) && !self.eat(&TokenKind::FatArrow) {
            return Err(self.err("expected `:` or `=>` before spec body"));
        }

        let body = self.parse_body_block()?;
        if !self.eat(&TokenKind::RBrace) {
            return Err(self.err("expected `}` to close spec"));
        }
        Ok(Spec { name, quantifiers, body })
    }

    fn parse_quantifier(&mut self) -> Result<Quantifier, SpecParseError> {
        let name = self.expect_ident("for quantifier var")?;
        if !self.eat(&TokenKind::ColonColon) {
            return Err(self.err("expected `::` after quantifier var"));
        }
        let ty = self.parse_spec_type()?;
        Ok(Quantifier { name, ty, constraint: None })
    }

    /// Parse a `SpecType` (#208). Scalar idents (`Int`, `Float`, `Bool`,
    /// `Str`) plus structured `{ name :: Ty, ... }` for records.
    /// Records nest, so `{ s :: { used :: Int, ceiling :: Int } }` works.
    fn parse_spec_type(&mut self) -> Result<SpecType, SpecParseError> {
        if self.eat(&TokenKind::LBrace) {
            // Record type: `{ name :: Ty, name :: Ty, ... }`
            let mut fields: Vec<(String, SpecType)> = Vec::new();
            self.skip_newlines();
            if !matches!(self.peek(), Some(TokenKind::RBrace)) {
                loop {
                    let field_name = self.expect_ident("for record field name")?;
                    if !self.eat(&TokenKind::ColonColon) {
                        return Err(self.err("expected `::` after record field name"));
                    }
                    let field_ty = self.parse_spec_type()?;
                    fields.push((field_name, field_ty));
                    self.skip_newlines();
                    if !self.eat(&TokenKind::Comma) { break; }
                    self.skip_newlines();
                }
            }
            if !self.eat(&TokenKind::RBrace) {
                return Err(self.err("expected `}` to close record type"));
            }
            return Ok(SpecType::Record { fields });
        }
        let ty_name = self.expect_ident("for quantifier type")?;
        match ty_name.as_str() {
            "Int" => Ok(SpecType::Int),
            "Float" => Ok(SpecType::Float),
            "Bool" => Ok(SpecType::Bool),
            "Str" => Ok(SpecType::Str),
            "List" => {
                // `List[T]` — element type in brackets.
                if !self.eat(&TokenKind::LBracket) {
                    return Err(self.err("expected `[` after `List`"));
                }
                let elem = self.parse_spec_type()?;
                if !self.eat(&TokenKind::RBracket) {
                    return Err(self.err("expected `]` to close `List[...]`"));
                }
                Ok(SpecType::List { element: Box::new(elem) })
            }
            // #208 slice 3: any other capitalized ident is treated as
            // a user-defined ADT name (`Message`, `Order`, …). The
            // gate path uses the value's runtime variant tag for
            // discrimination, so no compile-time variant table is
            // needed; the random-input prover doesn't enumerate these.
            other if other.chars().next().is_some_and(|c| c.is_ascii_uppercase()) => {
                Ok(SpecType::Named { name: other.to_string() })
            }
            other => Err(self.err(format!("unknown spec type `{other}`"))),
        }
    }

    /// The body may be a sequence of `let` bindings followed by a single
    /// expression, all inside the spec's outer `{ ... }`.
    fn parse_body_block(&mut self) -> Result<SpecExpr, SpecParseError> {
        // Optional sequence of lets, then the final expression.
        self.skip_newlines();
        let mut lets: Vec<(String, SpecExpr)> = Vec::new();
        while matches!(self.peek(), Some(TokenKind::Let)) {
            self.bump(); // 'let'
            let name = self.expect_ident("after `let`")?;
            // Allow `:: Type` annotation; ignored here.
            if self.eat(&TokenKind::ColonColon) {
                let _ = self.expect_ident("after `::`")?;
            }
            if !self.eat(&TokenKind::ColonEq) {
                return Err(self.err("expected `:=` in let"));
            }
            let value = self.parse_expr()?;
            lets.push((name, value));
            self.skip_newlines();
        }
        let mut body = self.parse_expr()?;
        // Wrap in nested Let from inside out.
        for (name, value) in lets.into_iter().rev() {
            body = SpecExpr::Let { name, value: Box::new(value), body: Box::new(body) };
        }
        Ok(body)
    }

    // ---- expressions ----

    fn parse_expr(&mut self) -> Result<SpecExpr, SpecParseError> {
        self.parse_or()
    }

    fn parse_or(&mut self) -> Result<SpecExpr, SpecParseError> {
        let mut lhs = self.parse_and()?;
        while self.peek_kw("or") {
            self.bump();
            let rhs = self.parse_and()?;
            lhs = SpecExpr::BinOp { op: SpecOp::Or, lhs: Box::new(lhs), rhs: Box::new(rhs) };
        }
        Ok(lhs)
    }

    fn parse_and(&mut self) -> Result<SpecExpr, SpecParseError> {
        let mut lhs = self.parse_cmp()?;
        while self.peek_kw("and") {
            self.bump();
            let rhs = self.parse_cmp()?;
            lhs = SpecExpr::BinOp { op: SpecOp::And, lhs: Box::new(lhs), rhs: Box::new(rhs) };
        }
        Ok(lhs)
    }

    fn parse_cmp(&mut self) -> Result<SpecExpr, SpecParseError> {
        let lhs = self.parse_add()?;
        let op = match self.peek() {
            Some(TokenKind::EqEq) => Some(SpecOp::Eq),
            Some(TokenKind::BangEq) => Some(SpecOp::Neq),
            Some(TokenKind::Lt) => Some(SpecOp::Lt),
            Some(TokenKind::LtEq) => Some(SpecOp::Le),
            Some(TokenKind::Gt) => Some(SpecOp::Gt),
            Some(TokenKind::GtEq) => Some(SpecOp::Ge),
            _ => None,
        };
        if let Some(op) = op {
            self.bump();
            let rhs = self.parse_add()?;
            return Ok(SpecExpr::BinOp { op, lhs: Box::new(lhs), rhs: Box::new(rhs) });
        }
        Ok(lhs)
    }

    fn parse_add(&mut self) -> Result<SpecExpr, SpecParseError> {
        let mut lhs = self.parse_mul()?;
        loop {
            let op = match self.peek() {
                Some(TokenKind::Plus) => SpecOp::Add,
                Some(TokenKind::Minus) => SpecOp::Sub,
                _ => break,
            };
            self.bump();
            let rhs = self.parse_mul()?;
            lhs = SpecExpr::BinOp { op, lhs: Box::new(lhs), rhs: Box::new(rhs) };
        }
        Ok(lhs)
    }

    fn parse_mul(&mut self) -> Result<SpecExpr, SpecParseError> {
        let mut lhs = self.parse_unary()?;
        loop {
            let op = match self.peek() {
                Some(TokenKind::Star) => SpecOp::Mul,
                Some(TokenKind::Slash) => SpecOp::Div,
                Some(TokenKind::Percent) => SpecOp::Mod,
                _ => break,
            };
            self.bump();
            let rhs = self.parse_unary()?;
            lhs = SpecExpr::BinOp { op, lhs: Box::new(lhs), rhs: Box::new(rhs) };
        }
        Ok(lhs)
    }

    fn parse_unary(&mut self) -> Result<SpecExpr, SpecParseError> {
        if self.peek_kw("not") {
            self.bump();
            let e = self.parse_unary()?;
            return Ok(SpecExpr::Not { expr: Box::new(e) });
        }
        self.parse_postfix()
    }

    fn parse_postfix(&mut self) -> Result<SpecExpr, SpecParseError> {
        let mut e = self.parse_primary()?;
        loop {
            match self.peek() {
                Some(TokenKind::LParen) => {
                    self.bump();
                    let mut args = Vec::new();
                    self.skip_newlines();
                    if !matches!(self.peek(), Some(TokenKind::RParen)) {
                        args.push(self.parse_expr()?);
                        while self.eat(&TokenKind::Comma) { args.push(self.parse_expr()?); }
                    }
                    if !self.eat(&TokenKind::RParen) {
                        return Err(self.err("expected `)` to close call"));
                    }
                    let func = match e {
                        SpecExpr::Var { name } => name,
                        other => return Err(self.err(format!("only ident-callable; got {other:?}"))),
                    };
                    e = SpecExpr::Call { func, args };
                }
                Some(TokenKind::Dot) => {
                    self.bump();
                    let field = self.expect_ident("after `.`")?;
                    e = SpecExpr::FieldAccess { value: Box::new(e), field };
                }
                Some(TokenKind::LBracket) => {
                    // `xs[i]` indexed access (#208). The expression
                    // inside brackets is parsed at full precedence so
                    // `xs[i + 1]`, `xs[length(xs) - 1]`, etc. work.
                    self.bump();
                    let index = self.parse_expr()?;
                    if !self.eat(&TokenKind::RBracket) {
                        return Err(self.err("expected `]` to close index"));
                    }
                    e = SpecExpr::Index { list: Box::new(e), index: Box::new(index) };
                }
                _ => break,
            }
        }
        Ok(e)
    }

    fn parse_primary(&mut self) -> Result<SpecExpr, SpecParseError> {
        self.skip_newlines();
        match self.peek() {
            Some(TokenKind::Int(_)) => match self.bump().unwrap().kind {
                TokenKind::Int(n) => Ok(SpecExpr::IntLit { value: n }), _ => unreachable!(),
            },
            Some(TokenKind::Float(_)) => match self.bump().unwrap().kind {
                TokenKind::Float(n) => Ok(SpecExpr::FloatLit { value: n }), _ => unreachable!(),
            },
            Some(TokenKind::Str(_)) => match self.bump().unwrap().kind {
                TokenKind::Str(s) => Ok(SpecExpr::StrLit { value: s }), _ => unreachable!(),
            },
            Some(TokenKind::True) => { self.bump(); Ok(SpecExpr::BoolLit { value: true }) }
            Some(TokenKind::False) => { self.bump(); Ok(SpecExpr::BoolLit { value: false }) }
            Some(TokenKind::Match) => self.parse_match(),
            Some(TokenKind::Ident(_)) => {
                let name = self.expect_ident("")?;
                Ok(SpecExpr::Var { name })
            }
            Some(TokenKind::LParen) => {
                self.bump();
                let e = self.parse_expr()?;
                if !self.eat(&TokenKind::RParen) {
                    return Err(self.err("expected `)`"));
                }
                Ok(e)
            }
            other => Err(self.err(format!("expected primary expression, got {other:?}"))),
        }
    }

    /// Parse `match scrutinee { Pattern => body, _ => body, ... }`
    /// (#208 slice 3). Arms are separated by `,`; trailing comma OK.
    fn parse_match(&mut self) -> Result<SpecExpr, SpecParseError> {
        self.bump(); // consume `match`
        let scrutinee = self.parse_expr()?;
        if !self.eat(&TokenKind::LBrace) {
            return Err(self.err("expected `{` after match scrutinee"));
        }
        let mut arms: Vec<crate::ast::MatchArm> = Vec::new();
        self.skip_newlines();
        while !matches!(self.peek(), Some(TokenKind::RBrace)) {
            let pattern = self.parse_pattern()?;
            if !self.eat(&TokenKind::FatArrow) {
                return Err(self.err("expected `=>` after match pattern"));
            }
            let body = self.parse_expr()?;
            arms.push(crate::ast::MatchArm { pattern, body });
            self.skip_newlines();
            if !self.eat(&TokenKind::Comma) { break; }
            self.skip_newlines();
        }
        if !self.eat(&TokenKind::RBrace) {
            return Err(self.err("expected `}` to close match"));
        }
        Ok(SpecExpr::Match {
            scrutinee: Box::new(scrutinee),
            arms,
        })
    }

    /// Parse a single pattern: `_`, `Variant`, `Variant(a, b, ...)`.
    fn parse_pattern(&mut self) -> Result<crate::ast::SpecPattern, SpecParseError> {
        self.skip_newlines();
        if self.eat(&TokenKind::Underscore) {
            return Ok(crate::ast::SpecPattern::Wildcard);
        }
        let name = self.expect_ident("for variant pattern")?;
        let bindings = if self.eat(&TokenKind::LParen) {
            let mut bs = Vec::new();
            self.skip_newlines();
            if !matches!(self.peek(), Some(TokenKind::RParen)) {
                bs.push(self.parse_variant_binding()?);
                while self.eat(&TokenKind::Comma) {
                    bs.push(self.parse_variant_binding()?);
                }
            }
            if !self.eat(&TokenKind::RParen) {
                return Err(self.err("expected `)` to close variant pattern"));
            }
            bs
        } else {
            Vec::new()
        };
        Ok(crate::ast::SpecPattern::Variant { name, bindings })
    }

    /// Parse a single binding inside a variant pattern: either an
    /// ident name or `_` for "ignore this positional arg". Internally
    /// `_` is bound to the literal name `_`; the body never references
    /// it, so collisions across `_` bindings are inert.
    fn parse_variant_binding(&mut self) -> Result<String, SpecParseError> {
        if self.eat(&TokenKind::Underscore) {
            return Ok("_".to_string());
        }
        self.expect_ident("for variant binding")
    }

    fn peek_kw(&self, name: &str) -> bool {
        match self.toks.get(self.idx).map(|t| &t.kind) {
            Some(TokenKind::Ident(n)) => n == name,
            Some(TokenKind::And) => name == "and",
            Some(TokenKind::Or) => name == "or",
            Some(TokenKind::Not) => name == "not",
            _ => false,
        }
    }
}