helo

2 files changed, 165 insertions, 43 deletions

diff --git a/lib/account.ml b/lib/account.ml
new file mode 100644
index 0000000..988b55c
--- /dev/null
+++ b/lib/account.ml
@@ -0,0 +1,96 @@
+open Prelude
+
+(** The 'kernel' of account types: a hierarchy of valid types. A valid type is a
+    path that leads to a node in the hierarchy. *)
+module Type_hierarchy : sig
+  type path
+
+  val children : path -> path list
+  val sub : path -> string -> path option
+  val super : path -> path option
+  val equal_path : path -> path -> bool
+  val is_prefix : path -> prefix:path -> bool
+  val root : path
+  val asset : path
+  val equity : path
+  val expense : path
+  val income : path
+  val liability : path
+end = struct
+  type tree = { car : tree String.Map.t }
+  type path = Root | Sub of string * path
+
+  let canonical : tree =
+    let mk alist = { car = String.Map.of_alist_exn alist } in
+    mk
+      [
+        ( "Asset",
+          mk
+            [
+              ("Accounts_receivable", mk []);
+              ("Bank", mk [ ("Savings", mk []); ("Checking", mk []) ]);
+              ("Cash", mk []);
+              ("Mutual_fund", mk []);
+              ("Stock", mk []);
+            ] );
+        ("Equity", mk []);
+        ("Expense", mk []);
+        ("Income", mk []);
+        ("Liability", mk [ ("Accounts_payable", mk []); ("Credit", mk []) ]);
+      ]
+
+  let rec get_node : path -> tree option = function
+    | Root -> Some canonical
+    | Sub (t, p) ->
+        let open Option.Let_syntax in
+        let%bind super = get_node p in
+        Map.find super.car t
+
+  (** Always gives [Some] under valid paths, giving a list of valid paths *)
+  let children (p : path) : path list =
+    let node = Option.value_exn (get_node p) in
+    List.map (Map.keys node.car) ~f:(fun sub -> Sub (sub, p))
+
+  let sub (p : path) (name : string) : path option =
+    let node = Option.value_exn (get_node p) in
+    if Map.mem node.car name then Some (Sub (name, p)) else None
+
+  let super : path -> path option = function
+    | Root -> None
+    | Sub (_, super) -> Some super
+
+  let rec equal_path p1 p2 =
+    match (p1, p2) with
+    | Root, Root -> true
+    | Sub (x1, p1'), Sub (x2, p2') -> String.(x1 = x2) && equal_path p1' p2'
+    | _, _ -> false
+
+  let rec is_prefix (p : path) ~(prefix : path) : bool =
+    match (prefix, p) with
+    | Root, Root | Root, Sub _ -> true
+    | Sub (x1, p'), Sub (x2, prefix') ->
+        String.(x1 = x2) && is_prefix p' ~prefix:prefix'
+    | _ -> false
+
+  let root = Root
+  let asset = sub root "Asset" |> Option.value_exn
+  let equity = sub root "Equity" |> Option.value_exn
+  let expense = sub root "Expense" |> Option.value_exn
+  let income = sub root "Income" |> Option.value_exn
+  let liability = sub root "Liability" |> Option.value_exn
+end
+
+module Type = struct
+  type t = Type_hierarchy.path [@@deriving equal]
+
+  let rec base (t : t) : t option =
+    match Type_hierarchy.super t with
+    | None -> (* [t] is the root type *) None
+    | Some t' ->
+        (* [t] is a base type if its supertype is the root type *)
+        Some (Option.value (base t') ~default:t)
+
+  (** [a] is a strict supertype of [b] *)
+  let is_strict_super a b =
+    Type_hierarchy.is_prefix b ~prefix:a && not ([%equal: t] a b)
+end
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