src/Reconstruction.bend

bend-edit-distance@0.2.2.0 · 253 lines · 26 definitions

1import Base
2import bend-edit-distance@0.2.2.0/src/Sequence.bend as S
3import bend-edit-distance@0.2.2.0/EdlibTypes.bend as E
4import bend-edit-distance@0.2.2.0/src/DistanceDP.bend as DP
5
6# Global reconstruction of one already selected target interval. Small DP
7# stores a score and predecessor operation per cell, never path prefixes.
8# Equal scores prefer up (Insert), left (Delete), then diagonal, as upstream
9# traceback does. Hirschberg only promises an optimal path, not upstream ties.
10
11def diagonal_operation(matches: Bool) -> E.Operation:
12  Bool.pick(E.Operation, matches, E.Match{}, E.Mismatch{})
13
14type TraceCell is Data:
15  TraceCell{score: U32, operation: E.Operation}
16
17def cell_score(cell: TraceCell) -> U32:
18  match cell:
19    case TraceCell{score, _}:
20      score
21
22def choose_operation(matches: Bool, up_is_best: Bool, left_is_best: Bool) -> E.Operation:
23  match up_is_best left_is_best:
24    case True{} _:
25      E.Insert{}
26    case False{} True{}:
27      E.Delete{}
28    case False{} False{}:
29      diagonal_operation(matches)
30
31def dp_cell(above: U32, left: U32, diagonal: U32, +matches: Bool) -> TraceCell:
32  +up = U32.add(above, 1)
33  +across = U32.add(left, 1)
34  +diagonal_cost = U32.add(diagonal, Bool.to_u32(Bool.not(matches)))
35  TraceCell{U32.min(U32.min(up, across), diagonal_cost),
36    choose_operation(matches,
37      U32.is_le(up, U32.min(across, diagonal_cost)), U32.is_le(across, diagonal_cost))}
38
39def dp_row_cells(columns: String, previous: List<&2, TraceCell>, diagonal: U32,
40                 left: U32, +symbol: Char, +pairs: List<&2, E.EqualityPair>) ->
41                 List<&2, TraceCell>:
42  match columns previous:
43    case SNil{} _:
44      Nil{}
45    case SCon{_, _} Nil{}:
46      Nil{}
47    case SCon{h, t} TraceCell{+above, _} <> rest:
48      +next = dp_cell(above, left, diagonal, E.equal(pairs, symbol, h))
49      next <> dp_row_cells(t, rest, above, cell_score(next), symbol, pairs)
50
51def dp_row(previous: List<&2, TraceCell>, columns: String, symbol: Char,
52           +index: U32, pairs: List<&2, E.EqualityPair>) ->
53           List<&2, TraceCell>:
54  match previous:
55    case Nil{}:
56      [TraceCell{index, E.Insert{}}]
57    case TraceCell{diagonal, _} <> rest:
58      TraceCell{index, E.Insert{}} <>
59        dp_row_cells(columns, rest, diagonal, index, symbol, pairs)
60
61def dp_initial(target: String, +index: U32) -> List<&2, TraceCell>:
62  match target:
63    case SNil{}:
64      [TraceCell{index, E.Delete{}}]
65    case SCon{_, tail}:
66      TraceCell{index, E.Delete{}} <> dp_initial(tail, U32.add(index, 1))
67
68def dp_rows(query: String, +target: String, previous: List<&2, TraceCell>,
69            +index: U32, +pairs: List<&2, E.EqualityPair>,
70            directions: List<&2, List<&2, TraceCell>>) -> List<&2, List<&2, TraceCell>>:
71  match query:
72    case SNil{}:
73      directions
74    case SCon{h, t}:
75      +row = dp_row(previous, target, h, index, pairs)
76      dp_rows(t, target, row, U32.add(index, 1), pairs,
77        List.reverse(&2, TraceCell, row) <> directions)
78
79# Rows and their operations are reversed so left moves consume the current
80# row directly. Entering the row above skips only its unused right suffix.
81# Including boundary rows/columns handles all-insertion/deletion paths.
82def traceback(fuel: Nat, current: List<&2, TraceCell>,
83               previous: List<&2, List<&2, TraceCell>>, +width: Nat,
84               +column: Nat, operations: List<&2, E.Operation>) -> List<&2, E.Operation>:
85  match fuel current previous:
86    case 0n _ _:
87      operations
88    case 1n+p _ Nil{}:
89      List.append(&2, E.Operation, List.replicate(E.Operation, column, E.Delete{}), operations)
90    case 1n+p Nil{} _:
91      operations
92    case 1n+p TraceCell{_, E.Insert{}} <> _ row <> rest:
93      traceback(p, List.drop(&2, TraceCell, row, Nat.sub(width, column)),
94        rest, width, column, E.Insert{} <> operations)
95    case 1n+p TraceCell{_, E.Delete{}} <> tail +rows:
96      traceback(p, tail, rows, width, Nat.sub(column, 1n), E.Delete{} <> operations)
97    case 1n+p TraceCell{_, operation} <> _ row <> rest:
98      +next_column = Nat.sub(column, 1n)
99      traceback(p, List.drop(&2, TraceCell, row, Nat.sub(width, next_column)),
100        rest, width, next_column, operation <> operations)
101
102def dp_trace(matrix: List<&2, List<&2, TraceCell>>, query_length: Nat,
103              +target_length: Nat) -> List<&2, E.Operation>:
104  match matrix:
105    case Nil{}:
106      Nil{}
107    case row <> rest:
108      traceback(Nat.add(query_length, target_length), row, rest,
109        target_length, target_length, Nil{})
110
111def dp(+query: String, +target: String, pairs: List<&2, E.EqualityPair>) -> List<&2, E.Operation>:
112  qn +tn = S.length(query) S.length(target)
113  +boundary = dp_initial(target, 0)
114  dp_trace(dp_rows(query, target, boundary, 1, pairs,
115    [List.reverse(&2, TraceCell, boundary)]), qn, tn)
116
117# Bound the complete small matrix, including its boundary cells. A one-row
118# problem is also a base case, with linear storage in the column width.
119def small_matrix(query_length: Nat, +target_length: Nat) -> Bool:
120  Nat.is_le(1n+query_length, Nat.div(4096n, 1n+target_length))
121
122def score_row(query: String, +target: String, pairs: List<&2, E.EqualityPair>) -> List<U32>:
123  DP.final_row(query, target, DP.initial_row(S.length(target), 0), 1, pairs)
124
125type Split is Data:
126  Split{index: Nat, score: U32}
127
128def split_index(split: Split) -> Nat:
129  match split:
130    case Split{index, _}:
131      index
132
133def split_score(split: Split) -> U32:
134  match split:
135    case Split{_, score}:
136      score
137
138def split_update(index: Nat, score: U32, old_index: Nat, old_score: U32,
139                  better: Bool) -> Split:
140  match better:
141    case True{}:
142      Split{index, score}
143    case False{}:
144      Split{old_index, old_score}
145
146# Forward[j] + reverse[n-j] is the cost through the midpoint at column j.
147# Strict comparison keeps the first minimum; this is our deterministic split
148# policy and is deliberately separate from upstream Hirschberg compatibility.
149def split_column(forward: List<U32>, backward: List<U32>, +index: Nat,
150                  +best_index: Nat, +best_score: U32) -> Nat:
151  match forward backward:
152    case Nil{} _:
153      best_index
154    case _ Nil{}:
155      best_index
156    case a <> at b <> bt:
157      +score = U32.add(a, b)
158      +next = split_update(index, score, best_index, best_score,
159        U32.is_lt(score, best_score))
160      split_column(at, bt, 1n+index, split_index(next), split_score(next))
161
162def midpoint(left: String, right: String, +target: String,
163               +pairs: List<&2, E.EqualityPair>) -> Nat:
164  # Prepare recursive arguments before the fork. Passing String.reverse(...)
165  # inside its second call drains the frontier in the pinned native runtime.
166  reverse_query = String.reverse(right)
167  reverse_target = String.reverse(target)
168  forward backward =
169    score_row(left, target, pairs)
170    score_row(reverse_query, reverse_target, pairs)
171  split_column(forward, List.reverse(&1, U32, backward), 0n, 0n, 4294967295)
172
173# The fuel is the initial row length; every nonbase call halves that length.
174# Score passes run in balanced parallel calls and are consumed before recursion.
175# Subpaths are evaluated sequentially to bound live score storage even on many
176# threads; the result lists themselves necessarily use O(m+n) space.
177def divide(fuel: Nat, +query: String, +target: String, +query_length: Nat,
178             +target_length: Nat, +pairs: List<&2, E.EqualityPair>, small: Bool) ->
179             List<&2, E.Operation>:
180  match fuel query target small:
181    case _ SNil{} _ _:
182      List.replicate(E.Operation, target_length, E.Delete{})
183    case _ _ SNil{} _:
184      List.replicate(E.Operation, query_length, E.Insert{})
185    case 0n _ _ _:
186      dp(query, target, pairs)
187    case 1n+p SCon{_, SNil{}} _ _:
188      dp(query, target, pairs)
189    case 1n+p _ _ True{}:
190      dp(query, target, pairs)
191    case Succ{+p} _ _ False{}:
192      +half = Nat.div(query_length, 2n)
193      +left = String.take(query, half)
194      +right = String.drop(query, half)
195      +cut = midpoint(left, right, target, pairs)
196      left_path = divide(p, left, String.take(target, cut), half, cut, pairs,
197        small_matrix(half, cut))
198      +right_length = Nat.sub(query_length, half)
199      +right_width = Nat.sub(target_length, cut)
200      right_path = divide(p, right, String.drop(target, cut), right_length, right_width, pairs,
201        small_matrix(right_length, right_width))
202      List.append(&2, E.Operation, left_path, right_path)
203
204def transpose(operations: List<&2, E.Operation>) -> List<&2, E.Operation>:
205  match operations:
206    case Nil{}:
207      Nil{}
208    case E.Insert{} <> tail:
209      E.Delete{} <> transpose(tail)
210    case E.Delete{} <> tail:
211      E.Insert{} <> transpose(tail)
212    case operation <> tail:
213      operation <> transpose(tail)
214
215def reconstruct_order(query: String, target: String, +query_length: Nat,
216                        +target_length: Nat, pairs: List<&2, E.EqualityPair>,
217                        small: Bool, query_is_taller: Bool) -> List<&2, E.Operation>:
218  match small query_is_taller:
219    case True{} _:
220      dp(query, target, pairs)
221    case False{} True{}:
222      divide(query_length, query, target, query_length, target_length, pairs, False{})
223    case False{} False{}:
224      transpose(divide(target_length, target, query, target_length, query_length, pairs, False{}))
225
226def reconstruct(+query: String, +target: String, pairs: List<&2, E.EqualityPair>) ->
227                  List<&2, E.Operation>:
228  +qn +tn = S.length(query) S.length(target)
229  reconstruct_order(query, target, qn, tn, pairs, small_matrix(qn, tn), Nat.is_ge(qn, tn))
230
231# Edlib codes: Match=0 consumes both, Insert=1 consumes query only,
232# Delete=2 consumes target only, Mismatch=3 consumes both. Validation requires
233# exact consumption, the configured (symmetric, nontransitive) predicate on
234# every diagonal operation, and an edit cost equal to the reported distance.
235def validate(operations: List<&2, E.Operation>, query: String, target: String,
236               +pairs: List<&2, E.EqualityPair>, +expected: Nat, cost: Nat) -> Bool:
237  match operations query target:
238    case Nil{} SNil{} SNil{}:
239      Nat.is_eq(cost, expected)
240    case E.Match{} <> tail SCon{q, qt} SCon{t, tt}:
241      Bool.and(E.equal(pairs, q, t), validate(tail, qt, tt, pairs, expected, cost))
242    case E.Mismatch{} <> tail SCon{q, qt} SCon{t, tt}:
243      Bool.and(Bool.not(E.equal(pairs, q, t)), validate(tail, qt, tt, pairs, expected, 1n+cost))
244    case E.Insert{} <> tail SCon{_, qt} _:
245      validate(tail, qt, target, pairs, expected, 1n+cost)
246    case E.Delete{} <> tail _ SCon{_, tt}:
247      validate(tail, query, tt, pairs, expected, 1n+cost)
248    case _ _ _:
249      False{}
250
251def valid_path(query: String, target: String, operations: List<&2, E.Operation>,
252                 pairs: List<&2, E.EqualityPair>, distance: Nat) -> Bool:
253  validate(operations, query, target, pairs, distance, 0n)