1	import Base
2	import ./Cli.bend as CLI
3	import ./EdlibTypes.bend as E
4	import ./Alignment.bend as A
5	import ./AlignmentCli.bend as AC
6	
7	# The NW/distance path from vendor/edlib/edlib/src/edlib.cpp. Bend's native
8	# integer is U32, so a Myers word covers 32 query characters instead of 64.
9	# The scalar helper is NW, unit costs, unlimited distance (k = -1). The
10	# configured API extends it with the typed modes/tasks in EdlibTypes.bend.
11	
12	type Peq is Data:
13	  Mask{bits: U32}
14	  Branch{low: Peq, high: Peq}
15	
16	type Block is Data:
17	  Block{positive: U32, negative: U32, score: U32}
18	
19	# Horizontal differences are -1, 0, or +1. Keep both signs explicit rather
20	# than relying on signed shifts or two's-complement arithmetic.
21	type Carry is Data:
22	  Carry{positive: U32, negative: U32}
23	
24	def split_query.prepend(h: Char, parts: String & String) -> String & String:
25	  match parts:
26	    case (prefix, suffix):
27	      (SCon{h, prefix}, suffix)
28	
29	def split_query(n: Nat, query: String) -> String & String:
30	  match n query:
31	    case 0n _:
32	      ("", query)
33	    case 1n+p SNil{}:
34	      ("", "")
35	    case 1n+p SCon{h, t}:
36	      split_query.prepend(h, split_query(p, t))
37	
38	# Bit i matches query[i]. Unused high bits are wildcards, as in buildPeq.
39	def build_peq_word(query: String, +symbol: Char, +bit: U32) -> U32:
40	  match query:
41	    case SNil{}:
42	      U32.not(U32.sub(bit, 1))
43	    case SCon{h, t}:
44	      (U32.mul(Bool.to_u32(Char.is_eq(h, symbol)), bit) .|.
45	        build_peq_word(t, symbol, U32.shl(bit)) : U32)
46	
47	# A byte-indexed immutable table replaces the C++ transformed alphabet.
48	# Low bits select branches, so lookup takes eight steps for every byte.
49	# Both halves of the table have equal, independent work and run in parallel.
50	def build_peq_table(+depth: Nat, +query: String, +symbol: U32, +step: U32) -> Peq:
51	  match depth:
52	    case 0n:
53	      Mask{build_peq_word(query, Chr{symbol}, 1)}
54	    case 1n+p:
55	      low high = build_peq_table(p, query, symbol, U32.shl(step)) build_peq_table(p, query, U32.add(symbol, step), U32.shl(step))
56	      Branch{low, high}
57	
58	def peq_get(table: Peq, +code: U32, bit: Bool) -> U32:
59	  match table bit:
60	    case Mask{bits} _:
61	      bits
62	    case Branch{low, high} True{}:
63	      +next = U32.shr(code)
64	      peq_get(high, next, U32.is_ne(U32.and(next, 1), 0))
65	    case Branch{low, high} False{}:
66	      +next = U32.shr(code)
67	      peq_get(low, next, U32.is_ne(U32.and(next, 1), 0))
68	
69	def build_peq.step(parts: String & String, tail: String -> List<&2, Peq>) -> List<&2, Peq>:
70	  match parts:
71	    case (word, rest):
72	      build_peq_table(8n, word, 0, 1) <> tail(rest)
73	
74	def build_peq(fuel: Nat, query: String) -> List<&2, Peq>:
75	  match fuel:
76	    case 0n:
77	      Nil{}
78	    case 1n+p:
79	      build_peq.step(split_query(32n, query), rest => build_peq(p, rest))
80	
81	def ceil_div(+x: Nat, +y: Nat) -> Nat:
82	  Nat.div(Nat.add(x, Nat.sub(y, 1n)), y)
83	
84	def init_blocks(n: Nat, +bottom: U32) -> List<Block>:
85	  match n:
86	    case 0n:
87	      Nil{}
88	    case 1n+p:
89	      Block{4294967295, 0, bottom} <> init_blocks(p, U32.add(bottom, 32))
90	
91	# calculateBlock / Advance_Block from Myers. Addition intentionally wraps
92	# at 32 bits. Xv uses the original Eq, before the negative carry is added.
93	def calculate_block(block: Block, +eq: U32, hin: Carry) -> Block & Carry:
94	  Block{+pv, +mv, score} = block
95	  Carry{hin_positive, +hin_negative} = hin
96	  +xv = U32.or(eq, mv)
97	  +eq_in = U32.or(eq, hin_negative)
98	  +xh = (U32.xor(U32.add(U32.and(eq_in, pv), pv), pv) .|. eq_in : U32)
99	  +ph = U32.or(mv, U32.not(U32.or(xh, pv)))
100	  +mh = U32.and(pv, xh)
101	  +hout_positive = U32.shrn(ph, 31n)
102	  +hout_negative = U32.shrn(mh, 31n)
103	  +ph_shifted = U32.or(U32.shl(ph), hin_positive)
104	  mh_shifted = U32.or(U32.shl(mh), hin_negative)
105	  next_pv = U32.or(mh_shifted, U32.not(U32.or(xv, ph_shifted)))
106	  next_mv = U32.and(ph_shifted, xv)
107	  (Block{next_pv, next_mv, U32.sub(U32.add(score, hout_positive), hout_negative)},
108	    Carry{hout_positive, hout_negative})
109	
110	# A column and consecutive target columns depend on their predecessor;
111	# parallelizing those calls would violate the Myers carry dependency.
112	def advance_column.step(result: Block & Carry,
113	                        tail: Carry -> List<Block>) -> List<Block>:
114	  match result:
115	    case (next, hout):
116	      next <> tail(hout)
117	
118	def advance_column(blocks: List<Block>, profiles: List<&2, Peq>, +symbol: U32,
119	                   hin: Carry) -> List<Block>:
120	  match blocks profiles:
121	    case Nil{} _:
122	      Nil{}
123	    case _ Nil{}:
124	      Nil{}
125	    case block <> rest profile <> tail:
126	      advance_column.step(
127	        calculate_block(block, peq_get(profile, symbol, U32.is_ne(U32.and(symbol, 1), 0)), hin),
128	        hout => advance_column(rest, tail, symbol, hout))
129	
130	# getBlockCellValues(block)[W], without constructing the other 31 cells.
131	# Walking up a row subtracts its positive vertical difference and adds its
132	# negative difference. W is always in 0..31, never a whole extra block.
133	def get_block_cell_value(padding: Nat, +pv: U32, +mv: U32, score: U32) -> U32:
134	  match padding:
135	    case 0n:
136	      score
137	    case 1n+p:
138	      get_block_cell_value(p, U32.shl(pv), U32.shl(mv),
139	        U32.add(U32.sub(score, U32.shrn(pv, 31n)), U32.shrn(mv, 31n)))
140	
141	def last_block_score(blocks: List<Block>, +padding: Nat) -> Nat:
142	  match blocks:
143	    case Nil{}:
144	      0n
145	    case Block{pv, mv, score} <> Nil{}:
146	      U32.to_nat(get_block_cell_value(padding, pv, mv, score))
147	    case _ <> rest:
148	      last_block_score(rest, padding)
149	
150	# Unlimited NW evaluates every block once per target character. It needs
151	# only the current column; Ukkonen threshold bands and traceback storage
152	# are unnecessary for the distance-only, k = -1 interface.
153	def myers_calc_edit_distance_nw(target: String, +profiles: List<&2, Peq>,
154	                               blocks: List<Block>, +padding: Nat) -> Nat:
155	  match target:
156	    case SNil{}:
157	      last_block_score(blocks, padding)
158	    case SCon{h, t}:
159	      myers_calc_edit_distance_nw(t, profiles,
160	        advance_column(blocks, profiles, Char.to_u32(h), Carry{1, 0}), padding)
161	
162	def align_query(+query: String, target: String, +length: Nat) -> Nat:
163	  +count = ceil_div(length, 32n)
164	  profiles = build_peq(count, query)
165	  blocks = init_blocks(count, 32)
166	  myers_calc_edit_distance_nw(target, profiles, blocks,
167	    Nat.sub(Nat.mul(count, 32n), length))
168	
169	def align_shorter(query: String, target: String, query_length: Nat,
170	                  target_length: Nat, query_is_shorter: Bool) -> Nat:
171	  match query_is_shorter:
172	    case True{}:
173	      align_query(query, target, query_length)
174	    case False{}:
175	      align_query(target, query, target_length)
176	
177	def distance(query: String, target: String) -> Nat:
178	  match query target:
179	    case SNil{} _:
180	      String.length(target)
181	    case SCon{h, t} SNil{}:
182	      1n+String.length(t)
183	    case SCon{qh, qt} SCon{th, tt}:
184	      +q = {SCon{qh, qt} : String}
185	      +t = {SCon{th, tt} : String}
186	      +qn +tn = String.length(q) String.length(t)
187	      align_shorter(q, t, qn, tn, Nat.is_le(qn, tn))
188	
189	def edlib_align(query: String, target: String, config: E.AlignConfig) -> E.AlignResult:
190	  match query target config:
191	    case SNil{} _ _:
192	      A.align("", target, config)
193	    case SCon{h, t} SNil{} _:
194	      A.align(SCon{h, t}, "", config)
195	    case +q +t E.AlignConfig{E.NW{}, E.Distance{}, +threshold, Nil{}, _}:
196	      +score = distance(q, t)
197	      E.distance_result(score, E.end_coordinate(String.length(t)),
198	        E.within_threshold(threshold, score))
199	    case _ _ _:
200	      A.align(query, target, config)
201	
202	def response(query: String, target: String) -> String:
203	  Nat.show(distance(query, target))
204	
205	def run(args: List<String>) -> IO(Unit):
206	  match args:
207	    case _ <> query <> target <> Nil{}:
208	      CLI.run(~distance, ["bend-edit-distance", query, target])
209	    case _:
210	      AC.run(~edlib_align, args)
211	
212	def main() -> IO(Unit):
213	  do IO<Unit>:
214	    args : List<String> <- IO.args()
215	    run(args)