Vectrum
Sign In
345 lines · 13.9 KB · Rust
//! Empryo's Genome: the ranked map of a repository Empryo keeps in
//! `.empryo/genome.db`, a SQLite file it updates as code changes. The Genome
//! map panel reads its most important files, what they export and which of
//! them import which, and lays them out as a graph.

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    time::{Duration, SystemTime},
};

use anyhow::{Context as _, Result};
use rusqlite::{Connection, OpenFlags, params_from_iter};

/// How many files the map shows, by rank.
pub const MAP_FILES: usize = 48;
/// How many symbols each file lists.
const FILE_SYMBOLS: usize = 12;
/// A file changed this recently is marked new, as Empryo does in its map.
const RECENT: Duration = Duration::from_secs(48 * 60 * 60);

pub struct Genome {
    pub total_files: u64,
    pub total_symbols: u64,
    pub total_edges: u64,
    /// The highest ranked files, highest first.
    pub files: Vec<GenomeFile>,
    /// Imports between those files: `from` depends on `to`.
    pub edges: Vec<GenomeEdge>,
    /// Where each file sits in the map, from 0 to 1 on both axes.
    pub layout: Vec<(f32, f32)>,
}

pub struct GenomeFile {
    pub path: String,
    pub language: String,
    pub lines: u64,
    pub symbol_count: u64,
    /// Its PageRank in the import graph; the files' ranks add up to 1.
    pub rank: f64,
    /// How many files import it, from anywhere in the repository.
    pub dependents: u64,
    /// How many files it imports.
    pub dependencies: u64,
    /// Changed in the last 48 hours.
    pub recent: bool,
    /// Its symbols, exported ones first.
    pub symbols: Vec<GenomeSymbol>,
}

pub struct GenomeSymbol {
    pub name: String,
    pub kind: String,
    pub line: u64,
    pub exported: bool,
    pub signature: Option<String>,
}

pub struct GenomeEdge {
    pub from: usize,
    pub to: usize,
    pub weight: f64,
}

/// The Genome database of the checkout `cwd` is in: Empryo keeps it at the
/// checkout's root, which may be a folder or two up.
pub fn db_path(cwd: &Path) -> Option<PathBuf> {
    cwd.ancestors().map(|dir| dir.join(".empryo").join("genome.db")).find(|path| path.is_file())
}

/// Reads the map of the checkout `cwd` is in.
pub fn load(cwd: &Path) -> Result<Option<Genome>> {
    let Some(path) = db_path(cwd) else { return Ok(None) };
    // Opened for writing so SQLite can share Empryo's write-ahead log, but
    // only ever queried.
    let db = Connection::open_with_flags(&path, OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX)
        .with_context(|| format!("couldn't open {}", path.display()))?;
    db.busy_timeout(Duration::from_secs(5))?;
    db.pragma_update(None, "query_only", true)?;

    let count = |table: &str| -> Result<u64> {
        Ok(db.query_row(&format!("SELECT count(*) FROM {table}"), [], |row| row.get::<_, i64>(0))? as u64)
    };
    let (total_files, total_symbols, total_edges) = (count("files")?, count("symbols")?, count("edges")?);

    let now_ms = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap_or_default().as_millis() as f64;
    let mut ids = Vec::new();
    let mut files = Vec::new();
    let mut statement = db.prepare(
        "SELECT id, path, language, line_count, symbol_count, pagerank, mtime_ms FROM files
         ORDER BY pagerank DESC, symbol_count DESC LIMIT ?1",
    )?;
    let rows = statement.query_map([MAP_FILES as i64], |row| {
        Ok((
            row.get::<_, i64>(0)?,
            GenomeFile {
                path: row.get(1)?,
                language: row.get(2)?,
                lines: row.get::<_, i64>(3)?.max(0) as u64,
                symbol_count: row.get::<_, i64>(4)?.max(0) as u64,
                rank: row.get(5)?,
                dependents: 0,
                dependencies: 0,
                recent: now_ms - row.get::<_, f64>(6)? < RECENT.as_millis() as f64,
                symbols: Vec::new(),
            },
        ))
    })?;
    for row in rows {
        let (id, file) = row?;
        ids.push(id);
        files.push(file);
    }
    drop(statement);
    let index: HashMap<i64, usize> = ids.iter().enumerate().map(|(ix, id)| (*id, ix)).collect();
    let placeholders = vec!["?"; ids.len()].join(",");

    let mut edges = Vec::new();
    if !ids.is_empty() {
        let mut statement = db.prepare(&format!(
            "SELECT target_file_id, count(*) FROM edges WHERE target_file_id IN ({placeholders}) GROUP BY target_file_id"
        ))?;
        let rows = statement.query_map(params_from_iter(&ids), |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)))?;
        for row in rows {
            let (id, count) = row?;
            if let Some(&ix) = index.get(&id) {
                files[ix].dependents = count as u64;
            }
        }
        drop(statement);
        let mut statement = db.prepare(&format!(
            "SELECT source_file_id, count(*) FROM edges WHERE source_file_id IN ({placeholders}) GROUP BY source_file_id"
        ))?;
        let rows = statement.query_map(params_from_iter(&ids), |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)))?;
        for row in rows {
            let (id, count) = row?;
            if let Some(&ix) = index.get(&id) {
                files[ix].dependencies = count as u64;
            }
        }
        drop(statement);

        let mut statement = db.prepare(&format!(
            "SELECT source_file_id, target_file_id, weight FROM edges
             WHERE source_file_id IN ({placeholders}) AND target_file_id IN ({placeholders})"
        ))?;
        let both: Vec<i64> = ids.iter().chain(ids.iter()).copied().collect();
        let rows = statement.query_map(params_from_iter(&both), |row| {
            Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?, row.get::<_, f64>(2)?))
        })?;
        for row in rows {
            let (from, to, weight) = row?;
            if let (Some(&from), Some(&to)) = (index.get(&from), index.get(&to))
                && from != to
            {
                edges.push(GenomeEdge { from, to, weight });
            }
        }
        drop(statement);

        let mut statement = db.prepare(&format!(
            "SELECT file_id, name, kind, line, is_exported, signature FROM symbols
             WHERE file_id IN ({placeholders}) ORDER BY is_exported DESC, line"
        ))?;
        let rows = statement.query_map(params_from_iter(&ids), |row| {
            Ok((
                row.get::<_, i64>(0)?,
                GenomeSymbol {
                    name: row.get(1)?,
                    kind: row.get(2)?,
                    line: row.get::<_, i64>(3)?.max(0) as u64,
                    exported: row.get::<_, i64>(4)? != 0,
                    signature: row.get(5)?,
                },
            ))
        })?;
        for row in rows {
            let (id, symbol) = row?;
            if let Some(&ix) = index.get(&id)
                && files[ix].symbols.len() < FILE_SYMBOLS
            {
                files[ix].symbols.push(symbol);
            }
        }
    }

    let layout = layout(&files, &edges);
    Ok(Some(Genome { total_files, total_symbols, total_edges, files, edges, layout }))
}

/// Places the files with a force-directed layout: imports pull files
/// together, every pair pushes apart, and higher ranked files start nearer
/// the middle. Deterministic, so the map holds still between reloads.
pub fn layout(files: &[GenomeFile], edges: &[GenomeEdge]) -> Vec<(f32, f32)> {
    let n = files.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![(0.5, 0.5)];
    }
    // A spiral by rank: the most important file in the middle.
    let golden = std::f32::consts::PI * (3. - 5f32.sqrt());
    let mut pos: Vec<(f32, f32)> = (0..n)
        .map(|ix| {
            let radius = ((ix as f32 + 0.5) / n as f32).sqrt() * 0.45;
            let angle = ix as f32 * golden;
            (radius * angle.cos(), radius * angle.sin())
        })
        .collect();
    let mut edges_of = vec![0usize; n];
    for edge in edges {
        edges_of[edge.from] += 1;
        edges_of[edge.to] += 1;
    }
    let k = (1. / n as f32).sqrt() * 1.1;
    let mut temperature = 0.1;
    for _ in 0..400 {
        let mut shift = vec![(0f32, 0f32); n];
        for a in 0..n {
            for b in a + 1..n {
                let (dx, dy) = (pos[a].0 - pos[b].0, pos[a].1 - pos[b].1);
                let distance = (dx * dx + dy * dy).sqrt().max(0.001);
                let force = k * k / distance;
                let (fx, fy) = (dx / distance * force, dy / distance * force);
                shift[a].0 += fx;
                shift[a].1 += fy;
                shift[b].0 -= fx;
                shift[b].1 -= fy;
            }
        }
        for edge in edges {
            let (a, b) = (edge.from, edge.to);
            let (dx, dy) = (pos[a].0 - pos[b].0, pos[a].1 - pos[b].1);
            let distance = (dx * dx + dy * dy).sqrt().max(0.001);
            // Hubs have many lines; each pulls less, so they don't swallow the map.
            let degree = (edges_of[a] + edges_of[b]) as f32;
            let force = distance * distance / k * (1. + (edge.weight as f32).ln_1p() * 0.2) / degree.sqrt().max(1.);
            let (fx, fy) = (dx / distance * force, dy / distance * force);
            shift[a].0 -= fx;
            shift[a].1 -= fy;
            shift[b].0 += fx;
            shift[b].1 += fy;
        }
        for (ix, (dx, dy)) in shift.into_iter().enumerate() {
            // A pull to the middle keeps files nothing imports from drifting off.
            let gravity = 2.5 * k;
            let (dx, dy) = (dx - pos[ix].0 * gravity, dy - pos[ix].1 * gravity);
            let length = (dx * dx + dy * dy).sqrt().max(0.0001);
            let step = length.min(temperature);
            pos[ix].0 += dx / length * step;
            pos[ix].1 += dy / length * step;
        }
        temperature *= 0.988;
    }
    // A few files far out would shrink everything else: bring the farthest
    // in to just outside the rest.
    let mut distances: Vec<f32> = pos.iter().map(|(x, y)| (x * x + y * y).sqrt()).collect();
    distances.sort_by(f32::total_cmp);
    let edge = distances[(n * 9 / 10).min(n - 1)] * 1.15;
    for (x, y) in &mut pos {
        let distance = (*x * *x + *y * *y).sqrt();
        if distance > edge {
            let squeeze = (edge + (distance - edge) * 0.2) / distance;
            *x *= squeeze;
            *y *= squeeze;
        }
    }
    // Fit into the unit square.
    let (min_x, max_x) = pos.iter().fold((f32::MAX, f32::MIN), |(lo, hi), (x, _)| (lo.min(*x), hi.max(*x)));
    let (min_y, max_y) = pos.iter().fold((f32::MAX, f32::MIN), |(lo, hi), (_, y)| (lo.min(*y), hi.max(*y)));
    let span = (max_x - min_x).max(max_y - min_y).max(0.0001);
    let (off_x, off_y) = ((1. - (max_x - min_x) / span) / 2., (1. - (max_y - min_y) / span) / 2.);
    pos.into_iter().map(|(x, y)| ((x - min_x) / span + off_x, (y - min_y) / span + off_y)).collect()
}

#[cfg(test)]
mod tests {
    use super::{GenomeEdge, GenomeFile, layout, load};

    fn file(path: &str, rank: f64) -> GenomeFile {
        GenomeFile {
            path: path.into(),
            language: "rust".into(),
            lines: 10,
            symbol_count: 1,
            rank,
            dependents: 0,
            dependencies: 0,
            recent: false,
            symbols: Vec::new(),
        }
    }

    #[test]
    fn layout_fits_the_square() {
        let files: Vec<GenomeFile> = (0..20).map(|ix| file(&format!("f{ix}.rs"), 1. / (ix as f64 + 1.))).collect();
        let edges: Vec<GenomeEdge> = (1..20).map(|ix| GenomeEdge { from: ix, to: 0, weight: 1. }).collect();
        let pos = layout(&files, &edges);
        assert_eq!(pos.len(), 20);
        assert!(pos.iter().all(|(x, y)| (0. ..=1.).contains(x) && (0. ..=1.).contains(y)));
        // No two files on top of each other.
        for a in 0..pos.len() {
            for b in a + 1..pos.len() {
                let distance = ((pos[a].0 - pos[b].0).powi(2) + (pos[a].1 - pos[b].1).powi(2)).sqrt();
                assert!(distance > 0.01, "{a} and {b} overlap");
            }
        }
        assert_eq!(layout(&files[..1], &[]), vec![(0.5, 0.5)]);
    }

    #[test]
    fn reads_a_genome() {
        let dir = std::env::temp_dir().join(format!("vectrum-genome-test-{}", std::process::id()));
        let sub = dir.join("src");
        std::fs::create_dir_all(dir.join(".empryo")).unwrap();
        std::fs::create_dir_all(&sub).unwrap();
        let db = rusqlite::Connection::open(dir.join(".empryo/genome.db")).unwrap();
        db.execute_batch(
            "PRAGMA journal_mode = WAL;
             CREATE TABLE files (id INTEGER PRIMARY KEY, path TEXT, mtime_ms REAL, language TEXT, line_count INTEGER,
               symbol_count INTEGER, pagerank REAL);
             CREATE TABLE symbols (id INTEGER PRIMARY KEY, file_id INTEGER, name TEXT, kind TEXT, line INTEGER,
               end_line INTEGER, is_exported INTEGER, signature TEXT);
             CREATE TABLE edges (source_file_id INTEGER, target_file_id INTEGER, weight REAL);
             INSERT INTO files VALUES (1, 'src/main.ts', 0, 'typescript', 4, 2, 0.35), (2, 'src/math.ts', 0, 'typescript', 3, 2, 0.65);
             INSERT INTO symbols VALUES (1, 1, 'main', 'function', 3, 3, 1, NULL), (2, 2, 'get', 'method', 2, 2, 0, NULL),
               (3, 2, 'add', 'function', 1, 1, 1, 'export function add(a, b)');
             INSERT INTO edges VALUES (1, 2, 4.5);",
        )
        .unwrap();
        let genome = load(&sub).unwrap().unwrap();
        assert_eq!((genome.total_files, genome.total_symbols, genome.total_edges), (2, 3, 1));
        assert_eq!(genome.files[0].path, "src/math.ts");
        assert_eq!(genome.files[0].dependents, 1);
        assert_eq!(genome.files[1].dependencies, 1);
        assert_eq!(genome.files[0].symbols[0].name, "add");
        assert!(genome.files[0].symbols[0].exported);
        assert_eq!(genome.edges.len(), 1);
        assert_eq!((genome.edges[0].from, genome.edges[0].to), (1, 0));
        drop(db);
        let _ = std::fs::remove_dir_all(&dir);
        assert!(load(&std::env::temp_dir().join("vectrum-genome-none")).unwrap().is_none());
    }
}