lysp/src/compile/syntax/function.rs

use std::rc::Rc;

use crate::{
    compile::{
        function::FunctionSignature,
        syntax::{
            CollectErrors, ExpectedWhat, ExpectedWhere, Expression, ParseError, ParseErrorKind,
        },
    },
    vm::value::{ConsCell, IdentifierValue, Keyword, Value},
};

#[derive(Debug, PartialEq)]
pub struct FunctionBody {
    pub head: Vec<Rc<Expression>>,
    pub tail: Rc<Expression>,
}

#[derive(Debug, PartialEq)]
pub struct DefunExpression {
    pub name: IdentifierValue,
    pub signature: FunctionSignature,
    pub body: FunctionBody,
}

#[derive(Debug, PartialEq)]
pub struct PrognExpression {
    pub body: FunctionBody,
}

impl FunctionSignature {
    pub(super) fn parse(mut value: &Value, input: &Value) -> Result<Self, ParseError> {
        enum Mode {
            Required,
            Optional,
            Rest,
        }

        let mut required_arguments = vec![];
        let mut optional_arguments = vec![];
        let mut rest_argument = None;
        let mut mode = Mode::Required;

        while !value.is_nil() {
            let Value::Cons(cons) = value else {
                return Err(ParseError {
                    input: input.clone(),
                    error: ParseErrorKind::Expected(
                        ExpectedWhat::ArgumentSpec,
                        ExpectedWhere::InArgumentList,
                    ),
                });
            };
            let ConsCell(car, cdr) = cons.as_ref();

            match (&mode, car) {
                (Mode::Required, Value::Identifier(arg)) => {
                    required_arguments.push(arg.clone());
                }
                (Mode::Optional, Value::Identifier(arg)) => {
                    optional_arguments.push(arg.clone());
                }
                (Mode::Rest, Value::Identifier(arg)) => {
                    if rest_argument.is_some() {
                        return Err(ParseError {
                            input: input.clone(),
                            error: ParseErrorKind::Expected(
                                ExpectedWhat::ExactlyOneArgument,
                                ExpectedWhere::AfterKeyword(Keyword::Rest),
                            ),
                        });
                    }
                    rest_argument = Some(arg.clone());
                }
                (Mode::Required, Value::Keyword(Keyword::Optional)) => {
                    mode = Mode::Optional;
                }
                (Mode::Required, Value::Keyword(Keyword::Rest)) => {
                    mode = Mode::Rest;
                }
                (Mode::Optional, Value::Keyword(Keyword::Rest)) => {
                    if optional_arguments.is_empty() {
                        return Err(ParseError {
                            input: input.clone(),
                            error: ParseErrorKind::Expected(
                                ExpectedWhat::AtLeastOneArgument,
                                ExpectedWhere::AfterKeyword(Keyword::Optional),
                            ),
                        });
                    }
                    mode = Mode::Rest;
                }
                _ => {
                    return Err(ParseError {
                        input: input.clone(),
                        error: ParseErrorKind::Expected(
                            ExpectedWhat::ArgumentSpec,
                            ExpectedWhere::InArgumentList,
                        ),
                    });
                }
            }

            value = cdr;
        }

        match mode {
            Mode::Required => (),
            Mode::Optional => {
                if optional_arguments.is_empty() {
                    return Err(ParseError {
                        input: input.clone(),
                        error: ParseErrorKind::Expected(
                            ExpectedWhat::AtLeastOneArgument,
                            ExpectedWhere::AfterKeyword(Keyword::Optional),
                        ),
                    });
                }
            }
            Mode::Rest => {
                if rest_argument.is_none() {
                    return Err(ParseError {
                        input: input.clone(),
                        error: ParseErrorKind::Expected(
                            ExpectedWhat::ExactlyOneArgument,
                            ExpectedWhere::AfterKeyword(Keyword::Rest),
                        ),
                    });
                }
            }
        }

        Ok(Self {
            required_arguments,
            optional_arguments,
            rest_argument,
        })
    }
}

impl FunctionBody {
    pub(super) fn parse(
        mut value: &Value,
        input: &Value,
        inside: Keyword,
    ) -> Result<Self, ParseError> {
        let mut expressions = vec![];
        while !value.is_nil() {
            let Value::Cons(cons) = value else {
                return Err(ParseError {
                    input: input.clone(),
                    error: ParseErrorKind::Expected(
                        ExpectedWhat::ProperList,
                        ExpectedWhere::AfterArgumentList(inside),
                    ),
                });
            };
            let ConsCell(car, cdr) = cons.as_ref();

            let expression = Expression::parse_inner(car);
            expressions.push(expression);

            value = cdr;
        }
        let Some(tail) = expressions.pop() else {
            return Err(ParseError {
                input: input.clone(),
                error: ParseErrorKind::Expected(
                    ExpectedWhat::Expression,
                    ExpectedWhere::AfterArgumentList(inside),
                ),
            });
        };
        Ok(Self {
            head: expressions,
            tail,
        })
    }
}

impl PrognExpression {
    pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
        let body = FunctionBody::parse(value, input, Keyword::Progn)?;
        Ok(Self { body })
    }
}

impl DefunExpression {
    pub(super) fn parse(value: &Value, input: &Value) -> Result<Self, ParseError> {
        let Value::Cons(value) = value else {
            return Err(ParseError {
                input: input.clone(),
                error: ParseErrorKind::Expected(
                    ExpectedWhat::ProperList,
                    ExpectedWhere::InFunctionDefinition,
                ),
            });
        };
        let ConsCell(identifier, cdr) = value.as_ref();

        let Value::Identifier(identifier) = identifier else {
            return Err(ParseError {
                input: input.clone(),
                error: ParseErrorKind::Expected(
                    ExpectedWhat::Identifier,
                    ExpectedWhere::AfterKeyword(Keyword::Defun),
                ),
            });
        };

        let Value::Cons(cdr) = cdr else {
            return Err(ParseError {
                input: input.clone(),
                error: ParseErrorKind::Expected(
                    ExpectedWhat::ProperList,
                    ExpectedWhere::InFunctionDefinition,
                ),
            });
        };
        let ConsCell(car, cdr) = cdr.as_ref();

        let signature = FunctionSignature::parse(car, input)?;
        let body = FunctionBody::parse(cdr, input, Keyword::Lambda)?;
        Ok(Self {
            name: identifier.clone(),
            signature,
            body,
        })
    }
}

impl CollectErrors<ParseError> for FunctionBody {
    fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
        let mut r = false;
        for expr in self.head.iter() {
            r |= expr.collect_errors(errors);
        }
        r |= self.tail.collect_errors(errors);
        r
    }
}

impl CollectErrors<ParseError> for DefunExpression {
    fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
        self.body.collect_errors(errors)
    }
}

impl CollectErrors<ParseError> for PrognExpression {
    fn collect_errors(&self, errors: &mut Vec<ParseError>) -> bool {
        self.body.collect_errors(errors)
    }
}

#[cfg(test)]
mod tests {
    use std::rc::Rc;

    use crate::{
        compile::{
            function::FunctionSignature,
            syntax::{
                ExpectedWhat, ExpectedWhere, Expression, FunctionBody, ParseError, ParseErrorKind,
            },
        },
        vm::value::{Keyword, Value},
    };

    #[test]
    fn test_parse_function_body() {
        let v = Value::list_or_nil([Value::Number(1.into())]);
        let e = FunctionBody::parse(&v, &v, Keyword::Lambda).unwrap();
        assert_eq!(
            e,
            FunctionBody {
                head: vec![],
                tail: Rc::new(Expression::IntegerLiteral(1.into()))
            }
        );

        let v = Value::list_or_nil([
            Value::Number(1.into()),
            Value::Number(2.into()),
            Value::Number(3.into()),
        ]);
        let e = FunctionBody::parse(&v, &v, Keyword::Lambda).unwrap();
        assert_eq!(
            e,
            FunctionBody {
                head: vec![
                    Rc::new(Expression::IntegerLiteral(1.into())),
                    Rc::new(Expression::IntegerLiteral(2.into()))
                ],
                tail: Rc::new(Expression::IntegerLiteral(3.into()))
            }
        );

        // syntax error
        let v = Value::list_or_nil([]);
        let e = FunctionBody::parse(&v, &v, Keyword::Lambda).unwrap_err();
        assert_eq!(
            e,
            ParseError {
                input: v,
                error: ParseErrorKind::Expected(
                    ExpectedWhat::Expression,
                    ExpectedWhere::AfterArgumentList(Keyword::Lambda)
                ),
            }
        );
        let v = Value::Identifier("a".into()).cons(Value::Identifier("b".into()));
        let e = FunctionBody::parse(&v, &v, Keyword::Lambda).unwrap_err();
        assert_eq!(
            e,
            ParseError {
                input: v,
                error: ParseErrorKind::Expected(
                    ExpectedWhat::ProperList,
                    ExpectedWhere::AfterArgumentList(Keyword::Lambda)
                ),
            }
        );
    }

    #[test]
    fn test_parse_function_signature() {
        let args = Value::list_or_nil([]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec![],
                optional_arguments: vec![],
                rest_argument: None
            }
        );
        let args =
            Value::list_or_nil([Value::Identifier("a".into()), Value::Identifier("b".into())]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec!["a".into(), "b".into()],
                optional_arguments: vec![],
                rest_argument: None
            }
        );
        let args = Value::list_or_nil([
            Value::Identifier("a".into()),
            Value::Keyword(Keyword::Optional),
            Value::Identifier("b".into()),
        ]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec!["a".into()],
                optional_arguments: vec!["b".into()],
                rest_argument: None
            }
        );
        let args = Value::list_or_nil([
            Value::Identifier("a".into()),
            Value::Keyword(Keyword::Rest),
            Value::Identifier("b".into()),
        ]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec!["a".into()],
                optional_arguments: vec![],
                rest_argument: Some("b".into())
            }
        );
        let args =
            Value::list_or_nil([Value::Keyword(Keyword::Rest), Value::Identifier("c".into())]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec![],
                optional_arguments: vec![],
                rest_argument: Some("c".into())
            }
        );
        let args = Value::list_or_nil([
            Value::Keyword(Keyword::Optional),
            Value::Identifier("b".into()),
            Value::Keyword(Keyword::Rest),
            Value::Identifier("c".into()),
        ]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec![],
                optional_arguments: vec!["b".into()],
                rest_argument: Some("c".into())
            }
        );
        let args = Value::list_or_nil([
            Value::Identifier("a".into()),
            Value::Keyword(Keyword::Rest),
            Value::Identifier("c".into()),
        ]);
        let v = FunctionSignature::parse(&args, &args).unwrap();
        assert_eq!(
            v,
            FunctionSignature {
                required_arguments: vec!["a".into()],
                optional_arguments: vec![],
                rest_argument: Some("c".into())
            }
        );

        // syntax errors
        let args = Value::list_or_nil([Value::Keyword(Keyword::Optional)]);
        let e = FunctionSignature::parse(&args, &args).unwrap_err();
        assert_eq!(
            e,
            ParseError {
                input: args,
                error: ParseErrorKind::Expected(
                    ExpectedWhat::AtLeastOneArgument,
                    ExpectedWhere::AfterKeyword(Keyword::Optional)
                )
            }
        );
        let args = Value::list_or_nil([
            Value::Keyword(Keyword::Rest),
            Value::Identifier("a".into()),
            Value::Identifier("b".into()),
        ]);
        let e = FunctionSignature::parse(&args, &args).unwrap_err();
        assert_eq!(
            e,
            ParseError {
                input: args,
                error: ParseErrorKind::Expected(
                    ExpectedWhat::ExactlyOneArgument,
                    ExpectedWhere::AfterKeyword(Keyword::Rest)
                )
            }
        );
        let args = Value::list_or_nil([Value::Boolean(false.into())]);
        let e = FunctionSignature::parse(&args, &args).unwrap_err();
        assert_eq!(
            e,
            ParseError {
                input: args,
                error: ParseErrorKind::Expected(
                    ExpectedWhat::ArgumentSpec,
                    ExpectedWhere::InArgumentList
                )
            }
        );
    }
}