lib: clarify transaction balancing & balance assertion checking

hledger-lib/Hledger/Data/Journal.hs

@@ -76,14 +76,13 @@ module Hledger.Data.Journal (
 )
 where
 import Control.Applicative (Const(..))
-import Control.Arrow
 import Control.Monad
 import Control.Monad.Except
-import qualified Control.Monad.Reader as R
+import Control.Monad.Reader as R
 import Control.Monad.ST
 import Data.Array.ST
 import Data.Functor.Identity (Identity(..))
-import qualified Data.HashTable.ST.Cuckoo as HT
+import qualified Data.HashTable.ST.Cuckoo as H
 import Data.List
 import Data.List.Extra (groupSort)
 import Data.Maybe
@@ -563,51 +562,223 @@ journalUntieTransactions t@Transaction{tpostings=ps} = t{tpostings=map (\p -> p{
 journalModifyTransactions :: Journal -> Journal
 journalModifyTransactions j = j{ jtxns = modifyTransactions (jtxnmodifiers j) (jtxns j) }
 
--- | Check any balance assertions in the journal and return an error
+-- | Check any balance assertions in the journal and return an error message
--- message if any of them fail.
+-- if any of them fail (or if the transaction balancing they require fails).
-journalCheckBalanceAssertions :: Journal -> Either String Journal
+journalCheckBalanceAssertions :: Journal -> Maybe String
-journalCheckBalanceAssertions j =
+journalCheckBalanceAssertions = either Just (const Nothing) . journalBalanceTransactions True
-  runST $ journalBalanceTransactionsST 
-    True 
-    j 
-    (return ())
-    (\_ _ -> return ()) 
-    (const $ return j)
 
--- | Check a posting's balance assertion and return an error if it
+-- | Infer any missing amounts (to satisfy balance assignments and
--- fails.
+-- to balance transactions) and check that all transactions balance 
+-- and (optional) all balance assertions pass. Or return an error message
+-- (just the first error encountered).
+--
+-- Assumes journalInferCommodityStyles has been called, since those affect transaction balancing.
+--
+-- This does multiple things because amount inferring, balance assignments, 
+-- balance assertions and posting dates are interdependent.
+-- 
+-- Overview, 20190216:
+-- @
+-- ****** parseAndFinaliseJournal['] [[Cli/Utils.hs]], journalAddForecast [[Common.hs]], budgetJournal [[BudgetReport.hs]], tests [[BalanceReport.hs]]
+-- ******* journalBalanceTransactions
+-- ******** runST
+-- ********* runExceptT
+-- ********** runReaderT
+-- *********** balanceNoAssignmentTransactionB
+-- ************ balanceTransactionB [[Transaction.hs]]
+-- ************* balanceTransactionHelper
+-- ************** inferBalancingAmount
+-- *********** balanceAssignmentTransactionAndOrCheckAssertionsB
+-- ************ addAmountAndCheckBalanceAssertionB
+-- ************* addToBalanceB
+-- ************ inferFromAssignmentB
+-- ************ balanceTransactionB [[Transaction.hs]]
+-- ************* balanceTransactionHelper
+-- ************ addToBalanceB
+-- ****** uiCheckBalanceAssertions d ui@UIState{aopts=UIOpts{cliopts_=copts}, ajournal=j} [[ErrorScreen.hs]]
+-- ******* journalCheckBalanceAssertions
+-- ******** journalBalanceTransactions
+-- ****** transactionWizard, postingsBalanced [[Add.hs]], tests [[Transaction.hs]]
+-- ******* balanceTransaction
+-- @
+journalBalanceTransactions :: Bool -> Journal -> Either String Journal
+journalBalanceTransactions assrt j' =
+  let
+    -- ensure transactions are numbered, so we can store them by number 
+    j@Journal{jtxns=ts} = journalNumberTransactions j'
+    styles = journalCommodityStyles j
+    -- balance assignments will not be allowed on these
+    txnmodifieraccts = S.fromList $ map paccount $ concatMap tmpostingrules $ jtxnmodifiers j 
+  in 
+    runST $ do 
+      bals <- H.newSized (length $ journalAccountNamesUsed j)
+      txns <- newListArray (1, genericLength ts) ts
+      runExceptT $ do
+        flip runReaderT (BalancingState styles txnmodifieraccts assrt bals txns) $ do
+          -- Fill in missing posting amounts, check transactions are balanced, 
+          -- and check balance assertions. This is done in two passes:
+          -- 1. Balance the transactions which don't have balance assignments,
+          -- and collect their postings, plus the still-unbalanced transactions, in date order.
+          sortedpsandts <- sortOn (either postingDate tdate) . concat <$>
+                           mapM' balanceNoAssignmentTransactionB (jtxns j)
+          -- 2. Step through these, keeping running account balances, 
+          -- performing balance assignments in and balancing the remaining transactions,
+          -- and checking balance assertions. This last could be a separate pass
+          -- but perhaps it's more efficient to do all at once.
+          void $ mapM' balanceAssignmentTransactionAndOrCheckAssertionsB sortedpsandts
+        ts' <- lift $ getElems txns
+        return j{jtxns=ts'} 
+
+-- | If this transaction has no balance assignments, balance and store it
+-- and return its postings. If it can't be balanced, an error will be thrown.
+--
+-- It it has balance assignments, return it unchanged. If any posting has both 
+-- a balance assignment and a custom date, an error will be thrown.
+--
+balanceNoAssignmentTransactionB :: Transaction -> Balancing s [Either Posting Transaction]
+balanceNoAssignmentTransactionB t
+  | null (assignmentPostings t) = do
+    styles <- R.reader bsStyles
+    t' <- lift $ ExceptT $ return $ balanceTransaction (Just styles) t
+    storeTransactionB t'
+    return [Left $ removePrices p | p <- tpostings t']
+
+  | otherwise = do
+    when (any (isJust . pdate) $ tpostings t) $  -- XXX check more carefully that date and assignment are on same posting ?
+      throwError $
+      unlines $
+      [ "postings may not have both a custom date and a balance assignment."
+      , "Write the posting amount explicitly, or remove the posting date:\n"
+      , showTransaction t
+      ]
+    return [Right $ t {tpostings = removePrices <$> tpostings t}]
+
+-- | This function is called in turn on each item in a date-ordered sequence 
+-- of postings (from already-balanced transactions) or transactions  
+-- (not yet balanced, because containing balance assignments).
+-- It applies balance assignments and balances the unbalanced transactions, 
+-- and checks any balance assertion(s).
+--
+-- For a posting: update the account's running balance, and 
+-- check the balance assertion if any.
+--
+-- For a transaction: for each posting, 
+-- 
+-- - if it has a missing amount and a balance assignment, infer the amount 
+--
+-- - update the account's running balance
+--
+-- - check the balance assertion if any
+--
+-- Then balance the transaction, so that any remaining missing amount is inferred. 
+-- And if that happened, also update *that* account's running balance.  XXX and check the assertion ? 
+-- And store the transaction.
+--
+-- Will throw an error if a transaction can't be balanced,
+-- or if an illegal balance assignment is found (cf checkIllegalBalanceAssignment). 
+--
+balanceAssignmentTransactionAndOrCheckAssertionsB :: Either Posting Transaction -> Balancing s ()
+balanceAssignmentTransactionAndOrCheckAssertionsB (Left p) = do
+  checkIllegalBalanceAssignmentB p
+  void $ addAmountAndCheckBalanceAssertionB return p
+balanceAssignmentTransactionAndOrCheckAssertionsB (Right t@Transaction{tpostings=ps}) = do
+  mapM_ checkIllegalBalanceAssignmentB ps
+  ps' <- forM ps $ addAmountAndCheckBalanceAssertionB inferFromAssignmentB
+  styles <- R.reader bsStyles
+  storeTransactionB =<< 
+    balanceTransactionB (fmap void . addToBalanceB) (Just styles) t{tpostings=ps'}
+
+-- | Throw an error if this posting is trying to do a balance assignment and
+-- the account does not allow balance assignments (because it is referenced
+-- by a transaction modifier).
+checkIllegalBalanceAssignmentB :: Posting -> Balancing s ()
+checkIllegalBalanceAssignmentB p = do
+  unassignable <- R.asks bsUnassignable
+  when (isAssignment p && paccount p `S.member` unassignable) $
+    throwError $
+    unlines $
+    [ "cannot assign amount to account "
+    , ""
+    , "    " ++ T.unpack (paccount p)
+    , ""
+    , "because it is also included in transaction modifiers."
+    ]
+
+-- | If this posting has a missing amount and a balance assignment, use
+-- the running account balance to infer the amount required to satisfy
+-- the assignment.
+inferFromAssignmentB :: Posting -> Balancing s Posting
+inferFromAssignmentB p@Posting{paccount=acc} =
+  case pbalanceassertion p of
+    Nothing -> return p
+    Just ba | batotal ba -> do
+      diff <- setAccountRunningBalance acc $ Mixed [baamount ba]
+      return $ setPostingAmount diff p
+    Just ba -> do
+      oldbal <- fromMaybe 0 <$> liftB (\bs -> H.lookup (bsBalances bs) acc)
+      let amt    = baamount ba
+          newbal = filterMixedAmount ((/=acommodity amt).acommodity) oldbal + Mixed [amt]
+      diff <- setAccountRunningBalance acc newbal
+      return $ setPostingAmount diff p
+  where
+    setPostingAmount a p = p{pamount=a, porigin=Just $ originalPosting p}
+    -- | Set the account's running balance, and return the difference from the old. 
+    setAccountRunningBalance :: AccountName -> MixedAmount -> Balancing s MixedAmount
+    setAccountRunningBalance acc amt = liftB $ \BalancingState{bsBalances=bals} -> do
+      old <- fromMaybe 0 <$> H.lookup bals acc
+      H.insert bals acc amt
+      return $ amt - old
+
+-- | Adds a posting's amount to the posting's account's running balance, and
+-- checks the posting's balance assertion if any. Or if the posting has no
+-- amount, runs the supplied fallback action.
+addAmountAndCheckBalanceAssertionB :: 
+     (Posting -> Balancing s Posting) -- ^ fallback action
+  -> Posting
+  -> Balancing s Posting
+addAmountAndCheckBalanceAssertionB _ p | hasAmount p = do
+  newAmt <- addToBalanceB (paccount p) (pamount p)
+  assrt <- R.reader bsAssrt
+  lift $ when assrt $ ExceptT $ return $ checkBalanceAssertion p newAmt
+  return p
+addAmountAndCheckBalanceAssertionB fallback p = fallback p
+
+-- | Check a posting's balance assertion against the given actual balance, and
+-- return an error if the assertion is not satisfied.
+-- If the assertion is partial, unasserted commodities in the actual balance
+-- are ignored; if it is total, they will cause the assertion to fail.
 checkBalanceAssertion :: Posting -> MixedAmount -> Either String ()
-checkBalanceAssertion p@Posting{pbalanceassertion=Just (BalanceAssertion{baamount,baexact})} actualbal =
+checkBalanceAssertion p@Posting{pbalanceassertion=Just (BalanceAssertion{baamount,batotal})} actualbal =
   foldl' f (Right ()) assertedamts
     where
-      f (Right _) assertedamt = checkBalanceAssertionCommodity p assertedamt actualbal
+      f (Right _) assertedamt = checkBalanceAssertionOneCommodity p assertedamt actualbal
       f err _                 = err
       assertedamts = baamount : otheramts
         where
           assertedcomm = acommodity baamount
-          otheramts | baexact   = map (\a -> a{ aquantity = 0 }) $ amounts $ filterMixedAmount (\a -> acommodity a /= assertedcomm) actualbal
+          otheramts | batotal   = map (\a -> a{ aquantity = 0 }) $ amounts $ filterMixedAmount (\a -> acommodity a /= assertedcomm) actualbal
                     | otherwise = []
 checkBalanceAssertion _ _ = Right ()
 
--- | Are the asserted balance and the actual balance
+-- | Does this (single commodity) expected balance match the amount of that
--- exactly equal (disregarding display precision) ?
+-- commodity in the given (multicommodity) actual balance ? If not, returns a
--- The posting is used for creating an error message.
+-- balance assertion failure message based on the provided posting.  To match,
-checkBalanceAssertionCommodity :: Posting -> Amount -> MixedAmount -> Either String ()
+-- the amounts must be exactly equal (display precision is ignored here).
-checkBalanceAssertionCommodity p assertedamt actualbal
+checkBalanceAssertionOneCommodity :: Posting -> Amount -> MixedAmount -> Either String ()
+checkBalanceAssertionOneCommodity p assertedamt actualbal
   | pass      = Right ()
-  | otherwise = Left err
+  | otherwise = Left errmsg
     where
       assertedcomm = acommodity assertedamt
       actualbalincommodity = fromMaybe nullamt $ find ((== assertedcomm) . acommodity) (amounts actualbal)
       pass =
         aquantity
-        -- traceWith (("asserted:"++).showAmountDebug)
+          -- traceWith (("asserted:"++).showAmountDebug)
-        assertedamt ==
+          assertedamt ==
         aquantity
-        -- traceWith (("actual:"++).showAmountDebug)
+          -- traceWith (("actual:"++).showAmountDebug)
-        actualbalincommodity
+          actualbalincommodity
-      diff = aquantity assertedamt - aquantity actualbalincommodity
+      errmsg = printf (unlines
-      err = printf (unlines
                     [ "balance assertion: %s",
                       "\nassertion details:",
                       "date:       %s",
@@ -635,208 +806,7 @@ checkBalanceAssertionCommodity p assertedamt actualbal
         -- (showAmount actualbalincommodity)
         (show $ aquantity assertedamt)
         -- (showAmount assertedamt)
-        (show diff)
+        (show $ aquantity assertedamt - aquantity actualbalincommodity)
-
--- | Fill in any missing amounts and check that all journal transactions
--- balance and all balance assertions pass, or return an error message.
--- This is done after parsing all amounts and applying canonical
--- commodity styles, since balancing depends on display precision.
--- Reports only the first error encountered.
-journalBalanceTransactions :: Bool -> Journal -> Either String Journal
-journalBalanceTransactions assrt j =
-  runST $ journalBalanceTransactionsST 
-    assrt -- check balance assertions also ?
-    (journalNumberTransactions j) -- journal to process
-    (newArray_ (1, genericLength $ jtxns j) :: forall s. ST s (STArray s Integer Transaction)) -- initialise state
-    (\arr tx -> writeArray arr (tindex tx) tx)    -- update state
-    (fmap (\txns -> j{ jtxns = txns}) . getElems) -- summarise state
-
--- | Helper used by 'journalBalanceTransactions' and 'journalCheckBalanceAssertions'.
--- Balances transactions, applies balance assignments, and checks balance assertions
--- at the same time.
-journalBalanceTransactionsST ::
-  Bool
-  -> Journal
-  -> ST s txns                        -- ^ initialise state
-  -> (txns -> Transaction -> ST s ()) -- ^ update state
-  -> (txns -> ST s a)                 -- ^ summarise state
-  -> ST s (Either String a)
-journalBalanceTransactionsST assrt j createStore storeIn extract =
-  runExceptT $ do
-    bals <- lift $ HT.newSized size
-    txStore <- lift $ createStore
-    let env = Env bals 
-                  (storeIn txStore) 
-                  assrt
-                  (Just $ journalCommodityStyles j)
-                  (getModifierAccountNames j)
-    flip R.runReaderT env $ do
-      dated <- fmap snd . sortOn fst . concat
-                <$> mapM' discriminateByDate (jtxns j)
-      mapM' checkInferAndRegisterAmounts dated
-    lift $ extract txStore
-    where 
-      size = genericLength $ journalPostings j
-
-
--- | Collect account names in account modifiers into a set
-getModifierAccountNames :: Journal -> S.Set AccountName
-getModifierAccountNames j = S.fromList $
-                            map paccount $
-                            concatMap tmpostingrules $
-                            jtxnmodifiers j
-
--- | Monad transformer stack with a reference to a mutable hashtable
--- of current account balances and a mutable array of finished
--- transactions in original parsing order.
-type CurrentBalancesModifier s = R.ReaderT (Env s) (ExceptT String (ST s))
-
--- | Environment for 'CurrentBalancesModifier'
-data Env s = Env { eBalances     :: HT.HashTable s AccountName MixedAmount
-                 , eStoreTx      :: Transaction -> ST s ()
-                 , eAssrt        :: Bool
-                 , eStyles       :: Maybe (M.Map CommoditySymbol AmountStyle)
-                 , eUnassignable :: S.Set AccountName
-                 }
-
--- | This converts a transaction into a list of transactions or
--- postings whose dates have to be considered when checking 
--- balance assertions and handled by 'checkInferAndRegisterAmounts'.
---
--- Transaction without balance assignments can be balanced and stored
--- immediately and their (possibly) dated postings are returned.
---
--- Transaction with balance assignments are only supported if no
--- posting has a 'pdate' value. Supported transactions will be
--- returned unchanged and balanced and stored later in 'checkInferAndRegisterAmounts'.
-discriminateByDate :: Transaction
-  -> CurrentBalancesModifier s [(Day, Either Posting Transaction)]
-discriminateByDate tx
-  | null (assignmentPostings tx) = do
-    styles <- R.reader $ eStyles
-    balanced <- lift $ ExceptT $ return $ balanceTransaction styles tx
-    storeTransaction balanced
-    return $ fmap (postingDate &&& (Left . removePrices)) $ tpostings $ balanced
-  | otherwise = do
-    when (any (isJust . pdate) $ tpostings tx) $
-      throwError $
-      unlines $
-      [ "postings may not have both a custom date and a balance assignment."
-      , "Write the posting amount explicitly, or remove the posting date:\n"
-      , showTransaction tx
-      ]
-    return [(tdate tx, Right $ tx {tpostings = removePrices <$> tpostings tx})]
-
--- | Throw an error if a posting is in the unassignable set.
-checkUnassignablePosting :: Posting -> CurrentBalancesModifier s ()
-checkUnassignablePosting p = do
-  unassignable <- R.asks eUnassignable
-  when (isAssignment p && paccount p `S.member` unassignable) $
-    throwError $
-    unlines $
-    [ "cannot assign amount to account "
-    , ""
-    , "    " ++ T.unpack (paccount p)
-    , ""
-    , "because it is also included in transaction modifiers."
-    ]
-
-
--- | This function takes an object describing changes to
--- account balances on a single day - either a single posting 
--- (from an already balanced transaction without assignments)
--- or a whole transaction with assignments (which is required to 
--- have no posting with pdate set).
---
--- For a single posting, there is not much to do. Only add its amount
--- to its account and check the assertion, if there is one. This
--- functionality is provided by 'addAmountAndCheckBalance'.
---
--- For a whole transaction, it loops over all postings, and performs
--- 'addAmountAndCheckBalance', if there is an amount. If there is no
--- amount, the amount is inferred by the assertion or left empty if
--- there is no assertion. Then, the transaction is balanced, the
--- inferred amount added to the balance (all in 'balanceTransactionUpdate') 
--- and the resulting transaction with no missing amounts is stored 
--- in the array, for later retrieval.
---
--- Again in short:
---
--- 'Left Posting': Check the balance assertion and update the
---  account balance. If the amount is empty do nothing.  this can be
---  the case e.g. for virtual postings
---
--- 'Right Transaction': Loop over all postings, infer their amounts
--- and then balance and store the transaction.
-checkInferAndRegisterAmounts :: Either Posting Transaction
-                             -> CurrentBalancesModifier s ()
-checkInferAndRegisterAmounts (Left p) = do
-  checkUnassignablePosting p
-  void $ addAmountAndCheckBalance return p
-checkInferAndRegisterAmounts (Right oldTx) = do
-  let ps = tpostings oldTx
-  mapM_ checkUnassignablePosting ps
-  styles <- R.reader $ eStyles
-  newPostings <- forM ps $ addAmountAndCheckBalance inferFromAssignment
-  storeTransaction =<< balanceTransactionUpdate
-    (fmap void . addToBalance) styles oldTx { tpostings = newPostings }
-  where
-    inferFromAssignment :: Posting -> CurrentBalancesModifier s Posting
-    inferFromAssignment p = do
-      let acc = paccount p
-      case pbalanceassertion p of
-        Just ba | baexact ba -> do
-          diff <- setMixedBalance acc $ Mixed [baamount ba]
-          fullPosting diff p
-        Just ba -> do
-          old <- liftModifier $ \Env{ eBalances = bals } -> HT.lookup bals acc
-          let amt = baamount ba
-              assertedcomm = acommodity amt
-          diff <- setMixedBalance acc $
-            Mixed [amt] + filterMixedAmount (\a -> acommodity a /= assertedcomm) (fromMaybe nullmixedamt old)
-          fullPosting diff p
-        Nothing -> return p
-    fullPosting amt p = return p
-      { pamount = amt
-      , porigin = Just $ originalPosting p
-      }
-
--- | Adds a posting's amount to the posting's account balance and
--- checks a possible balance assertion. Or if there is no amount,
--- runs the supplied fallback action.
-addAmountAndCheckBalance :: 
-     (Posting -> CurrentBalancesModifier s Posting) -- ^ action if posting has no amount
-  -> Posting
-  -> CurrentBalancesModifier s Posting
-addAmountAndCheckBalance _ p | hasAmount p = do
-  newAmt <- addToBalance (paccount p) $ pamount p
-  assrt <- R.reader eAssrt
-  lift $ when assrt $ ExceptT $ return $ checkBalanceAssertion p newAmt
-  return p
-addAmountAndCheckBalance fallback p = fallback p
-
--- | Sets all commodities comprising an account's balance to the given
--- amounts and returns the difference from the previous balance.
-setMixedBalance :: AccountName -> MixedAmount -> CurrentBalancesModifier s MixedAmount
-setMixedBalance acc amt = liftModifier $ \Env{ eBalances = bals } -> do
-  old <- HT.lookup bals acc
-  HT.insert bals acc amt
-  return $ maybe amt (amt -) old
-
--- | Adds an amount to an account's balance and returns the resulting balance.
-addToBalance :: AccountName -> MixedAmount -> CurrentBalancesModifier s MixedAmount
-addToBalance acc amt = liftModifier $ \Env{ eBalances = bals } -> do
-  new <- maybe amt (+ amt) <$> HT.lookup bals acc
-  HT.insert bals acc new
-  return new
-
--- | Stores a transaction in the transaction array in original parsing order.
-storeTransaction :: Transaction -> CurrentBalancesModifier s ()
-storeTransaction tx = liftModifier $ ($tx) . eStoreTx
-
--- | Helper function.
-liftModifier :: (Env s -> ST s a) -> CurrentBalancesModifier s a
-liftModifier f = R.ask >>= lift . lift . f
 
 -- | Choose and apply a consistent display format to the posting
 -- amounts in each commodity. Each commodity's format is specified by

hledger-lib/Hledger/Data/Posting.hs

@@ -104,7 +104,7 @@ nullsourcepos = JournalSourcePos "" (1,1)
 nullassertion, assertion :: BalanceAssertion
 nullassertion = BalanceAssertion
                   {baamount=nullamt
-                  ,baexact=False
+                  ,batotal=False
                   ,baposition=nullsourcepos
                   }
 assertion = nullassertion

hledger-lib/Hledger/Data/Transaction.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE FlexibleContexts #-}
 {-|
 
 A 'Transaction' represents a movement of some commodity(ies) between two
@@ -8,7 +7,11 @@ tags.
 
 -}
 
-{-# LANGUAGE OverloadedStrings, LambdaCase #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE Rank2Types #-}
+{-# LANGUAGE RecordWildCards #-}
 
 module Hledger.Data.Transaction (
   -- * Transaction
@@ -24,13 +27,18 @@ module Hledger.Data.Transaction (
   balancedVirtualPostings,
   transactionsPostings,
   isTransactionBalanced,
+  balanceTransaction,
+  Balancing,
+  BalancingState(..),
+  addToBalanceB,
+  storeTransactionB,
+  liftB,
+  balanceTransactionB,
   -- nonzerobalanceerror,
   -- * date operations
   transactionDate2,
   -- * arithmetic
   transactionPostingBalances,
-  balanceTransaction,
-  balanceTransactionUpdate,
   -- * rendering
   showTransaction,
   showTransactionUnelided,
@@ -47,7 +55,12 @@ module Hledger.Data.Transaction (
 where
 import Data.List
 import Control.Monad.Except
-import Control.Monad.Identity
+import Control.Monad.Reader (ReaderT, ask)
+import Control.Monad.ST
+import Data.Array.ST
+import qualified Data.HashTable.ST.Cuckoo as HT
+import qualified Data.Map as M
+import qualified Data.Set as S
 import Data.Maybe
 import Data.Text (Text)
 import qualified Data.Text as T
@@ -324,34 +337,78 @@ isTransactionBalanced styles t =
       bvsum' = canonicalise $ costOfMixedAmount bvsum
       canonicalise = maybe id canonicaliseMixedAmount styles
 
--- | Ensure this transaction is balanced, possibly inferring a missing
+-- | Monad used for statefully "balancing" a sequence of transactions.
--- amount or conversion price(s), or return an error message.
+type Balancing s = ReaderT (BalancingState s) (ExceptT String (ST s))
--- Balancing is affected by commodity display precisions, so those can
+
--- (optionally) be provided.
+-- | The state used while balancing a sequence of transactions.
--- 
+data BalancingState s = BalancingState {
--- this fails for example, if there are several missing amounts
+   -- read only
--- (possibly with balance assignments)
+   bsStyles       :: M.Map CommoditySymbol AmountStyle       -- ^ commodity display styles
-balanceTransaction :: Maybe (Map.Map CommoditySymbol AmountStyle)
+  ,bsUnassignable :: S.Set AccountName                       -- ^ accounts in which balance assignments may not be used
-                   -> Transaction -> Either String Transaction
+  ,bsAssrt        :: Bool                                    -- ^ whether to check balance assertions
-balanceTransaction stylemap = runIdentity . runExceptT
+   -- mutable
-  . balanceTransactionUpdate (\_ _ -> return ()) stylemap
+  ,bsBalances     :: HT.HashTable s AccountName MixedAmount  -- ^ running account balances, initially empty
+  ,bsTransactions :: STArray s Integer Transaction           -- ^ the transactions being balanced
+  }
+
+-- | Lift a BalancingState mutator through the Except and Reader 
+-- layers into the Balancing monad.
+liftB :: (BalancingState s -> ST s a) -> Balancing s a
+liftB f = ask >>= lift . lift . f
+
+-- | Add this amount to this account's running balance, 
+-- and return the new running balance.
+addToBalanceB :: AccountName -> MixedAmount -> Balancing s MixedAmount
+addToBalanceB acc amt = liftB $ \BalancingState{bsBalances=bals} -> do
+  b <- maybe amt (+amt) <$> HT.lookup bals acc
+  HT.insert bals acc b
+  return b
+
+-- | Update (overwrite) this transaction with a new one.
+storeTransactionB :: Transaction -> Balancing s ()
+storeTransactionB t = liftB $ \bs ->
+  void $ writeArray (bsTransactions bs) (tindex t) t
 
 
--- | More general version of 'balanceTransaction' that takes an update
+-- | Balance this transaction, ensuring that its postings sum to 0,
--- function
+-- by inferring a missing amount or conversion price(s) if needed. 
-balanceTransactionUpdate :: MonadError String m
+-- Or if balancing is not possible, because of unbalanced amounts or 
-  => (AccountName -> MixedAmount -> m ())
+-- more than one missing amount, returns an error message.
-     -- ^ update function
+-- Whether postings "sum to 0" depends on commodity display precisions,
-  -> Maybe (Map.Map CommoditySymbol AmountStyle)
+-- so those can optionally be provided.
-  -> Transaction -> m Transaction
+balanceTransaction ::
-balanceTransactionUpdate update mstyles t =
+     Maybe (Map.Map CommoditySymbol AmountStyle)  -- ^ commodity display styles
-  (finalize =<< inferBalancingAmount update (fromMaybe Map.empty mstyles) t)
+  -> Transaction
-    `catchError` (throwError . annotateErrorWithTxn t)
+  -> Either String Transaction
+balanceTransaction mstyles = fmap fst . balanceTransactionHelper mstyles 
+
+-- | Like balanceTransaction, but when inferring amounts it will also
+-- use the given state update function to update running account balances. 
+-- Used when balancing a sequence of transactions (see journalBalanceTransactions).
+balanceTransactionB ::
+     (AccountName -> MixedAmount -> Balancing s ())  -- ^ function to update running balances
+  -> Maybe (Map.Map CommoditySymbol AmountStyle)     -- ^ commodity display styles
+  -> Transaction
+  -> Balancing s Transaction
+balanceTransactionB updatebalsfn mstyles t = do
+  case balanceTransactionHelper mstyles t of
+    Left err -> throwError err 
+    Right (t', inferredacctsandamts) -> do
+      mapM_ (uncurry updatebalsfn) inferredacctsandamts
+      return t'
+
+balanceTransactionHelper ::
+     Maybe (Map.Map CommoditySymbol AmountStyle)  -- ^ commodity display styles
+  -> Transaction
+  -> Either String (Transaction, [(AccountName, MixedAmount)])
+balanceTransactionHelper mstyles t = do
+  (t', inferredamtsandaccts) <- 
+    inferBalancingAmount (fromMaybe Map.empty mstyles) $ inferBalancingPrices t 
+  if isTransactionBalanced mstyles t'
+  then Right (txnTieKnot t', inferredamtsandaccts)
+  else Left $ annotateErrorWithTxn t' $ nonzerobalanceerror t'
+
   where
-    finalize t' = let t'' = inferBalancingPrices t'
-                  in if isTransactionBalanced mstyles t''
-                     then return $ txnTieKnot t''
-                     else throwError $ nonzerobalanceerror t''
     nonzerobalanceerror :: Transaction -> String
     nonzerobalanceerror t = printf "could not balance this transaction (%s%s%s)" rmsg sep bvmsg
         where
@@ -364,45 +421,52 @@ balanceTransactionUpdate update mstyles t =
                   ++ showMixedAmount (costOfMixedAmount bvsum)
           sep = if not (null rmsg) && not (null bvmsg) then "; " else "" :: String
 
-    annotateErrorWithTxn t e = intercalate "\n" [showGenericSourcePos $ tsourcepos t, e, showTransactionUnelided t]
+annotateErrorWithTxn :: Transaction -> String -> String
+annotateErrorWithTxn t s = intercalate "\n" [showGenericSourcePos $ tsourcepos t, s, showTransactionUnelided t]    
 
 -- | Infer up to one missing amount for this transactions's real postings, and
 -- likewise for its balanced virtual postings, if needed; or return an error
--- message if we can't.
+-- message if we can't. Returns the updated transaction and any inferred posting amounts,
+-- with the corresponding accounts, in order).
 --
 -- We can infer a missing amount when there are multiple postings and exactly
 -- one of them is amountless. If the amounts had price(s) the inferred amount
 -- have the same price(s), and will be converted to the price commodity.
-inferBalancingAmount :: MonadError String m =>
+inferBalancingAmount :: 
-                        (AccountName -> MixedAmount -> m ()) -- ^ update function
+     Map.Map CommoditySymbol AmountStyle -- ^ commodity display styles
-                     -> Map.Map CommoditySymbol AmountStyle  -- ^ standard amount styles
+  -> Transaction
-                     -> Transaction
+  -> Either String (Transaction, [(AccountName, MixedAmount)])
-                     -> m Transaction
+inferBalancingAmount styles t@Transaction{tpostings=ps}
-inferBalancingAmount update styles t@Transaction{tpostings=ps}
   | length amountlessrealps > 1
-      = throwError "could not balance this transaction - can't have more than one real posting with no amount (remember to put 2 or more spaces before amounts)"
+      = Left $ annotateErrorWithTxn t "could not balance this transaction - can't have more than one real posting with no amount (remember to put 2 or more spaces before amounts)"
   | length amountlessbvps > 1
-      = throwError "could not balance this transaction - can't have more than one balanced virtual posting with no amount (remember to put 2 or more spaces before amounts)"
+      = Left $ annotateErrorWithTxn t "could not balance this transaction - can't have more than one balanced virtual posting with no amount (remember to put 2 or more spaces before amounts)"
   | otherwise
-      = do postings <- mapM inferamount ps
+      = let psandinferredamts = map inferamount ps
-           return t{tpostings=postings}
+            inferredacctsandamts = [(paccount p, amt) | (p, Just amt) <- psandinferredamts]
+        in Right (t{tpostings=map fst psandinferredamts}, inferredacctsandamts)
   where
     (amountfulrealps, amountlessrealps) = partition hasAmount (realPostings t)
     realsum = sumStrict $ map pamount amountfulrealps
     (amountfulbvps, amountlessbvps) = partition hasAmount (balancedVirtualPostings t)
     bvsum = sumStrict $ map pamount amountfulbvps
-    inferamount p@Posting{ptype=RegularPosting}
+
-     | not (hasAmount p) = updateAmount p realsum
+    inferamount :: Posting -> (Posting, Maybe MixedAmount)
-    inferamount p@Posting{ptype=BalancedVirtualPosting}
+    inferamount p =
-     | not (hasAmount p) = updateAmount p bvsum
+      let
-    inferamount p = return p
+        minferredamt = case ptype p of
-    updateAmount p amt = 
+          RegularPosting         | not (hasAmount p) -> Just realsum 
-      update (paccount p) amt' >> return p { pamount=amt', porigin=Just $ originalPosting p }
+          BalancedVirtualPosting | not (hasAmount p) -> Just bvsum 
-      where
+          _                                          -> Nothing 
-        -- Inferred amounts are converted to cost.
+      in
-        -- Also, ensure the new amount has the standard style for its commodity  
+        case minferredamt of
-        -- (the main amount styling pass happened before this balancing pass).   
+          Nothing -> (p, Nothing)
-        amt' = styleMixedAmount styles $ normaliseMixedAmount $ costOfMixedAmount (-amt)
+          Just a  -> (p{pamount=a', porigin=Just $ originalPosting p}, Just a') 
+            where
+              -- Inferred amounts are converted to cost.
+              -- Also ensure the new amount has the standard style for its commodity  
+              -- (since the main amount styling pass happened before this balancing pass);
+              a' = styleMixedAmount styles $ normaliseMixedAmount $ costOfMixedAmount (-a)
 
 -- | Infer prices for this transaction's posting amounts, if needed to make
 -- the postings balance, and if possible. This is done once for the real
@@ -627,17 +691,14 @@ tests_Transaction =
                in postingsAsLines False False t (tpostings t) `is`
                   ["    a          $-0.01", "    b           $0.005", "    c           $0.005"]
             ]
-    , do let inferTransaction :: Transaction -> Either String Transaction
+    , tests
-             inferTransaction = runIdentity . runExceptT . inferBalancingAmount (\_ _ -> return ()) Map.empty
+         "inferBalancingAmount"
-         tests
+         [ (fst <$> inferBalancingAmount Map.empty nulltransaction) `is` Right nulltransaction
-           "inferBalancingAmount"
+         , (fst <$> inferBalancingAmount Map.empty nulltransaction{tpostings = ["a" `post` usd (-5), "b" `post` missingamt]}) `is`
-           [ inferTransaction nulltransaction `is` Right nulltransaction
+           Right nulltransaction{tpostings = ["a" `post` usd (-5), "b" `post` usd 5]}
-           , inferTransaction nulltransaction {tpostings = ["a" `post` usd (-5), "b" `post` missingamt]} `is`
+         , (fst <$> inferBalancingAmount Map.empty nulltransaction{tpostings = ["a" `post` usd (-5), "b" `post` (eur 3 @@ usd 4), "c" `post` missingamt]}) `is`
-             Right nulltransaction {tpostings = ["a" `post` usd (-5), "b" `post` usd 5]}
+           Right nulltransaction{tpostings = ["a" `post` usd (-5), "b" `post` (eur 3 @@ usd 4), "c" `post` usd 1]}
-           , inferTransaction
+         ]
-               nulltransaction {tpostings = ["a" `post` usd (-5), "b" `post` (eur 3 @@ usd 4), "c" `post` missingamt]} `is`
-             Right nulltransaction {tpostings = ["a" `post` usd (-5), "b" `post` (eur 3 @@ usd 4), "c" `post` usd 1]}
-           ]
     , tests
         "showTransaction"
         [ test "show a balanced transaction, eliding last amount" $

hledger-lib/Hledger/Data/Types.hs

@@ -238,32 +238,47 @@ instance Show Status where -- custom show.. bad idea.. don't do it..
   show Pending   = "!"
   show Cleared   = "*"
 
--- | The amount to compare an account's balance to, to verify that the history
+-- | A balance assertion is a declaration about an account's expected balance
--- leading to a given point is correct or to set the account to a known value.
+-- at a certain point (posting date and parse order). They provide additional
+-- error checking and readability to a journal file.
 --
--- Different kinds of balance assertion (from #290):
+-- The 'BalanceAssertion' type is also used to represent balance assignments,
+-- which instruct hledger what an account's balance should become at a certain
+-- point.
 --
--- * simple assertions: single-commodity, non-total, subaccount-exclusive
+-- Different kinds of balance assertions are discussed eg on #290.
---   assertions, as in Ledger (syntax: `=`). See definitions below.
+-- Variables include:
 --
--- * subaccount-inclusive assertions: asserting the balance of an account
+-- - which postings are to be summed (real/virtual; unmarked/pending/cleared; this account/this account including subs)
---   including all its subaccounts' balances. Not implemented, proposed by #290.
 --
--- * multicommodity assertions: writing multiple amounts separated by + to
+-- - which commodities within the balance are to be checked
---   assert a multicommodity balance, in a single assertion. Not implemented,
---   proposed by #934.  In current hledger you can assert a multicommodity
---   balance by using multiple postings/assertions.  But in either case, the
---   balance might contain additional unasserted commodities. To disallow that
---   you need...
 --
--- * total assertions: asserting that the balance is as written, with no extra
+-- - whether to do a partial or a total check (disallowing other commodities)
---   commodities in the account. Added by #902, with syntax `==`. I sometimes
+--
---   wish this was the default behaviour, of `=`.
+-- I suspect we want:
+--
+-- 1. partial, subaccount-exclusive, Ledger-compatible assertions. Because
+--    they're what we've always had, and removing them would break some
+--    journals unnecessarily.  Implemented with = syntax.
+--
+-- 2. total assertions. Because otherwise assertions are a bit leaky.
+--    Implemented with == syntax.
+--
+-- 3. subaccount-inclusive assertions. Because that's something folks need.
+--    Not implemented.
+--
+-- 4. flexible assertions allowing custom criteria (perhaps arbitrary
+--    queries). Because power users have diverse needs and want to try out
+--    different schemes (assert cleared balances, assert balance from real or
+--    virtual postings, etc.). Not implemented.
+--
+-- 5. multicommodity assertions, asserting the balance of multiple commodities
+--    at once. Not implemented, requires #934.
 --
 data BalanceAssertion = BalanceAssertion {
-      baamount   :: Amount,             -- ^ the expected balance of a single commodity
+      baamount    :: Amount,             -- ^ the expected balance in a particular commodity
-      baexact    :: Bool,               -- ^ whether the assertion is total, ie disallowing amounts in other commodities
+      batotal     :: Bool,               -- ^ disallow additional non-asserted commodities ?
-      baposition :: GenericSourcePos
+      baposition  :: GenericSourcePos    -- ^ the assertion's file position, for error reporting
     } deriving (Eq,Typeable,Data,Generic,Show)
 
 instance NFData BalanceAssertion

hledger-lib/Hledger/Read/Common.hs

@@ -728,14 +728,14 @@ balanceassertionp :: JournalParser m BalanceAssertion
 balanceassertionp = do
   sourcepos <- genericSourcePos <$> lift getSourcePos
   char '='
-  exact <- optional $ try $ char '='
+  istotal <- fmap isJust $ optional $ try $ char '='
   lift (skipMany spacenonewline)
   -- this amount can have a price; balance assertions ignore it,
   -- but balance assignments will use it
   a <- amountp <?> "amount (for a balance assertion or assignment)"
   return BalanceAssertion
-    { baamount = a
+    { baamount   = a
-    , baexact = isJust exact
+    , batotal    = istotal
     , baposition = sourcepos
     }
 

hledger-ui/Hledger/UI/ErrorScreen.hs

@@ -176,8 +176,8 @@ uiCheckBalanceAssertions d ui@UIState{aopts=UIOpts{cliopts_=copts}, ajournal=j}
   | ignore_assertions_ $ inputopts_ copts = ui
   | otherwise =
     case journalCheckBalanceAssertions j of
-      Right _  -> ui
+      Nothing  -> ui
-      Left err ->
+      Just err ->
         case ui of
           UIState{aScreen=s@ErrorScreen{}} -> ui{aScreen=s{esError=err}}
           _                                -> screenEnter d errorScreen{esError=err} ui

tests/journal/parse-errors.test

@@ -63,10 +63,11 @@ $ hledger -f - print -x
   c
 
 $ hledger -f journal:- print
->2 /\<4\>/
+>2 /could not balance this transaction - can't have more than one real posting with no amount/
 >=1
 
 # 7. Two (or more) virtual postings with implicit amount cannot be balanced.
+# (And the error message contains line numbers).
 <
 2018/1/1
   [a]  1
@@ -74,5 +75,5 @@ $ hledger -f journal:- print
   [c]
 
 $ hledger -f journal:- print
->2 /\<4\>/
+>2 /lines 1-4/
 >=1