linkrs-storage 0.1.0

use crate::engine::graph_storage::context::GraphStorageContext;
use crate::engine::graph_storage::ops::edge_record_to_edge;
use linkrs_core::{Edge, StorageResult, Value};

use crate::engine::graph_storage::reader::utils::*;

/// Enable the per-table edge property index for `edge_type`.
pub(crate) fn enable_edge_property_index(
    ctx: &GraphStorageContext,
    space: &str,
    edge_type: &str,
    pool_capacity: u64,
) -> StorageResult<bool> {
    record_schema_read(ctx, space);
    let (src_label, dst_label, edge_label) = resolve_edge_table_labels(ctx, space, edge_type)?;
    if src_label != 0 && dst_label != 0 {
        ctx.enable_edge_property_index(src_label, dst_label, edge_label, pool_capacity)?;
    } else {
        // Unconstrained endpoint tags: enable on every table of this edge type.
        ctx.data_store()
            .with_edge_tables(|tables| -> StorageResult<()> {
                let matching: Vec<_> = tables
                    .values()
                    .filter(|arc| arc.read().label() == edge_label)
                    .cloned()
                    .collect();
                for arc in matching {
                    arc.write().enable_property_index(pool_capacity)?;
                }
                Ok(())
            })?;
    }
    Ok(true)
}

/// Whether the per-table edge property index is enabled for `edge_type`.
pub(crate) fn has_edge_property_index(
    ctx: &GraphStorageContext,
    space: &str,
    edge_type: &str,
) -> StorageResult<bool> {
    record_schema_read(ctx, space);
    let (src_label, dst_label, edge_label) = resolve_edge_table_labels(ctx, space, edge_type)?;
    if src_label != 0 && dst_label != 0 {
        Ok(ctx.has_edge_property_index(src_label, dst_label, edge_label))
    } else {
        Ok(ctx.data_store().with_edge_tables(|tables| {
            tables
                .values()
                .filter(|arc| arc.read().label() == edge_label)
                .any(|arc| arc.read().has_property_index())
        }))
    }
}

/// Drop the per-table edge property index for `edge_type`.
pub(crate) fn disable_edge_property_index(
    ctx: &GraphStorageContext,
    space: &str,
    edge_type: &str,
) -> StorageResult<()> {
    record_schema_read(ctx, space);
    let (src_label, dst_label, edge_label) = resolve_edge_table_labels(ctx, space, edge_type)?;
    if src_label != 0 && dst_label != 0 {
        ctx.disable_edge_property_index(src_label, dst_label, edge_label)?;
    } else {
        ctx.data_store()
            .with_edge_tables(|tables| -> StorageResult<()> {
                for arc in tables
                    .values()
                    .filter(|arc| arc.read().label() == edge_label)
                {
                    arc.write().disable_property_index();
                }
                Ok(())
            })?;
    }
    Ok(())
}

/// Look up edges of `edge_type` whose `prop_name` value falls in `[lower, upper)`.
///
/// Bounds are encoded with the ordered codec; the inclusion flags control
/// whether the boundary values themselves are part of the range.
#[allow(clippy::too_many_arguments)]
pub(crate) fn lookup_edges_by_property_range(
    ctx: &GraphStorageContext,
    space: &str,
    edge_type: &str,
    prop_name: &str,
    lower: Option<&Value>,
    upper: Option<&Value>,
    include_lower: bool,
    include_upper: bool,
) -> StorageResult<Vec<Edge>> {
    record_schema_read(ctx, space);
    let (src_label, dst_label, edge_label) = resolve_edge_table_labels(ctx, space, edge_type)?;
    let codec = linkrs_core::value::ordered_codec::OrderedCodec::new();
    // Degenerate range [v, v) with an exclusive upper bound is interpreted as
    // a prefix/equality bound: everything from v up to the next value boundary.
    let prefix_bounds = include_lower && !include_upper && lower.is_some() && upper == lower;
    let value_lower = match lower {
        Some(value) => {
            let encoded = codec.encode(value)?;
            if include_lower {
                encoded
            } else {
                linkrs_core::value::ordered_codec::OrderedCodec::prefix_upper_bound(&encoded)
            }
        }
        None => Vec::new(),
    };
    let value_upper = match upper {
        Some(value) => {
            let encoded = codec.encode(value)?;
            if prefix_bounds || include_upper {
                linkrs_core::value::ordered_codec::OrderedCodec::prefix_upper_bound(&encoded)
            } else {
                encoded
            }
        }
        None => Vec::new(),
    };

    let ts = ctx.get_read_timestamp();
    let mut edges = Vec::new();

    let records: Vec<(linkrs_core::types::LabelId, linkrs_core::types::LabelId, _)> =
        if src_label != 0 && dst_label != 0 {
            ctx.lookup_edges_by_property_range(
                src_label,
                dst_label,
                edge_label,
                prop_name,
                &value_lower,
                &value_upper,
                ts,
            )
            .into_iter()
            .map(|record| (src_label, dst_label, record))
            .collect()
        } else {
            ctx.data_store().with_edge_tables(|tables| {
                let matching: Vec<_> = tables
                    .values()
                    .filter(|arc| arc.read().label() == edge_label)
                    .cloned()
                    .collect();
                let mut records = Vec::new();
                for arc in matching {
                    let table = arc.read();
                    let table_src_label = table.src_label();
                    let table_dst_label = table.dst_label();
                    records.extend(
                        table
                            .lookup_edges_by_property_range(prop_name, &value_lower, &value_upper)
                            .into_iter()
                            .filter_map(|(src, dst, rank)| {
                                table
                                    .get_edge(src, dst, rank, ts)
                                    .map(|record| (table_src_label, table_dst_label, record))
                            }),
                    );
                }
                records
            })
        };

    for (record_src_label, record_dst_label, record) in records {
        let src_internal = record.src_vid.as_internal_u32();
        let dst_internal = record.dst_vid.as_internal_u32();
        let src_external = match src_internal {
            Some(internal) if record_src_label != 0 => {
                internal_to_external_vertex_id(ctx, record_src_label, internal, ts)
                    .unwrap_or(record.src_vid)
            }
            _ => record.src_vid,
        };
        let dst_external = match dst_internal {
            Some(internal) if record_dst_label != 0 => {
                internal_to_external_vertex_id(ctx, record_dst_label, internal, ts)
                    .unwrap_or(record.dst_vid)
            }
            _ => record.dst_vid,
        };
        edges.push(edge_record_to_edge(
            record,
            edge_type,
            src_external,
            dst_external,
        ));
    }

    Ok(edges)
}