From b363c81a24dcfd51aa63ca35a9be46291ae9b8fb Mon Sep 17 00:00:00 2001
From: Leo dev <klesti.selimay@icloud.com>
Date: Mon, 4 Aug 2025 12:16:23 -0500
Subject: [PATCH] Added a ElementRenderer trait for element renderers

---
 src/extensions/mod.rs |  6 +++++
 src/lib.rs            | 51 ++++++++++--------------------------------
 src/render_scope.rs   | 52 ++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 69 insertions(+), 40 deletions(-)

diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs
index 423962f..7624606 100644
--- a/src/extensions/mod.rs
+++ b/src/extensions/mod.rs
@@ -119,6 +119,12 @@ impl Context {
         components
     }
 
+    pub fn render_root(&self, scope: &mut RenderScope) {
+        for ext in self.screen.extensions.lock().unwrap().iter() {
+            ext.lock().unwrap().render(self, scope);
+        }
+    }
+
     pub fn render(&self, w: &Arc<Widget>, scope: &mut RenderScope) {
         for ext in self.screen.extensions.lock().unwrap().iter() {
             ext.lock().unwrap().render_widget(self, scope, w);
diff --git a/src/lib.rs b/src/lib.rs
index f3b4a82..3b7597c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -32,8 +32,8 @@
 use std::sync::{Arc, Mutex};
 
 use crate::{
-    extensions::{Extension, Handler},
-    prelude::Context,
+    extensions::Extension,
+    prelude::{Context, ElementRenderer},
     render_scope::RenderScope,
     widget::{BoxedElement, DynWidget, Element, StaticWidget, Widget, WidgetLoad},
 };
@@ -180,47 +180,19 @@ impl Screen {
     /// This method is called internally by `run`.
     pub fn render(self: &Arc<Self>, ctx: &Context) -> std::io::Result<()> {
         let mut scope = RenderScope::new();
+        let mut renderer = ScreenRenderer;
         let (w, h) = crossterm::terminal::size().unwrap();
         scope.set_parent_size(w, h);
-
-        for ext in self.extensions.lock().unwrap().iter() {
-            ext.lock().unwrap().render(ctx, &mut scope);
-        }
+        ctx.render_root(&mut scope);
 
         utils::clear()?;
-        for elem in self.widgets.lock().unwrap().iter() {
-            if elem.get::<NoRender>().is_some() || elem.get::<NoRenderRoot>().is_some() {
-                elem.auto_refresh();
+        for widget in self.widgets.lock().unwrap().iter() {
+            if widget.get::<NoRenderRoot>().is_some() {
+                widget.auto_refresh();
                 continue;
             }
 
-            if let Some(wrapper) = elem.get::<Handler<RenderWrapperEvent>>() {
-                wrapper.call(elem, &RenderWrapperEvent(&mut scope));
-            } else {
-                scope.clear();
-
-                if let Some(style) = elem.get() {
-                    scope.set_style(style);
-                }
-                if let Some(t) = elem.get() {
-                    scope.set_transform(&t);
-                }
-
-                for ext in self.extensions.lock().unwrap().iter() {
-                    ext.lock().unwrap().render_widget(ctx, &mut scope, elem);
-                }
-
-                elem.get_elem().render(&mut scope, ctx);
-
-                if let Some(t) = elem.get() {
-                    scope.set_transform(&t);
-                }
-                scope.draw();
-
-                elem.get_elem().after_render(&mut scope, ctx);
-            }
-
-            elem.auto_refresh();
+            scope.render_widget(&mut renderer, ctx, widget);
         }
         Ok(())
     }
@@ -228,13 +200,14 @@ impl Screen {
     /// Closes the loop and calls `on_close` in the extensions
     pub fn close(self: &Arc<Self>) {
         *self.running.lock().unwrap() = false;
-
         utils::show_cursor().unwrap();
-
         utils::clear().unwrap();
-
         for ext in self.extensions.lock().unwrap().iter() {
             ext.lock().unwrap().on_close();
         }
     }
 }
+
+struct ScreenRenderer;
+
+impl ElementRenderer for ScreenRenderer {}
diff --git a/src/render_scope.rs b/src/render_scope.rs
index 7f8c6f5..107574f 100644
--- a/src/render_scope.rs
+++ b/src/render_scope.rs
@@ -8,13 +8,22 @@
 //!
 //! Used internally by OSUI's layout and rendering system.
 
-use std::fmt::Debug;
+use std::{fmt::Debug, sync::Arc};
 
 use crate::{
+    prelude::{Context, Handler},
     style::{RawTransform, Style, Transform},
     utils::{self, hex_ansi_bg},
+    widget::Widget,
+    NoRender, RenderWrapperEvent,
 };
 
+pub trait ElementRenderer {
+    /// Called right after the `after_render` function is called
+    #[allow(unused)]
+    fn on_after_render(&mut self, scope: &mut RenderScope, widget: &Arc<Widget>) {}
+}
+
 /// Represents a single render instruction.
 #[derive(Clone)]
 enum RenderMethod {
@@ -291,6 +300,47 @@ impl RenderScope {
     pub fn get_style(&mut self) -> &mut Style {
         &mut self.style
     }
+
+    pub fn render_widget(
+        &mut self,
+        renderer: &mut dyn ElementRenderer,
+        ctx: &Context,
+        widget: &std::sync::Arc<crate::widget::Widget>,
+    ) -> bool {
+        if widget.get::<NoRender>().is_some() {
+            widget.auto_refresh();
+            return false;
+        }
+
+        if let Some(wrapper) = widget.get::<Handler<RenderWrapperEvent>>() {
+            wrapper.call(widget, &RenderWrapperEvent(self));
+        } else {
+            self.clear();
+
+            if let Some(style) = widget.get() {
+                self.set_style(style);
+            }
+            if let Some(t) = widget.get() {
+                self.set_transform(&t);
+            }
+
+            widget.get_elem().render(self, ctx);
+            ctx.render(widget, self);
+
+            if let Some(t) = widget.get() {
+                self.set_transform(&t);
+            }
+            self.draw();
+
+            widget.get_elem().after_render(self, ctx);
+            ctx.after_render(widget, self);
+            renderer.on_after_render(self, widget);
+        }
+
+        widget.auto_refresh();
+
+        true
+    }
 }
 
 impl Debug for RenderScope {