1 files changed, 83 insertions, 0 deletions

diff --git a/lib/account/type_hierarchy.ml b/lib/account/type_hierarchy.ml
new file mode 100644
index 0000000..7f27830
--- /dev/null
+++ b/lib/account/type_hierarchy.ml
@@ -0,0 +1,83 @@
+open Prelude
+
+type tree = { car : tree String.Map.t }
+type trunk = (Money.polarity * tree) String.Map.t
+type path = Base of string | Sub of string * path
+
+let rec path_to_list ?(suffix = []) p =
+  match p with
+  | Base x -> x :: suffix
+  | Sub (x, p') -> path_to_list ~suffix:(x :: suffix) p'
+
+let canonical : trunk =
+  let make alist : tree = { car = String.Map.of_alist_exn alist } in
+  String.Map.of_alist_exn
+    [
+      ( "Asset",
+        ( Money.Increase_on_debit,
+          make
+            [
+              ("Accounts_receivable", make []);
+              ("Bank", make [ ("Savings", make []); ("Checking", make []) ]);
+              ("Cash", make []);
+              ("Mutual_fund", make []);
+              ("Stock", make []);
+            ] ) );
+      ("Equity", (Money.Increase_on_credit, make []));
+      ("Expense", (Money.Increase_on_debit, make []));
+      ("Income", (Money.Increase_on_credit, make []));
+      ( "Liability",
+        ( Money.Increase_on_credit,
+          make [ ("Accounts_payable", make []); ("Credit", make []) ] ) );
+    ]
+
+(* In this module, only the following two function entertains the
+   option that the given path may not be valid (i.e., it does not
+   throw an exception for invalid paths). *)
+let rec get_node : path -> tree option =
+  let open Option.Let_syntax in
+  function
+  | Base x ->
+      let%map _, t = Map.find canonical x in
+      t
+  | Sub (t, p) ->
+      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
+  | Base _ -> None
+  | Sub (_, super) -> Some super
+
+let rec equal_path p1 p2 =
+  match (p1, p2) with
+  | Base x1, Base x2 -> String.(x1 = x2)
+  | Sub (x1, p1'), Sub (x2, p2') -> String.(x1 = x2) && equal_path p1' p2'
+  | _, _ -> false
+
+let is_prefix (p : path) ~(prefix : path) : bool =
+  List.is_prefix (path_to_list p) ~prefix:(path_to_list prefix)
+    ~equal:String.equal
+
+let rec polarity = function
+  | Base x ->
+      let pol, _ = Map.find_exn canonical x in
+      pol
+  | Sub (_, p') -> polarity p'
+
+let assert_valid acc =
+  match get_node acc with None -> failwith "invalid base account" | _ -> acc
+
+let asset = Base "Asset" |> assert_valid
+let equity = Base "Equity" |> assert_valid
+let expense = Base "Expense" |> assert_valid
+let income = Base "Income" |> assert_valid
+let liability = Base "Liability" |> assert_valid