1 files changed, 164 insertions, 92 deletions

diff --git a/lib/ledger.ml b/lib/ledger.ml
index 54a030e..7de131f 100644
--- a/lib/ledger.ml
+++ b/lib/ledger.ml
@@ -187,28 +187,34 @@ module Gh = struct
   type global
 
   (* The five top-level categories *)
-  type asset
-  type equity
-  type expense
-  type income
-  type liability
+  type asset [@@deriving sexp_of]
+  type equity [@@deriving sexp_of]
+  type expense [@@deriving sexp_of]
+  type income [@@deriving sexp_of]
+  type liability [@@deriving sexp_of]
+
+  (* Subcategories of assets *)
+  type bank [@@deriving sexp_of]
 
   (* No subcategories *)
-  type final
+  type final [@@deriving sexp_of]
 end
 
 module Account_structure0 (F : sig
-  type 'a t
+  type 'a t [@@deriving sexp_of]
 end) =
 struct
   type 'a f =
     | Accounts_payable : Gh.final f F.t -> Gh.liability f
     | Accounts_receivable : Gh.final f F.t -> Gh.asset f
-    | Bank : Gh.final f F.t -> Gh.asset f
+    | Bank : Gh.bank f F.t -> Gh.asset f
     | Cash : Gh.final f F.t -> Gh.asset f
     | Credit : Gh.final f F.t -> Gh.liability f
     | Mutual_fund : Gh.final f F.t -> Gh.asset f
     | Stock : Gh.final f F.t -> Gh.asset f
+    | Savings : Gh.final f F.t -> Gh.bank f
+    | Checking : Gh.final f F.t -> Gh.bank f
+  [@@deriving sexp_of]
 
   type t0 =
     | Asset of Gh.asset f F.t
@@ -216,18 +222,19 @@ struct
     | Expense of Gh.expense f F.t
     | Income of Gh.income f F.t
     | Liability of Gh.liability f F.t
+  [@@deriving sexp_of]
 end
 
 module Account_type = struct
-  type 'a elem = Leaf | Node of 'a
+  type 'a elem = Leaf | Node of 'a [@@deriving sexp_of]
 
   include Account_structure0 (struct
-    type 'a t = 'a elem
+    type 'a t = 'a elem [@@deriving sexp_of]
   end)
 end
 
 module Account_structure (F : sig
-  type 'a t
+  type 'a t [@@deriving sexp_of]
 end) =
 struct
   include Account_structure0 (F)
@@ -272,14 +279,36 @@ struct
       | Stock v ->
           let%map c, v' = f.car v (fun el -> mkt (Account_type.Stock el)) in
           (c, Stock v')
+      | Savings v ->
+          let%map c, v' = f.car v (fun el -> mkt (Account_type.Savings el)) in
+          (c, Savings v')
+      | Checking v ->
+          let%map c, v' = f.car v (fun el -> mkt (Account_type.Checking el)) in
+          (c, Checking v')
+  end
+
+  module Folder = struct
+    type nonrec 'b t = { car : 'a. 'a f F.t -> 'b }
+
+    let fold (type b c) (f : c t) : b f -> c = function
+      | Accounts_payable v -> f.car v
+      | Accounts_receivable v -> f.car v
+      | Bank v -> f.car v
+      | Cash v -> f.car v
+      | Credit v -> f.car v
+      | Mutual_fund v -> f.car v
+      | Stock v -> f.car v
+      | Savings v -> f.car v
+      | Checking v -> f.car v
   end
 end
 
 module Typed_account_path = struct
   type 'a elem = Leaf | Node of string * 'a | Ind of string * 'a elem
+  [@@deriving sexp_of]
 
   include Account_structure (struct
-    type 'a t = 'a elem
+    type 'a t = 'a elem [@@deriving sexp_of]
   end)
 end
 
@@ -300,78 +329,158 @@ module Account_hierarchy = struct
     | Ind of 'a account String.Map.t
     (* Has no subaccounts, has a balance in a certain commodity *)
     | Leaf of Commodity_id.t * Money.Diff.t
+  [@@deriving sexp_of]
 
-  and extra = { description : String.t }
+  and extra = { description : String.t } [@@deriving sexp_of]
   and 'a account = extra * 'a core
 
   include Account_structure (struct
-    type 'a t = 'a account
+    type 'a t = 'a account [@@deriving sexp_of]
   end)
 
   (* All accounts *)
   type world = t0 String.Map.t
 
-  let rec alter_aux (subaid : Account_path.t)
+  let rec unsafe_alter_aux (subaid : Account_path.t)
       (f :
-        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
-      : 'a Mapper.t =
+        Account_type.t0 ->
+        extra ->
+        Commodity_id.t ->
+        Money.Diff.t ->
+        'a * extra * 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, acc_bal' = f (mkt Account_type.Leaf) acc_comm acc_bal in
-              Some (x, (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')))
           | [], _ -> None
           | subaid0 :: subaid, (extra, Node subaccs) ->
-              let open Option.Let_syntax in
               let%bind subacc = Map.find subaccs subaid0 in
               let%map x, subacc' =
-                Mapper.map (alter_aux subaid f) (fun k -> mkt (Node k)) subacc
+                Mapper.map
+                  (unsafe_alter_aux subaid f)
+                  (fun k -> mkt (Node k))
+                  subacc
               in
               (x, (extra, Node (Map.set subaccs ~key:subaid0 ~data:subacc')))
           | subaid0 :: subaid, (extra, Ind subaccs) ->
-              let open Option.Let_syntax in
               let%bind subacc = Map.find subaccs subaid0 in
-              let%map x, subacc' = (alter_aux subaid f).car subacc mkt in
+              let%map x, subacc' = (unsafe_alter_aux subaid f).car subacc mkt in
               (x, (extra, Ind (Map.set subaccs ~key:subaid0 ~data:subacc')))
           | _ :: _, (_, Leaf _) -> None);
     }
 
-  let alter (aid : Account_path.t)
+  let unsafe_alter (aid : Account_path.t)
       (f :
-        Account_type.t0 -> Commodity_id.t -> Money.Diff.t -> 'a * Money.Diff.t)
-      (w : world) : ('a * world) option =
+        Account_type.t0 ->
+        extra ->
+        Commodity_id.t ->
+        Money.Diff.t ->
+        'a * extra * Money.Diff.t) (w : world) : ('a * world) option =
     match aid with
     | [] -> None
     | aid0 :: subaid -> (
         let open Option.Let_syntax in
         match%bind Map.find w aid0 with
         | Asset acc ->
-            let%map x, acc' = (alter_aux subaid f).car acc (fun k -> Asset k) in
+            let%map x, acc' =
+              (unsafe_alter_aux subaid f).car acc (fun k -> Asset k)
+            in
             (x, Map.set w ~key:aid0 ~data:(Asset acc'))
         | Expense acc ->
             let%map x, acc' =
-              (alter_aux subaid f).car acc (fun k -> Expense k)
+              (unsafe_alter_aux subaid f).car acc (fun k -> Expense k)
             in
             (x, Map.set w ~key:aid0 ~data:(Expense acc'))
         | Income acc ->
             let%map x, acc' =
-              (alter_aux subaid f).car acc (fun k -> Income k)
+              (unsafe_alter_aux subaid f).car acc (fun k -> Income k)
             in
             (x, Map.set w ~key:aid0 ~data:(Income acc'))
         | Liability acc ->
             let%map x, acc' =
-              (alter_aux subaid f).car acc (fun k -> Liability k)
+              (unsafe_alter_aux subaid f).car acc (fun k -> Liability k)
             in
             (x, Map.set w ~key:aid0 ~data:(Liability acc'))
         | Equity acc ->
             let%map x, acc' =
-              (alter_aux subaid f).car acc (fun k -> Equity k)
+              (unsafe_alter_aux subaid f).car acc (fun k -> Equity k)
             in
             (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 =
+    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)
+          else
+            let on_debit =
+              match acc_type with
+              | Asset _ -> `Incr
+              | Expense _ -> `Incr
+              | Income _ -> `Decr
+              | Liability _ -> `Decr
+              | Equity _ -> `Decr
+            in
+            let acc_bal' =
+              Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit)
+            in
+            (Some (acc_bal, acc_bal'), acc_extra, 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 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.
+      a 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))
+    | _extra, Node subaccs ->
+        Map.fold subaccs ~init:Commodity_id.Map.empty
+          ~f:(fun ~key:_ ~data:subacc comm_bal_sums ->
+            add_balance_maps comm_bal_sums
+              (Folder.fold { car = collect_balances } subacc))
+
+  type delete_error = Not_found | Nonzero_balance
+
+  (*
+  let rec delete_aux : type a. (Account_path.t * a f account) -> (a f account, delete_error) result = function
+    | [], (extra, Leaf (_acc_comm, acc_bal)) ->
+       if Money.Diff.(acc_bal = ~$0) then 
+  
+  let delete (aid : Account_path.t) (w : world) =
+   *)
+
   let world_inst : world =
     String.Map.of_alist_exn
       [
@@ -406,12 +515,21 @@ module Account_hierarchy = struct
       ]
 end
 
-type bal_assert = {
-  account : Account_path.t;
-  labels : Labels.t;
-  bal : Money.Diff.t;
-}
-[@@deriving sexp_of]
+module Bal_assert = struct
+  type t = { account : Account_path.t; labels : Labels.t; bal : Money.Diff.t }
+  [@@deriving sexp_of]
+end
+
+module Account_decl = struct
+  type t = {
+    type_ : Account_type.t0;
+    parent : Account_path.t;
+    name : string;
+    commodity : Commodity_id.t;
+    extra : Account_hierarchy.extra;
+  }
+  [@@deriving sexp_of]
+end
 
 module Tx : sig
   type entry = {
@@ -472,7 +590,8 @@ end
 
 type item =
   | Tx_item of Tx.t
-  | Bal_assert_item of bal_assert (*| Account_decl_item of account_decl*)
+  | Bal_assert_item of Bal_assert.t
+  | Account_decl_item of Account_decl.t
 [@@deriving sexp_of]
 
 type t = item list [@@deriving sexp_of]
@@ -482,77 +601,30 @@ module World = struct
 
   let empty : t = String.Map.empty
 
-  (** 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 (world : t) :
-      (Money.Diff.t * Money.Diff.t * t) option =
-    let open Option.Let_syntax in
-    let%bind mres, world' =
-      Account_hierarchy.alter aid
-        (fun acc_type acc_comm acc_bal ->
-          if not ([%equal: Commodity_id.t] acc_comm in_comm) then (None, acc_bal)
-          else
-            match acc_type with
-            | Asset _ ->
-                let acc_bal' =
-                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Incr)
-                in
-                (Some (acc_bal, acc_bal'), acc_bal')
-            | Expense _ ->
-                let acc_bal' =
-                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Incr)
-                in
-                (Some (acc_bal, acc_bal'), acc_bal')
-            | Income _ ->
-                let acc_bal' =
-                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
-                in
-                (Some (acc_bal, acc_bal'), acc_bal')
-            | Liability _ ->
-                let acc_bal' =
-                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
-                in
-                (Some (acc_bal, acc_bal'), acc_bal')
-            | Equity _ ->
-                let acc_bal' =
-                  Money.Diff.(acc_bal + of_amount by_amount dc ~on_debit:`Decr)
-                in
-                (Some (acc_bal, acc_bal'), acc_bal'))
-        world
-    in
-    let%map pre_bal, post_bal = mres in
-    (pre_bal, post_bal, world')
-
-  let get_bal aid (world : t) : (Commodity_id.t * Money.Diff.t) option =
-    let open Option.Let_syntax in
-    let%map cb, _world' =
-      Account_hierarchy.alter aid
-        (fun _acc_type acc_comm acc_bal -> ((acc_comm, acc_bal), acc_bal))
-        world
-    in
-    cb
-
   let apply_tx (tx : Tx.t) world : t option =
     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 =
-          update_bal aid data.dc data.amount data.commodity world
+          Account_hierarchy.update_bal aid data.dc data.amount data.commodity
+            world
         in
         match data.assertion with
         | None -> Some world
         | Some bal_ass ->
             if Money.Diff.(bal_ass = new_bal) then Some world else None)
 
-  let apply_ba (ba : bal_assert) world : t option =
+  let apply_ba (ba : Bal_assert.t) world : t option =
     let open Option.Let_syntax in
-    let%bind _comm, bal = get_bal ba.account world 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 apply_ad (_ad : Account_decl.t) _world : t option = None
+
   let apply : item -> t -> t option = function
     | Tx_item tx -> apply_tx tx
     | Bal_assert_item ba -> apply_ba ba
+    | Account_decl_item ad -> apply_ad ad
 end
 
 module Ctxd_item = struct end