1	# Shape-only groups, ordered folds from +0, then adjacent-pair trees. The
2	# reader is affine; the same traversal serves weight and bias gradients.
3	import Base
4	import ./traversal_loop.bend as Loop
5	
6	type Range is Data:
7	  Range{first: Nat,count: Nat}
8	
9	def rows(remaining: Nat,+first: Nat,+height: Nat,+width: Nat,+span: Nat) -> List<&2,Range>:
10	  match remaining:
11	    case 0n: []
12	    case 1n+rest:
13	      Range{Nat.mul(first,width),Nat.mul(Nat.min(span,Nat.sub(height,first)),width)} <> rows(rest,Nat.add(first,span),height,width,span)
14	
15	def columns(remaining: Nat,+row: Nat,+first: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
16	  match remaining:
17	    case 0n: []
18	    case 1n+rest:
19	      Range{Nat.add(Nat.mul(row,width),first),Nat.min(group,Nat.sub(width,first))} <> columns(rest,row,Nat.add(first,group),width,group)
20	
21	def cat(left: List<&2,Range>,+right: List<&2,Range>) -> List<&2,Range>:
22	  match left:
23	    case []: right
24	    case head <> tail: head <> cat(tail,right)
25	
26	def wide_rows(remaining: Nat,+row: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
27	  match remaining:
28	    case 0n: []
29	    case 1n+rest:
30	      cat(columns(Nat.div(Nat.add(width,Nat.sub(group,1n)),group),row,0n,width,group),wide_rows(rest,1n+row,width,group))
31	
32	def ranges_selected(wide: Bool,+height: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
33	  match wide:
34	    case True{}: wide_rows(height,0n,width,group)
35	    case False{}:
36	      +span = Nat.div(group,width)
37	      rows(Nat.div(Nat.add(height,Nat.sub(span,1n)),span),0n,height,width,span)
38	
39	def ranges(+height: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
40	  ranges_selected(Nat.is_lt(group,width),height,width,group)
41	
42	def pairs(~ownership: Quant,~Element: Kind(ownership),~Context: Data,
43	  ~combine: Context -> Element -> Element -> Element,+context: Context,values: List<ownership,Element>) -> List<ownership,Element>:
44	  match values:
45	    case []: []
46	    case a <> []: [a]
47	    case a <> b <> tail: combine(context,a,b) <> pairs(~ownership,~Element,~Context,~combine,context,tail)
48	
49	# The original length bounds the number of tree levels. A singleton is
50	# returned directly, including its sign bit; there is no final addition.
51	def levels(~ownership: Quant,~Element: Kind(ownership),~Context: Data,
52	  ~zero: Context -> Element,~combine: Context -> Element -> Element -> Element,
53	  fuel: Nat,+context: Context,values: List<ownership,Element>) -> Element:
54	  match fuel values:
55	    case 0n Nil{}: zero(context)
56	    case 0n head <> tail: head
57	    case 1n+rest Nil{}: zero(context)
58	    case 1n+rest head <> []: head
59	    case 1n+rest a <> b <> tail: levels(~ownership,~Element,~Context,~zero,~combine,rest,context,pairs(~ownership,~Element,~Context,~combine,context,a <> b <> tail))
60	
61	def tree(+values: List<&2,F32>) -> F32:
62	  levels(~&2,~F32,~Unit,~(context => 0.0),~(context => F32.add),List.length(&2,F32,values),Unit{},values)
63	
64	type State<-Reader: Type,-Context: Data> is Type:
65	  State{reader: Reader,context: Context,position: Nat,total: F32}
66	
67	def advanced(~Reader: Type,~Context: Data,+context: Context,position: Nat,total: F32,pair: Reader & F32) -> State<Reader,Context>:
68	  (reader,value) = pair
69	  State{reader,context,1n+position,(total + value : F32)}
70	
71	def step(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,state: State<Reader,Context>) -> State<Reader,Context>:
72	  State{reader,+context,+position,total} = state
73	  advanced(~Reader,~Context,context,position,total,read(reader,context,position))
74	
75	def finished(~Reader: Type,~Context: Data,state: State<Reader,Context>) -> Reader & F32:
76	  State{reader,context,position,total} = state
77	  (reader,total)
78	
79	def fold(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,count: Nat,first: Nat,reader: Reader,context: Context) -> Reader & F32:
80	  finished(~Reader,~Context,Loop.run(~State<Reader,Context>,~step(~Reader,~Context,~read),count,State{reader,context,first,0.0}))
81	
82	def prepend(~Reader: Type,value: F32,pair: Reader & List<&2,F32>) -> Reader & List<&2,F32>:
83	  (reader,values) = pair
84	  (reader,value <> values)
85	
86	def grouped_tail(~Reader: Type,pair: Reader & F32,tail: Reader -> Reader & List<&2,F32>) -> Reader & List<&2,F32>:
87	  (reader,value) = pair
88	  prepend(~Reader,value,tail(reader))
89	
90	def fold_ranges(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,ranges: List<&2,Range>,reader: Reader,+context: Context) -> Reader & List<&2,F32>:
91	  match ranges:
92	    case []: (reader,[])
93	    case Range{first,count} <> rest:
94	      grouped_tail(~Reader,fold(~Reader,~Context,~read,count,first,reader,context),reader => fold_ranges(~Reader,~Context,~read,rest,reader,context))
95	
96	def reduced(~Reader: Type,pair: Reader & List<&2,F32>) -> Reader & F32:
97	  (reader,values) = pair
98	  (reader,tree(values))
99	
100	def run(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,ranges: List<&2,Range>,reader: Reader,context: Context) -> Reader & F32:
101	  reduced(~Reader,fold_ranges(~Reader,~Context,~read,ranges,reader,context))