refactor: shard meta-schemas per draft

This introduces some new (internal-only and lazy) static entities in
order to shard meta-schemas according to the draft that they belong to.
Sharded meta-schemas allow to reduce the amount of registry resources
that need to be constructed and statically cached.

Author: lucab

crates/jsonschema-referencing/src/meta.rs

@@ -6,6 +6,8 @@ use std::sync::Arc;
 use once_cell::sync::Lazy;
 use serde_json::Value;
 
+use crate::Draft;
+
 macro_rules! schema {
     ($vis:vis $name:ident, $path:expr) => {
         $vis static $name: once_cell::sync::Lazy<Arc<serde_json::Value>> =
@@ -79,7 +81,7 @@ schema!(
     pub DRAFT202012_CONTENT,
     "../metaschemas/draft2020-12/meta/content.json"
 );
-pub(crate) static META_SCHEMAS: Lazy<[(&'static str, &'static Value); 18]> = Lazy::new(|| {
+pub(crate) static META_SCHEMAS_ALL: Lazy<[(&'static str, &'static Value); 18]> = Lazy::new(|| {
     [
         ("http://json-schema.org/draft-04/schema#", &*DRAFT4),
         ("http://json-schema.org/draft-06/schema#", &*DRAFT6),
@@ -146,3 +148,95 @@ pub(crate) static META_SCHEMAS: Lazy<[(&'static str, &'static Value); 18]> = Laz
         ),
     ]
 });
+
+pub(crate) static META_SCHEMAS_DRAFT4: Lazy<[(&'static str, &'static Value); 1]> =
+    Lazy::new(|| [("http://json-schema.org/draft-04/schema#", &*DRAFT4)]);
+
+pub(crate) static META_SCHEMAS_DRAFT6: Lazy<[(&'static str, &'static Value); 1]> =
+    Lazy::new(|| [("http://json-schema.org/draft-06/schema#", &*DRAFT6)]);
+
+pub(crate) static META_SCHEMAS_DRAFT7: Lazy<[(&'static str, &'static Value); 1]> =
+    Lazy::new(|| [("http://json-schema.org/draft-07/schema#", &*DRAFT7)]);
+
+pub(crate) static META_SCHEMAS_DRAFT2019: Lazy<[(&'static str, &'static Value); 7]> =
+    Lazy::new(|| {
+        [
+            (
+                "https://json-schema.org/draft/2019-09/schema",
+                &*DRAFT201909,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/applicator",
+                &*DRAFT201909_APPLICATOR,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/content",
+                &*DRAFT201909_CONTENT,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/core",
+                &*DRAFT201909_CORE,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/format",
+                &*DRAFT201909_FORMAT,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/meta-data",
+                &*DRAFT201909_META_DATA,
+            ),
+            (
+                "https://json-schema.org/draft/2019-09/meta/validation",
+                &*DRAFT201909_VALIDATION,
+            ),
+        ]
+    });
+
+pub(crate) static META_SCHEMAS_DRAFT2020: Lazy<[(&'static str, &'static Value); 8]> =
+    Lazy::new(|| {
+        [
+            (
+                "https://json-schema.org/draft/2020-12/schema",
+                &*DRAFT202012,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/core",
+                &*DRAFT202012_CORE,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/applicator",
+                &*DRAFT202012_APPLICATOR,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/unevaluated",
+                &*DRAFT202012_UNEVALUATED,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/validation",
+                &*DRAFT202012_VALIDATION,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/meta-data",
+                &*DRAFT202012_META_DATA,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/format-annotation",
+                &*DRAFT202012_FORMAT_ANNOTATION,
+            ),
+            (
+                "https://json-schema.org/draft/2020-12/meta/content",
+                &*DRAFT202012_CONTENT,
+            ),
+        ]
+    });
+
+/// Return all the meta-schemas which are part of a given draft.
+pub(crate) fn metas_for_draft(draft: Draft) -> &'static [(&'static str, &'static Value)] {
+    match draft {
+        Draft::Draft4 => &*META_SCHEMAS_DRAFT4,
+        Draft::Draft6 => &*META_SCHEMAS_DRAFT6,
+        Draft::Draft7 => &*META_SCHEMAS_DRAFT7,
+        Draft::Draft201909 => &*META_SCHEMAS_DRAFT2019,
+        Draft::Draft202012 => &*META_SCHEMAS_DRAFT2020,
+    }
+}

crates/jsonschema-referencing/src/registry.rs

@@ -15,7 +15,7 @@ use crate::{
     cache::{SharedUriCache, UriCache},
     hasher::BuildNoHashHasher,
     list::List,
-    meta,
+    meta::{self, metas_for_draft},
     resource::{unescape_segment, InnerResourcePtr, JsonSchemaResource},
     uri,
     vocabularies::{self, VocabularySet},
@@ -44,34 +44,8 @@ type DocumentStore = AHashMap<Arc<Uri<String>>, Pin<Arc<ValueWrapper>>>;
 type ResourceMap = AHashMap<Arc<Uri<String>>, InnerResourcePtr>;
 
 /// Pre-loaded registry containing all JSON Schema meta-schemas and their vocabularies
-pub static SPECIFICATIONS: Lazy<Registry> = Lazy::new(|| {
-    let pairs = meta::META_SCHEMAS.into_iter().map(|(uri, schema)| {
-        (
-            uri,
-            ResourceRef::from_contents(schema).expect("Invalid resource"),
-        )
-    });
-
-    // The capacity is known upfront
-    let mut documents = DocumentStore::with_capacity(18);
-    let mut resources = ResourceMap::with_capacity(18);
-    let mut anchors = AHashMap::with_capacity(8);
-    let mut resolution_cache = UriCache::with_capacity(35);
-    process_meta_schemas(
-        pairs,
-        &mut documents,
-        &mut resources,
-        &mut anchors,
-        &mut resolution_cache,
-    )
-    .expect("Failed to process meta schemas");
-    Registry {
-        documents,
-        resources,
-        anchors,
-        resolution_cache: resolution_cache.into_shared(),
-    }
-});
+pub static SPECIFICATIONS: Lazy<Registry> =
+    Lazy::new(|| Registry::build_from_meta_schemas(meta::META_SCHEMAS_ALL.as_slice()));
 
 /// A registry of JSON Schema resources, each identified by their canonical URIs.
 ///
@@ -563,6 +537,42 @@ impl Registry {
             _ => unreachable!(),
         }
     }
+
+    /// Build a registry with all the given meta-schemas from specs.
+    pub(crate) fn build_from_meta_schemas(schemas: &[(&'static str, &'static Value)]) -> Self {
+        let schemas_count = schemas.len();
+        let pairs = schemas.iter().map(|(uri, schema)| {
+            (
+                uri,
+                ResourceRef::from_contents(schema).expect("Invalid resource"),
+            )
+        });
+
+        let mut documents = DocumentStore::with_capacity(schemas_count);
+        let mut resources = ResourceMap::with_capacity(schemas_count);
+
+        // The actual number of anchors and cache-entries varies across
+        // drafts. We overshoot here to avoid reallocations, using the sum
+        // over all specifications.
+        let mut anchors = AHashMap::with_capacity(8);
+        let mut resolution_cache = UriCache::with_capacity(35);
+
+        process_meta_schemas(
+            pairs,
+            &mut documents,
+            &mut resources,
+            &mut anchors,
+            &mut resolution_cache,
+        )
+        .expect("Failed to process meta schemas");
+
+        Self {
+            documents,
+            resources,
+            anchors,
+            resolution_cache: resolution_cache.into_shared(),
+        }
+    }
 }
 
 fn process_meta_schemas(
@@ -706,15 +716,18 @@ fn handle_metaschemas(
     refers_metaschemas: bool,
     resources: &mut ResourceMap,
     anchors: &mut AHashMap<AnchorKey, Anchor>,
+    draft_version: Draft,
 ) {
     if refers_metaschemas {
-        resources.reserve(SPECIFICATIONS.resources.len());
-        for (key, resource) in &SPECIFICATIONS.resources {
-            resources.insert(Arc::clone(key), resource.clone());
+        let schemas = metas_for_draft(draft_version);
+        let draft_registry = Registry::build_from_meta_schemas(schemas);
+        resources.reserve(draft_registry.resources.len());
+        for (key, resource) in draft_registry.resources {
+            resources.insert(key, resource.clone());
         }
-        anchors.reserve(SPECIFICATIONS.anchors.len());
-        for (key, anchor) in &SPECIFICATIONS.anchors {
-            anchors.insert(key.clone(), anchor.clone());
+        anchors.reserve(draft_registry.anchors.len());
+        for (key, anchor) in draft_registry.anchors {
+            anchors.insert(key, anchor);
         }
     }
 }
@@ -783,7 +796,7 @@ fn process_resources(
         }
     }
 
-    handle_metaschemas(state.refers_metaschemas, resources, anchors);
+    handle_metaschemas(state.refers_metaschemas, resources, anchors, default_draft);
 
     Ok(())
 }
@@ -860,7 +873,7 @@ async fn process_resources_async(
         }
     }
 
-    handle_metaschemas(state.refers_metaschemas, resources, anchors);
+    handle_metaschemas(state.refers_metaschemas, resources, anchors, default_draft);
 
     Ok(())
 }