module Stdlib.Data.UnbalancedSet;

import Stdlib.Prelude open;

import Stdlib.Trait.Ord as Ord open using {Ord; module Ord};

import Stdlib.Data.BinaryTree as Tree;
open Tree using {BinaryTree; module BinaryTree};
open BinaryTree using {leaf; node};

type UnbalancedSet (A : Type) :=
  mk@{
    order : Ord A;
    tree : BinaryTree A;
  };

empty {A} {{order : Ord A}} : UnbalancedSet A := UnbalancedSet.mk order leaf;

isMember {A} (elem : A) (set : UnbalancedSet A) : Bool :=
  case set of
    UnbalancedSet.mk@{tree} :=
      let
        go : BinaryTree A -> Bool
          | leaf := false
          | node@{element; left; right} :=
            case Ord.compare elem element of
              | Ord.LessThan := go left
              | Ord.GreaterThan := go right
              | Ord.Equal := true;
      in go tree;

insert {A} (elem : A) (set : UnbalancedSet A) : UnbalancedSet A :=
  case set of
    UnbalancedSet.mk@{order; tree} :=
      let
        go : BinaryTree A -> BinaryTree A
          | leaf := node leaf elem leaf
          | tree@(node l y r) :=
            case Ord.compare elem y of
              | Ord.LessThan := node (go l) y r
              | Ord.Equal := tree
              | Ord.GreaterThan := node l y (go r);
      in UnbalancedSet.mk order (go tree);

length {A} (set : UnbalancedSet A) : Nat :=
  BinaryTree.length (UnbalancedSet.tree set);

toList {A} (set : UnbalancedSet A) : List A :=
  BinaryTree.toList (UnbalancedSet.tree set);

instance
ordUnbalancedSetI {A} {{Ord A}} : Ord (UnbalancedSet A) :=
  Ord.mk@{
    compare (s1 s2 : UnbalancedSet A) : Ordering :=
      Ord.compare (toList s1) (toList s2);
  };