1	# Indexed mapping with a runtime context and a closed index template. The
2	# source and destination traversal are shared with checked tensor operations.
3	import Base
4	import ./traversal_array.bend as Traversal
5	import ./traversal_loop.bend as Loop
6	import ./default_machine_limits.bend as Default
7	
8	def value(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
9	  context: Context,position: U32,result: Array<Element> & Element) -> Array<Element> & Element:
10	  (source,element) = result
11	  (source,operation(index(context,position),element))
12	
13	def read(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
14	  source: Array<Element>,context: Context,position: U32) -> Array<Element> & Element:
15	  +position = position
16	  value(~Element,~Context,~index,~operation,context,position,Array.get(Element,source,position))
17	
18	def destination(~Context: Data,context: Context,position: U32) -> U32:
19	  position
20	
21	type Call<-Context: Data> is Data:
22	  Call{count: Nat,context: Context}
23	
24	def run(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
25	  call: Call<Context>,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
26	  Call{count,context} = call
27	  Traversal.finish(~Array<Element>,~Element,~Context,
28	    Traversal.run(~Array<Element>,~Element,~Context,~read(~Element,~Context,~index,~operation),~destination(~Context),count,
29	      Traversal.State{input,output,0,context}))
30	
31	# The same mapping over blocks of consecutive positions. A block starts where
32	# the previous one ended and gives position p the index base + (p - start),
33	# so an index that is a quotient and a remainder of the position costs one
34	# base per block instead of a division per element.
35	type Block is Data:
36	  Block{base: U32,start: U32}
37	
38	def block_index(block: Block,position: U32) -> U32:
39	  Block{base,start} = block
40	  (base + (position - start) : U32)
41	
42	type Blocks<-Context: Data> is Data:
43	  Blocks{block: U32,size: U32,context: Context}
44	
45	def next_block(~Element: Data,~Context: Data,block: U32,size: U32,context: Context,
46	  state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Block>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:
47	  Traversal.State{input,output,position,inner} = state
48	  Traversal.State{input,output,position,Blocks{(block + 1 : U32),size,context}}
49	
50	def block_step(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
51	  state: Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>) -> Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>:
52	  Traversal.State{input,output,+position,Blocks{+block,+size,+context}} = state
53	  next_block(~Element,~Context,block,size,context,
54	    Traversal.run(~Array<Element>,~Element,~Block,~read(~Element,~Block,~block_index,~operation),~destination(~Block),U32.to_nat(size),
55	      Traversal.State{input,output,position,Block{base(context,block),position}}))
56	
57	def run_blocks(~Element: Data,~Context: Data,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
58	  blocks: Nat,size: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
59	  Traversal.finish(~Array<Element>,~Element,~Blocks<Context>,
60	    Loop.run(~Traversal.State<&1,&1,Array<Element>,Array<Element>,Blocks<Context>>,~block_step(~Element,~Context,~base,~operation),blocks,
61	      Traversal.State{input,output,0,Blocks{0,size,context}}))
62	
63	def run_whole(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~operation: U32 -> Element -> Element,
64	  whole: Bool,blocks: U32,size: U32,total: U32,context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
65	  match whole:
66	    case True{}: run_blocks(~Element,~Context,~base,~operation,U32.to_nat(blocks),size,context,input,output)
67	    case False{}: run(~Element,~Context,~index,~operation,Call{U32.to_nat(total),context},input,output)
68	
69	# Blocks are used only up to the machine capacity, where the block size
70	# divides a position in U32 without overflow. No storage is larger.
71	def block_limit() -> U32:
72	  Default.capacity()
73	
74	# total positions in blocks(context) blocks of size(context) positions. The
75	# blocks are used only when they are exactly the total and the total is
76	# within the limit; any other count keeps the index function.
77	def run_tiled(~Element: Data,~Context: Data,~index: Context -> U32 -> U32,~base: Context -> U32 -> U32,~blocks: Context -> U32,~size: Context -> U32,
78	  ~operation: U32 -> Element -> Element,+total: U32,+context: Context,input: Array<Element>,output: Array<Element>) -> Array<Element> & Array<Element>:
79	  run_whole(~Element,~Context,~index,~base,~operation,U32.is_le(total,block_limit()) && Nat.is_eq(Nat.mul(U32.to_nat(blocks(context)),U32.to_nat(size(context))),U32.to_nat(total)),
80	    blocks(context),size(context),total,context,input,output)