1 files changed, 121 insertions, 0 deletions

diff --git a/lib/account/account.ml b/lib/account/account.ml
new file mode 100644
index 0000000..0ef3d28
--- /dev/null
+++ b/lib/account/account.ml
@@ -0,0 +1,121 @@
+open Prelude
+
+(* TODO: Decide on public interface. Probably should not include
+   functions such as [unsafe_update_bal], but having [Balanced_batch]
+   under [Account] also feels a bit awkward. *)
+
+module Path = struct
+  type t = string list [@@deriving compare, sexp]
+end
+
+(** Ensures that only accounts with valid type hierarchies can be constructed.
+*)
+module Kernel : sig
+  type extra = { description : string }
+
+  type t = private { type_ : Type.t; extra : extra; core : core }
+
+  and core =
+    (* Balance in some commodity *)
+    | Leaf of Money.Commodity_id.t * Money.Diff.t
+    | Subtree of t String.Map.t
+
+  val make : Type.t -> extra -> core -> t option
+end = struct
+  type extra = { description : string }
+
+  type t = { type_ : Type.t; extra : extra; core : core }
+
+  and core =
+    (* Balance in some commodity *)
+    | Leaf of Money.Commodity_id.t * Money.Diff.t
+    | Subtree of t String.Map.t
+
+  let make type_ extra : core -> t option = function
+    | Leaf (comm, bal) -> Some { type_; extra; core = Leaf (comm, bal) }
+    | Subtree children ->
+        if
+          Map.for_all children ~f:(fun subacc ->
+              Type.is_super subacc.type_ ~super:type_ ~strict:false)
+        then Some { type_; extra; core = Subtree children }
+        else None
+end
+
+type t = Kernel.t
+
+type update_bal_error =
+  | Empty_path
+  | Unmatching_commodity_id of { in_account : Money.Commodity_id.t }
+  | Not_a_leaf_account
+  | Not_a_subtree_account
+  | Not_found (* TODO: report at which level *)
+
+(* We do not necessarily expect [aid] to be a valid path, as we
+   always do for paths in the type hierarchy. The difference is that
+   the type hierarchy is fixed, while the account hierarchy can
+   change over the course of a year. *)
+let rec unsafe_update_bal_aux (aid : Path.t) (dc : Money.Debit_credit.t)
+    (by_amount : Money.Amount.t) (in_comm : Money.Commodity_id.t) (w : t) :
+    (t, update_bal_error) result =
+  match (aid, w.core) with
+  | [], Leaf (comm, bal) ->
+      if [%equal: Money.Commodity_id.t] in_comm comm then (* slay! *)
+        let core' =
+          Kernel.Leaf
+            ( comm,
+              Money.Diff.(bal + of_amount by_amount dc (Type.polarity w.type_))
+            )
+        in
+        Ok (Option.value_exn (Kernel.make w.type_ w.extra core'))
+      else (* bruh *)
+        Error (Unmatching_commodity_id { in_account = comm })
+  | [], Subtree _ -> Error Not_a_leaf_account
+  | _ :: _, Leaf _ -> Error Not_a_subtree_account
+  | aid0 :: aid', Subtree subaccs -> (
+      match Map.find subaccs aid0 with
+      | None -> Error Not_found
+      | Some subacc ->
+          let open Result.Let_syntax in
+          let%bind subacc' =
+            (* TODO: when reporting at which level Not_found fails,
+               we want to make sure that we extend the information
+               in the error with the current aid0 (so we recover a
+               full path to where the account is missing *)
+            unsafe_update_bal_aux aid' dc by_amount in_comm subacc
+          in
+          let core' =
+            Kernel.Subtree (Map.set subaccs ~key:aid0 ~data:subacc')
+          in
+          Ok (Option.value_exn (Kernel.make w.type_ w.extra core')))
+
+(* Unfortunate but true, there has to be some kind of a root account
+   with no type :) *)
+type root = t String.Map.t
+
+let unsafe_update_bal (aid : Path.t) (dc : Money.Debit_credit.t)
+    (by_amount : Money.Amount.t) (in_comm : Money.Commodity_id.t) (w : root) :
+    (root, update_bal_error) result =
+  match aid with
+  | [] -> Error Empty_path
+  | aid0 :: aid' -> (
+      match Map.find w aid0 with
+      | None -> Error Not_found
+      | Some subacc ->
+          let open Result.Let_syntax in
+          (* TODO: when reporting at which level Not_found fails,
+               we want to make sure that we extend the information
+               in the error with the current aid0 (so we recover a
+               full path to where the account is missing *)
+          let%bind subacc' =
+            unsafe_update_bal_aux aid' dc by_amount in_comm subacc
+          in
+          Ok (Map.set w ~key:aid0 ~data:subacc'))
+
+module Balanced_batch_acc_paths = Balanced_batch.Make (Path)
+
+let apply_balanced_batch (b : Balanced_batch_acc_paths.t) (w : root) =
+  Map.fold_result (Balanced_batch_acc_paths.entries b) ~init:w
+    ~f:(fun ~key:aid ~(data : Balanced_batch_acc_paths.entry) w ->
+      let open Result.Let_syntax in
+      let%bind w = unsafe_update_bal aid data.dc data.amount data.commodity w in
+      Ok w)