Avoid losing focus when block decorations go offscreen (#14815)

Release Notes:

- Fixed a bug that caused focus to be lost when renames and inline
assists were scrolled offscreen.

---------

Co-authored-by: Nathan <nathan@zed.dev>

crates/assistant/src/assistant_panel.rs

@@ -23,7 +23,8 @@ use collections::{BTreeSet, HashMap, HashSet};
 use editor::{
     actions::{FoldAt, MoveToEndOfLine, Newline, ShowCompletions, UnfoldAt},
     display_map::{
-        BlockDisposition, BlockId, BlockProperties, BlockStyle, Crease, RenderBlock, ToDisplayPoint,
+        BlockDisposition, BlockProperties, BlockStyle, Crease, CustomBlockId, RenderBlock,
+        ToDisplayPoint,
     },
     scroll::{Autoscroll, AutoscrollStrategy, ScrollAnchor},
     Anchor, Editor, EditorEvent, ExcerptRange, MultiBuffer, RowExt, ToOffset as _, ToPoint,
@@ -984,11 +985,11 @@ pub struct ContextEditor {
     project: Model<Project>,
     lsp_adapter_delegate: Option<Arc<dyn LspAdapterDelegate>>,
     editor: View<Editor>,
-    blocks: HashSet<BlockId>,
+    blocks: HashSet<CustomBlockId>,
     scroll_position: Option<ScrollPosition>,
     remote_id: Option<workspace::ViewId>,
     pending_slash_command_creases: HashMap<Range<language::Anchor>, CreaseId>,
-    pending_slash_command_blocks: HashMap<Range<language::Anchor>, BlockId>,
+    pending_slash_command_blocks: HashMap<Range<language::Anchor>, CustomBlockId>,
     _subscriptions: Vec<Subscription>,
     active_edit_step: Option<ActiveEditStep>,
     assistant_panel: WeakView<AssistantPanel>,

crates/assistant/src/inline_assistant.rs

@@ -9,7 +9,7 @@ use collections::{hash_map, HashMap, HashSet, VecDeque};
 use editor::{
     actions::{MoveDown, MoveUp, SelectAll},
     display_map::{
-        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
+        BlockContext, BlockDisposition, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
         ToDisplayPoint,
     },
     Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle,
@@ -310,7 +310,7 @@ impl InlineAssistant {
         range: &Range<Anchor>,
         prompt_editor: &View<PromptEditor>,
         cx: &mut WindowContext,
-    ) -> [BlockId; 2] {
+    ) -> [CustomBlockId; 2] {
         let assist_blocks = vec![
             BlockProperties {
                 style: BlockStyle::Sticky,
@@ -1900,8 +1900,8 @@ impl InlineAssist {
         include_context: bool,
         editor: &View<Editor>,
         prompt_editor: &View<PromptEditor>,
-        prompt_block_id: BlockId,
+        prompt_block_id: CustomBlockId,
-        end_block_id: BlockId,
+        end_block_id: CustomBlockId,
         codegen: Model<Codegen>,
         workspace: Option<WeakView<Workspace>>,
         cx: &mut WindowContext,
@@ -1995,10 +1995,10 @@ impl InlineAssist {
 }
 
 struct InlineAssistDecorations {
-    prompt_block_id: BlockId,
+    prompt_block_id: CustomBlockId,
     prompt_editor: View<PromptEditor>,
-    removed_line_block_ids: HashSet<BlockId>,
+    removed_line_block_ids: HashSet<CustomBlockId>,
-    end_block_id: BlockId,
+    end_block_id: CustomBlockId,
 }
 
 #[derive(Debug)]

crates/diagnostics/src/diagnostics.rs

@@ -10,7 +10,7 @@ use anyhow::Result;
 use collections::{BTreeSet, HashSet};
 use editor::{
     diagnostic_block_renderer,
-    display_map::{BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock},
+    display_map::{BlockDisposition, BlockProperties, BlockStyle, CustomBlockId, RenderBlock},
     highlight_diagnostic_message,
     scroll::Autoscroll,
     Editor, EditorEvent, ExcerptId, ExcerptRange, MultiBuffer, ToOffset,
@@ -85,7 +85,7 @@ struct DiagnosticGroupState {
     primary_diagnostic: DiagnosticEntry<language::Anchor>,
     primary_excerpt_ix: usize,
     excerpts: Vec<ExcerptId>,
-    blocks: HashSet<BlockId>,
+    blocks: HashSet<CustomBlockId>,
     block_count: usize,
 }
 

crates/diagnostics/src/diagnostics_tests.rs

@@ -1,7 +1,7 @@
 use super::*;
 use collections::HashMap;
 use editor::{
-    display_map::{BlockContext, DisplayRow, TransformBlock},
+    display_map::{Block, BlockContext, DisplayRow},
     DisplayPoint, GutterDimensions,
 };
 use gpui::{px, AvailableSpace, Stateful, TestAppContext, VisualTestContext};
@@ -974,9 +974,9 @@ fn editor_blocks(
                 snapshot
                     .blocks_in_range(DisplayRow(0)..snapshot.max_point().row())
                     .filter_map(|(row, block)| {
-                        let transform_block_id = block.id();
+                        let block_id = block.id();
                         let name: SharedString = match block {
-                            TransformBlock::Custom(block) => {
+                            Block::Custom(block) => {
                                 let mut element = block.render(&mut BlockContext {
                                     context: cx,
                                     anchor_x: px(0.),
@@ -984,7 +984,7 @@ fn editor_blocks(
                                     line_height: px(0.),
                                     em_width: px(0.),
                                     max_width: px(0.),
-                                    transform_block_id,
+                                    block_id,
                                     editor_style: &editor::EditorStyle::default(),
                                 });
                                 let element = element.downcast_mut::<Stateful<Div>>().unwrap();
@@ -996,7 +996,7 @@ fn editor_blocks(
                                     .ok()?
                             }
 
-                            TransformBlock::ExcerptHeader {
+                            Block::ExcerptHeader {
                                 starts_new_buffer, ..
                             } => {
                                 if *starts_new_buffer {
@@ -1005,7 +1005,7 @@ fn editor_blocks(
                                     EXCERPT_HEADER.into()
                                 }
                             }
-                            TransformBlock::ExcerptFooter { .. } => EXCERPT_FOOTER.into(),
+                            Block::ExcerptFooter { .. } => EXCERPT_FOOTER.into(),
                         };
 
                         Some((row, name))

crates/diagnostics/src/grouped_diagnostics.rs

@@ -3,8 +3,8 @@ use collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 use editor::{
     diagnostic_block_renderer,
     display_map::{
-        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
+        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, CustomBlockId,
-        TransformBlockId,
+        RenderBlock,
     },
     scroll::Autoscroll,
     Bias, Editor, EditorEvent, ExcerptId, ExcerptRange, MultiBuffer, MultiBufferSnapshot, ToPoint,
@@ -71,7 +71,7 @@ struct PathState {
     path: ProjectPath,
     first_excerpt_id: Option<ExcerptId>,
     last_excerpt_id: Option<ExcerptId>,
-    diagnostics: Vec<(DiagnosticData, BlockId)>,
+    diagnostics: Vec<(DiagnosticData, CustomBlockId)>,
 }
 
 #[derive(Debug, Clone)]
@@ -657,10 +657,10 @@ fn compare_diagnostic_range_edges(
 struct PathUpdate {
     path_excerpts_borders: (Option<ExcerptId>, Option<ExcerptId>),
     latest_excerpt_id: ExcerptId,
-    new_diagnostics: Vec<(DiagnosticData, Option<BlockId>)>,
+    new_diagnostics: Vec<(DiagnosticData, Option<CustomBlockId>)>,
     diagnostics_by_row_label: BTreeMap<MultiBufferRow, (editor::Anchor, Vec<usize>)>,
-    blocks_to_remove: HashSet<BlockId>,
+    blocks_to_remove: HashSet<CustomBlockId>,
-    unchanged_blocks: HashMap<usize, BlockId>,
+    unchanged_blocks: HashMap<usize, CustomBlockId>,
     excerpts_with_new_diagnostics: HashSet<ExcerptId>,
     excerpts_to_remove: Vec<ExcerptId>,
     excerpt_expands: HashMap<(ExpandExcerptDirection, u32), Vec<ExcerptId>>,
@@ -749,7 +749,7 @@ impl PathUpdate {
         context: u32,
         multi_buffer_snapshot: MultiBufferSnapshot,
         buffer_snapshot: BufferSnapshot,
-        current_diagnostics: impl Iterator<Item = &'a (DiagnosticData, BlockId)> + 'a,
+        current_diagnostics: impl Iterator<Item = &'a (DiagnosticData, CustomBlockId)> + 'a,
     ) {
         let mut current_diagnostics = current_diagnostics.fuse().peekable();
         let mut excerpts_to_expand =
@@ -1234,7 +1234,10 @@ impl PathUpdate {
             .collect()
     }
 
-    fn new_blocks(mut self, new_block_ids: Vec<BlockId>) -> Vec<(DiagnosticData, BlockId)> {
+    fn new_blocks(
+        mut self,
+        new_block_ids: Vec<CustomBlockId>,
+    ) -> Vec<(DiagnosticData, CustomBlockId)> {
         let mut new_block_ids = new_block_ids.into_iter().fuse();
         for (_, (_, grouped_diagnostics)) in self.diagnostics_by_row_label {
             let mut created_block_id = None;
@@ -1285,8 +1288,8 @@ fn render_same_line_diagnostics(
     folded_block_height: u8,
 ) -> RenderBlock {
     Box::new(move |cx: &mut BlockContext| {
-        let block_id = match cx.transform_block_id {
+        let block_id = match cx.block_id {
-            TransformBlockId::Block(block_id) => block_id,
+            BlockId::Custom(block_id) => block_id,
             _ => {
                 debug_panic!("Expected a block id for the diagnostics block");
                 return div().into_any_element();
@@ -1320,7 +1323,7 @@ fn render_same_line_diagnostics(
             .child(v_flex().size_full().when_some_else(
                 toggle_expand_label,
                 |parent, label| {
-                    parent.child(Button::new(cx.transform_block_id, label).on_click({
+                    parent.child(Button::new(cx.block_id, label).on_click({
                         let diagnostics = Arc::clone(&diagnostics);
                         move |_, cx| {
                             let new_expanded = !expanded;

crates/editor/src/display_map.rs

@@ -28,9 +28,8 @@ use crate::{
     hover_links::InlayHighlight, movement::TextLayoutDetails, EditorStyle, InlayId, RowExt,
 };
 pub use block_map::{
-    BlockBufferRows, BlockChunks as DisplayChunks, BlockContext, BlockDisposition, BlockId,
+    Block, BlockBufferRows, BlockChunks as DisplayChunks, BlockContext, BlockDisposition, BlockId,
-    BlockMap, BlockPoint, BlockProperties, BlockStyle, RenderBlock, TransformBlock,
+    BlockMap, BlockPoint, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
-    TransformBlockId,
 };
 use block_map::{BlockRow, BlockSnapshot};
 use collections::{HashMap, HashSet};
@@ -270,7 +269,7 @@ impl DisplayMap {
         &mut self,
         blocks: impl IntoIterator<Item = BlockProperties<Anchor>>,
         cx: &mut ModelContext<Self>,
-    ) -> Vec<BlockId> {
+    ) -> Vec<CustomBlockId> {
         let snapshot = self.buffer.read(cx).snapshot(cx);
         let edits = self.buffer_subscription.consume().into_inner();
         let tab_size = Self::tab_size(&self.buffer, cx);
@@ -286,7 +285,7 @@ impl DisplayMap {
 
     pub fn replace_blocks(
         &mut self,
-        heights_and_renderers: HashMap<BlockId, (Option<u8>, RenderBlock)>,
+        heights_and_renderers: HashMap<CustomBlockId, (Option<u8>, RenderBlock)>,
         cx: &mut ModelContext<Self>,
     ) {
         //
@@ -307,8 +306,8 @@ impl DisplayMap {
         // directly and the new behavior separately.
         //
         //
-        let mut only_renderers = HashMap::<BlockId, RenderBlock>::default();
+        let mut only_renderers = HashMap::<CustomBlockId, RenderBlock>::default();
-        let mut full_replace = HashMap::<BlockId, (u8, RenderBlock)>::default();
+        let mut full_replace = HashMap::<CustomBlockId, (u8, RenderBlock)>::default();
         for (id, (height, render)) in heights_and_renderers {
             if let Some(height) = height {
                 full_replace.insert(id, (height, render));
@@ -335,7 +334,7 @@ impl DisplayMap {
         block_map.replace(full_replace);
     }
 
-    pub fn remove_blocks(&mut self, ids: HashSet<BlockId>, cx: &mut ModelContext<Self>) {
+    pub fn remove_blocks(&mut self, ids: HashSet<CustomBlockId>, cx: &mut ModelContext<Self>) {
         let snapshot = self.buffer.read(cx).snapshot(cx);
         let edits = self.buffer_subscription.consume().into_inner();
         let tab_size = Self::tab_size(&self.buffer, cx);
@@ -351,7 +350,7 @@ impl DisplayMap {
 
     pub fn row_for_block(
         &mut self,
-        block_id: BlockId,
+        block_id: CustomBlockId,
         cx: &mut ModelContext<Self>,
     ) -> Option<DisplayRow> {
         let snapshot = self.buffer.read(cx).snapshot(cx);
@@ -886,12 +885,16 @@ impl DisplaySnapshot {
     pub fn blocks_in_range(
         &self,
         rows: Range<DisplayRow>,
-    ) -> impl Iterator<Item = (DisplayRow, &TransformBlock)> {
+    ) -> impl Iterator<Item = (DisplayRow, &Block)> {
         self.block_snapshot
             .blocks_in_range(rows.start.0..rows.end.0)
             .map(|(row, block)| (DisplayRow(row), block))
     }
 
+    pub fn block_for_id(&self, id: BlockId) -> Option<Block> {
+        self.block_snapshot.block_for_id(id)
+    }
+
     pub fn intersects_fold<T: ToOffset>(&self, offset: T) -> bool {
         self.fold_snapshot.intersects_fold(offset)
     }

crates/editor/src/display_map/block_map.rs

@@ -18,7 +18,7 @@ use std::{
         Arc,
     },
 };
-use sum_tree::{Bias, SumTree};
+use sum_tree::{Bias, SumTree, TreeMap};
 use text::Edit;
 use ui::ElementId;
 
@@ -30,7 +30,8 @@ const NEWLINES: &[u8] = &[b'\n'; u8::MAX as usize];
 pub struct BlockMap {
     next_block_id: AtomicUsize,
     wrap_snapshot: RefCell<WrapSnapshot>,
-    blocks: Vec<Arc<Block>>,
+    custom_blocks: Vec<Arc<CustomBlock>>,
+    custom_blocks_by_id: TreeMap<CustomBlockId, Arc<CustomBlock>>,
     transforms: RefCell<SumTree<Transform>>,
     show_excerpt_controls: bool,
     buffer_header_height: u8,
@@ -39,7 +40,7 @@ pub struct BlockMap {
 }
 
 pub struct BlockMapReader<'a> {
-    blocks: &'a Vec<Arc<Block>>,
+    blocks: &'a Vec<Arc<CustomBlock>>,
     pub snapshot: BlockSnapshot,
 }
 
@@ -49,12 +50,13 @@ pub struct BlockMapWriter<'a>(&'a mut BlockMap);
 pub struct BlockSnapshot {
     wrap_snapshot: WrapSnapshot,
     transforms: SumTree<Transform>,
+    custom_blocks_by_id: TreeMap<CustomBlockId, Arc<CustomBlock>>,
 }
 
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct BlockId(usize);
+pub struct CustomBlockId(usize);
 
-impl Into<ElementId> for BlockId {
+impl Into<ElementId> for CustomBlockId {
     fn into(self) -> ElementId {
         ElementId::Integer(self.0)
     }
@@ -71,8 +73,8 @@ struct WrapRow(u32);
 
 pub type RenderBlock = Box<dyn Send + FnMut(&mut BlockContext) -> AnyElement>;
 
-pub struct Block {
+pub struct CustomBlock {
-    id: BlockId,
+    id: CustomBlockId,
     position: Anchor,
     height: u8,
     style: BlockStyle,
@@ -113,41 +115,41 @@ pub struct BlockContext<'a, 'b> {
     pub gutter_dimensions: &'b GutterDimensions,
     pub em_width: Pixels,
     pub line_height: Pixels,
-    pub transform_block_id: TransformBlockId,
+    pub block_id: BlockId,
     pub editor_style: &'b EditorStyle,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum TransformBlockId {
+pub enum BlockId {
-    Block(BlockId),
+    Custom(CustomBlockId),
     ExcerptHeader(ExcerptId),
     ExcerptFooter(ExcerptId),
 }
 
-impl From<TransformBlockId> for EntityId {
+impl From<BlockId> for EntityId {
-    fn from(value: TransformBlockId) -> Self {
+    fn from(value: BlockId) -> Self {
         match value {
-            TransformBlockId::Block(BlockId(id)) => EntityId::from(id as u64),
+            BlockId::Custom(CustomBlockId(id)) => EntityId::from(id as u64),
-            TransformBlockId::ExcerptHeader(id) => id.into(),
+            BlockId::ExcerptHeader(id) => id.into(),
-            TransformBlockId::ExcerptFooter(id) => id.into(),
+            BlockId::ExcerptFooter(id) => id.into(),
         }
     }
 }
 
-impl Into<ElementId> for TransformBlockId {
+impl Into<ElementId> for BlockId {
     fn into(self) -> ElementId {
         match self {
-            Self::Block(BlockId(id)) => ("Block", id).into(),
+            Self::Custom(CustomBlockId(id)) => ("Block", id).into(),
             Self::ExcerptHeader(id) => ("ExcerptHeader", EntityId::from(id)).into(),
             Self::ExcerptFooter(id) => ("ExcerptFooter", EntityId::from(id)).into(),
         }
     }
 }
 
-impl std::fmt::Display for TransformBlockId {
+impl std::fmt::Display for BlockId {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::Block(id) => write!(f, "Block({id:?})"),
+            Self::Custom(id) => write!(f, "Block({id:?})"),
             Self::ExcerptHeader(id) => write!(f, "ExcerptHeader({id:?})"),
             Self::ExcerptFooter(id) => write!(f, "ExcerptFooter({id:?})"),
         }
@@ -164,11 +166,11 @@ pub enum BlockDisposition {
 #[derive(Clone, Debug)]
 struct Transform {
     summary: TransformSummary,
-    block: Option<TransformBlock>,
+    block: Option<Block>,
 }
 
 pub(crate) enum BlockType {
-    Custom(BlockId),
+    Custom(CustomBlockId),
     Header,
     Footer,
 }
@@ -180,8 +182,8 @@ pub(crate) trait BlockLike {
 
 #[allow(clippy::large_enum_variant)]
 #[derive(Clone)]
-pub enum TransformBlock {
+pub enum Block {
-    Custom(Arc<Block>),
+    Custom(Arc<CustomBlock>),
     ExcerptHeader {
         id: ExcerptId,
         buffer: BufferSnapshot,
@@ -197,12 +199,12 @@ pub enum TransformBlock {
     },
 }
 
-impl BlockLike for TransformBlock {
+impl BlockLike for Block {
     fn block_type(&self) -> BlockType {
         match self {
-            TransformBlock::Custom(block) => BlockType::Custom(block.id),
+            Block::Custom(block) => BlockType::Custom(block.id),
-            TransformBlock::ExcerptHeader { .. } => BlockType::Header,
+            Block::ExcerptHeader { .. } => BlockType::Header,
-            TransformBlock::ExcerptFooter { .. } => BlockType::Footer,
+            Block::ExcerptFooter { .. } => BlockType::Footer,
         }
     }
 
@@ -211,33 +213,41 @@ impl BlockLike for TransformBlock {
     }
 }
 
-impl TransformBlock {
+impl Block {
-    pub fn id(&self) -> TransformBlockId {
+    pub fn id(&self) -> BlockId {
         match self {
-            TransformBlock::Custom(block) => TransformBlockId::Block(block.id),
+            Block::Custom(block) => BlockId::Custom(block.id),
-            TransformBlock::ExcerptHeader { id, .. } => TransformBlockId::ExcerptHeader(*id),
+            Block::ExcerptHeader { id, .. } => BlockId::ExcerptHeader(*id),
-            TransformBlock::ExcerptFooter { id, .. } => TransformBlockId::ExcerptFooter(*id),
+            Block::ExcerptFooter { id, .. } => BlockId::ExcerptFooter(*id),
         }
     }
 
     fn disposition(&self) -> BlockDisposition {
         match self {
-            TransformBlock::Custom(block) => block.disposition,
+            Block::Custom(block) => block.disposition,
-            TransformBlock::ExcerptHeader { .. } => BlockDisposition::Above,
+            Block::ExcerptHeader { .. } => BlockDisposition::Above,
-            TransformBlock::ExcerptFooter { disposition, .. } => *disposition,
+            Block::ExcerptFooter { disposition, .. } => *disposition,
         }
     }
 
     pub fn height(&self) -> u8 {
         match self {
-            TransformBlock::Custom(block) => block.height,
+            Block::Custom(block) => block.height,
-            TransformBlock::ExcerptHeader { height, .. } => *height,
+            Block::ExcerptHeader { height, .. } => *height,
-            TransformBlock::ExcerptFooter { height, .. } => *height,
+            Block::ExcerptFooter { height, .. } => *height,
+        }
+    }
+
+    pub fn style(&self) -> BlockStyle {
+        match self {
+            Block::Custom(block) => block.style,
+            Block::ExcerptHeader { .. } => BlockStyle::Sticky,
+            Block::ExcerptFooter { .. } => BlockStyle::Sticky,
         }
     }
 }
 
-impl Debug for TransformBlock {
+impl Debug for Block {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Custom(block) => f.debug_struct("Custom").field("block", block).finish(),
@@ -252,7 +262,7 @@ impl Debug for TransformBlock {
                 .field("path", &buffer.file().map(|f| f.path()))
                 .field("starts_new_buffer", &starts_new_buffer)
                 .finish(),
-            TransformBlock::ExcerptFooter {
+            Block::ExcerptFooter {
                 id, disposition, ..
             } => f
                 .debug_struct("ExcerptFooter")
@@ -296,7 +306,8 @@ impl BlockMap {
         let row_count = wrap_snapshot.max_point().row() + 1;
         let map = Self {
             next_block_id: AtomicUsize::new(0),
-            blocks: Vec::new(),
+            custom_blocks: Vec::new(),
+            custom_blocks_by_id: TreeMap::default(),
             transforms: RefCell::new(SumTree::from_item(Transform::isomorphic(row_count), &())),
             wrap_snapshot: RefCell::new(wrap_snapshot.clone()),
             show_excerpt_controls,
@@ -318,10 +329,11 @@ impl BlockMap {
         self.sync(&wrap_snapshot, edits);
         *self.wrap_snapshot.borrow_mut() = wrap_snapshot.clone();
         BlockMapReader {
-            blocks: &self.blocks,
+            blocks: &self.custom_blocks,
             snapshot: BlockSnapshot {
                 wrap_snapshot,
                 transforms: self.transforms.borrow().clone(),
+                custom_blocks_by_id: self.custom_blocks_by_id.clone(),
             },
         }
     }
@@ -443,7 +455,8 @@ impl BlockMap {
             let new_buffer_start =
                 wrap_snapshot.to_point(WrapPoint::new(new_start.0, 0), Bias::Left);
             let start_bound = Bound::Included(new_buffer_start);
-            let start_block_ix = match self.blocks[last_block_ix..].binary_search_by(|probe| {
+            let start_block_ix =
+                match self.custom_blocks[last_block_ix..].binary_search_by(|probe| {
                     probe
                         .position
                         .to_point(buffer)
@@ -456,12 +469,12 @@ impl BlockMap {
             let end_bound;
             let end_block_ix = if new_end.0 > wrap_snapshot.max_point().row() {
                 end_bound = Bound::Unbounded;
-                self.blocks.len()
+                self.custom_blocks.len()
             } else {
                 let new_buffer_end =
                     wrap_snapshot.to_point(WrapPoint::new(new_end.0, 0), Bias::Left);
                 end_bound = Bound::Excluded(new_buffer_end);
-                match self.blocks[start_block_ix..].binary_search_by(|probe| {
+                match self.custom_blocks[start_block_ix..].binary_search_by(|probe| {
                     probe
                         .position
                         .to_point(buffer)
@@ -474,10 +487,8 @@ impl BlockMap {
             last_block_ix = end_block_ix;
 
             debug_assert!(blocks_in_edit.is_empty());
-            blocks_in_edit.extend(
+            blocks_in_edit.extend(self.custom_blocks[start_block_ix..end_block_ix].iter().map(
-                self.blocks[start_block_ix..end_block_ix]
+                |block| {
-                    .iter()
-                    .map(|block| {
                     let mut position = block.position.to_point(buffer);
                     match block.disposition {
                         BlockDisposition::Above => position.column = 0,
@@ -486,12 +497,12 @@ impl BlockMap {
                         }
                     }
                     let position = wrap_snapshot.make_wrap_point(position, Bias::Left);
-                        (position.row(), TransformBlock::Custom(block.clone()))
+                    (position.row(), Block::Custom(block.clone()))
-                    }),
+                },
-            );
+            ));
 
             if buffer.show_headers() {
-                blocks_in_edit.extend(BlockMap::header_blocks(
+                blocks_in_edit.extend(BlockMap::header_and_footer_blocks(
                     self.show_excerpt_controls,
                     self.excerpt_footer_height,
                     self.buffer_header_height,
@@ -538,8 +549,8 @@ impl BlockMap {
         *transforms = new_transforms;
     }
 
-    pub fn replace_renderers(&mut self, mut renderers: HashMap<BlockId, RenderBlock>) {
+    pub fn replace_renderers(&mut self, mut renderers: HashMap<CustomBlockId, RenderBlock>) {
-        for block in &mut self.blocks {
+        for block in &mut self.custom_blocks {
             if let Some(render) = renderers.remove(&block.id) {
                 *block.render.lock() = render;
             }
@@ -550,7 +561,7 @@ impl BlockMap {
         self.show_excerpt_controls
     }
 
-    pub fn header_blocks<'a, 'b: 'a, 'c: 'a + 'b, R, T>(
+    pub fn header_and_footer_blocks<'a, 'b: 'a, 'c: 'a + 'b, R, T>(
         show_excerpt_controls: bool,
         excerpt_footer_height: u8,
         buffer_header_height: u8,
@@ -558,7 +569,7 @@ impl BlockMap {
         buffer: &'b multi_buffer::MultiBufferSnapshot,
         range: R,
         wrap_snapshot: &'c WrapSnapshot,
-    ) -> impl Iterator<Item = (u32, TransformBlock)> + 'b
+    ) -> impl Iterator<Item = (u32, Block)> + 'b
     where
         R: RangeBounds<T>,
         T: multi_buffer::ToOffset,
@@ -566,24 +577,36 @@ impl BlockMap {
         buffer
             .excerpt_boundaries_in_range(range)
             .flat_map(move |excerpt_boundary| {
-                let wrap_row = wrap_snapshot
+                let mut wrap_row = wrap_snapshot
                     .make_wrap_point(Point::new(excerpt_boundary.row.0, 0), Bias::Left)
                     .row();
 
                 [
                     show_excerpt_controls
                         .then(|| {
+                            let disposition;
+                            if excerpt_boundary.next.is_some() {
+                                disposition = BlockDisposition::Above;
+                            } else {
+                                wrap_row = wrap_snapshot
+                                    .make_wrap_point(
+                                        Point::new(
+                                            excerpt_boundary.row.0,
+                                            buffer.line_len(excerpt_boundary.row),
+                                        ),
+                                        Bias::Left,
+                                    )
+                                    .row();
+                                disposition = BlockDisposition::Below;
+                            }
+
                             excerpt_boundary.prev.as_ref().map(|prev| {
                                 (
                                     wrap_row,
-                                    TransformBlock::ExcerptFooter {
+                                    Block::ExcerptFooter {
                                         id: prev.id,
                                         height: excerpt_footer_height,
-                                        disposition: if excerpt_boundary.next.is_some() {
+                                        disposition,
-                                            BlockDisposition::Above
-                                        } else {
-                                            BlockDisposition::Below
-                                        },
                                     },
                                 )
                             })
@@ -596,7 +619,7 @@ impl BlockMap {
 
                         (
                             wrap_row,
-                            TransformBlock::ExcerptHeader {
+                            Block::ExcerptHeader {
                                 id: next.id,
                                 buffer: next.buffer,
                                 range: next.range,
@@ -692,7 +715,7 @@ impl<'a> DerefMut for BlockMapReader<'a> {
 }
 
 impl<'a> BlockMapReader<'a> {
-    pub fn row_for_block(&self, block_id: BlockId) -> Option<BlockRow> {
+    pub fn row_for_block(&self, block_id: CustomBlockId) -> Option<BlockRow> {
         let block = self.blocks.iter().find(|block| block.id == block_id)?;
         let buffer_row = block
             .position
@@ -737,14 +760,14 @@ impl<'a> BlockMapWriter<'a> {
     pub fn insert(
         &mut self,
         blocks: impl IntoIterator<Item = BlockProperties<Anchor>>,
-    ) -> Vec<BlockId> {
+    ) -> Vec<CustomBlockId> {
         let mut ids = Vec::new();
         let mut edits = Patch::default();
         let wrap_snapshot = &*self.0.wrap_snapshot.borrow();
         let buffer = wrap_snapshot.buffer_snapshot();
 
         for block in blocks {
-            let id = BlockId(self.0.next_block_id.fetch_add(1, SeqCst));
+            let id = CustomBlockId(self.0.next_block_id.fetch_add(1, SeqCst));
             ids.push(id);
 
             let position = block.position;
@@ -759,22 +782,21 @@ impl<'a> BlockMapWriter<'a> {
 
             let block_ix = match self
                 .0
-                .blocks
+                .custom_blocks
                 .binary_search_by(|probe| probe.position.cmp(&position, buffer))
             {
                 Ok(ix) | Err(ix) => ix,
             };
-            self.0.blocks.insert(
+            let new_block = Arc::new(CustomBlock {
-                block_ix,
-                Arc::new(Block {
                 id,
                 position,
                 height: block.height,
                 render: Mutex::new(block.render),
                 disposition: block.disposition,
                 style: block.style,
-                }),
+            });
-            );
+            self.0.custom_blocks.insert(block_ix, new_block.clone());
+            self.0.custom_blocks_by_id.insert(id, new_block);
 
             edits = edits.compose([Edit {
                 old: start_row..end_row,
@@ -786,16 +808,19 @@ impl<'a> BlockMapWriter<'a> {
         ids
     }
 
-    pub fn replace(&mut self, mut heights_and_renderers: HashMap<BlockId, (u8, RenderBlock)>) {
+    pub fn replace(
+        &mut self,
+        mut heights_and_renderers: HashMap<CustomBlockId, (u8, RenderBlock)>,
+    ) {
         let wrap_snapshot = &*self.0.wrap_snapshot.borrow();
         let buffer = wrap_snapshot.buffer_snapshot();
         let mut edits = Patch::default();
         let mut last_block_buffer_row = None;
 
-        for block in &mut self.0.blocks {
+        for block in &mut self.0.custom_blocks {
             if let Some((new_height, render)) = heights_and_renderers.remove(&block.id) {
                 if block.height != new_height {
-                    let new_block = Block {
+                    let new_block = CustomBlock {
                         id: block.id,
                         position: block.position,
                         height: new_height,
@@ -803,7 +828,9 @@ impl<'a> BlockMapWriter<'a> {
                         render: Mutex::new(render),
                         disposition: block.disposition,
                     };
-                    *block = Arc::new(new_block);
+                    let new_block = Arc::new(new_block);
+                    *block = new_block.clone();
+                    self.0.custom_blocks_by_id.insert(block.id, new_block);
 
                     let buffer_row = block.position.to_point(buffer).row;
                     if last_block_buffer_row != Some(buffer_row) {
@@ -828,12 +855,12 @@ impl<'a> BlockMapWriter<'a> {
         self.0.sync(wrap_snapshot, edits);
     }
 
-    pub fn remove(&mut self, block_ids: HashSet<BlockId>) {
+    pub fn remove(&mut self, block_ids: HashSet<CustomBlockId>) {
         let wrap_snapshot = &*self.0.wrap_snapshot.borrow();
         let buffer = wrap_snapshot.buffer_snapshot();
         let mut edits = Patch::default();
         let mut last_block_buffer_row = None;
-        self.0.blocks.retain(|block| {
+        self.0.custom_blocks.retain(|block| {
             if block_ids.contains(&block.id) {
                 let buffer_row = block.position.to_point(buffer).row;
                 if last_block_buffer_row != Some(buffer_row) {
@@ -850,6 +877,7 @@ impl<'a> BlockMapWriter<'a> {
                         new: start_row..end_row,
                     })
                 }
+                self.0.custom_blocks_by_id.remove(&block.id);
                 false
             } else {
                 true
@@ -934,10 +962,7 @@ impl BlockSnapshot {
         }
     }
 
-    pub fn blocks_in_range(
+    pub fn blocks_in_range(&self, rows: Range<u32>) -> impl Iterator<Item = (u32, &Block)> {
-        &self,
-        rows: Range<u32>,
-    ) -> impl Iterator<Item = (u32, &TransformBlock)> {
         let mut cursor = self.transforms.cursor::<BlockRow>();
         cursor.seek(&BlockRow(rows.start), Bias::Right, &());
         std::iter::from_fn(move || {
@@ -957,6 +982,60 @@ impl BlockSnapshot {
         })
     }
 
+    pub fn block_for_id(&self, block_id: BlockId) -> Option<Block> {
+        let buffer = self.wrap_snapshot.buffer_snapshot();
+
+        match block_id {
+            BlockId::Custom(custom_block_id) => {
+                let custom_block = self.custom_blocks_by_id.get(&custom_block_id)?;
+                Some(Block::Custom(custom_block.clone()))
+            }
+            BlockId::ExcerptHeader(excerpt_id) => {
+                let excerpt_range = buffer.range_for_excerpt::<Point>(excerpt_id)?;
+                let wrap_point = self
+                    .wrap_snapshot
+                    .make_wrap_point(excerpt_range.start, Bias::Left);
+                let mut cursor = self.transforms.cursor::<(WrapRow, BlockRow)>();
+                cursor.seek(&WrapRow(wrap_point.row()), Bias::Left, &());
+                while let Some(transform) = cursor.item() {
+                    if let Some(block) = transform.block.as_ref() {
+                        if block.id() == block_id {
+                            return Some(block.clone());
+                        }
+                    } else if cursor.start().0 > WrapRow(wrap_point.row()) {
+                        break;
+                    }
+
+                    cursor.next(&());
+                }
+
+                None
+            }
+            BlockId::ExcerptFooter(excerpt_id) => {
+                let excerpt_range = buffer.range_for_excerpt::<Point>(excerpt_id)?;
+                let wrap_point = self
+                    .wrap_snapshot
+                    .make_wrap_point(excerpt_range.end, Bias::Left);
+
+                let mut cursor = self.transforms.cursor::<(WrapRow, BlockRow)>();
+                cursor.seek(&WrapRow(wrap_point.row()), Bias::Left, &());
+                while let Some(transform) = cursor.item() {
+                    if let Some(block) = transform.block.as_ref() {
+                        if block.id() == block_id {
+                            return Some(block.clone());
+                        }
+                    } else if cursor.start().0 > WrapRow(wrap_point.row()) {
+                        break;
+                    }
+
+                    cursor.next(&());
+                }
+
+                None
+            }
+        }
+    }
+
     pub fn max_point(&self) -> BlockPoint {
         let row = self.transforms.summary().output_rows - 1;
         BlockPoint::new(row, self.line_len(BlockRow(row)))
@@ -1086,7 +1165,7 @@ impl Transform {
         }
     }
 
-    fn block(block: TransformBlock) -> Self {
+    fn block(block: Block) -> Self {
         Self {
             summary: TransformSummary {
                 input_rows: 0,
@@ -1235,7 +1314,7 @@ impl DerefMut for BlockContext<'_, '_> {
     }
 }
 
-impl Block {
+impl CustomBlock {
     pub fn render(&self, cx: &mut BlockContext) -> AnyElement {
         self.render.lock()(cx)
     }
@@ -1249,7 +1328,7 @@ impl Block {
     }
 }
 
-impl Debug for Block {
+impl Debug for CustomBlock {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Block")
             .field("id", &self.id)
@@ -1279,15 +1358,16 @@ fn offset_for_row(s: &str, target: u32) -> (u32, usize) {
 
 #[cfg(test)]
 mod tests {
-    use std::env;
-
     use super::*;
-    use crate::display_map::inlay_map::InlayMap;
+    use crate::display_map::{
-    use crate::display_map::{fold_map::FoldMap, tab_map::TabMap, wrap_map::WrapMap};
+        fold_map::FoldMap, inlay_map::InlayMap, tab_map::TabMap, wrap_map::WrapMap,
-    use gpui::{div, font, px, Element};
+    };
+    use gpui::{div, font, px, AppContext, Context as _, Element};
+    use language::{Buffer, Capability};
     use multi_buffer::MultiBuffer;
     use rand::prelude::*;
     use settings::SettingsStore;
+    use std::env;
     use util::RandomCharIter;
 
     #[gpui::test]
@@ -1474,6 +1554,89 @@ mod tests {
         assert_eq!(snapshot.text(), "aaa\n\nb!!!\n\n\nbb\nccc\nddd\n\n\n");
     }
 
+    #[gpui::test]
+    fn test_multibuffer_headers_and_footers(cx: &mut AppContext) {
+        init_test(cx);
+
+        let buffer1 = cx.new_model(|cx| Buffer::local("Buffer 1", cx));
+        let buffer2 = cx.new_model(|cx| Buffer::local("Buffer 2", cx));
+        let buffer3 = cx.new_model(|cx| Buffer::local("Buffer 3", cx));
+
+        let mut excerpt_ids = Vec::new();
+        let multi_buffer = cx.new_model(|cx| {
+            let mut multi_buffer = MultiBuffer::new(0, Capability::ReadWrite);
+            excerpt_ids.extend(multi_buffer.push_excerpts(
+                buffer1.clone(),
+                [ExcerptRange {
+                    context: 0..buffer1.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            ));
+            excerpt_ids.extend(multi_buffer.push_excerpts(
+                buffer2.clone(),
+                [ExcerptRange {
+                    context: 0..buffer2.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            ));
+            excerpt_ids.extend(multi_buffer.push_excerpts(
+                buffer3.clone(),
+                [ExcerptRange {
+                    context: 0..buffer3.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            ));
+
+            multi_buffer
+        });
+
+        let font = font("Helvetica");
+        let font_size = px(14.);
+        let font_id = cx.text_system().resolve_font(&font);
+        let mut wrap_width = px(0.);
+        for c in "Buff".chars() {
+            wrap_width += cx
+                .text_system()
+                .advance(font_id, font_size, c)
+                .unwrap()
+                .width;
+        }
+
+        let multi_buffer_snapshot = multi_buffer.read(cx).snapshot(cx);
+        let (_, inlay_snapshot) = InlayMap::new(multi_buffer_snapshot.clone());
+        let (_, fold_snapshot) = FoldMap::new(inlay_snapshot);
+        let (_, tab_snapshot) = TabMap::new(fold_snapshot, 4.try_into().unwrap());
+        let (_, wraps_snapshot) = WrapMap::new(tab_snapshot, font, font_size, Some(wrap_width), cx);
+
+        let block_map = BlockMap::new(wraps_snapshot.clone(), true, 1, 1, 1);
+        let snapshot = block_map.read(wraps_snapshot, Default::default());
+
+        // Each excerpt has a header above and footer below. Excerpts are also *separated* by a newline.
+        assert_eq!(
+            snapshot.text(),
+            "\nBuff\ner 1\n\n\nBuff\ner 2\n\n\nBuff\ner 3\n"
+        );
+
+        let blocks: Vec<_> = snapshot
+            .blocks_in_range(0..u32::MAX)
+            .map(|(row, block)| (row, block.id()))
+            .collect();
+        assert_eq!(
+            blocks,
+            vec![
+                (0, BlockId::ExcerptHeader(excerpt_ids[0])),
+                (3, BlockId::ExcerptFooter(excerpt_ids[0])),
+                (4, BlockId::ExcerptHeader(excerpt_ids[1])),
+                (7, BlockId::ExcerptFooter(excerpt_ids[1])),
+                (8, BlockId::ExcerptHeader(excerpt_ids[2])),
+                (11, BlockId::ExcerptFooter(excerpt_ids[2]))
+            ]
+        );
+    }
+
     #[gpui::test]
     fn test_replace_with_heights(cx: &mut gpui::TestAppContext) {
         let _update = cx.update(|cx| init_test(cx));
@@ -1807,7 +1970,7 @@ mod tests {
 
             // Note that this needs to be synced with the related section in BlockMap::sync
             expected_blocks.extend(
-                BlockMap::header_blocks(
+                BlockMap::header_and_footer_blocks(
                     true,
                     excerpt_footer_height,
                     buffer_start_header_height,
@@ -1911,6 +2074,16 @@ mod tests {
                 expected_block_positions
             );
 
+            for (_, expected_block) in
+                blocks_snapshot.blocks_in_range(0..(expected_row_count as u32))
+            {
+                let actual_block = blocks_snapshot.block_for_id(expected_block.id());
+                assert_eq!(
+                    actual_block.map(|block| block.id()),
+                    Some(expected_block.id())
+                );
+            }
+
             for (block_row, block) in expected_block_positions {
                 if let BlockType::Custom(block_id) = block.block_type() {
                     assert_eq!(
@@ -2007,7 +2180,7 @@ mod tests {
             },
             Custom {
                 disposition: BlockDisposition,
-                id: BlockId,
+                id: CustomBlockId,
                 height: u8,
             },
         }
@@ -2044,15 +2217,15 @@ mod tests {
             }
         }
 
-        impl From<TransformBlock> for ExpectedBlock {
+        impl From<Block> for ExpectedBlock {
-            fn from(block: TransformBlock) -> Self {
+            fn from(block: Block) -> Self {
                 match block {
-                    TransformBlock::Custom(block) => ExpectedBlock::Custom {
+                    Block::Custom(block) => ExpectedBlock::Custom {
                         id: block.id,
                         disposition: block.disposition,
                         height: block.height,
                     },
-                    TransformBlock::ExcerptHeader {
+                    Block::ExcerptHeader {
                         height,
                         starts_new_buffer,
                         ..
@@ -2060,7 +2233,7 @@ mod tests {
                         height,
                         starts_new_buffer,
                     },
-                    TransformBlock::ExcerptFooter {
+                    Block::ExcerptFooter {
                         height,
                         disposition,
                         ..
@@ -2080,12 +2253,12 @@ mod tests {
         assets::Assets.load_test_fonts(cx);
     }
 
-    impl TransformBlock {
+    impl Block {
-        fn as_custom(&self) -> Option<&Block> {
+        fn as_custom(&self) -> Option<&CustomBlock> {
             match self {
-                TransformBlock::Custom(block) => Some(block),
+                Block::Custom(block) => Some(block),
-                TransformBlock::ExcerptHeader { .. } => None,
+                Block::ExcerptHeader { .. } => None,
-                TransformBlock::ExcerptFooter { .. } => None,
+                Block::ExcerptFooter { .. } => None,
             }
         }
     }

crates/editor/src/editor.rs

@@ -568,6 +568,7 @@ pub struct Editor {
     previous_search_ranges: Option<Arc<[Range<Anchor>]>>,
     file_header_size: u8,
     breadcrumb_header: Option<String>,
+    focused_block: Option<FocusedBlock>,
 }
 
 #[derive(Clone)]
@@ -785,7 +786,7 @@ pub struct RenameState {
     pub range: Range<Anchor>,
     pub old_name: Arc<str>,
     pub editor: View<Editor>,
-    block_id: BlockId,
+    block_id: CustomBlockId,
 }
 
 struct InvalidationStack<T>(Vec<T>);
@@ -1537,7 +1538,7 @@ struct ActiveDiagnosticGroup {
     primary_range: Range<Anchor>,
     primary_message: String,
     group_id: usize,
-    blocks: HashMap<BlockId, Diagnostic>,
+    blocks: HashMap<CustomBlockId, Diagnostic>,
     is_valid: bool,
 }
 
@@ -1585,6 +1586,11 @@ impl InlayHintRefreshReason {
     }
 }
 
+pub(crate) struct FocusedBlock {
+    id: BlockId,
+    focus_handle: WeakFocusHandle,
+}
+
 impl Editor {
     pub fn single_line(cx: &mut ViewContext<Self>) -> Self {
         let buffer = cx.new_model(|cx| Buffer::local("", cx));
@@ -1908,6 +1914,7 @@ impl Editor {
             linked_edit_ranges: Default::default(),
             previous_search_ranges: None,
             breadcrumb_header: None,
+            focused_block: None,
         };
         this.tasks_update_task = Some(this.refresh_runnables(cx));
         this._subscriptions.extend(project_subscriptions);
@@ -10150,7 +10157,7 @@ impl Editor {
         blocks: impl IntoIterator<Item = BlockProperties<Anchor>>,
         autoscroll: Option<Autoscroll>,
         cx: &mut ViewContext<Self>,
-    ) -> Vec<BlockId> {
+    ) -> Vec<CustomBlockId> {
         let blocks = self
             .display_map
             .update(cx, |display_map, cx| display_map.insert_blocks(blocks, cx));
@@ -10162,7 +10169,7 @@ impl Editor {
 
     pub fn replace_blocks(
         &mut self,
-        blocks: HashMap<BlockId, (Option<u8>, RenderBlock)>,
+        blocks: HashMap<CustomBlockId, (Option<u8>, RenderBlock)>,
         autoscroll: Option<Autoscroll>,
         cx: &mut ViewContext<Self>,
     ) {
@@ -10175,7 +10182,7 @@ impl Editor {
 
     pub fn remove_blocks(
         &mut self,
-        block_ids: HashSet<BlockId>,
+        block_ids: HashSet<CustomBlockId>,
         autoscroll: Option<Autoscroll>,
         cx: &mut ViewContext<Self>,
     ) {
@@ -10189,13 +10196,21 @@ impl Editor {
 
     pub fn row_for_block(
         &self,
-        block_id: BlockId,
+        block_id: CustomBlockId,
         cx: &mut ViewContext<Self>,
     ) -> Option<DisplayRow> {
         self.display_map
             .update(cx, |map, cx| map.row_for_block(block_id, cx))
     }
 
+    pub(crate) fn set_focused_block(&mut self, focused_block: FocusedBlock) {
+        self.focused_block = Some(focused_block);
+    }
+
+    pub(crate) fn take_focused_block(&mut self) -> Option<FocusedBlock> {
+        self.focused_block.take()
+    }
+
     pub fn insert_creases(
         &mut self,
         creases: impl IntoIterator<Item = Crease>,
@@ -12830,7 +12845,7 @@ pub fn diagnostic_block_renderer(
         highlight_diagnostic_message(&diagnostic, max_message_rows);
 
     Box::new(move |cx: &mut BlockContext| {
-        let group_id: SharedString = cx.transform_block_id.to_string().into();
+        let group_id: SharedString = cx.block_id.to_string().into();
 
         let mut text_style = cx.text_style().clone();
         text_style.color = diagnostic_style(diagnostic.severity, cx.theme().status());
@@ -12842,7 +12857,7 @@ pub fn diagnostic_block_renderer(
 
         let multi_line_diagnostic = diagnostic.message.contains('\n');
 
-        let buttons = |diagnostic: &Diagnostic, block_id: TransformBlockId| {
+        let buttons = |diagnostic: &Diagnostic, block_id: BlockId| {
             if multi_line_diagnostic {
                 v_flex()
             } else {
@@ -12873,12 +12888,12 @@ pub fn diagnostic_block_renderer(
             )
         };
 
-        let icon_size = buttons(&diagnostic, cx.transform_block_id)
+        let icon_size = buttons(&diagnostic, cx.block_id)
             .into_any_element()
             .layout_as_root(AvailableSpace::min_size(), cx);
 
         h_flex()
-            .id(cx.transform_block_id)
+            .id(cx.block_id)
             .group(group_id.clone())
             .relative()
             .size_full()
@@ -12890,7 +12905,7 @@ pub fn diagnostic_block_renderer(
                     .w(cx.anchor_x - cx.gutter_dimensions.width - icon_size.width)
                     .flex_shrink(),
             )
-            .child(buttons(&diagnostic, cx.transform_block_id))
+            .child(buttons(&diagnostic, cx.block_id))
             .child(div().flex().flex_shrink_0().child(
                 StyledText::new(text_without_backticks.clone()).with_highlights(
                     &text_style,

crates/editor/src/element.rs

@@ -1,15 +1,11 @@
-use crate::editor_settings::ScrollBeyondLastLine;
-use crate::hunk_diff::ExpandedHunk;
-use crate::mouse_context_menu::MenuPosition;
-use crate::RangeToAnchorExt;
-use crate::TransformBlockId;
 use crate::{
     blame_entry_tooltip::{blame_entry_relative_timestamp, BlameEntryTooltip},
     display_map::{
-        BlockContext, BlockStyle, DisplaySnapshot, HighlightedChunk, ToDisplayPoint, TransformBlock,
+        Block, BlockContext, BlockStyle, DisplaySnapshot, HighlightedChunk, ToDisplayPoint,
     },
     editor_settings::{
-        CurrentLineHighlight, DoubleClickInMultibuffer, MultiCursorModifier, ShowScrollbar,
+        CurrentLineHighlight, DoubleClickInMultibuffer, MultiCursorModifier, ScrollBeyondLastLine,
+        ShowScrollbar,
     },
     git::{
         blame::{CommitDetails, GitBlame},
@@ -18,15 +14,17 @@ use crate::{
     hover_popover::{
         self, hover_at, HOVER_POPOVER_GAP, MIN_POPOVER_CHARACTER_WIDTH, MIN_POPOVER_LINE_HEIGHT,
     },
+    hunk_diff::ExpandedHunk,
     hunk_status,
     items::BufferSearchHighlights,
+    mouse_context_menu::MenuPosition,
     mouse_context_menu::{self, MouseContextMenu},
     scroll::scroll_amount::ScrollAmount,
-    CodeActionsMenu, CursorShape, DisplayPoint, DisplayRow, DocumentHighlightRead,
+    BlockId, CodeActionsMenu, CursorShape, DisplayPoint, DisplayRow, DocumentHighlightRead,
     DocumentHighlightWrite, Editor, EditorMode, EditorSettings, EditorSnapshot, EditorStyle,
-    ExpandExcerpts, GutterDimensions, HalfPageDown, HalfPageUp, HoveredCursor, HoveredHunk,
+    ExpandExcerpts, FocusedBlock, GutterDimensions, HalfPageDown, HalfPageUp, HoveredCursor,
-    LineDown, LineUp, OpenExcerpts, PageDown, PageUp, Point, RowExt, RowRangeExt, SelectPhase,
+    HoveredHunk, LineDown, LineUp, OpenExcerpts, PageDown, PageUp, Point, RangeToAnchorExt, RowExt,
-    Selection, SoftWrap, ToPoint, CURSORS_VISIBLE_FOR, MAX_LINE_LEN,
+    RowRangeExt, SelectPhase, Selection, SoftWrap, ToPoint, CURSORS_VISIBLE_FOR, MAX_LINE_LEN,
 };
 use client::ParticipantIndex;
 use collections::{BTreeMap, HashMap};
@@ -1526,7 +1524,7 @@ impl EditorElement {
         let mut block_offset = 0;
         let mut found_excerpt_header = false;
         for (_, block) in snapshot.blocks_in_range(prev_line..row_range.start) {
-            if matches!(block, TransformBlock::ExcerptHeader { .. }) {
+            if matches!(block, Block::ExcerptHeader { .. }) {
                 found_excerpt_header = true;
                 break;
             }
@@ -1543,7 +1541,7 @@ impl EditorElement {
         let mut block_height = 0;
         let mut found_excerpt_header = false;
         for (_, block) in snapshot.blocks_in_range(row_range.end..cons_line) {
-            if matches!(block, TransformBlock::ExcerptHeader { .. }) {
+            if matches!(block, Block::ExcerptHeader { .. }) {
                 found_excerpt_header = true;
             }
             block_height += block.height();
@@ -1921,35 +1919,25 @@ impl EditorElement {
     }
 
     #[allow(clippy::too_many_arguments)]
-    fn build_blocks(
+    fn render_block(
         &self,
-        rows: Range<DisplayRow>,
+        block: &Block,
+        available_space: Size<AvailableSpace>,
+        block_id: BlockId,
+        block_row_start: DisplayRow,
         snapshot: &EditorSnapshot,
-        hitbox: &Hitbox,
+        text_x: Pixels,
+        rows: &Range<DisplayRow>,
+        line_layouts: &[LineWithInvisibles],
+        gutter_dimensions: &GutterDimensions,
+        line_height: Pixels,
+        em_width: Pixels,
         text_hitbox: &Hitbox,
         scroll_width: &mut Pixels,
-        gutter_dimensions: &GutterDimensions,
-        em_width: Pixels,
-        text_x: Pixels,
-        line_height: Pixels,
-        line_layouts: &[LineWithInvisibles],
         cx: &mut WindowContext,
-    ) -> Vec<BlockLayout> {
+    ) -> (AnyElement, Size<Pixels>) {
-        let (fixed_blocks, non_fixed_blocks) = snapshot
-            .blocks_in_range(rows.clone())
-            .partition::<Vec<_>, _>(|(_, block)| match block {
-                TransformBlock::ExcerptHeader { .. } => false,
-                TransformBlock::Custom(block) => block.style() == BlockStyle::Fixed,
-                TransformBlock::ExcerptFooter { .. } => false,
-            });
-
-        let render_block = |block: &TransformBlock,
-                            available_space: Size<AvailableSpace>,
-                            block_id: TransformBlockId,
-                            block_row_start: DisplayRow,
-                            cx: &mut WindowContext| {
         let mut element = match block {
-                TransformBlock::Custom(block) => {
+            Block::Custom(block) => {
                 let align_to = block
                     .position()
                     .to_point(&snapshot.buffer_snapshot)
@@ -1969,13 +1957,13 @@ impl EditorElement {
                     gutter_dimensions,
                     line_height,
                     em_width,
-                        transform_block_id: block_id,
+                    block_id,
                     max_width: text_hitbox.size.width.max(*scroll_width),
                     editor_style: &self.style,
                 })
             }
 
-                TransformBlock::ExcerptHeader {
+            Block::ExcerptHeader {
                 buffer,
                 range,
                 starts_new_buffer,
@@ -2080,9 +2068,9 @@ impl EditorElement {
                                             )
                                             .when_some(parent_path, |then, path| {
                                                 then.child(
-                                                        div().child(path).text_color(
+                                                    div()
-                                                            cx.theme().colors().text_muted,
+                                                        .child(path)
-                                                        ),
+                                                        .text_color(cx.theme().colors().text_muted),
                                                 )
                                             }),
                                     ),
@@ -2091,11 +2079,7 @@ impl EditorElement {
                                     el.child(Icon::new(IconName::ArrowUpRight))
                                         .cursor_pointer()
                                         .tooltip(|cx| {
-                                                Tooltip::for_action(
+                                            Tooltip::for_action("Jump to File", &OpenExcerpts, cx)
-                                                    "Jump to File",
-                                                    &OpenExcerpts,
-                                                    cx,
-                                                )
                                         })
                                         .on_mouse_down(MouseButton::Left, |_, cx| {
                                             cx.stop_propagation()
@@ -2127,9 +2111,7 @@ impl EditorElement {
                                             svg()
                                                 .path(IconName::ArrowUpFromLine.path())
                                                 .size(IconSize::XSmall.rems())
-                                                    .text_color(
+                                                .text_color(cx.theme().colors().editor_line_number)
-                                                        cx.theme().colors().editor_line_number,
-                                                    )
                                                 .group("")
                                                 .hover(|style| {
                                                     style.text_color(
@@ -2189,7 +2171,8 @@ impl EditorElement {
                                 .w(icon_offset)
                                 .h_full()
                                 .child(
-                                        show_excerpt_controls.then(|| {
+                                    show_excerpt_controls
+                                        .then(|| {
                                             ButtonLike::new("expand-icon")
                                                 .style(ButtonStyle::Transparent)
                                                 .child(
@@ -2227,7 +2210,8 @@ impl EditorElement {
                                                         )
                                                     }
                                                 })
-                                        }).unwrap_or_else(|| {
+                                        })
+                                        .unwrap_or_else(|| {
                                             ButtonLike::new("jump-icon")
                                                 .style(ButtonStyle::Transparent)
                                                 .child(
@@ -2238,12 +2222,14 @@ impl EditorElement {
                                                             cx.theme().colors().border_variant,
                                                         )
                                                         .group("excerpt-jump-action")
-                                                        .group_hover("excerpt-jump-action", |style| {
+                                                        .group_hover(
+                                                            "excerpt-jump-action",
+                                                            |style| {
                                                                 style.text_color(
-                                                                cx.theme().colors().border
+                                                                    cx.theme().colors().border,
-
                                                                 )
-                                                        })
+                                                            },
+                                                        ),
                                                 )
                                                 .when_some(jump_data.clone(), |this, jump_data| {
                                                     this.on_click(cx.listener_for(&self.editor, {
@@ -2272,8 +2258,7 @@ impl EditorElement {
                                                         )
                                                     })
                                                 })
-                                        })
+                                        }),
-
                                 ),
                         )
                         .group("excerpt-jump-action")
@@ -2309,7 +2294,7 @@ impl EditorElement {
                 element.into_any()
             }
 
-                TransformBlock::ExcerptFooter { id, .. } => {
+            Block::ExcerptFooter { id, .. } => {
                 let element = v_flex()
                     .id(("excerpt footer", EntityId::from(block_id)))
                     .size_full()
@@ -2331,9 +2316,7 @@ impl EditorElement {
                                             .group("")
                                             .hover(|style| {
                                                 style.text_color(
-                                                        cx.theme()
+                                                    cx.theme().colors().editor_active_line_number,
-                                                            .colors()
-                                                            .editor_active_line_number,
                                                 )
                                             }),
                                     )
@@ -2364,8 +2347,30 @@ impl EditorElement {
 
         let size = element.layout_as_root(available_space, cx);
         (element, size)
-        };
+    }
 
+    #[allow(clippy::too_many_arguments)]
+    fn render_blocks(
+        &self,
+        rows: Range<DisplayRow>,
+        snapshot: &EditorSnapshot,
+        hitbox: &Hitbox,
+        text_hitbox: &Hitbox,
+        scroll_width: &mut Pixels,
+        gutter_dimensions: &GutterDimensions,
+        em_width: Pixels,
+        text_x: Pixels,
+        line_height: Pixels,
+        line_layouts: &[LineWithInvisibles],
+        cx: &mut WindowContext,
+    ) -> Vec<BlockLayout> {
+        let (fixed_blocks, non_fixed_blocks) = snapshot
+            .blocks_in_range(rows.clone())
+            .partition::<Vec<_>, _>(|(_, block)| block.style() == BlockStyle::Fixed);
+
+        let mut focused_block = self
+            .editor
+            .update(cx, |editor, _| editor.take_focused_block());
         let mut fixed_block_max_width = Pixels::ZERO;
         let mut blocks = Vec::new();
         for (row, block) in fixed_blocks {
@@ -2374,9 +2379,30 @@ impl EditorElement {
                 AvailableSpace::Definite(block.height() as f32 * line_height),
             );
             let block_id = block.id();
-            let (element, element_size) = render_block(block, available_space, block_id, row, cx);
+
+            if focused_block.as_ref().map_or(false, |b| b.id == block_id) {
+                focused_block = None;
+            }
+
+            let (element, element_size) = self.render_block(
+                block,
+                available_space,
+                block_id,
+                row,
+                snapshot,
+                text_x,
+                &rows,
+                line_layouts,
+                gutter_dimensions,
+                line_height,
+                em_width,
+                text_hitbox,
+                scroll_width,
+                cx,
+            );
             fixed_block_max_width = fixed_block_max_width.max(element_size.width + em_width);
             blocks.push(BlockLayout {
+                id: block_id,
                 row,
                 element,
                 available_space,
@@ -2384,11 +2410,7 @@ impl EditorElement {
             });
         }
         for (row, block) in non_fixed_blocks {
-            let style = match block {
+            let style = block.style();
-                TransformBlock::Custom(block) => block.style(),
-                TransformBlock::ExcerptHeader { .. } => BlockStyle::Sticky,
-                TransformBlock::ExcerptFooter { .. } => BlockStyle::Sticky,
-            };
             let width = match style {
                 BlockStyle::Sticky => hitbox.size.width,
                 BlockStyle::Flex => hitbox
@@ -2403,8 +2425,29 @@ impl EditorElement {
                 AvailableSpace::Definite(block.height() as f32 * line_height),
             );
             let block_id = block.id();
-            let (element, _) = render_block(block, available_space, block_id, row, cx);
+
+            if focused_block.as_ref().map_or(false, |b| b.id == block_id) {
+                focused_block = None;
+            }
+
+            let (element, _) = self.render_block(
+                block,
+                available_space,
+                block_id,
+                row,
+                snapshot,
+                text_x,
+                &rows,
+                line_layouts,
+                gutter_dimensions,
+                line_height,
+                em_width,
+                text_hitbox,
+                scroll_width,
+                cx,
+            );
             blocks.push(BlockLayout {
+                id: block_id,
                 row,
                 element,
                 available_space,
@@ -2412,6 +2455,56 @@ impl EditorElement {
             });
         }
 
+        if let Some(focused_block) = focused_block {
+            if let Some(focus_handle) = focused_block.focus_handle.upgrade() {
+                if focus_handle.is_focused(cx) {
+                    if let Some(block) = snapshot.block_for_id(focused_block.id) {
+                        let style = block.style();
+                        let width = match style {
+                            BlockStyle::Fixed => AvailableSpace::MinContent,
+                            BlockStyle::Flex => AvailableSpace::Definite(
+                                hitbox
+                                    .size
+                                    .width
+                                    .max(fixed_block_max_width)
+                                    .max(gutter_dimensions.width + *scroll_width),
+                            ),
+                            BlockStyle::Sticky => AvailableSpace::Definite(hitbox.size.width),
+                        };
+                        let available_space = size(
+                            width,
+                            AvailableSpace::Definite(block.height() as f32 * line_height),
+                        );
+
+                        let (element, _) = self.render_block(
+                            &block,
+                            available_space,
+                            focused_block.id,
+                            rows.end,
+                            snapshot,
+                            text_x,
+                            &rows,
+                            line_layouts,
+                            gutter_dimensions,
+                            line_height,
+                            em_width,
+                            text_hitbox,
+                            scroll_width,
+                            cx,
+                        );
+
+                        blocks.push(BlockLayout {
+                            id: block.id(),
+                            row: rows.end,
+                            element,
+                            available_space,
+                            style,
+                        });
+                    }
+                }
+            }
+        }
+
         *scroll_width = (*scroll_width).max(fixed_block_max_width - gutter_dimensions.width);
         blocks
     }
@@ -2433,9 +2526,19 @@ impl EditorElement {
             if !matches!(block.style, BlockStyle::Sticky) {
                 origin += point(-scroll_pixel_position.x, Pixels::ZERO);
             }
-            block
+
+            let focus_handle = block
                 .element
                 .prepaint_as_root(origin, block.available_space, cx);
+
+            if let Some(focus_handle) = focus_handle {
+                self.editor.update(cx, |editor, _cx| {
+                    editor.set_focused_block(FocusedBlock {
+                        id: block.id,
+                        focus_handle: focus_handle.downgrade(),
+                    });
+                });
+            }
         }
     }
 
@@ -3096,7 +3199,7 @@ impl EditorElement {
                     let end_row_in_current_excerpt = snapshot
                         .blocks_in_range(start_row..end_row)
                         .find_map(|(start_row, block)| {
-                            if matches!(block, TransformBlock::ExcerptHeader { .. }) {
+                            if matches!(block, Block::ExcerptHeader { .. }) {
                                 Some(start_row)
                             } else {
                                 None
@@ -4765,7 +4868,9 @@ impl Element for EditorElement {
             line_height: Some(self.style.text.line_height),
             ..Default::default()
         };
+        let focus_handle = self.editor.focus_handle(cx);
         cx.set_view_id(self.editor.entity_id());
+        cx.set_focus_handle(&focus_handle);
 
         let rem_size = self.rem_size(cx);
         cx.with_rem_size(rem_size, |cx| {
@@ -4994,7 +5099,7 @@ impl Element for EditorElement {
                         longest_line_width.max(max_visible_line_width) + overscroll.width;
 
                     let mut blocks = cx.with_element_namespace("blocks", |cx| {
-                        self.build_blocks(
+                        self.render_blocks(
                             start_row..end_row,
                             &snapshot,
                             &hitbox,
@@ -5405,7 +5510,6 @@ impl Element for EditorElement {
     ) {
         let focus_handle = self.editor.focus_handle(cx);
         let key_context = self.editor.read(cx).key_context(cx);
-        cx.set_focus_handle(&focus_handle);
         cx.set_key_context(key_context);
         cx.handle_input(
             &focus_handle,
@@ -5717,6 +5821,7 @@ impl PositionMap {
 }
 
 struct BlockLayout {
+    id: BlockId,
     row: DisplayRow,
     element: AnyElement,
     available_space: Size<AvailableSpace>,

crates/editor/src/hunk_diff.rs

@@ -23,7 +23,7 @@ use crate::{
     git::{diff_hunk_to_display, DisplayDiffHunk},
     hunk_status, hunks_for_selections,
     mouse_context_menu::MouseContextMenu,
-    BlockDisposition, BlockId, BlockProperties, BlockStyle, DiffRowHighlight, Editor,
+    BlockDisposition, BlockProperties, BlockStyle, CustomBlockId, DiffRowHighlight, Editor,
     EditorSnapshot, ExpandAllHunkDiffs, RangeToAnchorExt, RevertSelectedHunks, ToDisplayPoint,
     ToggleHunkDiff,
 };
@@ -58,7 +58,7 @@ impl ExpandedHunks {
 
 #[derive(Debug, Clone)]
 pub(super) struct ExpandedHunk {
-    pub block: Option<BlockId>,
+    pub block: Option<CustomBlockId>,
     pub hunk_range: Range<Anchor>,
     pub diff_base_byte_range: Range<usize>,
     pub status: DiffHunkStatus,
@@ -425,7 +425,7 @@ impl Editor {
         deleted_text_height: u8,
         hunk: &HoveredHunk,
         cx: &mut ViewContext<'_, Self>,
-    ) -> Option<BlockId> {
+    ) -> Option<CustomBlockId> {
         let deleted_hunk_color = deleted_hunk_color(cx);
         let (editor_height, editor_with_deleted_text) =
             editor_with_deleted_text(diff_base_buffer, deleted_hunk_color, hunk, cx);

crates/gpui/src/element.rs

@@ -32,8 +32,8 @@
 //! your own custom layout algorithm or rendering a code editor.
 
 use crate::{
-    util::FluentBuilder, ArenaBox, AvailableSpace, Bounds, DispatchNodeId, ElementId, LayoutId,
+    util::FluentBuilder, ArenaBox, AvailableSpace, Bounds, DispatchNodeId, ElementId, FocusHandle,
-    Pixels, Point, Size, Style, ViewContext, WindowContext, ELEMENT_ARENA,
+    LayoutId, Pixels, Point, Size, Style, ViewContext, WindowContext, ELEMENT_ARENA,
 };
 use derive_more::{Deref, DerefMut};
 pub(crate) use smallvec::SmallVec;
@@ -209,7 +209,7 @@ impl<C: RenderOnce> Element for Component<C> {
         _: &mut Self::PrepaintState,
         cx: &mut WindowContext,
     ) {
-        element.paint(cx)
+        element.paint(cx);
     }
 }
 
@@ -493,13 +493,23 @@ impl AnyElement {
 
     /// Prepares the element to be painted by storing its bounds, giving it a chance to draw hitboxes and
     /// request autoscroll before the final paint pass is confirmed.
-    pub fn prepaint(&mut self, cx: &mut WindowContext) {
+    pub fn prepaint(&mut self, cx: &mut WindowContext) -> Option<FocusHandle> {
-        self.0.prepaint(cx)
+        let focus_assigned = cx.window.next_frame.focus.is_some();
+
+        self.0.prepaint(cx);
+
+        if !focus_assigned {
+            if let Some(focus_id) = cx.window.next_frame.focus {
+                return FocusHandle::for_id(focus_id, &cx.window.focus_handles);
+            }
+        }
+
+        None
     }
 
     /// Paints the element stored in this `AnyElement`.
     pub fn paint(&mut self, cx: &mut WindowContext) {
-        self.0.paint(cx)
+        self.0.paint(cx);
     }
 
     /// Performs layout for this element within the given available space and returns its size.
@@ -512,19 +522,25 @@ impl AnyElement {
     }
 
     /// Prepaints this element at the given absolute origin.
-    pub fn prepaint_at(&mut self, origin: Point<Pixels>, cx: &mut WindowContext) {
+    /// If any element in the subtree beneath this element is focused, its FocusHandle is returned.
-        cx.with_absolute_element_offset(origin, |cx| self.0.prepaint(cx));
+    pub fn prepaint_at(
+        &mut self,
+        origin: Point<Pixels>,
+        cx: &mut WindowContext,
+    ) -> Option<FocusHandle> {
+        cx.with_absolute_element_offset(origin, |cx| self.prepaint(cx))
     }
 
     /// Performs layout on this element in the available space, then prepaints it at the given absolute origin.
+    /// If any element in the subtree beneath this element is focused, its FocusHandle is returned.
     pub fn prepaint_as_root(
         &mut self,
         origin: Point<Pixels>,
         available_space: Size<AvailableSpace>,
         cx: &mut WindowContext,
-    ) {
+    ) -> Option<FocusHandle> {
         self.layout_as_root(available_space, cx);
-        cx.with_absolute_element_offset(origin, |cx| self.0.prepaint(cx));
+        cx.with_absolute_element_offset(origin, |cx| self.prepaint(cx))
     }
 }
 
@@ -552,7 +568,7 @@ impl Element for AnyElement {
         _: &mut Self::RequestLayoutState,
         cx: &mut WindowContext,
     ) {
-        self.prepaint(cx)
+        self.prepaint(cx);
     }
 
     fn paint(
@@ -563,7 +579,7 @@ impl Element for AnyElement {
         _: &mut Self::PrepaintState,
         cx: &mut WindowContext,
     ) {
-        self.paint(cx)
+        self.paint(cx);
     }
 }
 

crates/gpui/src/elements/div.rs

@@ -1359,6 +1359,9 @@ impl Interactivity {
         f: impl FnOnce(&Style, Point<Pixels>, Option<Hitbox>, &mut WindowContext) -> R,
     ) -> R {
         self.content_size = content_size;
+        if let Some(focus_handle) = self.tracked_focus_handle.as_ref() {
+            cx.set_focus_handle(&focus_handle);
+        }
         cx.with_optional_element_state::<InteractiveElementState, _>(
             global_id,
             |element_state, cx| {
@@ -1998,9 +2001,6 @@ impl Interactivity {
         if let Some(context) = self.key_context.clone() {
             cx.set_key_context(context);
         }
-        if let Some(focus_handle) = self.tracked_focus_handle.as_ref() {
-            cx.set_focus_handle(focus_handle);
-        }
 
         for listener in key_down_listeners {
             cx.on_key_event(move |event: &KeyDownEvent, phase, cx| {

crates/gpui/src/key_dispatch.rs

@@ -92,6 +92,7 @@ pub(crate) struct DispatchNode {
 pub(crate) struct ReusedSubtree {
     old_range: Range<usize>,
     new_range: Range<usize>,
+    contains_focus: bool,
 }
 
 impl ReusedSubtree {
@@ -104,6 +105,10 @@ impl ReusedSubtree {
         );
         DispatchNodeId((node_id.0 - self.old_range.start) + self.new_range.start)
     }
+
+    pub fn contains_focus(&self) -> bool {
+        self.contains_focus
+    }
 }
 
 type KeyListener = Rc<dyn Fn(&dyn Any, DispatchPhase, &mut WindowContext)>;
@@ -246,9 +251,15 @@ impl DispatchTree {
         target.modifiers_changed_listeners = mem::take(&mut source.modifiers_changed_listeners);
     }
 
-    pub fn reuse_subtree(&mut self, old_range: Range<usize>, source: &mut Self) -> ReusedSubtree {
+    pub fn reuse_subtree(
+        &mut self,
+        old_range: Range<usize>,
+        source: &mut Self,
+        focus: Option<FocusId>,
+    ) -> ReusedSubtree {
         let new_range = self.nodes.len()..self.nodes.len() + old_range.len();
 
+        let mut contains_focus = false;
         let mut source_stack = vec![];
         for (source_node_id, source_node) in source
             .nodes
@@ -268,6 +279,9 @@ impl DispatchTree {
             }
 
             source_stack.push(source_node_id);
+            if source_node.focus_id.is_some() && source_node.focus_id == focus {
+                contains_focus = true;
+            }
             self.move_node(source_node);
         }
 
@@ -279,6 +293,7 @@ impl DispatchTree {
         ReusedSubtree {
             old_range,
             new_range,
+            contains_focus,
         }
     }
 

crates/gpui/src/window.rs

@@ -464,6 +464,7 @@ impl Frame {
         self.cursor_styles.clear();
         self.hitboxes.clear();
         self.deferred_draws.clear();
+        self.focus = None;
     }
 
     pub(crate) fn hit_test(&self, position: Point<Pixels>) -> HitTest {
@@ -1460,7 +1461,6 @@ impl<'a> WindowContext<'a> {
                 &mut self.window.rendered_frame.dispatch_tree,
                 self.window.focus,
             );
-        self.window.next_frame.focus = self.window.focus;
         self.window.next_frame.window_active = self.window.active.get();
 
         // Register requested input handler with the platform window.
@@ -1574,7 +1574,7 @@ impl<'a> WindowContext<'a> {
         self.paint_deferred_draws(&sorted_deferred_draws);
 
         if let Some(mut prompt_element) = prompt_element {
-            prompt_element.paint(self)
+            prompt_element.paint(self);
         } else if let Some(mut drag_element) = active_drag_element {
             drag_element.paint(self);
         } else if let Some(mut tooltip_element) = tooltip_element {
@@ -1730,7 +1730,13 @@ impl<'a> WindowContext<'a> {
         let reused_subtree = window.next_frame.dispatch_tree.reuse_subtree(
             range.start.dispatch_tree_index..range.end.dispatch_tree_index,
             &mut window.rendered_frame.dispatch_tree,
+            window.focus,
         );
+
+        if reused_subtree.contains_focus() {
+            window.next_frame.focus = window.focus;
+        }
+
         window.next_frame.deferred_draws.extend(
             window.rendered_frame.deferred_draws
                 [range.start.deferred_draws_index..range.end.deferred_draws_index]
@@ -2845,13 +2851,16 @@ impl<'a> WindowContext<'a> {
     /// Sets the focus handle for the current element. This handle will be used to manage focus state
     /// and keyboard event dispatch for the element.
     ///
-    /// This method should only be called as part of the paint phase of element drawing.
+    /// This method should only be called as part of the prepaint phase of element drawing.
     pub fn set_focus_handle(&mut self, focus_handle: &FocusHandle) {
         debug_assert_eq!(
             self.window.draw_phase,
-            DrawPhase::Paint,
+            DrawPhase::Prepaint,
-            "this method can only be called during paint"
+            "this method can only be called during prepaint"
         );
+        if focus_handle.is_focused(self) {
+            self.window.next_frame.focus = Some(focus_handle.id);
+        }
         self.window
             .next_frame
             .dispatch_tree

crates/markdown/src/markdown.rs

@@ -719,6 +719,9 @@ impl Element for MarkdownElement {
         rendered_markdown: &mut Self::RequestLayoutState,
         cx: &mut WindowContext,
     ) -> Self::PrepaintState {
+        let focus_handle = self.markdown.read(cx).focus_handle.clone();
+        cx.set_focus_handle(&focus_handle);
+
         let hitbox = cx.insert_hitbox(bounds, false);
         rendered_markdown.element.prepaint(cx);
         self.autoscroll(&rendered_markdown.text, cx);
@@ -733,9 +736,6 @@ impl Element for MarkdownElement {
         hitbox: &mut Self::PrepaintState,
         cx: &mut WindowContext,
     ) {
-        let focus_handle = self.markdown.read(cx).focus_handle.clone();
-        cx.set_focus_handle(&focus_handle);
-
         let mut context = KeyContext::default();
         context.add("Markdown");
         cx.set_key_context(context);

crates/multi_buffer/src/multi_buffer.rs

@@ -261,7 +261,7 @@ pub struct ExcerptRange<T> {
 }
 
 #[derive(Clone, Debug, Default)]
-struct ExcerptSummary {
+pub struct ExcerptSummary {
     excerpt_id: ExcerptId,
     /// The location of the last [`Excerpt`] being summarized
     excerpt_locator: Locator,
@@ -3744,6 +3744,21 @@ impl MultiBufferSnapshot {
         Some(&self.excerpt(excerpt_id)?.buffer)
     }
 
+    pub fn range_for_excerpt<'a, T: sum_tree::Dimension<'a, ExcerptSummary>>(
+        &'a self,
+        excerpt_id: ExcerptId,
+    ) -> Option<Range<T>> {
+        let mut cursor = self.excerpts.cursor::<(Option<&Locator>, T)>();
+        let locator = self.excerpt_locator_for_id(excerpt_id);
+        if cursor.seek(&Some(locator), Bias::Left, &()) {
+            let start = cursor.start().1.clone();
+            let end = cursor.end(&()).1;
+            Some(start..end)
+        } else {
+            None
+        }
+    }
+
     fn excerpt(&self, excerpt_id: ExcerptId) -> Option<&Excerpt> {
         let mut cursor = self.excerpts.cursor::<Option<&Locator>>();
         let locator = self.excerpt_locator_for_id(excerpt_id);

crates/repl/src/session.rs

@@ -5,7 +5,7 @@ use crate::{
 use collections::{HashMap, HashSet};
 use editor::{
     display_map::{
-        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
+        BlockContext, BlockDisposition, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
     },
     scroll::Autoscroll,
     Anchor, AnchorRangeExt as _, Editor, MultiBuffer, ToPoint,
@@ -37,7 +37,7 @@ struct EditorBlock {
     editor: WeakView<Editor>,
     code_range: Range<Anchor>,
     invalidation_anchor: Anchor,
-    block_id: BlockId,
+    block_id: CustomBlockId,
     execution_view: View<ExecutionView>,
 }
 
@@ -282,7 +282,7 @@ impl Session {
         if let multi_buffer::Event::Edited { .. } = event {
             let snapshot = buffer.read(cx).snapshot(cx);
 
-            let mut blocks_to_remove: HashSet<BlockId> = HashSet::default();
+            let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
 
             self.blocks.retain(|_id, block| {
                 if block.invalidation_anchor.is_valid(&snapshot) {
@@ -316,7 +316,7 @@ impl Session {
     }
 
     pub fn clear_outputs(&mut self, cx: &mut ViewContext<Self>) {
-        let blocks_to_remove: HashSet<BlockId> =
+        let blocks_to_remove: HashSet<CustomBlockId> =
             self.blocks.values().map(|block| block.block_id).collect();
 
         self.editor
@@ -346,7 +346,7 @@ impl Session {
 
         let message: JupyterMessage = execute_request.into();
 
-        let mut blocks_to_remove: HashSet<BlockId> = HashSet::default();
+        let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
 
         let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
 

crates/semantic_index/src/project_index_debug_view.rs

@@ -258,7 +258,9 @@ impl Render for ProjectIndexDebugView {
                     list.prepaint_as_root(bounds.origin, bounds.size.into(), cx);
                     list
                 },
-                |_, mut list, cx| list.paint(cx),
+                |_, mut list, cx| {
+                    list.paint(cx);
+                },
             )
             .size_full()
             .into_any_element()