1	# Owned inputs move through cells holding `width` consecutive output values.
2	# Each round gathers a cell, stores it in scratch, reads it back and scatters
3	# it into the owners. Storing the whole cell lets the compiler compute its
4	# elements together. The Array specialization refines Traversal.map_in_place.
5	import Base
6	import ./traversal_loop.bend as Loop
7	
8	type State<-Owner: Type,-Cell: Data,-Context: Data> is Type:
9	  State{values: Owner,scratch: Array<Cell>,position: U32,slot: U32,context: Context}
10	
11	def scattered(~Owner: Type,~Cell: Data,~Context: Data,~scatter: Owner -> U32 -> Cell -> Owner,~width: U32,
12	  +position: U32,+slot: U32,context: Context,values: Owner,read: Array<Cell> & Cell) -> State<Owner,Cell,Context>:
13	  (scratch,cell) = read
14	  State{scatter(values,position,cell),scratch,(position + width : U32),U32.and((slot + 1 : U32),63),context}
15	
16	def stored(~Owner: Type,~Cell: Data,~Context: Data,~scatter: Owner -> U32 -> Cell -> Owner,~width: U32,
17	  +position: U32,+slot: U32,context: Context,scratch: Array<Cell>,gathered: Owner & Cell) -> State<Owner,Cell,Context>:
18	  (values,cell) = gathered
19	  scattered(~Owner,~Cell,~Context,~scatter,~width,position,slot,context,values,Array.get(Cell,Array.set(Cell,scratch,slot,cell),slot))
20	
21	def step(~Owner: Type,~Cell: Data,~Context: Data,~gather: Owner -> U32 -> Context -> Owner & Cell,
22	  ~scatter: Owner -> U32 -> Cell -> Owner,~width: U32,state: State<Owner,Cell,Context>) -> State<Owner,Cell,Context>:
23	  State{values,scratch,+position,+slot,+context} = state
24	  stored(~Owner,~Cell,~Context,~scatter,~width,position,slot,context,scratch,gather(values,position,context))
25	
26	def output(~Owner: Type,~Cell: Data,~Context: Data,state: State<Owner,Cell,Context>) -> Owner:
27	  State{values,scratch,position,slot,context} = state
28	  values
29	
30	# `cells` rounds from `position`; the scratch array holds 64 cells.
31	def map(~Owner: Type,~Cell: Data,~Context: Data,~gather: Owner -> U32 -> Context -> Owner & Cell,
32	  ~scatter: Owner -> U32 -> Cell -> Owner,~width: U32,
33	  cells: Nat,position: U32,context: Context,values: Owner,scratch: Array<Cell>) -> Owner:
34	  output(~Owner,~Cell,~Context,Loop.run(~State<Owner,Cell,Context>,~step(~Owner,~Cell,~Context,~gather,~scatter,~width),
35	    cells,State{values,scratch,position,0,context}))