1	import Base
2	import ./EdlibTypes.bend as E
3	import ./main.bend as M
4	
5	type WeightedJob is Data:
6	  WeightedJob{index: Nat, work: F32, job: E.AlignJob}
7	
8	type Group is Data:
9	  Group{room: Nat, work: F32, jobs: List<&2, WeightedJob>}
10	
11	# Build the entire plan before entering the alignment fork tree. Recursive
12	# argument preparation in a parallel call can drain the native frontier.
13	type Plan is Type:
14	  Empty{}
15	  Leaf{index: Nat, job: E.AlignJob}
16	  Fork{left: Plan, right: Plan}
17	  Serial{jobs: List<&2, WeightedJob>}
18	
19	type IndexedResult is Data:
20	  IndexedResult{index: Nat, result: E.AlignResult}
21	
22	def axes(mode: E.Mode, +qn: Nat, +tn: Nat) -> Nat & Nat:
23	  match mode:
24	    case E.NW{}:
25	      (Nat.min(qn, tn), Nat.max(qn, tn))
26	    case _:
27	      (qn, tn)
28	
29	def extra_work(mode: E.Mode, task: E.Task, +q: F32, +t: F32,
30	               +blocks: F32, pairs: F32) -> F32:
31	  match mode task:
32	    case _ E.Distance{}:
33	      0.0
34	    case E.HW{} E.Path{}:
35	      (blocks * t * t + q * t * pairs : F32)
36	    case E.HW{} _:
37	      (blocks * t * t : F32)
38	    case _ E.Path{}:
39	      (q * t * pairs : F32)
40	    case _ _:
41	      0.0
42	
43	def estimate_axes(+mode: E.Mode, task: E.Task, pairs: List<&2, E.EqualityPair>,
44	                  sizes: Nat & Nat) -> F32:
45	  (+length, columns) = sizes
46	  +l = F32.from_nat(length)
47	  +c = F32.from_nat(columns)
48	  +blocks = F32.from_nat(M.ceil_div(length, 32n))
49	  +p = (1.0 + F32.from_nat(List.length(&2, E.EqualityPair, pairs)) : F32)
50	  (1.0 + l + c + 256.0 * l * p + blocks * c +
51	    extra_work(mode, task, l, c, blocks, p) : F32)
52	
53	# A scheduling heuristic, not a timing prediction: input scans, 256-symbol
54	# profiles, full Myers columns, worst-case HW start searches and PATH DP.
55	# F32 avoids wrapping a U32 product or allocating a quadratic unary Nat.
56	# Bands, retries, similarity and the selected interval can change actual work.
57	def estimate_sizes(qn: Nat, tn: Nat, config: E.AlignConfig) -> F32:
58	  match qn tn config:
59	    case 0n +t _:
60	      (1.0 + F32.from_nat(t) : F32)
61	    case +q 0n _:
62	      (1.0 + F32.from_nat(q) : F32)
63	    case +q +t E.AlignConfig{+mode, task, _, pairs, _}:
64	      estimate_axes(mode, task, pairs, axes(mode, q, t))
65	
66	def estimate(+job: E.AlignJob) -> F32:
67	  E.AlignJob{query, target, config} = job
68	  # Scans are sequential here; pairs of arbitrary lengths need not balance.
69	  estimate_sizes(String.length(query), String.length(target), config)
70	
71	def indexed(jobs: List<&2, E.AlignJob>, +index: Nat) -> List<&2, WeightedJob>:
72	  match jobs:
73	    case Nil{}:
74	      Nil{}
75	    case +job <> tail:
76	      WeightedJob{index, estimate(job), job} <> indexed(tail, 1n+index)
77	
78	def heavier(first: WeightedJob, second: WeightedJob) -> Bool:
79	  WeightedJob{_, a, _} = first
80	  WeightedJob{_, b, _} = second
81	  F32.is_ge(a, b)
82	
83	def add(group: Group, job: WeightedJob) -> Group:
84	  Group{room, work, jobs} = group
85	  WeightedJob{index, +cost, value} = job
86	  Group{Nat.sub(room, 1n), F32.add(work, cost), WeightedJob{index, cost, value} <> jobs}
87	
88	def choose(left: Group, right: Group) -> Bool:
89	  match left right:
90	    case Group{0n, _, _} _:
91	      False{}
92	    case _ Group{0n, _, _}:
93	      True{}
94	    case Group{_, a, _} Group{_, b, _}:
95	      F32.is_le(a, b)
96	
97	def place(job: WeightedJob, left: Group, right: Group, use_left: Bool) -> Group & Group:
98	  match use_left:
99	    case True{}:
100	      (add(left, job), right)
101	    case False{}:
102	      (left, add(right, job))
103	
104	# Largest estimated jobs go to the lighter branch. Capacity bounds keep both
105	# child counts within ceil(n/2), including zero/rounded/equal work estimates.
106	def partition(jobs: List<&2, WeightedJob>, groups: Group & Group) -> Group & Group:
107	  match jobs:
108	    case Nil{}:
109	      groups
110	    case job <> tail:
111	      (+left, +right) = groups
112	      partition(tail, place(job, left, right, choose(left, right)))
113	
114	def group_jobs(group: Group) -> List<&2, WeightedJob>:
115	  Group{_, _, jobs} = group
116	  List.reverse(&2, WeightedJob, jobs)
117	
118	def plan_step(groups: Group & Group,
119	              next: List<&2, WeightedJob> -> List<&2, WeightedJob> -> Plan) -> Plan:
120	  (left, right) = groups
121	  next(group_jobs(left), group_jobs(right))
122	
123	def plan(fuel: Nat, jobs: List<&2, WeightedJob>) -> Plan:
124	  match fuel jobs:
125	    case _ Nil{}:
126	      Empty{}
127	    case _ WeightedJob{index, _, job} <> Nil{}:
128	      Leaf{index, job}
129	    case 0n values:
130	      Serial{values}
131	    case Succ{+p} +values:
132	      +count = List.length(&2, WeightedJob, values)
133	      +half = Nat.div(count, 2n)
134	      plan_step(partition(values,
135	        (Group{half, 0.0, Nil{}}, Group{Nat.sub(count, half), 0.0, Nil{}})),
136	        left => right => Fork{plan(p, left), plan(p, right)})
137	
138	def prepare(+jobs: List<&2, E.AlignJob>) -> Plan:
139	  plan(List.length(&2, E.AlignJob, jobs),
140	    List.sort(~WeightedJob, ~heavier, indexed(jobs, 0n)))
141	
142	def single(index: Nat, job: E.AlignJob) -> IndexedResult:
143	  E.AlignJob{query, target, config} = job
144	  IndexedResult{index, M.edlib_align(query, target, config)}
145	
146	def serial(jobs: List<&2, WeightedJob>) -> List<&2, IndexedResult>:
147	  match jobs:
148	    case Nil{}:
149	      Nil{}
150	    case WeightedJob{index, _, job} <> tail:
151	      single(index, job) <> serial(tail)
152	
153	# Only ready, independent subplans enter parallel calls. A leaf uses the same
154	# single-pair implementation; column carries and successive columns stay serial.
155	def execute(plan: Plan) -> List<&2, IndexedResult>:
156	  match plan:
157	    case Empty{}:
158	      Nil{}
159	    case Leaf{index, job}:
160	      [single(index, job)]
161	    case Serial{jobs}:
162	      serial(jobs)
163	    case Fork{left, right}:
164	      a b = execute(left) execute(right)
165	      List.append(&2, IndexedResult, a, b)
166	
167	def earlier(first: IndexedResult, second: IndexedResult) -> Bool:
168	  IndexedResult{a, _} = first
169	  IndexedResult{b, _} = second
170	  Nat.is_le(a, b)
171	
172	def results(values: List<&2, IndexedResult>) -> List<&2, E.AlignResult>:
173	  match values:
174	    case Nil{}:
175	      Nil{}
176	    case IndexedResult{_, result} <> tail:
177	      result <> results(tail)
178	
179	def align_batch(jobs: List<&2, E.AlignJob>) -> List<&2, E.AlignResult>:
180	  match jobs:
181	    case Nil{}:
182	      Nil{}
183	    case E.AlignJob{query, target, config} <> Nil{}:
184	      [M.edlib_align(query, target, config)]
185	    case values:
186	      results(List.sort(~IndexedResult, ~earlier, execute(prepare(values))))