lean-tui 1.1.0

Standalone TUI infoview for Lean 4 theorem prover
use crossterm::event::KeyCode;
use tracing::trace;

use crate::{
    graph_layout::{Dag, DagLayout, NavNode, NodePosition},
    selection::{Region, Selection},
};

/// Direction for spatial navigation.
#[derive(Debug, Clone, Copy)]
pub enum Direction {
    Up,
    Down,
    Left,
    Right,
}

/// Convert vim-style hjkl and arrow keys to a Direction.
pub fn key_to_direction(code: KeyCode) -> Option<Direction> {
    match code {
        KeyCode::Char('j') | KeyCode::Down => Some(Direction::Down),
        KeyCode::Char('k') | KeyCode::Up => Some(Direction::Up),
        KeyCode::Char('h') | KeyCode::Left => Some(Direction::Left),
        KeyCode::Char('l') | KeyCode::Right => Some(Direction::Right),
        _ => None,
    }
}

impl Direction {
    /// Check if `other` is ahead of `cur` in this direction.
    fn is_ahead(self, cur: &Region, other: &Region) -> bool {
        match self {
            Self::Left => other.x + i32::from(other.width) <= cur.x,
            Self::Right => other.x >= cur.x + i32::from(cur.width),
            Self::Up => other.y + i32::from(other.height) <= cur.y,
            Self::Down => other.y >= cur.y + i32::from(cur.height),
        }
    }

    /// Calculate distance from `cur` to `other` in this direction.
    fn distance(self, cur: &Region, other: &Region) -> i32 {
        match self {
            Self::Left => (cur.x - (other.x + i32::from(other.width))).max(0),
            Self::Right => (other.x - (cur.x + i32::from(cur.width))).max(0),
            Self::Up => (cur.y - (other.y + i32::from(other.height))).max(0),
            Self::Down => (other.y - (cur.y + i32::from(cur.height))).max(0),
        }
    }

    /// Whether this is a horizontal direction.
    const fn is_horizontal(self) -> bool {
        matches!(self, Self::Left | Self::Right)
    }
}

/// Check if two 1D ranges overlap.
pub fn ranges_overlap(a_start: i32, a_len: u16, b_start: i32, b_len: u16) -> bool {
    let a_end = a_start + i32::from(a_len);
    let b_end = b_start + i32::from(b_len);
    a_start < b_end && b_start < a_end
}

/// Build navigation regions for all nodes in a layout.
pub fn build_navigation_regions<D>(dag: &D, layout: &DagLayout) -> Vec<Region>
where
    D: Dag,
    D::Node: NavNode,
{
    trace!(
        node_count = layout.nodes.len(),
        "Building navigation regions"
    );

    let regions: Vec<_> = layout
        .nodes
        .iter()
        .filter_map(|pos| dag.get_node(pos.node_id).map(|node| (pos, node)))
        .flat_map(|(pos, node)| {
            trace!(
                node_id = pos.node_id,
                x = pos.x,
                y = pos.y,
                "Processing node"
            );

            let hyp_count = node.hyp_count();

            let hyp_regions = build_hyp_regions(pos, hyp_count);
            let goal_regions = build_goal_regions(pos, node, hyp_count);

            hyp_regions.into_iter().chain(goal_regions)
        })
        .collect();

    trace!(total_regions = regions.len(), "Navigation regions built");
    regions
}

/// Build navigation regions for hypotheses/bindings in a node (stacked
/// vertically).
fn build_hyp_regions(pos: &NodePosition, hyp_count: usize) -> Vec<Region> {
    (0..hyp_count)
        .map(|hyp_idx| Region {
            x: pos.x + 1,
            y: pos.y + 1 + hyp_idx as i32,
            width: pos.width.saturating_sub(2),
            height: 1,
            selection: Selection::Hyp {
                node_id: pos.node_id,
                hyp_idx,
            },
        })
        .collect()
}

/// Build navigation regions for goals in a node (below the bindings).
fn build_goal_regions<N: NavNode>(pos: &NodePosition, node: &N, hyp_count: usize) -> Vec<Region> {
    if !node.has_selectable_goals() {
        return Vec::new();
    }

    let goal_count = node.goal_count();
    if goal_count == 0 {
        return Vec::new();
    }

    let goal_y = pos.y + 1 + hyp_count as i32;
    let goal_width = pos.width.saturating_sub(2);

    (0..goal_count)
        .map(|i| Region {
            x: pos.x + 1,
            y: goal_y + i as i32,
            width: goal_width,
            height: 1,
            selection: Selection::Goal {
                node_id: pos.node_id,
                goal_idx: i,
            },
        })
        .collect()
}

pub fn find_nearest_in_direction(
    regions: &[Region],
    current: Selection,
    direction: Direction,
) -> Option<Selection> {
    let cur = regions.iter().find(|r| r.selection == current)?;

    let is_aligned = |r: &Region| {
        let overlaps = if direction.is_horizontal() {
            ranges_overlap(cur.y, cur.height, r.y, r.height)
        } else {
            ranges_overlap(cur.x, cur.width, r.x, r.width)
        };
        overlaps && direction.is_ahead(cur, r)
    };

    regions
        .iter()
        .filter(|r| is_aligned(r))
        .min_by_key(|r| direction.distance(cur, r))
        .map(|r| r.selection)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn hyp_region(node_id: u64, hyp_idx: usize, x: i32, y: i32) -> Region {
        Region {
            x,
            y,
            width: 15,
            height: 1,
            selection: Selection::Hyp { node_id, hyp_idx },
        }
    }

    fn goal_region(node_id: u64, goal_idx: usize, x: i32, y: i32, width: u16) -> Region {
        Region {
            x,
            y,
            width,
            height: 1,
            selection: Selection::Goal { node_id, goal_idx },
        }
    }

    #[test]
    fn test_sibling_hypothesis_navigation() {
        let regions = vec![
            hyp_region(1, 0, 0, 10),
            hyp_region(1, 1, 16, 10),
            goal_region(1, 0, 1, 12, 30),
        ];

        assert_eq!(
            find_nearest_in_direction(
                &regions,
                Selection::Hyp {
                    node_id: 1,
                    hyp_idx: 0
                },
                Direction::Right
            ),
            Some(Selection::Hyp {
                node_id: 1,
                hyp_idx: 1
            })
        );

        assert_eq!(
            find_nearest_in_direction(
                &regions,
                Selection::Hyp {
                    node_id: 1,
                    hyp_idx: 1
                },
                Direction::Left
            ),
            Some(Selection::Hyp {
                node_id: 1,
                hyp_idx: 0
            })
        );

        assert_eq!(
            find_nearest_in_direction(
                &regions,
                Selection::Hyp {
                    node_id: 1,
                    hyp_idx: 0
                },
                Direction::Down
            ),
            Some(Selection::Goal {
                node_id: 1,
                goal_idx: 0
            })
        );
    }

    #[test]
    fn test_navigation_no_target() {
        let regions = vec![hyp_region(1, 0, 0, 10)];

        assert_eq!(
            find_nearest_in_direction(
                &regions,
                Selection::Hyp {
                    node_id: 1,
                    hyp_idx: 0
                },
                Direction::Right
            ),
            None
        );
    }

    #[test]
    fn test_ranges_overlap() {
        assert!(ranges_overlap(10, 1, 10, 1));
        assert!(!ranges_overlap(10, 1, 11, 1));
        assert!(ranges_overlap(10, 5, 12, 5));
        assert!(ranges_overlap(10, 10, 12, 2));
    }
}