JsonRPC-Rust/libjrpc/src/ir.rs

use std::{
    collections::{HashMap, HashSet},
    error::Error,
    fmt::{Debug, 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: Debug {
    fn get_position(&self) -> ParserPosition;
    fn get_name(&self) -> String;
}

#[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 BaseType {
    Int,
    Float,
    String,
    Bool,
    Bytes,
    Void,
    Custom(String),
}

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

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

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

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

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

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum TypeModifier {
    None,
    Array,
    Optional,
    Map(BaseType),
    OptionalArray,
    OptionalMap(BaseType),
    OptionalMapArray(BaseType),
    MapArray(BaseType),
}

impl TypeModifier {
    pub fn from_flags(optional: bool, array: bool, map: Option<BaseType>) -> Self {
        match (optional, array, map) {
            (false, false, None) => Self::None,
            (false, true, None) => Self::Array,
            (true, false, None) => Self::Optional,
            (true, true, None) => Self::OptionalArray,
            (false, false, Some(map)) => Self::Map(map),
            (false, true, Some(map)) => Self::MapArray(map),
            (true, false, Some(map)) => Self::OptionalMap(map),
            (true, true, Some(map)) => Self::OptionalMapArray(map),
        }
    }

    pub fn get_flags(&self) -> (bool, bool, Option<BaseType>) {
        match self.clone() {
            Self::None => (false, false, None),
            Self::Array => (false, true, None),
            Self::Optional => (true, false, None),
            Self::OptionalArray => (true, true, None),
            Self::Map(map) => (false, false, Some(map)),
            Self::MapArray(map) => (false, true, Some(map)),
            Self::OptionalMap(map) => (true, false, Some(map)),
            Self::OptionalMapArray(map) => (true, true, Some(map)),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Type(pub BaseType, pub TypeModifier);

impl Type {
    pub fn from_flags(base: BaseType, optional: bool, array: bool, map: Option<BaseType>) -> Self {
        Self(base, TypeModifier::from_flags(optional, array, map))
    }

    /// Returns flags in the order of: (OPTIONAL, ARRAY, MAP)
    pub fn into_flags(self) -> (BaseType, bool, bool, Option<BaseType>) {
        let b = self.0.clone();
        match self.1 {
            TypeModifier::None => (b, false, false, None),
            TypeModifier::Array => (b, false, true, None),
            TypeModifier::Optional => (b, true, false, None),
            TypeModifier::OptionalArray => (b, true, true, None),
            TypeModifier::Map(map) => (b, false, false, Some(map)),
            TypeModifier::MapArray(map) => (b, false, true, Some(map)),
            TypeModifier::OptionalMap(map) => (b, true, false, Some(map)),
            TypeModifier::OptionalMapArray(map) => (b, true, true, Some(map)),
        }
    }

    pub fn is_optional(&self) -> bool {
        self.1.get_flags().0
    }

    pub fn is_array(&self) -> bool {
        self.1.get_flags().1
    }

    pub fn is_map(&self) -> bool {
        self.1.get_flags().2.is_some()
    }

    pub fn get_map(&self) -> Option<BaseType> {
        self.1.get_flags().2.clone()
    }
}
// {
//     pub base: BaseType,
//     pub array: bool,
//     pub optional: bool,
//     pub map: Option<BaseType>,
// }

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

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

#[derive(Debug, Clone)]
pub struct Field {
    pub name: String,
    pub typ: 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()
    }
    fn get_name(&self) -> String {
        self.name.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<BaseType>,
    pub methods: Vec<Method>,
    pub position: ParserPosition,
}

impl Definition for ServiceDefinition {
    fn get_position(&self) -> ParserPosition {
        self.position.clone()
    }
    fn get_name(&self) -> String {
        self.name.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,
}

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

#[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 = BaseType::from(&field.fieldtype);
        if stmt.name != 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: Type::from_flags(
                typ.clone(),
                field.optional,
                field.array,
                field.map.as_ref().map(|s| BaseType::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_flags(BaseType::from(&rt.fieldtype), false, rt.array, None);
                if typ.0 != BaseType::Void {
                    servdef.depends.insert(typ.0.clone());
                }
                MethodOutput { typ }
            }),
            decorators: MethodDecorators {
                description: None,
                parameter_descriptions: HashMap::new(),
                return_description: None,
            },
        };

        let mut optional_starts = false;
        for inp in &method.inputs {
            let typ = BaseType::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: Type::from_flags(typ.clone(), inp.optional, inp.array, None),
            });
        }

        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 BaseType::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 BaseType::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)
    }
}