gigabrain / src / cypher / parser.rs

use nom::{
    IResult,
    branch::alt,
    combinator::{map, opt, value},
    multi::{many0, many1, separated_list0, separated_list1},
    sequence::{preceded, terminated, tuple, delimited, pair},
};
use std::collections::HashMap;
use crate::core::PropertyValue;
use crate::cypher::ast::*;
use crate::cypher::lexer::{Token, lex};

pub fn parse_cypher(input: &str) -> Result<CypherQuery, String> {
    let (_, tokens) = lex(input).map_err(|e| format!("Lexing error: {:?}", e))?;
    let (remaining, query) = cypher_query(&tokens).map_err(|e| format!("Parse error: {:?}", e))?;
    
    if !remaining.is_empty() {
        return Err(format!("Unexpected tokens after query: {:?}", remaining));
    }
    
    Ok(query)
}

type Tokens<'a> = &'a [Token];

fn cypher_query(input: Tokens) -> IResult<Tokens, CypherQuery> {
    alt((
        // Handle compound queries (MATCH ... RETURN ...)
        map(
            tuple((
                match_clause,
                opt(return_clause),
            )),
            |(match_clause, return_clause)| {
                if let Some(return_clause) = return_clause {
                    CypherQuery::Compound(vec![
                        CypherQuery::Match(match_clause),
                        CypherQuery::Return(return_clause),
                    ])
                } else {
                    CypherQuery::Match(match_clause)
                }
            }
        ),
        map(create_clause, CypherQuery::Create),
        map(return_clause, CypherQuery::Return),
        map(
            separated_list1(
                |i| expect_token(i, &Token::Semicolon),
                alt((
                    map(match_clause, CypherQuery::Match),
                    map(create_clause, CypherQuery::Create),
                    map(return_clause, CypherQuery::Return),
                ))
            ),
            CypherQuery::Compound
        ),
    ))(input)
}

fn match_clause(input: Tokens) -> IResult<Tokens, MatchClause> {
    map(
        tuple((
            opt(|i| expect_token(i, &Token::Optional)),
            |i| expect_token(i, &Token::Match),
            pattern,
            opt(where_clause),
        )),
        |(optional, _, pattern, where_clause)| MatchClause {
            pattern,
            where_clause,
            optional: optional.is_some(),
        },
    )(input)
}

fn create_clause(input: Tokens) -> IResult<Tokens, CreateClause> {
    map(
        preceded(
            |i| expect_token(i, &Token::Create),
            pattern,
        ),
        |pattern| CreateClause { pattern },
    )(input)
}

fn return_clause(input: Tokens) -> IResult<Tokens, ReturnClause> {
    map(
        tuple((
            |i| expect_token(i, &Token::Return),
            opt(|i| expect_token(i, &Token::Distinct)),
            separated_list1(|i| expect_token(i, &Token::Comma), return_item),
            opt(order_by),
            opt(preceded(|i| expect_token(i, &Token::Skip), expression)),
            opt(preceded(|i| expect_token(i, &Token::Limit), expression)),
        )),
        |(_, distinct, items, order_by, skip, limit)| ReturnClause {
            distinct: distinct.is_some(),
            items,
            order_by,
            skip,
            limit,
        },
    )(input)
}

fn return_item(input: Tokens) -> IResult<Tokens, ReturnItem> {
    map(
        tuple((
            expression,
            opt(preceded(|i| expect_token(i, &Token::As), identifier)),
        )),
        |(expression, alias)| ReturnItem { expression, alias },
    )(input)
}

fn order_by(input: Tokens) -> IResult<Tokens, Vec<OrderByItem>> {
    preceded(
        tuple((
            |i| expect_token(i, &Token::Order),
            |i| expect_token(i, &Token::By),
        )),
        separated_list1(|i| expect_token(i, &Token::Comma), order_by_item),
    )(input)
}

fn order_by_item(input: Tokens) -> IResult<Tokens, OrderByItem> {
    map(
        tuple((
            expression,
            opt(alt((
                value(true, |i| expect_token(i, &Token::Asc)),
                value(false, |i| expect_token(i, &Token::Desc)),
            ))),
        )),
        |(expression, ascending)| OrderByItem {
            expression,
            ascending: ascending.unwrap_or(true),
        },
    )(input)
}

fn where_clause(input: Tokens) -> IResult<Tokens, Expression> {
    preceded(
        |i| expect_token(i, &Token::Where),
        expression,
    )(input)
}

fn pattern(input: Tokens) -> IResult<Tokens, Pattern> {
    map(
        separated_list1(|i| expect_token(i, &Token::Comma), pattern_part),
        |parts| Pattern {
            elements: parts.into_iter().flatten().collect(),
        },
    )(input)
}

fn pattern_part(input: Tokens) -> IResult<Tokens, Vec<PatternElement>> {
    many1(pattern_element)(input)
}

fn pattern_element(input: Tokens) -> IResult<Tokens, PatternElement> {
    alt((
        map(node_pattern, PatternElement::Node),
        map(relationship_pattern, PatternElement::Relationship),
    ))(input)
}

fn node_pattern(input: Tokens) -> IResult<Tokens, NodePattern> {
    delimited(
        |i| expect_token(i, &Token::LeftParen),
        map(
            tuple((
                opt(identifier),
                many0(preceded(|i| expect_token(i, &Token::Colon), identifier)),
                opt(properties_map),
            )),
            |(variable, labels, properties)| NodePattern {
                variable,
                labels,
                properties,
            },
        ),
        |i| expect_token(i, &Token::RightParen),
    )(input)
}

fn relationship_pattern(input: Tokens) -> IResult<Tokens, RelationshipPattern> {
    alt((
        map(
            tuple((
                |i| expect_token(i, &Token::Minus),
                delimited(
                    |i| expect_token(i, &Token::LeftBracket),
                    relationship_detail,
                    |i| expect_token(i, &Token::RightBracket),
                ),
                |i| expect_token(i, &Token::Arrow),
            )),
            |(_, detail, _)| RelationshipPattern {
                direction: RelationshipDirection::Outgoing,
                ..detail
            },
        ),
        map(
            tuple((
                |i| expect_token(i, &Token::LeftArrow),
                delimited(
                    |i| expect_token(i, &Token::LeftBracket),
                    relationship_detail,
                    |i| expect_token(i, &Token::RightBracket),
                ),
                |i| expect_token(i, &Token::Minus),
            )),
            |(_, detail, _)| RelationshipPattern {
                direction: RelationshipDirection::Incoming,
                ..detail
            },
        ),
    ))(input)
}

fn relationship_detail(input: Tokens) -> IResult<Tokens, RelationshipPattern> {
    map(
        tuple((
            opt(identifier),
            many0(preceded(|i| expect_token(i, &Token::Colon), identifier)),
            opt(properties_map),
        )),
        |(variable, types, properties)| RelationshipPattern {
            variable,
            types,
            properties,
            direction: RelationshipDirection::Both,
            length: None,
        },
    )(input)
}

fn properties_map(input: Tokens) -> IResult<Tokens, HashMap<String, Expression>> {
    delimited(
        |i| expect_token(i, &Token::LeftBrace),
        map(
            separated_list0(
                |i| expect_token(i, &Token::Comma),
                tuple((
                    identifier,
                    |i| expect_token(i, &Token::Colon),
                    expression,
                )),
            ),
            |items| items.into_iter().map(|(k, _, v)| (k, v)).collect(),
        ),
        |i| expect_token(i, &Token::RightBrace),
    )(input)
}

fn expression(input: Tokens) -> IResult<Tokens, Expression> {
    or_expression(input)
}

fn or_expression(input: Tokens) -> IResult<Tokens, Expression> {
    let (input, left) = and_expression(input)?;
    
    let (input, rights) = many0(preceded(
        |i| expect_token(i, &Token::Or),
        and_expression,
    ))(input)?;
    
    Ok((input, rights.into_iter().fold(left, |acc, right| {
        Expression::Or(Box::new(acc), Box::new(right))
    })))
}

fn and_expression(input: Tokens) -> IResult<Tokens, Expression> {
    let (input, left) = comparison_expression(input)?;
    
    let (input, rights) = many0(preceded(
        |i| expect_token(i, &Token::And),
        comparison_expression,
    ))(input)?;
    
    Ok((input, rights.into_iter().fold(left, |acc, right| {
        Expression::And(Box::new(acc), Box::new(right))
    })))
}

fn comparison_expression(input: Tokens) -> IResult<Tokens, Expression> {
    let (input, left) = additive_expression(input)?;
    
    if let Ok((input, _)) = expect_token(input, &Token::Equal) {
        let (input, right) = additive_expression(input)?;
        return Ok((input, Expression::Equal(Box::new(left), Box::new(right))));
    }
    
    Ok((input, left))
}

fn additive_expression(input: Tokens) -> IResult<Tokens, Expression> {
    let (input, left) = multiplicative_expression(input)?;
    
    let (input, ops) = many0(tuple((
        alt((
            value(true, |i| expect_token(i, &Token::Plus)),
            value(false, |i| expect_token(i, &Token::Minus)),
        )),
        multiplicative_expression,
    )))(input)?;
    
    Ok((input, ops.into_iter().fold(left, |acc, (is_add, right)| {
        if is_add {
            Expression::Add(Box::new(acc), Box::new(right))
        } else {
            Expression::Subtract(Box::new(acc), Box::new(right))
        }
    })))
}

fn multiplicative_expression(input: Tokens) -> IResult<Tokens, Expression> {
    let (input, left) = primary_expression(input)?;
    Ok((input, left))
}

fn primary_expression(input: Tokens) -> IResult<Tokens, Expression> {
    alt((
        // Try property access first (a.name)
        map(
            tuple((
                identifier,
                |i| expect_token(i, &Token::Dot),
                identifier,
            )),
            |(obj, _, prop)| Expression::Property(PropertyExpression {
                expression: Box::new(Expression::Variable(obj)),
                property: prop,
            }),
        ),
        // Then literals and simple variables
        map(literal, |lit| Expression::Literal(lit)),
        map(identifier, Expression::Variable),
        delimited(
            |i| expect_token(i, &Token::LeftParen),
            expression,
            |i| expect_token(i, &Token::RightParen),
        ),
    ))(input)
}

fn literal(input: Tokens) -> IResult<Tokens, PropertyValue> {
    if input.is_empty() {
        return Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Eof)));
    }
    
    match &input[0] {
        Token::Integer(n) => Ok((&input[1..], PropertyValue::Integer(*n))),
        Token::Float(f) => Ok((&input[1..], PropertyValue::Float(*f))),
        Token::String(s) => Ok((&input[1..], PropertyValue::String(s.clone()))),
        Token::True => Ok((&input[1..], PropertyValue::Boolean(true))),
        Token::False => Ok((&input[1..], PropertyValue::Boolean(false))),
        Token::Null => Ok((&input[1..], PropertyValue::Null)),
        _ => Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Tag))),
    }
}

fn identifier(input: Tokens) -> IResult<Tokens, String> {
    if input.is_empty() {
        return Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Eof)));
    }
    
    match &input[0] {
        Token::Identifier(s) => Ok((&input[1..], s.clone())),
        _ => Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Tag))),
    }
}

fn expect_token<'a>(input: &'a [Token], expected: &Token) -> IResult<&'a [Token], &'a Token> {
    if input.is_empty() {
        return Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Eof)));
    }
    
    if std::mem::discriminant(&input[0]) == std::mem::discriminant(expected) {
        Ok((&input[1..], &input[0]))
    } else {
        Err(nom::Err::Error(nom::error::Error::new(input, nom::error::ErrorKind::Tag)))
    }
}