1 files changed, 69 insertions, 43 deletions

diff --git a/lib/ledger.ml b/lib/ledger.ml
index 9d315ae..d63dbed 100644
--- a/lib/ledger.ml
+++ b/lib/ledger.ml
@@ -395,21 +395,16 @@ module Account_hierarchy = struct
 
   let rec unsafe_alter_aux (subaid : Account_path.t)
       (f :
-        Account_type.t0 ->
-        extra ->
-        Commodity_id.t ->
-        Money.Diff.t ->
-        'a * extra * Money.Diff.t) : 'a Mapper.t =
+        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
+      : 'a Mapper.t =
     {
       car =
         (fun in_acc mkt ->
           let open Option.Let_syntax in
           match (subaid, in_acc) with
           | [], (extra, Leaf (acc_comm, acc_bal)) ->
-              let x, extra', acc_bal' =
-                f (mkt Account_type.Leaf) extra acc_comm acc_bal
-              in
-              Some (x, (extra', Leaf (acc_comm, acc_bal')))
+              let x, acc_bal' = f (mkt Account_type.Leaf) acc_comm acc_bal in
+              Some (x, (extra, Leaf (acc_comm, acc_bal')))
           | [], _ -> None
           | subaid0 :: subaid, (extra, Node subaccs) ->
               let%bind subacc = Map.find subaccs subaid0 in
@@ -429,11 +424,8 @@ module Account_hierarchy = struct
 
   let unsafe_alter (aid : Account_path.t)
       (f :
-        Account_type.t0 ->
-        extra ->
-        Commodity_id.t ->
-        Money.Diff.t ->
-        'a * extra * Money.Diff.t) (w : world) : ('a * world) option =
+        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
+      (w : world) : ('a * world) option =
     match aid with
     | [] -> None
     | aid0 :: subaid -> (
@@ -466,16 +458,18 @@ module Account_hierarchy = struct
             (x, Map.set w ~key:aid0 ~data:(Equity acc')))
 
   (** Update the balance (debit/credit ([dc])) of account [aid] [by_amount]
-      (commodity: [in_comm]) in [world], giving the updated world and the pre
-      and post balances for [aid] iff the account exists in [world]. *)
-  let update_bal aid dc by_amount in_comm (w : world) :
-      (Money.Diff.t * Money.Diff.t * world) option =
+      (commodity: [in_comm]) in [world], giving the updated world and the post
+      balance for [aid] iff the account exists in [world].
+
+      Unsafe because [by_amount] must be balanced with updates to other
+      accounts. *)
+  let unsafe_update_bal aid dc by_amount in_comm (w : world) :
+      (Money.Diff.t * world) option =
     let open Option.Let_syntax in
     let%bind mres, w' =
       unsafe_alter aid
-        (fun acc_type acc_extra acc_comm acc_bal ->
-          if not ([%equal: Commodity_id.t] acc_comm in_comm) then
-            (None, acc_extra, acc_bal)
+        (fun acc_type acc_comm acc_bal ->
+          if not ([%equal: Commodity_id.t] acc_comm in_comm) then (None, acc_bal)
           else
             let on_debit =
               match acc_type with
@@ -488,41 +482,68 @@ module Account_hierarchy = struct
             let acc_bal' =
               Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit)
             in
-            (Some (acc_bal, acc_bal'), acc_extra, acc_bal'))
+            (Some acc_bal', acc_bal'))
         w
     in
-    let%map pre_bal, post_bal = mres in
-    (pre_bal, post_bal, w')
-
-  let get_bal aid (w : world) : (Commodity_id.t * Money.Diff.t) option =
-    let open Option.Let_syntax in
-    let%map cb, _world' =
-      unsafe_alter aid
-        (fun _acc_type acc_extra acc_comm acc_bal ->
-          ((acc_comm, acc_bal), acc_extra, acc_bal))
-        w
-    in
-    cb
+    let%map post_bal = mres in
+    (post_bal, w')
 
   let add_balance_maps m1 m2 : Money.Diff.t Commodity_id.Map.t =
     Map.merge m1 m2 ~f:(fun ~key:_comm -> function
       | `Both (b1, b2) -> Some Money.Diff.(b1 + b2)
       | `Left b | `Right b -> Some b)
 
-  let rec collect_balances : type a.
+  module Account_visitor = struct
+    type 'res t = { car : 'a. 'a Structure.f account -> 'res }
+
+    let rec visit_aux : type a.
+        'res t -> Account_path.t * a Structure.f account -> 'res option =
+     fun visitor -> function
+      | [], in_acc -> Some (visitor.car in_acc)
+      | _ :: _, (_, Leaf _) -> None
+      | subaid0 :: subaid, (_, Ind subaccs) ->
+          let open Option.Let_syntax in
+          let%bind subacc = Map.find subaccs subaid0 in
+          visit_aux visitor (subaid, subacc)
+      | subaid0 :: subaid, (_, Node subaccs) ->
+          let open Option.Let_syntax in
+          let module Visitor = Account_structure.Basic_visitor (Structure) in
+          let%bind subacc = Map.find subaccs subaid0 in
+          Visitor.visit
+            { car = (fun subacc -> (visit_aux visitor) (subaid, subacc)) }
+            subacc
+
+    let visit (visitor : 'res t) (w : world) : Account_path.t -> 'res option =
+      function
+      | [] -> None
+      | aid0 :: subaid -> (
+          let open Option.Let_syntax in
+          match%bind Map.find w aid0 with
+          | Asset acc -> visit_aux visitor (subaid, acc)
+          | Expense acc -> visit_aux visitor (subaid, acc)
+          | Equity acc -> visit_aux visitor (subaid, acc)
+          | Liability acc -> visit_aux visitor (subaid, acc)
+          | Income acc -> visit_aux visitor (subaid, acc))
+  end
+
+  let rec collect_balances_aux : type a.
       a Structure.f account -> Money.Diff.t Commodity_id.Map.t = function
     | _extra, Leaf (acc_comm, acc_bal) ->
         Commodity_id.Map.singleton acc_comm acc_bal
     | _extra, Ind subaccs ->
         Map.fold subaccs ~init:Commodity_id.Map.empty
           ~f:(fun ~key:_ ~data:subacc comm_bal_sums ->
-            add_balance_maps comm_bal_sums (collect_balances subacc))
+            add_balance_maps comm_bal_sums (collect_balances_aux subacc))
     | _extra, Node subaccs ->
         Map.fold subaccs ~init:Commodity_id.Map.empty
           ~f:(fun ~key:_ ~data:subacc comm_bal_sums ->
             let module Visitor = Account_structure.Basic_visitor (Structure) in
             add_balance_maps comm_bal_sums
-              (Visitor.visit { car = collect_balances } subacc))
+              (Visitor.visit { car = collect_balances_aux } subacc))
+
+  let collect_balances :
+      world -> Account_path.t -> Money.Diff.t Commodity_id.Map.t option =
+    Account_visitor.visit { car = collect_balances_aux }
 
   type delete_error = Not_found | Nonzero_balance
 
@@ -569,7 +590,11 @@ module Account_hierarchy = struct
 end
 
 module Bal_assert = struct
-  type t = { account : Account_path.t; labels : Labels.t; bal : Money.Diff.t }
+  type t = {
+    account : Account_path.t;
+    labels : Labels.t;
+    bals : Money.Diff.t Commodity_id.Map.t;
+  }
   [@@deriving sexp_of]
 end
 
@@ -658,9 +683,9 @@ module World = struct
     Map.fold_option tx.entries ~init:world
       ~f:(fun ~key:aid ~(data : Tx.entry) world ->
         let open Option.Let_syntax in
-        let%bind _old_bal, new_bal, world =
-          Account_hierarchy.update_bal aid data.dc data.amount data.commodity
-            world
+        let%bind new_bal, world =
+          Account_hierarchy.unsafe_update_bal aid data.dc data.amount
+            data.commodity world
         in
         match data.assertion with
         | None -> Some world
@@ -669,8 +694,9 @@ module World = struct
 
   let apply_ba (ba : Bal_assert.t) world : t option =
     let open Option.Let_syntax in
-    let%bind _comm, bal = Account_hierarchy.get_bal ba.account world in
-    if not Money.Diff.(bal = ba.bal) then None else Some world
+    let%bind bals = Account_hierarchy.collect_balances world ba.account in
+    if not ([%equal: Money.Diff.t Commodity_id.Map.t] bals ba.bals) then None
+    else Some world
 
   let apply_ad (_ad : Account_decl.t) _world : t option = None