From c58ee1553b3c4298566ee8d041939189f7d7be74 Mon Sep 17 00:00:00 2001
From: Ali Abrar <aliabrar@gmail.com>
Date: Tue, 23 Mar 2021 01:14:46 -0400
Subject: [PATCH] Add mapImages to ImageWriter

---
 ChangeLog.md                        |  1 +
 src/Reflex/Vty/Widget.hs            | 23 ++++++++++++---------
 src/Reflex/Vty/Widget/Input/Text.hs |  2 +-
 src/Reflex/Vty/Widget/Layout.hs     | 32 +++++++++++++++++++++--------
 4 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/ChangeLog.md b/ChangeLog.md
index 258f5d1..6f524a7 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -51,6 +51,7 @@
 * _Breaking Change_: `HasVtyInput` now has a method `localInput` for filtering the input a child widget may receive
 * Add `anyChildFocused`, which provides information about whether subwidgets are focused
 * Add `filterKeys`, which is the same as `localInput` but only cares about keyboard events
+* _Breaking Change_: `ImageWriter` now has a method `mapImages` for transforming the images emitted by a child widget
 
 ## 0.1.4.1
 * Migrate to new dependent-sum / dependent-map (after the "some" package split)
diff --git a/src/Reflex/Vty/Widget.hs b/src/Reflex/Vty/Widget.hs
index e369eaf..a3fddea 100644
--- a/src/Reflex/Vty/Widget.hs
+++ b/src/Reflex/Vty/Widget.hs
@@ -132,14 +132,11 @@ class HasVtyWidgetCtx t m | m -> t where
   askCtx :: m (VtyWidgetCtx t)
   default askCtx :: (f m' ~ m, Monad m', MonadTrans f, HasVtyWidgetCtx t m') => m (VtyWidgetCtx t)
   askCtx = lift askCtx
-  localCtx :: (VtyWidgetCtx t -> VtyWidgetCtx t) -> (Behavior t [Image] -> Behavior t [Image]) -> m a -> m a
+  localCtx :: (VtyWidgetCtx t -> VtyWidgetCtx t) -> m a -> m a
 
 instance (Monad m, Reflex t) => HasVtyWidgetCtx t (VtyWidget t m) where
   askCtx = VtyWidget $ lift ask
-  localCtx f g (VtyWidget x) = VtyWidget $ do
-    (a, images) :: (a, Behavior t [Image]) <- lift $ local f $ runBehaviorWriterT x
-    tellImages $ g images
-    pure a
+  localCtx f (VtyWidget x) = VtyWidget $ local f x
 
 -- | Runs a 'VtyWidget' with a given context
 runVtyWidget
@@ -243,9 +240,15 @@ class (Reflex t, Monad m) => ImageWriter t m | m -> t where
   tellImages :: Behavior t [Image] -> m ()
   default tellImages :: (f m' ~ m, Monad m', MonadTrans f, ImageWriter t m') => Behavior t [Image] -> m ()
   tellImages = lift . tellImages
+  -- | Apply a transformation to the images produced by the child actions
+  mapImages :: (Behavior t [Image] -> Behavior t [Image]) -> m a -> m a
 
 instance (Monad m, Reflex t) => ImageWriter t (BehaviorWriterT t [Image] m) where
   tellImages = tellBehavior
+  mapImages f x = do
+    (a, images) <- lift $ runBehaviorWriterT x
+    tellImages $ f images
+    pure a
 
 -- | A chunk of the display area
 data Region = Region
@@ -283,7 +286,7 @@ withinImage (Region left top width height)
 -- * mouse inputs inside the region have their coordinates translated such
 --   that (0,0) is the top-left corner of the region
 pane
-  :: (Reflex t, Monad m, MonadNodeId m, HasVtyWidgetCtx t m, HasVtyInput t m)
+  :: (Reflex t, Monad m, MonadNodeId m, HasVtyWidgetCtx t m, HasVtyInput t m, ImageWriter t m)
   => Dynamic t Region
   -> Dynamic t Bool -- ^ Whether the widget should be focused when the parent is.
   -> m a
@@ -299,7 +302,7 @@ pane dr foc child = do
         attachWith (\(r,f) e -> filterInput r f e)
           (liftA2 (,) reg (current foc))
   let imagesWithinRegion images = liftA2 (\is r -> map (withinImage r) is) images reg
-  localInput inputFilter $ localCtx subContext imagesWithinRegion child
+  mapImages imagesWithinRegion $ localInput inputFilter $ localCtx subContext child
   where
     filterInput :: Region -> Bool -> VtyEvent -> Maybe VtyEvent
     filterInput (Region l t w h) focused e = case e of
@@ -444,7 +447,7 @@ keyCombos ks = do
 
 -- | A plain split of the available space into vertically stacked panes.
 -- No visual separator is built in here.
-splitV :: (Reflex t, Monad m, MonadNodeId m, HasVtyWidgetCtx t m, HasDisplaySize t m, HasVtyInput t m)
+splitV :: (Reflex t, Monad m, MonadNodeId m, HasVtyWidgetCtx t m, HasDisplaySize t m, HasVtyInput t m, ImageWriter t m)
        => Dynamic t (Int -> Int)
        -- ^ Function used to determine size of first pane based on available size
        -> Dynamic t (Bool, Bool)
@@ -465,7 +468,7 @@ splitV sizeFunD focD wA wB = do
 
 -- | A plain split of the available space into horizontally stacked panes.
 -- No visual separator is built in here.
-splitH :: (Reflex t, Monad m, MonadNodeId m, HasDisplaySize t m, HasVtyWidgetCtx t m, HasVtyInput t m)
+splitH :: (Reflex t, Monad m, MonadNodeId m, HasDisplaySize t m, HasVtyWidgetCtx t m, HasVtyInput t m, ImageWriter t m)
        => Dynamic t (Int -> Int)
        -- ^ Function used to determine size of first pane based on available size
        -> Dynamic t (Bool, Bool)
@@ -485,7 +488,7 @@ splitH sizeFunD focD wA wB = do
 -- | A split of the available space into two parts with a draggable separator.
 -- Starts with half the space allocated to each, and the first pane has focus.
 -- Clicking in a pane switches focus.
-splitVDrag :: (Reflex t, MonadFix m, MonadHold t m, MonadNodeId m, HasDisplaySize t m, HasVtyInput t m, HasVtyWidgetCtx t m)
+splitVDrag :: (Reflex t, MonadFix m, MonadHold t m, MonadNodeId m, HasDisplaySize t m, HasVtyInput t m, HasVtyWidgetCtx t m, ImageWriter t m)
   => m ()
   -> m a
   -> m b
diff --git a/src/Reflex/Vty/Widget/Input/Text.hs b/src/Reflex/Vty/Widget/Input/Text.hs
index 7e6f7bb..09b1355 100644
--- a/src/Reflex/Vty/Widget/Input/Text.hs
+++ b/src/Reflex/Vty/Widget/Input/Text.hs
@@ -99,7 +99,7 @@ multilineTextInput cfg = do
 -- the computed line count to greedily size the tile when vertically
 -- oriented, and uses the fallback width when horizontally oriented.
 textInputTile
-  :: (Monad m, MonadNodeId m, Reflex t, MonadFix m, MonadLayout t m, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m)
+  :: (Monad m, MonadNodeId m, Reflex t, MonadFix m, MonadLayout t m, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m, ImageWriter t m)
   => m (TextInput t)
   -> Dynamic t Int
   -> m (TextInput t)
diff --git a/src/Reflex/Vty/Widget/Layout.hs b/src/Reflex/Vty/Widget/Layout.hs
index 2190a3e..0e4ffab 100644
--- a/src/Reflex/Vty/Widget/Layout.hs
+++ b/src/Reflex/Vty/Widget/Layout.hs
@@ -129,14 +129,20 @@ instance (Reflex t, MonadFix m, HasVtyInput t m) => HasVtyInput t (Focus t m) wh
     return a
 
 instance (HasVtyWidgetCtx t m, Reflex t, MonadFix m) => HasVtyWidgetCtx t (Focus t m) where
-  localCtx f g (Focus w) = Focus $ do
+  localCtx f (Focus w) = Focus $ do
     d <- ask
-    ((a, fs), e) <- lift $ lift $ lift $ localCtx f g $ runEventWriterT $ flip runReaderT d $ runDynamicWriterT w
+    ((a, fs), e) <- lift $ lift $ lift $ localCtx f $ runEventWriterT $ flip runReaderT d $ runDynamicWriterT w
     tellEvent e
     tellDyn fs
     return a
 
-instance ImageWriter t m => ImageWriter t (Focus t m)
+instance (ImageWriter t m, MonadFix m) => ImageWriter t (Focus t m) where
+  mapImages f (Focus w) = Focus $ do
+    d <- ask
+    ((a, fs), e) <- lift $ lift $ lift $ mapImages f $ runEventWriterT $ flip runReaderT d $ runDynamicWriterT w
+    tellEvent e
+    tellDyn fs
+    return a
 
 instance (HasFocus t m, Monad m) => HasFocus t (Focus t m)
 
@@ -383,10 +389,10 @@ instance (Adjustable t m, MonadFix m, MonadHold t m) => Adjustable t (Layout t m
   traverseDMapWithKeyWithAdjustWithMove f m e = Layout $ traverseDMapWithKeyWithAdjustWithMove (\k v -> unLayout $ f k v) m e
 
 instance (HasVtyWidgetCtx t m, HasDisplaySize t m, Reflex t, MonadFix m) => HasVtyWidgetCtx t (Layout t m) where
-  localCtx f g x = do
+  localCtx f x = do
     solution <- Layout ask
     let orientation = snd . rootLT <$> solution
-    lift $ localCtx f g $ do
+    lift $ localCtx f $ do
       dw <- displayWidth
       dh <- displayHeight
       let reg = Region 0 0 <$> dw <*> dh
@@ -402,7 +408,15 @@ instance (HasVtyInput t m, HasDisplaySize t m, MonadFix m, Reflex t) => HasVtyIn
       let reg = Region 0 0 <$> dw <*> dh
       runLayout orientation reg x
 
-instance ImageWriter t m => ImageWriter t (Layout t m)
+instance (HasDisplaySize t m, ImageWriter t m, MonadFix m) => ImageWriter t (Layout t m) where
+  mapImages f x = do
+    solution <- Layout ask
+    let orientation = snd . rootLT <$> solution
+    lift $ mapImages f $ do
+      dw <- displayWidth
+      dh <- displayHeight
+      let reg = Region 0 0 <$> dw <*> dh
+      runLayout orientation reg x
 
 instance (HasFocus t m, Monad m) => HasFocus t (Layout t m)
 
@@ -495,7 +509,7 @@ initManager_ = fmap fst . initManager
 -- provided constraint. Returns the 'FocusId' allowing for manual focus
 -- management.
 tile'
-  :: (MonadNodeId m, MonadFix m, Reflex t, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m, MonadLayout t m)
+  :: (MonadNodeId m, MonadFix m, Reflex t, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m, MonadLayout t m, ImageWriter t m)
   => Dynamic t Constraint
   -> m a
   -> m (FocusId, a)
@@ -514,7 +528,7 @@ tile' c w = do
 -- | A widget that is focusable and occupies a layout region based on the
 -- provided constraint.
 tile
-  :: (MonadNodeId m, MonadFix m, Reflex t, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m, MonadLayout t m)
+  :: (MonadNodeId m, MonadFix m, Reflex t, HasVtyWidgetCtx t m, HasVtyInput t m, MonadFocus t m, MonadLayout t m, ImageWriter t m)
   => Dynamic t Constraint
   -> m a
   -> m a
@@ -525,7 +539,7 @@ tile c = fmap snd . tile' c
 -- | A widget that is not focusable and occupies a layout region based on the
 -- provided constraint.
 grout
-  :: (Reflex t, MonadNodeId m, HasVtyWidgetCtx t m, MonadLayout t m, HasVtyInput t m)
+  :: (Reflex t, MonadNodeId m, HasVtyWidgetCtx t m, MonadLayout t m, HasVtyInput t m, ImageWriter t m)
   => Dynamic t Constraint
   -> m a
   -> m a