(*------------------------------------------------------------------*)

#define ATS_DYNLOADFLAG 0

#include "share/atspre_staload.hats"

(*------------------------------------------------------------------*)
(* String hashing using XXH3_64bits from the xxHash suite.          *)

#define ATS_EXTERN_PREFIX "hashsets_postiats_"

%{^ /* Embedded C code. */

#include <xxhash.h>

ATSinline() atstype_uint64
hashsets_postiats_mem_hash (atstype_ptr data, atstype_size len)
{
  return (atstype_uint64) XXH3_64bits (data, len);
}

%}

extern fn mem_hash : (ptr, size_t) -<> uint64 = "mac#%"

fn
string_hash (s : string) :<> uint64 =
  let
    val len = string_length s
  in
    mem_hash ($UNSAFE.cast{ptr} s, len)
  end

(*------------------------------------------------------------------*)
(* A trimmed down version of the AVL trees from the AVL Tree task.  *)

datatype bal_t =
| bal_minus1
| bal_zero
| bal_plus1

datatype avl_t (key_t  : t@ype+,
                data_t : t@ype+,
                size   : int) =
| avl_t_nil (key_t, data_t, 0)
| {size_L, size_R : nat}
  avl_t_cons (key_t, data_t, size_L + size_R + 1) of
    (key_t, data_t, bal_t,
     avl_t (key_t, data_t, size_L),
     avl_t (key_t, data_t, size_R))
typedef avl_t (key_t  : t@ype+,
               data_t : t@ype+) =
  [size : int] avl_t (key_t, data_t, size)

extern fun {key_t : t@ype}
avl_t$compare (u : key_t, v : key_t) :<> int

#define NIL avl_t_nil ()
#define CONS avl_t_cons
#define LNIL list_nil ()
#define :: list_cons
#define F false
#define T true

typedef fixbal_t = bool

prfn
lemma_avl_t_param {key_t : t@ype} {data_t : t@ype} {size : int}
                  (avl : avl_t (key_t, data_t, size)) :<prf>
    [0 <= size] void =
  case+ avl of NIL => () | CONS _ => ()

fn {}
minus_neg_bal (bal : bal_t) :<> bal_t =
  case+ bal of
  | bal_minus1 () => bal_plus1
  | _ => bal_zero ()

fn {}
minus_pos_bal (bal : bal_t) :<> bal_t =
  case+ bal of
  | bal_plus1 () => bal_minus1
  | _ => bal_zero ()

fn
avl_t_is_empty {key_t : t@ype} {data_t : t@ype} {size   : int}
               (avl : avl_t (key_t, data_t, size)) :<>
    [b : bool | b == (size == 0)] bool b =
  case+ avl of
  | NIL => T
  | CONS _ => F

fn
avl_t_isnot_empty {key_t : t@ype} {data_t : t@ype} {size   : int}
                  (avl : avl_t (key_t, data_t, size)) :<>
    [b : bool | b == (size <> 0)] bool b =
  ~avl_t_is_empty avl

fn {key_t : t@ype} {data_t : t@ype}
avl_t_search_ref {size  : int}
                 (avl   : avl_t (key_t, data_t, size),
                  key   : key_t,
                  data  : &data_t? >> opt (data_t, found),
                  found : &bool? >> bool found) :<!wrt>
    #[found : bool] void =
  let
    fun
    search (p     : avl_t (key_t, data_t),
            data  : &data_t? >> opt (data_t, found),
            found : &bool? >> bool found) :<!wrt,!ntm>
        #[found : bool] void =
      case+ p of
      | NIL =>
        {
          prval _ = opt_none {data_t} data
          val _ = found := F
        }
      | CONS (k, d, _, left, right) =>
        begin
          case+ avl_t$compare<key_t> (key, k) of
          | cmp when cmp < 0 => search (left, data, found)
          | cmp when cmp > 0 => search (right, data, found)
          | _ =>
            {
              val _ = data := d
              prval _ = opt_some {data_t} data
              val _ = found := T
            }
        end
  in
    $effmask_ntm search (avl, data, found)
  end

fn {key_t : t@ype} {data_t : t@ype}
avl_t_search_opt {size : int}
                 (avl  : avl_t (key_t, data_t, size),
                  key  : key_t) :<>
    Option (data_t) =
  let
    var data : data_t?
    var found : bool?
    val _ = $effmask_wrt avl_t_search_ref (avl, key, data, found)
  in
    if found then
      let
        prval _ = opt_unsome data
      in
        Some {data_t} data
      end
    else
      let
        prval _ = opt_unnone data
      in
        None {data_t} ()
      end
  end

fn {key_t : t@ype} {data_t : t@ype}
avl_t_insert_or_replace {size : int}
                        (avl  : avl_t (key_t, data_t, size),
                         key  : key_t,
                         data : data_t) :<>
    [sz : pos] (avl_t (key_t, data_t, sz), bool) =
  let
    fun
    search {size   : nat}
           (p      : avl_t (key_t, data_t, size),
            fixbal : fixbal_t,
            found  : bool) :<!ntm>
        [sz : pos]
        (avl_t (key_t, data_t, sz), fixbal_t, bool) =
      case+ p of
      | NIL => (CONS (key, data, bal_zero, NIL, NIL), T, F)
      | CONS (k, d, bal, left, right) =>
        case+ avl_t$compare<key_t> (key, k) of
        | cmp when cmp < 0 =>
          let
            val (p1, fixbal, found) = search (left, fixbal, found)
          in
            case+ (fixbal, bal) of
            | (F, _) => (CONS (k, d, bal, p1, right), F, found)
            | (T, bal_plus1 ()) =>
              (CONS (k, d, bal_zero (), p1, right), F, found)
            | (T, bal_zero ()) =>
              (CONS (k, d, bal_minus1 (), p1, right), fixbal, found)
            | (T, bal_minus1 ()) =>
              let
                val+ CONS (k1, d1, bal1, left1, right1) = p1
              in
                case+ bal1 of
                | bal_minus1 () =>
                  let
                    val q = CONS (k, d, bal_zero (), right1, right)
                    val q1 = CONS (k1, d1, bal_zero (), left1, q)
                  in
                    (q1, F, found)
                  end
                | _ =>
                  let
                    val p2 = right1
                    val- CONS (k2, d2, bal2, left2, right2) = p2
                    val q = CONS (k, d, minus_neg_bal bal2,
                                  right2, right)
                    val q1 = CONS (k1, d1, minus_pos_bal bal2,
                                   left1, left2)
                    val q2 = CONS (k2, d2, bal_zero (), q1, q)
                  in
                    (q2, F, found)
                  end
              end
          end
        | cmp when cmp > 0 =>
          let
            val (p1, fixbal, found) = search (right, fixbal, found)
          in
            case+ (fixbal, bal) of
            | (F, _) => (CONS (k, d, bal, left, p1), F, found)
            | (T, bal_minus1 ()) =>
              (CONS (k, d, bal_zero (), left, p1), F, found)
            | (T, bal_zero ()) =>
              (CONS (k, d, bal_plus1 (), left, p1), fixbal, found)
            | (T, bal_plus1 ()) =>
              let
                val+ CONS (k1, d1, bal1, left1, right1) = p1
              in
                case+ bal1 of
                | bal_plus1 () =>
                  let
                    val q = CONS (k, d, bal_zero (), left, left1)
                    val q1 = CONS (k1, d1, bal_zero (), q, right1)
                  in
                    (q1, F, found)
                  end
                | _ =>
                  let
                    val p2 = left1
                    val- CONS (k2, d2, bal2, left2, right2) = p2
                    val q = CONS (k, d, minus_pos_bal bal2,
                                  left, left2)
                    val q1 = CONS (k1, d1, minus_neg_bal bal2,
                                   right2, right1)
                    val q2 = CONS (k2, d2, bal_zero (), q, q1)
                  in
                    (q2, F, found)
                  end
              end
          end
        | _ => (CONS (key, data, bal, left, right), F, T)
  in
    if avl_t_is_empty avl then
      (CONS (key, data, bal_zero, NIL, NIL), F)
    else
      let
        prval _ = lemma_avl_t_param avl
        val (avl, _, found) = $effmask_ntm search (avl, F, F)
      in
        (avl, found)
      end
  end

fn {key_t : t@ype} {data_t : t@ype}
avl_t_insert {size : int}
             (avl  : avl_t (key_t, data_t, size),
              key  : key_t,
              data : data_t) :<>
    [sz : pos] avl_t (key_t, data_t, sz) =
  (avl_t_insert_or_replace<key_t><data_t> (avl, key, data)).0

fun {key_t : t@ype} {data_t : t@ype}
push_all_the_way_left (stack : List (avl_t (key_t, data_t)),
                       p     : avl_t (key_t, data_t)) :
    List0 (avl_t (key_t, data_t)) =
  let
    prval _ = lemma_list_param stack
  in
    case+ p of
    | NIL => stack
    | CONS (_, _, _, left, _) =>
      push_all_the_way_left (p :: stack, left)
  end

fun {key_t : t@ype} {data_t : t@ype}
update_generator_stack (stack     : List (avl_t (key_t, data_t)),
                        right     : avl_t (key_t, data_t)) :
    List0 (avl_t (key_t, data_t)) =
  let
    prval _ = lemma_list_param stack
  in
    if avl_t_is_empty right then
      stack
    else
      push_all_the_way_left<key_t><data_t> (stack, right)
  end

fn {key_t : t@ype} {data_t : t@ype}
avl_t_make_data_generator {size : int}
                          (avl  : avl_t (key_t, data_t, size)) :
    () -<cloref1> Option data_t =
  let
    typedef avl_t = avl_t (key_t, data_t)

    val stack = push_all_the_way_left<key_t><data_t> (LNIL, avl)
    val stack_ref = ref stack

    (* Cast stack_ref to its (otherwise untyped) pointer, so it can be
       enclosed within ‘generate’. *)
    val p_stack_ref = $UNSAFE.castvwtp0{ptr} stack_ref

    fun
    generate () :<cloref1> Option data_t =
      let
        (* Restore the type information for stack_ref. *)
        val stack_ref =
          $UNSAFE.castvwtp0{ref (List avl_t)} p_stack_ref

        var stack : List0 avl_t = !stack_ref
        var retval : Option data_t
      in
        begin
          case+ stack of
          | LNIL => retval := None ()
          | p :: tail =>
            let
              val- CONS (_, d, _, left, right) = p
            in
              retval := Some d;
              stack :=
                update_generator_stack<key_t><data_t> (tail, right)
            end
        end;
        !stack_ref := stack;
        retval
      end
  in
    generate
  end

(*------------------------------------------------------------------*)
(* Sets implemented with a hash function, AVL trees and association *)
(* lists.                                                           *)

(* The interface  - - - - - - - - - - - - - - - - - - - - - - - - - *)

(* For simplicity, let us support only 64-bit hashes. *)

typedef hashset_t (key_t : t@ype+) =
  avl_t (uint64, List1 key_t)

extern fun {key_t : t@ype}  (* Implement a hash function with this. *)
hashset_t$hashfunc : key_t -<> uint64

extern fun {key_t : t@ype}     (* Implement key equality with this. *)
hashset_t$key_eq : (key_t, key_t) -<> bool

extern fun
hashset_t_nil :
  {key_t : t@ype}
  () -<> hashset_t key_t

extern fun {key_t : t@ype}
hashset_t_add_member :
  (hashset_t key_t, key_t) -<> hashset_t key_t

(*
"remove_member" is not implemented here, because the trimmed down AVL
tree implementation above does not include deletion. We shall
implement everything else without using a member deletion routine.

extern fun {key_t : t@ype}
hashset_t_remove_member :
  (hashset_t key_t, key_t) -<> hashset_t key_t

Of course you can remove a member by using hashset_t_difference.
*)

extern fun {key_t : t@ype}
hashset_t_has_member :
  (hashset_t key_t, key_t) -<> bool

typedef hashset_t_binary_operation (key_t : t@ype) =
  (hashset_t key_t, hashset_t key_t) -> hashset_t key_t

extern fun {key_t : t@ype}
hashset_t_union : hashset_t_binary_operation key_t

extern fun {key_t : t@ype}
hashset_t_intersection : hashset_t_binary_operation key_t

extern fun {key_t : t@ype}
hashset_t_difference : hashset_t_binary_operation key_t

extern fun {key_t : t@ype}
hashset_t_subset :
  (hashset_t key_t, hashset_t key_t) -> bool

extern fun {key_t : t@ype}
hashset_t_equal :
  (hashset_t key_t, hashset_t key_t) -> bool

(* Note: generators for hashset_t produce their output in unspecified
         order. *)
extern fun {key_t : t@ype}
hashset_t_make_generator :
  hashset_t key_t -> () -<cloref1> Option key_t

(* The implementation - - - - - - - - - - - - - - - - - - - - - - - *)

(* I make no promises that these are the most efficient
   implementations I could devise. They certainly are not! But they
   were easy to write and will work. *)

implement
hashset_t_nil () =
  avl_t_nil ()

fun {key_t  : t@ype}
find_key {n : nat} .<n>.
         (lst : list (key_t, n),
          key : key_t) :<>
    List0 key_t =
  (* This implementation is tail recursive. It will not build up the
     stack. *)
  case+ lst of
  | list_nil () => lst
  | list_cons (head, tail) =>
    if hashset_t$key_eq<key_t> (key, head) then
      lst
    else
      find_key (tail, key)

implement {key_t}
hashset_t_add_member (set, key) =
  (* The following implementation assumes equal keys are
     interchangeable. *)
  let
    implement
    avl_t$compare<uint64> (u, v) =
      if u < v then ~1 else if v < u then 1 else 0
    typedef lst_t = List1 key_t
    val hash = hashset_t$hashfunc<key_t> key
    val lst_opt = avl_t_search_opt<uint64><lst_t> (set, hash)
  in
    case+ lst_opt of
    | Some lst =>
      begin
        case+ find_key<key_t> (lst, key) of
        | list_cons _ => set
        | list_nil () =>
          avl_t_insert<uint64><lst_t>
            (set, hash, list_cons (key, lst))
      end
    | None () =>
      avl_t_insert<uint64><lst_t>
        (set, hash, list_cons (key, list_nil ()))
  end

implement {key_t}
hashset_t_has_member (set, key) =
  let
    implement
    avl_t$compare<uint64> (u, v) =
      if u < v then ~1 else if v < u then 1 else 0
    typedef lst_t = List1 key_t
    val hash = hashset_t$hashfunc<key_t> key
    val lst_opt = avl_t_search_opt<uint64><lst_t> (set, hash)
  in
    case+ lst_opt of
    | None () => false
    | Some lst =>
      begin
        case+ find_key<key_t> (lst, key) of
        | list_nil () => false
        | list_cons _ => true
      end
  end

implement {key_t}
hashset_t_union (u, v) =
  let
    val gen_u = hashset_t_make_generator<key_t> u
    val gen_v = hashset_t_make_generator<key_t> v
    var w : hashset_t key_t = hashset_t_nil ()
    var k_opt : Option key_t
  in
    for (k_opt := gen_u (); option_is_some k_opt; k_opt := gen_u ())
      w := hashset_t_add_member (w, option_unsome k_opt);
    for (k_opt := gen_v (); option_is_some k_opt; k_opt := gen_v ())
      w := hashset_t_add_member (w, option_unsome k_opt);
    w
  end

implement {key_t}
hashset_t_intersection (u, v) =
  let
    val gen_u = hashset_t_make_generator<key_t> u
    var w : hashset_t key_t = hashset_t_nil ()
    var k_opt : Option key_t
  in
    for (k_opt := gen_u (); option_is_some k_opt; k_opt := gen_u ())
      let
        val+ Some k = k_opt
      in
        if hashset_t_has_member<key_t> (v, k) then
          w := hashset_t_add_member (w, k)
      end;
    w
  end

implement {key_t}
hashset_t_difference (u, v) =
  let
    val gen_u = hashset_t_make_generator<key_t> u
    var w : hashset_t key_t = hashset_t_nil ()
    var k_opt : Option key_t
  in
    for (k_opt := gen_u (); option_is_some k_opt; k_opt := gen_u ())
      let
        val+ Some k = k_opt
      in
        if ~hashset_t_has_member<key_t> (v, k) then
          w := hashset_t_add_member (w, k)
      end;
    w
  end

implement {key_t}
hashset_t_subset (u, v) =
  let
    val gen_u = hashset_t_make_generator<key_t> u
    var subset : bool = true
    var done : bool = false
  in
    while (~done)
      case+ gen_u () of
      | None () => done := true
      | Some k =>
        if ~hashset_t_has_member<key_t> (v, k) then
          begin
            subset := false;
            done := true
          end;
    subset
  end

implement {key_t}
hashset_t_equal (u, v) =
  hashset_t_subset<key_t> (u, v)
    && hashset_t_subset<key_t> (v, u)

implement {key_t}
hashset_t_make_generator (set) =
  let
    typedef lst_t = List1 key_t
    typedef lst_t_0 = List0 key_t

    val avl_gen = avl_t_make_data_generator<uint64><lst_t> (set)

    val current_list_ref : ref lst_t_0 = ref (list_nil ())
    val current_list_ptr =
      $UNSAFE.castvwtp0{ptr} current_list_ref
  in
    lam () =>
      let
        val current_list_ref =
          $UNSAFE.castvwtp0{ref lst_t_0} current_list_ptr
      in
        case+ !current_list_ref of
        | list_nil () =>
          begin
            case+ avl_gen () of
            | None () => None ()
            | Some lst =>
              begin
                case+ lst of
                | list_cons (head, tail) =>
                  begin
                    !current_list_ref := tail;
                    Some head
                  end
              end
          end
        | list_cons (head, tail) =>
          begin
            !current_list_ref := tail;
            Some head
          end
      end
  end

(*------------------------------------------------------------------*)

implement
hashset_t$hashfunc<string> (s) =
  string_hash s

implement
hashset_t$key_eq<string> (s, t) =
  s = t

typedef strset_t = hashset_t string

fn {}
strset_t_nil () :<> strset_t =
  hashset_t_nil ()

fn
strset_t_add_member (set    : strset_t,
                     member : string) :<> strset_t =
  hashset_t_add_member<string> (set, member)

fn {}
strset_t_member_add (member : string,
                     set    : strset_t) :<> strset_t =
  strset_t_add_member (set, member)

#define SNIL strset_t_nil ()
infixr ( :: ) ++        (* Right associative, same precedence as :: *)
overload ++ with strset_t_member_add

fn
strset_t_has_member (set    : strset_t,
                     member : string) :<> bool =
  hashset_t_has_member<string> (set, member)
overload [] with strset_t_has_member

fn
strset_t_union (u : strset_t, v : strset_t) : strset_t =
  hashset_t_union<string> (u, v)
overload + with strset_t_union

fn
strset_t_intersection (u : strset_t, v : strset_t) : strset_t =
  hashset_t_intersection<string> (u, v)
infixl ( + ) ^
overload ^ with strset_t_intersection

fn
strset_t_difference (u : strset_t, v : strset_t) : strset_t =
  hashset_t_difference<string> (u, v)
overload - with strset_t_difference

fn
strset_t_subset (u : strset_t, v : strset_t) : bool =
  hashset_t_subset<string> (u, v)
overload <= with strset_t_subset

fn
strset_t_equal (u : strset_t, v : strset_t) : bool =
  hashset_t_equal<string> (u, v)
overload = with strset_t_equal

fn
strset_t_make_generator (set : strset_t) :
    () -<cloref1> Option string =
  hashset_t_make_generator<string> set

fn
strset_t_print (set : strset_t) : void =
  let
    val gen = strset_t_make_generator set
    var s_opt : Option string
    var separator : string = ""
  in
    print! ("#<strset_t ");
    for (s_opt := gen (); option_is_some s_opt; s_opt := gen ())
      case+ s_opt of
      | Some s =>
        begin
          (* The following quick and dirty implemenetation does not
             insert escape sequences. *)
          print! (separator, "\"", s, "\"");
          separator := " "
        end;
    print! (">")
  end

implement
main0 () =
  let
    val set1 =
      "one" ++ "two" ++ "three" ++ "guide" ++ "design" ++ SNIL
    val set2 =
      "ett" ++ "två" ++ "tre" ++ "guide" ++ "design" ++ SNIL
  in
    print! ("set1 = ");
    strset_t_print set1;

    println! ();
    println! ();
    println! ("set1[\"one\"] = ", set1["one"]);
    println! ("set1[\"two\"] = ", set1["two"]);
    println! ("set1[\"three\"] = ", set1["three"]);
    println! ("set1[\"four\"] = ", set1["four"]);

    println! ();
    print! ("set2 = ");
    strset_t_print set2;

    println! ();
    println! ();
    println! ("set2[\"ett\"] = ", set2["ett"]);
    println! ("set2[\"två\"] = ", set2["två"]);
    println! ("set2[\"tre\"] = ", set2["tre"]);
    println! ("set2[\"fyra\"] = ", set2["fyra"]);

    println! ();
    print! ("Union\nset1 + set2 = ");
    strset_t_print (set1 + set2);
    println! ();

    println! ();
    print! ("Intersection\nset1 ^ set2 = ");
    strset_t_print (set1 ^ set2);
    println! ();

    println! ();
    print! ("Difference\nset1 - set2 = ");
    strset_t_print (set1 - set2);
    println! ();

    println! ();
    println! ("Subset");
    println! ("set1 <= set1: ", set1 <= set1);
    println! ("set2 <= set2: ", set2 <= set2);
    println! ("set1 <= set2: ", set1 <= set2);
    println! ("set2 <= set1: ", set2 <= set1);
    println! ("(set1 ^ set2) <= set1: ", (set1 ^ set2) <= set1);
    println! ("(set1 ^ set2) <= set2: ", (set1 ^ set2) <= set2);

    println! ();
    println! ("Equal");
    println! ("set1 = set1: ", set1 = set1);
    println! ("set2 = set2: ", set2 = set2);
    println! ("set1 = set2: ", set1 = set2);
    println! ("set2 = set1: ", set2 = set1);
    println! ("(set1 ^ set2) = (set2 ^ set1): ",
              (set1 ^ set2) = (set2 ^ set1));
    println! ("(set1 ^ set2) = set1: ", (set1 ^ set2) = set1);
    println! ("(set1 ^ set2) = set2: ", (set1 ^ set2) = set2)
  end

(*------------------------------------------------------------------*)