1	import Base
2	import ./EdlibTypes.bend as E
3	
4	# Configured alignment uses the scalar recurrence. Only two score rows are
5	# retained; PATH also carries persistent, shared operation prefixes. The
6	# ordinary NW/DISTANCE API in main.bend keeps its Myers fast path.
7	type Cell is Data:
8	  Cell{score: U32, start: Nat, operations: List<&2, E.Operation>}
9	
10	type Location is Data:
11	  Location{end: E.Coordinate, start: Nat}
12	
13	type Best is Data:
14	  Best{score: U32, locations: List<&2, Location>}
15	
16	# Additional pairs are symmetric, with no transitive closure.
17	def pair_matches(pair: E.EqualityPair, +a: Char, +b: Char) -> Bool:
18	  match pair:
19	    case E.EqualityPair{+first, +second}:
20	      Bool.or(Bool.and(Char.is_eq(a, first), Char.is_eq(b, second)),
21	        Bool.and(Char.is_eq(a, second), Char.is_eq(b, first)))
22	
23	def equal(pairs: List<&2, E.EqualityPair>, +a: Char, +b: Char) -> Bool:
24	  match pairs:
25	    case Nil{}:
26	      Char.is_eq(a, b)
27	    case pair <> tail:
28	      Bool.or(pair_matches(pair, a, b), equal(tail, a, b))
29	
30	def trace_step(operations: List<&2, E.Operation>, operation: E.Operation,
31	               tracing: Bool) -> List<&2, E.Operation>:
32	  match tracing:
33	    case False{}:
34	      Nil{}
35	    case True{}:
36	      operation <> operations
37	
38	def extend(cell: Cell, cost: U32, operation: E.Operation, tracing: Bool) -> Cell:
39	  match cell:
40	    case Cell{score, start, operations}:
41	      Cell{U32.add(score, cost), start, trace_step(operations, operation, tracing)}
42	
43	def prefer_cells(first: Cell, second: Cell, prefer: Bool) -> Cell:
44	  match prefer:
45	    case True{}:
46	      first
47	    case False{}:
48	      second
49	
50	# Earliest optimal start reproduces Edlib's last reverse-SHW location.
51	# Equal starts retain above, left, diagonal priority for NW traceback.
52	def prefer(+first: Cell, +second: Cell) -> Cell:
53	  Cell{+a, sa, _} = first
54	  Cell{+b, sb, _} = second
55	  prefer_cells(first, second,
56	    Bool.or(U32.is_lt(a, b), Bool.and(U32.is_eq(a, b), Nat.is_le(sa, sb))))
57	
58	def step(above: Cell, left: Cell, diagonal: Cell, +tracing: Bool,
59	         +matches: Bool) -> Cell:
60	  prefer(prefer(extend(above, 1, E.Insert{}, tracing),
61	    extend(left, 1, E.Delete{}, tracing)),
62	    extend(diagonal, Bool.to_u32(Bool.not(matches)),
63	      Bool.pick(E.Operation, matches, E.Match{}, E.Mismatch{}), tracing))
64	
65	def row_cells(columns: String, previous: List<&2, Cell>, diagonal: Cell,
66	              left: Cell, +symbol: Char, +pairs: List<&2, E.EqualityPair>,
67	              +tracing: Bool) -> List<&2, Cell>:
68	  match columns previous:
69	    case SNil{} _:
70	      Nil{}
71	    case SCon{h, t} Nil{}:
72	      Nil{}
73	    case SCon{h, t} +above <> rest:
74	      +next = step(above, left, diagonal, tracing, equal(pairs, symbol, h))
75	      next <> row_cells(t, rest, above, next, symbol, pairs, tracing)
76	
77	def next_row(previous: List<&2, Cell>, columns: String, symbol: Char,
78	             pairs: List<&2, E.EqualityPair>, tracing: Bool,
79	             +first: Cell) -> List<&2, Cell>:
80	  match previous:
81	    case Nil{}:
82	      [first]
83	    case diagonal <> rest:
84	      first <> row_cells(columns, rest, diagonal, first, symbol, pairs, tracing)
85	
86	def initial_cell(+mode: E.Mode, +index: U32,
87	                 operations: List<&2, E.Operation>) -> Cell:
88	  match mode:
89	    case E.HW{}:
90	      Cell{0, U32.to_nat(index), Nil{}}
91	    case _:
92	      Cell{index, 0n, operations}
93	
94	def initial_row(columns: String, +mode: E.Mode, +index: U32, +tracing: Bool,
95	                +operations: List<&2, E.Operation>) -> List<&2, Cell>:
96	  match columns:
97	    case SNil{}:
98	      [initial_cell(mode, index, operations)]
99	    case SCon{_, t}:
100	      initial_cell(mode, index, operations) <>
101	        initial_row(t, mode, U32.add(index, 1), tracing,
102	          trace_step(operations, E.Delete{}, tracing))
103	
104	def rows(query: String, +columns: String, previous: List<&2, Cell>,
105	         +pairs: List<&2, E.EqualityPair>, +index: U32, +tracing: Bool,
106	         +operations: List<&2, E.Operation>) -> List<&2, Cell>:
107	  match query:
108	    case SNil{}:
109	      previous
110	    case SCon{h, t}:
111	      +path = trace_step(operations, E.Insert{}, tracing)
112	      rows(t, columns, next_row(previous, columns, h, pairs, tracing,
113	        Cell{index, 0n, path}), pairs, U32.add(index, 1), tracing, path)
114	
115	def scores(query: String, +target: String, mode: E.Mode,
116	           pairs: List<&2, E.EqualityPair>, +tracing: Bool) -> List<&2, Cell>:
117	  rows(query, target, initial_row(target, mode, 0, tracing, Nil{}),
118	    pairs, 1, tracing, Nil{})
119	
120	def last_cell(cells: List<&2, Cell>) -> Cell:
121	  match cells:
122	    case Nil{}:
123	      Cell{0, 0n, Nil{}}
124	    case cell <> Nil{}:
125	      cell
126	    case _ <> tail:
127	      last_cell(tail)
128	
129	def update_best(best: Best, score: U32, location: Location, comparison: Cmp) -> Best:
130	  match best comparison:
131	    case Best{_, _} LT{}:
132	      Best{score, [location]}
133	    case Best{old, locations} EQ{}:
134	      Best{old, location <> locations}
135	    case _ GT{}:
136	      best
137	
138	def collect(cells: List<&2, Cell>, +index: Nat, +best: Best) -> Best:
139	  match cells:
140	    case Nil{}:
141	      best
142	    case Cell{+score, start, _} <> tail:
143	      Best{old, _} = best
144	      collect(tail, 1n+index, update_best(best, score,
145	        Location{E.Position{index}, start}, U32.cmp(score, old)))
146	
147	def ordered(best: Best) -> Best:
148	  match best:
149	    case Best{score, locations}:
150	      Best{score, List.reverse(&2, Location, locations)}
151	
152	def nw_endpoint(cell: Cell, target_length: Nat) -> Best:
153	  match cell:
154	    case Cell{score, start, _}:
155	      Best{score, [Location{E.end_coordinate(target_length), start}]}
156	
157	def semiglobal_endpoints(cells: List<&2, Cell>, include_before: Bool) -> Best:
158	  match cells include_before:
159	    case Nil{} _:
160	      Best{0, Nil{}}
161	    case Cell{score, start, _} <> tail True{}:
162	      ordered(collect(tail, 0n, Best{score, [Location{E.BeforeStart{}, start}]}))
163	    case _ <> tail False{}:
164	      ordered(collect(tail, 0n, Best{4294967295, Nil{}}))
165	
166	def endpoints(cells: List<&2, Cell>, mode: E.Mode, target_length: Nat,
167	              query_length: Nat) -> Best:
168	  match cells mode:
169	    case _ E.NW{}:
170	      nw_endpoint(last_cell(cells), target_length)
171	    case _ _:
172	      # Pinned Edlib reports column zero through 64-bit wildcard padding.
173	      # With a full final word (W=0), it omits that end location.
174	      semiglobal_endpoints(cells, Nat.is_ne(Nat.mod(query_length, 64n), 0n))
175	
176	def ends(locations: List<&2, Location>) -> List<&2, E.Coordinate>:
177	  match locations:
178	    case Nil{}:
179	      Nil{}
180	    case Location{end, _} <> tail:
181	      end <> ends(tail)
182	
183	def starts(locations: List<&2, Location>) -> List<&2, E.Coordinate>:
184	  match locations:
185	    case Nil{}:
186	      Nil{}
187	    case Location{_, start} <> tail:
188	      E.Position{start} <> starts(tail)
189	
190	def span(target: String, location: Location) -> String:
191	  match location:
192	    case Location{E.BeforeStart{}, _}:
193	      ""
194	    case Location{E.Position{index}, start}:
195	      String.drop(String.take(target, 1n+index), start)
196	
197	def path_result(score: U32, +locations: List<&2, Location>,
198	                +operations: List<&2, E.Operation>, format: E.CigarFormat) -> E.AlignResult:
199	  E.AlignResult{E.Ok{}, E.Found{U32.to_nat(score)}, ends(locations),
200	    Some{starts(locations)}, Some{operations}, Some{E.alignment_to_cigar(operations, format)}}
201	
202	def path_cell(score: U32, locations: List<&2, Location>, format: E.CigarFormat,
203	              cell: Cell) -> E.AlignResult:
204	  match cell:
205	    case Cell{_, _, operations}:
206	      path_result(score, locations, List.reverse(&2, E.Operation, operations), format)
207	
208	def result(query: String, target: String, task: E.Task,
209	           pairs: List<&2, E.EqualityPair>, format: E.CigarFormat,
210	           best: Best, accepted: Bool) -> E.AlignResult:
211	  match task best accepted:
212	    case _ _ False{}:
213	      E.not_found()
214	    case E.Distance{} Best{score, locations} True{}:
215	      E.AlignResult{E.Ok{}, E.Found{U32.to_nat(score)}, ends(locations), None{}, None{}, None{}}
216	    case E.Locations{} Best{score, +locations} True{}:
217	      E.AlignResult{E.Ok{}, E.Found{U32.to_nat(score)}, ends(locations), Some{starts(locations)}, None{}, None{}}
218	    case E.Path{} Best{score, +locations} True{}:
219	      location = Maybe.default(&2, Location, List.head(&2, Location, locations),
220	        Location{E.BeforeStart{}, 0n})
221	      path_cell(score, locations, format,
222	        last_cell(scores(query, span(target, location), E.NW{}, pairs, True{})))
223	
224	def finish(query: String, target: String, task: E.Task,
225	           threshold: E.Threshold, pairs: List<&2, E.EqualityPair>,
226	           format: E.CigarFormat, +best: Best) -> E.AlignResult:
227	  Best{score, _} = best
228	  result(query, target, task, pairs, format, best,
229	    E.within_threshold(threshold, U32.to_nat(score)))
230	
231	# Pinned Edlib returns before applying k or calculating starts/path on empty
232	# inputs, including PATH. Preserve that observable oracle behavior.
233	def align(query: String, target: String, config: E.AlignConfig) -> E.AlignResult:
234	  match query target config:
235	    case SNil{} +t E.AlignConfig{E.NW{}, _, _, _, _}:
236	      +length = String.length(t)
237	      E.distance_result(length, E.end_coordinate(length), True{})
238	    case SNil{} _ _:
239	      E.distance_result(0n, E.BeforeStart{}, True{})
240	    case SCon{_, qt} SNil{} _:
241	      E.distance_result(1n+String.length(qt), E.BeforeStart{}, True{})
242	    case +q +t E.AlignConfig{+mode, task, threshold, +pairs, format}:
243	      qn tn = String.length(q) String.length(t)
244	      finish(q, t, task, threshold, pairs, format,
245	        endpoints(scores(q, t, mode, pairs, False{}), mode, tn, qn))