1	import Base
2	import ../../shape/core.bend as Shape
3	import ../core.bend as Core
4	import ../matmul.bend as Matrix
5	import ./fold.bend as Fold
6	
7	# What a fold result promises, stated only through the result: it succeeds
8	# with `total`, and the returned input still holds every word. `total` and
9	# `word` are the specification; an instance fills them in.
10	type FoldContract<-input: Shape.Shape, -total: F32, -word: U32 -> F32,
11	  -result: Core.Dynamic & Result<&2, &1, Core.Error, F32>> is Type:
12	  FoldContract{
13	    value: {Fold.value(result) == Some{total} : Maybe<&2, F32>},
14	    kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
15	      {Fold.input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}}
16	
17	# The same for a pair fold: the total, and every word of both returned inputs.
18	type PairFoldContract<-left: Shape.Shape, -right: Shape.Shape, -total: F32, -left_word: U32 -> F32, -right_word: U32 -> F32,
19	  -result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>> is Type:
20	  PairFoldContract{
21	    value: {Fold.pair_value(result) == Some{total} : Maybe<&2, F32>},
22	    left_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
23	      {Fold.left_point(result, other) == Some{left_word(other)} : Maybe<&2, F32>},
24	    right_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
25	      {Fold.right_point(result, other) == Some{right_word(other)} : Maybe<&2, F32>}}
26	
27	def product_succeeded(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Bool:
28	  Result.is_done(&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic, result)
29	
30	def product_extents(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Maybe<&2, List<&2, Nat>>:
31	  match result:
32	    case Fail{error}:
33	      None{}
34	    case Done{values}:
35	      (left, right, product) = values
36	      Core.Packed{shape, tensor} = product
37	      Some{Shape.extents(shape)}
38	
39	def product_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
40	  match result:
41	    case Fail{error}:
42	      None{}
43	    case Done{values}:
44	      (left, right, product) = values
45	      Core.dynamic_read_value(Core.read_dynamic(product, index))
46	
47	def product_left_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
48	  match result:
49	    case Fail{error}:
50	      None{}
51	    case Done{values}:
52	      (left, right, product) = values
53	      Core.dynamic_read_value(Core.read_dynamic(left, index))
54	
55	def product_right_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
56	  match result:
57	    case Fail{error}:
58	      None{}
59	    case Done{values}:
60	      (left, right, product) = values
61	      Core.dynamic_read_value(Core.read_dynamic(right, index))
62	
63	# What a matrix product promises, stated only through the result: it
64	# succeeds with a rows x columns output whose cell (row, column), at index
65	# row * columns + column, holds `cell(row, column)`, and both returned inputs
66	# still hold every word.
67	type MatmulContract<-left: Shape.Shape, -right: Shape.Shape, -rows: Nat, -columns: Nat, -cell: Nat -> Nat -> F32,
68	  -left_word: U32 -> F32, -right_word: U32 -> F32,
69	  -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Type:
70	  MatmulContract{
71	    successful: {product_succeeded(result) == True{} : Bool},
72	    extents: {product_extents(result) == Some{[rows, columns]} : Maybe<&2, List<&2, Nat>>},
73	    value: @+row: Nat -> @row_below: {Nat.is_lt(row, rows) == True{} : Bool} ->
74	      @+column: Nat -> @column_below: {Nat.is_lt(column, columns) == True{} : Bool} ->
75	      {product_point(result, U32.from_nat(Nat.add(Nat.mul(row, columns), column))) == Some{cell(row, column)} : Maybe<&2, F32>},
76	    left_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
77	      {product_left_point(result, other) == Some{left_word(other)} : Maybe<&2, F32>},
78	    right_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
79	      {product_right_point(result, other) == Some{right_word(other)} : Maybe<&2, F32>}}