jsonschema 0.58.5

//! Shared state for one canonicalization run: draft, pattern engine, and a compiled-regex cache.
use std::{
    cell::{Cell, RefCell},
    cmp::Ordering,
    collections::BTreeSet,
    sync::{Arc, Mutex},
};

use ahash::AHashMap;

use referencing::Draft;

use crate::{
    canonical::{ir::Schema, schema::DefinitionMap},
    options::PatternEngineOptions,
};

/// Past this many remembered pairs a run keeps recomputing rather than grow without end.
const INTERSECTION_CACHE_CAPACITY: usize = 1 << 20;

/// Past this many remembered nodes a run walks them again rather than grow without end.
const FACET_CACHE_CAPACITY: usize = 1 << 16;

/// Patterns compiled under one document's engine, shared by every run over it.
pub(crate) type SharedRegexes = Arc<Mutex<AHashMap<Arc<str>, Option<Arc<CompiledMatcher>>>>>;

pub(crate) enum CompiledMatcher {
    Regex(regex::Regex),
    FancyRegex(fancy_regex::Regex),
}

impl CompiledMatcher {
    /// A match error (e.g. `fancy_regex` hitting its backtrack limit) counts as no match, matching
    /// the runtime `pattern` validator's `is_valid`.
    pub(crate) fn is_match(&self, text: &str) -> bool {
        match self {
            Self::Regex(regex) => regex.is_match(text),
            Self::FancyRegex(regex) => regex.is_match(text).unwrap_or(false),
        }
    }
}

pub(crate) struct CanonicalizationContext {
    draft: Draft,
    pattern_options: PatternEngineOptions,
    /// When false `format` is an annotation, constrains nothing, and is dropped.
    validate_formats: bool,
    /// `None` caches a rejected pattern so callers don't recompile it.
    regex_cache: RefCell<AHashMap<Arc<str>, Option<Arc<CompiledMatcher>>>>,
    /// The compiled patterns of the document this run reads, which outlive it: every run over that
    /// document translates and compiles the same texts under the same engine.
    shared_regexes: Option<SharedRegexes>,
    /// An `allOf` over unions takes the product of their branches, which reaches the same pair
    /// of nodes over and over - on a schema of five such `allOf`s, 431 times per distinct pair.
    intersections: RefCell<AHashMap<(Schema, Schema), Remembered>>,
    /// Every containment query reads the facets of the node it asks about, and the object walks ask
    /// about the same handful of nodes once per piece they cut out.
    uncheckable_facets: RefCell<AHashMap<Schema, Arc<BTreeSet<Arc<str>>>>>,
    /// An intersection reached during this run that the canonical form cannot express exactly.
    /// Nodes built around it may already be wrong, so the whole run is discarded rather than the site.
    inexact_intersection: Cell<bool>,
    /// The targets an intersection may read through. Absent while a document is being canonicalized,
    /// where bodies are still arriving.
    definitions: Option<Arc<DefinitionMap>>,
    /// The targets that lie on a reference cycle, which no walk reads through: it would not
    /// terminate. Every other target of the same map is read through as usual.
    cyclic: BTreeSet<Arc<str>>,
    /// Intersections this run may still take before giving up. An `allOf` over unions multiplies them.
    intersections_left: Cell<u64>,
    /// Variants the conditional splits of this run may still produce. Nesting multiplies them.
    variants_left: Cell<u64>,
    /// Set where a conditional split asked for more cases than were left.
    outgrew_cases: Cell<bool>,
    /// Address of the first subschema whose parse declined, for the pointer naming it. Compared
    /// against the document's nodes, never read through.
    declined_at: Cell<Option<usize>>,
}

/// Intersections one run may take before giving up and leaving the document `Raw`. Above what the
/// most demanding document written in earnest needs; a row of `allOf`s over unions passes it in
/// a fraction of a second.
const INTERSECTION_BUDGET: u64 = 1_000_000;

/// Variants the conditional splits of one run may produce before the document stays `Raw`. Two
/// nested nodes at the per-node cap fit; a third would spend 4096 more.
const VARIANT_BUDGET: u64 = 4096;

impl CanonicalizationContext {
    pub(crate) fn new(
        draft: Draft,
        pattern_options: PatternEngineOptions,
        validate_formats: bool,
    ) -> Self {
        Self {
            draft,
            pattern_options,
            validate_formats,
            regex_cache: RefCell::new(AHashMap::new()),
            shared_regexes: None,
            intersections: RefCell::new(AHashMap::new()),
            uncheckable_facets: RefCell::new(AHashMap::new()),
            inexact_intersection: Cell::new(false),
            intersections_left: Cell::new(INTERSECTION_BUDGET),
            variants_left: Cell::new(VARIANT_BUDGET),
            outgrew_cases: Cell::new(false),
            declined_at: Cell::new(None),
            definitions: None,
            cyclic: BTreeSet::new(),
        }
    }

    /// The same context, keeping what it compiles for the next run over the same document.
    pub(crate) fn sharing_regexes(mut self, regexes: SharedRegexes) -> Self {
        self.shared_regexes = Some(regexes);
        self
    }

    /// The same context, reading intersections through `definitions`. The caller passes a map only
    /// when it is complete, and names the targets on a cycle, which stay unread.
    pub(crate) fn resolving(
        mut self,
        definitions: Arc<DefinitionMap>,
        cyclic: BTreeSet<Arc<str>>,
    ) -> Self {
        self.definitions = Some(definitions);
        self.cyclic = cyclic;
        self
    }

    /// The same context, allowed `budget` intersections rather than a whole document's worth.
    pub(crate) fn within(mut self, budget: u64) -> Self {
        self.intersections_left = Cell::new(budget);
        self
    }

    /// The targets this run reads through.
    pub(crate) fn targets(&self) -> &DefinitionMap {
        self.definitions
            .as_deref()
            .expect("a settling run reads targets")
    }

    /// The targets, for a settling pass to move on. What the run already remembers stays: a pass
    /// settles a body only after every body it reads met those bodies in their final form.
    pub(crate) fn targets_mut(&mut self) -> &mut DefinitionMap {
        let definitions = self
            .definitions
            .as_mut()
            .expect("a settling run reads targets");
        // Held here alone, so the edit lands in place instead of copying the map per body.
        debug_assert_eq!(
            Arc::strong_count(definitions),
            1,
            "a settling run holds its targets alone"
        );
        Arc::make_mut(definitions)
    }

    pub(crate) fn into_targets(self) -> DefinitionMap {
        Arc::unwrap_or_clone(self.definitions.expect("a settling run reads targets"))
    }

    pub(crate) fn pattern_options(&self) -> PatternEngineOptions {
        self.pattern_options
    }

    /// The body `uri` names, or `None` where this run reads no targets and where reading through
    /// this one would not terminate.
    pub(crate) fn definition(&self, uri: &str) -> Option<&Schema> {
        if self.cyclic.contains(uri) {
            return None;
        }
        self.definitions.as_ref()?.get(uri)
    }

    pub(crate) fn draft(&self) -> Draft {
        self.draft
    }

    pub(crate) fn record_inexact_intersection(&self) {
        self.inexact_intersection.set(true);
    }

    pub(crate) fn saw_inexact_intersection(&self) -> bool {
        self.inexact_intersection.get()
    }

    /// Run `probe`, reporting whether it reached an intersection the canonical form cannot express
    /// exactly. The flag is left as it was, so what a probe reaches decides nothing beyond its own
    /// answer. The intersection budget is not restored: it bounds the work one run may do, and
    /// speculative work is work.
    pub(crate) fn probe<T>(&self, probe: impl FnOnce() -> T) -> (T, bool) {
        let before = self.inexact_intersection.replace(false);
        let answer = probe();
        let inexact = self.inexact_intersection.replace(before);
        (answer, inexact)
    }

    /// Take one intersection from what this run may still spend, reporting whether it had any. Once
    /// it runs out every later intersection is refused too and the walk unwinds.
    pub(crate) fn take_intersection(&self) -> bool {
        let left = self.intersections_left.get();
        if left == 0 {
            return false;
        }
        self.intersections_left.set(left - 1);
        true
    }

    /// Take `count` conditional-split variants from what this run may still spend, reporting
    /// whether it had them.
    pub(crate) fn take_variants(&self, count: u64) -> bool {
        let left = self.variants_left.get();
        if left < count {
            self.note_outgrew_cases();
            return false;
        }
        self.variants_left.set(left - count);
        true
    }

    /// Run `work` against at most `cap` intersections, reporting whether it ran out. What it spends
    /// counts against the run's budget, so a bounded side question cannot exhaust the main answer.
    pub(crate) fn capped<T>(&self, cap: u64, work: impl FnOnce() -> T) -> (T, bool) {
        let before = self.intersections_left.get();
        let allowance = before.min(cap);
        self.intersections_left.set(allowance);
        let answer = work();
        let left = self.intersections_left.get();
        self.intersections_left.set(before - (allowance - left));
        (answer, left == 0)
    }

    /// Run `work` without letting it change the run: the budget it spends is given back and the
    /// inexactness it reaches is forgotten. A recording parse must spend exactly what the plain
    /// parse spends. What it remembers stays: a pair the probe could only approximate had less
    /// budget than any later computation of the same pair, so the run reads the same answer.
    pub(crate) fn speculate<T>(&self, work: impl FnOnce() -> T) -> T {
        let budget = self.intersections_left.get();
        let (answer, _) = self.probe(work);
        self.intersections_left.set(budget);
        answer
    }

    pub(crate) fn outgrew_distribution(&self) -> bool {
        self.intersections_left.get() == 0
    }

    pub(crate) fn note_outgrew_cases(&self) {
        self.outgrew_cases.set(true);
    }

    pub(crate) fn outgrew_cases(&self) -> bool {
        self.outgrew_cases.get()
    }

    /// Remember the subschema whose parse declined, unless one is already remembered: the walk is
    /// depth first and every caller passes a decline on, so the first is the one that caused it.
    pub(crate) fn note_declined(&self, address: usize) {
        if self.declined_at.get().is_none() {
            self.declined_at.set(Some(address));
        }
    }

    /// Forget the decline of an earlier parse attempt, whose nodes this one re-reads.
    pub(crate) fn forget_decline(&self) {
        self.declined_at.set(None);
    }

    pub(crate) fn declined_at(&self) -> Option<usize> {
        self.declined_at.get()
    }

    /// Parse a schema this run wrote itself, keeping what it declines on off the record: those
    /// nodes are not the document's, and the node standing in for them is the one being rewritten.
    pub(crate) fn over_rewritten<T>(&self, parse: impl FnOnce() -> T) -> T {
        let before = self.declined_at.replace(None);
        let parsed = parse();
        self.declined_at.set(before);
        parsed
    }

    pub(crate) fn validate_formats(&self) -> bool {
        self.validate_formats
    }

    /// The pattern compiled under the configured engine, or `None` if the engine rejects it. Compiled
    /// once per run and cached, so parse-time validation and membership share the same matcher.
    pub(crate) fn compile_regex(&self, pattern: &Arc<str>) -> Option<Arc<CompiledMatcher>> {
        if let Some(cached) = self.regex_cache.borrow().get(pattern) {
            return cached.clone();
        }
        // Reached once per pattern per run, so the shared map is locked that often rather than
        // once per use. The document fixes the engine, so its text alone names the matcher.
        let shared = self.shared_regexes.as_ref();
        let held = shared.and_then(|shared| {
            let shared = shared.lock().ok()?;
            shared.get(pattern).cloned()
        });
        let compiled = if let Some(compiled) = held {
            compiled
        } else {
            // Compiled outside the lock: holding it across a translation would serialize every
            // other run over the document, and a panic there would poison the cache.
            let compiled = compile(self.pattern_options, pattern).map(Arc::new);
            if let Some(mut shared) = shared.and_then(|shared| shared.lock().ok()) {
                shared.insert(Arc::clone(pattern), compiled.clone());
            }
            compiled
        };
        self.regex_cache
            .borrow_mut()
            .insert(Arc::clone(pattern), compiled.clone());
        compiled
    }

    /// The facets of `schema` no checker covers, walking the node the first time it is asked about.
    pub(crate) fn uncheckable_facets(
        &self,
        schema: &Schema,
        walk: impl FnOnce() -> BTreeSet<Arc<str>>,
    ) -> Arc<BTreeSet<Arc<str>>> {
        if let Some(cached) = self.uncheckable_facets.borrow().get(schema) {
            return Arc::clone(cached);
        }
        let found = Arc::new(walk());
        let mut cache = self.uncheckable_facets.borrow_mut();
        if cache.len() < FACET_CACHE_CAPACITY {
            cache.insert(schema.clone(), Arc::clone(&found));
        }
        found
    }

    /// The intersection of these two, from an earlier run of the same pair. One the form could only
    /// approximate is recorded again here: a walk reading it is as approximate as the one that
    /// first took it.
    pub(crate) fn recall_intersection(&self, left: &Schema, right: &Schema) -> Option<Schema> {
        let key = intersection_key(left.clone(), right.clone());
        let remembered = self.intersections.borrow().get(&key).cloned()?;
        if remembered.inexact {
            self.record_inexact_intersection();
        }
        Some(remembered.result)
    }

    /// Remember this pair's intersection, and whether taking it approximated.
    pub(crate) fn remember_intersection(
        &self,
        left: Schema,
        right: Schema,
        result: &Schema,
        inexact: bool,
    ) {
        let mut intersections = self.intersections.borrow_mut();
        if intersections.len() < INTERSECTION_CACHE_CAPACITY {
            intersections.insert(
                intersection_key(left, right),
                Remembered {
                    result: result.clone(),
                    inexact,
                },
            );
        }
    }
}

/// One pair's intersection, beside whether the form could only approximate it.
#[derive(Clone)]
struct Remembered {
    result: Schema,
    inexact: bool,
}

fn intersection_key(left: Schema, right: Schema) -> (Schema, Schema) {
    match left.cached_hash().cmp(&right.cached_hash()) {
        Ordering::Greater => (right, left),
        Ordering::Equal if left > right => (right, left),
        Ordering::Less | Ordering::Equal => (left, right),
    }
}

fn compile(options: PatternEngineOptions, pattern: &str) -> Option<CompiledMatcher> {
    let translated = jsonschema_regex::to_rust_regex(pattern).ok()?;
    match options {
        PatternEngineOptions::Regex {
            size_limit,
            dfa_size_limit,
        } => crate::regex::build_standard_regex(&translated, size_limit, dfa_size_limit)
            .ok()
            .map(CompiledMatcher::Regex),
        PatternEngineOptions::FancyRegex {
            backtrack_limit,
            size_limit,
            dfa_size_limit,
        } => crate::regex::build_fancy_regex(
            &translated,
            backtrack_limit,
            size_limit,
            dfa_size_limit,
        )
        .ok()
        .map(CompiledMatcher::FancyRegex),
    }
}