use std::collections::{BTreeMap, BTreeSet};
use std::convert::Infallible;
use helm_schema_json_schema_walk::{
SchemaTraversalContext, ref_points_inside, try_map_schema_context,
};
use serde_json::{Map, Value};
use crate::schema_doc::strip_ref;
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub(crate) struct SourceBundleLocation {
pub(crate) document: String,
pub(crate) pointer: String,
}
impl SourceBundleLocation {
pub(crate) fn new(document: impl Into<String>, pointer: impl Into<String>) -> Self {
Self {
document: document.into(),
pointer: pointer.into(),
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct SourceBundleNode {
pub(crate) location: SourceBundleLocation,
pub(crate) schema: Value,
}
impl SourceBundleNode {
pub(crate) fn new(
document: impl Into<String>,
pointer: impl Into<String>,
schema: Value,
) -> Self {
Self {
location: SourceBundleLocation::new(document, pointer),
schema,
}
}
}
pub(crate) fn bundle_source_definition(
document: &str,
pointer: &str,
source_schema: &Value,
mut resolve_external_ref: impl FnMut(&SourceBundleLocation, &str) -> Option<SourceBundleNode>,
) -> Value {
if helm_schema_json_schema_walk::schema_refs_point_inside(source_schema) {
return source_schema.clone();
}
bundle_source_schema(
SourceBundleNode::new(document, pointer, source_schema.clone()),
|current_location, reference| {
if let Some(pointer) = reference.strip_prefix('#')
&& source_schema.pointer(pointer).is_some()
{
return None;
}
resolve_external_ref(current_location, reference)
},
)
}
pub(crate) fn bundle_source_schema(
root: SourceBundleNode,
resolve_external_ref: impl FnMut(&SourceBundleLocation, &str) -> Option<SourceBundleNode>,
) -> Value {
let mut bundler = SourceSchemaBundler::new(root.schema.clone(), resolve_external_ref);
bundler.bundle_root(root)
}
struct SourceSchemaBundler<R> {
root_schema: Value,
resolve_external_ref: R,
definition_names_by_location: BTreeMap<SourceBundleLocation, String>,
definitions_by_name: BTreeMap<String, Value>,
used_definition_names: BTreeSet<String>,
root_location: Option<SourceBundleLocation>,
}
impl<R> SourceSchemaBundler<R>
where
R: FnMut(&SourceBundleLocation, &str) -> Option<SourceBundleNode>,
{
fn new(root_schema: Value, resolve_external_ref: R) -> Self {
Self {
root_schema,
resolve_external_ref,
definition_names_by_location: BTreeMap::new(),
definitions_by_name: BTreeMap::new(),
used_definition_names: BTreeSet::new(),
root_location: None,
}
}
fn bundle_root(&mut self, root: SourceBundleNode) -> Value {
self.seed_existing_definition_names(&root.schema);
self.root_location = Some(root.location.clone());
let mut schema = self.bundle_root_node(root, 0);
if !self.definitions_by_name.is_empty() {
let definitions = schema.as_object_mut().and_then(|object| {
object
.entry("$defs".to_string())
.or_insert_with(|| Value::Object(Map::new()))
.as_object_mut()
});
if let Some(definitions) = definitions {
for (name, definition) in std::mem::take(&mut self.definitions_by_name) {
definitions.insert(name, definition);
}
}
}
schema
}
fn seed_existing_definition_names(&mut self, schema: &Value) {
for key in ["$defs", "definitions"] {
let Some(definitions) = schema.get(key).and_then(Value::as_object) else {
continue;
};
self.used_definition_names
.extend(definitions.keys().cloned());
}
}
fn bundle_root_node(&mut self, node: SourceBundleNode, depth: usize) -> Value {
if depth > 64 {
return node.schema;
}
if let Some(reference) = node.schema.get("$ref").and_then(Value::as_str) {
if ref_points_inside(&self.root_schema, reference) {
return self.bundle_schema_value(
&node.schema,
SchemaTraversalContext::Schema,
&node.location,
depth,
);
}
if let Some(target) = (self.resolve_external_ref)(&node.location, reference) {
if self.root_location.as_ref() == Some(&target.location) {
return ref_object("#".to_string());
}
return self.bundle_root_node(target, depth + 1);
}
return strip_ref(&node.schema);
}
self.bundle_schema_value(
&node.schema,
SchemaTraversalContext::Schema,
&node.location,
depth,
)
}
fn bundle_schema_value(
&mut self,
value: &Value,
context: SchemaTraversalContext,
current_location: &SourceBundleLocation,
depth: usize,
) -> Value {
let mapped: Result<Value, Infallible> =
try_map_schema_context(value, context, |value, context, relative_depth| {
if matches!(
context,
SchemaTraversalContext::Data | SchemaTraversalContext::Ref
) {
return Ok(None);
}
let Some(reference) = value.get("$ref").and_then(Value::as_str) else {
return Ok(None);
};
Ok(Some(self.bundle_schema_ref(
value,
reference,
current_location,
depth + relative_depth,
)))
});
match mapped {
Ok(value) => value,
Err(err) => match err {},
}
}
fn bundle_schema_ref(
&mut self,
schema: &Value,
reference: &str,
current_location: &SourceBundleLocation,
depth: usize,
) -> Value {
if ref_points_inside(&self.root_schema, reference) {
return schema.clone();
}
let Some(target) = (self.resolve_external_ref)(current_location, reference) else {
return strip_ref(schema);
};
if self.root_location.as_ref() == Some(&target.location) {
return ref_object("#".to_string());
}
let definition_name = self.definition_name_for_target(target, depth + 1);
ref_object(format!("#/$defs/{definition_name}"))
}
fn definition_name_for_target(&mut self, target: SourceBundleNode, depth: usize) -> String {
if let Some(name) = self.definition_names_by_location.get(&target.location) {
return name.clone();
}
let name = self.next_definition_name(&target.location);
self.definition_names_by_location
.insert(target.location.clone(), name.clone());
let definition_schema = self.bundle_root_node(target, depth);
self.definitions_by_name
.insert(name.clone(), definition_schema);
name
}
fn next_definition_name(&mut self, location: &SourceBundleLocation) -> String {
let base_name = definition_base_name(&location.pointer);
if self.used_definition_names.insert(base_name.clone()) {
return base_name;
}
let mut suffix = 2;
loop {
let candidate = format!("{base_name}_{suffix}");
suffix += 1;
if self.used_definition_names.insert(candidate.clone()) {
return candidate;
}
}
}
}
fn ref_object(reference: String) -> Value {
Value::Object(
[("$ref".to_string(), Value::String(reference))]
.into_iter()
.collect(),
)
}
fn definition_base_name(pointer: &str) -> String {
let base = pointer
.rsplit('/')
.find(|segment| !segment.is_empty())
.unwrap_or("definition");
let sanitized: String = base
.chars()
.map(|character| {
if character.is_ascii_alphanumeric() || character == '_' {
character
} else {
'_'
}
})
.collect();
if sanitized.is_empty() {
"definition".to_string()
} else {
sanitized
}
}
#[cfg(test)]
#[path = "tests/source_bundle.rs"]
mod tests;