1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise_totality.bend as Totality
5	import ../tensor/projection_LAWS.bend as Projection
6	import ./reduce.bend as Reduce
7	
8	def owner(pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Core.Dynamic:
9	  (input, result) = pair
10	  input
11	
12	def succeeded(pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Bool:
13	  (input, result) = pair
14	  Result.is_done(&2, &1, Reduce.Error, Core.Dynamic, result)
15	
16	def word_at(values: List<&2, F32>, +position: Nat) -> Maybe<&2, F32>:
17	  match values position:
18	    case Nil{} _:
19	      None{}
20	    case value <> rest 0n:
21	      Some{value}
22	    case value <> rest 1n+next:
23	      word_at(rest, next)
24	
25	def add_word(total: Maybe<&2, F32>, value: Maybe<&2, F32>) -> Maybe<&2, F32>:
26	  match total value:
27	    case None{} _:
28	      None{}
29	    case Some{sum} None{}:
30	      None{}
31	    case Some{sum} Some{word}:
32	      Some{F32.add(sum, word)}
33	
34	def matched(total: Maybe<&2, F32>, value: Maybe<&2, F32>, same: Bool) -> Maybe<&2, F32>:
35	  match same:
36	    case False{}:
37	      total
38	    case True{}:
39	      add_word(total, value)
40	
41	def contribute(+values: List<&2, F32>, +target_index: U32, total: Maybe<&2, F32>, offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
42	  match offsets:
43	    case None{}:
44	      None{}
45	    case Some{pair}:
46	      (target_offset, source_offset) = pair
47	      matched(total, word_at(values, U32.to_nat(source_offset)), Nat.is_eq(U32.to_nat(target_offset), U32.to_nat(target_index)))
48	
49	# The specified accumulation: start at +0.0 and visit contribution positions
50	# 0, 1, ..., count - 1 in ascending order; each position whose independent
51	# broadcast projection, Projection.offset_spec, has target_offset equal to
52	# target_index adds the word at its source_offset with one F32 add.
53	def fiber(+plan: Shape.BroadcastPlan, +values: List<&2, F32>, +target_index: U32, +count: Nat) -> Maybe<&2, F32>:
54	  match count:
55	    case 0n:
56	      Some{0.0}
57	    case 1n+position:
58	      contribute(values, target_index, fiber(plan, values, target_index, position), Projection.offset_spec(plan, position))
59	
60	def planned_fiber(+values: List<&2, F32>, +target_index: U32, +count: Nat,
61	  plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Maybe<&2, F32>:
62	  match plan:
63	    case Fail{error}:
64	      None{}
65	    case Done{+plan}:
66	      fiber(plan, values, target_index, count)
67	
68	# values lists the contribution's words: the checked observation at every
69	# U32 position is the word at that position.
70	def Listed(-input: Core.Dynamic, +values: List<&2, F32>) -> Type:
71	  @+source: U32 -> {Core.checked_observe_dynamic(input, source) == word_at(values, U32.to_nat(source)) : Maybe<&2, F32>}
72	
73	def MappingWith(-input: Core.Dynamic, +target: Shape.Shape, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>,
74	  pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
75	  (retained, result) = pair
76	  match result:
77	    case Fail{error}:
78	      Unit
79	    case Done{output}:
80	      {Totality.dimensions(output) == Shape.dimensions(target) : List<&2, U32>} &
81	        (@+values: List<&2, F32> -> @-listed: Listed(input, values) ->
82	          @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
83	            {Core.checked_observe_dynamic(output, index) ==
84	              planned_fiber(values, index, Shape.logical_numel(Totality.dimensions(input)), plan) : Maybe<&2, F32>})
85	
86	def Mapping(input: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
87	  MappingWith(input, target, Shape.broadcast_checked(target, Totality.dynamic_shape(input)), pair)
88	
89	def ZeroValues(+target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
90	  (retained, result) = pair
91	  match result:
92	    case Fail{error}:
93	      Unit
94	    case Done{output}:
95	      @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
96	        {Core.checked_observe_dynamic(output, index) == Some{0.0} : Maybe<&2, F32>}
97	
98	type Retained<-input: Core.Dynamic, -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
99	  Retained{-kept: Totality.SameSource(input, owner(pair))}
100	
101	type MappingModel<-input: Core.Dynamic, -target: Shape.Shape,
102	  -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
103	  MappingModel{-mapping: Mapping(input, target, pair)}
104	
105	type Success<-pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
106	  Success{-successful: {succeeded(pair) == True{} : Bool}}
107	
108	type Length<-input: Core.Dynamic, -values: List<&2, F32>> is Data:
109	  Length{-exact: {List.length(&2, F32, values) == Shape.logical_numel(Totality.dimensions(input)) : Nat}}
110	
111	type Zeros<-target: Shape.Shape, -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
112	  Zeros{-zeros: ZeroValues(target, pair)}
113	
114	# The contribution owner comes back with the same dimensions and observations.
115	law reduce_retains_input:
116	  for input: Core.Dynamic
117	  for +target: Shape.Shape
118	  Retained<input, Reduce.reduce_to(input, target)>
119	
120	# On success the output has the target's dimensions, and for every list of
121	# words that lists the contribution, each in-bounds target position holds the
122	# ordered fiber sum of those words.
123	law reduce_mapping:
124	  for input: Core.Dynamic
125	  for +target: Shape.Shape
126	  MappingModel<input, target, Reduce.reduce_to(input, target)>
127	
128	# A list of words that lists the contribution has exactly the contribution's
129	# logical length.
130	law listed_length:
131	  for input: Core.Dynamic
132	  for +values: List<&2, F32>
133	  for listed: Listed(input, values)
134	  Length<input, values>
135	
136	# A target whose right-aligned broadcast with the contribution shape is the
137	# contribution shape itself always succeeds.
138	law reduce_valid_success:
139	  for input: Core.Dynamic
140	  for +target: Shape.Shape
141	  for +alignment: Shape.Broadcast<Shape.extents(target), Shape.extents(Totality.dynamic_shape(input)), Shape.extents(Totality.dynamic_shape(input))>
142	  Success<Reduce.reduce_to(input, target)>
143	
144	# Reducing an empty contribution gives zeros at every target position, not an
145	# empty tensor: for example [0, 3] reduced to [1, 3] is three zeros.
146	law reduce_empty_zeros:
147	  for input: Core.Dynamic
148	  for +target: Shape.Shape
149	  for empty: {Shape.logical_numel(Totality.dimensions(input)) == 0n : Nat}
150	  Zeros<target, Reduce.reduce_to(input, target)>
151	
152	# When the broadcast planner rejects the target and contribution shapes,
153	# reduce_to fails with exactly that planner error and returns the contribution
154	# owner. For example [2] to [0] fails with IncompatibleBroadcast instead of
155	# returning an empty tensor.
156	law reduce_shape_error:
157	  for input: Core.Dynamic
158	  for +target: Shape.Shape
159	  for error: Shape.Error
160	  for rejected: {Shape.broadcast_checked(target, Totality.dynamic_shape(input)) == Fail{error} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
161	  {Reduce.reduce_to(input, target) == (input, Fail{Reduce.TensorError{Core.ShapeError{error}}}) : Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>}
162	
163	# When the planned broadcast of the target is not the contribution's shape
164	# (for example a [2, 3] target for a [3] contribution), reduce_to fails with
165	# NotReducible and returns the contribution owner.
166	law reduce_not_reducible:
167	  for input: Core.Dynamic
168	  for +target: Shape.Shape
169	  for +plan: Shape.BroadcastPlan
170	  for planned: {Shape.broadcast_checked(target, Totality.dynamic_shape(input)) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
171	  for differ: {Shape.dimensions(Shape.output_shape(plan)) != Totality.dimensions(input) : List<&2, U32>}
172	  {Reduce.reduce_to(input, target) == (input, Fail{Reduce.NotReducible{}}) : Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>}