1	import Base
2	import ../shape/core.bend as Shape
3	import ../shape/LAWS.bend as ShapeLaws
4	import ./core.bend as Core
5	import ./elementwise.bend as Elementwise
6	
7	def pair_offsets(left: Maybe<&2, U32>, right: Maybe<&2, U32>) -> Maybe<&1, U32 & U32>:
8	  match left right:
9	    case None{} _:
10	      None{}
11	    case Some{left} None{}:
12	      None{}
13	    case Some{left} Some{right}:
14	      Some{(left, right)}
15	
16	def offset_spec(+plan: Shape.BroadcastPlan, +position: Nat) -> Maybe<&1, U32 & U32>:
17	  +left = Shape.extents(Shape.left_shape(plan))
18	  +right = Shape.extents(Shape.right_shape(plan))
19	  +output = Shape.extents(Shape.output_shape(plan))
20	  +index = ShapeLaws.fin_of(Shape.numel(output), position)
21	  pair_offsets(
22	    ShapeLaws.flatten_spec(left, ShapeLaws.project_left_coordinate(left, right, output, Shape.alignment(plan), index)),
23	    ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), index)))
24	
25	def indices_result(result: Result<&2, &1, Core.Error, U32 & U32>) -> Maybe<&1, U32 & U32>:
26	  match result:
27	    case Fail{error}:
28	      None{}
29	    case Done{indices}:
30	      Some{indices}
31	
32	def indices_accepted(result: Result<&2, &1, Core.Error, U32 & U32>) -> Bool:
33	  match result:
34	    case Fail{error}:
35	      False{}
36	    case Done{indices}:
37	      True{}
38	
39	def specified_expected(operation: Elementwise.Operation, left: Core.Dynamic, right: Core.Dynamic,
40	  offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
41	  match offsets:
42	    case None{}:
43	      None{}
44	    case Some{indices}:
45	      (left_index, right_index) = indices
46	      Elementwise.expected_left(operation, right, right_index, Core.checked_observe_dynamic(left, left_index))
47	
48	law actual_flat_projection:
49	  for +plan: Shape.BroadcastPlan
50	  for +index: Shape.FlatIndex<Shape.output_shape(plan)>
51	  {indices_result(Elementwise.flat_input_indices(plan, Done{index})) ==
52	    offset_spec(plan, ShapeLaws.position(Shape.output_shape(plan), index)) : Maybe<&1, U32 & U32>}
53	
54	law input_indices_refines:
55	  for +plan: Shape.BroadcastPlan
56	  for +index: U32
57	  {indices_result(Elementwise.input_indices(plan, index)) ==
58	    offset_spec(plan, U32.to_nat(index)) : Maybe<&1, U32 & U32>}
59	
60	law input_indices_acceptance:
61	  for +plan: Shape.BroadcastPlan
62	  for +index: U32
63	  {indices_accepted(Elementwise.input_indices(plan, index)) ==
64	    Nat.is_lt(U32.to_nat(index), Shape.numel(Shape.extents(Shape.output_shape(plan)))) : Bool}
65	
66	law expected_projection_refines:
67	  for +operation: Elementwise.Operation
68	  for +plan: Shape.BroadcastPlan
69	  for left: Core.Dynamic
70	  for right: Core.Dynamic
71	  for +index: U32
72	  {Elementwise.expected(operation, plan, left, right, index) ==
73	    specified_expected(operation, left, right, offset_spec(plan, U32.to_nat(index))) : Maybe<&2, F32>}