import Base

# Set.bend — U32 sets as sorted deduped lists.
#
# A set is a wrapper around an ascending, duplicate-free List<&2, U32>.
# Insertion keeps order and drops duplicates via the put-pattern: sput
# never calls back into sinsert_go, so the call graph stays acyclic and
# each insertion visits every element exactly once. Membership is a
# linear scan with inline Bool.pick (the lfind_go shape from listx);
# union folds insert, intersection/difference filter on membership.

type USet is Data:
  US{elems: List<&2, U32>}

def sempty() -> USet:
  US{Nil{}}

# sput: insert-branch. lt/eq compare x against head h; t is the original
# tail, r the recursively-inserted tail. x<h -> x<>h<>t (drop r);
# x==h -> h<>t (drop r, dedup); else h<>r (drop t).
def sput(lt: Bool, eq: Bool, x: U32, h: U32, t: List<&2, U32>, r: List<&2, U32>) -> List<&2, U32>:
  match lt eq:
    case True{} _:
      x <> h <> t
    case False{} True{}:
      h <> t
    case False{} False{}:
      h <> r

def sinsert_go(xs: List<&2, U32>, +x: U32) -> List<&2, U32>:
  match xs:
    case Nil{}:
      x <> Nil{}
    case +h <> +t:
      sput(U32.is_lt(x, h), U32.is_eq(x, h), x, h, t, sinsert_go(t, x))

def sinsert(s: USet, +x: U32) -> USet:
  match s:
    case US{xs}:
      US{sinsert_go(xs, x)}

def smem_list(xs: List<&2, U32>, +x: U32) -> Bool:
  match xs:
    case Nil{}:
      False{}
    case h <> t:
      Bool.pick(Bool, U32.is_eq(x, h), True{}, smem_list(t, x))

def smem_b(+s: USet, +x: U32) -> Bool:
  match s:
    case US{xs}:
      smem_list(xs, x)

def smember(+s: USet, +x: U32) -> USet & Bool:
  (s, smem_b(s, x))

def sunion_go(ys: List<&2, U32>, acc: USet) -> USet:
  match ys:
    case Nil{}:
      acc
    case h <> t:
      sunion_go(t, sinsert(acc, h))

def sunion(+a: USet, +b: USet) -> USet:
  match a b:
    case US{xs} US{ys}:
      sunion_go(ys, US{xs})

def sinter_put(keep: Bool, h: U32, r: List<&2, U32>) -> List<&2, U32>:
  match keep:
    case False{}:
      r
    case True{}:
      h <> r

def sinter_go(xs: List<&2, U32>, +b: USet) -> List<&2, U32>:
  match xs:
    case Nil{}:
      Nil{}
    case +h <> t:
      sinter_put(smem_b(b, h), h, sinter_go(t, b))

def sinter(a: USet, +b: USet) -> USet:
  match a:
    case US{xs}:
      US{sinter_go(xs, b)}

def sdiff_put(found: Bool, h: U32, r: List<&2, U32>) -> List<&2, U32>:
  match found:
    case False{}:
      h <> r
    case True{}:
      r

def sdiff_go(xs: List<&2, U32>, +b: USet) -> List<&2, U32>:
  match xs:
    case Nil{}:
      Nil{}
    case +h <> t:
      sdiff_put(smem_b(b, h), h, sdiff_go(t, b))

def sdiff(a: USet, +b: USet) -> USet:
  match a:
    case US{xs}:
      US{sdiff_go(xs, b)}

def sllen(xs: List<&2, U32>, acc: U32) -> U32:
  match xs:
    case Nil{}:
      acc
    case h <> t:
      sllen(t, (acc + 1 : U32))

def ssize(s: USet) -> U32:
  match s:
    case US{xs}:
      sllen(xs, 0)

law set_mem_hit:
  {smember(sinsert(sempty(), 5), 5) == (US{5 <> Nil{}}, True{}) : USet & Bool}

def set_mem_hit():
  {==}

law set_size_two:
  {ssize(sinsert(sinsert(sempty(), 1), 2)) == 2 : U32}

def set_size_two():
  {==}

law set_size_dedup:
  {ssize(sinsert(sinsert(sempty(), 1), 1)) == 1 : U32}

def set_size_dedup():
  {==}

law ssize_empty:
  {ssize(sempty()) == 0 : U32}

def ssize_empty():
  {==}

# NOTE (dropped ssize_insert_bound, 2 attempts): the bound
# {U32.is_le(ssize(sinsert(s, x)), (ssize(s) + 1 : U32)) == True{}} needs a
# sinsert_go helper lemma, but (1) the direct IH rewrite misses: the goal
# wraps the IH term in sput with computed lt/eq plus an accumulator shift
# (sllen(t, 1) vs sllen(t, 0)); (2) the sput-level split mis-states the
# bound (x<h branch yields len+2, so the lemma must thread the IH as a
# `where` assumption through the False/False branch). Dropping per
# drop-twice rule; function kept, law omitted.