JsonRPC-Rust/libjrpc/src/targets/rust.rs

use anyhow::Result;
use log::warn;
use std::collections::{BTreeMap, BTreeSet};

use crate::compile::{Compile, CompileContext, FileGenerator};
use crate::ir::{BaseType, EnumDefinition, ServiceDefinition, Step, Type, TypeDefinition};
use crate::shared::Keywords;
use crate::IR;

pub struct RustCompiler {
    crate_name: String,
}

static RUST_KEYWORDS: [&'static str; 6] = ["type", "return", "static", "pub", "enum", "self"];

impl RustCompiler {
    fn type_to_rust(typ: &BaseType) -> String {
        match typ {
            BaseType::String => "String".to_string(),
            BaseType::Int => "i64".to_string(),
            BaseType::Float => "f64".to_string(),
            BaseType::Bool => "bool".to_string(),
            BaseType::Bytes => "Vec<u8>".to_string(),
            BaseType::Void => "()".to_string(),
            BaseType::Custom(name) => name.clone(),
        }
    }

    fn type_to_rust_ext(typ: &Type) -> String {
        let mut result = Self::type_to_rust(&typ.0);

        let (optional, array, map) = typ.1.get_flags();

        if array {
            result = format!("Vec<{}>", result);
        }
        if let Some(map) = &map {
            result = format!("HashMap<{}, {}>", Self::type_to_rust(&map), result);
        }
        if optional {
            result = format!("Option<{}>", result);
        }
        result
    }

    fn add_dependencies(
        &mut self,
        file: &mut FileGenerator,
        depends: &BTreeSet<BaseType>,
    ) -> Result<()> {
        for dep in depends {
            match dep {
                BaseType::Custom(name) => {
                    file.a0(&format!("use crate::{};", name));
                }
                _ => {}
            }
        }
        file.a0("");
        file.a0("");
        Ok(())
    }

    fn fix_keyword_name(name: &str) -> String {
        if RUST_KEYWORDS.contains(&name) {
            format!("{}_", name)
        } else {
            name.to_string()
        }
    }

    fn generate_service_lib(ctx: &CompileContext, ir: &IR) -> Result<()> {
        let mut f = FileGenerator::new();
        let mut fc = FileGenerator::new();
        let mut fs = FileGenerator::new();

        f.a0("pub mod base_lib;");
        f.a0("pub use base_lib::{JRPCServer, JRPCClient, Result};");

        for step in ir.steps.iter() {
            match step {
                Step::Type(def) => {
                    f.a0(format!("mod {};", CompileContext::to_snake(&def.name)));
                    f.a(
                        0,
                        format!(
                            "pub use {}::{};",
                            CompileContext::to_snake(&def.name),
                            def.name
                        ),
                    );
                }
                Step::Enum(def) => {
                    f.a0(format!("mod {};", CompileContext::to_snake(&def.name)));
                    f.a(
                        0,
                        format!(
                            "pub use {}::{};",
                            CompileContext::to_snake(&def.name),
                            def.name
                        ),
                    );
                }
                Step::Service(def) => {
                    fs.a0(format!("mod {};", CompileContext::to_snake(&def.name)));
                    fs.a0(format!(
                        "pub use {}::{{ {}, {}Handler }};",
                        CompileContext::to_snake(&def.name),
                        def.name,
                        def.name
                    ));

                    fc.a0(format!("mod {};", CompileContext::to_snake(&def.name)));
                    fc.a0(format!(
                        "pub use {}::{};",
                        CompileContext::to_snake(&def.name),
                        def.name,
                    ));
                }
            }
        }

        f.a0("pub mod server;");
        f.a0("pub mod client;");

        ctx.write_file("src/lib.rs", &f.get_content())?;
        ctx.write_file("src/server/mod.rs", &fs.get_content())?;
        ctx.write_file("src/client/mod.rs", &fc.get_content())?;

        Ok(())
    }

    fn generate_service_server(
        &mut self,
        ctx: &mut CompileContext,
        definition: &ServiceDefinition,
    ) -> anyhow::Result<()> {
        let mut f = FileGenerator::new();

        self.add_dependencies(&mut f, &definition.depends)?;

        f.a0("use crate::base_lib::{JRPCServerService, JRPCRequest, Result};");
        f.a0("use serde_json::Value;");
        f.a0("use std::sync::Arc;");
        f.a0("use async_trait::async_trait;");

        f.a0("#[async_trait]");
        f.a0(format!("pub trait {} {{", definition.name));
        f.a1("type Context: Clone + Sync + Send + 'static;");
        for method in definition.methods.iter() {
            let mut params = method
                .inputs
                .iter()
                .map(|arg| {
                    format!(
                        "{}: {}",
                        Self::fix_keyword_name(&arg.name),
                        Self::type_to_rust_ext(&arg.typ)
                    )
                })
                .collect::<Vec<String>>();
            params.push("ctx: Self::Context".to_string());
            let params = params.join(", ");

            let ret = method
                .output
                .as_ref()
                .map_or_else(|| "()".to_owned(), |r| Self::type_to_rust_ext(&r.typ));

            f.a1("#[allow(non_snake_case)]");
            f.a(
                1,
                format!(
                    "async fn {}(&self, {}) -> Result<{}>;",
                    Self::fix_keyword_name(&method.name),
                    params,
                    ret
                ),
            );
        }
        f.a0("}");

        f.a0("");
        f.a0(format!("pub struct {}Handler<Context> {{", definition.name));
        f.a1(format!(
            "implementation: Box<dyn {}<Context = Context> + Sync + Send + 'static>,",
            definition.name
        ));
        f.a0("}");
        f.a0("");

        f.a0(format!(
            "impl<Context: Clone + Sync + Send + 'static> {}Handler<Context> {{",
            definition.name
        ));
        f.a1(format!(
            "pub fn new(implementation: Box<dyn {}<Context = Context> + Sync + Send + 'static>) -> Arc<Self> {{",
            definition.name,
        ));
        f.a2("Arc::from(Self { implementation })");
        f.a1("}");
        f.a0("}");

        f.a0("");

        f.a0("#[async_trait]");
        f.a0(format!(
            "impl<Context: Clone + Sync + Send + 'static> JRPCServerService for {}Handler<Context> {{",
            definition.name
        ));
        f.a1("type Context = Context;");
        f.a1(format!(
            "fn get_id(&self) -> String {{ \"{}\".to_owned() }} ",
            definition.name
        ));

        f.a1("");

        f.a1("#[allow(non_snake_case)]");
        f.a1(
            "async fn handle(&self, msg: &JRPCRequest, function: &str, ctx: Self::Context) -> Result<(bool, Value)> {",
        );
        f.a2("match function {");

        // TODO: Implement optional arguments!

        for method in &definition.methods {
            f.a3(format!(
                "\"{}\" => {{",
                Self::fix_keyword_name(&method.name)
            ));
            if method.inputs.len() < 1 {
                f.a5(format!(
                    "let res = self.implementation.{}(ctx).await?;",
                    method.name
                ));
                f.a5("Ok((true, serde_json::to_value(res)?))");
            } else {
                f.a4("if msg.params.is_array() {");
                if method.inputs.len() > 0 {
                    f.a5(
                        "let arr = msg.params.as_array().unwrap(); //TODO: Check if this can fail.",
                    );
                }
                f.a5(format!("let res = self.implementation.{}(", method.name));
                for (i, arg) in method.inputs.iter().enumerate() {
                    f.a6(format!("serde_json::from_value(arr[{}].clone())", i));
                    f.a(
                        7,
                        format!(
                            ".map_err(|_| \"Parameter for field '{}' should be of type '{}'!\")?{}",
                            arg.name,
                            arg.typ.0.to_string(),
                            if i < method.inputs.len() - 1 { "," } else { "" }
                        ),
                    );
                }
                f.a5(",ctx).await?;");

                if let Some(_output) = &method.output {
                    f.a5("Ok((true, serde_json::to_value(res)?))");
                } else {
                    f.a5("_ = res;");
                    f.a5("Ok((true, Value::Null))");
                }

                f.a4("} else if msg.params.is_object() {");
                f.a5("let obj = msg.params.as_object().unwrap();  //TODO: Check if this can fail.");
                f.a5(format!("let res = self.implementation.{}(", method.name));
                for (i, arg) in method.inputs.iter().enumerate() {
                    f.a6(format!(
                    "serde_json::from_value(obj.get(\"{}\").ok_or(\"Parameter of field '{}' missing!\")?.clone())",
                    arg.name,
                    arg.name
                ));
                    f.a(
                        7,
                        format!(
                            ".map_err(|_| \"Parameter for field {} should be of type '{}'!\")?{}",
                            arg.name,
                            arg.typ.0.to_string(),
                            if i < method.inputs.len() - 1 { "," } else { "" }
                        ),
                    );
                }
                f.a5(", ctx).await?;");
                if let Some(_output) = &method.output {
                    f.a5("Ok((true, serde_json::to_value(res)?))");
                } else {
                    f.a5("Ok((false, Value::Null))");
                }
                f.a4("} else {");
                f.a5("Err(Box::from(\"Invalid parameters??\".to_owned()))");
                f.a4("}");
            }
            f.a3("}");
        }
        f.a3("_ => { Err(Box::from(format!(\"Invalid function {}\", function).to_owned())) },");

        f.a2("}");
        f.a1("}");
        f.a0("}");

        ctx.write_file(
            &format!(
                "src/server/{}.rs",
                CompileContext::to_snake(&definition.name)
            ),
            &f.into_content(),
        )?;

        Ok(())
    }

    fn generate_service_client(
        &mut self,
        ctx: &mut CompileContext,
        definition: &ServiceDefinition,
    ) -> anyhow::Result<()> {
        let mut f = FileGenerator::new();

        self.add_dependencies(&mut f, &definition.depends)?;

        f.a0("use crate::base_lib::{JRPCClient, JRPCRequest, Result};");
        f.a0("use serde_json::json;");
        f.a0("");
        f.a0(format!("pub struct {} {{", definition.name));
        f.a1("client: JRPCClient,");
        f.a0("}");
        f.a0("");

        f.a0(format!("impl {} {{", definition.name));
        f.a1("pub fn new(client: JRPCClient) -> Self {");
        f.a2(format!("Self {{ client }}"));
        f.a1("}");

        f.a0("");

        for method in &definition.methods {
            let params = method
                .inputs
                .iter()
                .map(|arg| {
                    format!(
                        "{}: {}",
                        Self::fix_keyword_name(&arg.name),
                        Self::type_to_rust_ext(&arg.typ)
                    )
                })
                .collect::<Vec<String>>()
                .join(", ");
            let ret = method.output.as_ref().map_or("()".to_string(), |output| {
                Self::type_to_rust_ext(&output.typ)
            });

            f.a1("#[allow(non_snake_case)]");
            f.a1(format!(
                "pub async fn {}(&self, {}) -> Result<{}> {{",
                method.name, params, ret
            ));

            f.a2("let l_req = JRPCRequest {");
            f.a3("jsonrpc: \"2.0\".to_owned(),");
            f.a3("id: None, // 'id' will be set by the send_request function");
            f.a3(format!(
                "method: \"{}.{}\".to_owned(),",
                definition.name, method.name
            ));
            f.a3(format!(
                "params: json!([{}])",
                method
                    .inputs
                    .iter()
                    .map(|e| Self::fix_keyword_name(&e.name))
                    .collect::<Vec<String>>()
                    .join(", ")
            ));
            f.a2("};");

            if let Some(output) = &method.output {
                f.a2("let l_res = self.client.send_request(l_req).await;");
                f.a2("match l_res {");
                f.a3("Err(e) => Err(e),");
                if output.typ.0 == BaseType::Void {
                    f.a3("Ok(_) => Ok(())");
                } else {
                    f.a3("Ok(o) => serde_json::from_value(o).map_err(|e| Box::from(e))");
                }
                f.a2("}");
            } else {
                f.a2("self.client.send_notification(l_req).await;");
                f.a2("Ok(())");
            }
            f.a1("}");
        }

        f.a0("}");

        ctx.write_file(
            &format!(
                "src/client/{}.rs",
                CompileContext::to_snake(&definition.name)
            ),
            &f.into_content(),
        )?;

        Ok(())
    }
}

impl Compile for RustCompiler {
    fn new(options: &BTreeMap<String, String>) -> anyhow::Result<Self> {
        let crate_name = if let Some(crate_name) = options.get("rust_crate") {
            crate_name.to_string()
        } else {
            anyhow::bail!("crate option is required for rust compiler");
        };

        if let Some(allow_bytes) = options.get("allow_bytes") {
            if allow_bytes == "true" {
                anyhow::bail!("allow_bytes option is not supported for rust compiler");
            }
        }

        Ok(RustCompiler { crate_name })
    }

    fn name(&self) -> String {
        "rust".to_string()
    }

    fn start(&mut self, ctx: &mut CompileContext) -> anyhow::Result<()> {
        ctx.write_file(
            "Cargo.toml",
            &include_str!("../../templates/Rust/Cargo.toml")
                .to_owned()
                .replace("__name__", &self.crate_name),
        )?;

        ctx.write_file(
            "src/base_lib.rs",
            &include_str!("../../templates/Rust/src/lib.rs").to_owned(),
        )?;

        Ok(())
    }

    fn generate_type(
        &mut self,
        ctx: &mut CompileContext,
        definition: &TypeDefinition,
    ) -> anyhow::Result<()> {
        let mut f = FileGenerator::new();

        if definition
            .fields
            .iter()
            .any(|e| e.typ.1.get_flags().2.is_some())
        {
            f.a0("use std::collections::hash_map::HashMap;")
        }
        f.a0("use serde::{Deserialize, Serialize};");

        self.add_dependencies(&mut f, &definition.depends)?;

        f.a0("#[derive(Clone, Debug, Serialize, Deserialize)]");
        f.a0(format!("pub struct {} {{", definition.name));
        for field in definition.fields.iter() {
            f.a(1, "#[allow(non_snake_case)]");

            if Keywords::is_keyword(&field.name) {
                warn!(
                    "[RUST] Warning: Field name '{}' is not allowed in Rust. Renaming to '{}_'",
                    field.name, field.name
                );

                f.a(1, format!("#[serde(rename = \"{}\")]", field.name));
            }

            f.a1(format!(
                "pub {}: {},",
                Self::fix_keyword_name(&field.name),
                Self::type_to_rust_ext(&field.typ)
            ));
        }

        f.a0("}");

        ctx.write_file(
            &format!("src/{}.rs", CompileContext::to_snake(&definition.name)),
            &f.into_content(),
        )?;

        Ok(())
    }

    fn generate_enum(
        &mut self,
        ctx: &mut CompileContext,
        definition: &EnumDefinition,
    ) -> anyhow::Result<()> {
        let mut f = FileGenerator::new();
        f.a0("use int_enum::IntEnum;");
        // f.a0("use serde::{Deserialize, Serialize};");
        f.a0("");
        f.a0("");
        f.a0("#[repr(i64)]");
        f.a0("#[derive(Clone, Copy, Debug, Eq, PartialEq, IntEnum)]");
        f.a0(format!("pub enum {} {{", definition.name));
        for val in definition.values.iter() {
            f.a(1, format!("{} = {},", val.name, val.value));
        }
        f.a0("}");

        ctx.write_file(
            &format!("src/{}.rs", CompileContext::to_snake(&definition.name)),
            &f.into_content(),
        )?;

        Ok(())
    }

    fn generate_service(
        &mut self,
        ctx: &mut CompileContext,
        definition: &ServiceDefinition,
    ) -> anyhow::Result<()> {
        self.generate_service_client(ctx, definition)?;
        self.generate_service_server(ctx, definition)?;
        Ok(())
    }

    fn finalize(&mut self, ctx: &mut CompileContext, ir: &IR) -> anyhow::Result<()> {
        Self::generate_service_lib(ctx, ir)?;
        Ok(())
    }
}