JsonRPC-Rust/libjrpc/src/ir.rs

use std::{
    collections::{HashMap, HashSet},
    error::Error,
    fmt::Display,
    hash::{Hash, Hasher},
};

use anyhow::Result;

use crate::parser::{
    EnumStatement, Node, ParserPosition, RootNode, ServiceStatement, TypeStatement,
};

static BUILT_INS: [&str; 6] = ["int", "float", "string", "boolean", "bytes", "void"];

pub trait Definition {
    fn get_position(&self) -> ParserPosition;
}

#[derive(Debug, Clone)]
pub struct IR {
    pub options: HashMap<String, String>,
    pub steps: Vec<Step>,
}

#[derive(Debug, Clone)]
pub enum Step {
    Type(TypeDefinition),
    Enum(EnumDefinition),
    Service(ServiceDefinition),
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum Type {
    Int,
    Float,
    String,
    Bool,
    Bytes,
    Void,
    Custom(String),
}

impl ToString for Type {
    fn to_string(&self) -> String {
        match self {
            Type::Int => "int".to_string(),
            Type::Float => "float".to_string(),
            Type::String => "string".to_string(),
            Type::Bool => "bool".to_string(),
            Type::Bytes => "bytes".to_string(),
            Type::Void => "void".to_string(),
            Type::Custom(name) => name.clone(),
        }
    }
}

impl Hash for Type {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.to_string().hash(state);
    }
}

impl From<&String> for Type {
    fn from(value: &String) -> Self {
        Self::from(value.as_str())
    }
}

impl From<String> for Type {
    fn from(value: String) -> Self {
        Self::from(value.as_str())
    }
}

impl From<&str> for Type {
    fn from(s: &str) -> Self {
        match s {
            "int" => Type::Int,
            "float" => Type::Float,
            "string" => Type::String,
            "bool" => Type::Bool,
            "boolean" => Type::Bool,
            "bytes" => Type::Bytes,
            "void" => Type::Void,
            _ => Type::Custom(s.to_string()),
        }
    }
}

#[derive(Debug, Clone)]
pub struct TypeDefinition {
    pub name: String,
    pub depends: HashSet<Type>,
    pub fields: Vec<Field>,
    pub position: ParserPosition,
}

impl Definition for TypeDefinition {
    fn get_position(&self) -> ParserPosition {
        self.position.clone()
    }
}

#[derive(Debug, Clone)]
pub struct Field {
    pub name: String,
    pub typ: Type,
    pub array: bool,
    pub optional: bool,
    pub map: Option<Type>,
}

#[derive(Debug, Clone)]
pub struct EnumDefinition {
    pub name: String,
    pub values: Vec<EnumField>,
    pub position: ParserPosition,
}

impl Definition for EnumDefinition {
    fn get_position(&self) -> ParserPosition {
        self.position.clone()
    }
}

#[derive(Debug, Clone)]
pub struct EnumField {
    pub name: String,
    pub value: i32,
}

#[derive(Debug, Clone)]
pub struct ServiceDefinition {
    pub name: String,
    pub depends: HashSet<Type>,
    pub methods: Vec<Method>,
    pub position: ParserPosition,
}

impl Definition for ServiceDefinition {
    fn get_position(&self) -> ParserPosition {
        self.position.clone()
    }
}

#[derive(Debug, Clone)]
pub struct Method {
    pub name: String,
    pub inputs: Vec<MethodInput>,
    pub output: Option<MethodOutput>,
    pub decorators: MethodDecorators,
}

#[derive(Debug, Clone)]
pub struct MethodInput {
    pub name: String,
    pub typ: Type,
    pub array: bool,
    pub optional: bool,
}

#[derive(Debug, Clone)]
pub struct MethodOutput {
    pub typ: Type,
    pub array: bool,
}

#[derive(Debug, Clone)]
pub struct MethodDecorators {
    pub description: Option<String>,
    pub parameter_descriptions: HashMap<String, String>,
    pub return_description: Option<String>,
}

fn build_type(stmt: &TypeStatement) -> Result<TypeDefinition> {
    let mut typedef = TypeDefinition {
        position: stmt.position.clone(),
        name: stmt.name.clone(),
        depends: HashSet::new(),
        fields: Vec::new(),
    };

    for field in &stmt.fields {
        let typ = Type::from(&field.fieldtype);
        typedef.depends.insert(typ.clone());

        if let Some(maptype) = &field.map {
            if maptype != "string" && maptype != "int" {
                return Err(IRError::new("Map type must be string or int", field).into());
            }
        }

        typedef.fields.push(Field {
            name: field.name.clone(),
            typ: typ.clone(),
            array: field.array,
            optional: field.optional,
            map: field.map.as_ref().map(|s| Type::from(s)),
        });
    }

    Ok(typedef)
}

fn build_enum(stmt: &EnumStatement) -> Result<EnumDefinition> {
    let mut enumdef = EnumDefinition {
        position: stmt.position.clone(),
        name: stmt.name.clone(),
        values: Vec::new(),
    };

    let mut last = -1 as i32;
    for field in &stmt.values {
        let value = if let Some(value) = field.value {
            if value > std::i32::MAX as i64 {
                return Err(IRError::new("Enum value too large", field).into());
            }
            let value = value as i32;
            if value <= last {
                return Err(IRError::new("Enum values must be increasing", field).into());
            }
            last = value;
            value
        } else {
            last = last + 1;
            last
        };
        enumdef.values.push(EnumField {
            name: field.name.clone(),
            value,
        });
    }
    Ok(enumdef)
}

fn build_service(stmt: &ServiceStatement) -> Result<ServiceDefinition> {
    let mut servdef = ServiceDefinition {
        position: stmt.position.clone(),
        name: stmt.name.clone(),
        depends: HashSet::new(),
        methods: Vec::new(),
    };

    for method in &stmt.methods {
        let mut methoddef = Method {
            name: method.name.clone(),
            inputs: Vec::new(),
            output: method.return_type.as_ref().map(|rt| {
                let typ = Type::from(&rt.fieldtype);
                if typ != Type::Void {
                    servdef.depends.insert(typ.clone());
                }
                MethodOutput {
                    typ,
                    array: rt.array,
                }
            }),
            decorators: MethodDecorators {
                description: None,
                parameter_descriptions: HashMap::new(),
                return_description: None,
            },
        };

        let mut optional_starts = false;
        for inp in &method.inputs {
            let typ = Type::from(&inp.fieldtype);
            servdef.depends.insert(typ.clone());

            if optional_starts && !inp.optional {
                return Err(
                    IRError::new("Optional fields must come after required fields", inp).into(),
                );
            }

            if inp.optional {
                optional_starts = true;
            }

            methoddef.inputs.push(MethodInput {
                name: inp.name.clone(),
                typ,
                array: inp.array,
                optional: inp.optional,
            });
        }

        for decorator in &method.decorators {
            match decorator.name.as_str() {
                "Description" => {
                    if methoddef.decorators.description.is_some() {
                        return Err(IRError::new("Duplicate description", decorator).into());
                    }

                    if decorator.args.len() != 1 {
                        return Err(IRError::new(
                            "Description must have exactly one argument",
                            decorator,
                        )
                        .into());
                    }

                    methoddef.decorators.description = Some(decorator.args[0].clone());
                }
                "Returns" => {
                    if methoddef.decorators.return_description.is_some() {
                        return Err(IRError::new("Duplicate return description", decorator).into());
                    }

                    if decorator.args.len() != 1 {
                        return Err(IRError::new(
                            "Returns must have exactly one argument",
                            decorator,
                        )
                        .into());
                    }

                    methoddef.decorators.return_description = Some(decorator.args[0].clone());
                }
                "Param" => {
                    if decorator.args.len() != 2 {
                        return Err(IRError::new(
                            "Param must have exactly two arguments",
                            decorator,
                        )
                        .into());
                    }

                    let name = decorator.args[0].clone();
                    let description = decorator.args[1].clone();

                    if methoddef
                        .decorators
                        .parameter_descriptions
                        .contains_key(&name)
                    {
                        return Err(
                            IRError::new("Duplicate parameter description", decorator).into()
                        );
                    }

                    if methoddef.inputs.iter().find(|i| i.name == name).is_none() {
                        return Err(IRError::new("Parameter not found", decorator).into());
                    }

                    methoddef
                        .decorators
                        .parameter_descriptions
                        .insert(name, description);
                }
                _ => {
                    return Err(IRError::new("Unknown decorator", decorator).into());
                }
            }
        }
        servdef.methods.push(methoddef);
    }

    Ok(servdef)
}

pub fn build_ir(root: &Vec<RootNode>) -> Result<IR> {
    let mut options = HashMap::<String, String>::new();
    let mut steps = Vec::new();

    for node in root {
        match node {
            RootNode::Type(stmt) => steps.push(Step::Type(build_type(stmt)?)),
            RootNode::Enum(stmt) => steps.push(Step::Enum(build_enum(stmt)?)),
            RootNode::Service(stmt) => steps.push(Step::Service(build_service(stmt)?)),
            RootNode::Define(stmt) => {
                if options.contains_key(&stmt.key) {
                    return Err(IRError::new("Duplicate define", stmt).into());
                }

                if (stmt.key == "use_messagepack" || stmt.key == "allow_bytes")
                    && stmt.value == "true"
                {
                    options.insert("allow_bytes".to_owned(), "true".to_owned());
                }
                options.insert(stmt.key.clone(), stmt.value.clone());
            }
            RootNode::Import(_) => {
                panic!("Import not supported at this stage!");
            }
        }
    }

    let mut all_types = HashSet::<String>::new();
    let mut serv_types = HashSet::<String>::new();

    for bi in &BUILT_INS {
        all_types.insert(bi.to_string());
    }

    for step in &steps {
        match step {
            Step::Type(typedef) => {
                if all_types.contains(&typedef.name) {
                    return Err(IRError::new_from_def("Duplicate type", typedef).into());
                }

                all_types.insert(typedef.name.clone());
            }
            Step::Enum(enumdef) => {
                if all_types.contains(&enumdef.name) {
                    return Err(IRError::new_from_def("Duplicate type", enumdef).into());
                }

                all_types.insert(enumdef.name.clone());
            }
            Step::Service(servdef) => {
                if serv_types.contains(&servdef.name) {
                    return Err(IRError::new_from_def("Duplicate type", servdef).into());
                }

                serv_types.insert(servdef.name.clone());
            }
        }
    }

    // Verify dependencies

    for step in &steps {
        match step {
            Step::Type(typedef) => {
                for dep in &typedef.depends {
                    if let Type::Custom(dep) = dep {
                        if !all_types.contains(dep) {
                            return Err(IRError::new_from_def(
                                &format!("Type {} depends on unknown type {}", typedef.name, dep),
                                typedef,
                            )
                            .into());
                        }
                    }
                }
            }
            Step::Service(servdef) => {
                for dep in &servdef.depends {
                    if let Type::Custom(dep) = dep {
                        if !all_types.contains(dep) {
                            return Err(IRError::new_from_def(
                                &format!(
                                    "Service {} depends on unknown type {}",
                                    servdef.name, dep
                                ),
                                servdef,
                            )
                            .into());
                        }
                    }
                }
            }
            _ => {}
        }
    }

    Ok(IR { options, steps })
}

#[derive(Debug, Clone)]
pub struct IRError {
    pub message: String,
    pub position: ParserPosition,
}

impl IRError {
    fn new(msg: &str, node: &impl Node) -> IRError {
        IRError {
            message: format!("{}: {}", msg, node.get_name()),
            position: node.get_position(),
        }
    }

    fn new_from_def(msg: &str, def: &impl Definition) -> IRError {
        IRError {
            message: msg.to_string(),
            position: def.get_position(),
        }
    }
}

impl Error for IRError {}
impl Display for IRError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ParserError: {} at {:?}", self.message, self.position)
    }
}