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