Conditionally apply pre-compilation optimisations (#531)

Signed-off-by: Juan Cruz Viotti <jv@jviotti.com>

Author: jviotti

benchmark/micro/draft4.cc

@@ -1114,6 +1114,50 @@ static void Micro_Draft4_Ref_Single_100(benchmark::State &state) {
   }
 }
 
+// Pathological Case: AdaptiveCard-like Pattern
+static void Micro_Draft4_Ref_Many_Nested(benchmark::State &state) {
+  const sourcemeta::core::JSON schema{sourcemeta::core::parse_json(R"JSON({
+    "$schema": "http://json-schema.org/draft-04/schema#",
+    "definitions": {
+      "def0": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def1" }, "b": { "$ref": "#/definitions/def2" } } },
+      "def1": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def2" }, "b": { "$ref": "#/definitions/def3" } } },
+      "def2": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def3" }, "b": { "$ref": "#/definitions/def4" } } },
+      "def3": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def4" }, "b": { "$ref": "#/definitions/def5" } } },
+      "def4": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def5" }, "b": { "$ref": "#/definitions/def6" } } },
+      "def5": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def6" }, "b": { "$ref": "#/definitions/def7" } } },
+      "def6": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def7" }, "b": { "$ref": "#/definitions/def8" } } },
+      "def7": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def8" }, "b": { "$ref": "#/definitions/def9" } } },
+      "def8": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def9" }, "b": { "$ref": "#/definitions/def10" } } },
+      "def9": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def10" }, "b": { "$ref": "#/definitions/def11" } } },
+      "def10": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def11" }, "b": { "$ref": "#/definitions/def12" } } },
+      "def11": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def12" }, "b": { "$ref": "#/definitions/def13" } } },
+      "def12": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def13" }, "b": { "$ref": "#/definitions/def14" } } },
+      "def13": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def14" }, "b": { "$ref": "#/definitions/def0" } } },
+      "def14": { "type": "object", "properties": { "a": { "$ref": "#/definitions/def0" }, "b": { "$ref": "#/definitions/def1" } } }
+    },
+    "properties": {
+      "p0": { "$ref": "#/definitions/def0" },
+      "p1": { "$ref": "#/definitions/def0" },
+      "p2": { "$ref": "#/definitions/def0" },
+      "p3": { "$ref": "#/definitions/def0" },
+      "p4": { "$ref": "#/definitions/def0" },
+      "p5": { "$ref": "#/definitions/def1" },
+      "p6": { "$ref": "#/definitions/def1" },
+      "p7": { "$ref": "#/definitions/def1" },
+      "p8": { "$ref": "#/definitions/def1" },
+      "p9": { "$ref": "#/definitions/def1" }
+    }
+  })JSON")};
+
+  for (auto _ : state) {
+    auto result{sourcemeta::blaze::compile(
+        schema, sourcemeta::core::schema_official_walker,
+        sourcemeta::core::schema_official_resolver,
+        sourcemeta::blaze::default_schema_compiler)};
+    benchmark::DoNotOptimize(result.instructions);
+  }
+}
+
 BENCHMARK(Micro_Draft4_Meta_1_No_Callback);
 BENCHMARK(Micro_Draft4_Required_Properties);
 BENCHMARK(Micro_Draft4_Many_Optional_Properties_Minimal_Match);
@@ -1134,3 +1178,4 @@ BENCHMARK(Micro_Draft4_Long_Enum);
 BENCHMARK(Micro_Draft4_Long_Enum_Short_Strings);
 BENCHMARK(Micro_Draft4_Type_Object);
 BENCHMARK(Micro_Draft4_Ref_Single_100);
+BENCHMARK(Micro_Draft4_Ref_Many_Nested);

src/compiler/compile.cc

@@ -6,7 +6,8 @@
 #include <algorithm> // std::move, std::sort, std::unique
 #include <cassert>   // assert
 #include <iterator>  // std::back_inserter
-#include <utility>   // std::move
+#include <tuple>     // std::tuple, std::get
+#include <utility>   // std::move, std::pair
 
 #include "compile_helpers.h"
 
@@ -48,7 +49,6 @@ auto compile_subschema(const sourcemeta::blaze::Context &context,
               .base = schema_context.base,
               // TODO: This represents a copy
               .labels = schema_context.labels,
-              .references = schema_context.references,
               .is_property_name = schema_context.is_property_name},
              {.keyword = keyword,
               .base_schema_location = dynamic_context.base_schema_location,
@@ -91,7 +91,6 @@ auto precompile(
       .vocabularies = std::move(nested_vocabularies),
       .base = entry.second.base,
       .labels = {},
-      .references = {},
       .is_property_name = schema_context.is_property_name};
 
   return {make(sourcemeta::blaze::InstructionIndex::ControlMark, context,
@@ -117,37 +116,24 @@ auto compile(const sourcemeta::core::JSON &schema,
              const std::optional<std::string> &default_id) -> Template {
   assert(is_schema(schema));
 
+  ///////////////////////////////////////////////////////////////////
+  // (1) Determine the root frame entry
+  ///////////////////////////////////////////////////////////////////
+
   const std::string base{sourcemeta::core::URI::canonicalize(
       sourcemeta::core::identify(
           schema, resolver,
           sourcemeta::core::SchemaIdentificationStrategy::Strict,
           default_dialect, default_id)
           .value_or(""))};
-
   assert(frame.locations().contains(
       {sourcemeta::core::SchemaReferenceType::Static, base}));
   const auto root_frame_entry{frame.locations().at(
       {sourcemeta::core::SchemaReferenceType::Static, base})};
 
-  // Check whether dynamic referencing takes places in this schema. If not,
-  // we can avoid the overhead of keeping track of dynamics scopes, etc
-  bool uses_dynamic_scopes{false};
-  for (const auto &reference : frame.references()) {
-    if (reference.first.first ==
-        sourcemeta::core::SchemaReferenceType::Dynamic) {
-      uses_dynamic_scopes = true;
-      break;
-    }
-  }
-
-  SchemaContext schema_context{
-      .relative_pointer = sourcemeta::core::empty_pointer,
-      .schema = schema,
-      .vocabularies = vocabularies(schema, resolver, root_frame_entry.dialect),
-      .base = sourcemeta::core::URI::canonicalize(root_frame_entry.base),
-      .labels = {},
-      .references = {},
-      .is_property_name = false};
+  ///////////////////////////////////////////////////////////////////
+  // (2) Determine all the schema resources in the schema
+  ///////////////////////////////////////////////////////////////////
 
   std::vector<std::string> resources;
   for (const auto &entry : frame.locations()) {
@@ -165,50 +151,108 @@ auto compile(const sourcemeta::core::JSON &schema,
   assert(resources.size() ==
          std::set<std::string>(resources.cbegin(), resources.cend()).size());
 
-  // Calculate the top static reference destinations for precompilation purposes
-  // TODO: Replace this logic with `.frame()` `destination_of` information
-  std::set<std::string> precompiled_static_schemas;
-  // As a workaround, we avoid pre-compiling schemas on schemas
-  // that look like they are just wrapping other schemas
-  if (schema.is_object() && !schema.defines("$ref")) {
-    std::map<std::string, std::size_t> static_references_count;
+  ///////////////////////////////////////////////////////////////////
+  // (3) Check if the schema relies on dynamic scopes
+  ///////////////////////////////////////////////////////////////////
+
+  bool uses_dynamic_scopes{false};
+  for (const auto &reference : frame.references()) {
+    // Check whether dynamic referencing takes places in this schema. If not,
+    // we can avoid the overhead of keeping track of dynamics scopes, etc
+    if (reference.first.first ==
+        sourcemeta::core::SchemaReferenceType::Dynamic) {
+      uses_dynamic_scopes = true;
+      break;
+    }
+  }
+
+  ///////////////////////////////////////////////////////////////////
+  // (4) Plan which static references we will precompile
+  ///////////////////////////////////////////////////////////////////
+
+  // Use string views to avoid copying the actual strings, as we know
+  // that the frame survives the entire compilation process
+  std::vector<std::tuple<std::string_view, std::size_t, std::size_t>>
+      sorted_references;
+
+  constexpr auto PRECOMPILED_SCHEMAS_MAXIMUM{10};
+  constexpr auto PRECOMPILED_SCHEMAS_MINIMUM_COUNT{10};
+
+  {
+    std::unordered_map<std::string_view, std::pair<std::size_t, std::size_t>>
+        static_reference_destinations;
     for (const auto &reference : frame.references()) {
-      if (reference.first.first !=
-              sourcemeta::core::SchemaReferenceType::Static ||
-          !frame.locations().contains(
+      if (reference.first.first ==
+              sourcemeta::core::SchemaReferenceType::Static &&
+          frame.locations().contains(
               {sourcemeta::core::SchemaReferenceType::Static,
                reference.second.destination})) {
-        continue;
+        // TODO: Maybe try circular references or non-circular with >100 inbound
+        // locations or something like that?
+        std::unordered_set<std::string> visited;
+        if (!is_circular(frame, reference.first.second, reference.second,
+                         visited)) {
+          continue;
+        }
+
+        const auto label{Evaluator{}.hash(
+            schema_resource_id(resources, reference.second.base.value_or("")),
+            reference.second.fragment.value_or(""))};
+        auto [iterator, inserted] = static_reference_destinations.try_emplace(
+            reference.second.destination, std::make_pair(label, 0));
+        iterator->second.second++;
       }
+    }
 
-      const auto &entry{
-          frame.locations().at({sourcemeta::core::SchemaReferenceType::Static,
-                                reference.second.destination})};
-      for (const auto &subreference : frame.references()) {
-        if (subreference.first.second.starts_with(entry.pointer)) {
-          static_references_count[reference.second.destination] += 1;
-        }
+    sorted_references.reserve(static_reference_destinations.size());
+    for (const auto &reference : static_reference_destinations) {
+      if (reference.second.second >= PRECOMPILED_SCHEMAS_MINIMUM_COUNT) {
+        sorted_references.emplace_back(reference.first, reference.second.first,
+                                       reference.second.second);
       }
     }
-    std::vector<std::pair<std::string, std::size_t>> top_static_destinations(
-        static_references_count.cbegin(), static_references_count.cend());
-    std::ranges::sort(top_static_destinations,
+    std::ranges::sort(sorted_references,
                       [](const auto &left, const auto &right) {
-                        return left.second > right.second;
+                        return std::get<2>(left) > std::get<2>(right);
                       });
-    constexpr auto MAXIMUM_NUMBER_OF_SCHEMAS_TO_PRECOMPILE{5};
-    for (auto iterator = top_static_destinations.cbegin();
-         iterator != top_static_destinations.cend() &&
-         iterator != top_static_destinations.cbegin() +
-                         MAXIMUM_NUMBER_OF_SCHEMAS_TO_PRECOMPILE;
-         ++iterator) {
-      // Only consider highly referenced schemas
-      if (iterator->second > 100) {
-        precompiled_static_schemas.insert(iterator->first);
-      }
+
+    if (sorted_references.size() > PRECOMPILED_SCHEMAS_MAXIMUM) {
+      sorted_references.erase(sorted_references.begin() +
+                                  PRECOMPILED_SCHEMAS_MAXIMUM,
+                              sorted_references.end());
     }
+
+    // We do not apply this pre-compilation optimisation on meta-schemas
+    if (sourcemeta::core::schema_official_resolver(base).has_value() ||
+        (uses_dynamic_scopes && schema.is_object() &&
+         schema.defines("$vocabulary"))) {
+      sorted_references.clear();
+    }
+  }
+
+  assert(sorted_references.size() <= PRECOMPILED_SCHEMAS_MAXIMUM);
+  std::unordered_set<std::size_t> precompiled_labels;
+  for (const auto &reference : sorted_references) {
+    assert(std::get<2>(reference) >= PRECOMPILED_SCHEMAS_MINIMUM_COUNT);
+    precompiled_labels.emplace(std::get<1>(reference));
   }
 
+  ///////////////////////////////////////////////////////////////////
+  // (5) Build the starting schema context
+  ///////////////////////////////////////////////////////////////////
+
+  SchemaContext schema_context{
+      .relative_pointer = sourcemeta::core::empty_pointer,
+      .schema = schema,
+      .vocabularies = vocabularies(schema, resolver, root_frame_entry.dialect),
+      .base = sourcemeta::core::URI::canonicalize(root_frame_entry.base),
+      .labels = {},
+      .is_property_name = false};
+
+  ///////////////////////////////////////////////////////////////////
+  // (6) Build the gloal compilation context
+  ///////////////////////////////////////////////////////////////////
+
   auto unevaluated{
       sourcemeta::blaze::unevaluated(schema, frame, walker, resolver)};
 
@@ -221,22 +265,19 @@ auto compile(const sourcemeta::core::JSON &schema,
                         .mode = mode,
                         .uses_dynamic_scopes = uses_dynamic_scopes,
                         .unevaluated = std::move(unevaluated),
-                        .precompiled_static_schemas =
-                            std::move(precompiled_static_schemas)};
+                        .precompiled_labels = std::move(precompiled_labels)};
+
+  ///////////////////////////////////////////////////////////////////
+  // (7) Build the initial dynamic context
+  ///////////////////////////////////////////////////////////////////
+
   const DynamicContext dynamic_context{relative_dynamic_context()};
-  Instructions compiler_template;
 
-  for (const auto &destination : context.precompiled_static_schemas) {
-    assert(context.frame.locations().contains(
-        {sourcemeta::core::SchemaReferenceType::Static, destination}));
-    const auto match{context.frame.locations().find(
-        {sourcemeta::core::SchemaReferenceType::Static, destination})};
-    for (auto &&substep :
-         precompile(context, schema_context, dynamic_context, *match)) {
-      compiler_template.push_back(std::move(substep));
-    }
-  }
+  ///////////////////////////////////////////////////////////////////
+  // (8) Pre compile dynamic reference locations
+  ///////////////////////////////////////////////////////////////////
 
+  Instructions compiler_template;
   if (uses_dynamic_scopes &&
       (schema_context.vocabularies.contains(
            "https://json-schema.org/draft/2019-09/vocab/core") ||
@@ -257,9 +298,59 @@ auto compile(const sourcemeta::core::JSON &schema,
     }
   }
 
+  ///////////////////////////////////////////////////////////////////
+  // (9) Pre compile static reference locations
+  ///////////////////////////////////////////////////////////////////
+
+  // Attempt to precompile static destinations to avoid explosive compilation
+  Instructions static_reference_template;
+  for (const auto &reference : sorted_references) {
+    const auto entry{context.frame.locations().find(
+        {sourcemeta::core::SchemaReferenceType::Static,
+         std::string{std::get<0>(reference)}})};
+    assert(entry != context.frame.locations().cend());
+    auto subschema{sourcemeta::core::get(context.root, entry->second.pointer)};
+    if (!sourcemeta::core::is_schema(subschema)) {
+      continue;
+    }
+
+    auto nested_vocabularies{sourcemeta::core::vocabularies(
+        subschema, context.resolver, entry->second.dialect)};
+    const sourcemeta::blaze::SchemaContext nested_schema_context{
+        .relative_pointer = entry->second.relative_pointer,
+        .schema = std::move(subschema),
+        .vocabularies = std::move(nested_vocabularies),
+        // TODO: I think this is hiding a framing bug that we should later
+        // investigate
+        .base = entry->second.base.starts_with('#') ? "" : entry->second.base,
+        .labels = {},
+        .is_property_name = schema_context.is_property_name};
+    static_reference_template.push_back(
+        make(sourcemeta::blaze::InstructionIndex::ControlMark, context,
+             nested_schema_context, dynamic_context,
+             sourcemeta::blaze::ValueUnsignedInteger{std::get<1>(reference)},
+             sourcemeta::blaze::compile(
+                 context, nested_schema_context,
+                 sourcemeta::blaze::relative_dynamic_context(),
+                 sourcemeta::core::empty_pointer,
+                 sourcemeta::core::empty_pointer, entry->first.second)));
+  }
+
+  for (auto &&substep : static_reference_template) {
+    compiler_template.push_back(std::move(substep));
+  }
+
+  ///////////////////////////////////////////////////////////////////
+  // (10) Compile the actual schema
+  ///////////////////////////////////////////////////////////////////
+
   auto children{compile_subschema(context, schema_context, dynamic_context,
                                   root_frame_entry.dialect)};
 
+  ///////////////////////////////////////////////////////////////////
+  // (11) Return final template
+  ///////////////////////////////////////////////////////////////////
+
   const bool track{
       context.mode != Mode::FastValidation ||
       requires_evaluation(context, schema_context) ||
@@ -353,7 +444,6 @@ auto compile(const Context &context, const SchemaContext &schema_context,
                    .value_or(""),
        // TODO: This represents a copy
        .labels = schema_context.labels,
-       .references = schema_context.references,
        .is_property_name = schema_context.is_property_name},
       {.keyword = dynamic_context.keyword,
        .base_schema_location = destination_pointer,