rocksgraph 0.1.0

// Copyright (c) 2026 Austin Han <austinhan1024@gmail.com>
//
// This file is part of RocksGraph.
//
// RocksGraph is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 2 of the License, or
// (at your option) any later version.
//
// RocksGraph is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with RocksGraph.  If not, see <https://www.gnu.org/licenses/>.

//! Terminal execution types: [`BuiltTraversal`] (the lazy iterator) and the
//! internal [`materialize`] function that converts engine values into user-facing
//! [`Value`]s.

use std::collections::HashMap;

use smol_str::SmolStr;

use crate::{
    engine::volcano::builder::PhysicalPlan,
    gremlin::{
        type_bridge::primitive_to_value,
        value::{Edge as UserEdge, Map, Path, Property as UserProperty, Value, Vertex as UserVertex},
    },
    schema::Schema,
    types::{
        gvalue::{GValue, Primitive},
        keys::{CanonicalKey, LabelId, VertexKey},
        prop_key::{ID_KEY_ID, LABEL_KEY_ID, RANK_KEY_ID},
        StoreError,
    },
};

/// Pre-built lookup table: label_id → SmolStr, avoiding per-result BiHashMap lookups.
/// Pre-built lookup table: label_id/prop_key_id ↔ SmolStr, avoiding per-result BiHashMap lookups.
pub(crate) struct SchemaCache {
    vertex_labels: HashMap<LabelId, SmolStr>,
    edge_labels: HashMap<LabelId, SmolStr>,
    prop_key_str: HashMap<u16, SmolStr>,
    prop_key_id: HashMap<SmolStr, u16>,
}

impl SchemaCache {
    pub(crate) fn from_schema(schema: &Schema) -> Self {
        // Iterate all known vertex label ids and resolve eagerly.
        let mut vertex_labels = HashMap::new();
        for id in 1..=schema.vertex_labels_count() as LabelId {
            if let Some(name) = schema.vertex_label_str(id) {
                vertex_labels.insert(id, name.clone());
            }
        }
        let mut edge_labels = HashMap::new();
        for id in 1..=schema.edge_labels_count() as LabelId {
            if let Some(name) = schema.edge_label_str(id) {
                edge_labels.insert(id, name.clone());
            }
        }
        let mut prop_key_str = HashMap::new();
        let mut prop_key_id = HashMap::new();
        for (&id, key) in &schema.prop_keys {
            prop_key_str.insert(id, key.clone());
            prop_key_id.insert(key.clone(), id);
        }
        Self { vertex_labels, edge_labels, prop_key_str, prop_key_id }
    }

    #[inline]
    pub(crate) fn vertex_label(&self, label_id: LabelId) -> &SmolStr {
        self.vertex_labels.get(&label_id).unwrap_or_else(|| {
            static EMPTY: SmolStr = SmolStr::new_inline("");
            &EMPTY
        })
    }

    #[inline]
    pub(crate) fn edge_label(&self, label_id: LabelId) -> &SmolStr {
        self.edge_labels.get(&label_id).unwrap_or_else(|| {
            static EMPTY: SmolStr = SmolStr::new_inline("");
            &EMPTY
        })
    }

    #[inline]
    pub(crate) fn prop_key_str(&self, id: u16) -> Option<&SmolStr> {
        self.prop_key_str.get(&id)
    }

    #[inline]
    pub(crate) fn prop_key_id(&self, name: &str) -> Option<u16> {
        self.prop_key_id.get(name).copied()
    }
}

/// Materialize an internal [`GValue`] into a user-facing [`Value`].
///
/// `prop_keys` controls property fetching:
/// - `None` → default: return id + label only, no property reads.
/// - `Some([])` → fetch and return all properties (existing behavior).
/// - `Some(keys)` → fetch only named properties.
///
/// `cache` is a pre-populated cache of the schema registry.
pub(crate) fn materialize(
    gv: &GValue,
    ctx: &mut dyn crate::engine::GraphCtx,
    cache: &SchemaCache,
    prop_keys: Option<&[SmolStr]>,
) -> Result<Value, StoreError> {
    match gv {
        GValue::Scalar(ref p) => Ok(primitive_to_value(p.clone())),
        GValue::Vertex(vk) => materialize_vertex(*vk, ctx, cache, prop_keys),
        GValue::Edge(ek) => materialize_edge(*ek, ctx, cache, prop_keys),
        GValue::Property(ref p) => {
            let key = cache.prop_key_str(p.key).cloned().unwrap_or_else(|| SmolStr::from(format!("key_{}", p.key)));
            Ok(Value::Property(UserProperty { key, value: Box::new(primitive_to_value(p.value.clone())) }))
        }
        GValue::List(ref list) => {
            let mut out = Vec::with_capacity(list.len());
            for item in list {
                out.push(materialize(item, ctx, cache, prop_keys)?);
            }
            Ok(Value::List(out))
        }
        GValue::Map(ref map) => {
            let mut out = Map::new();
            for (k, v) in map {
                out.entries.push((materialize(k, ctx, cache, prop_keys)?, materialize(v, ctx, cache, prop_keys)?));
            }
            Ok(Value::Map(out))
        }
        GValue::Path(ref path) => {
            let mut objects = Vec::with_capacity(path.len());
            let mut labels: Vec<Vec<String>> = Vec::with_capacity(path.len());
            for (val, step_labels) in path {
                objects.push(materialize(val, ctx, cache, prop_keys)?);
                labels.push(match step_labels {
                    Some(ls) => ls.iter().map(|s| s.to_string()).collect(),
                    None => vec![],
                });
            }
            Ok(Value::Path(Path { objects, labels }))
        }
    }
}

/// Materialize a vertex, respecting the property fetch hint.
fn materialize_vertex(
    vk: VertexKey,
    ctx: &mut dyn crate::engine::GraphCtx,
    cache: &SchemaCache,
    prop_keys: Option<&[SmolStr]>,
) -> Result<Value, StoreError> {
    match prop_keys {
        // Default: id + label only — read the synthesized LABEL_KEY_ID, no blob decode.
        None => match ctx.get_value(&CanonicalKey::Vertex(vk), LABEL_KEY_ID)? {
            None => Err(StoreError::NotFound),
            Some(Primitive::Int32(label_id)) => Ok(Value::Vertex(UserVertex {
                id: vk,
                label: cache.vertex_label(label_id).clone(),
                properties: HashMap::new(),
            })),
            _ => Err(StoreError::CorruptData("vertex label_id not Int32")),
        },
        // All properties — full decode path.
        Some([]) => match ctx.get_all_props(&CanonicalKey::Vertex(vk))? {
            None => Err(StoreError::NotFound),
            Some((label_id, props)) => {
                let label = cache.vertex_label(label_id).clone();
                let mut properties: HashMap<SmolStr, Vec<Value>> = HashMap::new();
                for (key, prim) in props {
                    properties.entry(key).or_default().push(primitive_to_value(prim));
                }
                Ok(Value::Vertex(UserVertex { id: vk, label, properties }))
            }
        },
        // Named properties only — targeted per-key lookup (O(log P) each, no HashMap build).
        // Reserved keys (id=1, label=2, rank=3) are synthesized, not stored in the blob; skip them
        // to preserve parity with get_all_props which only returns user-defined properties.
        Some(keys) => match ctx.get_value(&CanonicalKey::Vertex(vk), LABEL_KEY_ID)? {
            None => Err(StoreError::NotFound),
            Some(Primitive::Int32(label_id)) => {
                let label = cache.vertex_label(label_id).clone();
                let mut properties: HashMap<SmolStr, Vec<Value>> = HashMap::new();
                for key in keys {
                    let Some(prop_key_id) = cache.prop_key_id(key) else { continue };
                    if matches!(prop_key_id, ID_KEY_ID | LABEL_KEY_ID | RANK_KEY_ID) {
                        continue;
                    }
                    if let Some(val) = ctx.get_value(&CanonicalKey::Vertex(vk), prop_key_id)? {
                        properties.entry(key.clone()).or_default().push(primitive_to_value(val));
                    }
                }
                Ok(Value::Vertex(UserVertex { id: vk, label, properties }))
            }
            _ => Err(StoreError::CorruptData("vertex label_id not Int32")),
        },
    }
}

/// Materialize an edge, respecting the property fetch hint.
fn materialize_edge(
    ek: crate::types::keys::EdgeKey,
    ctx: &mut dyn crate::engine::GraphCtx,
    cache: &SchemaCache,
    prop_keys: Option<&[SmolStr]>,
) -> Result<Value, StoreError> {
    let cek = ek.canonical_edge_key();
    match prop_keys {
        // Default: id + label only — label_id is in the key, zero store reads.
        None => Ok(Value::Edge(UserEdge {
            id: ek.to_id_string(),
            out_v: cek.src_id,
            in_v: cek.dst_id,
            label: cache.edge_label(ek.label_id).clone(),
            rank: cek.rank,
            properties: HashMap::new(),
        })),
        // All properties — full decode path.
        Some([]) => match ctx.get_all_props(&CanonicalKey::Edge(cek))? {
            None => Err(StoreError::NotFound),
            Some((label_id, props)) => {
                let label = cache.edge_label(label_id).clone();
                let mut properties: HashMap<SmolStr, Value> = HashMap::new();
                for (key, prim) in props {
                    properties.insert(key, primitive_to_value(prim));
                }
                Ok(Value::Edge(UserEdge {
                    id: ek.to_id_string(),
                    out_v: cek.src_id,
                    in_v: cek.dst_id,
                    label,
                    rank: cek.rank,
                    properties,
                }))
            }
        },
        // Named properties only — targeted per-key lookup.
        // Edge label comes from the key (zero store reads); properties are looked up individually.
        Some(keys) => {
            // Ensure the edge is in the overlay before calling get_value (overlay-only for edges).
            if ctx.get_edge(&ek)?.is_none() {
                return Err(StoreError::NotFound);
            }
            let label = cache.edge_label(ek.label_id).clone();
            let mut properties: HashMap<SmolStr, Value> = HashMap::new();
            for key in keys {
                let Some(prop_key_id) = cache.prop_key_id(key) else { continue };
                if matches!(prop_key_id, ID_KEY_ID | LABEL_KEY_ID | RANK_KEY_ID) {
                    continue;
                }
                if let Some(val) = ctx.get_value(&CanonicalKey::Edge(cek), prop_key_id)? {
                    properties.insert(key.clone(), primitive_to_value(val));
                }
            }
            Ok(Value::Edge(UserEdge {
                id: ek.to_id_string(),
                out_v: cek.src_id,
                in_v: cek.dst_id,
                label,
                rank: cek.rank,
                properties,
            }))
        }
    }
}

// ── BuiltTraversal ────────────────────────────────────────────────────────────

/// The result of building a traversal — a pull-based lazy iterator over results.
///
/// Obtained from [`ReadTraversal::iter`] or [`WriteTraversal::iter`].
/// Implements `Iterator<Item = Result<Value, StoreError>>`.
///
/// [`ReadTraversal::iter`]: super::ReadTraversal::iter
/// [`WriteTraversal::iter`]: super::WriteTraversal::iter
pub struct BuiltTraversal<'g> {
    pub(super) graph: &'g mut dyn crate::engine::GraphCtx,
    pub(super) plan: PhysicalPlan,
    /// Property fetch hint set by `withProperties()`.
    /// - `None` → default: id + label only, no properties.
    /// - `Some(vec![])` → all properties.
    /// - `Some(vec!["name"])` → named properties only.
    pub(super) prop_keys: Option<Vec<SmolStr>>,
    /// Pre-built label and property registry cache.
    pub(super) cache: SchemaCache,
}

impl<'g> Iterator for BuiltTraversal<'g> {
    type Item = Result<Value, StoreError>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.plan.next(self.graph) {
            Err(e) => Some(Err(e)),
            Ok(None) => None,
            Ok(Some(t)) => {
                // Scalar values (from count(), values(), is(), etc.) carry no
                // label-id or property-key-id that needs schema decoding — skip
                // any schema lookups.
                if let GValue::Scalar(ref p) = &t.value {
                    return Some(Ok(primitive_to_value(p.clone())));
                }
                Some(materialize(&t.value, self.graph, &self.cache, self.prop_keys.as_deref()))
            }
        }
    }
}