datui_lib/inspector/

inspector_drill.rs

//! Drilling into a nested value from the row inspector: a struct's fields, a
//! list's items, or the objects and arrays of JSON stored as text.
//!
//! A level holds where it is, never a copy of what is there: a column's value is
//! a one-row slice of the buffer's series, and a JSON value is a path into a
//! document parsed once. Only the items on screen are ever turned into rows, so a
//! list of a million items costs as much to show as a list of ten.

use polars::prelude::*;
use serde_json::Value;
use std::sync::Arc;

/// Text up to this long is parsed as JSON on the key that asks; longer text is
/// parsed off the event thread.
pub const JSON_INLINE_BYTES: usize = 64 * 1024;
/// Text longer than this is not parsed: its tree would take several times its size.
pub const JSON_MAX_BYTES: usize = 4 * 1024 * 1024;
/// Keys whose widths size a column of names: an object's keys past this are fitted
/// to the width found, not measured.
const MEASURED: usize = 1000;

/// One place in a value: what a level of the drill shows, or one of its items.
#[derive(Debug, Clone)]
pub enum Node {
    /// One value of a column, as a one-row slice: drilling copies nothing.
    Native(Series),
    /// A value in a JSON document parsed from text: the steps from its root.
    Json { root: Arc<Value>, path: Vec<Step> },
}

/// One step into a JSON document. An object's step is its key, not its position:
/// serde_json's map has no lookup by position, and walking to the 500,000th key of
/// a large object for every item drawn, on every frame, is what a key avoids.
#[derive(Debug, Clone)]
pub enum Step {
    Key(Arc<str>),
    Index(usize),
}

/// What a node holds, as the drill sees it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Shape {
    Struct,
    List,
    Object,
    Array,
    /// A scalar, text, bytes or a null: nothing to drill into.
    Leaf,
}

impl Shape {
    /// What the rule over a level's items calls them.
    pub fn items_title(self) -> &'static str {
        match self {
            Shape::Struct => "Fields",
            Shape::Object => "Keys",
            _ => "Items",
        }
    }
}

impl Node {
    /// The JSON value at this node's path.
    pub fn json(&self) -> Option<&Value> {
        let Node::Json { root, path } = self else {
            return None;
        };
        let mut at: &Value = root;
        for step in path {
            at = match (at, step) {
                (Value::Object(map), Step::Key(key)) => map.get(&**key)?,
                (Value::Array(items), Step::Index(i)) => items.get(*i)?,
                _ => return None,
            };
        }
        Some(at)
    }

    fn is_null(&self) -> bool {
        match self {
            Node::Native(s) => s.null_count() == s.len(),
            Node::Json { .. } => matches!(self.json(), None | Some(Value::Null)),
        }
    }

    pub fn shape(&self) -> Shape {
        if self.is_null() {
            return Shape::Leaf;
        }
        match self {
            Node::Native(s) => match s.dtype() {
                DataType::Struct(_) => Shape::Struct,
                DataType::List(_) | DataType::Array(..) => Shape::List,
                _ => Shape::Leaf,
            },
            Node::Json { .. } => match self.json() {
                Some(Value::Object(_)) => Shape::Object,
                Some(Value::Array(_)) => Shape::Array,
                _ => Shape::Leaf,
            },
        }
    }

    /// A list's items, as one series.
    fn items(&self) -> Option<Series> {
        let Node::Native(s) = self else {
            return None;
        };
        match s.get(0).ok()? {
            AnyValue::List(inner) | AnyValue::Array(inner, _) => Some(inner),
            _ => None,
        }
    }

    /// A struct's fields, each a one-row series named for its field.
    fn fields(&self) -> Option<Vec<Series>> {
        let Node::Native(s) = self else {
            return None;
        };
        Some(s.struct_().ok()?.fields_as_series())
    }

    /// How many fields, items or keys a level of this node lists.
    pub fn len(&self) -> usize {
        match self.shape() {
            Shape::Struct => self.fields().map_or(0, |f| f.len()),
            Shape::List => self.items().map_or(0, |s| s.len()),
            Shape::Object | Shape::Array => match self.json() {
                Some(Value::Object(map)) => map.len(),
                Some(Value::Array(items)) => items.len(),
                _ => 0,
            },
            Shape::Leaf => 0,
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn json_child(&self, step: Step) -> Node {
        let Node::Json { root, path } = self else {
            unreachable!("only a JSON node has JSON children")
        };
        let mut path = path.clone();
        path.push(step);
        Node::Json {
            root: Arc::clone(root),
            path,
        }
    }

    /// `count` items from `start`, each with its label: a field's or key's name, or
    /// `[i]`. Only these are made, whatever the length.
    pub fn children(&self, start: usize, count: usize) -> Vec<(String, Node)> {
        match self.shape() {
            Shape::Struct => self
                .fields()
                .unwrap_or_default()
                .into_iter()
                .skip(start)
                .take(count)
                .map(|s| (s.name().to_string(), Node::Native(s)))
                .collect(),
            Shape::List => {
                let Some(items) = self.items() else {
                    return Vec::new();
                };
                let end = items.len().min(start.saturating_add(count));
                (start..end)
                    .map(|i| (format!("[{i}]"), Node::Native(items.slice(i as i64, 1))))
                    .collect()
            }
            Shape::Object => match self.json() {
                Some(Value::Object(map)) => page(map.keys(), start, count)
                    .into_iter()
                    .map(|key| (key.clone(), self.json_child(Step::Key(key.as_str().into()))))
                    .collect(),
                _ => Vec::new(),
            },
            Shape::Array => {
                let end = self.len().min(start.saturating_add(count));
                (start..end)
                    .map(|i| (format!("[{i}]"), self.json_child(Step::Index(i))))
                    .collect()
            }
            Shape::Leaf => Vec::new(),
        }
    }

    pub fn child(&self, i: usize) -> Option<(String, Node)> {
        self.children(i, 1).pop()
    }

    /// The widest label a level of this node lists, without walking a long list:
    /// a list's last index is its widest, and only the first keys are measured.
    pub fn label_width(&self) -> usize {
        match self.shape() {
            Shape::List | Shape::Array => format!("[{}]", self.len().saturating_sub(1)).len(),
            Shape::Struct => self
                .fields()
                .unwrap_or_default()
                .iter()
                .map(|s| crate::glyphs::cell_width(s.name()))
                .max()
                .unwrap_or(0),
            Shape::Object => match self.json() {
                Some(Value::Object(map)) => map
                    .keys()
                    .take(MEASURED)
                    .map(|k| crate::glyphs::cell_width(k))
                    .max()
                    .unwrap_or(0),
                _ => 0,
            },
            Shape::Leaf => 0,
        }
    }

    /// The text of a text value, for `f`; None for any other value.
    pub fn with_text<R>(&self, f: impl FnOnce(&str) -> R) -> Option<R> {
        match self {
            Node::Native(s) => match s.get(0).ok()? {
                AnyValue::String(text) => Some(f(text)),
                AnyValue::StringOwned(text) => Some(f(text.as_str())),
                _ => None,
            },
            Node::Json { .. } => match self.json()? {
                Value::String(text) => Some(f(text)),
                _ => None,
            },
        }
    }

    /// Whether Enter opens this node as a level: a struct, a list, an object or an
    /// array, or text that reads as a JSON object or array.
    pub fn opens(&self) -> bool {
        self.shape() != Shape::Leaf || self.with_text(opens_as_json).unwrap_or(false)
    }

    /// The type as the item list names it.
    pub fn type_label(&self) -> String {
        match self {
            Node::Native(s) => crate::formats::column_types::dtype_label(s.dtype()),
            Node::Json { .. } => json_kind(self.json().unwrap_or(&Value::Null)).to_string(),
        }
    }

    /// The type the item list colors this node's name by.
    pub fn color_dtype(&self) -> DataType {
        match self {
            Node::Native(s) => s.dtype().clone(),
            Node::Json { .. } => json_dtype(self.json().unwrap_or(&Value::Null)),
        }
    }

    /// The columns a list of structs, or an array of objects, shows its items in,
    /// with the type that colors each name: the struct's fields, or the first
    /// object's keys. None for any other level.
    pub fn table_columns(&self) -> Option<Vec<(String, DataType)>> {
        match self.shape() {
            Shape::List => match self.items()?.dtype() {
                DataType::Struct(fields) if !fields.is_empty() => Some(
                    fields
                        .iter()
                        .map(|f| (f.name().to_string(), f.dtype().clone()))
                        .collect(),
                ),
                _ => None,
            },
            Shape::Array => match self.json()? {
                Value::Array(items) => match items.first()? {
                    Value::Object(first) if !first.is_empty() => Some(
                        first
                            .iter()
                            .map(|(k, v)| (k.clone(), json_dtype(v)))
                            .collect(),
                    ),
                    _ => None,
                },
                _ => None,
            },
            _ => None,
        }
    }

    /// One item's value under `column` in a level shown as a table; None when the
    /// item has no such field or key.
    pub fn cell(&self, column: &str) -> Option<Node> {
        match self.shape() {
            // One field made, not all of them for each cell of a row.
            Shape::Struct => match self {
                Node::Native(s) => s
                    .struct_()
                    .ok()?
                    .field_by_name(column)
                    .ok()
                    .map(Node::Native),
                Node::Json { .. } => None,
            },
            Shape::Object => match self.json()? {
                Value::Object(map) if map.contains_key(column) => {
                    Some(self.json_child(Step::Key(column.into())))
                }
                _ => None,
            },
            _ => None,
        }
    }
}

/// Items `start..start + count` of `items`, walked to from the nearer end: an
/// object's keys can only be stepped through, and `End` on a large one is then as
/// quick as `Home`.
fn page<T>(
    items: impl DoubleEndedIterator<Item = T> + ExactSizeIterator,
    start: usize,
    count: usize,
) -> Vec<T> {
    let len = items.len();
    let end = len.min(start.saturating_add(count));
    if start >= end {
        return Vec::new();
    }
    if start <= len - end {
        return items.skip(start).take(end - start).collect();
    }
    let mut back: Vec<T> = items.rev().skip(len - end).take(end - start).collect();
    back.reverse();
    back
}

/// The type a JSON value's name is colored by, as a column of that type is.
fn json_dtype(value: &Value) -> DataType {
    match value {
        Value::String(_) => DataType::String,
        Value::Bool(_) => DataType::Boolean,
        Value::Number(n) if n.is_f64() => DataType::Float64,
        Value::Number(_) => DataType::Int64,
        Value::Array(_) => DataType::List(Box::new(DataType::Null)),
        Value::Object(_) => DataType::Struct(Vec::new()),
        Value::Null => DataType::Null,
    }
}

/// A JSON value's kind, as the item list names it.
pub fn json_kind(value: &Value) -> &'static str {
    match value {
        Value::Null => "null",
        Value::Bool(_) => "bool",
        Value::Number(_) => "number",
        Value::String(_) => "str",
        Value::Array(_) => "array",
        Value::Object(_) => "object",
    }
}

/// Whether `text` reads as a JSON object or array: it starts and ends with the
/// brackets of one. Whether it parses is learned only by parsing it.
pub fn looks_like_json(text: &str) -> bool {
    let t = text.trim();
    matches!(
        (t.as_bytes().first(), t.as_bytes().last()),
        (Some(b'{'), Some(b'}')) | (Some(b'['), Some(b']'))
    )
}

/// Whether Enter offers to open `text` as JSON: it reads as an object or array and
/// is not over [`JSON_MAX_BYTES`], past which Enter shows more of it instead.
pub fn opens_as_json(text: &str) -> bool {
    text.len() <= JSON_MAX_BYTES && looks_like_json(text)
}

/// Parse `text` as JSON, refusing text over [`JSON_MAX_BYTES`]. serde_json stops at
/// 128 levels of nesting, so a deep document is an error, not a stack overflow.
pub fn parse_json(text: &str) -> Result<Value, String> {
    if text.len() > JSON_MAX_BYTES {
        return Err(format!(
            "{} MiB of text is over the {} MiB that opens as JSON",
            text.len().div_ceil(1024 * 1024),
            JSON_MAX_BYTES / (1024 * 1024)
        ));
    }
    serde_json::from_str(text).map_err(|e| format!("not JSON: {e}"))
}

/// Writes up to `cap` bytes, then fails, so serializing a huge value stops there.
struct Capped {
    out: Vec<u8>,
    cap: usize,
    cut: bool,
}

impl std::io::Write for Capped {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        let room = self.cap.saturating_sub(self.out.len());
        if bytes.len() > room {
            self.out.extend_from_slice(&bytes[..room]);
            self.cut = true;
            return Err(std::io::Error::other("cap reached"));
        }
        self.out.extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

/// JSON on one line with a space after each comma and colon, as datui writes a
/// list or struct on one line.
struct Spaced;

impl serde_json::ser::Formatter for Spaced {
    fn begin_array_value<W: ?Sized + std::io::Write>(
        &mut self,
        writer: &mut W,
        first: bool,
    ) -> std::io::Result<()> {
        if first {
            Ok(())
        } else {
            writer.write_all(b", ")
        }
    }

    fn begin_object_key<W: ?Sized + std::io::Write>(
        &mut self,
        writer: &mut W,
        first: bool,
    ) -> std::io::Result<()> {
        if first {
            Ok(())
        } else {
            writer.write_all(b", ")
        }
    }

    fn begin_object_value<W: ?Sized + std::io::Write>(
        &mut self,
        writer: &mut W,
    ) -> std::io::Result<()> {
        writer.write_all(b": ")
    }
}

/// `value` as JSON text, indented or on one line, stopping at `cap` bytes. True
/// when it was cut there.
pub fn json_text(value: &Value, pretty: bool, cap: usize) -> (String, bool) {
    use serde::Serialize;
    let mut w = Capped {
        out: Vec::new(),
        cap,
        cut: false,
    };
    // An error here is only the cap: serializing a parsed value cannot fail otherwise.
    let _ = if pretty {
        serde_json::to_writer_pretty(&mut w, value)
    } else {
        value.serialize(&mut serde_json::Serializer::with_formatter(&mut w, Spaced))
    };
    let text = match String::from_utf8(w.out) {
        Ok(text) => text,
        // Cut inside a character: keep what is whole.
        Err(e) => {
            let valid = e.utf8_error().valid_up_to();
            let mut bytes = e.into_bytes();
            bytes.truncate(valid);
            String::from_utf8(bytes).unwrap_or_default()
        }
    };
    (text, w.cut)
}

/// A JSON value as copy text: text as itself, a null as empty, anything else as
/// JSON, indented when it is an object or array. None when it is over `cap` bytes.
pub fn json_copy_text(value: &Value, cap: usize) -> Option<String> {
    match value {
        Value::String(s) => (s.len() <= cap).then(|| s.clone()),
        Value::Null => Some(String::new()),
        v => {
            let (text, cut) = json_text(v, true, cap);
            (!cut).then_some(text)
        }
    }
}

/// One level of a drill: the value it lists, and the item focused in it.
#[derive(Debug, Clone)]
pub struct Level {
    /// Its step in the breadcrumb: a field's or key's name, or `[i]`.
    pub label: String,
    pub node: Node,
    pub selected: usize,
}

impl Level {
    /// The focused item, and its label.
    pub fn focused(&self) -> Option<(String, Node)> {
        self.node.child(self.selected)
    }
}

/// The levels opened under one field of one row, outermost first.
#[derive(Debug, Clone)]
pub struct Drill {
    pub frame: u64,
    pub row: usize,
    pub levels: Vec<Level>,
}

impl Drill {
    pub fn level(&self) -> &Level {
        self.levels.last().expect("a drill has a level")
    }

    pub fn level_mut(&mut self) -> &mut Level {
        self.levels.last_mut().expect("a drill has a level")
    }
}

/// Text being parsed as JSON off the event thread, and where its level opens.
#[derive(Debug, Clone)]
pub struct JsonWait {
    pub token: u64,
    pub frame: u64,
    pub row: usize,
    pub label: String,
    /// The text's place, as [`path_key`] names it.
    pub path: String,
}

/// A place in a row, for telling one item from another: the field, then each step
/// opened and the item. Unit separators cannot be typed into a name, so two places
/// never meet.
pub fn path_key<'a>(steps: impl IntoIterator<Item = &'a str>) -> String {
    steps.into_iter().collect::<Vec<_>>().join("\u{1f}")
}

impl Drill {
    /// The place of the item `label` in the level shown.
    pub fn item_key(&self, label: &str) -> String {
        path_key(
            self.levels
                .iter()
                .map(|l| l.label.as_str())
                .chain(std::iter::once(label)),
        )
    }
}

#[cfg(test)]
mod tests;