From 9371754942d16ea00e08c00513e52a7a196f0f46 Mon Sep 17 00:00:00 2001
From: Nathan Sobo <nathan@zed.dev>
Date: Mon, 28 Aug 2023 15:13:59 -0600
Subject: [PATCH] Provide themes to subtrees via context

Co-Authored-By: Mikayla Maki <mikayla@zed.dev>
---
 crates/gpui/playground/src/div.rs           |  3 +-
 crates/gpui/playground/src/element.rs       | 15 +++-
 crates/gpui/playground/src/paint_context.rs | 19 -----
 crates/gpui/playground/src/playground.rs    | 80 ++++++++++++++-------
 crates/gpui/playground/src/style.rs         | 32 +++++++++
 crates/gpui/playground/src/themes.rs        | 63 +++++++++++++++-
 crates/gpui/src/app.rs                      | 27 +------
 crates/gpui/src/app/window.rs               | 41 ++++++++++-
 8 files changed, 206 insertions(+), 74 deletions(-)

diff --git a/crates/gpui/playground/src/div.rs b/crates/gpui/playground/src/div.rs
index ca71ede6d4..79d539b221 100644
--- a/crates/gpui/playground/src/div.rs
+++ b/crates/gpui/playground/src/div.rs
@@ -65,7 +65,8 @@ impl<V: 'static> Element<V> for Div<V> {
     {
         let style = &self.computed_style();
         let pop_text_style = style.text_style().map_or(false, |style| {
-            cx.push_text_style(cx.text_style().clone().refined(&style));
+            let style = cx.text_style().clone().refined(&style);
+            cx.push_text_style(style);
             true
         });
         style.paint_background(layout.bounds, cx);
diff --git a/crates/gpui/playground/src/element.rs b/crates/gpui/playground/src/element.rs
index 741ba68317..2ebe07f5fd 100644
--- a/crates/gpui/playground/src/element.rs
+++ b/crates/gpui/playground/src/element.rs
@@ -1,5 +1,8 @@
+use std::marker::PhantomData;
+
 pub use crate::layout_context::LayoutContext;
 pub use crate::paint_context::PaintContext;
+use crate::themes::{Theme, Themed};
 use anyhow::Result;
 use gpui::geometry::vector::Vector2F;
 pub use gpui::{Layout, LayoutId};
@@ -34,6 +37,17 @@ pub trait Element<V: 'static>: 'static {
             phase: ElementPhase::Init,
         }))
     }
+
+    fn themed(self, theme: Theme) -> Themed<V, Self>
+    where
+        Self: Sized,
+    {
+        crate::themes::Themed {
+            child: self,
+            theme,
+            view_type: PhantomData,
+        }
+    }
 }
 
 /// Used to make ElementState<V, E> into a trait object, so we can wrap it in AnyElement<V>.
@@ -85,7 +99,6 @@ impl<V, E: Element<V>> AnyStatefulElement<V> for StatefulElement<V, E> {
                 ElementPhase::Error(message)
             }
         };
-
         result
     }
 
diff --git a/crates/gpui/playground/src/paint_context.rs b/crates/gpui/playground/src/paint_context.rs
index 9cd59143cf..0760b13cd3 100644
--- a/crates/gpui/playground/src/paint_context.rs
+++ b/crates/gpui/playground/src/paint_context.rs
@@ -3,7 +3,6 @@ use derive_more::{Deref, DerefMut};
 pub use gpui::taffy::tree::NodeId;
 use gpui::{
     scene::EventHandler, EventContext, Layout, LayoutId, PaintContext as LegacyPaintContext,
-    RenderContext, ViewContext,
 };
 use std::{any::TypeId, rc::Rc};
 
@@ -15,24 +14,6 @@ pub struct PaintContext<'a, 'b, 'c, 'd, V> {
     pub(crate) scene: &'d mut gpui::SceneBuilder,
 }
 
-impl<'a, 'b, V> RenderContext<'a, 'b, V> for PaintContext<'a, 'b, '_, '_, V> {
-    fn text_style(&self) -> gpui::fonts::TextStyle {
-        self.legacy_cx.text_style()
-    }
-
-    fn push_text_style(&mut self, style: gpui::fonts::TextStyle) {
-        self.legacy_cx.push_text_style(style)
-    }
-
-    fn pop_text_style(&mut self) {
-        self.legacy_cx.pop_text_style()
-    }
-
-    fn as_view_context(&mut self) -> &mut ViewContext<'a, 'b, V> {
-        &mut self.view_context
-    }
-}
-
 impl<'a, 'b, 'c, 'd, V: 'static> PaintContext<'a, 'b, 'c, 'd, V> {
     pub fn new(
         legacy_cx: &'d mut LegacyPaintContext<'a, 'b, 'c, V>,
diff --git a/crates/gpui/playground/src/playground.rs b/crates/gpui/playground/src/playground.rs
index 58b9f9438f..59c9113fd4 100644
--- a/crates/gpui/playground/src/playground.rs
+++ b/crates/gpui/playground/src/playground.rs
@@ -1,16 +1,15 @@
 #![allow(dead_code, unused_variables)]
-use crate::{
-    color::black, element::ParentElement, style::StyleHelpers, themes::rose_pine::RosePinePalette,
-};
-use element::Element;
+use crate::{element::ParentElement, style::StyleHelpers};
+use element::{Element, IntoElement};
 use gpui::{
     geometry::{pixels, rect::RectF, vector::vec2f},
     platform::WindowOptions,
     ViewContext,
 };
 use log::LevelFilter;
+use playground_macros::Element;
 use simplelog::SimpleLogger;
-use themes::{rose_pine, ThemeColors};
+use themes::{current_theme, rose_pine, Theme, ThemeColors};
 use view::view;
 
 mod adapter;
@@ -41,30 +40,63 @@ fn main() {
                 center: true,
                 ..Default::default()
             },
-            |_| view(|cx| workspace(&rose_pine::moon(), cx)),
+            |_| {
+                view(|cx| {
+                    playground(Theme {
+                        colors: rose_pine::dawn(),
+                    })
+                })
+            },
         );
         cx.platform().activate(true);
     });
 }
 
-fn playground<V: 'static>(theme: &ThemeColors) -> impl Element<V> {
-    use div::div;
-    let p = RosePinePalette::dawn();
-
-    div()
-        .text_color(black())
-        .h_full()
-        .w_full()
-        .fill(p.rose)
-        .child(div().fill(p.pine).child(div().fill(p.love).w_6().h_3()))
-        .child(div().fill(p.gold).child(div().fill(p.iris).w_3().h_3()))
+fn playground<V: 'static>(theme: Theme) -> impl Element<V> {
+    workspace().themed(theme)
 }
 
-fn workspace<V: 'static>(theme: &ThemeColors, cx: &mut ViewContext<V>) -> impl Element<V> {
-    use div::div;
-    // one line change1!
-    div()
-        .full()
-        .fill(theme.base(0.5))
-        .child(div().h(pixels(cx.titlebar_height())).fill(theme.base(0.)))
+fn workspace<V: 'static>() -> impl Element<V> {
+    WorkspaceElement
+}
+
+use crate as playground;
+#[derive(Element)]
+struct WorkspaceElement;
+
+impl WorkspaceElement {
+    fn render<V: 'static>(&mut self, _: &mut V, cx: &mut ViewContext<V>) -> impl Element<V> {
+        use div::div;
+        let theme = &cx.theme::<Theme>().colors;
+        // one line change1!
+        div()
+            .full()
+            .flex()
+            .flex_col()
+            .fill(theme.base(0.5))
+            .child(self.title_bar(cx))
+            .child(self.stage(cx))
+            .child(self.status_bar(cx))
+    }
+
+    fn title_bar<V: 'static>(&mut self, cx: &mut ViewContext<V>) -> impl IntoElement<V> {
+        use div::div;
+
+        let theme = &current_theme(cx).colors;
+        div().h(pixels(cx.titlebar_height())).fill(theme.base(0.))
+    }
+
+    fn status_bar<V: 'static>(&mut self, cx: &mut ViewContext<V>) -> impl IntoElement<V> {
+        use div::div;
+
+        let theme = &current_theme(cx).colors;
+        div().h(pixels(cx.titlebar_height())).fill(theme.base(0.))
+    }
+
+    fn stage<V: 'static>(&mut self, cx: &mut ViewContext<V>) -> impl IntoElement<V> {
+        use div::div;
+
+        let theme = &current_theme(cx).colors;
+        div().flex_grow()
+    }
 }
diff --git a/crates/gpui/playground/src/style.rs b/crates/gpui/playground/src/style.rs
index f01f269aa6..d2f639fdc2 100644
--- a/crates/gpui/playground/src/style.rs
+++ b/crates/gpui/playground/src/style.rs
@@ -327,6 +327,38 @@ pub trait StyleHelpers: Styleable<Style = Style> {
         self
     }
 
+    fn flex(mut self) -> Self
+    where
+        Self: Sized,
+    {
+        self.declared_style().display = Some(Display::Flex);
+        self
+    }
+
+    fn flex_col(mut self) -> Self
+    where
+        Self: Sized,
+    {
+        self.declared_style().flex_direction = Some(FlexDirection::Column);
+        self
+    }
+
+    fn flex_row(mut self) -> Self
+    where
+        Self: Sized,
+    {
+        self.declared_style().flex_direction = Some(FlexDirection::Row);
+        self
+    }
+
+    fn flex_grow(mut self) -> Self
+    where
+        Self: Sized,
+    {
+        self.declared_style().flex_grow = Some(1.);
+        self
+    }
+
     fn fill<F>(mut self, fill: F) -> Self
     where
         F: Into<Fill>,
diff --git a/crates/gpui/playground/src/themes.rs b/crates/gpui/playground/src/themes.rs
index 320d2c5813..ce57c8c147 100644
--- a/crates/gpui/playground/src/themes.rs
+++ b/crates/gpui/playground/src/themes.rs
@@ -1,8 +1,23 @@
-use crate::color::{Hsla, Lerp};
-use std::ops::Range;
+use crate::{
+    color::{Hsla, Lerp},
+    element::{Element, PaintContext},
+    layout_context::LayoutContext,
+};
+use gpui::{AppContext, WindowContext};
+use std::{marker::PhantomData, ops::Range};
 
 pub mod rose_pine;
 
+#[derive(Clone, Debug)]
+pub struct Theme {
+    pub colors: ThemeColors,
+}
+
+pub fn current_theme<'a>(cx: &'a WindowContext) -> &'a Theme {
+    cx.theme::<Theme>()
+}
+
+#[derive(Clone, Debug)]
 pub struct ThemeColors {
     pub base: Range<Hsla>,
     pub surface: Range<Hsla>,
@@ -22,6 +37,12 @@ pub struct ThemeColors {
 }
 
 impl ThemeColors {
+    fn current(cx: &AppContext) -> &Self {
+        cx.global::<Vec<Self>>()
+            .last()
+            .expect("must call within a theme provider")
+    }
+
     pub fn base(&self, level: f32) -> Hsla {
         self.base.lerp(level)
     }
@@ -82,3 +103,41 @@ impl ThemeColors {
         self.modified.lerp(level)
     }
 }
+
+pub struct Themed<V: 'static, E> {
+    pub(crate) theme: Theme,
+    pub(crate) child: E,
+    pub(crate) view_type: PhantomData<V>,
+}
+
+impl<V: 'static, E: Element<V>> Element<V> for Themed<V, E> {
+    type PaintState = E::PaintState;
+
+    fn layout(
+        &mut self,
+        view: &mut V,
+        cx: &mut LayoutContext<V>,
+    ) -> anyhow::Result<(gpui::LayoutId, Self::PaintState)>
+    where
+        Self: Sized,
+    {
+        cx.push_theme(self.theme.clone());
+        let result = self.child.layout(view, cx);
+        cx.pop_theme();
+        result
+    }
+
+    fn paint(
+        &mut self,
+        view: &mut V,
+        layout: &gpui::Layout,
+        state: &mut Self::PaintState,
+        cx: &mut PaintContext<V>,
+    ) where
+        Self: Sized,
+    {
+        cx.push_theme(self.theme.clone());
+        self.child.paint(view, layout, state, cx);
+        cx.pop_theme();
+    }
+}
diff --git a/crates/gpui/src/app.rs b/crates/gpui/src/app.rs
index a2536e0b4a..4750c71ea2 100644
--- a/crates/gpui/src/app.rs
+++ b/crates/gpui/src/app.rs
@@ -3619,36 +3619,11 @@ impl<V> BorrowWindowContext for LayoutContext<'_, '_, '_, V> {
 
 pub struct PaintContext<'a, 'b, 'c, V> {
     pub view_context: &'c mut ViewContext<'a, 'b, V>,
-    text_style_stack: Vec<TextStyle>,
 }
 
 impl<'a, 'b, 'c, V> PaintContext<'a, 'b, 'c, V> {
     pub fn new(view_context: &'c mut ViewContext<'a, 'b, V>) -> Self {
-        Self {
-            view_context,
-            text_style_stack: Vec::new(),
-        }
-    }
-}
-
-impl<'a, 'b, 'c, V> RenderContext<'a, 'b, V> for PaintContext<'a, 'b, 'c, V> {
-    fn text_style(&self) -> TextStyle {
-        self.text_style_stack
-            .last()
-            .cloned()
-            .unwrap_or(TextStyle::default(&self.font_cache))
-    }
-
-    fn push_text_style(&mut self, style: TextStyle) {
-        self.text_style_stack.push(style);
-    }
-
-    fn pop_text_style(&mut self) {
-        self.text_style_stack.pop();
-    }
-
-    fn as_view_context(&mut self) -> &mut ViewContext<'a, 'b, V> {
-        &mut self.view_context
+        Self { view_context }
     }
 }
 
diff --git a/crates/gpui/src/app/window.rs b/crates/gpui/src/app/window.rs
index e1492eda99..e1811b6655 100644
--- a/crates/gpui/src/app/window.rs
+++ b/crates/gpui/src/app/window.rs
@@ -1,5 +1,6 @@
 use crate::{
     elements::AnyRootElement,
+    fonts::TextStyle,
     geometry::{rect::RectF, Size},
     json::ToJson,
     keymap_matcher::{Binding, KeymapContext, Keystroke, MatchResult},
@@ -30,7 +31,7 @@ use sqlez::{
     statement::Statement,
 };
 use std::{
-    any::TypeId,
+    any::{type_name, Any, TypeId},
     mem,
     ops::{Deref, DerefMut, Range, Sub},
 };
@@ -53,6 +54,8 @@ pub struct Window {
     pub(crate) invalidation: Option<WindowInvalidation>,
     pub(crate) platform_window: Box<dyn platform::Window>,
     pub(crate) rendered_views: HashMap<usize, Box<dyn AnyRootElement>>,
+    pub(crate) text_style_stack: Vec<TextStyle>,
+    pub(crate) theme_stack: Vec<Box<dyn Any>>,
     titlebar_height: f32,
     appearance: Appearance,
     cursor_regions: Vec<CursorRegion>,
@@ -100,6 +103,8 @@ impl Window {
             clicked_region: None,
             titlebar_height,
             appearance,
+            text_style_stack: Vec::new(),
+            theme_stack: Vec::new(),
         };
 
         let mut window_context = WindowContext::mutable(cx, &mut window, handle);
@@ -1265,6 +1270,40 @@ impl<'a> WindowContext<'a> {
         };
         handle
     }
+
+    pub fn text_style(&self) -> TextStyle {
+        self.window
+            .text_style_stack
+            .last()
+            .cloned()
+            .unwrap_or(TextStyle::default(&self.font_cache))
+    }
+
+    pub fn push_text_style(&mut self, style: TextStyle) {
+        self.window.text_style_stack.push(style);
+    }
+
+    pub fn pop_text_style(&mut self) {
+        self.window.text_style_stack.pop();
+    }
+
+    pub fn theme<T: 'static>(&self) -> &T {
+        self.window
+            .theme_stack
+            .iter()
+            .rev()
+            .find_map(|theme| theme.downcast_ref())
+            .ok_or_else(|| anyhow!("no theme provided of type {}", type_name::<T>()))
+            .unwrap()
+    }
+
+    pub fn push_theme<T: 'static>(&mut self, theme: T) {
+        self.window.theme_stack.push(Box::new(theme));
+    }
+
+    pub fn pop_theme(&mut self) {
+        self.window.theme_stack.pop();
+    }
 }
 
 #[derive(Default)]