ccarp 0.1.2

//! C Statements
//! 
//! This module defines types associated with C statements.
//! Importand definitions include:
//! - `Statement`, which signifies a C99 statement
use std::mem::swap;

use pest::iterators::Pair;

use crate::ccarp::error::{rule_err, rule_mismatch, safe_unwrap, CCErr, Result};

use super::{decl::{Declaration, NonEmptyVec}, defs::{list_ast, Rule, AST}, expr::{ConstExpr, Expression}, tt::Identifier};

/// C Statement
/// 
/// These can be one of:
/// - Labeled Statement - `label`, `case`, `default`
/// - Compound Statement - `{ decl statement decl statement ... }`
/// - Expression Statement - `expr? ;`
/// - Selection Statement - `if`, `if-else`, `switch`
/// - Iteration Statement - `while`, `do-while`, `for`
/// - Jump Statement - `goto`, `continue`, `break`, `return`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Statement {
    Labeled(LabeledStmt),
    Compound(CompoundStmt),
    Expr(ExprStmt),
    Selection(SelectionStmt),
    Iter(IterStmt),
    Jump(JumpStmt)
}
impl AST for Statement {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::stmt => {
                let pair=safe_unwrap!(pair.into_inner().next(),"Statement");
                match pair.as_rule() {
                    Rule::labeled_stmt => Ok(Self::Labeled(LabeledStmt::take(pair)?)),
                    Rule::compound_stmt => Ok(Self::Compound(CompoundStmt::take(pair)?)),
                    Rule::expr_stmt => Ok(Self::Expr(ExprStmt::take(pair)?)),
                    Rule::select_stmt => Ok(Self::Selection(SelectionStmt::take(pair)?)),
                    Rule::iter_stmt => Ok(Self::Iter(IterStmt::take(pair)?)),
                    Rule::jump_stmt => Ok(Self::Jump(JumpStmt::take(pair)?)),
                    _ => Err(rule_mismatch!(pair))
                }
            },
            _ => Err(rule_mismatch!(pair,stmt))
        }
    }
}

/// C Labeled Statement
/// 
/// These can be:
/// - Label (an Identifier followed by a Statement) - `ident : statement`
/// - Case (a Constant expression followed by a Statement) - `case const-expr : statement`
/// - Default (a Statement) - `default : statement`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum LabeledStmt {
    Label(Identifier,Box<Statement>),   // identifier : statement
    Case(ConstExpr,Box<Statement>),     // case constant-expression : statement
    Default(Box<Statement>)             // default : statement
}
impl AST for LabeledStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::labeled_stmt => {
                let mut inner=pair.into_inner();
                let pair=safe_unwrap!(inner.next(),"Labeled Statement");
                match pair.as_rule() {
                    Rule::tok_ident => Ok(Self::Label(Identifier::take(pair)?, Box::new(Statement::take(safe_unwrap!(inner.next(),"Label Statement"))?))),
                    Rule::const_expr => Ok(Self::Case(ConstExpr::take(pair)?, Box::new(Statement::take(safe_unwrap!(inner.next(),"Case Statement"))?))),
                    Rule::stmt => Ok(Self::Default(Box::new(Statement::take(pair)?))),
                    _ => Err(rule_mismatch!(pair))
                }
            },
            _ => Err(rule_mismatch!(pair,labeled_stmt))
        }
    }
}
impl LabeledStmt {
    pub fn push_statement(&mut self,s: Statement) {
        match self {
            Self::Label(_, statement)|Self::Case(_, statement)|Self::Default(statement) => {
                if let Statement::Compound(compound_stmt) = statement.as_mut() {
                    match &mut compound_stmt.0 {
                        Some(list) => list.0.push(BlockItem::Stmt(Box::new(s))),
                        None => compound_stmt.0=Some(BlockItemList(vec![BlockItem::Stmt(Box::new(s))])),
                    }
                } else {
                    let mut tmp=Box::new(Statement::Compound(CompoundStmt(None)));
                    swap(&mut tmp, statement);
                    *statement=Box::new(Statement::Compound(CompoundStmt(Some(BlockItemList(vec![BlockItem::Stmt(tmp),BlockItem::Stmt(Box::new(s))])))));
                }
            },
        }
    }
    pub fn push_declaration(&mut self, d: Declaration) {
        match self {
            Self::Label(_, statement)|Self::Case(_, statement)|Self::Default(statement) => {
                if let Statement::Compound(compound_stmt) = statement.as_mut() {
                    match &mut compound_stmt.0 {
                        Some(list) => list.0.push(BlockItem::Decl(d)),
                        None => compound_stmt.0=Some(BlockItemList(vec![BlockItem::Decl(d)])),
                    }
                } else {
                    let mut tmp=Box::new(Statement::Compound(CompoundStmt(None)));
                    swap(&mut tmp, statement);
                    *statement=Box::new(Statement::Compound(CompoundStmt(Some(BlockItemList(vec![BlockItem::Stmt(tmp),BlockItem::Decl(d)])))));
                }
            },
        }
    }
}

/// C Compound Statement
/// 
/// These consist of:
/// - An optional block item list - `{ block-item-list? }`
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct CompoundStmt(pub Option<BlockItemList>); // { block-item-list? }
impl AST for CompoundStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::compound_stmt => {
                if let Some(pair)=pair.into_inner().next() {
                    match pair.as_rule() {
                        Rule::block_item_list => Ok(Self(Some(BlockItemList::take(pair)?))),
                        _ => Err(rule_mismatch!(pair,block_item_list))
                    }
                }
                else { Ok(Self(None)) }
            },
            _ => Err(rule_mismatch!(pair,compound_stmt))
        }
    }
}

/// C Block Item List
/// 
/// These consist of:
/// - Block Items - `block-item block-item ...`
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct BlockItemList(pub NonEmptyVec<BlockItem>); // block-item block-item ...
list_ast!(BlockItemList : BlockItem where block_item_list => block_item);

/// C Block Item
/// 
/// These can be:
/// - Declarations - `declaration`
/// - or Statements - `statement`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum BlockItem {
    Decl(Declaration),
    Stmt(Box<Statement>)
}
impl AST for BlockItem {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::block_item => {
                let pair=safe_unwrap!(pair.into_inner().next(),"Block Item");
                match pair.as_rule() {
                    Rule::decl => Ok(Self::Decl(Declaration::take(pair)?)),
                    Rule::stmt => Ok(Self::Stmt(Box::new(Statement::take(pair)?))),
                    _ => Err(rule_mismatch!(pair))
                }
            },
            _ => Err(rule_mismatch!(pair,block_item))
        }
    }
}

/// C Expression Statement
/// 
/// These consist of:
/// - Optional Expressions - `expression? ;`
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct ExprStmt(pub Option<Expression>);    // expression? ;
impl AST for ExprStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::expr_stmt => {
                if let Some(pair)=pair.into_inner().next() {
                    match pair.as_rule() {
                        Rule::expr => Ok(Self(Some(Expression::take(pair)?))),
                        _ => Err(rule_mismatch!(pair))
                    }
                }
                else { Ok(Self(None)) }
            },
            _ => Err(rule_mismatch!(pair,expr_stmt))
        }
    }
}

/// C Selection Statement
/// 
/// These can be one of:
/// - If Statement (Expression followed by a Statement) - `if ( expression ) statement`
/// - If-Else Statement (Expression followed by two Statements) - `if ( expression ) statement else statement`
/// - Switch Statement (Expression followed by Statement) - `switch ( expression ) statement`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum SelectionStmt {
    If(Expression,Box<Statement>),                      // if ( expression ) statement
    IfElse(Expression,Box<Statement>,Box<Statement>),   // if ( expression ) statement else statement
    Switch(Expression,Box<Statement>)                   // switch ( expression ) statement
}
impl AST for SelectionStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::select_stmt => {
                let mut inner=pair.into_inner();
                let pair=safe_unwrap!(inner.next(),"Selection Statement");
                match pair.as_rule() {
                    Rule::expr => {
                        let expr=Expression::take(pair)?;   // expression
                        let stmt=Statement::take(safe_unwrap!(inner.next(),"If/If Else Statement"))?;   // statement
                        if let Some(pair)=inner.next() { Ok(Self::IfElse(expr, Box::new(stmt), Box::new(Statement::take(pair)?))) } // if ( expression ) statement else statement
                        else { Ok(Self::If(expr, Box::new(stmt))) }     // if ( expression ) statement
                    },
                    Rule::switch => Ok(Self::Switch(    // switch
                        Expression::take(safe_unwrap!(inner.next(),"Switch Statement"))?,           // ( expression )
                        Box::new(Statement::take(safe_unwrap!(inner.next(),"Switch Statement"))?)   // statement
                    )),
                    _ => Err(rule_mismatch!(pair))
                }
            },
            _ => Err(rule_mismatch!(pair,select_stmt))
        }
    }
}

/// C Iteration Statement
/// 
/// These can be:
/// - While Statement (Expression followed by a Statement) - `while ( expression ) statement`
/// - Do-While Statement (Statement followed by an Expression) - `do statement while ( expression ) ;`
/// - For Statement (3 optional Expressions followed by a Statement) - `for ( expression? ; expression? ; expression? ) statement`
/// - For Statement II (a Declaration followed by 2 optional Expressions and a Statement) - `for ( declaration expression? ; expression? ) statement`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum IterStmt {
    While(Expression,Box<Statement>),                                                              // while ( expression ) statement
    DoWhile(Box<Statement>,Expression),                                                            // do statement while ( expression ) ;
    For(Option<Box<Expression>>,Option<Box<Expression>>,Option<Box<Expression>>,Box<Statement>),   // for ( expression? ; expression? ; expression? ) statement
    ForDecl(Box<Declaration>,Option<Box<Expression>>,Option<Box<Expression>>,Box<Statement>)       // for ( declaration expression? ; expression? ) statement
}
impl AST for IterStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::iter_stmt => {
                let mut is_do=false;    // do
                let (mut expr1,mut expr2,mut expr3,mut decl,mut stmt)=(None,None,None,None,None);
                for pair in pair.into_inner() {
                    match pair.as_rule() {
                        Rule::do2 => is_do=true,    // do
                        Rule::expr => {
                            if expr1.is_none() { expr1=Some(Expression::take(pair)?) }      // expression_1
                            else if expr2.is_none() { expr2=Some(Expression::take(pair)?) } // expression_2
                            else if expr3.is_none() { expr3=Some(Expression::take(pair)?) } // expression_3
                        },
                        Rule::stmt => stmt=Some(Statement::take(pair)?),    // statement
                        Rule::decl => decl=Some(Declaration::take(pair)?),  // declaration
                        _ => return Err(rule_mismatch!(pair))
                    }
                }
                // do statement while ( expression ) ;
                if is_do { Ok(Self::DoWhile(Box::new(safe_unwrap!(stmt,"Do While Statement")), safe_unwrap!(expr1,"Do While Statement"))) }
                // for ( expression? ; expression? ; expression? ) statement
                else if expr3.is_some() { Ok(Self::For(expr1.map(Box::new), expr2.map(Box::new), expr3.map(Box::new), Box::new(safe_unwrap!(stmt,"For Statement")))) }
                // for ( declaration expression? ; expression? ) statement
                else if let Some(item)=decl { Ok(Self::ForDecl(Box::new(item), expr1.map(Box::new), expr2.map(Box::new), Box::new(safe_unwrap!(stmt,"For Statement")))) }
                // while ( expression ) statement
                else { Ok(Self::While(safe_unwrap!(expr1,"While Statement"), Box::new(safe_unwrap!(stmt,"While Statement")))) }
            },
            _ => Err(rule_mismatch!(pair,iter_stmt))
        }
    }
}

/// C Jump Statement
/// 
/// These can be:
/// - Goto Statements (Identifier) - `goto identifier ;`
/// - Continue - `continue ;`
/// - Break - `break ;`
/// - Return Statement (optional Expression) - `return expression? ;`
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum JumpStmt {
    Goto(Identifier),           // goto identifier ;
    Continue,                   // continue ;
    Break,                      // break ;
    Return(Option<Expression>)  // return expression? ;
}
impl AST for JumpStmt {
    fn take(pair: Pair<Rule>) -> Result<Self> {
        match pair.as_rule() {
            Rule::jump_stmt => {
                if let Some(pair)=pair.into_inner().next() {
                    match pair.as_rule() {
                        Rule::tok_ident => Ok(Self::Goto(Identifier::take(pair)?)),
                        Rule::cont => Ok(Self::Continue),
                        Rule::br => Ok(Self::Break),
                        Rule::expr => Ok(Self::Return(Some(Expression::take(pair)?))),
                        _ => Err(rule_mismatch!(pair))
                    }
                }
                else { Ok(Self::Return(None)) }
            },
            _ => Err(rule_mismatch!(pair,jump_stmt))
        }
    }
}