sausage/sdk/src/discovery.rs

//! Schema contracts and OpenRPC document generation.
use crate::*;
use serde_json::{Value, json};

/// Types within the supported schema subset.
#[diagnostic::on_unimplemented(message = "this type is outside the Sausage schema subset")]
pub trait SchemaType: schemars::JsonSchema {}
/// Named parameter objects.
#[diagnostic::on_unimplemented(message = "RPC params must be a #[schema] struct with named fields")]
pub trait Params: SchemaType {}
macro_rules! primitive { ($($t:ty),*) => { $(impl SchemaType for $t {})* }; }
primitive!(
    String,
    bool,
    u8,
    u16,
    u32,
    u64,
    usize,
    i8,
    i16,
    i32,
    i64,
    isize,
    f32,
    f64,
    ()
);
impl<T: SchemaType> SchemaType for Vec<T> {}
impl<T: SchemaType> SchemaType for Option<T> {}

/// Create an OpenRPC document builder for generated interfaces.
#[doc(hidden)]
pub fn document(title: &str, version: &str) -> Discovery {
    fn generator(serialize: bool) -> schemars::SchemaGenerator {
        let mut settings = schemars::generate::SchemaSettings::draft07();
        settings.definitions_path = if serialize {
            "/sdk-output"
        } else {
            "/sdk-input"
        }
        .into();
        if serialize {
            settings.for_serialize()
        } else {
            settings.for_deserialize()
        }
        .into_generator()
    }
    Discovery {
        value: sausage_interface::Document::new(title, version).into_value(),
        input: generator(false),
        output: generator(true),
    }
}

/// Shared generators for all methods and endpoint types in one document.
#[doc(hidden)]
pub struct Discovery {
    /// The OpenRPC document before reference resolution.
    pub value: Value,
    input: schemars::SchemaGenerator,
    output: schemars::SchemaGenerator,
}
impl Discovery {
    /// Add an endpoint with the same input and output generators.
    pub fn endpoint<T: Interface>(&mut self, name: &str) {
        let methods = std::mem::replace(&mut self.value["methods"], json!([]));
        T::add_methods(self);
        let endpoint = std::mem::replace(&mut self.value["methods"], methods);
        self.value["x-sausage-endpoints"][name] = json!({"methods":endpoint});
    }
}

/// Add a method with independent deserialize and serialize contracts.
#[doc(hidden)]
pub fn method<P: Params, R: SchemaType>(doc: &mut Discovery, name: &str, docs: &str) {
    let params_ref = if P::schema_id() == <Empty as schemars::JsonSchema>::schema_id() {
        Value::Null
    } else {
        serde_json::to_value(doc.input.subschema_for::<P>()).unwrap()
    };
    let schema = params_ref["$ref"]
        .as_str()
        .and_then(|reference| reference.strip_prefix("#/sdk-input/"))
        .and_then(|name| doc.input.definitions().get(name))
        .unwrap_or(&Value::Null);
    let params:Vec<_>=schema["properties"].as_object().into_iter().flatten().map(|(name,field)| {
        json!({
            "name": name,
            "description": field.get("description").cloned().unwrap_or(Value::Null),
            "required": schema["required"].as_array().is_some_and(|required|required.contains(&json!(name))),
            "schema": field,
        })
    }).collect();
    let result_schema = serde_json::to_value(doc.output.subschema_for::<R>()).unwrap();
    let has_params = !params.is_empty();
    let mut method = json!({
        "name":name,
        "summary":docs.split("\n\n").next().unwrap_or_default(),
        "description":docs,
        "paramStructure":"by-name",
        "params":params,
        "result":{"name":"result","schema":result_schema},
    });
    if has_params {
        method["x-sausage-params"] = params_ref;
    }
    doc.value["methods"].as_array_mut().unwrap().push(method);
}

/// Resolve shared schema names after all methods contribute their contracts.
#[doc(hidden)]
pub fn finish_document(mut discovery: Discovery) -> Value {
    let input = serde_json::to_value(discovery.input.take_definitions(true)).unwrap();
    let output = serde_json::to_value(discovery.output.take_definitions(true)).unwrap();
    let mut doc = discovery.value;
    fn refs(value: &mut Value, names: &std::collections::HashMap<String, String>) {
        match value {
            Value::Object(map) => {
                if let Some(Value::String(reference)) = map.get_mut("$ref") {
                    if let Some(name) = reference.strip_prefix("#/sdk-input/") {
                        *reference = format!("#/components/schemas/{name}");
                    } else if let Some(name) = reference.strip_prefix("#/sdk-output/") {
                        *reference = format!(
                            "#/components/schemas/{}",
                            names.get(name).map(String::as_str).unwrap_or(name)
                        );
                    }
                }
                for child in map.values_mut() {
                    refs(child, names);
                }
            }
            Value::Array(values) => {
                for child in values {
                    refs(child, names);
                }
            }
            _ => {}
        }
    }
    let mut names = std::collections::HashMap::new();
    for (name, schema) in output.as_object().unwrap() {
        if let Some(other) = input.get(name) {
            let mut a = schema.clone();
            let mut b = other.clone();
            refs(&mut a, &names);
            refs(&mut b, &names);
            if a != b {
                names.insert(name.clone(), format!("{name}Result"));
            }
        }
    }
    // Propagate differences through references to nested contracts.
    loop {
        let before = names.len();
        for (name, schema) in output.as_object().unwrap() {
            if let Some(other) = input.get(name) {
                let mut a = schema.clone();
                let mut b = other.clone();
                refs(&mut a, &names);
                refs(&mut b, &names);
                if a != b {
                    names.insert(name.clone(), format!("{name}Result"));
                }
            }
        }
        if names.len() == before {
            break;
        }
    }
    doc["components"]["schemas"] = input;
    for (name, schema) in output.as_object().unwrap() {
        let name = names.get(name).unwrap_or(name);
        doc["components"]["schemas"]
            .as_object_mut()
            .unwrap()
            .insert(name.clone(), schema.clone());
    }
    refs(&mut doc, &names);
    for import in inventory::iter::<crate::Import> {
        doc["x-sausage-imports"][import.name] =
            serde_json::from_str(import.document).expect("validated import document");
    }
    doc
}

/// Empty parameters for methods without arguments.
#[derive(serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
pub struct Empty {}
impl SchemaType for Empty {}
impl Params for Empty {}