1	# Parallel map and reduce over lists and arrays, a worker pool, select over channels, and timeouts.
2	import Base
3	
4	#   import bend-kit-concurrency@0.2.0.0/concurrency.bend as Conc
5	#   ys = Conc.par_map(~U32, ~U32, ~work, 8, xs)
6	#
7	# par_map and par_reduce are pure. They split the input into at most `workers`
8	# balanced parts (rounded down to a power of two) and run the parts as parallel
9	# calls, so a native build spreads them over its `--threads`. Results keep the
10	# input order. par_reduce needs an associative f; the list form also needs z to
11	# be f's identity, since each part starts from z.
12	#
13	# pool, select, and timeout are IO, on Base's Chan, IO.spawn and IO.fork. Base
14	# has no cancellation: an abandoned computation keeps running, and the program
15	# exits only when every computation is done.
16	
17	# The length of xs, beside xs.
18	def count.cons(-A: Type, h: A, r: Nat & List<A>) -> Nat & List<A>:
19	  (n, t) = r
20	  (1n+n, h <> t)
21	
22	def count(-A: Type, xs: List<A>) -> Nat & List<A>:
23	  match xs:
24	    case Nil{}:
25	      (0n, Nil{})
26	    case h <> t:
27	      count.cons(A, h, count(A, t))
28	
29	# The first n elements of xs, and the rest.
30	def split.cons(-A: Type, h: A, r: List<A> & List<A>) -> List<A> & List<A>:
31	  (l, rest) = r
32	  (h <> l, rest)
33	
34	def split(-A: Type, xs: List<A>, n: Nat) -> List<A> & List<A>:
35	  match xs n:
36	    case Nil{} _:
37	      (Nil{}, Nil{})
38	    case h <> t 0n:
39	      (Nil{}, h <> t)
40	    case h <> t 1n+p:
41	      split.cons(A, h, split(A, t, p))
42	
43	# lr holds two halves of nl and nr elements. Depth d runs 2^(d+1) parallel parts.
44	def par_map.list.go(
45	  ~A: Type, ~B: Type, ~f: A -> B, d: Nat, +nl: Nat, +nr: Nat, lr: List<A> & List<A>
46	) -> List<B>:
47	  match d:
48	    case 0n:
49	      (l, r) = lr
50	      lo hi = List.map(~A, ~B, ~f, l) List.map(~A, ~B, ~f, r)
51	      List.append(&1, B, lo, hi)
52	    case 1n+ +p:
53	      (l, r) = lr
54	      +hl = Nat.div(nl, 2n)
55	      +hr = Nat.div(nr, 2n)
56	      ls = split(A, l, hl)
57	      rs = split(A, r, hr)
58	      lo hi = par_map.list.go(~A, ~B, ~f, p, hl, Nat.sub(nl, hl), ls) par_map.list.go(~A, ~B, ~f, p, hr, Nat.sub(nr, hr), rs)
59	      List.append(&1, B, lo, hi)
60	
61	def par_map.list.top(~A: Type, ~B: Type, ~f: A -> B, d: Nat, nx: Nat & List<A>) -> List<B>:
62	  match d:
63	    case 0n:
64	      (n, xs) = nx
65	      List.map(~A, ~B, ~f, xs)
66	    case 1n+p:
67	      (+n, xs) = nx
68	      +h = Nat.div(n, 2n)
69	      par_map.list.go(~A, ~B, ~f, p, h, Nat.sub(n, h), split(A, xs, h))
70	
71	# List.map(f, xs), on at most `workers` parallel parts.
72	def par_map(~A: Type, ~B: Type, ~f: A -> B, workers: U32, xs: List<A>) -> List<B>:
73	  par_map.list.top(~A, ~B, ~f, U32.log2(workers), count(A, xs))
74	
75	def fold(~A: Data, ~f: A -> A -> A, xs: List<A>, acc: A) -> A:
76	  match xs:
77	    case Nil{}:
78	      acc
79	    case h <> t:
80	      fold(~A, ~f, t, f(acc, h))
81	
82	def par_reduce.list.go(
83	  ~A: Data, ~f: A -> A -> A, d: Nat, +nl: Nat, +nr: Nat, +z: A, lr: List<A> & List<A>
84	) -> A:
85	  match d:
86	    case 0n:
87	      (l, r) = lr
88	      lo hi = fold(~A, ~f, l, z) fold(~A, ~f, r, z)
89	      f(lo, hi)
90	    case 1n+ +p:
91	      (l, r) = lr
92	      +hl = Nat.div(nl, 2n)
93	      +hr = Nat.div(nr, 2n)
94	      ls = split(A, l, hl)
95	      rs = split(A, r, hr)
96	      lo hi = par_reduce.list.go(~A, ~f, p, hl, Nat.sub(nl, hl), z, ls) par_reduce.list.go(~A, ~f, p, hr, Nat.sub(nr, hr), z, rs)
97	      f(lo, hi)
98	
99	def par_reduce.list.top(~A: Data, ~f: A -> A -> A, d: Nat, +z: A, nx: Nat & List<A>) -> A:
100	  match d:
101	    case 0n:
102	      (n, xs) = nx
103	      fold(~A, ~f, xs, z)
104	    case 1n+p:
105	      (+n, xs) = nx
106	      +h = Nat.div(n, 2n)
107	      par_reduce.list.go(~A, ~f, p, h, Nat.sub(n, h), z, split(A, xs, h))
108	
109	# The left fold of xs from z, on at most `workers` parallel parts. f must be
110	# associative and z its identity.
111	def par_reduce(~A: Data, ~f: A -> A -> A, workers: U32, +z: A, xs: List<A>) -> A:
112	  par_reduce.list.top(~A, ~f, U32.log2(workers), z, count(A, xs))
113	
114	# An Array is a balanced tree, so each ANode above depth d is a parallel call.
115	def par_map.array.go(~A: Type, ~B: Type, ~f: A -> B, a: Array<A>, d: Nat) -> Array<B>:
116	  match a d:
117	    case ALeaf{x} _:
118	      ALeaf{f(x)}
119	    case ANode{xs, ys} 0n:
120	      ANode{par_map.array.go(~A, ~B, ~f, xs, 0n), par_map.array.go(~A, ~B, ~f, ys, 0n)}
121	    case ANode{xs, ys} 1n+ +p:
122	      lo hi = par_map.array.go(~A, ~B, ~f, xs, p) par_map.array.go(~A, ~B, ~f, ys, p)
123	      ANode{lo, hi}
124	
125	# f on every slot of a, on at most `workers` parallel parts. Slots keep their index.
126	def par_map.array(~A: Type, ~B: Type, ~f: A -> B, workers: U32, a: Array<A>) -> Array<B>:
127	  par_map.array.go(~A, ~B, ~f, a, U32.log2(workers))
128	
129	def par_reduce.array.go(~A: Type, ~f: A -> A -> A, a: Array<A>, d: Nat) -> A:
130	  match a d:
131	    case ALeaf{x} _:
132	      x
133	    case ANode{xs, ys} 0n:
134	      f(par_reduce.array.go(~A, ~f, xs, 0n), par_reduce.array.go(~A, ~f, ys, 0n))
135	    case ANode{xs, ys} 1n+ +p:
136	      lo hi = par_reduce.array.go(~A, ~f, xs, p) par_reduce.array.go(~A, ~f, ys, p)
137	      f(lo, hi)
138	
139	# The slots of a combined in index order with an associative f, on at most
140	# `workers` parallel parts. An Array is never empty, so no identity is needed.
141	def par_reduce.array(~A: Type, ~f: A -> A -> A, workers: U32, a: Array<A>) -> A:
142	  par_reduce.array.go(~A, ~f, a, U32.log2(workers))
143	
144	# A pool worker between channel steps. It owns its state s alone.
145	type Worker<-S: Type, -A: Type, -B: Type> is Type:
146	  WIdle{s: S}
147	  WGot{s: S, got: Maybe<&1, A & Chan(B)>}
148	  WDone{sb: S & B, reply: Chan(B)}
149	
150	# Takes jobs until the queue is closed and empty, then answers its state.
151	# Each job takes three steps, so 3 * (jobs + 1) fuel is enough.
152	def pool.worker(
153	  ~S: Type, ~A: Type, ~B: Type, ~work: S -> A -> IO(S & B), fuel: Nat,
154	  +q: Chan(A & Chan(B)), w: Worker<S, A, B>
155	) -> IO(S):
156	  match fuel:
157	    case 0n:
158	      IO.die(S, 1, "Conc.pool: a worker ran out of fuel")
159	    case 1n+p:
160	      match w:
161	        case WIdle{s}:
162	          do IO<S>:
163	            got : Maybe<&1, A & Chan(B)> <- Chan.recv(A & Chan(B), q)
164	            pool.worker(~S, ~A, ~B, ~work, p, q, WGot{s, got})
165	        case WGot{s, got}:
166	          match got:
167	            case None{}:
168	              IO.pure(S, s)
169	            case Some{(a, reply)}:
170	              do IO<S>:
171	                sb : S & B <- work(s, a)
172	                pool.worker(~S, ~A, ~B, ~work, p, q, WDone{sb, reply})
173	        case WDone{sb, reply}:
174	          (s, b) = sb
175	          do IO<S>:
176	            sent : Result<&1, &1, B, Unit> <- Chan.send(B, reply, b)
177	            pool.worker(~S, ~A, ~B, ~work, p, q, WIdle{s})
178	
179	# Queues each job with its own reply channel, and answers the replies in order.
180	def pool.put(~A: Type, ~B: Type, xs: List<A>, +q: Chan(A & Chan(B))) -> IO(List<Chan(B)>):
181	  match xs:
182	    case Nil{}:
183	      IO.pure(List<Chan(B)>, Nil{})
184	    case x <> rest:
185	      do IO<List<Chan(B)>>:
186	        reply : Chan(B) <- Chan.new(B, 1)
187	        +r : Chan(B) = reply
188	        sent : Result<&1, &1, A & Chan(B), Unit> <- Chan.send(A & Chan(B), q, (x, r))
189	        replies : List<Chan(B)> <- pool.put(~A, ~B, rest, q)
190	        return r <> replies
191	
192	def pool.start(
193	  ~S: Type, ~A: Type, ~B: Type, ~work: S -> A -> IO(S & B), states: List<S>,
194	  +fuel: Nat, +q: Chan(A & Chan(B))
195	) -> IO(List<Chan(S)>):
196	  match states:
197	    case Nil{}:
198	      IO.pure(List<Chan(S)>, Nil{})
199	    case s <> rest:
200	      do IO<List<Chan(S)>>:
201	        c : Chan(S) <- IO.fork(S, pool.worker(~S, ~A, ~B, ~work, fuel, q, WIdle{s}))
202	        cs : List<Chan(S)> <- pool.start(~S, ~A, ~B, ~work, rest, fuel, q)
203	        return c <> cs
204	
205	def pool.join(-A: Type, cs: List<Chan(A)>) -> IO(List<A>):
206	  match cs:
207	    case Nil{}:
208	      IO.pure(List<A>, Nil{})
209	    case c <> rest:
210	      do IO<List<A>>:
211	        x : A <- IO.join(A, c)
212	        xs : List<A> <- pool.join(A, rest)
213	        return x <> xs
214	
215	def pool.run(
216	  ~S: Type, ~A: Type, ~B: Type, ~work: S -> A -> IO(S & B), states: List<S>,
217	  nx: Nat & List<A>
218	) -> IO(List<S> & List<B>):
219	  (+n, xs) = nx
220	  do IO<List<S> & List<B>>:
221	    queue : Chan(A & Chan(B)) <- Chan.new(A & Chan(B), U32.from_nat(n))
222	    +q : Chan(A & Chan(B)) = queue
223	    replies : List<Chan(B)> <- pool.put(~A, ~B, xs, q)
224	    Chan.close(A & Chan(B), q)
225	    workers : List<Chan(S)> <- pool.start(~S, ~A, ~B, ~work, states, Nat.mul(3n, 1n+n), q)
226	    outs : List<B> <- pool.join(B, replies)
227	    done : List<S> <- pool.join(S, workers)
228	    return (done, outs)
229	
230	# Runs work(s, x) for every x in xs, on one worker per state in states. A
231	# worker owns its state and threads it through its jobs, so an affine handle
232	# such as a connection pool is never shared. Answers the final states in
233	# worker order and the results in the order of xs.
234	def pool(
235	  ~S: Type, ~A: Type, ~B: Type, ~work: S -> A -> IO(S & B), states: List<S>, xs: List<A>
236	) -> IO(List<S> & List<B>):
237	  match states:
238	    case Nil{}:
239	      IO.die(List<S> & List<B>, 1, "Conc.pool: no workers")
240	    case s <> rest:
241	      pool.run(~S, ~A, ~B, ~work, s <> rest, count(A, xs))
242	
243	# A select relay between channel steps.
244	type Relay<-A: Type> is Type:
245	  RWait{}
246	  RGot{got: Maybe<&1, A>}
247	  RSent{ok: Result<&1, &1, U32 & A, Unit>}
248	
249	# Forwards each value of src to out as (i, value), then reports on done.
250	# It stops when src closes, or when out is closed under it.
251	def select.relay(
252	  -A: Type, fuel: Nat, +i: U32, +src: Chan(A), +out: Chan(U32 & A),
253	  +done: Chan(Unit), r: Relay<A>
254	) -> IO(Result<&1, &1, Unit, Unit>):
255	  match fuel:
256	    case 0n:
257	      Chan.send(Unit, done, Unit{})
258	    case 1n+p:
259	      match r:
260	        case RWait{}:
261	          do IO<Result<&1, &1, Unit, Unit>>:
262	            got : Maybe<&1, A> <- Chan.recv(A, src)
263	            select.relay(A, p, i, src, out, done, RGot{got})
264	        case RGot{got}:
265	          match got:
266	            case None{}:
267	              Chan.send(Unit, done, Unit{})
268	            case Some{x}:
269	              do IO<Result<&1, &1, Unit, Unit>>:
270	                ok : Result<&1, &1, U32 & A, Unit> <- Chan.send(U32 & A, out, (i, x))
271	                select.relay(A, p, i, src, out, done, RSent{ok})
272	        case RSent{ok}:
273	          match ok:
274	            case Done{u}:
275	              select.relay(A, p, i, src, out, done, RWait{})
276	            case Fail{held}:
277	              Chan.send(Unit, done, Unit{})
278	
279	# ponytail: relay fuel is the largest Nat literal, and a value takes three
280	# steps, so a relay stops after about 1.4 billion values. Drop the fuel if
281	# Base gains an unbounded loop.
282	def select.start(
283	  -A: Type, chans: List<Chan(A)>, +i: U32, +out: Chan(U32 & A), +done: Chan(Unit)
284	) -> IO(Nat):
285	  match chans:
286	    case Nil{}:
287	      IO.pure(Nat, 0n)
288	    case src <> rest:
289	      do IO<Nat>:
290	        IO.spawn(Result<&1, &1, Unit, Unit>, select.relay(A, 4294967295n, i, src, out, done, RWait{}))
291	        n : Nat <- select.start(A, rest, (i + 1 : U32), out, done)
292	        return 1n+n
293	
294	# Closes out once all n relays have reported.
295	def select.close(-A: Type, n: Nat, +out: Chan(U32 & A), +done: Chan(Unit)) -> IO(Unit):
296	  match n:
297	    case 0n:
298	      do IO<Unit>:
299	        Chan.close(U32 & A, out)
300	        Chan.close(Unit, done)
301	    case 1n+p:
302	      do IO<Unit>:
303	        got : Maybe<&1, Unit> <- Chan.recv(Unit, done)
304	        select.close(A, p, out, done)
305	
306	def select.go(-A: Type, chans: List<Chan(A)>, +out: Chan(U32 & A), +done: Chan(Unit)) ->
307	  IO(Chan(U32 & A)):
308	  do IO<Chan(U32 & A)>:
309	    n : Nat <- select.start(A, chans, 0, out, done)
310	    IO.spawn(Unit, select.close(A, n, out, done))
311	    return out
312	
313	# One channel that carries (i, value) for each value received on chans[i], in
314	# the order they arrive. Chan.recv on it waits for whichever source is ready
315	# first, and answers None once every source is closed. The selector owns the
316	# receiving side of its sources. Close it to stop early: each relay then drops
317	# the one value it holds, if any, and stops.
318	def select(-A: Type, chans: List<Chan(A)>) -> IO(Chan(U32 & A)):
319	  do IO<Chan(U32 & A)>:
320	    out : Chan(U32 & A) <- Chan.new(U32 & A, 0)
321	    done : Chan(Unit) <- Chan.new(Unit, 0)
322	    select.go(A, chans, out, done)
323	
324	def timeout.won(-A: Type, +c: Chan(Maybe<&1, A>), got: Maybe<&1, Maybe<&1, A>>) ->
325	  IO(Maybe<&1, A>):
326	  match got:
327	    case None{}:
328	      IO.pure(Maybe<&1, A>, None{})
329	    case Some{m}:
330	      do IO<Maybe<&1, A>>:
331	        Chan.close(Maybe<&1, A>, c)
332	        return m
333	
334	# Some{result} if act finishes within ms milliseconds, else None at the
335	# deadline. act is not cancelled: it runs on, its late result is dropped, and
336	# the program waits for it (and for the timer) before it exits.
337	def timeout(-A: Type, ms: U32, act: IO(A)) -> IO(Maybe<&1, A>):
338	  do IO<Maybe<&1, A>>:
339	    chan : Chan(Maybe<&1, A>) <- Chan.new(Maybe<&1, A>, 0)
340	    +c : Chan(Maybe<&1, A>) = chan
341	    IO.spawn(Result<&1, &1, Maybe<&1, A>, Unit>, IO.bind(A, Result<&1, &1, Maybe<&1, A>, Unit>, act, x => Chan.send(Maybe<&1, A>, c, Some{x})))
342	    IO.spawn(Result<&1, &1, Maybe<&1, A>, Unit>, IO.bind(Unit, Result<&1, &1, Maybe<&1, A>, Unit>, IO.sleep(ms), u => Chan.send(Maybe<&1, A>, c, None{})))
343	    got : Maybe<&1, Maybe<&1, A>> <- Chan.recv(Maybe<&1, A>, c)
344	    timeout.won(A, c, got)