From 3b3b1c079aa441c02c9b5d893df63e68c7376487 Mon Sep 17 00:00:00 2001
From: Lucas Pickering <lucas@lucaspickering.me>
Date: Sun, 27 Oct 2024 11:59:39 -0400
Subject: [PATCH] Optimize display of large text bodies

It turns out calculating the width of a text body is expensive for large bodies, because we have to count every grapheme. This change caches the text dimensions so we only have to calculate it when the text changes.

Further progress on #356
---
 CHANGELOG.md                                  |   3 +
 .../tui/src/view/common/template_preview.rs   |  20 ++--
 crates/tui/src/view/common/text_window.rs     | 112 +++++++++++-------
 .../tui/src/view/component/queryable_body.rs  |  67 +++++------
 crates/tui/src/view/component/request_view.rs |  11 +-
 crates/tui/src/view/state.rs                  |  29 +++++
 6 files changed, 155 insertions(+), 87 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 97eed30a..8c37ed7d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 ### Fixed
 
 - Fixed `ignore_certificate_hosts` and `large_body_size` fields not being loaded from config
+- Improve performance of large response bodies [#356](https://github.com/LucasPickering/slumber/issues/356)
+  - This includes disabling prettyification and syntax highlighting on bodies over 1 MB (this size is configurable, via the `large_body_size` [config field](https://slumber.lucaspickering.me/book/api/configuration/index.html))
+  - Loading a large response body should no longer cause the UI to freeze or low framerate
 
 ## [2.2.0] - 2024-10-22
 
diff --git a/crates/tui/src/view/common/template_preview.rs b/crates/tui/src/view/common/template_preview.rs
index de016b6d..76dcd111 100644
--- a/crates/tui/src/view/common/template_preview.rs
+++ b/crates/tui/src/view/common/template_preview.rs
@@ -1,7 +1,7 @@
 use crate::{
     context::TuiContext,
     message::Message,
-    view::{draw::Generate, util::highlight, ViewContext},
+    view::{draw::Generate, state::Identified, util::highlight, ViewContext},
 };
 use ratatui::{
     buffer::Buffer,
@@ -35,7 +35,7 @@ pub struct TemplatePreview {
     /// because it also gets an initial value. There should be effectively zero
     /// contention on the mutex because of the single write, and reads being
     /// single-threaded.
-    text: Arc<Mutex<Text<'static>>>,
+    text: Arc<Mutex<Identified<Text<'static>>>>,
 }
 
 impl TemplatePreview {
@@ -46,7 +46,7 @@ impl TemplatePreview {
     /// unrendered and rendered content.
     pub fn new(template: Template, content_type: Option<ContentType>) -> Self {
         // Calculate raw text
-        let text = highlight::highlight_if(
+        let text: Identified<Text> = highlight::highlight_if(
             content_type,
             // We have to clone the template to detach the lifetime. We're
             // choosing to pay one upfront cost here so we don't have to
@@ -54,7 +54,8 @@ impl TemplatePreview {
             // the template and have this text reference it, but that would be
             // self-referential
             template.display().into_owned().into(),
-        );
+        )
+        .into();
         let text = Arc::new(Mutex::new(text));
 
         // Trigger a task to render the preview and write the answer back into
@@ -78,17 +79,17 @@ impl TemplatePreview {
     /// mutex
     fn calculate_rendered_text(
         chunks: Vec<TemplateChunk>,
-        destination: &Mutex<Text<'static>>,
+        destination: &Mutex<Identified<Text<'static>>>,
         content_type: Option<ContentType>,
     ) {
         let text = TextStitcher::stitch_chunks(&chunks);
         let text = highlight::highlight_if(content_type, text);
         *destination
             .lock()
-            .expect("Template preview text lock is poisoned") = text;
+            .expect("Template preview text lock is poisoned") = text.into();
     }
 
-    pub fn text(&self) -> impl '_ + Deref<Target = Text<'static>> {
+    pub fn text(&self) -> impl '_ + Deref<Target = Identified<Text<'static>>> {
         self.text
             .lock()
             .expect("Template preview text lock is poisoned")
@@ -103,7 +104,8 @@ impl From<Template> for TemplatePreview {
 
 /// Clone internal text. Only call this for small pieces of text
 impl Generate for &TemplatePreview {
-    type Output<'this> =  Text<'this>
+    type Output<'this>
+        = Text<'this>
     where
         Self: 'this;
 
@@ -117,7 +119,7 @@ impl Generate for &TemplatePreview {
 
 impl Widget for &TemplatePreview {
     fn render(self, area: Rect, buf: &mut Buffer) {
-        self.text().deref().render(area, buf)
+        (&**self.text()).render(area, buf)
     }
 }
 
diff --git a/crates/tui/src/view/common/text_window.rs b/crates/tui/src/view/common/text_window.rs
index ffff6d0b..deafd719 100644
--- a/crates/tui/src/view/common/text_window.rs
+++ b/crates/tui/src/view/common/text_window.rs
@@ -5,6 +5,7 @@ use crate::{
         context::UpdateContext,
         draw::{Draw, DrawMetadata},
         event::{Event, EventHandler, Update},
+        state::{Identified, StateCell},
     },
 };
 use ratatui::{
@@ -19,6 +20,7 @@ use ratatui::{
 use slumber_config::Action;
 use std::{cell::Cell, cmp};
 use unicode_width::UnicodeWidthStr;
+use uuid::Uuid;
 
 /// A scrollable (but not editable) block of text. Internal state will be
 /// updated on each render, to adjust to the text's width/height. Generally the
@@ -27,14 +29,13 @@ use unicode_width::UnicodeWidthStr;
 /// (especially if it includes syntax highlighting).
 #[derive(derive_more::Debug, Default)]
 pub struct TextWindow {
+    /// Cache the size of the text window, because it's expensive to calculate.
+    /// Checking the width of a text requires counting all its graphemes.
+    text_size: StateCell<Uuid, TextSize>,
     /// Horizontal scroll
     offset_x: Cell<usize>,
     /// Vertical scroll
     offset_y: Cell<usize>,
-    /// How wide is the full text content?
-    text_width: Cell<usize>,
-    /// How tall is the full text content?
-    text_height: Cell<usize>,
     /// How wide is the visible text area, excluding gutter/scrollbars?
     window_width: Cell<usize>,
     /// How tall is the visible text area, exluding gutter/scrollbars?
@@ -45,7 +46,7 @@ pub struct TextWindow {
 pub struct TextWindowProps<'a> {
     /// Text to render. We take a reference because this component tends to
     /// contain a lot of text, and we don't want to force a clone on render
-    pub text: &'a Text<'a>,
+    pub text: &'a Identified<Text<'a>>,
     /// Extra text to render below the text window
     pub footer: Option<Text<'a>>,
     pub margins: ScrollbarMargins,
@@ -70,21 +71,35 @@ impl Default for ScrollbarMargins {
     }
 }
 
+#[derive(Debug, Default)]
+struct TextSize {
+    /// Number of graphemes in the longest line in the text
+    width: usize,
+    /// Number of lines in the text
+    height: usize,
+}
+
 impl TextWindow {
     /// Get the final line that we can't scroll past. This will be the first
     /// line of the last page of text
     fn max_scroll_line(&self) -> usize {
-        self.text_height
+        let text_height = self
+            .text_size
             .get()
-            .saturating_sub(self.window_height.get())
+            .map(|state| state.height)
+            .unwrap_or_default();
+        text_height.saturating_sub(self.window_height.get())
     }
 
     /// Get the final column that we can't scroll (horizontally) past. This will
     /// be the left edge of the rightmost "page" of text
     fn max_scroll_column(&self) -> usize {
-        self.text_width
+        let text_width = self
+            .text_size
             .get()
-            .saturating_sub(self.window_width.get())
+            .map(|state| state.width)
+            .unwrap_or_default();
+        text_width.saturating_sub(self.window_width.get())
     }
 
     fn scroll_up(&mut self, lines: usize) {
@@ -137,6 +152,11 @@ impl TextWindow {
             .take(self.window_height.get())
             .enumerate();
         for (y, line) in lines {
+            // This could be expensive if we're skipping a lot of graphemes,
+            // i.e. scrolled far to the right in a wide body. Fortunately that's
+            // a niche use case so not optimized for yet. To fix this we would
+            // have to map grapheme number -> byte offset and cache that,
+            // because skipping bytes is O(1) instead of O(n)
             let graphemes = line
                 .styled_graphemes(Style::default())
                 .skip(self.offset_x.get())
@@ -185,31 +205,39 @@ impl<'a> Draw<TextWindowProps<'a>> for TextWindow {
     ) {
         let styles = &TuiContext::get().styles;
 
-        // Assume no line wrapping when calculating line count
-        // Note: Paragraph has methods for this, but that requires an owned copy
-        // of Text, which involves a lot of cloning
-        let text_height = props.text.lines.len();
-        let text_width = props
-            .text
-            .lines
-            .iter()
-            .map(Line::width)
-            .max()
-            .unwrap_or_default();
+        let text_state = self.text_size.get_or_update(&props.text.id(), || {
+            // Note: Paragraph has methods for this, but that requires an
+            // owned copy of Text, which involves a lot of cloning
+
+            // This counts _graphemes_, not bytes, so it's O(n)
+            let text_width = props
+                .text
+                .lines
+                .iter()
+                .map(Line::width)
+                .max()
+                .unwrap_or_default();
+            // Assume no line wrapping when calculating line count
+            let text_height = props.text.lines.len();
+            TextSize {
+                width: text_width,
+                height: text_height,
+            }
+        });
 
         let [gutter_area, _, text_area] = Layout::horizontal([
             // Size gutter based on width of max line number
-            Constraint::Length((text_height as f32).log10().floor() as u16 + 1),
+            Constraint::Length(
+                (text_state.height as f32).log10().floor() as u16 + 1,
+            ),
             Constraint::Length(1), // Spacer
             Constraint::Min(0),
         ])
         .areas(metadata.area());
-        let has_vertical_scroll = text_height > text_area.height as usize;
-        let has_horizontal_scroll = text_width > text_area.width as usize;
+        let has_vertical_scroll = text_state.height > text_area.height as usize;
+        let has_horizontal_scroll = text_state.width > text_area.width as usize;
 
         // Store text and window sizes for calculations in the update code
-        self.text_width.set(text_width);
-        self.text_height.set(text_height);
         self.window_width.set(text_area.width as usize);
         self.window_height.set(text_area.height as usize);
 
@@ -218,8 +246,10 @@ impl<'a> Draw<TextWindowProps<'a>> for TextWindow {
 
         // Draw line numbers in the gutter
         let first_line = self.offset_y.get() + 1;
-        let last_line =
-            cmp::min(first_line + self.window_height.get() - 1, text_height);
+        let last_line = cmp::min(
+            first_line + self.window_height.get() - 1,
+            text_state.height,
+        );
         frame.render_widget(
             Paragraph::new(
                 (first_line..=last_line)
@@ -244,7 +274,7 @@ impl<'a> Draw<TextWindowProps<'a>> for TextWindow {
                     x: text_area.x,
                     y: text_area.y
                         + (cmp::min(
-                            self.text_height.get(),
+                            text_state.height,
                             self.window_height.get(),
                         )) as u16,
                     width: text_area.width,
@@ -257,7 +287,7 @@ impl<'a> Draw<TextWindowProps<'a>> for TextWindow {
         if has_vertical_scroll {
             frame.render_widget(
                 Scrollbar {
-                    content_length: self.text_height.get(),
+                    content_length: text_state.height,
                     offset: self.offset_y.get(),
                     // We substracted the margin from the text area before, so
                     // we have to add that back now
@@ -270,7 +300,7 @@ impl<'a> Draw<TextWindowProps<'a>> for TextWindow {
         if has_horizontal_scroll {
             frame.render_widget(
                 Scrollbar {
-                    content_length: self.text_width.get(),
+                    content_length: text_state.width,
                     offset: self.offset_x.get(),
                     orientation: ScrollbarOrientation::HorizontalBottom,
                     // See note on other scrollbar for +1
@@ -299,7 +329,8 @@ mod tests {
         harness: TestHarness,
     ) {
         let text =
-            Text::from("line 1\nline 2 is longer\nline 3\nline 4\nline 5");
+            Text::from("line 1\nline 2 is longer\nline 3\nline 4\nline 5")
+                .into();
         let mut component = TestComponent::new(
             &harness,
             &terminal,
@@ -383,7 +414,8 @@ mod tests {
         #[with(35, 3)] terminal: TestTerminal,
         harness: TestHarness,
     ) {
-        let text = Text::from("intro\n💚💙💜 this is a longer line\noutro");
+        let text =
+            Text::from("intro\n💚💙💜 this is a longer line\noutro").into();
         TestComponent::new(
             &harness,
             &terminal,
@@ -410,7 +442,7 @@ mod tests {
         #[with(10, 2)] terminal: TestTerminal,
         harness: TestHarness,
     ) {
-        let text = Text::from("💚💙💜💚💙💜");
+        let text = Text::raw("💚💙💜💚💙💜").into();
         TestComponent::new(
             &harness,
             &terminal,
@@ -438,9 +470,9 @@ mod tests {
         #[with(10, 3)] terminal: TestTerminal,
         harness: TestHarness,
     ) {
-        let text = ["1 this is a long line", "2", "3", "4", "5"]
-            .join("\n")
-            .into();
+        let text =
+            Text::from_iter(["1 this is a long line", "2", "3", "4", "5"])
+                .into();
         let mut component = TestComponent::new(
             &harness,
             &terminal,
@@ -462,7 +494,7 @@ mod tests {
         assert_eq!(component.data().offset_x.get(), 10);
         assert_eq!(component.data().offset_y.get(), 2);
 
-        let text = ["1 less long line", "2", "3", "4"].join("\n").into();
+        let text = Text::from_iter(["1 less long line", "2", "3", "4"]).into();
         component.set_props(TextWindowProps {
             text: &text,
             margins: ScrollbarMargins {
@@ -481,9 +513,9 @@ mod tests {
     /// automatically be clamped to match
     #[rstest]
     fn test_grow_window(terminal: TestTerminal, harness: TestHarness) {
-        let text = ["1 this is a long line", "2", "3", "4", "5"]
-            .join("\n")
-            .into();
+        let text =
+            Text::from_iter(["1 this is a long line", "2", "3", "4", "5"])
+                .into();
         let mut component = TestComponent::new(
             &harness,
             &terminal,
diff --git a/crates/tui/src/view/component/queryable_body.rs b/crates/tui/src/view/component/queryable_body.rs
index 34b0c2a2..7c67b2ac 100644
--- a/crates/tui/src/view/component/queryable_body.rs
+++ b/crates/tui/src/view/component/queryable_body.rs
@@ -10,7 +10,7 @@ use crate::{
         context::UpdateContext,
         draw::{Draw, DrawMetadata},
         event::{Child, Event, EventHandler, Update},
-        state::StateCell,
+        state::{Identified, StateCell},
         util::{highlight, str_to_text},
         Component, ViewContext,
     },
@@ -37,7 +37,7 @@ use std::cell::Cell;
 pub struct QueryableBody {
     /// Visible text state. This needs to be in a cell because it's initialized
     /// from the body passed in via props
-    filtered_text: StateCell<Option<Query>, Text<'static>>,
+    filtered_text: StateCell<Option<Query>, Identified<Text<'static>>>,
     /// Store whether the body can be queried. True only if it's a recognized
     /// and parsed format
     query_available: Cell<bool>,
@@ -224,42 +224,43 @@ fn init_text(
     content_type: Option<ContentType>,
     body: &ResponseBody,
     query: Option<&Query>,
-) -> Text<'static> {
-    // For bodies over the "large" size, skip prettification and
-    // highlighting because it's slow. We could try to push this work
-    // into a background thread instead, but there's way to kill those
-    // threads so we could end up piling up a lot of work. It also burns
-    // a lot of CPU, regardless of where it's run
-    //
-    // We don't show a hint to the user in this case because it's not
-    // worth the screen real estate
-    if TuiContext::get().config.http.is_large(body.size()) {
-        return if let Some(text) = body.text() {
+) -> Identified<Text<'static>> {
+    let text = if TuiContext::get().config.http.is_large(body.size()) {
+        // For bodies over the "large" size, skip prettification and
+        // highlighting because it's slow. We could try to push this work
+        // into a background thread instead, but there's way to kill those
+        // threads so we could end up piling up a lot of work. It also burns
+        // a lot of CPU, regardless of where it's run
+        //
+        // We don't show a hint to the user in this case because it's not
+        // worth the screen real estate
+        if let Some(text) = body.text() {
             str_to_text(text)
         } else {
             // Showing binary content is a bit of a novelty, there's not much
             // value in it. For large bodies it's not worth the CPU cycles
             "<binary>".into()
-        };
-    }
-
-    // Query and prettify text if possible. This involves a lot of cloning
-    // because it makes stuff easier. If it becomes a bottleneck on large
-    // responses it's fixable.
-    let body = body
-        .parsed()
-        .map(|parsed_body| {
-            // Body is a known content type so we parsed it - apply a query if
-            // necessary and prettify the output
-            query
-                .map(|query| query.query_content(parsed_body).prettify())
-                .unwrap_or_else(|| parsed_body.prettify())
-        })
-        // Content couldn't be parsed, fall back to the raw text
-        // If the text isn't UTF-8, we'll show a placeholder instead
-        .unwrap_or_else(|| format!("{:#}", MaybeStr(body.bytes())));
-    // Apply syntax highlighting
-    highlight::highlight_if(content_type, body.into())
+        }
+    } else {
+        // Query and prettify text if possible. This involves a lot of cloning
+        // because it makes stuff easier. If it becomes a bottleneck on large
+        // responses it's fixable.
+        let body = body
+            .parsed()
+            .map(|parsed_body| {
+                // Body is a known content type so we parsed it - apply a query
+                // if necessary and prettify the output
+                query
+                    .map(|query| query.query_content(parsed_body).prettify())
+                    .unwrap_or_else(|| parsed_body.prettify())
+            })
+            // Content couldn't be parsed, fall back to the raw text
+            // If the text isn't UTF-8, we'll show a placeholder instead
+            .unwrap_or_else(|| format!("{:#}", MaybeStr(body.bytes())));
+        // Apply syntax highlighting
+        highlight::highlight_if(content_type, body.into())
+    };
+    text.into()
 }
 
 #[cfg(test)]
diff --git a/crates/tui/src/view/component/request_view.rs b/crates/tui/src/view/component/request_view.rs
index f072ce16..18336902 100644
--- a/crates/tui/src/view/component/request_view.rs
+++ b/crates/tui/src/view/component/request_view.rs
@@ -10,7 +10,7 @@ use crate::{
         context::UpdateContext,
         draw::{Draw, DrawMetadata, Generate, ToStringGenerate},
         event::{Child, Event, EventHandler, Update},
-        state::StateCell,
+        state::{Identified, StateCell},
         util::highlight,
         Component, ViewContext,
     },
@@ -44,7 +44,7 @@ struct State {
     request: Arc<RequestRecord>,
     /// Persist the visible body, because it may vary from the actual body.
     /// `None` iff the request has no body
-    body: Option<Text<'static>>,
+    body: Option<Identified<Text<'static>>>,
 }
 
 /// Items in the actions popup menu
@@ -157,7 +157,7 @@ impl Draw<RequestViewProps> for RequestView {
 
 /// Calculate body text, including syntax highlighting. We have to clone the
 /// body to prevent a self-reference
-fn init_body(request: &RequestRecord) -> Option<Text<'static>> {
+fn init_body(request: &RequestRecord) -> Option<Identified<Text<'static>>> {
     let content_type = ContentType::from_headers(&request.headers).ok();
     request
         .body()
@@ -166,6 +166,7 @@ fn init_body(request: &RequestRecord) -> Option<Text<'static>> {
                 content_type,
                 format!("{:#}", MaybeStr(body)).into(),
             )
+            .into()
         })
         .or_else(|| {
             // No body available: check if it's because the recipe has no body,
@@ -179,11 +180,11 @@ fn init_body(request: &RequestRecord) -> Option<Text<'static>> {
                 let config = &TuiContext::get().config;
 
                 Some(
-                    format!(
+                    Text::raw(format!(
                         "Body not available. Streamed bodies, or bodies over \
                         {}, are not persisted",
                         format_byte_size(config.http.large_body_size)
-                    )
+                    ))
                     .into(),
                 )
             } else {
diff --git a/crates/tui/src/view/state.rs b/crates/tui/src/view/state.rs
index f88cd6b5..4f84ea7b 100644
--- a/crates/tui/src/view/state.rs
+++ b/crates/tui/src/view/state.rs
@@ -6,6 +6,7 @@ pub mod select;
 use chrono::{DateTime, Utc};
 use derive_more::Deref;
 use std::cell::{Ref, RefCell};
+use uuid::Uuid;
 
 /// An internally mutable cell for UI state. Certain state needs to be updated
 /// during the draw phase, typically because it's derived from parent data
@@ -84,6 +85,34 @@ impl<K, V> Default for StateCell<K, V> {
     }
 }
 
+/// A uniquely identified immutable value. Useful for detecting changes in
+/// values that are expensive to do full comparisons on.
+#[derive(Debug, Deref)]
+pub struct Identified<T> {
+    id: Uuid,
+    #[deref]
+    value: T,
+}
+
+impl<T> Identified<T> {
+    pub fn new(value: T) -> Self {
+        Self {
+            id: Uuid::new_v4(),
+            value,
+        }
+    }
+
+    pub fn id(&self) -> Uuid {
+        self.id
+    }
+}
+
+impl<T> From<T> for Identified<T> {
+    fn from(value: T) -> Self {
+        Self::new(value)
+    }
+}
+
 /// A notification is an ephemeral informational message generated by some async
 /// action. It doesn't grab focus, but will be useful to the user nonetheless.
 /// It should be shown for a short period of time, then disappear on its own.