(* The levelization of a netlist given in blocks (Lemma 2.5 applied): the depth of every signal, the last layer that
   reads it and the width of every layer, from claims that are checked block by block. *)
From Stdlib Require Import List Bool Arith ZArith Lia.
Import ListNotations.
From TH Require Import TwoHead Net Lev.

(* ---- finite sums -------------------------------------------------------------------------------------------- *)
Fixpoint nsum (k : nat) (g : nat -> nat) : nat := match k with O => 0 | S k' => nsum k' g + g k' end.

Lemma nsum_ext : forall k g h, (forall i, i < k -> g i = h i) -> nsum k g = nsum k h.
Proof. induction k as [|k IH]; intros g h H; cbn [nsum]; [reflexivity|]. rewrite (IH g h), (H k); [reflexivity | lia |]. intros i Hi. apply H. lia. Qed.

Lemma nsum_const : forall k c, nsum k (fun _ => c) = k * c.
Proof. induction k as [|k IH]; intro c; cbn [nsum]; [reflexivity|]. rewrite IH. lia. Qed.

Lemma nsum_add : forall k g h, nsum k (fun i => g i + h i) = nsum k g + nsum k h.
Proof. induction k as [|k IH]; intros g h; cbn [nsum]; [reflexivity|]. rewrite IH. lia. Qed.

Lemma nsum_split : forall a b g, nsum (a + b) g = nsum a g + nsum b (fun i => g (a + i)).
Proof.
  intros a b g. induction b as [|b IH]; [rewrite Nat.add_0_r; cbn [nsum]; lia|].
  replace (a + S b) with (S (a + b)) by lia. cbn [nsum]. rewrite IH. lia.
Qed.

Lemma nsum_mul : forall c s g, nsum (c * s) g = nsum c (fun j => nsum s (fun r => g (s * j + r))).
Proof.
  induction c as [|c IH]; intros s g; cbn [nsum]; [reflexivity|]. replace (S c * s) with (c * s + s) by lia.
  rewrite nsum_split, IH. f_equal. apply nsum_ext. intros r Hr. f_equal. lia.
Qed.

Lemma nsum_le : forall k g h, (forall i, i < k -> g i <= h i) -> nsum k g <= nsum k h.
Proof. induction k as [|k IH]; intros g h H; cbn [nsum]; [lia|]. pose proof (IH g h (fun i Hi => H i ltac:(lia))). pose proof (H k ltac:(lia)). lia. Qed.

Lemma nsum_zero : forall k g, (forall i, i < k -> g i = 0) -> nsum k g = 0.
Proof. intros k g H. rewrite (nsum_ext k g (fun _ => 0)) by exact H. rewrite nsum_const. lia. Qed.

Lemma nsum_one : forall k g j, j < k -> (forall i, i < k -> i <> j -> g i = 0) -> nsum k g = g j.
Proof.
  induction k as [|k IH]; intros g j Hj H; [lia|]. cbn [nsum]. destruct (Nat.eq_dec j k) as [->|E].
  - rewrite nsum_zero; [lia|]. intros i Hi. apply H; lia.
  - rewrite (IH g j) by (lia || (intros i Hi Hne; apply H; lia)). rewrite (H k) by lia. lia.
Qed.

Lemma nsum_swap : forall a b (f : nat -> nat -> nat), nsum a (fun i => nsum b (f i)) = nsum b (fun j => nsum a (fun i => f i j)).
Proof.
  induction a as [|a IH]; intros b f; cbn [nsum].
  - symmetry. apply nsum_zero. reflexivity.
  - rewrite IH, <- nsum_add. reflexivity.
Qed.

Lemma nsum_eq_split : forall k a u v, a < k -> nsum k (fun j => if Nat.eqb j a then u else v) = (k - 1) * v + u.
Proof.
  induction k as [|k IH]; intros a u v Ha; [lia|]. cbn [nsum]. destruct (Nat.eqb_spec k a) as [<-|E].
  - rewrite (nsum_ext k _ (fun _ => v)) by (intros i Hi; destruct (Nat.eqb_spec i k); [lia | reflexivity]).
    rewrite nsum_const. lia.
  - rewrite IH by lia. destruct k; [lia|]. lia.
Qed.

Lemma nsum_lt_split : forall k a u v, a <= k -> nsum k (fun j => if j <? a then u else v) = a * u + (k - a) * v.
Proof.
  intros k a u v Ha. replace k with (a + (k - a)) at 1 by lia. rewrite nsum_split.
  rewrite (nsum_ext a _ (fun _ => u)) by (intros i Hi; destruct (Nat.ltb_spec i a); [reflexivity | lia]).
  rewrite (nsum_ext (k - a) _ (fun _ => v)) by (intros i Hi; destruct (Nat.ltb_spec (a + i) a); [lia | reflexivity]).
  rewrite !nsum_const. lia.
Qed.

Lemma nsum_succ : forall k g, nsum (S k) g = nsum k g + g k.
Proof. reflexivity. Qed.

Lemma nsum_mulc : forall k f c, nsum k (fun i => f i * c) = nsum k f * c.
Proof. induction k as [|k IH]; intros f c; cbn [nsum]; [reflexivity|]. rewrite IH. lia. Qed.

Lemma fold_nsum_seq : forall m (f : nat -> nat), fold_right Nat.add 0 (map f (seq 0 m)) = nsum m f.
Proof.
  assert (G : forall l c, fold_right Nat.add c l = fold_right Nat.add 0 l + c)
    by (induction l as [|y l IHl]; intro c; cbn [fold_right]; [reflexivity | rewrite IHl; lia]).
  induction m as [|m IH]; intro f; [reflexivity|]. rewrite seq_S, map_app, fold_right_app. cbn [map fold_right nsum].
  rewrite G, IH. cbn [Nat.add]. lia.
Qed.

Lemma length_filter_seq : forall (f : nat -> bool) a k, length (filter f (seq a k)) = nsum k (fun i => Nat.b2n (f (a + i))).
Proof.
  intros f a k. revert a. induction k as [|k IH]; intro a; [reflexivity|].
  rewrite seq_S, filter_app, length_app, IH. cbn [nsum filter]. destruct (f (a + k)); cbn [length Nat.b2n]; lia.
Qed.

Definition usu (bs : list block) (x : sref) : su := match x with SIn _ => su0 | SRef b j r => sunit bs b j r end.

(* ---- the rows of a layered map by their sources ------------------------------------------------------------ *)
Definition isUn (c : src) : bool := match c with Un _ => true | Cp _ => false end.
Definition rowlen (us : list tu) (n : nat) (c : src) : nat := match c with Cp _ => 1 | Un s => length (tin (nth (s - n) us tu0)) end.

Lemma count_Un : forall (L : list src) n m, (forall s, In (Un s) L -> n <= s < n + m) ->
  length (filter isUn L) = nsum m (fun k => length (filter (isun (n + k)) L)).
Proof.
  induction L as [|c t IH]; intros n m H; cbn [filter length].
  - symmetry. apply nsum_zero. reflexivity.
  - rewrite (nsum_ext m _ (fun k => (if isun (n + k) c then 1 else 0) + length (filter (isun (n + k)) t)))
      by (intros k _; cbn [filter]; destruct (isun (n + k) c); reflexivity).
    rewrite nsum_add, <- IH by (intros s Hs; apply H; right; exact Hs).
    destruct c as [s|s]; cbn [isUn isun].
    + rewrite nsum_zero by reflexivity. lia.
    + specialize (H s (or_introl eq_refl)). rewrite (nsum_one m _ (s - n)); [| lia |].
      * replace (n + (s - n)) with s by lia. rewrite Nat.eqb_refl. cbn [length]. lia.
      * intros k Hk E. destruct (Nat.eqb_spec s (n + k)); [lia | reflexivity].
Qed.

Lemma sum_rowlen : forall (L : list src) us n m, (forall s, In (Un s) L -> n <= s < n + m) ->
  fold_right Nat.add 0 (map (rowlen us n) L)
  = length (filter (fun c => negb (isUn c)) L) + nsum m (fun k => length (filter (isun (n + k)) L) * length (tin (nth k us tu0))).
Proof.
  induction L as [|c t IH]; intros us n m H; cbn [map fold_right filter length].
  - rewrite nsum_zero by reflexivity. reflexivity.
  - rewrite (nsum_ext m _ (fun k => (if isun (n + k) c then length (tin (nth k us tu0)) else 0)
                                     + length (filter (isun (n + k)) t) * length (tin (nth k us tu0))))
      by (intros k _; cbn [filter]; destruct (isun (n + k) c); cbn [length]; lia).
    rewrite nsum_add, (IH us n m) by (intros s Hs; apply H; right; exact Hs).
    destruct c as [s|s]; cbn [isUn isun negb rowlen].
    + rewrite (nsum_zero m (fun _ => 0)) by reflexivity. cbn [length]. lia.
    + specialize (H s (or_introl eq_refl)).
      rewrite (nsum_one m (fun k => if Nat.eqb s (n + k) then length (tin (nth k us tu0)) else 0) (s - n)); [| lia |].
      * replace (n + (s - n)) with s by lia. rewrite Nat.eqb_refl. lia.
      * intros k Hk E. destruct (Nat.eqb_spec s (n + k)); [lia | reflexivity].
Qed.

Lemma fold_add : forall (l : list nat) c, fold_right Nat.add c l = fold_right Nat.add 0 l + c.
Proof. induction l as [|a t IH]; intro c; cbn [fold_right]; [reflexivity|]. rewrite IH. lia. Qed.

Lemma fold_add_app : forall (l1 l2 : list nat), fold_right Nat.add 0 (l1 ++ l2) = fold_right Nat.add 0 l1 + fold_right Nat.add 0 l2.
Proof. intros l1 l2. rewrite fold_right_app, fold_add. lia. Qed.

(* the rows of Lev(N) and their entries, by their sources *)
Lemma lev_rowcount : forall n us nx, length (concat (lev n us nx)) = length (concat (map snd (lsrc n us nx))).
Proof.
  intros n us nx. unfold lev. induction (lsrc n us nx) as [|[l cs] P IH]; [reflexivity|].
  cbn [map concat fst snd]. rewrite !length_app, length_map, IH. reflexivity.
Qed.

Lemma lev_rowlens : forall n us nx,
  fold_right Nat.add 0 (map (fun row => length (tin row)) (concat (lev n us nx)))
  = fold_right Nat.add 0 (map (rowlen us n) (concat (map snd (lsrc n us nx)))).
Proof.
  intros n us nx. unfold lev. induction (lsrc n us nx) as [|[l cs] P IH]; [reflexivity|].
  cbn [map concat fst snd]. rewrite !map_app, !fold_add_app, IH. f_equal. f_equal. rewrite map_map.
  apply map_ext. intros [s|s]; cbn [mkrow rowlen idr reidx tin]; [reflexivity|]. apply length_map.
Qed.

Lemma boff_S_gen : forall bs b, b < length bs -> boff bs (S b) = boff bs b + bcount (nth b bs blk0) * bsize (nth b bs blk0).
Proof.
  induction bs as [|blk t IH]; intros b Hb; cbn [length] in Hb; [lia|].
  destruct b as [|b]; cbn [boff nth]; [assert (boff t 0 = 0) by (destruct t; reflexivity); lia|]. rewrite IH by lia. lia.
Qed.

Definition isref (x : sref) : bool := match x with SIn _ => false | SRef _ _ _ => true end.

Section BL.
Variable bs : list block.
Variable n : nat.
Variable out : nat -> sref.
Hypothesis Hwf : wf bs n.
Hypothesis Hout : forall i, i < n -> valid bs n (out i).

Definition fus : list tu := compile n bs.
Definition fnx : list nat := map (fun i => idx n bs (out i)) (seq 0 n).
Definition nun : nat := boff bs (length bs).

Definition cnt (b : nat) : nat := bcount (nth b bs blk0).
Definition sz (b : nat) : nat := bsize (nth b bs blk0).

(* the sum of g over the units, block by block *)
Definition bsum (g : sref -> nat) : nat := nsum (length bs) (fun b => nsum (cnt b) (fun j => nsum (sz b) (fun r => g (SRef b j r)))).

Lemma shaped_all : Forall shaped bs.
Proof. exact (proj1 Hwf). Qed.

Lemma fus_len : length fus = nun.
Proof. unfold fus, nun. apply length_compile. exact shaped_all. Qed.

Lemma fwfn : wfn n fus.
Proof. intros k Hk. exact (wf_reads_below bs n Hwf k Hk). Qed.

Lemma fnx_len : length fnx = n.
Proof. unfold fnx. rewrite length_map, length_seq. reflexivity. Qed.

Lemma boff_S : forall b, b < length bs -> boff bs (S b) = boff bs b + cnt b * sz b.
Proof. intros b Hb. unfold cnt, sz. apply boff_S_gen. exact Hb. Qed.

Lemma nsum_units : forall (g : nat -> nat), nsum nun g = bsum (fun x => match x with SRef b j r => g (boff bs b + sz b * j + r) | SIn _ => 0 end).
Proof.
  intro g. unfold nun, bsum.
  assert (G : forall k, k <= length bs -> nsum (boff bs k) g
            = nsum k (fun b => nsum (cnt b) (fun j => nsum (sz b) (fun r => g (boff bs b + sz b * j + r))))).
  { induction k as [|k IH]; intro Hk; [destruct bs; reflexivity|].
    rewrite boff_S by lia. cbn [nsum]. rewrite nsum_split, IH by lia. f_equal.
    rewrite (Nat.mul_comm (cnt k) (sz k)) at 1. rewrite Nat.mul_comm, nsum_mul. apply nsum_ext. intros j Hj.
    apply nsum_ext. intros r Hr. f_equal. lia. }
  apply G. lia.
Qed.

Lemma nsum_bsum_swap : forall m (g : nat -> sref -> nat), nsum m (fun l => bsum (g l)) = bsum (fun x => nsum m (fun l => g l x)).
Proof.
  intros m g. unfold bsum. rewrite nsum_swap. apply nsum_ext. intros b _. rewrite nsum_swap. apply nsum_ext.
  intros j _. apply nsum_swap.
Qed.

Lemma bsum_ext : forall g1 g2, (forall b j r, valid bs n (SRef b j r) -> g1 (SRef b j r) = g2 (SRef b j r)) ->
  bsum g1 = bsum g2.
Proof.
  intros g1 g2 H. unfold bsum. apply nsum_ext. intros b Hb. apply nsum_ext. intros j Hj. apply nsum_ext. intros r Hr.
  apply H. split; [exact Hb | split; [exact Hj | exact Hr]].
Qed.

(* locating a unit *)
Lemma loc : forall k, k < nun -> exists b j r, valid bs n (SRef b j r) /\ k = boff bs b + sz b * j + r.
Proof.
  unfold nun, sz, valid. pose proof shaped_all as Hsh. clear -Hsh. intros k Hk. revert k Hk.
  induction Hsh as [|blk t Hs Ht IH]; intros k Hk; cbn [boff length] in Hk; [lia|].
  destruct (Nat.lt_ge_cases k (bcount blk * bsize blk)) as [Hl|Hg].
  - exists 0, (k / bsize blk), (k mod bsize blk).
    assert (bsize blk <> 0) by (intro E; rewrite E in Hl; lia).
    cbn [nth length boff]. split; [split; [lia|]; split; [apply Nat.Div0.div_lt_upper_bound; nia | apply Nat.mod_upper_bound; exact H]|].
    rewrite (Nat.div_mod k (bsize blk)) at 1 by exact H. lia.
  - destruct (IH (k - bcount blk * bsize blk) ltac:(lia)) as (b & j & r & (Hb & Hj & Hr) & Hk').
    exists (S b), j, r. cbn [nth length boff]. split; [split; [lia|]; split; assumption|]. lia.
Qed.

Lemma idx_ref : forall b j r, idx n bs (SRef b j r) = n + (boff bs b + sz b * j + r).
Proof. intros b j r. cbn [idx]. unfold sz. lia. Qed.

Lemma fus_nth : forall b j r, valid bs n (SRef b j r) ->
  nth (boff bs b + sz b * j + r) fus tu0 = flat_unit n bs (sunit bs b j r).
Proof. intros b j r (Hb & Hj & Hr). unfold fus, sz. apply compile_nth; [exact shaped_all | exact Hb | exact Hj | exact Hr]. Qed.

Lemma unit_lt : forall b j r, valid bs n (SRef b j r) -> boff bs b + sz b * j + r < nun.
Proof.
  intros b j r Hv. pose proof (idx_lt_total bs n (SRef b j r) Hv) as H. rewrite idx_ref in H. unfold nun. lia.
Qed.

(* the units and the edges, block by block *)
Lemma units_bsum : length fus = bsum (fun _ => 1).
Proof.
  rewrite fus_len. rewrite <- (Nat.mul_1_r nun) at 1. rewrite <- nsum_const, nsum_units.
  apply bsum_ext. reflexivity.
Qed.

Lemma edges_bsum : nsum (length fus) (fun k => length (tin (nth k fus tu0))) = bsum (fun x => length (sin (usu bs x))).
Proof.
  rewrite fus_len, nsum_units. apply bsum_ext. intros b j r Hv. cbn [usu].
  rewrite fus_nth by exact Hv. unfold flat_unit. cbn [tin]. apply length_map.
Qed.

(* the predecessors of a unit are valid signals with smaller indices *)
Lemma pred_ok : forall b j r x z, valid bs n (SRef b j r) -> In (x, z) (sin (sunit bs b j r)) ->
  valid bs n x /\ idx n bs x < idx n bs (SRef b j r).
Proof.
  intros b j r x z Hv Hin. destruct Hv as (Hb & Hj & Hr). destruct Hwf as [Hsh Hrd].
  pose proof (Hrd b j r x z Hb Hj Hr Hin) as Hbk. destruct x as [i|b' j' r']; cbn [reads_back valid] in *.
  - split; [exact Hbk|]. cbn [idx]. lia.
  - destruct Hbk as [(Hb' & Hj' & Hr')|(-> & -> & Hr')].
    + split; [split; [lia|]; split; assumption|].
      pose proof (boff_mono bs b b' Hb' ltac:(lia)) as Hm. cbn [idx]. nia.
    + split; [split; [lia|]; split; lia|]. cbn [idx]. lia.
Qed.

Lemma idx_inj : forall x y, valid bs n x -> valid bs n y -> idx n bs x = idx n bs y -> x = y.
Proof.
  assert (L : forall b j r b' j' r', valid bs n (SRef b j r) -> valid bs n (SRef b' j' r') -> b < b' ->
            idx n bs (SRef b j r) < idx n bs (SRef b' j' r')).
  { intros b j r b' j' r' (_ & Hj & Hr) (Hb' & _ & _) Hlt.
    pose proof (boff_mono bs b' b Hlt ltac:(lia)) as Hm. cbn [idx]. nia. }
  intros [i|b j r] [i'|b' j' r'] Hx Hy E; cbn [idx valid] in E, Hx, Hy.
  - congruence.
  - lia.
  - lia.
  - destruct (Nat.lt_trichotomy b b') as [Hlt|[<-|Hgt]].
    + pose proof (L b j r b' j' r' Hx Hy Hlt). cbn [idx] in *. lia.
    + assert (Hs : bsize (nth b bs blk0) * j + r = bsize (nth b bs blk0) * j' + r') by lia.
      destruct Hx as (_ & _ & Hr), Hy as (_ & _ & Hr').
      destruct (Nat.lt_trichotomy j j') as [Hj|[<-|Hj]]; [nia | | nia].
      assert (r = r') by lia. subst. reflexivity.
    + pose proof (L b' j' r' b j r Hy Hx Hgt). cbn [idx] in *. lia.
Qed.

(* ---- depths ----------------------------------------------------------------------------------------- *)
Lemma dep_su : forall b j r, valid bs n (SRef b j r) ->
  dep n fus (idx n bs (SRef b j r)) = S (maxl (map (fun p => dep n fus (idx n bs (fst p))) (sin (sunit bs b j r)))).
Proof.
  intros b j r Hv. rewrite idx_ref. rewrite dep_unit; [| exact fwfn | rewrite fus_len; exact (unit_lt b j r Hv)].
  rewrite fus_nth by exact Hv. unfold flat_unit. cbn [tin]. rewrite map_map. reflexivity.
Qed.

(* the depth recursion from bounds on each unit: every predecessor lies below it, one of them just below *)
Lemma cd_rec : forall (cd : sref -> nat) (u : su) c,
  (forall p z, In (p, z) (sin u) -> S (cd p) <= c) -> (exists p z, In (p, z) (sin u) /\ S (cd p) = c) ->
  c = S (maxl (map (fun p => cd (fst p)) (sin u))).
Proof.
  intros cd u c Hle (p0 & z0 & Hin0 & E0). apply Nat.le_antisymm.
  - rewrite <- E0. apply le_n_S. apply maxl_ge. apply in_map_iff. exists (p0, z0). split; [reflexivity | exact Hin0].
  - assert (maxl (map (fun p => cd (fst p)) (sin u)) <= c - 1); [|lia].
    apply maxl_le. intros a Ha. apply in_map_iff in Ha as ([p z] & <- & Hin). specialize (Hle p z Hin). cbn [fst]. lia.
Qed.

Section Depth.
Variable cd : sref -> nat.
Hypothesis Hcd0 : forall i, cd (SIn i) = 0.
Hypothesis Hcd : forall b j r, valid bs n (SRef b j r) ->
  cd (SRef b j r) = S (maxl (map (fun p => cd (fst p)) (sin (sunit bs b j r)))).

Theorem dep_cd : forall x, valid bs n x -> dep n fus (idx n bs x) = cd x.
Proof.
  assert (G : forall k x, idx n bs x < k -> valid bs n x -> dep n fus (idx n bs x) = cd x).
  { induction k as [|k IH]; intros x Hk Hv; [lia|]. destruct x as [i|b j r].
    - cbn [idx]. rewrite Hcd0. apply dep_state. exact Hv.
    - rewrite dep_su by exact Hv. rewrite Hcd by exact Hv. f_equal. f_equal. apply map_ext_in. intros [p z] Hin.
      cbn [fst]. destruct (pred_ok b j r p z Hv Hin) as [Hvp Hlt]. apply IH; [lia | exact Hvp]. }
  intros x Hv. exact (G (S (idx n bs x)) x (Nat.lt_succ_diag_r _) Hv).
Qed.

Theorem dmax_cd : forall D, (forall b j r, valid bs n (SRef b j r) -> cd (SRef b j r) <= D) ->
  (exists b j r, valid bs n (SRef b j r) /\ cd (SRef b j r) = D) -> dmax n fus = D.
Proof.
  intros D Hle (b0 & j0 & r0 & Hv0 & E0). apply Nat.le_antisymm.
  - apply maxl_le. intros a Ha. apply in_map_iff in Ha as (k & <- & Hk). apply in_seq in Hk. rewrite fus_len in Hk.
    destruct (loc k ltac:(lia)) as (b & j & r & Hv & ->). rewrite <- idx_ref, dep_cd by exact Hv. apply Hle. exact Hv.
  - rewrite <- E0, <- (dep_cd _ Hv0), idx_ref. apply maxl_ge. apply in_map_iff. exists (boff bs b0 + sz b0 * j0 + r0).
    split; [reflexivity|]. apply in_seq. rewrite fus_len. pose proof (unit_lt b0 j0 r0 Hv0). lia.
Qed.

End Depth.

(* ---- every unit contributes to a next value ---------------------------------------------------------- *)
Section Need.
Hypothesis Hrd : forall b j r, valid bs n (SRef b j r) -> (exists i, i < n /\ out i = SRef b j r) \/
  (exists b' j' r' z, valid bs n (SRef b' j' r') /\ In (SRef b j r, z) (sin (sunit bs b' j' r'))).

Theorem need_all : forall b j r, valid bs n (SRef b j r) -> need n fus fnx (idx n bs (SRef b j r)) = true.
Proof.
  assert (G : forall k b j r, n + nun - idx n bs (SRef b j r) < k -> valid bs n (SRef b j r) ->
            need n fus fnx (idx n bs (SRef b j r)) = true).
  { induction k as [|k IH]; intros b j r Hk Hv; [lia|].
    destruct (Hrd b j r Hv) as [(i & Hi & Ei)|(b' & j' & r' & z & Hvr & Hin)].
    - apply need_nx. unfold fnx. apply in_map_iff. exists i. split; [rewrite Ei; reflexivity | apply in_seq; lia].
    - destruct (pred_ok b' j' r' (SRef b j r) z Hvr Hin) as [_ Hlt].
      assert (Hn' : need n fus fnx (idx n bs (SRef b' j' r')) = true).
      { apply IH; [| exact Hvr]. pose proof (idx_lt_total bs n _ Hvr). unfold nun in *. lia. }
      rewrite idx_ref in Hn'. apply (need_pred n fus fnx fwfn (boff bs b' + sz b' * j' + r') (idx n bs (SRef b j r)) z).
      + rewrite fus_len. exact (unit_lt b' j' r' Hvr).
      + exact Hn'.
      + rewrite fus_nth by exact Hvr. unfold flat_unit. cbn [tin]. apply in_map_iff. exists (SRef b j r, z).
        split; [reflexivity | exact Hin]. }
  intros b j r Hv. apply (G (S (n + nun)) b j r); [lia | exact Hv].
Qed.

End Need.

(* ---- the last layer that reads each signal ------------------------------------------------------------- *)
Lemma reads_su : forall b j r x, valid bs n (SRef b j r) -> valid bs n x ->
  reads fus (boff bs b + sz b * j + r) (idx n bs x) = true <-> exists z, In (x, z) (sin (sunit bs b j r)).
Proof.
  intros b j r x Hv Hx. rewrite reads_in. rewrite fus_nth by exact Hv. unfold flat_unit. cbn [tin]. split.
  - intros (z & Hin). apply in_map_iff in Hin as ([p z'] & E & Hin). cbn [fst snd] in E. injection E as Ep <-.
    exists z'. destruct (pred_ok b j r p z' Hv Hin) as [Hvp _]. rewrite (idx_inj p x Hvp Hx Ep) in Hin. exact Hin.
  - intros (z & Hin). exists z. apply in_map_iff. exists (x, z). split; [reflexivity | exact Hin].
Qed.

Section Lam.
Variable cd : sref -> nat.
Hypothesis Hdep : forall x, valid bs n x -> dep n fus (idx n bs x) = cd x.
Hypothesis Hneed : forall b j r, valid bs n (SRef b j r) -> need n fus fnx (idx n bs (SRef b j r)) = true.
Variable D : nat.
Hypothesis HD : dmax n fus = D.
Variable cl : sref -> nat.
Hypothesis Hup : forall b j r p z, valid bs n (SRef b j r) -> In (p, z) (sin (sunit bs b j r)) -> cd (SRef b j r) <= cl p.
Hypothesis Hcl_out : forall i, i < n -> cl (out i) = S D.
Hypothesis Hlow : forall x, valid bs n x -> (forall i, i < n -> out i <> x) ->
  cl x = 0 \/ exists b j r z, valid bs n (SRef b j r) /\ In (x, z) (sin (sunit bs b j r)) /\ cd (SRef b j r) = cl x.

Lemma unit_dep_le : forall k, k < nun -> dep n fus (n + k) <= D.
Proof. intros k Hk. rewrite <- HD. apply dep_le_dmax. rewrite fus_len. exact Hk. Qed.

Theorem lam_cl : forall x, valid bs n x -> lam n fus fnx (idx n bs x) = cl x.
Proof.
  intros x Hx. unfold lam.
  (* the readers *)
  assert (Hr : maxl (map (fun k => if need n fus fnx (n + k) && reads fus k (idx n bs x) then dep n fus (n + k) else 0)
                         (seq 0 (length fus))) <= cl x).
  { apply maxl_le. intros a Ha. apply in_map_iff in Ha as (k & <- & Hk). apply in_seq in Hk. rewrite fus_len in Hk.
    destruct (loc k ltac:(lia)) as (b & j & r & Hv & ->).
    destruct (need n fus fnx _ && reads fus _ (idx n bs x)) eqn:E; [|lia].
    apply andb_true_iff in E as [_ E]. apply (reads_su b j r x Hv Hx) in E as (z & Hin).
    rewrite <- idx_ref, Hdep by exact Hv. exact (Hup b j r x z Hv Hin). }
  destruct (existsb (Nat.eqb (idx n bs x)) fnx) eqn:Eo.
  - apply existsb_exists in Eo as (s & Hs & E). apply Nat.eqb_eq in E. unfold fnx in Hs.
    apply in_map_iff in Hs as (i & Ei & Hi). apply in_seq in Hi. rewrite <- Ei in E.
    apply (idx_inj x (out i) Hx (Hout i ltac:(lia))) in E. subst x. rewrite Hcl_out by lia. rewrite HD.
    apply Nat.max_l. eapply Nat.le_trans; [|apply Nat.le_succ_diag_r]. apply maxl_le. intros a Ha.
    apply in_map_iff in Ha as (k & <- & Hk). apply in_seq in Hk. rewrite fus_len in Hk.
    destruct (_ && _); [apply unit_dep_le; lia | lia].
  - cbn [Nat.max]. apply Nat.le_antisymm; [exact Hr|].
    assert (Hno : forall i, i < n -> out i <> x).
    { intros i Hi E. subst x. assert (existsb (Nat.eqb (idx n bs (out i))) fnx = true) as C.
      { apply existsb_exists. exists (idx n bs (out i)). split; [|apply Nat.eqb_refl].
        unfold fnx. apply in_map_iff. exists i. split; [reflexivity | apply in_seq; lia]. }
      rewrite C in Eo. discriminate. }
    destruct (Hlow x Hx Hno) as [Z|(b & j & r & z & Hv & Hin & E)]; [rewrite Z; lia|].
    rewrite <- E. apply maxl_ge. apply in_map_iff. exists (boff bs b + sz b * j + r).
    split; [|apply in_seq; rewrite fus_len; pose proof (unit_lt b j r Hv); lia].
    rewrite <- idx_ref, Hneed by exact Hv.
    replace (reads fus (boff bs b + sz b * j + r) (idx n bs x)) with true
      by (symmetry; apply (reads_su b j r x Hv Hx); exists z; exact Hin).
    cbn [andb]. apply Hdep. exact Hv.
Qed.

(* the rows of layer l: the units of depth l and the copies of the signals of smaller depth read after l *)
Theorem live_cl : forall l x, valid bs n x ->
  live n fus fnx l (idx n bs x) = isref x && Nat.eqb (cd x) l || Nat.ltb (cd x) l && Nat.ltb l (cl x).
Proof.
  intros l x Hx. unfold live, isunit. rewrite Hdep, lam_cl by exact Hx. destruct x as [i|b j r].
  - cbn [idx isref valid] in *. destruct (Nat.leb_spec n i); [lia|]. reflexivity.
  - rewrite Hneed by exact Hx. rewrite idx_ref. destruct (Nat.leb_spec n (n + (boff bs b + sz b * j + r))); [|lia].
    cbn [isref andb]. rewrite andb_true_r. reflexivity.
Qed.

Theorem width_bsum : forall l, 1 <= l ->
  width n fus fnx l = nsum n (fun i => Nat.b2n (Nat.ltb (cd (SIn i)) l && Nat.ltb l (cl (SIn i))))
                      + bsum (fun x => Nat.b2n (Nat.eqb (cd x) l || Nat.ltb (cd x) l && Nat.ltb l (cl x))).
Proof.
  intros [|l] Hl; [lia|]. cbn [width]. unfold NS. rewrite fus_len, seq_app, filter_app, length_app.
  rewrite !length_filter_seq. f_equal.
  - apply nsum_ext. intros i Hi. cbn [Nat.add]. replace i with (idx n bs (SIn i)) at 1 by reflexivity.
    rewrite live_cl by (cbn [valid]; exact Hi). reflexivity.
  - rewrite nsum_units. unfold bsum. apply nsum_ext. intros b Hb. apply nsum_ext. intros j Hj. apply nsum_ext.
    intros r Hr. assert (Hv : valid bs n (SRef b j r)) by (split; [exact Hb | split; [exact Hj | exact Hr]]).
    replace (0 + n + (boff bs b + sz b * j + r)) with (idx n bs (SRef b j r)) by (rewrite idx_ref; lia).
    rewrite live_cl by exact Hv. reflexivity.
Qed.

End Lam.

(* ---- netlists with units that contribute to nothing ------------------------------------------------------ *)
Section Need2.
Variable nd : sref -> bool.
Hypothesis Hnd_up : forall b j r, valid bs n (SRef b j r) -> nd (SRef b j r) = true ->
  (exists i, i < n /\ out i = SRef b j r) \/
  (exists b' j' r' z, valid bs n (SRef b' j' r') /\ nd (SRef b' j' r') = true /\ In (SRef b j r, z) (sin (sunit bs b' j' r'))).
Hypothesis Hnd_down : forall b j r, valid bs n (SRef b j r) -> nd (SRef b j r) = false ->
  (forall i, i < n -> out i <> SRef b j r) /\
  forall b' j' r' z, valid bs n (SRef b' j' r') -> In (SRef b j r, z) (sin (sunit bs b' j' r')) -> nd (SRef b' j' r') = false.

Theorem need_eq : forall b j r, valid bs n (SRef b j r) -> need n fus fnx (idx n bs (SRef b j r)) = nd (SRef b j r).
Proof.
  assert (G : forall m b j r, n + nun - idx n bs (SRef b j r) < m -> valid bs n (SRef b j r) ->
            need n fus fnx (idx n bs (SRef b j r)) = nd (SRef b j r)).
  { induction m as [|m IH]; intros b j r Hm Hv; [lia|]. pose proof (idx_lt_total bs n _ Hv) as Hlt.
    destruct (nd (SRef b j r)) eqn:Hnd.
    - destruct (Hnd_up b j r Hv Hnd) as [(i & Hi & Ei)|(b' & j' & r' & z & Hv' & Hn' & Hin)].
      + apply need_nx. unfold fnx. apply in_map_iff. exists i. split; [rewrite Ei; reflexivity | apply in_seq; lia].
      + destruct (pred_ok b' j' r' (SRef b j r) z Hv' Hin) as [_ Hl].
        pose proof (idx_lt_total bs n _ Hv').
        assert (Hn'' : need n fus fnx (idx n bs (SRef b' j' r')) = true)
          by (rewrite (IH b' j' r'); [exact Hn' | unfold nun in *; lia | exact Hv']).
        rewrite idx_ref in Hn''. apply (need_pred n fus fnx fwfn (boff bs b' + sz b' * j' + r') _ z).
        * rewrite fus_len. exact (unit_lt b' j' r' Hv').
        * exact Hn''.
        * rewrite fus_nth by exact Hv'. unfold flat_unit. cbn [tin]. apply in_map_iff. exists (SRef b j r, z).
          split; [reflexivity | exact Hin].
    - destruct (need n fus fnx (idx n bs (SRef b j r))) eqn:E; [|reflexivity]. exfalso.
      destruct (Hnd_down b j r Hv Hnd) as [Hno Hrd].
      destruct (needf_step fus n _ _ E) as [H0|(k & z & Hk & Hk2 & Hin)].
      + apply existsb_exists in H0 as (s & Hs & Es). apply Nat.eqb_eq in Es. unfold fnx in Hs.
        apply in_map_iff in Hs as (i & Ei & Hi). apply in_seq in Hi. rewrite <- Ei in Es.
        apply (Hno i ltac:(lia)). symmetry. exact (idx_inj _ _ Hv (Hout i ltac:(lia)) Es).
      + rewrite fus_len in Hk. destruct (loc k Hk) as (b' & j' & r' & Hv' & ->).
        rewrite fus_nth in Hin by exact Hv'. unfold flat_unit in Hin. cbn [tin] in Hin.
        apply in_map_iff in Hin as ([p z'] & Ep & Hin). cbn [fst snd] in Ep. injection Ep as Ep _.
        destruct (pred_ok b' j' r' p z' Hv' Hin) as [Hvp Hl].
        rewrite (idx_inj _ _ Hvp Hv Ep) in Hin, Hl.
        pose proof (Hrd b' j' r' z' Hv' Hin) as Hf. pose proof (idx_lt_total bs n _ Hv').
        change (need n fus fnx (n + (boff bs b' + sz b' * j' + r')) = true) in Hk2.
        rewrite <- idx_ref, (IH b' j' r') in Hk2; [rewrite Hf in Hk2; discriminate | unfold nun in *; lia | exact Hv']. }
  intros b j r Hv. apply (G (S (n + nun)) b j r); [lia | exact Hv].
Qed.

End Need2.

Section Lam2.
Variable cd : sref -> nat.
Hypothesis Hdep : forall x, valid bs n x -> dep n fus (idx n bs x) = cd x.
Variable nd : sref -> bool.
Hypothesis Hneed : forall b j r, valid bs n (SRef b j r) -> need n fus fnx (idx n bs (SRef b j r)) = nd (SRef b j r).
Variable D : nat.
Hypothesis HD : dmax n fus = D.
Variable cl : sref -> nat.
Hypothesis Hup : forall b j r p z, valid bs n (SRef b j r) -> nd (SRef b j r) = true ->
  In (p, z) (sin (sunit bs b j r)) -> cd (SRef b j r) <= cl p.
Hypothesis Hcl_out : forall i, i < n -> cl (out i) = S D.
Hypothesis Hlow : forall x, valid bs n x -> (forall i, i < n -> out i <> x) ->
  cl x = 0 \/ exists b j r z, valid bs n (SRef b j r) /\ nd (SRef b j r) = true /\ In (x, z) (sin (sunit bs b j r))
                             /\ cd (SRef b j r) = cl x.

Theorem lam_cl2 : forall x, valid bs n x -> lam n fus fnx (idx n bs x) = cl x.
Proof.
  intros x Hx. unfold lam.
  assert (Hr : maxl (map (fun k => if need n fus fnx (n + k) && reads fus k (idx n bs x) then dep n fus (n + k) else 0)
                         (seq 0 (length fus))) <= cl x).
  { apply maxl_le. intros a Ha. apply in_map_iff in Ha as (k & <- & Hk). apply in_seq in Hk. rewrite fus_len in Hk.
    destruct (loc k ltac:(lia)) as (b & j & r & Hv & ->).
    destruct (need n fus fnx _ && reads fus _ (idx n bs x)) eqn:E; [|lia].
    apply andb_true_iff in E as [E1 E]. apply (reads_su b j r x Hv Hx) in E as (z & Hin).
    rewrite <- idx_ref, Hneed in E1 by exact Hv.
    rewrite <- idx_ref, Hdep by exact Hv. exact (Hup b j r x z Hv E1 Hin). }
  destruct (existsb (Nat.eqb (idx n bs x)) fnx) eqn:Eo.
  - apply existsb_exists in Eo as (s & Hs & E). apply Nat.eqb_eq in E. unfold fnx in Hs.
    apply in_map_iff in Hs as (i & Ei & Hi). apply in_seq in Hi. rewrite <- Ei in E.
    apply (idx_inj x (out i) Hx (Hout i ltac:(lia))) in E. subst x. rewrite Hcl_out by lia. rewrite HD.
    apply Nat.max_l. eapply Nat.le_trans; [|apply Nat.le_succ_diag_r]. apply maxl_le. intros a Ha.
    apply in_map_iff in Ha as (k & <- & Hk). apply in_seq in Hk. rewrite fus_len in Hk.
    destruct (_ && _); [rewrite <- HD; apply dep_le_dmax; rewrite fus_len; lia | lia].
  - cbn [Nat.max]. apply Nat.le_antisymm; [exact Hr|].
    assert (Hno : forall i, i < n -> out i <> x).
    { intros i Hi E. subst x. assert (existsb (Nat.eqb (idx n bs (out i))) fnx = true) as C.
      { apply existsb_exists. exists (idx n bs (out i)). split; [|apply Nat.eqb_refl].
        unfold fnx. apply in_map_iff. exists i. split; [reflexivity | apply in_seq; lia]. }
      rewrite C in Eo. discriminate. }
    destruct (Hlow x Hx Hno) as [Z|(b & j & r & z & Hv & Hn & Hin & E)]; [rewrite Z; lia|].
    rewrite <- E. apply maxl_ge. apply in_map_iff. exists (boff bs b + sz b * j + r).
    split; [|apply in_seq; rewrite fus_len; pose proof (unit_lt b j r Hv); lia].
    rewrite <- idx_ref, Hneed, Hn by exact Hv.
    replace (reads fus (boff bs b + sz b * j + r) (idx n bs x)) with true
      by (symmetry; apply (reads_su b j r x Hv Hx); exists z; exact Hin).
    cbn [andb]. apply Hdep. exact Hv.
Qed.

Theorem live_cl2 : forall l x, valid bs n x ->
  live n fus fnx l (idx n bs x) = isref x && nd x && Nat.eqb (cd x) l || Nat.ltb (cd x) l && Nat.ltb l (cl x).
Proof.
  intros l x Hx. unfold live, isunit. rewrite Hdep, lam_cl2 by exact Hx. destruct x as [i|b j r].
  - cbn [idx isref valid] in *. destruct (Nat.leb_spec n i); [lia|]. reflexivity.
  - rewrite Hneed by exact Hx. rewrite idx_ref. destruct (Nat.leb_spec n (n + (boff bs b + sz b * j + r))); [|lia].
    cbn [isref andb]. destruct (nd (SRef b j r)), (Nat.eqb (cd (SRef b j r)) l); reflexivity.
Qed.

(* the rows of the layers 1 .. d - 1, signal by signal *)
Theorem rows_bsum : forall m, fold_right Nat.add 0 (map (width n fus fnx) (seq 1 m))
  = nsum n (fun i => nsum m (fun l => Nat.b2n (Nat.ltb (cd (SIn i)) (S l) && Nat.ltb (S l) (cl (SIn i)))))
    + bsum (fun x => nsum m (fun l => Nat.b2n (nd x && Nat.eqb (cd x) (S l) || Nat.ltb (cd x) (S l) && Nat.ltb (S l) (cl x)))).
Proof.
  intro m. rewrite <- (seq_shift m 0), map_map.
  rewrite (fold_nsum_seq _ (fun l => width n fus fnx (S l))).
  rewrite (nsum_ext m _ (fun l => nsum n (fun i => Nat.b2n (Nat.ltb (cd (SIn i)) (S l) && Nat.ltb (S l) (cl (SIn i))))
                                  + bsum (fun x => Nat.b2n (nd x && Nat.eqb (cd x) (S l) || Nat.ltb (cd x) (S l) && Nat.ltb (S l) (cl x))))).
  - rewrite nsum_add, nsum_swap, nsum_bsum_swap. reflexivity.
  - intros l _. cbn [width]. unfold NS. rewrite fus_len, seq_app, filter_app, length_app, !length_filter_seq. f_equal.
    + apply nsum_ext. intros i Hi. cbn [Nat.add]. replace i with (idx n bs (SIn i)) at 1 by reflexivity.
      rewrite live_cl2 by (cbn [valid]; exact Hi). reflexivity.
    + rewrite nsum_units. unfold bsum. apply nsum_ext. intros b Hb. apply nsum_ext. intros j Hj. apply nsum_ext.
      intros r Hr. assert (Hv : valid bs n (SRef b j r)) by (split; [exact Hb | split; [exact Hj | exact Hr]]).
      replace (0 + n + (boff bs b + sz b * j + r)) with (idx n bs (SRef b j r)) by (rewrite idx_ref; lia).
      rewrite live_cl2 by exact Hv. cbn [isref andb]. reflexivity.
Qed.

End Lam2.

End BL.