use crate::schema::metadata::{MetadataEnvelope, TypeId};
use crate::schema::schema_type::SchemaType;
use crate::schema::schema_value::SchemaValue;
use golem_schema_derive::{FromSchema, IntoSchema};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, IntoSchema, FromSchema)]
#[cfg_attr(feature = "full", derive(desert_rust::BinaryCodec))]
#[cfg_attr(feature = "full", desert(evolution()))]
#[serde(rename_all = "camelCase")]
#[schema(named = "schema-graph")]
#[cfg_attr(feature = "full", derive(golem_schema_derive::PoemSchema))]
pub struct SchemaGraph {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub defs: Vec<SchemaTypeDef>,
pub root: SchemaType,
}
impl SchemaGraph {
pub fn anonymous(root: SchemaType) -> Self {
Self {
defs: Vec::new(),
root,
}
}
pub fn empty() -> Self {
Self {
defs: Vec::new(),
root: SchemaType::Record {
fields: Vec::new(),
metadata: MetadataEnvelope::default(),
},
}
}
pub fn lookup(&self, id: &TypeId) -> Option<&SchemaTypeDef> {
self.defs.iter().find(|d| &d.id == id)
}
pub fn contains_stream_type(&self) -> bool {
let mut pending = vec![&self.root];
let mut visited = HashSet::new();
while let Some(ty) = pending.pop() {
match ty {
SchemaType::Stream { .. } => return true,
SchemaType::Ref { id, .. } => {
if visited.insert(id)
&& let Some(def) = self.lookup(id)
{
pending.push(&def.body);
}
}
SchemaType::Record { fields, .. } => {
pending.extend(fields.iter().map(|field| &field.body));
}
SchemaType::Variant { cases, .. } => {
pending.extend(cases.iter().filter_map(|case| case.payload.as_ref()));
}
SchemaType::Tuple { elements, .. } => pending.extend(elements),
SchemaType::List { element, .. } | SchemaType::FixedList { element, .. } => {
pending.push(element);
}
SchemaType::Map { key, value, .. } => {
pending.extend([key.as_ref(), value.as_ref()])
}
SchemaType::Option { inner, .. } => pending.push(inner),
SchemaType::Result { spec, .. } => {
pending.extend(spec.ok.as_deref());
pending.extend(spec.err.as_deref());
}
SchemaType::Union { spec, .. } => {
pending.extend(spec.branches.iter().map(|branch| &branch.body));
}
SchemaType::Future { inner, .. } => pending.extend(inner.as_deref()),
SchemaType::Secret { spec, .. } => pending.push(&spec.inner),
SchemaType::Bool { .. }
| SchemaType::S8 { .. }
| SchemaType::S16 { .. }
| SchemaType::S32 { .. }
| SchemaType::S64 { .. }
| SchemaType::U8 { .. }
| SchemaType::U16 { .. }
| SchemaType::U32 { .. }
| SchemaType::U64 { .. }
| SchemaType::F32 { .. }
| SchemaType::F64 { .. }
| SchemaType::Char { .. }
| SchemaType::String { .. }
| SchemaType::Enum { .. }
| SchemaType::Flags { .. }
| SchemaType::Text { .. }
| SchemaType::Binary { .. }
| SchemaType::Path { .. }
| SchemaType::Url { .. }
| SchemaType::Uuid { .. }
| SchemaType::Datetime { .. }
| SchemaType::Duration { .. }
| SchemaType::Quantity { .. }
| SchemaType::QuotaToken { .. }
| SchemaType::PermissionCard { .. } => {}
}
}
false
}
pub fn resolve_ref<'a>(
&'a self,
ty: &'a SchemaType,
) -> Result<&'a SchemaType, RefResolutionError> {
let mut visiting: Vec<TypeId> = Vec::new();
let mut current = ty;
loop {
match current {
SchemaType::Ref { id, .. } => {
if visiting.iter().any(|x| x == id) {
return Err(RefResolutionError::RecursiveRef(id.clone()));
}
let def = self
.lookup(id)
.ok_or_else(|| RefResolutionError::DanglingRef(id.clone()))?;
visiting.push(id.clone());
current = &def.body;
}
other => return Ok(other),
}
}
}
}
pub struct GraphIndex<'a> {
defs: &'a [SchemaTypeDef],
index: Option<HashMap<&'a str, &'a SchemaTypeDef>>,
}
impl<'a> GraphIndex<'a> {
const INDEX_THRESHOLD: usize = 8;
pub fn new(graph: &'a SchemaGraph) -> Self {
let defs = graph.defs.as_slice();
let index = (defs.len() > Self::INDEX_THRESHOLD).then(|| {
let mut map = HashMap::with_capacity(defs.len());
for def in defs {
map.entry(def.id.as_str()).or_insert(def);
}
map
});
Self { defs, index }
}
pub fn defs_len(&self) -> usize {
self.defs.len()
}
pub fn lookup(&self, id: &TypeId) -> Option<&'a SchemaTypeDef> {
match &self.index {
Some(map) => map.get(id.as_str()).copied(),
None => self.defs.iter().find(|d| &d.id == id),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RefResolutionError {
RecursiveRef(TypeId),
DanglingRef(TypeId),
}
impl std::fmt::Display for RefResolutionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::RecursiveRef(id) => {
write!(f, "schema graph contains a recursive reference: {id}")
}
Self::DanglingRef(id) => write!(f, "schema graph contains dangling reference: {id}"),
}
}
}
impl std::error::Error for RefResolutionError {}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, IntoSchema, FromSchema)]
#[cfg_attr(feature = "full", derive(desert_rust::BinaryCodec))]
#[cfg_attr(feature = "full", desert(evolution()))]
#[serde(rename_all = "camelCase")]
#[schema(named = "schema-type-def")]
#[cfg_attr(feature = "full", derive(golem_schema_derive::PoemSchema))]
pub struct SchemaTypeDef {
pub id: TypeId,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
pub body: SchemaType,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, IntoSchema, FromSchema)]
#[cfg_attr(feature = "full", derive(desert_rust::BinaryCodec))]
#[cfg_attr(feature = "full", desert(evolution()))]
#[serde(rename_all = "camelCase")]
#[schema(named = "typed-schema-value")]
#[cfg_attr(feature = "full", derive(golem_schema_derive::PoemSchema))]
pub struct TypedSchemaValue {
graph: SchemaGraph,
value: SchemaValue,
}
impl TypedSchemaValue {
pub fn new(graph: SchemaGraph, value: SchemaValue) -> Self {
Self { graph, value }
}
pub fn graph(&self) -> &SchemaGraph {
&self.graph
}
pub fn value(&self) -> &SchemaValue {
&self.value
}
pub fn root_type(&self) -> &SchemaType {
&self.graph.root
}
pub fn into_parts(self) -> (SchemaGraph, SchemaValue) {
(self.graph, self.value)
}
}
pub fn reachable_defs<'a>(graph: &'a SchemaGraph, root: &'a SchemaType) -> Vec<SchemaTypeDef> {
let mut stack: Vec<&'a TypeId> = Vec::new();
collect_refs(root, &mut stack);
if stack.is_empty() {
return Vec::new();
}
let index = GraphIndex::new(graph);
let mut reachable: HashSet<&'a str> = HashSet::new();
while let Some(id) = stack.pop() {
if !reachable.insert(id.as_str()) {
continue;
}
if let Some(def) = index.lookup(id) {
collect_refs(&def.body, &mut stack);
}
}
graph
.defs
.iter()
.filter(|d| reachable.remove(d.id.as_str()))
.cloned()
.collect()
}
fn collect_refs<'a>(ty: &'a SchemaType, out: &mut Vec<&'a TypeId>) {
match ty {
SchemaType::Ref { id, .. } => out.push(id),
SchemaType::Record { fields, .. } => {
for field in fields {
collect_refs(&field.body, out);
}
}
SchemaType::Variant { cases, .. } => {
for case in cases {
if let Some(payload) = &case.payload {
collect_refs(payload, out);
}
}
}
SchemaType::Tuple { elements, .. } => {
for element in elements {
collect_refs(element, out);
}
}
SchemaType::List { element, .. } | SchemaType::FixedList { element, .. } => {
collect_refs(element, out);
}
SchemaType::Map { key, value, .. } => {
collect_refs(key, out);
collect_refs(value, out);
}
SchemaType::Option { inner, .. } => collect_refs(inner, out),
SchemaType::Result { spec, .. } => {
if let Some(ok) = &spec.ok {
collect_refs(ok, out);
}
if let Some(err) = &spec.err {
collect_refs(err, out);
}
}
SchemaType::Union { spec, .. } => {
for branch in &spec.branches {
collect_refs(&branch.body, out);
}
}
SchemaType::Future { inner, .. } | SchemaType::Stream { inner, .. } => {
if let Some(inner) = inner {
collect_refs(inner, out);
}
}
SchemaType::Secret { spec, .. } => collect_refs(&spec.inner, out),
SchemaType::Bool { .. }
| SchemaType::S8 { .. }
| SchemaType::S16 { .. }
| SchemaType::S32 { .. }
| SchemaType::S64 { .. }
| SchemaType::U8 { .. }
| SchemaType::U16 { .. }
| SchemaType::U32 { .. }
| SchemaType::U64 { .. }
| SchemaType::F32 { .. }
| SchemaType::F64 { .. }
| SchemaType::Char { .. }
| SchemaType::String { .. }
| SchemaType::Enum { .. }
| SchemaType::Flags { .. }
| SchemaType::Text { .. }
| SchemaType::Binary { .. }
| SchemaType::Path { .. }
| SchemaType::Url { .. }
| SchemaType::Uuid { .. }
| SchemaType::Datetime { .. }
| SchemaType::Duration { .. }
| SchemaType::Quantity { .. }
| SchemaType::QuotaToken { .. }
| SchemaType::PermissionCard { .. } => {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use test_r::test;
#[test]
fn stream_detection_follows_reachable_recursive_definitions() {
let recursive = TypeId::new("recursive");
let output = TypeId::new("output");
let mut graph = SchemaGraph {
root: SchemaType::ref_to(recursive.clone()),
defs: vec![
SchemaTypeDef {
id: recursive.clone(),
name: None,
body: SchemaType::tuple(vec![SchemaType::ref_to(recursive)]),
},
SchemaTypeDef {
id: output.clone(),
name: None,
body: SchemaType::option(SchemaType::stream(Some(SchemaType::u32()))),
},
],
};
assert!(
!graph.contains_stream_type(),
"unreachable streams do not count"
);
let SchemaType::Tuple { elements, .. } = &mut graph.defs[0].body else {
unreachable!()
};
elements.push(SchemaType::ref_to(output));
assert!(
graph.contains_stream_type(),
"a cycle must not hide a sibling stream"
);
assert!(SchemaGraph::anonymous(SchemaType::stream(None)).contains_stream_type());
}
fn def(id: &str, name: Option<&str>) -> SchemaTypeDef {
SchemaTypeDef {
id: TypeId::new(id),
name: name.map(|n| n.to_string()),
body: SchemaType::bool(),
}
}
fn assert_index_matches_linear(graph: &SchemaGraph) {
let index = GraphIndex::new(graph);
let mut ids: Vec<TypeId> = graph.defs.iter().map(|d| d.id.clone()).collect();
ids.push(TypeId::new("definitely-absent-id"));
for id in &ids {
assert_eq!(
index.lookup(id),
graph.lookup(id),
"GraphIndex and SchemaGraph disagree on id {id:?}"
);
}
}
#[test]
fn graph_index_matches_linear_small_with_duplicates() {
let graph = SchemaGraph {
defs: vec![
def("a", Some("first-a")),
def("b", Some("b")),
def("a", Some("second-a")),
],
root: SchemaType::bool(),
};
let index = GraphIndex::new(&graph);
assert!(
index.index.is_none(),
"small graph must use linear fallback"
);
assert_eq!(
index
.lookup(&TypeId::new("a"))
.and_then(|d| d.name.as_deref()),
Some("first-a"),
"first definition must win on duplicate ids"
);
assert_index_matches_linear(&graph);
}
#[test]
fn graph_index_matches_linear_wide_with_duplicates() {
let mut defs: Vec<SchemaTypeDef> = (0..20)
.map(|i| def(&format!("t{i:02}"), Some(&format!("name-{i:02}"))))
.collect();
defs.push(def("t05", Some("duplicate-of-t05")));
let graph = SchemaGraph {
defs,
root: SchemaType::bool(),
};
let index = GraphIndex::new(&graph);
assert!(index.index.is_some(), "wide graph must build an index");
assert_eq!(
index
.lookup(&TypeId::new("t05"))
.and_then(|d| d.name.as_deref()),
Some("name-05"),
"first definition must win on duplicate ids in the indexed branch"
);
assert_index_matches_linear(&graph);
}
}