From eca112f7d4fcfebd0b98c0175168fd807c24b4cf Mon Sep 17 00:00:00 2001
From: Simon Michael <simon@joyful.com>
Date: Fri, 5 Dec 2008 05:09:54 +0000
Subject: [PATCH] extract common balance-calculating code

---
 Ledger/Ledger.hs   | 39 ++++++++++++++++++++++++---------------
 RegisterCommand.hs | 34 +++++-----------------------------
 2 files changed, 29 insertions(+), 44 deletions(-)

diff --git a/Ledger/Ledger.hs b/Ledger/Ledger.hs
index 6858b203c..ece810628 100644
--- a/Ledger/Ledger.hs
+++ b/Ledger/Ledger.hs
@@ -33,28 +33,37 @@ instance Show Ledger where
 cacheLedger :: [String] -> RawLedger -> Ledger
 cacheLedger apats l = Ledger{rawledgertext="",rawledger=l,accountnametree=ant,accountmap=acctmap}
     where
+      acctmap = Map.fromList [(a, mkacct a) | a <- flatten ant]
+      mkacct a = Account a (txnsof a) (inclbalof a)
+      (ant,txnsof,_,inclbalof) = groupTransactions ts
       ts = filtertxns apats $ rawLedgerTransactions l
-      ant = rawLedgerAccountNameTree l
 
+-- | Given a list of transactions, return an account name tree and three
+-- query functions that fetch transactions, balance, and
+-- subaccount-including balance by account name. This is to factor out
+-- common logic from cacheLedger and summariseTransactionsInDateSpan.
+groupTransactions :: [Transaction] -> (Tree AccountName, (AccountName -> [Transaction]), (AccountName -> MixedAmount), (AccountName -> MixedAmount))
+groupTransactions ts = (ant,txnsof,exclbalof,inclbalof)
+    where
+      ant = accountNameTreeFrom $ expandAccountNames $ sort $ nub $ map account ts
       anames = flatten ant
       txnmap = Map.union (transactionsByAccount ts) (Map.fromList [(a,[]) | a <- anames])
       txnsof = (txnmap !)
+      balmap = Map.fromList $ flatten $ calculateBalances ant txnsof
+      exclbalof = fst . (balmap !)
+      inclbalof = snd . (balmap !)
 
-      -- add subaccount-including balances to a tree of account names
-      -- somewhat efficiently
-      addbalances :: Tree AccountName -> Tree (AccountName, MixedAmount)
-      addbalances (Node a []) = Node (a,sumTransactions $ txnsof a) []
-      addbalances (Node a subs) = Node (a,sumtxns + sumsubaccts) subbals
+-- | Add subaccount-excluding and subaccount-including balances to a tree
+-- of account names somewhat efficiently, given a function that looks up
+-- transactions by account name.
+calculateBalances :: Tree AccountName -> (AccountName -> [Transaction]) -> Tree (AccountName, (MixedAmount, MixedAmount))
+calculateBalances ant txnsof = addbalances ant
+    where 
+      addbalances (Node a subs) = Node (a,(bal,bal+subsbal)) subs'
           where
-            sumtxns = sumTransactions $ txnsof a
-            sumsubaccts = sum $ map (snd . root) subbals
-            subbals = map addbalances subs
-      balmap = Map.fromList $ flatten $ addbalances ant
-      balof = (balmap !)
-
-      mkacct a = Account a (txnsof a) (balof a)
-      acctmap = Map.fromList [(a, mkacct a) | a <- anames]
-
+            bal         = sumTransactions $ txnsof a
+            subsbal     = sum $ map (snd . snd . root) subs'
+            subs'       = map addbalances subs
 
 -- | Convert a list of transactions to a map from account name to the list
 -- of all transactions in that account.
diff --git a/RegisterCommand.hs b/RegisterCommand.hs
index 800d49aa8..bfed49741 100644
--- a/RegisterCommand.hs
+++ b/RegisterCommand.hs
@@ -76,37 +76,13 @@ summariseTransactionsInDateSpan (DateSpan b e) entryno depth showempty ts
       summaryts'
           | showempty = summaryts
           | otherwise = filter (not . isZeroMixedAmount . amount) summaryts
-      txnanames = sort $ nub $ map account ts
-
-      -- aggregate balances by account like cacheLedger
-      ant = accountNameTreeFrom txnanames
-      anames = flatten ant
-      txnmap = Map.union (transactionsByAccount ts) (Map.fromList [(a,[]) | a <- anames])
-      txnsof = (txnmap !)
-      addbalances :: Tree AccountName -> Tree (AccountName, MixedAmount)
-      addbalances (Node a []) = Node (a,sumTransactions $ txnsof a) []
-      addbalances (Node a subs) = Node (a,sumtxns + sumsubaccts) subbals
-          where
-            sumtxns = sumTransactions $ txnsof a
-            sumsubaccts = sum $ map (snd . root) subbals
-            subbals = map addbalances subs
-      inclbalmap = Map.fromList $ flatten $ addbalances ant
-      inclbalof = (inclbalmap !)
-
-      -- and do depth-clipping
-      addexclbalances :: Tree AccountName -> Tree (AccountName, MixedAmount)
-      addexclbalances (Node a subs) = Node (a,sumtxns) subbals
-          where
-            sumtxns = sumTransactions $ txnsof a
-            subbals = map addexclbalances subs
-      exclbalmap = Map.fromList $ flatten $ addexclbalances ant
-      exclbalof = (exclbalmap !)
-
-      clippedanames = clipAccountNames depth txnanames
-      isclipped a = accountNameLevel a >= fromMaybe 9999 depth
+      summaryts = [txn{account=a,amount=balancetoshowfor a} | a <- clippedanames]
       balancetoshowfor a =
           (if isclipped a then inclbalof else exclbalof) (if null a then "top" else a)
-      summaryts = [txn{account=a,amount=balancetoshowfor a} | a <- clippedanames]
+      (_,_,exclbalof,inclbalof) = groupTransactions ts
+      isclipped a = accountNameLevel a >= fromMaybe 9999 depth
+      clippedanames = clipAccountNames depth txnanames
+      txnanames = sort $ nub $ map account ts
 
 clipAccountNames :: Maybe Int -> [AccountName] -> [AccountName]
 clipAccountNames Nothing as = as