use crate::utils::ast::{AstNode, AstNodeType};
use anyhow::{Result, anyhow};
use pest::Parser;
use pest::iterators::Pair;
#[derive(pest_derive::Parser)]
#[grammar = "src/vibelang.pest"]
pub struct VibeParser;
pub fn parse_source(source: &str) -> Result<AstNode> {
let pairs = VibeParser::parse(Rule::Program, source)?
.next()
.ok_or_else(|| anyhow!("Failed to parse program: no pairs found"))?;
let mut program_node = AstNode::new(AstNodeType::Program);
for pair in pairs.into_inner() {
if let Rule::Declaration = pair.as_rule() {
if let Some(declaration) = pair.into_inner().next() {
let decl_node = build_ast_from_pair(declaration)?;
program_node.add_child(decl_node);
}
}
}
Ok(program_node)
}
fn build_ast_from_pair(pair: Pair<Rule>) -> Result<AstNode> {
match pair.as_rule() {
Rule::FunctionDecl => {
let mut inner = pair.into_inner();
let name = inner.next().unwrap().as_str();
let mut func = AstNode::new(AstNodeType::FunctionDecl);
func.set_string("name", name);
for item in inner {
match item.as_rule() {
Rule::ParamList => {
let mut params_node = AstNode::new(AstNodeType::ParamList);
for param_pair in item.into_inner() {
let mut param_inner = param_pair.into_inner();
let param_name = param_inner.next().unwrap().as_str();
let param_type = build_ast_from_pair(param_inner.next().unwrap())?;
let mut param_node = AstNode::new(AstNodeType::Parameter);
param_node.set_string("name", param_name);
param_node.add_child(param_type);
params_node.add_child(param_node);
}
func.add_child(params_node);
}
Rule::Type => {
let type_node = build_ast_from_pair(item)?;
func.add_child(type_node);
}
Rule::Block => {
let body = build_ast_from_pair(item)?;
func.add_child(body);
}
_ => {}
}
}
Ok(func)
}
Rule::Block => {
let mut block_node = AstNode::new(AstNodeType::Block);
for stmt_pair in pair.into_inner() {
if let Rule::Statement = stmt_pair.as_rule() {
if let Some(statement) = stmt_pair.into_inner().next() {
let stmt_node = build_ast_from_pair(statement)?;
block_node.add_child(stmt_node);
}
}
}
Ok(block_node)
}
Rule::PromptStmt => {
let template = pair.into_inner().next().unwrap().as_str();
let mut prompt_node = AstNode::new(AstNodeType::PromptBlock);
prompt_node.set_string("template", &template[1..template.len() - 1]);
Ok(prompt_node)
}
Rule::TypeDecl => {
let mut inner = pair.into_inner();
let name = inner.next().unwrap().as_str();
let type_def = build_ast_from_pair(inner.next().unwrap())?;
let mut type_decl_node = AstNode::new(AstNodeType::TypeDecl);
type_decl_node.set_string("name", name);
type_decl_node.add_child(type_def);
Ok(type_decl_node)
}
Rule::Type => {
build_ast_from_pair(pair.into_inner().next().unwrap())
}
Rule::MeaningType => {
let mut inner = pair.into_inner();
let base_type = build_ast_from_pair(inner.next().unwrap())?;
let meaning_str = inner.next().unwrap().as_str();
let mut meaning_node = AstNode::new(AstNodeType::MeaningType);
meaning_node.set_string("meaning", &meaning_str[1..meaning_str.len() - 1]);
meaning_node.add_child(base_type);
Ok(meaning_node)
}
Rule::BasicType => {
let type_name = pair.as_str();
let mut basic_type_node = AstNode::new(AstNodeType::BasicType);
basic_type_node.set_string("type", type_name);
Ok(basic_type_node)
}
_ => Err(anyhow!("Unhandled grammar rule: {:?}", pair.as_rule())),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::utils::ast::AstNodeType;
#[test]
fn test_parse_simple_function() {
let source = "fn get_greeting() { prompt \"Say hello\"; }";
let ast = parse_source(source).expect("Parsing failed");
assert_eq!(ast.node_type, AstNodeType::Program);
assert_eq!(ast.children.len(), 1, "Program should have one child node");
let func_node = &ast.children[0];
assert_eq!(func_node.node_type, AstNodeType::FunctionDecl);
assert_eq!(func_node.get_string("name").unwrap(), "get_greeting");
let block_node = func_node
.children
.iter()
.find(|n| n.node_type == AstNodeType::Block)
.unwrap();
let prompt_node = &block_node.children[0];
assert_eq!(prompt_node.node_type, AstNodeType::PromptBlock);
assert_eq!(prompt_node.get_string("template").unwrap(), "Say hello");
}
#[test]
fn test_parse_function_with_params_and_return_type() {
let source = r#"
fn get_age(birth_year: Int) -> Int {
prompt "Calculate age from {birth_year}";
}
"#;
let ast = parse_source(source).expect("Parsing failed");
let func_node = &ast.children[0];
assert_eq!(func_node.get_string("name").unwrap(), "get_age");
let params_node = func_node
.children
.iter()
.find(|n| n.node_type == AstNodeType::ParamList)
.unwrap();
assert_eq!(params_node.children.len(), 1);
let param_node = ¶ms_node.children[0];
assert_eq!(param_node.get_string("name").unwrap(), "birth_year");
let param_type_node = ¶m_node.children[0];
assert_eq!(param_type_node.get_string("type").unwrap(), "Int");
let return_type_node = func_node
.children
.iter()
.find(|n| n.node_type == AstNodeType::BasicType)
.unwrap();
assert_eq!(return_type_node.get_string("type").unwrap(), "Int");
}
#[test]
fn test_parse_type_alias_with_meaning() {
let source = r#"type Population = Meaning<Int>("population count in millions");"#;
let ast = parse_source(source).expect("Parsing failed");
assert_eq!(ast.children.len(), 1);
let type_decl_node = &ast.children[0];
assert_eq!(type_decl_node.node_type, AstNodeType::TypeDecl);
assert_eq!(type_decl_node.get_string("name").unwrap(), "Population");
let meaning_node = &type_decl_node.children[0];
assert_eq!(meaning_node.node_type, AstNodeType::MeaningType);
assert_eq!(
meaning_node.get_string("meaning").unwrap(),
"population count in millions"
);
let base_type_node = &meaning_node.children[0];
assert_eq!(base_type_node.get_string("type").unwrap(), "Int");
}
#[test]
fn test_parse_full_payload_with_multiple_definitions() {
let source = r#"
type Year = Int;
type City = Meaning<String>("a major city");
fn get_year() -> Year { prompt "What year is it?"; }
"#;
let ast = parse_source(source).expect("Parsing failed");
assert_eq!(ast.children.len(), 3);
assert_eq!(ast.children[0].node_type, AstNodeType::TypeDecl);
assert_eq!(ast.children[1].node_type, AstNodeType::TypeDecl);
assert_eq!(ast.children[2].node_type, AstNodeType::FunctionDecl);
}
#[test]
fn test_parser_rejects_invalid_syntax() {
let source = "fn my_func( -> ) { }"; let result = parse_source(source);
assert!(result.is_err(), "Parser should fail on invalid syntax");
}
}