WIP EnclosingLinesInfo and calculating previous context linessyntactic_context

Expand syntactic padding logic

3 files changed, 162 insertions, 29 deletions

diff --git a/src/display/context.rs b/src/display/context.rs
index 74148758e..eb94e271d 100644
--- a/src/display/context.rs
+++ b/src/display/context.rs
@@ -8,7 +8,7 @@ use rustc_hash::FxHashSet;
 
 use crate::{
     hash::DftHashMap,
-    parse::syntax::{zip_repeat_shorter, MatchKind, MatchedPos},
+    parse::syntax::{zip_repeat_shorter, EnclosingLinesInfo, MatchKind, MatchedPos},
 };
 
 pub(crate) fn all_matched_lines_filled(
@@ -417,6 +417,95 @@ fn before_with_opposites(
     res
 }
 
+fn pad_before_enclosing(
+    ln: LineNumber,
+    num_context_lines: usize,
+    is_lhs: bool,
+    enclosing_lines: &EnclosingLinesInfo,
+) -> Vec<LineNumber> {
+    let mut res = vec![];
+
+    let starts = if is_lhs {
+        &enclosing_lines.lhs_starts
+    } else {
+        &enclosing_lines.rhs_starts
+    };
+
+    let earliest = if ln.0 > num_context_lines as u32 {
+        ln.0 - num_context_lines as u32
+    } else {
+        0
+    };
+
+    // Previous enclosing block that's still within the context limit.
+    let mut prev_line_enclosing: LineNumber = if let Some(enclosing_starts) = starts.get(&ln) {
+        enclosing_starts
+            .iter()
+            .copied()
+            .find(|enc_ln| enc_ln.0 >= earliest)
+            .unwrap_or(ln.clone())
+    } else {
+        ln
+    };
+
+    // TODO: any previous lines (up to `earliest`) that have the same enclosing_starts.
+
+    // Any previous blank lines that are still within the context limit.
+    while prev_line_enclosing.0 > earliest && !starts.contains_key(&prev_line_enclosing) {
+        prev_line_enclosing = (prev_line_enclosing.0 - 1).into();
+    }
+
+    for n in prev_line_enclosing.0..ln.0 {
+        res.push(n.into());
+    }
+
+    res
+}
+
+fn pad_after_enclosing(
+    ln: LineNumber,
+    num_context_lines: usize,
+    max_line: LineNumber,
+    is_lhs: bool,
+    enclosing_lines: &EnclosingLinesInfo,
+) -> Vec<LineNumber> {
+    let mut res = vec![];
+
+    let ends = if is_lhs {
+        &enclosing_lines.lhs_ends
+    } else {
+        &enclosing_lines.rhs_ends
+    };
+
+    let latest = if ln.0 + num_context_lines as u32 > max_line.0 {
+        max_line.0
+    } else {
+        ln.0 + num_context_lines as u32
+    };
+
+    // Trailing enclosing block that's still within the context limit.
+    let mut next_line_enclosing: LineNumber = if let Some(enclosing_ends) = ends.get(&ln) {
+        enclosing_ends
+            .iter()
+            .copied()
+            .find(|enc_ln| enc_ln.0 <= latest)
+            .unwrap_or(ln.clone())
+    } else {
+        ln
+    };
+
+    // Any following blank lines that are still within the context limit.
+    while next_line_enclosing.0 < latest && !ends.contains_key(&next_line_enclosing) {
+        next_line_enclosing = (next_line_enclosing.0 - 1).into();
+    }
+
+    for n in ln.0 + 1..next_line_enclosing.0 {
+        res.push(n.into());
+    }
+
+    res
+}
+
 fn pad_before(ln: LineNumber, num_context_lines: usize) -> Vec<LineNumber> {
     let mut line_nums = vec![];
 
@@ -515,6 +604,31 @@ fn after_with_opposites(
     res
 }
 
+pub(crate) fn calculate_before_context_enclosing(
+    lines: &[(Option<LineNumber>, Option<LineNumber>)],
+    opposite_to_lhs: &DftHashMap<LineNumber, HashSet<LineNumber>>,
+    opposite_to_rhs: &DftHashMap<LineNumber, HashSet<LineNumber>>,
+    num_context_lines: usize,
+    enclosing_lines: &EnclosingLinesInfo,
+) -> Vec<(Option<LineNumber>, Option<LineNumber>)> {
+    match lines.first() {
+        Some(first_line) => match *first_line {
+            (Some(lhs_line), _) => {
+                let padded_lines =
+                    pad_before_enclosing(lhs_line, num_context_lines, true, enclosing_lines);
+                before_with_opposites(&padded_lines, opposite_to_lhs)
+            }
+            (_, Some(rhs_line)) => {
+                let padded_lines =
+                    pad_before_enclosing(rhs_line, num_context_lines, false, enclosing_lines);
+                flip_tuples(&before_with_opposites(&padded_lines, opposite_to_rhs))
+            }
+            (None, None) => vec![],
+        },
+        None => vec![],
+    }
+}
+
 pub(crate) fn calculate_before_context(
     lines: &[(Option<LineNumber>, Option<LineNumber>)],
     opposite_to_lhs: &DftHashMap<LineNumber, HashSet<LineNumber>>,
diff --git a/src/main.rs b/src/main.rs
index 16bb3231e..9c5d0c6bd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -76,7 +76,7 @@ use std::{env, thread};
 
 use humansize::{format_size, BINARY};
 use owo_colors::OwoColorize;
-use parse::syntax::enclosing_start_lines;
+use parse::syntax::EnclosingLinesInfo;
 use rayon::prelude::*;
 use strum::IntoEnumIterator;
 use typed_arena::Arena;
@@ -501,6 +501,8 @@ fn diff_file_content(
         };
     }
 
+    let mut enclosing_info = None;
+
     let (file_format, lhs_positions, rhs_positions) = match lang_config {
         None => {
             let file_format = FileFormat::PlainText;
@@ -586,9 +588,8 @@ fn diff_file_content(
                                 let mut lhs_positions = syntax::change_positions(&lhs, &change_map);
                                 let mut rhs_positions = syntax::change_positions(&rhs, &change_map);
 
-                                let enclosing_starts = enclosing_start_lines(&lhs);
-                                let enclosing_ends = enclosing_start_lines(&rhs);
-                                dbg!(enclosing_starts);
+                                enclosing_info = Some(EnclosingLinesInfo::new(&lhs, &rhs));
+                                dbg!(enclosing_info);
 
                                 if diff_options.ignore_comments {
                                     let lhs_comments =
diff --git a/src/parse/syntax.rs b/src/parse/syntax.rs
index 56d433346..1875fac0e 100644
--- a/src/parse/syntax.rs
+++ b/src/parse/syntax.rs
@@ -4,6 +4,7 @@
 
 use std::{cell::Cell, env, fmt, hash::Hash, num::NonZeroU32};
 
+use line_numbers::LineNumber;
 use line_numbers::LinePositions;
 use line_numbers::SingleLineSpan;
 use typed_arena::Arena;
@@ -333,17 +334,41 @@ impl<'a> Syntax<'a> {
     }
 }
 
-pub fn enclosing_start_lines<'a>(nodes: &[&'a Syntax<'a>]) -> HashMap<LineNumber, Vec<LineNumber>> {
-    let mut res = HashMap::new();
-    walk_enclosing_start_lines(nodes, &[], &mut res);
+#[derive(Debug)]
+pub(crate) struct EnclosingLinesInfo {
+    pub(crate) lhs_starts: DftHashMap<LineNumber, Vec<LineNumber>>,
+    pub(crate) rhs_starts: DftHashMap<LineNumber, Vec<LineNumber>>,
+    pub(crate) lhs_ends: DftHashMap<LineNumber, Vec<LineNumber>>,
+    pub(crate) rhs_ends: DftHashMap<LineNumber, Vec<LineNumber>>,
+}
 
-    res
+impl EnclosingLinesInfo {
+    pub(crate) fn new<'a>(lhs_nodes: &[&'a Syntax<'a>], rhs_nodes: &[&'a Syntax<'a>]) -> Self {
+        let mut lhs_starts = DftHashMap::default();
+        walk_enclosing_start_lines(lhs_nodes, &[], &mut lhs_starts);
+
+        let mut rhs_starts = DftHashMap::default();
+        walk_enclosing_start_lines(rhs_nodes, &[], &mut rhs_starts);
+
+        let mut lhs_ends = DftHashMap::default();
+        walk_enclosing_end_lines(lhs_nodes, &[], &mut lhs_ends);
+
+        let mut rhs_ends = DftHashMap::default();
+        walk_enclosing_end_lines(rhs_nodes, &[], &mut rhs_ends);
+
+        EnclosingLinesInfo {
+            lhs_starts,
+            rhs_starts,
+            lhs_ends,
+            rhs_ends,
+        }
+    }
 }
 
 fn walk_enclosing_start_lines<'a>(
     nodes: &[&'a Syntax<'a>],
-    parent_enclosing: &[LineNumber],
-    blocks_including: &mut HashMap<LineNumber, Vec<LineNumber>>,
+    enclosing: &[LineNumber],
+    blocks_including: &mut DftHashMap<LineNumber, Vec<LineNumber>>,
 ) {
     for node in nodes {
         match node {
@@ -353,7 +378,11 @@ fn walk_enclosing_start_lines<'a>(
                 ..
             } => {
                 if let Some(position) = open_position.first() {
-                    let mut enclosing = parent_enclosing.to_vec();
+                    let mut enclosing = enclosing.to_vec();
+
+                    if !blocks_including.contains_key(&position.line) {
+                        blocks_including.insert(position.line, enclosing.clone());
+                    }
 
                     if let Some(last) = enclosing.last() {
                         if *last != position.line {
@@ -363,17 +392,13 @@ fn walk_enclosing_start_lines<'a>(
                         enclosing.push(position.line);
                     }
 
-                    if !blocks_including.contains_key(&position.line) {
-                        blocks_including.insert(position.line, enclosing.clone());
-                    }
-
                     walk_enclosing_start_lines(children, &enclosing, blocks_including)
                 }
             }
             Atom { position, .. } => {
                 for line_span in position {
                     if !blocks_including.contains_key(&line_span.line) {
-                        blocks_including.insert(line_span.line, parent_enclosing.into());
+                        blocks_including.insert(line_span.line, enclosing.into());
                     }
                 }
             }
@@ -381,17 +406,10 @@ fn walk_enclosing_start_lines<'a>(
     }
 }
 
-pub fn enclosing_end_lines<'a>(nodes: &[&'a Syntax<'a>]) -> HashMap<LineNumber, Vec<LineNumber>> {
-    let mut res = HashMap::new();
-    walk_enclosing_end_lines(nodes, &[], &mut res);
-
-    res
-}
-
 fn walk_enclosing_end_lines<'a>(
     nodes: &[&'a Syntax<'a>],
     parent_enclosing: &[LineNumber],
-    blocks_including: &mut HashMap<LineNumber, Vec<LineNumber>>,
+    blocks_including: &mut DftHashMap<LineNumber, Vec<LineNumber>>,
 ) {
     for node in nodes {
         match node {
@@ -403,6 +421,10 @@ fn walk_enclosing_end_lines<'a>(
                 if let Some(position) = close_position.last() {
                     let mut enclosing = parent_enclosing.to_vec();
 
+                    if !blocks_including.contains_key(&position.line) {
+                        blocks_including.insert(position.line, enclosing.clone());
+                    }
+
                     if let Some(last) = enclosing.last() {
                         if *last != position.line {
                             enclosing.push(position.line);
@@ -411,10 +433,6 @@ fn walk_enclosing_end_lines<'a>(
                         enclosing.push(position.line);
                     }
 
-                    if !blocks_including.contains_key(&position.line) {
-                        blocks_including.insert(position.line, enclosing.clone());
-                    }
-
                     walk_enclosing_end_lines(children, &enclosing, blocks_including)
                 }
             }