ICFP 2007 Contest: https://web.archive.org/web/20090301164728/https://save-endo.cs.uu.nl/
1use std::collections::VecDeque;
2use std::ops::Add;
3use std::ops::AddAssign;
4use std::ops::Deref;
5use std::ops::Index;
6use std::rc::Rc;
7
8#[derive(Copy, Clone, PartialEq, Debug)]
9pub enum Base {
10 I,
11 C,
12 F,
13 P,
14}
15
16#[derive(Clone, Debug)]
17pub enum Dna {
18 Empty,
19 Leaf(Base),
20 TwoNode(TwoNode),
21 ThreeNode(ThreeNode),
22}
23
24#[derive(Clone, Debug)]
25// Wrap Rc in a single-value enum so I can implement traits for it.
26pub enum DnaRef {
27 DnaRef(Rc<Dna>),
28}
29
30#[derive(Clone, Debug)]
31pub struct TwoNode {
32 len: usize,
33 depth: usize,
34 children: (DnaRef, DnaRef),
35}
36
37#[derive(Clone, Debug)]
38pub struct ThreeNode {
39 len: usize,
40 depth: usize,
41 children: (DnaRef, DnaRef, DnaRef),
42}
43
44impl Base {
45 pub fn from_char(c: char) -> Option<Base> {
46 match c {
47 'I' => Some(Base::I),
48 'C' => Some(Base::C),
49 'F' => Some(Base::F),
50 'P' => Some(Base::P),
51 _ => None,
52 }
53 }
54
55 pub fn to_char(&self) -> char {
56 match self {
57 Base::I => 'I',
58 Base::C => 'C',
59 Base::F => 'F',
60 Base::P => 'P',
61 }
62 }
63}
64
65impl DnaRef {
66 pub fn new() -> Self {
67 Self::DnaRef(Rc::new(Dna::Empty))
68 }
69
70 pub fn from_base(b: Base) -> Self {
71 Self::DnaRef(Rc::new(Dna::Leaf(b)))
72 }
73
74 fn from_two_children(a: DnaRef, b: DnaRef) -> Self {
75 debug_assert_eq!(a.depth(), b.depth());
76 Self::DnaRef(Rc::new(Dna::TwoNode(TwoNode {
77 len: a.len() + b.len(),
78 depth: a.depth() + 1,
79 children: (a, b),
80 })))
81 }
82
83 fn from_three_children(a: DnaRef, b: DnaRef, c: DnaRef) -> Self {
84 debug_assert_eq!(a.depth(), b.depth());
85 debug_assert_eq!(a.depth(), c.depth());
86 Self::DnaRef(Rc::new(Dna::ThreeNode(ThreeNode {
87 len: a.len() + b.len() + c.len(),
88 depth: a.depth() + 1,
89 children: (a, b, c),
90 })))
91 }
92
93 pub fn from_string(s: &str) -> Self {
94 s.chars()
95 .filter_map(|c| Base::from_char(c))
96 .map(|b| DnaRef::from_base(b))
97 .collect::<DnaRef>()
98 }
99
100 pub fn len(&self) -> usize {
101 match &**self {
102 Dna::Empty => 0,
103 Dna::Leaf(_) => 1,
104 Dna::TwoNode(node) => node.len,
105 Dna::ThreeNode(node) => node.len,
106 }
107 }
108
109 fn depth(&self) -> usize {
110 match &**self {
111 Dna::Empty => 0,
112 Dna::Leaf(_) => 1,
113 Dna::TwoNode(node) => node.depth,
114 Dna::ThreeNode(node) => node.depth,
115 }
116 }
117
118 pub fn iter<'a>(&'a self) -> DnaIterator<'a> {
119 let mut stack = Vec::new();
120 stack.push(self);
121 DnaIterator { stack }
122 }
123
124 pub fn split(&self, n: usize) -> (DnaRef, DnaRef) {
125 debug_assert!(n <= self.len());
126 match &**self {
127 Dna::Empty => (DnaRef::new(), DnaRef::new()),
128 Dna::Leaf(b) => {
129 if n == 0 {
130 (DnaRef::new(), DnaRef::from_base(*b))
131 } else {
132 (DnaRef::from_base(*b), DnaRef::new())
133 }
134 }
135 Dna::TwoNode(node) => {
136 let (a, b) = &node.children;
137 if n < a.len() {
138 let (x, y) = a.split(n);
139 (x, y + b)
140 } else if n == a.len() {
141 (a.clone(), b.clone())
142 } else {
143 let (x, y) = b.split(n - a.len());
144 (a + x, y)
145 }
146 }
147 Dna::ThreeNode(node) => {
148 let (a, b, c) = &node.children;
149 if n < a.len() {
150 let (x, y) = a.split(n);
151 (x, y + b + c)
152 } else if n == a.len() {
153 (a.clone(), b + c)
154 } else if n - a.len() < b.len() {
155 let (x, y) = b.split(n - a.len());
156 (a + x, y + c)
157 } else if n == a.len() + b.len() {
158 (a + b, c.clone())
159 } else {
160 let (x, y) = c.split(n - a.len() - b.len());
161 (a + b + x, y)
162 }
163 }
164 }
165 }
166
167 #[cfg(test)]
168 pub fn to_string(&self) -> String {
169 self.iter().map(|b| Base::to_char(&b)).collect::<String>()
170 }
171}
172
173impl Deref for DnaRef {
174 type Target = Dna;
175
176 fn deref(&self) -> &Self::Target {
177 match self {
178 DnaRef::DnaRef(r) => &*r,
179 }
180 }
181}
182
183struct DnaRefIter<I>
184where
185 I: Iterator<Item = DnaRef>,
186{
187 iter: I,
188 buf: VecDeque<DnaRef>,
189}
190
191impl<I> DnaRefIter<I>
192where
193 I: Iterator<Item = DnaRef>,
194{
195 fn new(iter: I) -> DnaRefIter<I> {
196 DnaRefIter {
197 iter: iter,
198 buf: VecDeque::with_capacity(5),
199 }
200 }
201}
202
203impl<I> Iterator for DnaRefIter<I>
204where
205 I: Iterator<Item = DnaRef>,
206{
207 type Item = DnaRef;
208 fn next(&mut self) -> Option<Self::Item> {
209 while self.buf.len() < 5 {
210 let Some(x) = self.iter.next() else {
211 break;
212 };
213
214 self.buf.push_back(x);
215 }
216
217 if self.buf.len() == 5 || self.buf.len() == 3 {
218 let a = self.buf.pop_front().unwrap();
219 let b = self.buf.pop_front().unwrap();
220 let c = self.buf.pop_front().unwrap();
221
222 return Some(DnaRef::from_three_children(a, b, c));
223 } else if self.buf.len() == 4 || self.buf.len() == 2 {
224 let a = self.buf.pop_front().unwrap();
225 let b = self.buf.pop_front().unwrap();
226
227 return Some(DnaRef::from_two_children(a, b));
228 } else {
229 return None;
230 }
231 }
232}
233
234impl FromIterator<DnaRef> for DnaRef {
235 fn from_iter<I: IntoIterator<Item = DnaRef>>(iter: I) -> DnaRef {
236 // We need some kind of buffer no matter what, even if we aggregate as we go.
237 // This implementation does the easy thing and just repeatedly collects into a Vec.
238 let mut cur: Vec<DnaRef> = iter.into_iter().collect::<Vec<_>>();
239 while cur.len() > 1 {
240 cur = DnaRefIter::new(cur.into_iter()).collect();
241 }
242
243 match cur.pop() {
244 Some(dna) => dna,
245 None => DnaRef::new(),
246 }
247 }
248}
249
250pub struct DnaIterator<'a> {
251 stack: Vec<&'a DnaRef>,
252}
253
254impl<'a> Iterator for DnaIterator<'a> {
255 type Item = Base;
256
257 fn next(&mut self) -> Option<Self::Item> {
258 while let Some(node) = self.stack.pop() {
259 match &**node {
260 Dna::Empty => return None,
261 Dna::Leaf(c) => return Some(*c),
262 Dna::TwoNode(node) => {
263 let (a, b) = &node.children;
264 self.stack.push(b);
265 self.stack.push(a);
266 }
267 Dna::ThreeNode(node) => {
268 let (a, b, c) = &node.children;
269 self.stack.push(c);
270 self.stack.push(b);
271 self.stack.push(a);
272 }
273 }
274 }
275
276 return None;
277 }
278}
279
280enum AddState {
281 One(DnaRef),
282 Two(DnaRef, DnaRef),
283}
284
285fn add_helper(lhs: &DnaRef, rhs: &DnaRef) -> AddState {
286 if lhs.depth() == rhs.depth() {
287 AddState::Two(lhs.clone(), rhs.clone())
288 } else if lhs.depth() < rhs.depth() {
289 match &**rhs {
290 Dna::Empty => unreachable!(),
291 Dna::Leaf(_) => AddState::One(rhs.clone()),
292 Dna::TwoNode(node) => {
293 let (a, b) = &node.children;
294 match add_helper(lhs, a) {
295 AddState::One(x) => AddState::One(DnaRef::from_two_children(x, b.clone())),
296 AddState::Two(x, y) => {
297 AddState::One(DnaRef::from_three_children(x, y, b.clone()))
298 }
299 }
300 }
301 Dna::ThreeNode(node) => {
302 let (a, b, c) = &node.children;
303 match add_helper(lhs, a) {
304 AddState::One(x) => {
305 AddState::One(DnaRef::from_three_children(x, b.clone(), c.clone()))
306 }
307 AddState::Two(x, y) => AddState::Two(
308 DnaRef::from_two_children(x, y),
309 DnaRef::from_two_children(b.clone(), c.clone()),
310 ),
311 }
312 }
313 }
314 } else {
315 match &**lhs {
316 Dna::Empty => unreachable!(),
317 Dna::Leaf(_) => AddState::One(lhs.clone()),
318 Dna::TwoNode(node) => {
319 let (a, b) = &node.children;
320 match add_helper(b, rhs) {
321 AddState::One(x) => AddState::One(DnaRef::from_two_children(a.clone(), x)),
322 AddState::Two(x, y) => {
323 AddState::One(DnaRef::from_three_children(a.clone(), x, y))
324 }
325 }
326 }
327 Dna::ThreeNode(node) => {
328 let (a, b, c) = &node.children;
329 match add_helper(c, rhs) {
330 AddState::One(x) => {
331 AddState::One(DnaRef::from_three_children(a.clone(), b.clone(), x))
332 }
333 AddState::Two(x, y) => AddState::Two(
334 DnaRef::from_two_children(a.clone(), b.clone()),
335 DnaRef::from_two_children(x, y),
336 ),
337 }
338 }
339 }
340 }
341}
342
343impl Add<DnaRef> for DnaRef {
344 type Output = DnaRef;
345
346 fn add(self, other: DnaRef) -> DnaRef {
347 match add_helper(&self, &other) {
348 AddState::One(a) => a,
349 AddState::Two(a, b) => DnaRef::from_two_children(a, b),
350 }
351 }
352}
353
354impl Add<&DnaRef> for DnaRef {
355 type Output = DnaRef;
356
357 fn add(self, other: &DnaRef) -> DnaRef {
358 match add_helper(&self, other) {
359 AddState::One(a) => a,
360 AddState::Two(a, b) => DnaRef::from_two_children(a, b),
361 }
362 }
363}
364
365impl Add<DnaRef> for &DnaRef {
366 type Output = DnaRef;
367
368 fn add(self, other: DnaRef) -> DnaRef {
369 match add_helper(self, &other) {
370 AddState::One(a) => a,
371 AddState::Two(a, b) => DnaRef::from_two_children(a, b),
372 }
373 }
374}
375
376impl Add<&DnaRef> for &DnaRef {
377 type Output = DnaRef;
378
379 fn add(self, other: &DnaRef) -> DnaRef {
380 match add_helper(self, other) {
381 AddState::One(a) => a,
382 AddState::Two(a, b) => DnaRef::from_two_children(a, b),
383 }
384 }
385}
386
387impl AddAssign<DnaRef> for DnaRef {
388 fn add_assign(&mut self, other: DnaRef) {
389 match add_helper(self, &other) {
390 AddState::One(a) => {
391 *self = a;
392 }
393 AddState::Two(a, b) => {
394 *self = Self::from_two_children(a, b);
395 }
396 }
397 }
398}
399
400impl AddAssign<&DnaRef> for DnaRef {
401 fn add_assign(&mut self, other: &DnaRef) {
402 match add_helper(self, other) {
403 AddState::One(a) => {
404 *self = a;
405 }
406 AddState::Two(a, b) => {
407 *self = Self::from_two_children(a, b);
408 }
409 }
410 }
411}
412
413impl Index<usize> for DnaRef {
414 type Output = Base;
415
416 fn index(&self, n: usize) -> &Self::Output {
417 debug_assert!(n < self.len());
418
419 match &**self {
420 Dna::Empty => unreachable!(),
421 Dna::Leaf(c) => c,
422 Dna::TwoNode(node) => {
423 let (a, b) = &node.children;
424 if n < a.len() {
425 a.index(n)
426 } else {
427 b.index(n - a.len())
428 }
429 }
430 Dna::ThreeNode(node) => {
431 let (a, b, c) = &node.children;
432 if n < a.len() {
433 a.index(n)
434 } else if n < a.len() + b.len() {
435 b.index(n - a.len())
436 } else {
437 c.index(n - a.len() - b.len())
438 }
439 }
440 }
441 }
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447
448 #[test]
449 fn test_empty() {
450 let dna: DnaRef = DnaRef::new();
451 assert_eq!(dna.len(), 0);
452 }
453
454 #[test]
455 fn test_to_from_string() {
456 let dna: DnaRef = DnaRef::from_string("ICFP");
457 assert_eq!(dna.to_string(), "ICFP");
458 }
459
460 #[test]
461 fn test_index() {
462 let s = "ICFP";
463 let dna: DnaRef = DnaRef::from_string(s);
464 assert_eq!(dna[0], Base::I);
465 assert_eq!(dna[1], Base::C);
466 assert_eq!(dna[2], Base::F);
467 assert_eq!(dna[3], Base::P);
468 }
469
470 #[test]
471 fn test_add() {
472 let lhs = DnaRef::from_string("IIIC");
473 let rhs = DnaRef::from_string("PFFF");
474 assert_eq!((lhs + rhs).to_string(), "IIICPFFF");
475 }
476
477 #[test]
478 fn test_split() {
479 let dna = DnaRef::from_string("IIICPFFF");
480 let (lhs, rhs) = dna.split(4);
481 assert_eq!(lhs.to_string(), "IIIC");
482 assert_eq!(rhs.to_string(), "PFFF");
483 }
484}