monty 0.0.19-beta.2

//! JSON serialization support for `json.dumps()`.
//!
//! This module owns encoder keyword parsing, CPython-compatible string/float
//! formatting, and recursive serialization of Monty values.

use std::{
    cmp::Ordering,
    fmt::{Display, Write},
};

use crate::{
    args::{ArgValues, FromArgs},
    bytecode::{ContainsVM, VM},
    defer_drop_mut, defer_drop_vm, defer_drop_vm_mut,
    exception_private::{ExcType, RunResult},
    heap::{ContainsHeap, Heap, HeapData, HeapGuard, HeapId, HeapRead, HeapReadOutput},
    resource::ResourceTracker,
    sorting::{apply_permutation, sort_indices},
    types::{Dict, PyTrait, long_int::check_bigint_str_digits_limit, str::allocate_string},
    value::Value,
};

/// Serializer configuration derived from `json.dumps()` keyword arguments.
///
/// The struct stores only the subset of encoder configuration that this module
/// actually uses while serializing. Unsupported or not-yet-implemented kwargs
/// still raise during parsing so call sites do not silently lose behavior.
struct JsonDumpsConfig {
    indent: Option<String>,
    item_separator: String,
    key_separator: String,
    flags: u8,
}

impl Default for JsonDumpsConfig {
    /// Returns the CPython default `json.dumps()` configuration.
    ///
    /// Compact output uses `", "` between items and `": "` between keys and
    /// values, ASCII escaping is enabled, NaN and infinity are emitted as
    /// `NaN`/`Infinity`, and invalid dict keys raise immediately.
    fn default() -> Self {
        Self {
            indent: None,
            item_separator: ", ".to_owned(),
            key_separator: ": ".to_owned(),
            flags: Self::ENSURE_ASCII | Self::ALLOW_NAN,
        }
    }
}

impl JsonDumpsConfig {
    /// Bit flag storing the `sort_keys` option.
    const SORT_KEYS: u8 = 1 << 0;
    /// Bit flag storing the `ensure_ascii` option.
    const ENSURE_ASCII: u8 = 1 << 1;
    /// Bit flag storing the `allow_nan` option.
    const ALLOW_NAN: u8 = 1 << 2;
    /// Bit flag storing the `skipkeys` option.
    const SKIPKEYS: u8 = 1 << 3;

    /// Returns whether `sort_keys=True` is enabled.
    fn sort_keys(&self) -> bool {
        self.flags & Self::SORT_KEYS != 0
    }

    /// Returns whether non-ASCII characters must be escaped.
    fn ensure_ascii(&self) -> bool {
        self.flags & Self::ENSURE_ASCII != 0
    }

    /// Returns whether NaN and infinity may be emitted as JSON tokens.
    fn allow_nan(&self) -> bool {
        self.flags & Self::ALLOW_NAN != 0
    }

    /// Returns whether unsupported dict keys should be skipped.
    fn skipkeys(&self) -> bool {
        self.flags & Self::SKIPKEYS != 0
    }

    /// Builds the encoder config from the macro-extracted [`JsonDumpsArgs`].
    ///
    /// `JsonDumpsArgs` already performed the positional/kwargs split, the
    /// per-name duplicate detection, and the unknown-keyword rejection (with
    /// the "JSONEncoder.__init__" prefix that matches CPython's error
    /// wording). This method only translates the raw `Value` slots into the
    /// concrete encoder options used by the serializer.
    fn from_macro_args(args: JsonDumpsArgs, vm: &mut VM<'_, impl ResourceTracker>) -> RunResult<(Value, Self)> {
        let JsonDumpsArgs {
            obj,
            indent,
            sort_keys,
            ensure_ascii,
            allow_nan,
            separators,
            skipkeys,
        } = args;

        // Keep `obj` alive across kwarg processing — early errors below must
        // not leak the heap reference.
        let mut obj_guard = HeapGuard::new(obj, vm);
        let vm = obj_guard.heap();

        let mut config = Self::default();

        let indent = parse_indent_value(indent, vm)?;
        config.indent = indent;

        config.flags = apply_bool_flag(config.flags, Self::SORT_KEYS, sort_keys, vm);
        config.flags = apply_bool_flag(config.flags, Self::ENSURE_ASCII, ensure_ascii, vm);
        config.flags = apply_bool_flag(config.flags, Self::ALLOW_NAN, allow_nan, vm);
        config.flags = apply_bool_flag(config.flags, Self::SKIPKEYS, skipkeys, vm);

        // `separators=None` is documented as equivalent to "use the indent-
        // aware defaults", so we only override the per-instance separators
        // when `parse_separators_value` actually returned a pair.
        let separators_were_set = if let Some((item, key)) = parse_separators_value(separators, vm)? {
            config.item_separator = item;
            config.key_separator = key;
            true
        } else {
            false
        };

        if config.indent.is_some() && !separators_were_set {
            ",".clone_into(&mut config.item_separator);
            ": ".clone_into(&mut config.key_separator);
        }

        Ok((obj_guard.into_inner(), config))
    }
}

/// Implements `json.dumps(obj, **kwargs)`.
///
/// Only the first argument may be positional. Supported keyword arguments mirror
/// the high-value subset of CPython's encoder configuration: `indent`,
/// `sort_keys`, `ensure_ascii`, `allow_nan`, `separators`, and `skipkeys`.
///
/// CPython kwargs `cls`, `default`, and `check_circular` are intentionally
/// unsupported and will raise `TypeError` if passed.
pub(super) fn call_dumps(vm: &mut VM<'_, impl ResourceTracker>, args: ArgValues) -> RunResult<Value> {
    let macro_args = JsonDumpsArgs::from_args(args, vm)?;
    let (obj, config) = JsonDumpsConfig::from_macro_args(macro_args, vm)?;

    let mut obj_guard = HeapGuard::new(obj, vm);
    let mut output = String::new();
    let mut active_containers = Vec::new();
    {
        let (obj, vm) = obj_guard.as_parts_mut();
        let mut encoder = Encoder {
            out: &mut output,
            config: &config,
            active_containers: &mut active_containers,
            vm,
        };
        encoder.serialize_value(obj, 0)?;
    }

    let (obj, vm) = obj_guard.into_parts();
    obj.drop_with_heap(vm);
    Ok(allocate_string(output, vm.heap)?)
}

/// Argument shape for `json.dumps(obj, *, indent=None, sort_keys=False,
/// ensure_ascii=True, allow_nan=True, separators=None, skipkeys=False)`.
///
/// Arity and missing-arg errors use the `dumps()` descriptor with `style =
/// def` wording (CPython's `json.dumps` is a pure-Python `def`), but the
/// unknown-kwarg error uses `JSONEncoder.__init__()` — CPython's `json.dumps`
/// forwards unknown kwargs straight to the encoder constructor, which is what
/// surfaces in the error. `kwarg_error_name` overrides the function name used
/// in the unexpected-keyword message without affecting other error paths.
/// Every field is a raw `Value` so the encoder can apply its own truth-test
/// (`py_bool`) or shape coercion (`parse_indent_value` /
/// `parse_separators_value`) on the way through.
#[derive(FromArgs)]
#[from_args(name = "dumps", style = def, kwarg_error_name = "JSONEncoder.__init__")]
struct JsonDumpsArgs {
    obj: Value,
    #[from_args(kw_only, default = Value::None)]
    indent: Value,
    #[from_args(kw_only, default = Value::Bool(false))]
    sort_keys: Value,
    #[from_args(kw_only, default = Value::Bool(true))]
    ensure_ascii: Value,
    #[from_args(kw_only, default = Value::Bool(true))]
    allow_nan: Value,
    #[from_args(kw_only, default = Value::None)]
    separators: Value,
    #[from_args(kw_only, default = Value::Bool(false))]
    skipkeys: Value,
}

/// Sets `bit` in `flags` when `value` is truthy, clearing it otherwise. The
/// value is dropped afterwards. Used by the json.dumps kwarg pipeline so each
/// boolean-style flag is handled with a single line.
fn apply_bool_flag(flags: u8, bit: u8, value: Value, vm: &mut VM<'_, impl ResourceTracker>) -> u8 {
    let new_flags = if value.py_bool(vm) { flags | bit } else { flags & !bit };
    value.drop_with_heap(vm);
    new_flags
}

/// Parses the `indent=` value for `json.dumps()`.
///
/// `None` keeps compact mode, integers switch to pretty mode using that many
/// spaces per nesting level (with zero and negative values enabling newline-
/// only pretty printing), and
/// strings are repeated once per depth level exactly like CPython.
fn parse_indent_value(value: Value, vm: &mut VM<'_, impl ResourceTracker>) -> RunResult<Option<String>> {
    let mut value_guard = HeapGuard::new(value, vm);
    let (value, vm) = value_guard.as_parts_mut();

    match value {
        Value::None => Ok(None),
        Value::Bool(flag) => Ok(Some(" ".repeat(usize::from(*flag)))),
        Value::Int(count) => spaces_from_indent_count(*count),
        Value::InternString(string_id) => Ok(Some(vm.interns.get_str(*string_id).to_owned())),
        Value::Ref(heap_id) => match vm.heap.read(*heap_id) {
            HeapReadOutput::Str(string) => Ok(Some(string.get(vm.heap).as_str().to_owned())),
            HeapReadOutput::LongInt(long_int) => {
                spaces_from_indent_count(long_int.get(vm.heap).to_i64().ok_or_else(ExcType::overflow_c_ssize_t)?)
            }
            _ => Err(ExcType::type_error("indent must be None, an integer or a string")),
        },
        _ => Err(ExcType::type_error("indent must be None, an integer or a string")),
    }
}

/// Converts an integer indent width into the repeated-space string used per level.
///
/// Zero and negative values return an empty indent string, which keeps pretty
/// printing enabled while omitting leading spaces on each line like CPython.
fn spaces_from_indent_count(count: i64) -> RunResult<Option<String>> {
    if count <= 0 {
        Ok(Some(String::new()))
    } else {
        match usize::try_from(count) {
            Ok(count) => Ok(Some(" ".repeat(count))),
            Err(_) => Err(ExcType::overflow_c_ssize_t()),
        }
    }
}

/// Parses the `separators=` value for `json.dumps()`.
///
/// `None` leaves the default separators intact. Otherwise the value must be a
/// two-item list or tuple of strings representing the item and key separators.
fn parse_separators_value(value: Value, vm: &mut VM<'_, impl ResourceTracker>) -> RunResult<Option<(String, String)>> {
    let mut value_guard = HeapGuard::new(value, vm);
    let (value, vm) = value_guard.as_parts_mut();

    if matches!(value, Value::None) {
        return Ok(None);
    }

    let pair = match value {
        Value::Ref(heap_id) => match vm.heap.read(*heap_id) {
            HeapReadOutput::Tuple(tuple) => {
                let items = tuple.get(vm.heap).as_slice();
                check_separators_length(items.len())?;
                (
                    json_separator_to_string(&items[0], "item_separator", vm)?,
                    json_separator_to_string(&items[1], "key_separator", vm)?,
                )
            }
            HeapReadOutput::List(list) => {
                let items = list.get(vm.heap).as_slice();
                check_separators_length(items.len())?;
                (
                    json_separator_to_string(&items[0], "item_separator", vm)?,
                    json_separator_to_string(&items[1], "key_separator", vm)?,
                )
            }
            _ => {
                return Err(ExcType::type_error(format!(
                    "cannot unpack non-iterable {} object",
                    value.py_type_name(vm)
                )));
            }
        },
        _ => {
            return Err(ExcType::type_error(format!(
                "cannot unpack non-iterable {} object",
                value.py_type_name(vm)
            )));
        }
    };

    Ok(Some(pair))
}

/// Validates that the separators sequence has exactly two elements.
///
/// Raises `ValueError` with the same unpacking-style message as CPython when
/// the length does not match the expected two elements.
fn check_separators_length(len: usize) -> RunResult<()> {
    match len.cmp(&2) {
        Ordering::Greater => Err(ExcType::value_error(format!(
            "too many values to unpack (expected 2, got {len})"
        ))),
        Ordering::Less => Err(ExcType::value_error(format!(
            "not enough values to unpack (expected 2, got {len})"
        ))),
        Ordering::Equal => Ok(()),
    }
}

/// Converts a Monty value to a string for use as a JSON separator.
///
/// CPython's C encoder validates separators as strings and refers to them by
/// their positional argument index in `make_encoder()`. The `role` parameter
/// selects the matching CPython argument number (6 for `item_separator`,
/// 5 for `key_separator`) so the error message matches CPython exactly.
fn json_separator_to_string(value: &Value, role: &str, vm: &VM<'_, impl ResourceTracker>) -> RunResult<String> {
    let arg_num = if role == "item_separator" { 6 } else { 5 };
    match value {
        Value::InternString(string_id) => Ok(vm.interns.get_str(*string_id).to_owned()),
        Value::Ref(heap_id) => match vm.heap.get(*heap_id) {
            HeapData::Str(string) => Ok(string.as_str().to_owned()),
            _ => Err(ExcType::type_error(format!(
                "make_encoder() argument {arg_num} must be str, not {}",
                value.py_type_name(vm)
            ))),
        },
        _ => Err(ExcType::type_error(format!(
            "make_encoder() argument {arg_num} must be str, not {}",
            value.py_type_name(vm)
        ))),
    }
}

/// Mutable serialization context shared across the recursive `serialize_*`
/// methods.
///
/// Groups the four pieces of per-call state that every recursive level needs
/// — the JSON output buffer, the encoder configuration, the cycle-detection
/// stack, and the VM (heap + interns) — so the recursive methods can be
/// invoked as `encoder.serialize_value(value, depth)` rather than threading
/// five separate parameters through every signature.
///
/// `depth` is intentionally **not** a field of this struct: every recursive
/// call needs to bump it by one for the descent and restore it on return,
/// which is exactly what a stack-passed parameter does for free.
struct Encoder<'a, 'h, R: ResourceTracker> {
    out: &'a mut String,
    config: &'a JsonDumpsConfig,
    active_containers: &'a mut Vec<HeapId>,
    vm: &'a mut VM<'h, R>,
}

/// Lets the encoder participate in the [`HeapGuard`] / [`defer_drop_mut!`]
/// pattern: passing the encoder as the "heap" argument re-borrows the whole
/// encoder (including its `vm`) into the guard, which is exactly what we need
/// so the rebound iter and the rebound encoder share a lifetime.
impl<R: ResourceTracker> ContainsHeap for Encoder<'_, '_, R> {
    type ResourceTracker = R;

    fn heap(&self) -> &Heap<R> {
        self.vm.heap()
    }

    fn heap_mut(&mut self) -> &mut Heap<R> {
        self.vm.heap_mut()
    }
}

/// Lets a [`RecursionToken`](crate::bytecode::RecursionToken) (and the container
/// iterators that hold one) be released through the encoder via `defer_drop_vm!`,
/// reaching the VM-side recursion counter while the encoder itself stays borrowable.
impl<'h, R: ResourceTracker> ContainsVM<'h> for Encoder<'_, 'h, R> {
    type Tracker = R;

    fn vm(&mut self) -> &mut VM<'h, Self::Tracker> {
        self.vm
    }
}

impl<'h, R: ResourceTracker> Encoder<'_, 'h, R> {
    /// Serializes a Monty value into JSON text.
    ///
    /// Handles immediate primitives directly and delegates to type-specific
    /// helpers for strings, long integers, lists, tuples, and dicts.
    fn serialize_value(&mut self, value: &Value, depth: usize) -> RunResult<()> {
        match value {
            Value::None => {
                self.out.push_str("null");
                Ok(())
            }
            Value::Bool(true) => {
                self.out.push_str("true");
                Ok(())
            }
            Value::Bool(false) => {
                self.out.push_str("false");
                Ok(())
            }
            Value::Int(value) => {
                write!(self.out, "{value}").expect("writing to String cannot fail");
                Ok(())
            }
            Value::Float(value) => serialize_float(*value, self.out, self.config),
            Value::InternString(string_id) => {
                write_json_string(
                    self.vm.interns.get_str(*string_id),
                    self.out,
                    self.config.ensure_ascii(),
                );
                Ok(())
            }
            Value::InternLongInt(long_int_id) => {
                let value = self.vm.interns.get_long_int(*long_int_id);
                check_bigint_str_digits_limit(value)?;
                write!(self.out, "{value}").expect("writing to String cannot fail");
                Ok(())
            }
            Value::Ref(heap_id) => match self.vm.heap.read(*heap_id) {
                HeapReadOutput::Str(string) => {
                    write_json_string(string.get(self.vm.heap).as_str(), self.out, self.config.ensure_ascii());
                    Ok(())
                }
                HeapReadOutput::LongInt(long_int) => {
                    long_int.get(self.vm.heap).check_str_digits_limit()?;
                    write!(self.out, "{}", long_int.get(self.vm.heap).inner()).expect("writing to String cannot fail");
                    Ok(())
                }
                HeapReadOutput::List(list) => self.with_entered_container(*heap_id, |enc| {
                    let iter = list.iter(enc.vm)?;
                    defer_drop_vm_mut!(iter, enc);
                    enc.serialize_array(depth, |enc, depth| {
                        if let Some(item) = iter.next(enc.vm)? {
                            enc.serialize_value(item, depth)?;
                            Ok(true)
                        } else {
                            Ok(false)
                        }
                    })
                }),
                HeapReadOutput::Tuple(tuple) => self.with_entered_container(*heap_id, |enc| {
                    let iter = tuple.iter(enc.vm)?;
                    defer_drop_vm_mut!(iter, enc);
                    enc.serialize_array(depth, |enc, depth| {
                        if let Some(item) = iter.next(enc.vm)? {
                            enc.serialize_value(item, depth)?;
                            Ok(true)
                        } else {
                            Ok(false)
                        }
                    })
                }),
                HeapReadOutput::Dict(dict) => {
                    // Dict pre-materializes entries because `sort_keys` and
                    // `skipkeys` need to mutate the entries vector before
                    // output.
                    let entries = self.collect_dict_entries(&dict);
                    let this = self;
                    defer_drop_mut!(entries, this);
                    // Need to explicitly acquire a recursion token for the dict as we don't go
                    // via the default dict iterator.
                    let token = this.vm.recursion_token()?;
                    defer_drop_vm!(token, this);
                    this.with_entered_container(*heap_id, |enc| enc.serialize_dict(entries, depth))
                }
                _ => Err(ExcType::json_not_serializable_error(&value.py_type_name(self.vm))),
            },
            _ => Err(ExcType::json_not_serializable_error(&value.py_type_name(self.vm))),
        }
    }

    /// Streams items from a list or tuple into a JSON array, formatting
    /// separators and indentation between calls to `write_next`.
    ///
    /// `write_next` is invoked once per item slot and should write the JSON
    /// for the next item, returning `Ok(true)` if it wrote one or `Ok(false)`
    /// when the underlying iterator is exhausted. This shape lets the caller
    /// plug in either a [`ListIter`](crate::types::list::ListIter) or a
    /// [`TupleIter`](crate::types::tuple::TupleIter) without having to unify
    /// their lending-iterator types.
    ///
    /// Circular-reference tracking still happens at the container level (via
    /// [`with_entered_container`](Self::with_entered_container)) before this
    /// helper is called; the recursion-depth bound is enforced by the
    /// iterator's `RecursionToken`.
    fn serialize_array(
        &mut self,
        depth: usize,
        mut write_next: impl FnMut(&mut Self, usize) -> RunResult<bool>,
    ) -> RunResult<()> {
        self.out.push('[');
        let pretty = self.config.indent.is_some();
        let mut wrote_any = false;
        loop {
            // Reserve the separator + indent BEFORE we know whether the next
            // item exists; if `write_next` reports exhaustion we roll the
            // cursor back.
            let prefix_start = self.out.len();
            if wrote_any {
                self.out.push_str(&self.config.item_separator);
            }
            if pretty {
                self.out.push('\n');
                write_indent(self.out, self.config, depth + 1);
            }
            let body_start = self.out.len();
            if !write_next(self, depth + 1)? {
                self.out.truncate(prefix_start);
                break;
            }
            debug_assert!(
                self.out.len() > body_start,
                "write_next reported true but wrote nothing"
            );
            wrote_any = true;
        }
        if pretty && wrote_any {
            self.out.push('\n');
            write_indent(self.out, self.config, depth);
        }
        self.out.push(']');
        Ok(())
    }

    /// Copies a dict's `(key, value)` pairs into an owned `Vec` via
    /// [`DictIter`](crate::types::dict::DictIter), so subsequent passes
    /// (`skipkeys`, `sort_keys`) can mutate the buffer in place. The
    /// `DictIter`'s recursion token is acquired and released during the
    /// copy, bounding the depth of the *enclosing* `serialize_value` call.
    fn collect_dict_entries(&mut self, dict: &HeapRead<'h, Dict>) -> Vec<(Value, Value)> {
        dict.get(self.vm.heap)
            .iter()
            .map(|(k, v)| (k.clone_with_heap(self.vm.heap), v.clone_with_heap(self.vm.heap)))
            .collect::<Vec<_>>()
    }

    /// Serializes a dict as a JSON object.
    ///
    /// Dict keys are validated and optionally skipped before serialization.
    /// When `sort_keys=True`, entries are sorted using Python comparison
    /// semantics on the original keys so mixed incomparable key types raise
    /// the same style of `TypeError` as CPython.
    fn serialize_dict(&mut self, entries: &mut Vec<(Value, Value)>, depth: usize) -> RunResult<()> {
        if self.config.skipkeys() {
            skip_disallowed_dict_keys(entries, self.vm);
        } else if let Some((key, _)) = entries.iter().find(|(key, _)| !is_json_key_allowed(key, self.vm)) {
            return Err(ExcType::json_invalid_key_error(&key.py_type_name(self.vm)));
        }

        if self.config.sort_keys() {
            sort_dict_entries(entries, self.vm)?;
        }

        self.out.push('{');

        let pretty = self.config.indent.is_some();
        for (index, (key, value)) in entries.iter().enumerate() {
            if index != 0 {
                self.out.push_str(&self.config.item_separator);
            }
            if pretty {
                self.out.push('\n');
                write_indent(self.out, self.config, depth + 1);
            }
            write_json_key(key, self.out, self.config, self.vm)?;
            self.out.push_str(&self.config.key_separator);
            self.serialize_value(value, depth + 1)?;
        }
        if pretty && !entries.is_empty() {
            self.out.push('\n');
            write_indent(self.out, self.config, depth);
        }
        self.out.push('}');
        Ok(())
    }

    /// Runs `f` while `heap_id` is marked active for cycle detection.
    ///
    /// Centralizes the push/pop bookkeeping so every serialization path pops
    /// the container again regardless of whether recursive serialization
    /// succeeds or returns early with an error.
    fn with_entered_container<T>(
        &mut self,
        heap_id: HeapId,
        f: impl FnOnce(&mut Self) -> RunResult<T>,
    ) -> RunResult<T> {
        if self.active_containers.contains(&heap_id) {
            return Err(ExcType::json_circular_reference_error());
        }
        self.active_containers.push(heap_id);
        let result = f(self);
        self.active_containers
            .pop()
            .expect("entered container missing from JSON serialization stack");
        result
    }
}

/// Sorts dict entries in-place using Python comparison semantics on the keys.
///
/// The implementation mirrors the error style used by `sorted()` and
/// `list.sort()`: when two keys are not orderable, it raises
/// `TypeError: '<' not supported between instances of ...`.
fn sort_dict_entries(entries: &mut Vec<(Value, Value)>, vm: &mut VM<'_, impl ResourceTracker>) -> RunResult<()> {
    let mut indices: Vec<usize> = (0..entries.len()).collect();
    let compare_values: Vec<Value> = entries.iter().map(|(key, _)| key.clone_with_heap(vm)).collect();
    let mut compare_values_guard = HeapGuard::new(compare_values, vm);
    let (compare_values, vm) = compare_values_guard.as_parts_mut();
    sort_indices(&mut indices, compare_values.as_slice(), false, vm)?;
    apply_permutation(entries.as_mut_slice(), &mut indices);
    Ok(())
}

/// Removes dict entries whose keys are not JSON-serializable, preserving order.
///
/// `skipkeys=True` must drop invalid entries without disturbing the relative
/// order of the retained pairs. A two-pointer compaction avoids the repeated
/// shifting cost of `Vec::remove(i)` while still cleaning up skipped `Value`
/// references with `drop_with_heap`.
fn skip_disallowed_dict_keys(entries: &mut Vec<(Value, Value)>, vm: &mut VM<'_, impl ResourceTracker>) {
    let mut write = 0;
    for read in 0..entries.len() {
        if is_json_key_allowed(&entries[read].0, vm) {
            if write != read {
                entries.swap(write, read);
            }
            write += 1;
        }
    }

    for (key, value) in entries.drain(write..) {
        key.drop_with_heap(vm);
        value.drop_with_heap(vm);
    }
}

/// Returns whether a value is an allowed JSON object key type.
///
/// CPython accepts strings, integers, floats, booleans, and `None`, then
/// coerces the non-string cases to JSON strings during output.
fn is_json_key_allowed(value: &Value, vm: &VM<'_, impl ResourceTracker>) -> bool {
    matches!(
        value,
        Value::None | Value::Bool(_) | Value::Int(_) | Value::Float(_) | Value::InternString(_)
    ) || matches!(value, Value::Ref(heap_id) if matches!(vm.heap.get(*heap_id), HeapData::Str(_) | HeapData::LongInt(_)))
}

/// Serializes a dict key by applying CPython's JSON key coercions.
///
/// Non-string supported key types are rendered to their JSON string form first,
/// then escaped as a JSON string token.
fn write_json_key(
    key: &Value,
    out: &mut String,
    config: &JsonDumpsConfig,
    vm: &VM<'_, impl ResourceTracker>,
) -> RunResult<()> {
    let ensure_ascii = config.ensure_ascii();
    match key {
        Value::None => write_json_ascii_key("null", out),
        Value::Bool(true) => write_json_ascii_key("true", out),
        Value::Bool(false) => write_json_ascii_key("false", out),
        Value::Int(value) => write_json_display_key(value, out),
        Value::Float(value) => {
            serialize_float_key(*value, out, config)?;
        }
        Value::InternString(string_id) => write_json_string(vm.interns.get_str(*string_id), out, ensure_ascii),
        Value::Ref(heap_id) => match vm.heap.get(*heap_id) {
            HeapData::Str(string) => write_json_string(string.as_str(), out, ensure_ascii),
            HeapData::LongInt(long_int) => {
                long_int.check_str_digits_limit()?;
                write_json_display_key(long_int.inner(), out);
            }
            _ => return Err(ExcType::json_invalid_key_error(&key.py_type_name(vm))),
        },
        _ => return Err(ExcType::json_invalid_key_error(&key.py_type_name(vm))),
    }
    Ok(())
}

/// Writes an already-ASCII JSON object key without going through the string
/// escaper.
///
/// Coerced keys such as `None`, booleans, and numeric reprs are always ASCII
/// and require no escaping, so this avoids building intermediate `String`
/// values on the dict-key hot path.
fn write_json_ascii_key(value: &str, out: &mut String) {
    out.push('"');
    out.push_str(value);
    out.push('"');
}

/// Writes a displayable value as a quoted JSON object key.
///
/// The caller is responsible for ensuring the formatted output is ASCII-safe
/// and does not require JSON string escaping.
fn write_json_display_key(value: impl Display, out: &mut String) {
    out.push('"');
    write!(out, "{value}").expect("writing to String cannot fail");
    out.push('"');
}

/// Serializes a float value as a quoted JSON object key, respecting `allow_nan`.
///
/// Non-finite values (NaN, +/-Infinity) are emitted as their Python repr
/// (`NaN`, `Infinity`, `-Infinity`) when `allow_nan` is enabled. When disabled,
/// the same `ValueError` raised for non-finite float *values* applies to keys
/// too, matching CPython's behavior.
fn serialize_float_key(value: f64, out: &mut String, config: &JsonDumpsConfig) -> RunResult<()> {
    out.push('"');
    if value.is_nan() {
        if !config.allow_nan() {
            return Err(ExcType::json_nan_error("nan"));
        }
        out.push_str("NaN");
    } else if value == f64::INFINITY {
        if !config.allow_nan() {
            return Err(ExcType::json_nan_error("inf"));
        }
        out.push_str("Infinity");
    } else if value == f64::NEG_INFINITY {
        if !config.allow_nan() {
            return Err(ExcType::json_nan_error("-inf"));
        }
        out.push_str("-Infinity");
    } else {
        write_json_float_text(value, out);
    }
    out.push('"');
    Ok(())
}

/// Serializes a float using JSON's number and NaN rules.
///
/// Finite floats use CPython-compatible `json` float formatting, including the
/// switch to exponent notation for very small or very large magnitudes while
/// still preserving a decimal point for whole-valued non-exponent outputs.
fn serialize_float(value: f64, out: &mut String, config: &JsonDumpsConfig) -> RunResult<()> {
    if value.is_nan() {
        if config.allow_nan() {
            out.push_str("NaN");
            Ok(())
        } else {
            Err(ExcType::json_nan_error("nan"))
        }
    } else if value == f64::INFINITY {
        if config.allow_nan() {
            out.push_str("Infinity");
            Ok(())
        } else {
            Err(ExcType::json_nan_error("inf"))
        }
    } else if value == f64::NEG_INFINITY {
        if config.allow_nan() {
            out.push_str("-Infinity");
            Ok(())
        } else {
            Err(ExcType::json_nan_error("-inf"))
        }
    } else {
        write_json_float_text(value, out);
        Ok(())
    }
}

/// Writes a finite float using CPython-compatible JSON float repr rules.
///
/// Python switches to scientific notation when the magnitude is `>= 1e16` or
/// `< 1e-4` (and non-zero). We decide notation by comparing the absolute value
/// directly against these thresholds rather than using `log10().floor()`, which
/// has precision errors at boundary values (e.g. `9999999999999998.0` whose
/// `log10` rounds up to `16.0`). Direct comparison is exact because `1e16` is
/// exactly representable as f64 and `1e-4` as an f64 constant aligns with
/// Python's notation boundary.
///
/// Each path formats the float exactly once: the scientific path writes via
/// `{:e}` and post-processes the exponent to Python style (`e+XX` / `e-XX`),
/// while the fixed path uses `Display` with a `.0` suffix for whole numbers.
fn write_json_float_text(value: f64, out: &mut String) {
    let abs = value.abs();
    if abs != 0.0 && !(1e-4..1e16).contains(&abs) {
        // Python-style scientific notation: format via `{:e}`, then rewrite the
        // exponent from Rust's bare `e<N>` to Python's `e+XX` / `e-XX`.
        let start = out.len();
        write!(out, "{value:e}").expect("writing to String cannot fail");
        let e_pos = out[start..].find('e').expect("scientific format must contain 'e'") + start;
        let exponent: i32 = out[e_pos + 1..].parse().expect("exponent must be a valid integer");
        out.truncate(e_pos);
        let exp_sign = if exponent >= 0 { '+' } else { '-' };
        write!(out, "e{exp_sign}{:02}", exponent.unsigned_abs()).expect("writing to String cannot fail");
    } else {
        // Fixed notation: single `Display` write, appending `.0` for whole numbers.
        let start = out.len();
        write!(out, "{value}").expect("writing to String cannot fail");
        if !out[start..].contains('.') {
            out.push_str(".0");
        }
    }
}

/// Writes indentation for pretty-printed JSON output.
///
/// The `indent` string is repeated once for each nesting level, matching
/// CPython's behavior for both numeric and string indentation.
fn write_indent(out: &mut String, config: &JsonDumpsConfig, depth: usize) {
    if let Some(indent) = &config.indent {
        for _ in 0..depth {
            out.push_str(indent);
        }
    }
}

/// Writes a Rust string as a JSON string token.
///
/// Uses a byte-oriented batch strategy inspired by serde_json: a 256-entry
/// lookup table classifies each byte in O(1), and contiguous runs of safe bytes
/// are flushed with a single `push_str` rather than character-by-character.
///
/// When `ensure_ascii` is enabled, non-ASCII code points (bytes >= 0x80) are
/// emitted as `\uXXXX` escapes using surrogate pairs for supplementary-plane
/// characters.
fn write_json_string(value: &str, out: &mut String, ensure_ascii: bool) {
    out.push('"');
    let bytes = value.as_bytes();
    let mut start = 0;
    let mut i = 0;

    while i < bytes.len() {
        let byte = bytes[i];

        if ensure_ascii && byte >= 0x7F {
            // Flush the safe ASCII run accumulated so far.
            out.push_str(&value[start..i]);
            if byte == 0x7F {
                // DEL (0x7F) is a control character that CPython escapes.
                out.push_str("\\u007f");
                i += 1;
            } else {
                // Decode the full character at this position and emit \uXXXX escapes.
                let ch = value[i..].chars().next().expect("valid UTF-8");
                write_json_escape_for_non_ascii(ch, out);
                i += ch.len_utf8();
            }
            start = i;
            continue;
        }

        let escape = ESCAPE_TABLE[byte as usize];
        if escape == 0 {
            // Safe byte — keep scanning.
            i += 1;
            continue;
        }

        // Flush the safe run before this byte.
        out.push_str(&value[start..i]);

        // Write the escape sequence.
        match escape {
            b'b' => out.push_str("\\b"),
            b't' => out.push_str("\\t"),
            b'n' => out.push_str("\\n"),
            b'f' => out.push_str("\\f"),
            b'r' => out.push_str("\\r"),
            b'"' => out.push_str("\\\""),
            b'\\' => out.push_str("\\\\"),
            b'u' => {
                write!(out, "\\u{:04x}", u32::from(byte)).expect("writing to String cannot fail");
            }
            _ => unreachable!(),
        }

        i += 1;
        start = i;
    }

    // Flush the final safe run.
    out.push_str(&value[start..]);
    out.push('"');
}

/// Byte lookup table for JSON string escaping.
///
/// Each entry is either 0 (byte is safe, no escaping needed) or a shorthand
/// character that indicates which escape to emit:
/// - `b'"'`  → `\"`
/// - `b'\\'` → `\\`
/// - `b'b'`  → `\b` (backspace, 0x08)
/// - `b't'`  → `\t` (tab, 0x09)
/// - `b'n'`  → `\n` (newline, 0x0A)
/// - `b'f'`  → `\f` (form feed, 0x0C)
/// - `b'r'`  → `\r` (carriage return, 0x0D)
/// - `b'u'`  → `\u00XX` (other control characters, 0x00–0x1F)
#[rustfmt::skip]
static ESCAPE_TABLE: [u8; 256] = {
    let mut table = [0u8; 256];
    // Control characters 0x00–0x1F default to \u00XX escapes.
    let mut i = 0;
    while i < 0x20 {
        table[i] = b'u';
        i += 1;
    }
    // Override the named escapes.
    table[0x08] = b'b';  // backspace
    table[0x09] = b't';  // tab
    table[0x0A] = b'n';  // newline
    table[0x0C] = b'f';  // form feed
    table[0x0D] = b'r';  // carriage return
    table[0x22] = b'"';  // quote
    table[0x5C] = b'\\'; // backslash
    table
};

/// Writes a non-ASCII character using JSON `\uXXXX` escapes.
///
/// Code points above `U+FFFF` are encoded as UTF-16 surrogate pairs to match
/// CPython's `ensure_ascii=True` behavior.
fn write_json_escape_for_non_ascii(ch: char, out: &mut String) {
    let code = ch as u32;
    if code <= 0xFFFF {
        write!(out, "\\u{code:04x}").expect("writing to String cannot fail");
    } else {
        let code = code - 0x1_0000;
        let high = 0xD800 + (code >> 10);
        let low = 0xDC00 + (code & 0x3FF);
        write!(out, "\\u{high:04x}\\u{low:04x}").expect("writing to String cannot fail");
    }
}