From 6c31393dd37dfaab8e902baad8667aac1a4f3103 Mon Sep 17 00:00:00 2001
From: Samuel May <ag.eitilt@gmail.com>
Date: Fri, 12 Oct 2018 21:28:46 -0700
Subject: [PATCH] lib: Groundwork allowing multi-commodity assertions

---
 hledger-lib/Hledger/Data/Journal.hs | 32 +++++++++++++++++++----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/hledger-lib/Hledger/Data/Journal.hs b/hledger-lib/Hledger/Data/Journal.hs
index a3ac29729..f3b19cc5b 100644
--- a/hledger-lib/Hledger/Data/Journal.hs
+++ b/hledger-lib/Hledger/Data/Journal.hs
@@ -717,9 +717,21 @@ checkInferAndRegisterAmounts (Right oldTx) = do
     (fmap void . addToBalance) styles oldTx { tpostings = newPostings }
   where
     inferFromAssignment :: Posting -> CurrentBalancesModifier s Posting
-    inferFromAssignment p = maybe (return p)
-      (fmap (\a -> p { pamount = a, porigin = Just $ originalPosting p }) . setBalance (paccount p) . baamount)
-      $ pbalanceassertion p
+    inferFromAssignment p = do
+      let acc = paccount p
+      case pbalanceassertion p of
+        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,
@@ -735,15 +747,13 @@ addAmountAndCheckBalance _ p | hasAmount p = do
   return p
 addAmountAndCheckBalance fallback p = fallback p
 
--- | Sets an account's balance to a given amount and returns the
--- difference of new and old amount.
-setBalance :: AccountName -> Amount -> CurrentBalancesModifier s MixedAmount
-setBalance acc amt = liftModifier $ \Env{ eBalances = bals } -> do
+-- | 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
-  let new = Mixed $ (amt :) $ maybe []
-              (filter ((/= acommodity amt) . acommodity) . amounts) old
-  HT.insert bals acc new
-  return $ maybe new (new -) old
+  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