cc6502/
assemble.rs

1/*
2    cc6502 - a subset of C compiler for the 6502 processor
3    Copyright (C) 2023-2024 Bruno STEUX
4
5    This program is free software: you can redistribute it and/or modify
6    it under the terms of the GNU General Public License as published by
7    the Free Software Foundation, either version 3 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU General Public License for more details.
14
15    You should have received a copy of the GNU General Public License
16    along with this program.  If not, see <https://www.gnu.org/licenses/>.
17
18    Contact info: bruno.steux@gmail.com
19*/
20
21use crate::compile::Operation;
22use crate::generate::FlagsState;
23use log::{debug, error};
24use std::fmt::{self, Debug};
25use std::io::Write;
26
27#[derive(Debug, Copy, Clone, PartialEq)]
28pub enum AsmMnemonic {
29    LDA,
30    LDX,
31    LDY,
32    STA,
33    STX,
34    STY,
35    TAX,
36    TAY,
37    TXA,
38    TYA,
39    ADC,
40    SBC,
41    EOR,
42    AND,
43    ORA,
44    LSR,
45    ASL,
46    ROL,
47    ROR,
48    CLC,
49    SEC,
50    CMP,
51    CPX,
52    CPY,
53    BCC,
54    BCS,
55    BEQ,
56    BMI,
57    BNE,
58    BPL,
59    INC,
60    INX,
61    INY,
62    DEC,
63    DEX,
64    DEY,
65    JMP,
66    JSR,
67    RTS,
68    RTI,
69    PHA,
70    PLA,
71    PHP,
72    PLP,
73    NOP,
74}
75
76impl fmt::Display for AsmMnemonic {
77    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
78        fmt::Debug::fmt(self, f)
79    }
80}
81
82#[derive(Debug, Clone)]
83pub struct AsmInstruction {
84    pub mnemonic: AsmMnemonic,
85    pub dasm_operand: String,
86    pub cycles: u32,
87    pub cycles_alt: Option<u32>,
88    pub nb_bytes: u32,
89    pub protected: bool,
90}
91
92#[derive(Debug, Clone)]
93enum AsmLine {
94    Label(String),
95    Instruction(AsmInstruction),
96    Inline(String, u32),
97    Comment(String),
98    Dummy,
99}
100
101impl AsmLine {
102    fn write(&self, writer: &mut dyn Write, cycles: bool) -> Result<usize, std::io::Error> {
103        let mut s = 0;
104        match self {
105            AsmLine::Label(string) => {
106                s += writer.write(string.as_bytes())?;
107                s += writer.write("\n".as_bytes())?;
108            }
109            AsmLine::Instruction(inst) => {
110                if cycles {
111                    let c = if let Some(alt) = inst.cycles_alt {
112                        format!("\t; {}/{}", inst.cycles, alt)
113                    } else {
114                        format!("\t; {}", inst.cycles)
115                    };
116                    if !inst.dasm_operand.is_empty() {
117                        s += writer.write(
118                            format!("\t{} {:19}{}\n", inst.mnemonic, &inst.dasm_operand, c)
119                                .as_bytes(),
120                        )?;
121                    } else {
122                        s += writer.write(format!("\t{:23}{}\n", inst.mnemonic, c).as_bytes())?;
123                    }
124                } else {
125                    if !inst.dasm_operand.is_empty() {
126                        s += writer.write(
127                            format!("\t{} {}\n", inst.mnemonic, &inst.dasm_operand).as_bytes(),
128                        )?;
129                    } else {
130                        s += writer.write(format!("\t{}\n", inst.mnemonic).as_bytes())?;
131                    }
132                }
133            }
134            AsmLine::Inline(inst, _) => {
135                s += writer.write(format!("\t{}\n", inst).as_bytes())?;
136            }
137            AsmLine::Comment(comment) => {
138                s += writer.write(format!(";{}\n", comment).as_bytes())?;
139            }
140            AsmLine::Dummy => (),
141        }
142        Ok(s)
143    }
144}
145
146#[derive(Debug, Clone)]
147pub struct AssemblyCode {
148    code: Vec<AsmLine>,
149}
150
151impl AssemblyCode {
152    pub fn new() -> AssemblyCode {
153        AssemblyCode {
154            code: Vec::<AsmLine>::new(),
155        }
156    }
157
158    pub fn size_bytes(&self) -> u32 {
159        let mut size = 0;
160        for c in self.code.iter() {
161            match c {
162                AsmLine::Instruction(i) => {
163                    size += i.nb_bytes;
164                }
165                AsmLine::Inline(_, s) => {
166                    size += s;
167                }
168                _ => {}
169            }
170        }
171        size
172    }
173
174    pub fn append_asm(&mut self, inst: AsmInstruction) {
175        self.code.push(AsmLine::Instruction(inst));
176    }
177
178    pub fn append_inline(&mut self, s: String, size: Option<u32>) {
179        self.code.push(AsmLine::Inline(s, size.unwrap_or(3)));
180    }
181
182    pub fn append_label(&mut self, s: String) {
183        self.code.push(AsmLine::Label(s));
184    }
185
186    pub fn append_comment(&mut self, s: String) {
187        self.code.push(AsmLine::Comment(s));
188    }
189
190    pub fn append_dummy(&mut self) -> usize {
191        self.code.push(AsmLine::Dummy);
192        self.code.len() - 1
193    }
194
195    pub fn set(&mut self, line: usize, inst: AsmInstruction) {
196        *self.code.get_mut(line).unwrap() = AsmLine::Instruction(inst);
197    }
198
199    // For inlined code: modify local labels in each instruction
200    pub fn append_code(&mut self, code: &AssemblyCode, inline_counter: u32) {
201        for i in &code.code {
202            match i {
203                AsmLine::Label(l) => self
204                    .code
205                    .push(AsmLine::Label(format!("{}inline{}", l, inline_counter))),
206                AsmLine::Instruction(inst) => match inst.mnemonic {
207                    AsmMnemonic::BCC
208                    | AsmMnemonic::BCS
209                    | AsmMnemonic::BEQ
210                    | AsmMnemonic::BMI
211                    | AsmMnemonic::BNE
212                    | AsmMnemonic::BPL
213                    | AsmMnemonic::JMP => {
214                        self.code.push(AsmLine::Instruction(AsmInstruction {
215                            mnemonic: inst.mnemonic,
216                            dasm_operand: format!("{}inline{}", inst.dasm_operand, inline_counter),
217                            cycles: inst.cycles,
218                            cycles_alt: inst.cycles_alt,
219                            nb_bytes: inst.nb_bytes,
220                            protected: false,
221                        }));
222                    }
223                    _ => self.code.push(i.clone()),
224                },
225                _ => self.code.push(i.clone()),
226            }
227        }
228        //self.code.extend_from_slice(&code.code);
229    }
230
231    pub fn write(&self, writer: &mut dyn Write, cycles: bool) -> Result<usize, std::io::Error> {
232        let mut s = 0;
233        for i in &self.code {
234            s += i.write(writer, cycles)?;
235        }
236        Ok(s)
237    }
238
239    pub fn optimize(&mut self) -> u32 {
240        let mut removed_instructions = 0u32;
241        let mut accumulator = None;
242        let mut x_register = None;
243        let mut y_register = None;
244        let mut iter = itertools::multipeek(self.code.iter_mut());
245        let mut first = iter.next();
246        let mut flags = FlagsState::Unknown;
247
248        loop {
249            match &first {
250                None => return removed_instructions,
251                Some(AsmLine::Instruction(_)) => break,
252                _ => first = iter.next(),
253            }
254        }
255        let mut second = iter.next();
256
257        // Analyze the first instruction to check for a load
258        if let Some(AsmLine::Instruction(inst)) = &first {
259            if inst.mnemonic == AsmMnemonic::LDA {
260                accumulator = Some(inst.dasm_operand.clone());
261                flags = FlagsState::A;
262            } else if inst.mnemonic == AsmMnemonic::LDX {
263                x_register = Some(inst.dasm_operand.clone());
264                flags = FlagsState::X;
265            } else if inst.mnemonic == AsmMnemonic::LDY {
266                y_register = Some(inst.dasm_operand.clone());
267                flags = FlagsState::Y;
268            }
269        } else {
270            unreachable!();
271        }
272
273        loop {
274            // For each iteration of this loop, first must point to an Instruction
275            // and second point to the next asm line
276            let mut remove_both = false;
277            let mut remove_first = false;
278            let mut remove_second = false;
279            let mut swap_both = false;
280
281            // Remove JMP to the following label
282            if let Some(AsmLine::Instruction(i1)) = &first {
283                if let Some(AsmLine::Label(l)) = &second {
284                    if i1.mnemonic == AsmMnemonic::JMP && &i1.dasm_operand == l && !i1.protected {
285                        *first.unwrap() = AsmLine::Dummy;
286                        removed_instructions += 1;
287                        first = second;
288                        loop {
289                            match &first {
290                                None => return removed_instructions,
291                                Some(AsmLine::Instruction(_)) => break,
292                                _ => first = iter.next(),
293                            }
294                        }
295                        second = iter.next();
296                    }
297                }
298            }
299
300            // Make sure second points also to an instruction
301            loop {
302                match &second {
303                    None => return removed_instructions,
304                    Some(AsmLine::Instruction(_)) => break,
305                    Some(AsmLine::Label(_)) => {
306                        // If this is a label, restart
307                        first = iter.next();
308                        loop {
309                            match &first {
310                                None => return removed_instructions,
311                                Some(AsmLine::Instruction(_)) => break,
312                                _ => first = iter.next(),
313                            }
314                        }
315                        second = iter.next();
316                        // Reset known register values
317                        accumulator = None;
318                        x_register = None;
319                        y_register = None;
320                        // Analyze the first instruction to check for a load
321                        if let Some(AsmLine::Instruction(inst)) = &first {
322                            if inst.mnemonic == AsmMnemonic::LDA {
323                                accumulator = Some(inst.dasm_operand.clone());
324                                flags = FlagsState::A;
325                            } else if inst.mnemonic == AsmMnemonic::LDX {
326                                x_register = Some(inst.dasm_operand.clone());
327                                flags = FlagsState::X;
328                            } else if inst.mnemonic == AsmMnemonic::LDY {
329                                y_register = Some(inst.dasm_operand.clone());
330                                flags = FlagsState::Y;
331                            } else {
332                                flags = FlagsState::Unknown;
333                            }
334                        } else {
335                            unreachable!();
336                        }
337                    }
338                    _ => second = iter.next(),
339                }
340            }
341
342            // Analyze pairs of instructions
343            if let Some(AsmLine::Instruction(i1)) = &first {
344                if let Some(AsmLine::Instruction(i2)) = &second {
345                    // Remove PLA/PHA pairs
346                    if i1.mnemonic == AsmMnemonic::PLA
347                        && i2.mnemonic == AsmMnemonic::PHA
348                        && !i1.protected
349                        && !i2.protected
350                    {
351                        remove_both = true;
352                    }
353                    // Remove JMP followed by JMP
354                    if i1.mnemonic == AsmMnemonic::JMP
355                        && i2.mnemonic == AsmMnemonic::JMP
356                        && !i1.protected
357                        && !i2.protected
358                    {
359                        remove_second = true;
360                    }
361                    // Remove STA followed by LDA
362                    if i1.mnemonic == AsmMnemonic::STA
363                        && i2.mnemonic == AsmMnemonic::LDA
364                        && i1.dasm_operand == i2.dasm_operand
365                        && !i2.protected
366                    {
367                        remove_second = true;
368                    }
369                    // Remove LDA followed by STA
370                    if i1.mnemonic == AsmMnemonic::LDA
371                        && i2.mnemonic == AsmMnemonic::STA
372                        && i1.dasm_operand == i2.dasm_operand
373                        && !i2.protected
374                    {
375                        remove_second = true;
376                    }
377                    // Remove LDY followed by STY
378                    if i1.mnemonic == AsmMnemonic::LDY
379                        && i2.mnemonic == AsmMnemonic::STY
380                        && i1.dasm_operand == i2.dasm_operand
381                        && !i2.protected
382                    {
383                        remove_second = true;
384                    }
385                    // Remove LDX followed by STX
386                    if i1.mnemonic == AsmMnemonic::LDX
387                        && i2.mnemonic == AsmMnemonic::STX
388                        && i1.dasm_operand == i2.dasm_operand
389                        && !i2.protected
390                    {
391                        remove_second = true;
392                    }
393                    // Remove TAX followed by TXA
394                    if i1.mnemonic == AsmMnemonic::TAX
395                        && i2.mnemonic == AsmMnemonic::TXA
396                        && !i2.protected
397                    {
398                        remove_second = true;
399                    }
400                    // Remove TXA followed by TAX
401                    if i1.mnemonic == AsmMnemonic::TXA
402                        && i2.mnemonic == AsmMnemonic::TAX
403                        && !i2.protected
404                    {
405                        remove_second = true;
406                    }
407                    // Remove TAY followed by TYA
408                    if i1.mnemonic == AsmMnemonic::TAY
409                        && i2.mnemonic == AsmMnemonic::TYA
410                        && !i2.protected
411                    {
412                        remove_second = true;
413                    }
414                    // Remove TYA followed by TAY
415                    if i1.mnemonic == AsmMnemonic::TYA
416                        && i2.mnemonic == AsmMnemonic::TAY
417                        && !i2.protected
418                    {
419                        remove_second = true;
420                    }
421                    if i1.mnemonic == AsmMnemonic::LDA
422                        && i2.mnemonic == AsmMnemonic::LDA
423                        && !i1.protected
424                    {
425                        remove_first = true;
426                    }
427                    if i1.mnemonic == AsmMnemonic::LDY
428                        && i2.mnemonic == AsmMnemonic::LDY
429                        && !i1.protected
430                    {
431                        remove_first = true;
432                    }
433                    if i1.mnemonic == AsmMnemonic::LDX
434                        && i2.mnemonic == AsmMnemonic::LDX
435                        && !i1.protected
436                    {
437                        remove_first = true;
438                    }
439                    if i2.mnemonic == AsmMnemonic::ORA && i2.dasm_operand == "#0" && !i2.protected {
440                        remove_second = true;
441                    }
442                    if i1.mnemonic == AsmMnemonic::LDA
443                        && (i2.mnemonic == AsmMnemonic::SEC || i2.mnemonic == AsmMnemonic::CLC)
444                    {
445                        swap_both = true;
446                    }
447                    // Check CMP and remove the branck if the result is obvious
448                    if let Some(r) = &accumulator {
449                        if r.starts_with("#")
450                            && i1.mnemonic == AsmMnemonic::CMP
451                            && i1.dasm_operand.starts_with('#')
452                        {
453                            // The result IS obvious
454                            match i2.mnemonic {
455                                AsmMnemonic::BNE => {
456                                    if *r == i1.dasm_operand && !i2.protected{
457                                        remove_both = true;
458                                    }
459                                }
460                                AsmMnemonic::BEQ => {
461                                    if *r != i1.dasm_operand && !i2.protected {
462                                        remove_both = true;
463                                    }
464                                }
465                                _ => (),
466                            }
467                        }
468                    }
469                    if let Some(r) = &x_register {
470                        if r.starts_with("#")
471                            && i1.mnemonic == AsmMnemonic::CPX
472                            && i1.dasm_operand.starts_with('#')
473                        {
474                            // The result IS obvious
475                            match i2.mnemonic {
476                                AsmMnemonic::BNE => {
477                                    if *r == i1.dasm_operand && !i2.protected {
478                                        remove_both = true;
479                                    }
480                                }
481                                AsmMnemonic::BEQ => {
482                                    if *r != i1.dasm_operand && !i2.protected {
483                                        remove_both = true;
484                                    }
485                                }
486                                _ => (),
487                            }
488                        }
489                    }
490                    if let Some(r) = &y_register {
491                        if r.starts_with("#")
492                            && i1.mnemonic == AsmMnemonic::CPY
493                            && i1.dasm_operand.starts_with('#')
494                        {
495                            // The result IS obvious
496                            match i2.mnemonic {
497                                AsmMnemonic::BNE => {
498                                    if *r == i1.dasm_operand && !i2.protected {
499                                        remove_both = true;
500                                    }
501                                }
502                                AsmMnemonic::BEQ => {
503                                    if *r != i1.dasm_operand && !i2.protected {
504                                        remove_both = true;
505                                    }
506                                }
507                                _ => (),
508                            }
509                        }
510                    }
511                } else {
512                    unreachable!()
513                };
514            } else {
515                unreachable!()
516            };
517
518            if !remove_second && !remove_both {
519                // Analyze the second instruction to check for a load
520                if let Some(AsmLine::Instruction(inst)) = &second {
521                    match inst.mnemonic {
522                        AsmMnemonic::LDA => {
523                            if let Some(v) = &accumulator {
524                                if v.eq(&inst.dasm_operand) {
525                                    if flags == FlagsState::A {
526                                        // Remove this instruction
527                                        remove_second = !inst.protected;
528                                    } else {
529                                        // OK. It's a little more complex. Let's check the
530                                        // following instructions.
531                                        // If we can find a CMP or a STA/LDA pair in the next 2
532                                        // instructions, it's OK, we can remove the instruction
533                                        if let Some(AsmLine::Instruction(i1)) = iter.peek() {
534                                            if i1.mnemonic == AsmMnemonic::CMP {
535                                                remove_second = !inst.protected;
536                                            } else if i1.mnemonic == AsmMnemonic::STA {
537                                                if let Some(line) = iter.peek() {
538                                                    if let AsmLine::Instruction(i2) = line {
539                                                        if i2.mnemonic == AsmMnemonic::LDA
540                                                            || i2.mnemonic == AsmMnemonic::LDX
541                                                            || i2.mnemonic == AsmMnemonic::LDY
542                                                        {
543                                                            remove_second = !inst.protected;
544                                                        }
545                                                    } else if let AsmLine::Dummy = line {
546                                                        if let Some(AsmLine::Instruction(i3)) =
547                                                            iter.peek()
548                                                        {
549                                                            if i3.mnemonic == AsmMnemonic::LDA
550                                                                || i3.mnemonic == AsmMnemonic::LDX
551                                                                || i3.mnemonic == AsmMnemonic::LDY
552                                                            {
553                                                                remove_second = !inst.protected;
554                                                            }
555                                                        }
556                                                    }
557                                                }
558                                            }
559                                        }
560                                    }
561                                }
562                            }
563                            accumulator = Some(inst.dasm_operand.clone());
564                            flags = FlagsState::A;
565                        }
566                        AsmMnemonic::LDX => {
567                            if let Some(v) = &accumulator {
568                                if v.ends_with(",X") {
569                                    accumulator = None;
570                                }
571                            }
572                            if let Some(v) = &y_register {
573                                if v.ends_with(",X") {
574                                    y_register = None;
575                                }
576                            }
577                            if let Some(v) = &x_register {
578                                if v.eq(&inst.dasm_operand) {
579                                    // Remove this instruction
580                                    remove_second = !inst.protected;
581                                }
582                            }
583                            x_register = Some(inst.dasm_operand.clone());
584                            flags = FlagsState::X;
585                        }
586                        AsmMnemonic::LDY => {
587                            if let Some(v) = &accumulator {
588                                if v.ends_with(",Y") {
589                                    accumulator = None;
590                                }
591                            }
592                            if let Some(v) = &x_register {
593                                if v.ends_with(",Y") {
594                                    x_register = None;
595                                }
596                            }
597                            if let Some(v) = &y_register {
598                                if v.eq(&inst.dasm_operand) {
599                                    // Remove this instruction
600                                    remove_second = !inst.protected;
601                                }
602                            }
603                            y_register = Some(inst.dasm_operand.clone());
604                            flags = FlagsState::Y;
605                        }
606                        AsmMnemonic::DEC | AsmMnemonic::INC => {
607                            if let Some(v) = &accumulator {
608                                if v.eq(&inst.dasm_operand) {
609                                    accumulator = None;
610                                }
611                            }
612                            if let Some(v) = &x_register {
613                                if v.eq(&inst.dasm_operand) {
614                                    x_register = None;
615                                }
616                            }
617                            if let Some(v) = &y_register {
618                                if v.eq(&inst.dasm_operand) {
619                                    y_register = None;
620                                }
621                            }
622                        }
623                        AsmMnemonic::INX | AsmMnemonic::DEX => {
624                            if let Some(v) = &accumulator {
625                                if v.ends_with(",X") {
626                                    accumulator = None;
627                                }
628                            }
629                            if let Some(v) = &y_register {
630                                if v.ends_with(",X") {
631                                    y_register = None;
632                                }
633                            }
634                            x_register = None;
635                        }
636                        AsmMnemonic::INY | AsmMnemonic::DEY => {
637                            if let Some(v) = &accumulator {
638                                if v.ends_with(",Y") {
639                                    accumulator = None;
640                                }
641                            }
642                            if let Some(v) = &x_register {
643                                if v.ends_with(",Y") {
644                                    x_register = None;
645                                }
646                            }
647                            y_register = None;
648                        }
649                        AsmMnemonic::TAX => {
650                            x_register = accumulator.clone();
651                            if let Some(v) = &accumulator {
652                                if v.ends_with(",X") {
653                                    accumulator = None;
654                                    x_register = None;
655                                }
656                            }
657                            if let Some(v) = &y_register {
658                                if v.ends_with(",X") {
659                                    y_register = None;
660                                }
661                            }
662                        }
663                        AsmMnemonic::TAY => {
664                            y_register = accumulator.clone();
665                            if let Some(v) = &accumulator {
666                                if v.ends_with(",Y") {
667                                    accumulator = None;
668                                    y_register = None;
669                                }
670                            }
671                            if let Some(v) = &x_register {
672                                if v.ends_with(",Y") {
673                                    x_register = None;
674                                }
675                            }
676                        }
677                        AsmMnemonic::TXA => {
678                            accumulator = x_register.clone();
679                        }
680                        AsmMnemonic::TYA => {
681                            accumulator = y_register.clone();
682                        }
683                        AsmMnemonic::STA | AsmMnemonic::STX | AsmMnemonic::STY => {
684                            if let Some(v) = &accumulator {
685                                if !v.starts_with("#") {
686                                    accumulator = None;
687                                }
688                            }
689                            if let Some(v) = &x_register {
690                                if !v.starts_with("#") {
691                                    x_register = None;
692                                }
693                            }
694                            if let Some(v) = &y_register {
695                                if !v.starts_with("#") {
696                                    y_register = None;
697                                }
698                            }
699                        }
700                        AsmMnemonic::ADC
701                        | AsmMnemonic::SBC
702                        | AsmMnemonic::EOR
703                        | AsmMnemonic::AND
704                        | AsmMnemonic::ORA => accumulator = None,
705                        AsmMnemonic::LSR | AsmMnemonic::ASL => accumulator = None,
706                        AsmMnemonic::PLA | AsmMnemonic::PHA => accumulator = None,
707                        AsmMnemonic::JSR | AsmMnemonic::JMP => {
708                            accumulator = None;
709                            x_register = None;
710                            y_register = None;
711                        }
712                        AsmMnemonic::CPX | AsmMnemonic::CPY | AsmMnemonic::CMP => {
713                            flags = FlagsState::Unknown;
714                        }
715                        _ => (),
716                    }
717                } else {
718                    unreachable!();
719                }
720            }
721
722            if swap_both {
723                let tmp1 = if let Some(f) = &first {
724                    (**f).clone()
725                } else {
726                    AsmLine::Dummy
727                };
728                let tmp2 = if let Some(f) = &second {
729                    (**f).clone()
730                } else {
731                    AsmLine::Dummy
732                };
733                if let Some(f) = &mut first {
734                    (**f) = tmp2;
735                }
736                if let Some(f) = &mut second {
737                    (**f) = tmp1;
738                }
739                accumulator = None;
740            } else if remove_both {
741                *first.unwrap() = AsmLine::Dummy;
742                *second.unwrap() = AsmLine::Dummy;
743                removed_instructions += 2;
744                first = iter.next();
745                loop {
746                    match &first {
747                        None => return removed_instructions,
748                        Some(AsmLine::Instruction(_)) => break,
749                        _ => first = iter.next(),
750                    }
751                }
752                // Reset known register values
753                // This is not optimal, since theoretically this instruction is in the flow of instructions.
754                accumulator = None;
755                x_register = None;
756                y_register = None;
757                // Analyze the first instruction to check for a load
758                if let Some(AsmLine::Instruction(inst)) = &first {
759                    if inst.mnemonic == AsmMnemonic::LDA {
760                        accumulator = Some(inst.dasm_operand.clone());
761                    }
762                    if inst.mnemonic == AsmMnemonic::LDX {
763                        x_register = Some(inst.dasm_operand.clone());
764                    }
765                    if inst.mnemonic == AsmMnemonic::LDY {
766                        y_register = Some(inst.dasm_operand.clone());
767                    }
768                } else {
769                    unreachable!();
770                }
771                second = iter.next();
772            } else if remove_second {
773                *second.unwrap() = AsmLine::Dummy;
774                removed_instructions += 1;
775                second = iter.next();
776            } else if remove_first {
777                *first.unwrap() = AsmLine::Dummy;
778                removed_instructions += 1;
779                first = second;
780                second = iter.next();
781            } else {
782                first = second;
783                second = iter.next();
784            }
785        }
786    }
787
788    pub fn check_branches(&mut self) -> u32 {
789        // Loop until there is no problematic branch instruction
790        let mut restart = true;
791        let mut nb_fixes = 0;
792        debug!("Code: {:?}", self);
793        while restart {
794            // Check each branch instruction one after each other
795            // Let's find the first one
796            let mut position = 0;
797            let mut i = self.code.iter();
798            let mut repair = false;
799            loop {
800                let j = i.next();
801                if j.is_none() {
802                    restart = false;
803                    break;
804                }
805                if let Some(AsmLine::Instruction(inst)) = j {
806                    match inst.mnemonic {
807                        AsmMnemonic::BEQ
808                        | AsmMnemonic::BNE
809                        | AsmMnemonic::BMI
810                        | AsmMnemonic::BPL
811                        | AsmMnemonic::BCS
812                        | AsmMnemonic::BCC => {
813                            // Ok, let's try to find the label above and under and try to count the bytes
814                            let mut bytes_above = 0;
815                            let mut bytes_below = 0;
816                            let mut index_above = position;
817                            let mut index_below = position + 1;
818                            let mut reached_above = false;
819                            let above;
820                            let mut notfound = 0;
821                            loop {
822                                if !reached_above {
823                                    match &self.code[index_above] {
824                                        AsmLine::Label(l) => {
825                                            debug!("Iter above: {:?}", l);
826                                            if *l == inst.dasm_operand {
827                                                above = true;
828                                                break;
829                                            }
830                                        }
831                                        AsmLine::Inline(_, s) => {
832                                            bytes_above += s;
833                                        }
834                                        AsmLine::Instruction(k) => {
835                                            debug!("Iter above: {:?}", k);
836                                            bytes_above += k.nb_bytes;
837                                        }
838                                        _ => (),
839                                    }
840                                }
841                                match self.code.get(index_below) {
842                                    Some(AsmLine::Label(l)) => {
843                                        debug!("Iter below: {:?}", l);
844                                        if *l == inst.dasm_operand {
845                                            above = false;
846                                            break;
847                                        }
848                                    }
849                                    Some(AsmLine::Inline(_, s)) => {
850                                        bytes_below += s;
851                                    }
852                                    Some(AsmLine::Instruction(k)) => {
853                                        debug!("Iter below: {:?}", k);
854                                        bytes_below += k.nb_bytes;
855                                    }
856                                    None => notfound |= 2,
857                                    _ => (),
858                                }
859                                if index_above == 0 {
860                                    reached_above = true;
861                                    notfound |= 1;
862                                } else {
863                                    index_above -= 1;
864                                }
865                                index_below += 1;
866                                if notfound == 3 {
867                                    error!("Label {} not found", inst.dasm_operand);
868                                    unreachable!()
869                                };
870                            }
871                            // Ok, now we have the distance in bytes
872                            let distance = if above { bytes_above } else { bytes_below };
873                            //error!("distance = {:?}", distance);
874                            //if above {unreachable!();}
875                            if distance > 127 {
876                                // OK. We have a problem here
877                                // This branch should be changed for a jump
878                                repair = true;
879                                break;
880                            }
881                        }
882                        _ => (),
883                    }
884                }
885                position += 1;
886            }
887            if repair {
888                nb_fixes += 1;
889                // Identifies branch operation
890                let signed;
891                let mut remove = 1;
892                let label;
893                let operation = if let AsmLine::Instruction(inst) = &self.code[position] {
894                    signed =
895                        (inst.mnemonic == AsmMnemonic::BPL) || (inst.mnemonic == AsmMnemonic::BMI);
896                    label = inst.dasm_operand.clone();
897                    match inst.mnemonic {
898                        AsmMnemonic::BNE => Operation::Neq,
899                        AsmMnemonic::BEQ => Operation::Eq,
900                        AsmMnemonic::BMI => {
901                            if let Some(AsmLine::Instruction(inst2)) = self.code.get(position + 1) {
902                                if inst2.mnemonic == AsmMnemonic::BEQ
903                                    && inst2.dasm_operand == inst.dasm_operand
904                                {
905                                    remove = 2;
906                                    Operation::Lte
907                                } else {
908                                    Operation::Lt
909                                }
910                            } else {
911                                Operation::Lt
912                            }
913                        }
914                        AsmMnemonic::BCC => {
915                            if let Some(AsmLine::Instruction(inst2)) = self.code.get(position + 1) {
916                                if inst2.mnemonic == AsmMnemonic::BEQ
917                                    && inst2.dasm_operand == inst.dasm_operand
918                                {
919                                    remove = 2;
920                                    Operation::Lte
921                                } else {
922                                    Operation::Lt
923                                }
924                            } else {
925                                Operation::Lt
926                            }
927                        }
928                        AsmMnemonic::BPL => Operation::Gte,
929                        AsmMnemonic::BCS => Operation::Gte,
930                        _ => unreachable!(),
931                    }
932                } else {
933                    unreachable!();
934                };
935                // Negate the operation
936                let operation2 = match operation {
937                    Operation::Eq => Operation::Neq,
938                    Operation::Neq => Operation::Eq,
939                    Operation::Gt => Operation::Lte,
940                    Operation::Gte => Operation::Lt,
941                    Operation::Lt => Operation::Gte,
942                    Operation::Lte => Operation::Gt,
943                    _ => unreachable!(),
944                };
945                let mut tail = self.code.split_off(position + remove);
946                self.code.truncate(position);
947                let label2 = format!(".fix{}", nb_fixes);
948                match operation2 {
949                    Operation::Eq => self.code.push(AsmLine::Instruction(AsmInstruction {
950                        mnemonic: AsmMnemonic::BEQ,
951                        dasm_operand: label2.clone(),
952                        cycles: 2,
953                        cycles_alt: Some(3),
954                        nb_bytes: 2,
955                        protected: false,
956                    })),
957                    Operation::Neq => self.code.push(AsmLine::Instruction(AsmInstruction {
958                        mnemonic: AsmMnemonic::BNE,
959                        dasm_operand: label2.clone(),
960                        cycles: 2,
961                        cycles_alt: Some(3),
962                        nb_bytes: 2,
963                        protected: false,
964                    })),
965                    Operation::Gt => {
966                        let label3 = format!(".fixup{}", nb_fixes);
967                        self.code.push(AsmLine::Instruction(AsmInstruction {
968                            mnemonic: AsmMnemonic::BEQ,
969                            dasm_operand: label3.clone(),
970                            cycles: 2,
971                            cycles_alt: Some(3),
972                            nb_bytes: 2,
973                            protected: false,
974                        }));
975                        if signed {
976                            self.code.push(AsmLine::Instruction(AsmInstruction {
977                                mnemonic: AsmMnemonic::BPL,
978                                dasm_operand: label2.clone(),
979                                cycles: 2,
980                                cycles_alt: Some(3),
981                                nb_bytes: 2,
982                                protected: false,
983                            }));
984                        } else {
985                            self.code.push(AsmLine::Instruction(AsmInstruction {
986                                mnemonic: AsmMnemonic::BCS,
987                                dasm_operand: label2.clone(),
988                                cycles: 2,
989                                cycles_alt: Some(3),
990                                nb_bytes: 2,
991                                protected: false,
992                            }));
993                        }
994                        self.code.push(AsmLine::Label(label3));
995                    }
996                    Operation::Gte => {
997                        if signed {
998                            self.code.push(AsmLine::Instruction(AsmInstruction {
999                                mnemonic: AsmMnemonic::BPL,
1000                                dasm_operand: label2.clone(),
1001                                cycles: 2,
1002                                cycles_alt: Some(3),
1003                                nb_bytes: 2,
1004                                protected: false,
1005                            }));
1006                        } else {
1007                            self.code.push(AsmLine::Instruction(AsmInstruction {
1008                                mnemonic: AsmMnemonic::BCS,
1009                                dasm_operand: label2.clone(),
1010                                cycles: 2,
1011                                cycles_alt: Some(3),
1012                                nb_bytes: 2,
1013                                protected: false,
1014                            }));
1015                        }
1016                    }
1017                    Operation::Lt => {
1018                        if signed {
1019                            self.code.push(AsmLine::Instruction(AsmInstruction {
1020                                mnemonic: AsmMnemonic::BMI,
1021                                dasm_operand: label2.clone(),
1022                                cycles: 2,
1023                                cycles_alt: Some(3),
1024                                nb_bytes: 2,
1025                                protected: false,
1026                            }));
1027                        } else {
1028                            self.code.push(AsmLine::Instruction(AsmInstruction {
1029                                mnemonic: AsmMnemonic::BCC,
1030                                dasm_operand: label2.clone(),
1031                                cycles: 2,
1032                                cycles_alt: Some(3),
1033                                nb_bytes: 2,
1034                                protected: false,
1035                            }));
1036                        }
1037                    }
1038                    _ => unreachable!(),
1039                }
1040                self.code.push(AsmLine::Instruction(AsmInstruction {
1041                    mnemonic: AsmMnemonic::JMP,
1042                    dasm_operand: label,
1043                    cycles: 3,
1044                    cycles_alt: None,
1045                    nb_bytes: 3,
1046                    protected: false,
1047                }));
1048                self.code.push(AsmLine::Label(label2));
1049                self.code.append(&mut tail);
1050            }
1051        }
1052        nb_fixes
1053    }
1054}