import Base
import ../../lib/logic.bend as L
import ../../lib/nat.bend as N
import ../../lib/u32.bend as U
import ../../lib/array.bend as AR
import ../../lib/array_ext.bend as AX
import ../../lib/list.bend as LL
import ../../../spec/lib/common.bend as SC
import ../dynamic_array/layout.bend as LY

# The canonical block of a list: a perfect tree of depth d whose slots are
# the list's items (Some) followed by empty slots (None).


# the first k slots of the items xs
def fill(-T: Data, +k: Nat, xs: List<&2, T>) -> List<&2, Maybe<&2, T>>:
  match k xs:
    case 0n _:
      Nil{}
    case 1n+j Nil{}:
      Con{None{}, fill(T, j, Nil{})}
    case 1n+j Con{x, r}:
      Con{Some{x}, fill(T, j, r)}

def headm(-T: Data, xs: List<&2, T>) -> Maybe<&2, T>:
  match xs:
    case Nil{}:
      None{}
    case Con{x, r}:
      Some{x}

def mk(-T: Data, +d: Nat, +xs: List<&2, T>) -> AR.Tree<Maybe<&2, T>>:
  match d:
    case 0n:
      AR.TLeaf{headm(T, xs)}
    case 1n+p:
      AR.TNode{mk(T, p, xs), mk(T, p, SC.drop(T, xs, SC.pow2(p)))}

def mk_perfect(-T: Data, +d: Nat, +xs: List<&2, T>) -> {AR.perfect(Maybe<&2, T>, d, mk(T, d, xs)) == True{} : Bool}:
  match d:
    case 0n:
      {==}
    case 1n+p:
      L.and_intro(AR.perfect(Maybe<&2, T>, p, mk(T, p, xs)), AR.perfect(Maybe<&2, T>, p, mk(T, p, SC.drop(T, xs, SC.pow2(p)))), mk_perfect(T, p, xs), mk_perfect(T, p, SC.drop(T, xs, SC.pow2(p))))

def fill_add(-T: Data, +a: Nat, +b: Nat, +xs: List<&2, T>) -> {fill(T, Nat.add(a, b), xs) == SC.append(Maybe<&2, T>, fill(T, a, xs), fill(T, b, SC.drop(T, xs, a))) : List<&2, Maybe<&2, T>>}:
  match a xs:
    case 0n _:
      Equal.cong(List<&2, T>, List<&2, Maybe<&2, T>>, z => fill(T, b, z), xs, SC.drop(T, xs, 0n), Equal.sym(List<&2, T>, SC.drop(T, xs, 0n), xs, AX.drop_zero(T, xs)))
    case 1n+j Nil{}:
      LL.cons_cong(Maybe<&2, T>, None{}, fill(T, Nat.add(j, b), Nil{}), SC.append(Maybe<&2, T>, fill(T, j, Nil{}), fill(T, b, Nil{})), fill_add(T, j, b, Nil{}))
    case 1n+j Con{x, +r}:
      LL.cons_cong(Maybe<&2, T>, Some{x}, fill(T, Nat.add(j, b), r), SC.append(Maybe<&2, T>, fill(T, j, r), fill(T, b, SC.drop(T, r, j))), fill_add(T, j, b, r))

def dbl(+p: Nat) -> {SC.pow2(1n+p) == Nat.add(SC.pow2(p), SC.pow2(p)) : Nat}:
  Equal.trans(Nat, Nat.double(SC.pow2(p)), Nat.add(SC.pow2(p), Nat.add(SC.pow2(p), 0n)), Nat.add(SC.pow2(p), SC.pow2(p)), U.double_pow(SC.pow2(p)), Equal.cong(Nat, Nat, z => Nat.add(SC.pow2(p), z), Nat.add(SC.pow2(p), 0n), SC.pow2(p), N.add_zero(SC.pow2(p))))

def fill_one(-T: Data, +xs: List<&2, T>) -> {fill(T, 1n, xs) == Con{headm(T, xs), Nil{}} : List<&2, Maybe<&2, T>>}:
  match xs:
    case Nil{}:
      {==}
    case Con{x, r}:
      {==}

# the slots of the canonical block
def mk_slots(-T: Data, +d: Nat, +xs: List<&2, T>) -> {AR.slots(Maybe<&2, T>, mk(T, d, xs)) == fill(T, SC.pow2(d), xs) : List<&2, Maybe<&2, T>>}:
  match d:
    case 0n:
      Equal.sym(List<&2, Maybe<&2, T>>, fill(T, 1n, xs), Con{headm(T, xs), Nil{}}, fill_one(T, xs))
    case 1n+p:
      +P = SC.pow2(p)
      %Equal.sym(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, mk(T, p, xs)), fill(T, P, xs), mk_slots(T, p, xs)) : {SC.append(Maybe<&2, T>, _, AR.slots(Maybe<&2, T>, mk(T, p, SC.drop(T, xs, P)))) == fill(T, SC.pow2(1n+p), xs) : List<&2, Maybe<&2, T>>}
      %Equal.sym(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, mk(T, p, SC.drop(T, xs, P))), fill(T, P, SC.drop(T, xs, P)), mk_slots(T, p, SC.drop(T, xs, P))) : {SC.append(Maybe<&2, T>, fill(T, P, xs), _) == fill(T, SC.pow2(1n+p), xs) : List<&2, Maybe<&2, T>>}
      %Equal.sym(Nat, SC.pow2(1n+p), Nat.add(P, P), dbl(p)) : {SC.append(Maybe<&2, T>, fill(T, P, xs), fill(T, P, SC.drop(T, xs, P))) == fill(T, _, xs) : List<&2, Maybe<&2, T>>}
      Equal.sym(List<&2, Maybe<&2, T>>, fill(T, Nat.add(P, P), xs), SC.append(Maybe<&2, T>, fill(T, P, xs), fill(T, P, SC.drop(T, xs, P))), fill_add(T, P, P, xs))

# ---- the layout facts ----

def fill_somes(-T: Data, +k: Nat, +xs: List<&2, T>, +h: {Nat.is_le(SC.length(T, xs), k) == True{} : Bool}) -> {LY.somes(T, fill(T, k, xs)) == xs : List<&2, T>}:
  match k xs:
    case 0n Nil{}:
      {==}
    case 0n Con{x, r}:
      Empty.absurd({Nil{} == Con{x, r} : List<&2, T>}, L.false_true(h))
    case 1n+j Nil{}:
      fill_somes(T, j, Nil{}, N.zero_le(j))
    case 1n+j Con{x, +r}:
      LL.cons_cong(T, x, LY.somes(T, fill(T, j, r)), r, fill_somes(T, j, r, h))

def fill_nones(-T: Data, +k: Nat) -> {LY.nones(T, fill(T, k, Nil{})) == True{} : Bool}:
  match k:
    case 0n:
      {==}
    case 1n+j:
      fill_nones(T, j)

def fill_lay(-T: Data, +k: Nat, +xs: List<&2, T>, +h: {Nat.is_le(SC.length(T, xs), k) == True{} : Bool}) -> {LY.lay(T, fill(T, k, xs), SC.length(T, xs)) == True{} : Bool}:
  match k xs:
    case 0n Nil{}:
      {==}
    case 0n Con{x, r}:
      Empty.absurd({LY.lay(T, Nil{}, 1n+SC.length(T, r)) == True{} : Bool}, L.false_true(h))
    case 1n+j Nil{}:
      fill_nones(T, j)
    case 1n+j Con{x, +r}:
      fill_lay(T, j, r, h)

def fill_len(-T: Data, +k: Nat, +xs: List<&2, T>) -> {SC.length(Maybe<&2, T>, fill(T, k, xs)) == k : Nat}:
  match k xs:
    case 0n _:
      {==}
    case 1n+j Nil{}:
      N.succ_cong(SC.length(Maybe<&2, T>, fill(T, j, Nil{})), j, fill_len(T, j, Nil{}))
    case 1n+j Con{x, +r}:
      N.succ_cong(SC.length(Maybe<&2, T>, fill(T, j, r)), j, fill_len(T, j, r))

def fill_upd(-T: Data, +k: Nat, +xs: List<&2, T>, +i: Nat, +v: T, +h: {Nat.is_lt(i, SC.length(T, xs)) == True{} : Bool}) -> {SC.update(Maybe<&2, T>, fill(T, k, xs), i, Some{v}) == fill(T, k, SC.update(T, xs, i, v)) : List<&2, Maybe<&2, T>>}:
  match k xs i:
    case 0n _ _:
      {==}
    case 1n+j Nil{} _:
      Empty.absurd({SC.update(Maybe<&2, T>, fill(T, 1n+j, Nil{}), i, Some{v}) == fill(T, 1n+j, Nil{}) : List<&2, Maybe<&2, T>>}, N.lt_zero_absurd(i, h))
    case 1n+j Con{x, r} 0n:
      {==}
    case 1n+j Con{x, +r} 1n+q:
      LL.cons_cong(Maybe<&2, T>, Some{x}, SC.update(Maybe<&2, T>, fill(T, j, r), q, Some{v}), fill(T, j, SC.update(T, r, q, v)), fill_upd(T, j, r, q, v, h))

def fill_snoc(-T: Data, +k: Nat, +xs: List<&2, T>, +v: T, +h: {Nat.is_lt(SC.length(T, xs), k) == True{} : Bool}) -> {SC.update(Maybe<&2, T>, fill(T, k, xs), SC.length(T, xs), Some{v}) == fill(T, k, SC.snoc(T, xs, v)) : List<&2, Maybe<&2, T>>}:
  match k xs:
    case 0n _:
      Empty.absurd({SC.update(Maybe<&2, T>, Nil{}, SC.length(T, xs), Some{v}) == Nil{} : List<&2, Maybe<&2, T>>}, N.lt_zero_absurd(SC.length(T, xs), h))
    case 1n+j Nil{}:
      {==}
    case 1n+j Con{x, +r}:
      LL.cons_cong(Maybe<&2, T>, Some{x}, SC.update(Maybe<&2, T>, fill(T, j, r), SC.length(T, r), Some{v}), fill(T, j, SC.snoc(T, r, v)), fill_snoc(T, j, r, v, h))

# ---- trees equal to canonical blocks ----

def mk_eq(-T: Data, +d: Nat, +u: AR.Tree<Maybe<&2, T>>, +xs: List<&2, T>, +pu: {AR.perfect(Maybe<&2, T>, d, u) == True{} : Bool}, +e: {AR.slots(Maybe<&2, T>, u) == fill(T, SC.pow2(d), xs) : List<&2, Maybe<&2, T>>}) -> {u == mk(T, d, xs) : AR.Tree<Maybe<&2, T>>}:
  AX.tree_ext(Maybe<&2, T>, d, u, mk(T, d, xs), pu, mk_perfect(T, d, xs), Equal.trans(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, u), fill(T, SC.pow2(d), xs), AR.slots(Maybe<&2, T>, mk(T, d, xs)), e, Equal.sym(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, mk(T, d, xs)), fill(T, SC.pow2(d), xs), mk_slots(T, d, xs))))

# a slot written below the length
def mk_set(-T: Data, +d: Nat, +xs: List<&2, T>, +i: Nat, +v: T, +h: {Nat.is_lt(i, SC.length(T, xs)) == True{} : Bool}, +hc: {Nat.is_le(SC.length(T, xs), SC.pow2(d)) == True{} : Bool}) -> {AR.upd(Maybe<&2, T>, d, mk(T, d, xs), i, Some{v}) == mk(T, d, SC.update(T, xs, i, v)) : AR.Tree<Maybe<&2, T>>}:
  +hi = N.lt_le_trans(i, SC.length(T, xs), SC.pow2(d), h, hc)
  +es = Equal.trans(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, AR.upd(Maybe<&2, T>, d, mk(T, d, xs), i, Some{v})), SC.update(Maybe<&2, T>, AR.slots(Maybe<&2, T>, mk(T, d, xs)), i, Some{v}), fill(T, SC.pow2(d), SC.update(T, xs, i, v)), AR.upd_slots(Maybe<&2, T>, d, mk(T, d, xs), i, Some{v}, hi, mk_perfect(T, d, xs)), Equal.trans(List<&2, Maybe<&2, T>>, SC.update(Maybe<&2, T>, AR.slots(Maybe<&2, T>, mk(T, d, xs)), i, Some{v}), SC.update(Maybe<&2, T>, fill(T, SC.pow2(d), xs), i, Some{v}), fill(T, SC.pow2(d), SC.update(T, xs, i, v)), Equal.cong(List<&2, Maybe<&2, T>>, List<&2, Maybe<&2, T>>, z => SC.update(Maybe<&2, T>, z, i, Some{v}), AR.slots(Maybe<&2, T>, mk(T, d, xs)), fill(T, SC.pow2(d), xs), mk_slots(T, d, xs)), fill_upd(T, SC.pow2(d), xs, i, v, h)))
  mk_eq(T, d, AR.upd(Maybe<&2, T>, d, mk(T, d, xs), i, Some{v}), SC.update(T, xs, i, v), AR.upd_perfect(Maybe<&2, T>, d, mk(T, d, xs), i, Some{v}, mk_perfect(T, d, xs)), es)

# the slot at the length written (a push with room)
def mk_push(-T: Data, +d: Nat, +xs: List<&2, T>, +v: T, +h: {Nat.is_lt(SC.length(T, xs), SC.pow2(d)) == True{} : Bool}) -> {AR.upd(Maybe<&2, T>, d, mk(T, d, xs), SC.length(T, xs), Some{v}) == mk(T, d, SC.snoc(T, xs, v)) : AR.Tree<Maybe<&2, T>>}:
  +es = Equal.trans(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, AR.upd(Maybe<&2, T>, d, mk(T, d, xs), SC.length(T, xs), Some{v})), SC.update(Maybe<&2, T>, AR.slots(Maybe<&2, T>, mk(T, d, xs)), SC.length(T, xs), Some{v}), fill(T, SC.pow2(d), SC.snoc(T, xs, v)), AR.upd_slots(Maybe<&2, T>, d, mk(T, d, xs), SC.length(T, xs), Some{v}, h, mk_perfect(T, d, xs)), Equal.trans(List<&2, Maybe<&2, T>>, SC.update(Maybe<&2, T>, AR.slots(Maybe<&2, T>, mk(T, d, xs)), SC.length(T, xs), Some{v}), SC.update(Maybe<&2, T>, fill(T, SC.pow2(d), xs), SC.length(T, xs), Some{v}), fill(T, SC.pow2(d), SC.snoc(T, xs, v)), Equal.cong(List<&2, Maybe<&2, T>>, List<&2, Maybe<&2, T>>, z => SC.update(Maybe<&2, T>, z, SC.length(T, xs), Some{v}), AR.slots(Maybe<&2, T>, mk(T, d, xs)), fill(T, SC.pow2(d), xs), mk_slots(T, d, xs)), fill_snoc(T, SC.pow2(d), xs, v, h)))
  mk_eq(T, d, AR.upd(Maybe<&2, T>, d, mk(T, d, xs), SC.length(T, xs), Some{v}), SC.snoc(T, xs, v), AR.upd_perfect(Maybe<&2, T>, d, mk(T, d, xs), SC.length(T, xs), Some{v}, mk_perfect(T, d, xs)), es)

def fill_rep(-T: Data, +k: Nat) -> {fill(T, k, Nil{}) == SC.replicate(Maybe<&2, T>, k, None{}) : List<&2, Maybe<&2, T>>}:
  match k:
    case 0n:
      {==}
    case 1n+j:
      LL.cons_cong(Maybe<&2, T>, None{}, fill(T, j, Nil{}), SC.replicate(Maybe<&2, T>, j, None{}), fill_rep(T, j))

def drop_all(-T: Data, +xs: List<&2, T>, +k: Nat, +h: {Nat.is_le(SC.length(T, xs), k) == True{} : Bool}) -> {SC.drop(T, xs, k) == Nil{} : List<&2, T>}:
  match xs k:
    case Nil{} _:
      {==}
    case Con{x, r} 0n:
      Empty.absurd({Con{x, r} == Nil{} : List<&2, T>}, L.false_true(h))
    case Con{x, +r} 1n+j:
      drop_all(T, r, j, h)

# the doubled block of a full list (old block and an empty half)
def mk_grow(-T: Data, +d: Nat, +xs: List<&2, T>, +hc: {Nat.is_le(SC.length(T, xs), SC.pow2(d)) == True{} : Bool}) -> {AR.TNode{mk(T, d, xs), AR.trep(Maybe<&2, T>, d, None{})} == mk(T, 1n+d, xs) : AR.Tree<Maybe<&2, T>>}:
  +e1 = Equal.trans(AR.Tree<Maybe<&2, T>>, AR.trep(Maybe<&2, T>, d, None{}), mk(T, d, Nil{}), mk(T, d, SC.drop(T, xs, SC.pow2(d))), mk_eq(T, d, AR.trep(Maybe<&2, T>, d, None{}), Nil{}, AR.trep_perfect(Maybe<&2, T>, d, None{}), Equal.trans(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, AR.trep(Maybe<&2, T>, d, None{})), SC.replicate(Maybe<&2, T>, SC.pow2(d), None{}), fill(T, SC.pow2(d), Nil{}), AR.trep_slots(Maybe<&2, T>, d, None{}), Equal.sym(List<&2, Maybe<&2, T>>, fill(T, SC.pow2(d), Nil{}), SC.replicate(Maybe<&2, T>, SC.pow2(d), None{}), fill_rep(T, SC.pow2(d))))), Equal.cong(List<&2, T>, AR.Tree<Maybe<&2, T>>, z => mk(T, d, z), Nil{}, SC.drop(T, xs, SC.pow2(d)), Equal.sym(List<&2, T>, SC.drop(T, xs, SC.pow2(d)), Nil{}, drop_all(T, xs, SC.pow2(d), hc))))
  Equal.cong(AR.Tree<Maybe<&2, T>>, AR.Tree<Maybe<&2, T>>, z => AR.TNode{mk(T, d, xs), z}, AR.trep(Maybe<&2, T>, d, None{}), mk(T, d, SC.drop(T, xs, SC.pow2(d))), e1)

def mk_empty(-T: Data, +d: Nat) -> {AR.trep(Maybe<&2, T>, d, None{}) == mk(T, d, Nil{}) : AR.Tree<Maybe<&2, T>>}:
  mk_eq(T, d, AR.trep(Maybe<&2, T>, d, None{}), Nil{}, AR.trep_perfect(Maybe<&2, T>, d, None{}), Equal.trans(List<&2, Maybe<&2, T>>, AR.slots(Maybe<&2, T>, AR.trep(Maybe<&2, T>, d, None{})), SC.replicate(Maybe<&2, T>, SC.pow2(d), None{}), fill(T, SC.pow2(d), Nil{}), AR.trep_slots(Maybe<&2, T>, d, None{}), Equal.sym(List<&2, Maybe<&2, T>>, fill(T, SC.pow2(d), Nil{}), SC.replicate(Maybe<&2, T>, SC.pow2(d), None{}), fill_rep(T, SC.pow2(d)))))