1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/main.bend as Storage
4	import ../storage/model.bend as Model
5	import ../storage/refinement.bend as Refinement
6	import ./core.bend as Core
7	
8	def scalar_shape() -> Shape.Shape:
9	  Shape.SShape{Nil{}, 1n, {==}, Unit{}, Unit{}}
10	
11	def sequential_sum(values: List<&2, F32>, accumulator: F32) -> F32:
12	  match values:
13	    case Nil{}:
14	      accumulator
15	    case value <> tail:
16	      sequential_sum(tail, F32.add(accumulator, value))
17	
18	def read_step(-shape: Shape.Shape,
19	  rec: Core.Tensor<shape> -> F32 -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>,
20	  accumulator: F32, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) ->
21	  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
22	  (tensor, result) = pair
23	  match result:
24	    case Fail{error}:
25	      Fail{error}
26	    case Done{value}:
27	      rec(tensor, F32.add(accumulator, value))
28	
29	def sum_decided(-shape: Shape.Shape,
30	  rec: Core.Tensor<shape> -> F32 -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>, tensor: Core.Tensor<shape>, +cursor: Nat,
31	  +index: U32, accumulator: F32, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) ->
32	  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
33	  match decision:
34	    case Shape.BFalse{equal}:
35	      Fail{Core.StorageError{Storage.InvalidIndex{}}}
36	    case Shape.BTrue{equal}:
37	      read_step(shape, rec, accumulator, Core.read_flat(shape, tensor, index))
38	
39	def natural_sum_loop(fuel: Nat, -shape: Shape.Shape, tensor: Core.Tensor<shape>, +cursor: Nat,
40	  accumulator: F32) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
41	  match fuel:
42	    case 0n:
43	      Done{(tensor, accumulator)}
44	    case 1n+tail:
45	      +index = U32.from_nat(cursor)
46	      sum_decided(shape, tensor => value => natural_sum_loop(tail, shape, tensor, 1n+cursor, value), tensor, cursor, index, accumulator,
47	        Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
48	
49	def sum_loop(fuel: Nat, -shape: Shape.Shape, tensor: Core.Tensor<shape>, +index: U32,
50	  accumulator: F32) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
51	  natural_sum_loop(fuel, shape, tensor, U32.to_nat(index), accumulator)
52	
53	def scalar_result(-input_shape: Shape.Shape, tensor: Core.Tensor<input_shape>,
54	  result: Result<&2, &1, Core.Error, Core.Tensor<scalar_shape()>>) ->
55	  Result<&2, &1, Core.Error, Core.Tensor<input_shape> & Core.Tensor<scalar_shape()>>:
56	  match result:
57	    case Fail{error}:
58	      Fail{error}
59	    case Done{scalar}:
60	      Done{(tensor, scalar)}
61	
62	def sum_result(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>) ->
63	  Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>:
64	  match result:
65	    case Fail{error}:
66	      Fail{error}
67	    case Done{pair}:
68	      (tensor, value) = pair
69	      scalar_result(shape, tensor, Core.new(scalar_shape(), value))
70	
71	# The fuel + 1 words of the array from slot `cursor`, whose read `pair` is
72	# already made, added to `accumulator` left to right. Each step reads the
73	# next word and hands it on as an argument: one tail call with no checked
74	# read and no closure per word.
75	def array_run(fuel: Nat, +cursor: Nat, accumulator: F32, pair: Array<F32> & F32) -> F32 & Array<F32>:
76	  match fuel pair:
77	    case 0n Tuple{array, word}:
78	      (F32.add(accumulator, word), array)
79	    case 1n+rest Tuple{array, word}:
80	      array_run(rest, 1n+cursor, F32.add(accumulator, word), Array.get(F32, array, U32.from_nat(1n+cursor)))
81	
82	# The first `fuel` words of the array from slot `cursor`, added to
83	# `accumulator` left to right: the F32 adds of `natural_sum_loop`, in the
84	# same order.
85	def array_sum(fuel: Nat, +cursor: Nat, accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
86	  match fuel:
87	    case 0n:
88	      (accumulator, array)
89	    case 1n+rest:
90	      array_run(rest, cursor, accumulator, Array.get(F32, array, U32.from_nat(cursor)))
91	
92	def array_from(pair: Array<F32> & F32) -> Array<F32>:
93	  Refinement.first(F32, pair)
94	
95	type Fact<-P: Type> is Data:
96	  Fact{-proof: P}
97	
98	def kept_balanced(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), -kept: Array<F32>,
99	  -same: {kept == original : Array<F32>}) -> Fact<Model.Balanced(F32, depth, kept)>:
100	  Fact{
101	    %Equal.sym(Array<F32>, kept, original, same) : Model.Balanced(F32, depth, _)
102	    balanced}
103	
104	# The array a read hands back is the array it read.
105	def read_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), -array: Array<F32>,
106	  -index: U32, -pair: Array<F32> & F32, -exact: {pair == Array.get(F32, array, index) : Array<F32> & F32},
107	  -same: {array == original : Array<F32>}) -> Fact<{array_from(pair) == original : Array<F32>}>:
108	  Fact{Equal.trans(Array<F32>, array_from(pair), array_from(Array.get(F32, array, index)), original,
109	    Equal.cong(Array<F32> & F32, Array<F32>, p => array_from(p), pair, Array.get(F32, array, index), exact),
110	    Equal.trans(Array<F32>, array_from(Array.get(F32, array, index)), array, original,
111	      Refinement.law_get_retains(F32, depth, array, index,
112	        {\{Fact: p => p} : Fact<(Model.Balanced(F32, depth, array))> -> Model.Balanced(F32, depth, array)}(
113	          kept_balanced(depth, original, balanced, array, same))),
114	      same))}
115	
116	# A run hands back the array it read.
117	def array_run_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), fuel: Nat, +cursor: Nat,
118	  -accumulator: F32, -array: Array<F32>, pair: Array<F32> & F32,
119	  -exact: {pair == Array.get(F32, array, U32.from_nat(cursor)) : Array<F32> & F32}, -same: {array == original : Array<F32>}) ->
120	  Fact<{Pair.snd(F32, Array<F32>, array_run(fuel, cursor, accumulator, pair)) == original : Array<F32>}>:
121	  match fuel pair:
122	    case 0n Tuple{kept, word}:
123	      read_keeps(depth, original, balanced, array, U32.from_nat(cursor), (kept, word), exact, same)
124	    case 1n+rest Tuple{kept, word}:
125	      array_run_keeps(depth, original, balanced, rest, 1n+cursor, F32.add(accumulator, word), kept,
126	        Array.get(F32, kept, U32.from_nat(1n+cursor)), {==},
127	        {\{Fact: p => p} : Fact<({array_from((kept, word)) == original : Array<F32>})> -> {kept == original : Array<F32>}}(
128	          read_keeps(depth, original, balanced, array, U32.from_nat(cursor), (kept, word), exact, same)))
129	
130	def array_sum_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), fuel: Nat, +cursor: Nat,
131	  -accumulator: F32, array: Array<F32>, -same: {array == original : Array<F32>}) ->
132	  Fact<{Pair.snd(F32, Array<F32>, array_sum(fuel, cursor, accumulator, array)) == original : Array<F32>}>:
133	  match fuel:
134	    case 0n:
135	      Fact{same}
136	    case 1n+rest:
137	      array_run_keeps(depth, original, balanced, rest, cursor, accumulator, array, Array.get(F32, array, U32.from_nat(cursor)), {==}, same)
138	
139	def owned_sum(-shape: Shape.Shape, +depth: Nat, +length: Nat, -array: Array<F32>,
140	  -balanced: Model.Balanced(F32, depth, array), -valid: {Storage.bounds(depth, length) == True{} : Bool},
141	  -exact: {length == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
142	  pair: F32 & Array<F32>, -kept: Fact<{Pair.snd(F32, Array<F32>, pair) == array : Array<F32>}>) ->
143	  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
144	  (value, retained) = pair
145	  Done{(Core.TTensor{Storage.Owned{depth, length, retained,
146	    {\{Fact: p => p} : Fact<(Model.Balanced(F32, depth, retained))> -> Model.Balanced(F32, depth, retained)}(
147	      kept_balanced(depth, array, balanced, retained,
148	        {\{Fact: p => p} : Fact<({Pair.snd(F32, Array<F32>, (value, retained)) == array : Array<F32>})> -> {retained == array : Array<F32>}}(kept))),
149	    valid}, exact}, value)}
150	
151	# The first `count` words of the tensor added left to right from +0.0, next
152	# to the tensor, which keeps every word.
153	def tensor_sum(-shape: Shape.Shape, +count: Nat, tensor: Core.Tensor<shape>) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
154	  Core.TTensor{storage, exact} = tensor
155	  match storage:
156	    case Storage.Empty{}:
157	      Done{(Core.TTensor{Storage.Empty{}, exact}, 0.0)}
158	    case Storage.Owned{+depth, +length, array, balanced, valid}:
159	      owned_sum(shape, depth, length, array, balanced, valid, exact, array_sum(count, 0n, 0.0, array),
160	        array_sum_keeps(depth, array, balanced, count, 0n, 0.0, array, {==}))
161	
162	def sum(+shape: Shape.Shape, tensor: Core.Tensor<shape>) ->
163	  Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>:
164	  sum_result(shape, tensor_sum(shape, Shape.elements(shape), tensor))
165	
166	def pack_result(+shape: Shape.Shape,
167	  result: Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>) ->
168	  Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>:
169	  match result:
170	    case Fail{error}:
171	      Fail{error}
172	    case Done{pair}:
173	      (tensor, scalar) = pair
174	      Done{(Core.pack(shape, tensor), Core.pack(scalar_shape(), scalar))}
175	
176	def sum_dynamic(dynamic: Core.Dynamic) -> Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>:
177	  Core.Packed{+shape, tensor} = dynamic
178	  pack_result(shape, sum(shape, tensor))
179	
180	def sum_observation(-shape: Shape.Shape,
181	  result: Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>) -> Maybe<&2, F32>:
182	  match result:
183	    case Fail{error}:
184	      None{}
185	    case Done{pair}:
186	      (tensor, scalar) = pair
187	      Core.observe(scalar_shape(), scalar, 0)