lean-ctx 3.1.3

//! Node types and CRUD operations for graph nodes.

use rusqlite::{params, Connection, OptionalExtension};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NodeKind {
    File,
    Symbol,
    Module,
}

impl NodeKind {
    pub fn as_str(&self) -> &'static str {
        match self {
            NodeKind::File => "file",
            NodeKind::Symbol => "symbol",
            NodeKind::Module => "module",
        }
    }

    pub fn parse(s: &str) -> Self {
        match s {
            "symbol" => NodeKind::Symbol,
            "module" => NodeKind::Module,
            _ => NodeKind::File,
        }
    }
}

#[derive(Debug, Clone)]
pub struct Node {
    pub id: Option<i64>,
    pub kind: NodeKind,
    pub name: String,
    pub file_path: String,
    pub line_start: Option<usize>,
    pub line_end: Option<usize>,
    pub metadata: Option<String>,
}

impl Node {
    pub fn file(path: &str) -> Self {
        Self {
            id: None,
            kind: NodeKind::File,
            name: path.to_string(),
            file_path: path.to_string(),
            line_start: None,
            line_end: None,
            metadata: None,
        }
    }

    pub fn symbol(name: &str, file_path: &str, kind: NodeKind) -> Self {
        Self {
            id: None,
            kind,
            name: name.to_string(),
            file_path: file_path.to_string(),
            line_start: None,
            line_end: None,
            metadata: None,
        }
    }

    pub fn with_lines(mut self, start: usize, end: usize) -> Self {
        self.line_start = Some(start);
        self.line_end = Some(end);
        self
    }

    pub fn with_metadata(mut self, meta: &str) -> Self {
        self.metadata = Some(meta.to_string());
        self
    }
}

pub fn upsert(conn: &Connection, node: &Node) -> anyhow::Result<i64> {
    conn.execute(
        "INSERT INTO nodes (kind, name, file_path, line_start, line_end, metadata)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6)
         ON CONFLICT(kind, name, file_path) DO UPDATE SET
            line_start = excluded.line_start,
            line_end = excluded.line_end,
            metadata = excluded.metadata",
        params![
            node.kind.as_str(),
            node.name,
            node.file_path,
            node.line_start.map(|v| v as i64),
            node.line_end.map(|v| v as i64),
            node.metadata,
        ],
    )?;

    let id: i64 = conn.query_row(
        "SELECT id FROM nodes WHERE kind = ?1 AND name = ?2 AND file_path = ?3",
        params![node.kind.as_str(), node.name, node.file_path],
        |row| row.get(0),
    )?;

    Ok(id)
}

pub fn get_by_path(conn: &Connection, file_path: &str) -> anyhow::Result<Option<Node>> {
    let result = conn
        .query_row(
            "SELECT id, kind, name, file_path, line_start, line_end, metadata
             FROM nodes WHERE kind = 'file' AND file_path = ?1",
            params![file_path],
            |row| {
                Ok(Node {
                    id: Some(row.get(0)?),
                    kind: NodeKind::parse(&row.get::<_, String>(1)?),
                    name: row.get(2)?,
                    file_path: row.get(3)?,
                    line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
                    line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
                    metadata: row.get(6)?,
                })
            },
        )
        .optional()?;
    Ok(result)
}

pub fn get_by_symbol(
    conn: &Connection,
    name: &str,
    file_path: &str,
) -> anyhow::Result<Option<Node>> {
    let result = conn
        .query_row(
            "SELECT id, kind, name, file_path, line_start, line_end, metadata
             FROM nodes WHERE name = ?1 AND file_path = ?2 AND kind != 'file'",
            params![name, file_path],
            |row| {
                Ok(Node {
                    id: Some(row.get(0)?),
                    kind: NodeKind::parse(&row.get::<_, String>(1)?),
                    name: row.get(2)?,
                    file_path: row.get(3)?,
                    line_start: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
                    line_end: row.get::<_, Option<i64>>(5)?.map(|v| v as usize),
                    metadata: row.get(6)?,
                })
            },
        )
        .optional()?;
    Ok(result)
}

pub fn remove_by_file(conn: &Connection, file_path: &str) -> anyhow::Result<()> {
    conn.execute(
        "DELETE FROM edges WHERE source_id IN (SELECT id FROM nodes WHERE file_path = ?1)
         OR target_id IN (SELECT id FROM nodes WHERE file_path = ?1)",
        params![file_path],
    )?;
    conn.execute("DELETE FROM nodes WHERE file_path = ?1", params![file_path])?;
    Ok(())
}

pub fn count(conn: &Connection) -> anyhow::Result<usize> {
    let c: i64 = conn.query_row("SELECT COUNT(*) FROM nodes", [], |row| row.get(0))?;
    Ok(c as usize)
}