From 16b4702dce8282a2756571de58b69df2d7a793f4 Mon Sep 17 00:00:00 2001
From: Stephen Morgan <morgan.stephen99@gmail.com>
Date: Fri, 27 Aug 2021 12:43:37 +1000
Subject: [PATCH] fix: ui: Ensure that forecast_ argument gets restored to the
 startup opts when toggling.

---
 hledger-ui/Hledger/UI/ErrorScreen.hs | 15 ---------------
 hledger-ui/Hledger/UI/UIState.hs     | 21 ++++++++++++++++++---
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/hledger-ui/Hledger/UI/ErrorScreen.hs b/hledger-ui/Hledger/UI/ErrorScreen.hs
index 5f661f671..60bd45441 100644
--- a/hledger-ui/Hledger/UI/ErrorScreen.hs
+++ b/hledger-ui/Hledger/UI/ErrorScreen.hs
@@ -28,7 +28,6 @@ import Hledger.UI.UITypes
 import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.Editor
-import Data.Foldable (asum)
 
 errorScreen :: Screen
 errorScreen = ErrorScreen{
@@ -178,20 +177,6 @@ uiReloadJournalIfChanged copts d j ui = do
         UIState{aScreen=s@ErrorScreen{}} -> ui{aScreen=s{esError=err}}
         _                                -> screenEnter d errorScreen{esError=err} ui
 
--- | Ensure this CliOpts enables forecasted transactions.
--- If a forecast period was specified in the old CliOpts,
--- or in the provided UIState's startup options,
--- it is preserved.
-enableForecastPreservingPeriod :: UIState -> CliOpts -> CliOpts
-enableForecastPreservingPeriod ui copts@CliOpts{inputopts_=iopts} =
-    copts{inputopts_=iopts{forecast_=mforecast}}
-  where
-    mforecast = asum [mprovidedforecastperiod, mstartupforecastperiod, mdefaultforecastperiod]
-      where
-        mprovidedforecastperiod = forecast_ $ inputopts_ copts
-        mstartupforecastperiod  = forecast_ $ inputopts_ $ cliopts_ $ astartupopts ui
-        mdefaultforecastperiod  = Just nulldatespan
-
 -- Re-check any balance assertions in the current journal, and if any
 -- fail, enter (or update) the error screen. Or if balance assertions
 -- are disabled, do nothing.
diff --git a/hledger-ui/Hledger/UI/UIState.hs b/hledger-ui/Hledger/UI/UIState.hs
index 27817a0f8..1c4470fda 100644
--- a/hledger-ui/Hledger/UI/UIState.hs
+++ b/hledger-ui/Hledger/UI/UIState.hs
@@ -7,6 +7,7 @@ module Hledger.UI.UIState
 where
 
 import Brick.Widgets.Edit
+import Data.Foldable (asum)
 import Data.List ((\\), foldl', sort)
 import Data.Semigroup (Max(..))
 import qualified Data.Text as T
@@ -156,11 +157,25 @@ toggleHistorical ui@UIState{aopts=uopts@UIOpts{cliopts_=copts@CliOpts{reportspec
 -- (which are usually but not necessarily future-dated).
 -- In normal mode, both of these are hidden.
 toggleForecast :: Day -> UIState -> UIState
-toggleForecast d ui@UIState{aopts=UIOpts{cliopts_=CliOpts{inputopts_=iopts}}} =
+toggleForecast _d ui@UIState{aopts=UIOpts{cliopts_=copts}} =
   uiSetForecast ui $
-    case forecast_ iopts of
+    case forecast_ $ inputopts_ copts of
       Just _  -> Nothing
-      Nothing -> forecastPeriod d iopts{forecast_=Just nulldatespan} (ajournal ui)
+      Nothing -> forecast_ . inputopts_ $ enableForecastPreservingPeriod ui copts
+
+-- | Ensure this CliOpts enables forecasted transactions.
+-- If a forecast period was specified in the old CliOpts,
+-- or in the provided UIState's startup options,
+-- it is preserved.
+enableForecastPreservingPeriod :: UIState -> CliOpts -> CliOpts
+enableForecastPreservingPeriod ui copts@CliOpts{inputopts_=iopts} =
+    copts{inputopts_=iopts{forecast_=mforecast}}
+  where
+    mforecast = asum [mprovidedforecastperiod, mstartupforecastperiod, mdefaultforecastperiod]
+      where
+        mprovidedforecastperiod = forecast_ $ inputopts_ copts
+        mstartupforecastperiod  = forecast_ $ inputopts_ $ cliopts_ $ astartupopts ui
+        mdefaultforecastperiod  = Just nulldatespan
 
 -- | Helper: set forecast mode (with the given forecast period) on or off in the UI state.
 uiSetForecast :: UIState -> Maybe DateSpan -> UIState