1 files changed, 115 insertions, 1 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>>,