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