use super::{node, unquote};
use crate::ast::source::*;
use crate::error::{NbclError, Result, Span};
use crate::parser::Rule;
use pest::iterators::Pair;
use std::sync::atomic::{AtomicU64, Ordering};
pub fn build_stmt(pair: Pair<Rule>) -> Result<Stmt> {
let inner = match pair.as_rule() {
Rule::stmt | Rule::node_stmt => pair.into_inner().next().unwrap(),
_ => pair,
};
let span = Span::from_pair(&inner);
match inner.as_rule() {
Rule::let_stmt | Rule::const_stmt => {
let mut ii = inner.clone().into_inner();
let name = ii.next().unwrap().as_str().to_string();
let next = ii.next().unwrap();
let value = build_expr(next)?;
if inner.as_rule() == Rule::let_stmt {
Ok(Stmt::Let(name, value))
} else {
Ok(Stmt::Const(name, value))
}
}
Rule::assign_stmt => {
let span = Span::from_pair(&inner);
let mut ii = inner.clone().into_inner();
let name = build_expr(ii.next().unwrap())?;
let pair = ii.next().unwrap();
let inner = pair.into_inner().next().unwrap();
let assign_op = match inner.as_rule() {
Rule::equal => AssignOp::Equal,
Rule::plus_equal => AssignOp::PlusEqual,
Rule::min_equal => AssignOp::MinEqual,
Rule::mult_equal => AssignOp::MultEqual,
Rule::div_equal => AssignOp::DivEqual,
_ => {
return Err(NbclError::Parse {
message: format!(
"Unkown assign operator ('{}') encountered.",
inner.as_str().to_string()
),
hint: Some(
"Replace the assign operator with one of these: '=', '+=', '-='."
.into(),
),
span: Some(span.clone()),
});
}
};
let next = ii.next().unwrap();
let value = build_expr(next)?;
Ok(Stmt::Assign(name, assign_op, value, span))
}
Rule::for_stmt => {
let mut ii = inner.into_inner();
let pattern_pair = ii.next().unwrap();
let mut patterns = Vec::new();
for ident in pattern_pair.into_inner() {
patterns.push(ident.as_str().to_string());
}
let _ = ii.next().unwrap();
let iter_expr = build_expr(ii.next().unwrap())?;
let block_pair = ii.next().unwrap();
let body = build_block(block_pair)?;
Ok(Stmt::For(patterns, iter_expr, body))
}
Rule::while_stmt => {
let mut ii = inner.into_inner();
let condition = build_expr(ii.next().unwrap())?;
let block_pair = ii.next().unwrap();
let body = build_block(block_pair)?;
Ok(Stmt::While(condition, body))
}
Rule::return_stmt => {
let mut ii = inner.into_inner();
let return_type = if let Some(e_pair) = ii.next() {
match e_pair.as_rule() {
Rule::node_invocation => {
let node_inv = node::build_node_invocation(e_pair)?;
Some(ReturnType::Node(node_inv))
}
Rule::expr => {
let expr = build_expr(e_pair)?;
Some(ReturnType::Expr(expr))
}
rt => {
return Err(NbclError::Ast {
message: format!("unknown return type: {:?}", rt),
hint: None,
span: Some(span),
});
}
}
} else {
None
};
Ok(Stmt::Return(return_type, span))
}
Rule::expr_stmt => Ok(Stmt::Expr(build_expr(inner.into_inner().next().unwrap())?)),
_ => Err(NbclError::Ast {
message: format!("unknown Statement: {:?}", inner.as_rule()),
hint: None,
span: Some(span),
}),
}
}
pub fn build_expr(pair: Pair<Rule>) -> Result<Expr> {
let span = Span::from_pair(&pair);
match pair.as_rule() {
Rule::expr | Rule::or_expr | Rule::and_expr | Rule::add_expr | Rule::mul_expr => {
build_binop(pair)
}
Rule::cmp_expr => {
let mut inner = pair.into_inner();
let lhs = build_expr(inner.next().unwrap())?;
if let Some(op_pair) = inner.next() {
let op = op_pair.as_str().to_string();
let rhs = build_expr(inner.next().unwrap())?;
Ok(Expr { kind: ExprKind::Binary(Box::new(lhs), op, Box::new(rhs)), span })
} else {
Ok(lhs)
}
}
Rule::unary_expr => {
let mut inner = pair.into_inner();
let first = inner.next().ok_or_else(|| NbclError::Ast {
message: "Empty unary expression".to_string(),
hint: None,
span: Some(span.clone()),
})?;
if first.as_rule() == Rule::postfix_expr {
build_expr(first)
} else {
let op = first.as_str().to_string();
let operand_pair = inner.next().ok_or_else(|| NbclError::Ast {
message: format!("Expected operand after unary operator '{}'", op),
hint: None,
span: Some(span.clone()),
})?;
let operand = build_expr(operand_pair)?;
Ok(Expr { kind: ExprKind::Unary(op, Box::new(operand)), span })
}
}
Rule::assignable_lhs => {
let mut inner = pair.into_inner();
let mut res = build_expr(inner.next().unwrap())?;
let mut it = inner.peekable();
while let Some(suffix) = it.next() {
res = match suffix.as_rule() {
Rule::accessor => {
let is_safe = suffix.as_str() == "?.";
let ident = it.next().unwrap().as_str().to_string();
Expr {
kind: ExprKind::Field(Box::new(res), ident, is_safe),
span: span.clone(),
}
}
Rule::expr => Expr {
kind: ExprKind::Index(Box::new(res), Box::new(build_expr(suffix)?)),
span: span.clone(),
},
_ => res,
};
}
Ok(res)
}
Rule::postfix_expr => {
let mut inner = pair.into_inner();
let mut res = build_expr(inner.next().unwrap())?;
let mut it = inner.peekable();
while let Some(suffix) = it.next() {
res = match suffix.as_rule() {
Rule::accessor => {
let is_safe = suffix.as_str() == "?.";
let ident = it.next().unwrap().as_str().to_string();
Expr {
kind: ExprKind::Field(Box::new(res), ident, is_safe),
span: span.clone(),
}
}
Rule::expr => Expr {
kind: ExprKind::Index(Box::new(res), Box::new(build_expr(suffix)?)),
span: span.clone(),
},
Rule::call_args => {
let args = suffix
.into_inner()
.map(|arg_pair| build_expr(arg_pair))
.collect::<Result<Vec<_>>>()?;
Expr { kind: ExprKind::Call(Box::new(res), args), span: span.clone() }
}
_ => res,
};
}
Ok(res)
}
Rule::range_expr => {
let mut inner = pair.into_inner();
let start = build_expr(inner.next().unwrap())?;
let op = inner.next().unwrap();
let inclusive = op.as_str() == "..=";
let end = build_expr(inner.next().unwrap())?;
Ok(Expr { kind: ExprKind::Range(Box::new(start), Box::new(end), inclusive), span })
}
Rule::id_expression => build_expr(pair.into_inner().next().unwrap()),
Rule::string_lit => {
Ok(Expr { kind: ExprKind::Literal(Literal::Str(unquote(pair.as_str()))), span })
}
Rule::primary_expr => build_expr(pair.into_inner().next().unwrap()),
Rule::literal => Ok(Expr { kind: ExprKind::Literal(build_literal(pair)?), span }),
Rule::if_expr => Ok(Expr { kind: ExprKind::If(Box::new(build_if(pair)?)), span }),
Rule::match_expr => Ok(Expr { kind: build_match(pair)?, span }),
Rule::lambda_expr => Ok(Expr { kind: build_lambda(pair)?, span }),
Rule::snake_ident => Ok(Expr { kind: ExprKind::Variable(pair.as_str().to_string()), span }),
_ => Err(NbclError::Ast {
message: format!("unknown expr: {:?}", pair.as_rule()),
hint: None,
span: Some(span),
}),
}
}
fn build_binop(pair: Pair<Rule>) -> Result<Expr> {
let span = Span::from_pair(&pair);
let mut inner = pair.into_inner();
let mut lhs = build_expr(inner.next().unwrap())?;
while let Some(op_pair) = inner.next() {
let op_str = op_pair.as_str().to_string();
let rhs_pair = inner.next().ok_or_else(|| NbclError::Ast {
message: "expected operand after operator".to_string(),
hint: Some(
"An operator like '+' must be followed by a value, variable, or '('.".to_string(),
),
span: Some(span.clone()),
})?;
let rhs = build_expr(rhs_pair)?;
lhs = Expr {
kind: ExprKind::Binary(Box::new(lhs), op_str, Box::new(rhs)),
span: span.clone(),
};
}
Ok(lhs)
}
fn build_literal(pair: Pair<Rule>) -> Result<Literal> {
let inner = pair.into_inner().next().unwrap();
match inner.as_rule() {
Rule::int_lit => Ok(Literal::Int(inner.as_str().parse().unwrap())),
Rule::float_lit => Ok(Literal::Float(inner.as_str().parse().unwrap())),
Rule::bool_lit => Ok(Literal::Bool(inner.as_str() == "true")),
Rule::string_lit => Ok(Literal::Str(unquote(inner.as_str()))),
Rule::list_lit => {
let mut exprs = Vec::new();
for p in inner.into_inner() {
exprs.push(build_expr(p)?);
}
Ok(Literal::List(exprs))
}
Rule::map_lit => {
let mut pairs = Vec::new();
for p in inner.into_inner() {
let mut inner_pair = p.into_inner();
let key = inner_pair.next().unwrap().as_str().to_string();
let _ = inner_pair.next().unwrap();
let value = build_expr(inner_pair.next().unwrap())?;
pairs.push((key, value));
}
Ok(Literal::Map(pairs))
}
_ => Ok(Literal::Null),
}
}
fn build_if(pair: Pair<Rule>) -> Result<IfExpr> {
let mut inner = pair.into_inner();
let condition = build_expr(inner.next().unwrap())?;
let then_branch = build_branch(inner.next().unwrap())?;
let mut else_ifs = Vec::new();
let mut else_branch = None;
while let Some(next) = inner.next() {
match next.as_rule() {
Rule::else_if_branch => {
let mut ei_inner = next.into_inner();
let cond = build_expr(ei_inner.next().unwrap())?;
let body = build_branch(ei_inner.next().unwrap())?;
else_ifs.push((cond, body));
}
Rule::else_branch => {
else_branch = Some(build_branch(next.into_inner().next().unwrap())?);
}
_ => unreachable!("Unexpected rule in if_expr: {:?}", next.as_rule()),
}
}
Ok(IfExpr { condition, then_branch, else_ifs, else_branch })
}
fn build_branch(pair: Pair<Rule>) -> Result<(Vec<Stmt>, Option<Expr>)> {
let mut stmts = Vec::new();
let mut final_expr = None;
for p in pair.into_inner() {
match p.as_rule() {
Rule::stmt => stmts.push(build_stmt(p)?),
Rule::expr => final_expr = Some(build_expr(p)?),
_ => {}
}
}
if final_expr.is_none() {
if let Some(last_stmt) = stmts.last() {
if let Stmt::Expr(expr) = last_stmt {
final_expr = Some(expr.clone());
stmts.pop();
}
}
}
Ok((stmts, final_expr))
}
fn build_match(pair: Pair<Rule>) -> Result<ExprKind> {
let mut inner = pair.into_inner();
let condition = build_expr(inner.next().unwrap())?;
let mut arms = Vec::new();
for arm_pair in inner {
arms.push(build_match_arm(arm_pair)?);
}
Ok(ExprKind::Match(Box::new(condition), arms))
}
fn build_match_arm(pair: Pair<Rule>) -> Result<MatchArm> {
let mut inner = pair.into_inner();
let pattern_pair = inner.next().unwrap();
let pattern_inner = pattern_pair.into_inner().next().unwrap();
let is_var = match pattern_inner.as_rule() {
Rule::snake_ident => true,
_ => false,
};
let pattern = unquote(pattern_inner.as_str());
let body_pair = inner.next().unwrap();
let body = match body_pair.as_rule() {
Rule::block_body => {
let block = build_block(body_pair)?;
MatchBody::Block(block.stmts, block.terminator)
}
_ => {
let expr = build_expr(body_pair)?;
MatchBody::Expr(expr)
}
};
Ok(MatchArm { pattern, body, is_var })
}
pub fn generate_anon_fn_name() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(0);
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("<lambda:{}>", id)
}
fn build_lambda(pair: Pair<Rule>) -> Result<ExprKind> {
let span = Span::from_pair(&pair);
let mut inner = pair.clone().into_inner();
let mut params = Vec::new();
while let Some(next) = inner.peek() {
if next.as_rule() == Rule::lambda_param {
let mut param_inner = inner.next().unwrap().into_inner();
let name = param_inner.next().unwrap().as_str().to_string();
params.push(name);
} else {
break;
}
}
let body_pair = inner.next().ok_or_else(|| NbclError::Ast {
message: "Lambda must have a body".into(),
hint: None,
span: Some(span.clone()),
})?;
let mut body_items = Vec::new();
let actual_body = if body_pair.as_rule() == Rule::lambda_body {
body_pair.into_inner().next().unwrap()
} else {
body_pair
};
match actual_body.as_rule() {
Rule::block_body => {
let block = build_block(actual_body)?;
for stmt in block.stmts {
body_items.push(FnItem::Stmt(stmt));
}
if let Some(terminator_expr) = block.terminator {
body_items.push(FnItem::Stmt(Stmt::Expr(terminator_expr)));
}
}
_ => {
let expr = build_expr(actual_body)?;
body_items.push(FnItem::Stmt(Stmt::Expr(expr)));
}
}
let fn_def = FnDef { name: generate_anon_fn_name(), params, body: body_items, span };
Ok(ExprKind::Lambda(fn_def))
}
pub fn build_block(pair: Pair<Rule>) -> Result<Block> {
let mut stmts = Vec::new();
let mut terminator = None;
for inner_pair in pair.into_inner() {
match inner_pair.as_rule() {
Rule::stmt => {
stmts.push(build_stmt(inner_pair)?);
}
Rule::expr => {
terminator = Some(build_expr(inner_pair)?);
}
_ => {}
}
}
Ok(Block { stmts, terminator })
}