# BitSet — a small packed bit set for natural-number indices.
#
# Words are stored little-endian in a List of U32 values: bit 0 is the low
# bit of the first word.  The list is canonical (zero words at the high end
# are removed), while the public representation remains simple and inspectable.
# Base already provides Word and U32 bit operations; this package only wraps
# them in a growable List facade.
import Base

type BitSet is Data:
  B{words: List<&2, U32>}

# ---- representation helpers ----------------------------------------------

def BitSet.trim.arm(+h: U32, t: List<&2, U32>) -> List<&2, U32>:
  match t:
    case Nil{}:
      Bool.pick(List<&2, U32>, U32.is_zero(h), Nil{}, h <> Nil{})
    case _ <> _:
      h <> t

def BitSet.trim(xs: List<&2, U32>) -> List<&2, U32>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      BitSet.trim.arm(h, BitSet.trim(t))

def BitSet.word.got(got: Maybe<&2, U32>) -> U32:
  match got:
    case None{}:
      0
    case Some{x}:
      x

def BitSet.word(xs: List<&2, U32>, n: Nat) -> U32:
  BitSet.word.got(List.get(&2, U32, xs, n))

def BitSet.mask(n: Nat) -> U32:
  U32.shln(1, n)

def BitSet.location(i: Nat) -> Nat & Nat:
  Nat.divmod(i, 32n)

# ---- constructors / conversion -------------------------------------------

def BitSet.empty() -> BitSet:
  B{Nil{}}

def BitSet.from_words(xs: List<&2, U32>) -> BitSet:
  B{BitSet.trim(xs)}

def BitSet.to_words(bs: BitSet) -> List<&2, U32>:
  match bs:
    case B{xs}:
      xs

# ---- membership and updates ----------------------------------------------

def BitSet.contains.loc(xs: List<&2, U32>, loc: Nat & Nat) -> Bool:
  (wi, bi) = loc
  U32.is_ne(U32.and(BitSet.word(xs, wi), BitSet.mask(bi)), 0)

def BitSet.contains(bs: BitSet, i: Nat) -> Bool:
  match bs:
    case B{xs}:
      BitSet.contains.loc(xs, BitSet.location(i))

def BitSet.contains_u32(bs: BitSet, i: U32) -> Bool:
  BitSet.contains(bs, U32.to_nat(i))

def BitSet.insert.go(
  xs: List<&2, U32>, n: Nat, mask: U32
) -> List<&2, U32>:
  match xs n:
    case Nil{} 0n:
      U32.or(0, mask) <> Nil{}
    case Nil{} 1n+p:
      0 <> BitSet.insert.go(Nil{}, p, mask)
    case h <> t 0n:
      U32.or(h, mask) <> t
    case h <> t 1n+p:
      h <> BitSet.insert.go(t, p, mask)

def BitSet.insert.loc(xs: List<&2, U32>, loc: Nat & Nat) -> BitSet:
  (wi, bi) = loc
  B{BitSet.insert.go(xs, wi, BitSet.mask(bi))}

def BitSet.insert(bs: BitSet, i: Nat) -> BitSet:
  match bs:
    case B{xs}:
      BitSet.insert.loc(xs, BitSet.location(i))

def BitSet.insert_u32(bs: BitSet, i: U32) -> BitSet:
  BitSet.insert(bs, U32.to_nat(i))

def BitSet.singleton(i: Nat) -> BitSet:
  BitSet.insert(BitSet.empty(), i)

def BitSet.singleton_u32(i: U32) -> BitSet:
  BitSet.singleton(U32.to_nat(i))

def BitSet.remove.go(
  xs: List<&2, U32>, n: Nat, mask: U32
) -> List<&2, U32>:
  match xs n:
    case Nil{} _:
      Nil{}
    case h <> t 0n:
      U32.and(h, U32.not(mask)) <> t
    case h <> t 1n+p:
      h <> BitSet.remove.go(t, p, mask)

def BitSet.remove.loc(xs: List<&2, U32>, loc: Nat & Nat) -> BitSet:
  (wi, bi) = loc
  B{BitSet.trim(BitSet.remove.go(xs, wi, BitSet.mask(bi)))}

def BitSet.remove(bs: BitSet, i: Nat) -> BitSet:
  match bs:
    case B{xs}:
      BitSet.remove.loc(xs, BitSet.location(i))

def BitSet.remove_u32(bs: BitSet, i: U32) -> BitSet:
  BitSet.remove(bs, U32.to_nat(i))

# ---- set operations -------------------------------------------------------

def BitSet.union.go(
  xs: List<&2, U32>, ys: List<&2, U32>
) -> List<&2, U32>:
  match xs ys:
    case Nil{} _:
      ys
    case _ Nil{}:
      xs
    case x <> xt y <> yt:
      U32.or(x, y) <> BitSet.union.go(xt, yt)

def BitSet.union(a: BitSet, b: BitSet) -> BitSet:
  match a b:
    case B{x} B{y}:
      B{BitSet.trim(BitSet.union.go(x, y))}

def BitSet.intersect.go(
  xs: List<&2, U32>, ys: List<&2, U32>
) -> List<&2, U32>:
  match xs ys:
    case Nil{} _:
      Nil{}
    case _ Nil{}:
      Nil{}
    case x <> xt y <> yt:
      U32.and(x, y) <> BitSet.intersect.go(xt, yt)

def BitSet.intersect(a: BitSet, b: BitSet) -> BitSet:
  match a b:
    case B{x} B{y}:
      B{BitSet.trim(BitSet.intersect.go(x, y))}

# ---- size / enumeration ---------------------------------------------------

def BitSet.word_size.word(n: Nat, w: Word(n), acc: Nat) -> Nat:
  match n:
    case 0n:
      acc
    case 1n+p:
      match w:
        case WCon{False{}, t}:
          BitSet.word_size.word(p, t, acc)
        case WCon{True{}, t}:
          BitSet.word_size.word(p, t, 1n+acc)

def BitSet.word_size(w: U32) -> Nat:
  match w:
    case U32{x}:
      BitSet.word_size.word(32n, x, 0n)

def BitSet.size.go(xs: List<&2, U32>, acc: Nat) -> Nat:
  match xs:
    case Nil{}:
      acc
    case h <> t:
      BitSet.size.go(t, Nat.add(acc, BitSet.word_size(h)))

def BitSet.size(bs: BitSet) -> Nat:
  match bs:
    case B{xs}:
      BitSet.size.go(xs, 0n)

def BitSet.is_empty.go(xs: List<&2, U32>) -> Bool:
  match xs:
    case Nil{}:
      True{}
    case h <> t:
      Bool.and(U32.is_zero(h), BitSet.is_empty.go(t))

def BitSet.is_empty(bs: BitSet) -> Bool:
  match bs:
    case B{xs}:
      BitSet.is_empty.go(xs)

def BitSet.to_list.word(n: Nat, +base: Nat, w: Word(n)) -> List<&2, Nat>:
  match n:
    case 0n:
      Nil{}
    case 1n+p:
      match w:
        case WCon{False{}, t}:
          BitSet.to_list.word(p, 1n+base, t)
        case WCon{True{}, t}:
          base <> BitSet.to_list.word(p, 1n+base, t)

def BitSet.to_list.u32(base: Nat, w: U32) -> List<&2, Nat>:
  match w:
    case U32{x}:
      BitSet.to_list.word(32n, base, x)

def BitSet.to_list.go(xs: List<&2, U32>, +base: Nat) -> List<&2, Nat>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      List.append(&2, Nat, BitSet.to_list.u32(base, h),
        BitSet.to_list.go(t, Nat.add(base, 32n)))

def BitSet.to_list(bs: BitSet) -> List<&2, Nat>:
  match bs:
    case B{xs}:
      BitSet.to_list.go(xs, 0n)