1	import Base
2	import ../shape/nat_proofs.bend as Nats
3	import ../shape/index_proofs.bend as Index
4	import ../storage/model.bend as Model
5	import ../storage/refinement.bend as Refinement
6	import ../storage/clone.bend as Clone
7	import ./parallel.bend as Build
8	
9	# Proofs that `Build.build` holds, at every output cell, what the
10	# one-thread tree computes there from the input state `start`, whatever the
11	# split. They need two facts of the operation: a cell hands `start` back
12	# unchanged (`Keeps`), and `split` copies `start` into both halves. Each
13	# operation then proves only what its cell function computes from `start`.
14	
15	type Fact<-P: Type> is Data:
16	  Fact{-proof: P}
17	
18	# What cell `position` holds: the F32 its cell function computes from
19	# `start`.
20	def CellIs(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
21	  +p: P, -start: (A & B), +count: Nat, +position: Nat, -observed: Maybe<&2, F32>) -> Type:
22	  {observed == Some{Pair.fst(F32, (A & B), cell(p, position, start, Nat.is_lt(position, count)))} : Maybe<&2, F32>}
23	
24	# Every slot of `out`, the cells `base .. base + 2^depth`, holds what CellIs
25	# names for its index.
26	def Cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
27	  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -out: Array<F32>) -> Type:
28	  @+path: Model.Path(depth) ->
29	    CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(depth, path), base), Model.observe(F32, depth, out, path))
30	
31	# A cell hands its input state back unchanged.
32	def Keeps(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
33	  +p: P, -start: (A & B)) -> Type:
34	  @+position: Nat -> @inside: Bool -> {Pair.snd(F32, (A & B), cell(p, position, start, inside)) == start : (A & B)}
35	
36	# A run of the one-thread tree: it hands `start` back and fills its cells.
37	def TreeFacts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
38	  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B)) -> Type:
39	  {Pair.snd(Array<F32>, (A & B), result) == start : (A & B)} & Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result))
40	
41	def facts_kept(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
42	  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B),
43	  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)>) ->
44	  Fact<{Pair.snd(Array<F32>, (A & B), result) == start : (A & B)}>:
45	  Fact{Pair.fst({Pair.snd(Array<F32>, (A & B), result) == start : (A & B)},
46	    Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result)),
47	    {\{Fact: q => q} : Fact<(TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result))> ->
48	      TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)}(facts))}
49	
50	def facts_cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
51	  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, -result: Array<F32> & (A & B),
52	  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)>) ->
53	  Fact<Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result))>:
54	  Fact{Pair.snd({Pair.snd(Array<F32>, (A & B), result) == start : (A & B)},
55	    Cells(~P, ~A, ~B, ~cell, p, start, depth, count, base, Pair.fst(Array<F32>, (A & B), result)),
56	    {\{Fact: q => q} : Fact<(TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result))> ->
57	      TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, result)}(facts))}
58	
59	def leaf_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
60	  +p: P, -start: (A & B), +count: Nat, +base: Nat, pair: F32 & (A & B),
61	  -exact: {pair == cell(p, base, start, Nat.is_lt(base, count)) : F32 & (A & B)},
62	  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
63	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), pair))>:
64	  (value, rest) = pair
65	  Fact{(Equal.trans((A & B), rest, Pair.snd(F32, (A & B), cell(p, base, start, Nat.is_lt(base, count))), start,
66	      Equal.cong(F32 & (A & B), (A & B), q => Pair.snd(F32, (A & B), q), (value, rest), cell(p, base, start, Nat.is_lt(base, count)), exact),
67	      keeps(base)(Nat.is_lt(base, count))),
68	    path => Equal.cong(F32 & (A & B), Maybe<&2, F32>, q => Some{Pair.fst(F32, (A & B), q)}, (value, rest),
69	      cell(p, base, start, Nat.is_lt(base, count)), exact))}
70	
71	def position_step(+capacity: Nat, +offset: Nat, +base: Nat) ->
72	  {Nat.add(offset, Nat.add(capacity, base)) == Nat.add(Nat.add(capacity, offset), base) : Nat}:
73	  Equal.trans(Nat, Nat.add(offset, Nat.add(capacity, base)), Nat.add(capacity, Nat.add(offset, base)), Nat.add(Nat.add(capacity, offset), base),
74	    Nats.add_swap(offset, capacity, base), Index.add_assoc(capacity, offset, base))
75	
76	def node_cell(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
77	  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, -b: Array<F32>, +path: Model.Path(1n+inner),
78	  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
79	  -high: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)>) ->
80	  Fact<CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path))>:
81	  Model.Step{+bit, +rest} = path
82	  match bit:
83	    case False{}:
84	      Fact{{\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a))> ->
85	        Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)}(low)(rest)}
86	    case True{}:
87	      Fact{
88	        %position_step(Model.capacity(inner), Model.offset(inner, rest), base) :
89	          CellIs(~P, ~A, ~B, ~cell, p, start, count, _, Model.observe(F32, inner, b, rest))
90	        {\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b))> ->
91	          Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)}(high)(rest)}
92	
93	def node_cells(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
94	  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, -b: Array<F32>,
95	  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
96	  -high: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), b)>) ->
97	  Fact<Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b})>:
98	  Fact{path => {\{Fact: q => q} :
99	    Fact<(CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path)))> ->
100	    CellIs(~P, ~A, ~B, ~cell, p, start, count, Nat.add(Model.offset(1n+inner, path), base), Model.observe(F32, 1n+inner, ANode{a, b}, path))}(
101	      node_cell(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, b, path, low, high))}
102	
103	def join_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
104	  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -a: Array<F32>, high: Array<F32> & (A & B),
105	  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a)>,
106	  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), high)>) ->
107	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, Build.join(~(A & B), a, high))>:
108	  (b, rest) = high
109	  Fact{(
110	    {\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), (b, rest)) == start : (A & B)})> -> {Pair.snd(Array<F32>, (A & B), (b, rest)) == start : (A & B)}}(
111	      facts_kept(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (b, rest), facts)),
112	    {\{Fact: q => q} : Fact<(Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b}))> ->
113	      Cells(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, ANode{a, b})}(
114	      node_cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, b, low,
115	        facts_cells(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (b, rest), facts))))}
116	
117	def node_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
118	  +p: P, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, rec: A -> B -> Array<F32> & (A & B), low: Array<F32> & (A & B),
119	  -facts: Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, base, low)>,
120	  -next: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} ->
121	    Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), rec(l, r))>) ->
122	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n+inner, count, base, Build.node(~A, ~B, rec, low))>:
123	  (a, s) = low
124	  (l, r) = s
125	  join_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base, a, rec(l, r),
126	    facts_cells(~P, ~A, ~B, ~cell, p, start, inner, count, base, (a, (l, r)), facts),
127	    next(l)(r)({\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), (a, (l, r))) == start : (A & B)})> -> {Pair.snd(Array<F32>, (A & B), (a, (l, r))) == start : (A & B)}}(
128	      facts_kept(~P, ~A, ~B, ~cell, p, start, inner, count, base, (a, (l, r)), facts))))
129	
130	# The depth-1 tree: the leaf of cell `base` next to the leaf of the cell
131	# after it, which reads the state the first hands back.
132	def pair_join_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
133	  +p: P, -start: (A & B), +count: Nat, +base: Nat, x: F32, high: F32 & (A & B),
134	  -low: Fact<Cells(~P, ~A, ~B, ~cell, p, start, 0n, count, base, ALeaf{x})>,
135	  -exact: {high == cell(p, Nat.add(1n, base), start, Nat.is_lt(Nat.add(1n, base), count)) : F32 & (A & B)},
136	  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
137	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n, count, base, Build.pair_join(~(A & B), x, high))>:
138	  (y, rest) = high
139	  join_facts(~P, ~A, ~B, ~cell, p, start, 0n, count, base, ALeaf{x}, Build.leaf(~(A & B), (y, rest)), low,
140	    leaf_facts(~P, ~A, ~B, ~cell, p, start, count, Build.upper(0n, base), (y, rest), exact, keeps))
141	
142	def pair_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
143	  +p: P, -start: (A & B), +count: Nat, +base: Nat, low: F32 & (A & B),
144	  -exact: {low == cell(p, base, start, Nat.is_lt(base, count)) : F32 & (A & B)},
145	  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
146	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, 1n, count, base, Build.pair(~P, ~A, ~B, ~cell, p, count, base, low))>:
147	  (x, s) = low
148	  (l, r) = s
149	  pair_join_facts(~P, ~A, ~B, ~cell, p, start, count, base, x, cell(p, Nat.add(1n, base), (l, r), Nat.is_lt(Nat.add(1n, base), count)),
150	    facts_cells(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), (x, (l, r))),
151	      leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, (x, (l, r)), exact, keeps)),
152	    Equal.cong((A & B), F32 & (A & B), q => cell(p, Nat.add(1n, base), q, Nat.is_lt(Nat.add(1n, base), count)), (l, r), start,
153	      {\{Fact: q => q} : Fact<({Pair.snd(Array<F32>, (A & B), Build.leaf(~(A & B), (x, (l, r)))) == start : (A & B)})> ->
154	        {Pair.snd(Array<F32>, (A & B), Build.leaf(~(A & B), (x, (l, r)))) == start : (A & B)}}(
155	        facts_kept(~P, ~A, ~B, ~cell, p, start, 0n, count, base, Build.leaf(~(A & B), (x, (l, r))),
156	          leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, (x, (l, r)), exact, keeps)))),
157	    keeps)
158	
159	# The one-thread tree hands `start` back and fills every cell as CellIs
160	# says.
161	def tree_facts(~P: Data, ~A: Type, ~B: Type, ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
162	  +p: P, -start: (A & B), +depth: Nat, +count: Nat, +base: Nat, state: (A & B), -at: {state == start : (A & B)},
163	  -keeps: Keeps(~P, ~A, ~B, ~cell, p, start)) ->
164	  Fact<TreeFacts(~P, ~A, ~B, ~cell, p, start, depth, count, base, Build.tree(~P, ~A, ~B, ~cell, p, depth, count, base, state))>:
165	  match depth:
166	    case 0n:
167	      leaf_facts(~P, ~A, ~B, ~cell, p, start, count, base, cell(p, base, state, Nat.is_lt(base, count)),
168	        Equal.cong((A & B), F32 & (A & B), s => cell(p, base, s, Nat.is_lt(base, count)), state, start, at), keeps)
169	    case 1n+inner:
170	      match inner:
171	        case 0n:
172	          pair_facts(~P, ~A, ~B, ~cell, p, start, count, base, cell(p, base, state, Nat.is_lt(base, count)),
173	            Equal.cong((A & B), F32 & (A & B), s => cell(p, base, s, Nat.is_lt(base, count)), state, start, at), keeps)
174	        case 1n+below:
175	          node_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base,
176	            l => r => Build.tree(~P, ~A, ~B, ~cell, p, inner, count, Build.upper(inner, base), (l, r)),
177	            Build.tree(~P, ~A, ~B, ~cell, p, inner, count, base, state),
178	            tree_facts(~P, ~A, ~B, ~cell, p, start, inner, count, base, state, at, keeps),
179	            l => r => at => tree_facts(~P, ~A, ~B, ~cell, p, start, inner, count, Build.upper(inner, base), (l, r), at, keeps))
180	
181	def half_equal(-S: Type, -start: S, -state: S, -half: S, -split: S -> S & S, -pick: (S & S) -> S,
182	  picked: {half == pick(split(state)) : S}, at: {state == start : S}, copied: {pick(split(start)) == start : S}) ->
183	  {half == start : S}:
184	  Equal.trans(S, half, pick(split(start)), start,
185	    Equal.trans(S, half, pick(split(state)), pick(split(start)), picked, Equal.cong(S, S, s => pick(split(s)), state, start, at)),
186	    copied)
187	
188	def fork_cells(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: (A & B) -> (A & B) & (A & B),
189	  ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
190	  +q: Q, -start: (A & B), +inner: Nat, +count: Nat, +base: Nat, -state: (A & B), pair: (A & B) & (A & B),
191	  -exact: {pair == split(state) : (A & B) & (A & B)},
192	  -low: A -> B -> Array<F32>, -high: A -> B -> Array<F32>, -at: {state == start : (A & B)},
193	  -first: {Pair.fst((A & B), (A & B), split(start)) == start : (A & B)}, -second: {Pair.snd((A & B), (A & B), split(start)) == start : (A & B)},
194	  -low_cells: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} -> Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, base, low(l, r))>,
195	  -high_cells: @l: A -> @r: B -> @-at: {(l, r) == start : (A & B)} ->
196	    Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, Build.upper(inner, base), high(l, r))>) ->
197	  Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, Build.fork(~A, ~B, pair, low, high))>:
198	  (one, two) = pair
199	  (l1, r1) = one
200	  (l2, r2) = two
201	  node_cells(~P, ~A, ~B, ~cell, prepare(q), start, inner, count, base, low(l1, r1), high(l2, r2),
202	    low_cells(l1)(r1)(half_equal((A & B), start, state, (l1, r1), split, h => Pair.fst((A & B), (A & B), h),
203	      Equal.cong((A & B) & (A & B), (A & B), h => Pair.fst((A & B), (A & B), h), ((l1, r1), (l2, r2)), split(state), exact), at, first)),
204	    high_cells(l2)(r2)(half_equal((A & B), start, state, (l2, r2), split, h => Pair.snd((A & B), (A & B), h),
205	      Equal.cong((A & B) & (A & B), (A & B), h => Pair.snd((A & B), (A & B), h), ((l1, r1), (l2, r2)), split(state), exact), at, second)))
206	
207	# Whatever the split, every cell of the parallel build holds what the
208	# one-thread tree computes there from `start`.
209	def build_cells(~Q: Data, ~P: Data, ~A: Type, ~B: Type, ~prepare: @+q: Q -> P, ~split: (A & B) -> (A & B) & (A & B), ~cell: @+p: P -> @+position: Nat -> (A & B) -> Bool -> F32 & (A & B),
210	  +q: Q, -start: (A & B), +depth: Nat, +splits: Nat, +count: Nat, +base: Nat, state: (A & B), -at: {state == start : (A & B)},
211	  -keeps: Keeps(~P, ~A, ~B, ~cell, prepare(q), start),
212	  -first: {Pair.fst((A & B), (A & B), split(start)) == start : (A & B)}, -second: {Pair.snd((A & B), (A & B), split(start)) == start : (A & B)}) ->
213	  Fact<Cells(~P, ~A, ~B, ~cell, prepare(q), start, depth, count, base, Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, depth, splits, count, base, state))>:
214	  match depth splits:
215	    case 1n+inner 1n+rest:
216	      fork_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, count, base, state, split(state), {==},
217	        l => r => Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, base, (l, r)),
218	        l => r => Build.build(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, inner, rest, count, Build.upper(inner, base), (l, r)),
219	        at, first, second,
220	        l => r => at => build_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, rest, count, base, (l, r), at, keeps, first, second),
221	        l => r => at => build_cells(~Q, ~P, ~A, ~B, ~prepare, ~split, ~cell, q, start, inner, rest, count, Build.upper(inner, base), (l, r), at, keeps, first, second))
222	    case 0n _:
223	      facts_cells(~P, ~A, ~B, ~cell, prepare(q), start, 0n, count, base, Build.tree(~P, ~A, ~B, ~cell, prepare(q), 0n, count, base, state),
224	        tree_facts(~P, ~A, ~B, ~cell, prepare(q), start, 0n, count, base, state, at, keeps))
225	    case 1n+inner 0n:
226	      facts_cells(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, Build.tree(~P, ~A, ~B, ~cell, prepare(q), 1n+inner, count, base, state),
227	        tree_facts(~P, ~A, ~B, ~cell, prepare(q), start, 1n+inner, count, base, state, at, keeps))