1	# One owned read/transform/write traversal. Reader and address mapping are static
2	# templates, so composition does not create a closure for each element.
3	import Base
4	import ./traversal_loop.bend as Loop
5	
6	def write_list(~Element: Data,values: List<&2,Element>,+offset: U32,storage: Array<Element>) -> Array<Element>:
7	  match values:
8	    case []: storage
9	    case head <> tail:
10	      write_list(~Element,tail,(offset + 1 : U32),Array.set(Element,storage,offset,head))
11	
12	type State<reader_kind,output_kind,-Reader: Kind(reader_kind),-Output: Kind(output_kind),-Context: Data> is Kind(reader_kind <&> output_kind):
13	  State{reader: Reader,output: Output,position: U32,context: Context}
14	
15	# Indexed state updates share the traversal counter without requiring an array
16	# destination. The caller establishes any no-wrap bound on the U32 position.
17	def indexed_step(~Value: Type,~Context: Data,~update: U32 -> Context -> Value -> Value,
18	  state: State<&2,&1,Unit,Value,Context>) -> State<&2,&1,Unit,Value,Context>:
19	  State{Unit{},value,+position,+context} = state
20	  State{Unit{},update(position,context,value),(position + 1 : U32),context}
21	
22	def indexed_output(~Value: Type,~Context: Data,state: State<&2,&1,Unit,Value,Context>) -> Value:
23	  State{Unit{},value,position,context} = state
24	  value
25	
26	def indexed_fold(~Value: Type,~Context: Data,~update: U32 -> Context -> Value -> Value,
27	  count: Nat,position: U32,context: Context,value: Value) -> Value:
28	  indexed_output(~Value,~Context,
29	    Loop.run(~State<&2,&1,Unit,Value,Context>,~indexed_step(~Value,~Context,~update),count,
30	      State{Unit{},value,position,context}))
31	
32	def write_state(~Reader: Type,~Element: Data,~Context: Data,reader: Reader,output: Array<Element>,position: U32,context: Context,
33	  destination: U32,value: Element) -> State<&1,&1,Reader,Array<Element>,Context>:
34	  State{reader,Array.set(Element,output,destination,value),(position + 1 : U32),context}
35	
36	def store(~Reader: Type,~Element: Data,~Context: Data,~destination: Context -> U32 -> U32,
37	  +position: U32,+context: Context,output: Array<Element>,result: Reader & Element) -> State<&1,&1,Reader,Array<Element>,Context>:
38	  (reader,value) = result
39	  write_state(~Reader,~Element,~Context,reader,output,position,context,destination(context,position),value)
40	
41	def step(~Reader: Type,~Element: Data,~Context: Data,
42	  ~read: Reader -> Context -> U32 -> Reader & Element,~destination: Context -> U32 -> U32,
43	  state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
44	  State{reader,output,+position,+context} = state
45	  store(~Reader,~Element,~Context,~destination,position,context,output,read(reader,context,position))
46	
47	def run(~Reader: Type,~Element: Data,~Context: Data,
48	  ~read: Reader -> Context -> U32 -> Reader & Element,~destination: Context -> U32 -> U32,
49	  count: Nat,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
50	  Loop.run(~State<&1,&1,Reader,Array<Element>,Context>,~step(~Reader,~Element,~Context,~read,~destination),count,state)
51	
52	# Return both owners when the caller can reuse the source allocation later.
53	def finish(~Reader: Type,~Element: Data,~Context: Data,state: State<&1,&1,Reader,Array<Element>,Context>) -> Reader & Array<Element>:
54	  State{reader,output,position,context} = state
55	  (reader,output)
56	
57	def index(context: Unit,position: U32) -> U32:
58	  position
59	
60	def read_index(~Element: Data,storage: Array<Element>,context: Unit,position: U32) -> Array<Element> & Element:
61	  Array.get(Element,storage,position)
62	
63	type Offsets is Data:
64	  Offsets{source: U32,destination: U32}
65	
66	def destination_offset(context: Offsets,position: U32) -> U32:
67	  Offsets{source,destination} = context
68	  (destination + position : U32)
69	
70	def source_offset(~Element: Data,input: Array<Element>,context: Offsets,position: U32) -> Array<Element> & Element:
71	  Offsets{source,destination} = context
72	  Array.get(Element,input,(source + position : U32))
73	
74	def copy(~Element: Data,count: Nat,source: U32,destination: U32,input: Array<Element>,target: Array<Element>) -> Array<Element> & Array<Element>:
75	  finish(~Array<Element>,~Element,~Offsets,
76	    run(~Array<Element>,~Element,~Offsets,~source_offset(~Element),~destination_offset,count,
77	      State{input,target,0,Offsets{source,destination}}))
78	
79	# Ordered folds share the same state/counter and preserve the combine order.
80	# Element and Accumulator can differ; no associativity assumption is required.
81	def fold_value(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,
82	  ~combine: Accumulator -> Element -> Accumulator,accumulator: Accumulator,
83	  position: U32,context: Context,result: Reader & Element) -> State<&1,&2,Reader,Accumulator,Context>:
84	  (reader,value) = result
85	  State{reader,combine(accumulator,value),(position + 1 : U32),context}
86	
87	def fold_step(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,
88	  ~read: Reader -> Context -> U32 -> Reader & Element,~combine: Accumulator -> Element -> Accumulator,
89	  state: State<&1,&2,Reader,Accumulator,Context>) -> State<&1,&2,Reader,Accumulator,Context>:
90	  State{reader,accumulator,+position,+context} = state
91	  fold_value(~Reader,~Element,~Accumulator,~Context,~combine,accumulator,position,context,read(reader,context,position))
92	
93	def finish_fold(~Reader: Type,~Accumulator: Data,~Context: Data,state: State<&1,&2,Reader,Accumulator,Context>) -> Reader & Accumulator:
94	  State{reader,accumulator,position,context} = state
95	  (reader,accumulator)
96	
97	def fold(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,
98	  ~read: Reader -> Context -> U32 -> Reader & Element,~combine: Accumulator -> Element -> Accumulator,
99	  count: Nat,position: U32,reader: Reader,initial: Accumulator,context: Context) -> Reader & Accumulator:
100	  finish_fold(~Reader,~Accumulator,~Context,
101	    Loop.run(~State<&1,&2,Reader,Accumulator,Context>,~fold_step(~Reader,~Element,~Accumulator,~Context,~read,~combine),
102	      count,State{reader,initial,position,context}))
103	
104	# Reuse the output owner for each read/write; no second array or clone.
105	def mapped_value(~Element: Data,~Context: Data,~map: Context -> Element -> Element,
106	  +position: U32,+context: Context,result: Array<Element> & Element) -> State<&2,&1,Unit,Array<Element>,Context>:
107	  (output,value) = result
108	  State{Unit{},Array.set(Element,output,position,map(context,value)),(position + 1 : U32),context}
109	
110	def map_step(~Element: Data,~Context: Data,~map: Context -> Element -> Element,
111	  state: State<&2,&1,Unit,Array<Element>,Context>) -> State<&2,&1,Unit,Array<Element>,Context>:
112	  State{Unit{},output,+position,context} = state
113	  mapped_value(~Element,~Context,~map,position,context,Array.get(Element,output,position))
114	
115	def mapped_output(~Element: Data,~Context: Data,state: State<&2,&1,Unit,Array<Element>,Context>) -> Array<Element>:
116	  State{Unit{},output,position,context} = state
117	  output
118	
119	def map_in_place(~Element: Data,~Context: Data,~map: Context -> Element -> Element,
120	  count: Nat,position: U32,context: Context,output: Array<Element>) -> Array<Element>:
121	  mapped_output(~Element,~Context,Loop.run(~State<&2,&1,Unit,Array<Element>,Context>,
122	    ~map_step(~Element,~Context,~map),count,State{Unit{},output,position,context}))
123	
124	# Read the destination's current element, combine it with an owned reader, then
125	# store through the same traversal state. Distinct Array owners cannot alias.
126	def update_sample(~Reader: Type,~Element: Data,~Sample: Data,~Context: Data,
127	  ~combine: Element -> Sample -> Element,~destination: Context -> U32 -> U32,
128	  position: U32,context: Context,output: Array<Element>,current: Element,result: Reader & Sample) -> State<&1,&1,Reader,Array<Element>,Context>:
129	  (reader,sample) = result
130	  store(~Reader,~Element,~Context,~destination,position,context,output,(reader,combine(current,sample)))
131	
132	def update_current(~Reader: Type,~Element: Data,~Sample: Data,~Context: Data,
133	  ~read: Reader -> Context -> U32 -> Reader & Sample,~combine: Element -> Sample -> Element,
134	  ~destination: Context -> U32 -> U32,reader: Reader,+position: U32,+context: Context,result: Array<Element> & Element) -> State<&1,&1,Reader,Array<Element>,Context>:
135	  (output,current) = result
136	  update_sample(~Reader,~Element,~Sample,~Context,~combine,~destination,position,context,output,current,read(reader,context,position))
137	
138	def update_step(~Reader: Type,~Element: Data,~Sample: Data,~Context: Data,
139	  ~read: Reader -> Context -> U32 -> Reader & Sample,~combine: Element -> Sample -> Element,
140	  ~destination: Context -> U32 -> U32,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
141	  State{reader,output,+position,+context} = state
142	  update_current(~Reader,~Element,~Sample,~Context,~read,~combine,~destination,reader,position,context,
143	    Array.get(Element,output,destination(context,position)))
144	
145	def update(~Reader: Type,~Element: Data,~Sample: Data,~Context: Data,
146	  ~read: Reader -> Context -> U32 -> Reader & Sample,~combine: Element -> Sample -> Element,
147	  ~destination: Context -> U32 -> U32,count: Nat,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
148	  Loop.run(~State<&1,&1,Reader,Array<Element>,Context>,
149	    ~update_step(~Reader,~Element,~Sample,~Context,~read,~combine,~destination),count,state)
150	
151	# A cursor is consumed by read and returned advanced. The state position is the
152	# absolute contiguous destination, so it need not retain a second cursor copy.
153	type Visited<-Context: Data,-Element: Data> is Data:
154	  Visited{context: Context,value: Element}
155	
156	def visit_store(~Reader: Type,~Element: Data,~Context: Data,+position: U32,output: Array<Element>,result: Reader & Visited<Context,Element>) ->
157	  State<&1,&1,Reader,Array<Element>,Context>:
158	  (reader,Visited{context,value}) = result
159	  write_state(~Reader,~Element,~Context,reader,output,position,context,position,value)
160	
161	def visit_step(~Reader: Type,~Element: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
162	  state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
163	  State{reader,output,position,context} = state
164	  visit_store(~Reader,~Element,~Context,position,output,read(reader,context))
165	
166	def visit(~Reader: Type,~Element: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
167	  count: Nat,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
168	  Loop.run(~State<&1,&1,Reader,Array<Element>,Context>,~visit_step(~Reader,~Element,~Context,~read),count,state)
169	
170	# A scatter reads the source at the state position, which counts up, and
171	# stores at the address its context names; target yields that address with
172	# the context of the next element.
173	type Target<-Context: Data> is Data:
174	  Target{context: Context,address: U32}
175	
176	def spread_store(~Element: Data,~Context: Data,position: U32,output: Array<Element>,target: Target<Context>,result: Array<Element> & Element) ->
177	  State<&1,&1,Array<Element>,Array<Element>,Context>:
178	  Target{context,address} = target
179	  (input,value) = result
180	  State{input,Array.set(Element,output,address,value),(position + 1 : U32),context}
181	
182	def spread_step(~Element: Data,~Context: Data,~target: Context -> Target<Context>,
183	  state: State<&1,&1,Array<Element>,Array<Element>,Context>) -> State<&1,&1,Array<Element>,Array<Element>,Context>:
184	  State{input,output,+position,context} = state
185	  spread_store(~Element,~Context,position,output,target(context),Array.get(Element,input,position))
186	
187	def visit_update_sample(~Reader: Type,~Element: Data,~Context: Data,~combine: Element -> Element -> Element,
188	  +position: U32,output: Array<Element>,current: Element,result: Reader & Visited<Context,Element>) -> State<&1,&1,Reader,Array<Element>,Context>:
189	  (reader,Visited{context,value}) = result
190	  write_state(~Reader,~Element,~Context,reader,output,position,context,position,combine(current,value))
191	
192	def visit_update_current(~Reader: Type,~Element: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
193	  ~combine: Element -> Element -> Element,position: U32,context: Context,reader: Reader,result: Array<Element> & Element) -> State<&1,&1,Reader,Array<Element>,Context>:
194	  (output,current) = result
195	  visit_update_sample(~Reader,~Element,~Context,~combine,position,output,current,read(reader,context))
196	
197	def visit_update_step(~Reader: Type,~Element: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
198	  ~combine: Element -> Element -> Element,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
199	  State{reader,output,+position,context} = state
200	  visit_update_current(~Reader,~Element,~Context,~read,~combine,position,context,reader,Array.get(Element,output,position))
201	
202	def visit_update(~Reader: Type,~Element: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
203	  ~combine: Element -> Element -> Element,count: Nat,state: State<&1,&1,Reader,Array<Element>,Context>) -> State<&1,&1,Reader,Array<Element>,Context>:
204	  Loop.run(~State<&1,&1,Reader,Array<Element>,Context>,~visit_update_step(~Reader,~Element,~Context,~read,~combine),count,state)
205	
206	def visit_fold_value(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,~combine: Accumulator -> Element -> Accumulator,
207	  position: U32,accumulator: Accumulator,result: Reader & Visited<Context,Element>) -> State<&1,&2,Reader,Accumulator,Context>:
208	  (reader,Visited{context,value}) = result
209	  State{reader,combine(accumulator,value),(position + 1 : U32),context}
210	
211	def visit_fold_step(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
212	  ~combine: Accumulator -> Element -> Accumulator,state: State<&1,&2,Reader,Accumulator,Context>) -> State<&1,&2,Reader,Accumulator,Context>:
213	  State{reader,accumulator,position,context} = state
214	  visit_fold_value(~Reader,~Element,~Accumulator,~Context,~combine,position,accumulator,read(reader,context))
215	
216	def visit_fold(~Reader: Type,~Element: Data,~Accumulator: Data,~Context: Data,~read: Reader -> Context -> Reader & Visited<Context,Element>,
217	  ~combine: Accumulator -> Element -> Accumulator,count: Nat,state: State<&1,&2,Reader,Accumulator,Context>) -> State<&1,&2,Reader,Accumulator,Context>:
218	  Loop.run(~State<&1,&2,Reader,Accumulator,Context>,~visit_fold_step(~Reader,~Element,~Accumulator,~Context,~read,~combine),count,state)