fix:ui: re-check balance assertions properly when --pivot is used [#2451]

When hledger-ui is started with --pivot, re-enabling balance
assertions with the I key now does a full journal reload, to check
balance assertions accurately. It means that in pivot mode, the I key
can also show other data changes (as if you pressed the g key).

hledger-ui/Hledger/UI/AccountsScreen.hs

@@ -48,7 +48,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.RegisterScreen (rsCenterSelection)
 import Data.Either (fromRight)
 import Control.Arrow ((>>>))
@@ -257,7 +257,7 @@ asHandleNormalMode (ALS scons ass) ev = do
     VtyEvent (EvKey (KLeft)     [MShift]) -> modify' (previousReportPeriod journalspan >>> regenerateScreens j d)
 
     -- various toggles and settings:
-    VtyEvent (EvKey (KChar 'I') []) -> modify' (toggleIgnoreBalanceAssertions >>> uiCheckBalanceAssertions d)
+    VtyEvent (EvKey (KChar 'I') []) -> get' >>= uiToggleBalanceAssertions d
     VtyEvent (EvKey (KChar 'F') []) -> modify' (toggleForecast d   >>> regenerateScreens j d)
     VtyEvent (EvKey (KChar 'B') []) -> modify' (toggleConversionOp >>> regenerateScreens j d)
     VtyEvent (EvKey (KChar 'V') []) -> modify' (toggleValue        >>> regenerateScreens j d)

hledger-ui/Hledger/UI/ErrorScreen.hs

@@ -13,6 +13,7 @@ module Hledger.UI.ErrorScreen
  ,uiCheckBalanceAssertions
  ,uiReload
  ,uiReloadIfFileChanged
+ ,uiToggleBalanceAssertions
  )
 where
 
@@ -102,7 +103,7 @@ esHandle ev = do
                 runEditor pos f
               esReloadIfFileChanged copts d j ui
 
-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (popScreen $ toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d (popScreen ui)
             VtyEvent (EvKey (KChar 'l') [MCtrl]) -> redraw
             VtyEvent (EvKey (KChar 'z') [MCtrl]) -> suspend ui
             _ -> return ()
@@ -150,6 +151,10 @@ hledgerparseerrorpositionp = do
 
 -- Defined here so it can reference the error screen:
 
+-- | Modify some input options for hledger-ui (enable --forecast).
+uiAdjustOpts :: UIOpts -> CliOpts -> CliOpts
+uiAdjustOpts uopts = enableForecast uopts
+
 -- | Reload the journal from its input files, then update the ui app state accordingly.
 -- This means regenerate the entire screen stack from top level down to the current screen, using the provided today-date.
 -- As a convenience (usually), if journal reloading fails, this enters the error screen, or if already there, updates its message.
@@ -163,8 +168,8 @@ hledgerparseerrorpositionp = do
 uiReload :: CliOpts -> Day -> UIState -> EventM Name UIState UIState
 uiReload copts d ui = liftIO $ do
   ej <-
-    let copts' = enableForecast (astartupopts ui) copts
+    let copts1 = uiAdjustOpts (astartupopts ui) copts
-    in runExceptT $ journalTransform copts' <$> journalReload copts'
+    in runExceptT $ journalTransform copts1 <$> journalReload copts1
   -- dbg1IO "uiReload before reload" (map tdescription $ jtxns $ ajournal ui)
   return $ case ej of
     Right j  ->
@@ -188,20 +193,24 @@ uiReload copts d ui = liftIO $ do
 -- Also, this one runs in IO, suitable for suspendAndResume.
 uiReloadIfFileChanged :: CliOpts -> Day -> Journal -> UIState -> IO UIState
 uiReloadIfFileChanged copts d j ui = do
-  let copts' = enableForecast (astartupopts ui) copts
+  ej <-
-  ej <- runExceptT $ journalReloadIfChanged copts' d j
+    let copts1 = uiAdjustOpts (astartupopts ui) copts
+    in runExceptT $ journalReloadIfChanged copts1 d j
   return $ case ej of
     Right (j', _) -> regenerateScreens j' d ui
     Left err -> case aScreen ui of
         ES _ -> ui{aScreen=esNew err}
         _    -> pushScreen (esNew err) ui
 
--- Re-check any balance assertions in the current journal, and if any
+-- Re-check any balance assertions in the current journal,
--- fail, enter (or update) the error screen. Or if balance assertions
+-- and if any fail, enter (or update) the error screen.
--- are disabled, do nothing.
+-- Or if balance assertions are disabled or pivot is active, do nothing.
+-- (When pivot is active, assertions have already been checked on the pre-pivot journal,
+-- and the current post-pivot journal's account names don't match the original assertions.)
 uiCheckBalanceAssertions :: Day -> UIState -> UIState
-uiCheckBalanceAssertions _d ui@UIState{ajournal=j}
+uiCheckBalanceAssertions _d ui@UIState{ajournal=j, aopts=UIOpts{uoCliOpts=CliOpts{inputopts_=InputOpts{pivot_=pval}}}}
-  | ui^.ignore_assertions = ui
+  | ui^.ignore_assertions = ui        -- user disabled checks
+  | not (null pval) = ui              -- post-pivot journal, assertions already checked pre-pivot
   | otherwise =
     case journalCheckBalanceAssertions j of
       Right () -> ui
@@ -209,3 +218,16 @@ uiCheckBalanceAssertions _d ui@UIState{ajournal=j}
         case ui of
           UIState{aScreen=ES sst} -> ui{aScreen=ES sst{_essError=err}}
           _                        -> pushScreen (esNew err) ui
+
+-- | Toggle ignoring balance assertions (when user presses I), and if no longer ignoring, recheck them.
+-- Normally the recheck is done quickly on the in-memory journal.
+-- But if --pivot is active, a full journal reload is done instead
+-- (because we can't check balance assertions after pivoting has occurred).
+-- In that case, this operation could be slower and could reveal other data changes (not just balance assertion failures).
+uiToggleBalanceAssertions :: Day -> UIState -> EventM Name UIState ()
+uiToggleBalanceAssertions d ui@UIState{aopts=UIOpts{uoCliOpts=copts@CliOpts{inputopts_=InputOpts{pivot_=pivotval}}}} =
+  let ui' = toggleIgnoreBalanceAssertions ui
+  in case (ui'^.ignore_assertions, null pivotval) of
+    (True, _)      -> put' ui'                                -- ignoring enabled, no check needed
+    (False, True)  -> put' $ uiCheckBalanceAssertions d ui'   -- unpivoted journal, can check in memory
+    (False, False) -> uiReload copts d ui' >>= put'           -- pivoted journal, must reload to check it

hledger-ui/Hledger/UI/MenuScreen.hs

@@ -33,7 +33,7 @@ import Hledger.UI.UITypes
 import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.Editor (runIadd, runEditor, endPosition)
 import Brick.Widgets.Edit (getEditContents, handleEditorEvent)
 import Control.Arrow ((>>>))
@@ -156,7 +156,7 @@ msHandle ev = do
               where
                 p = reportPeriod ui
             e | e `elem` [VtyEvent (EvKey (KChar 'g') []), AppEvent FileChange] -> uiReload copts d ui >>= put'
-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui
             VtyEvent (EvKey (KChar 'a') []) -> suspendAndResume $ clearScreen >> setCursorPosition 0 0 >> add (cliOptsDropArgs copts) j >> uiReloadIfFileChanged copts d j ui
             VtyEvent (EvKey (KChar 'A') []) -> suspendAndResume $ void (runIadd (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui
             VtyEvent (EvKey (KChar 'E') []) -> suspendAndResume $ void (runEditor endPosition (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui

hledger-ui/Hledger/UI/RegisterScreen.hs

@@ -43,7 +43,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 
 rsDraw :: UIState -> [Widget Name]
 rsDraw UIState{aopts=_uopts@UIOpts{uoCliOpts=copts@CliOpts{reportspec_=rspec}}
@@ -247,7 +247,7 @@ rsHandle ev = do
 
             e | e `elem` [AppEvent FileChange, VtyEvent (EvKey (KChar 'g') [])] -> uiReload copts d ui >>= put'
 
-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui
             VtyEvent (EvKey (KChar 'a') []) -> suspendAndResume $ clearScreen >> setCursorPosition 0 0 >> add (cliOptsDropArgs copts) j >> uiReloadIfFileChanged copts d j ui
             VtyEvent (EvKey (KChar 'A') []) -> suspendAndResume $ void (runIadd (journalFilePath j)) >> uiReloadIfFileChanged copts d j ui
             VtyEvent (EvKey (KChar 'T') []) -> put' $ regenerateScreens j d $ setReportPeriod (DayPeriod d) ui

hledger-ui/Hledger/UI/TransactionScreen.hs

@@ -30,7 +30,7 @@ import Hledger.UI.UIState
 import Hledger.UI.UIUtils
 import Hledger.UI.UIScreens
 import Hledger.UI.Editor
-import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged)
+import Hledger.UI.ErrorScreen (uiCheckBalanceAssertions, uiReload, uiReloadIfFileChanged, uiToggleBalanceAssertions)
 import Hledger.UI.RegisterScreen (rsHandle)
 
 tsDraw :: UIState -> [Widget Name]
@@ -162,7 +162,7 @@ tsHandle ev = do
               where
                 p = reportPeriod ui
 
-            VtyEvent (EvKey (KChar 'I') []) -> put' $ uiCheckBalanceAssertions d (toggleIgnoreBalanceAssertions ui)
+            VtyEvent (EvKey (KChar 'I') []) -> uiToggleBalanceAssertions d ui
 
             -- for toggles that may change the current/prev/next transactions,
             -- we must regenerate the transaction list, like the g handler above ? with regenerateTransactions ? TODO WIP