celestial_hub_compass/codegen/mips/
mod.rs

1use celestial_hub_astrolabe::{
2  ast::{
3    DataSection, Instruction, InstructionArgument, Program, Register, Statement, TextSection,
4    Variable,
5  },
6  lexer::{
7    tokens::{Type, Value},
8    traits::iterator,
9  },
10};
11
12use crate::ast::{
13  context, BinaryOperation, Condition, Expr, Function, FunctionCall, Operand, Operator,
14  Statement as CompassStatement, VarType,
15};
16use std::collections::HashMap;
17
18use super::{context::Context, Codegen};
19
20pub struct MipsCodegen;
21
22macro_rules! create_instruction {
23  ($instruction:path, $register:expr, $lhs_register:expr, $rhs:expr) => {
24    Statement::Instruction($instruction(
25      [
26        InstructionArgument::Register(Register { name: $register }),
27        InstructionArgument::Register(Register {
28          name: $lhs_register,
29        }),
30        $rhs,
31      ]
32      .into(),
33    ))
34  };
35}
36
37impl Codegen for MipsCodegen {
38  fn generate(&self, ast: Vec<CompassStatement>, context: &mut Context) -> Result<String, String> {
39    if context.scope_level == 0 {
40      context
41        .text_section
42        .statements
43        .push(Statement::Label("main".to_string()));
44    }
45
46    for statement in ast {
47      match statement {
48        CompassStatement::VariableDeclaration(var) => {
49          let register = find_or_create_reg(&mut context.register_map, var.name.clone());
50
51          match var.value {
52            Expr::Operand(op) => match op {
53              Operand::LiteralI8(val) => {
54                load_immediate(&mut context.text_section, register, val as i32)
55              }
56              Operand::LiteralI16(val) => {
57                load_immediate(&mut context.text_section, register, val as i32)
58              }
59              Operand::LiteralI32(val) => load_immediate(&mut context.text_section, register, val),
60              Operand::LiteralI64(val) => {
61                load_immediate(&mut context.text_section, register, val as i32)
62              }
63              Operand::LiteralU8(val) => {
64                load_immediate(&mut context.text_section, register, val as i32)
65              }
66              Operand::LiteralU16(val) => {
67                load_immediate(&mut context.text_section, register, val as i32)
68              }
69              Operand::LiteralU32(val) => {
70                load_immediate(&mut context.text_section, register, val as i32)
71              }
72              Operand::LiteralStr(val) => load_string(
73                &mut context.text_section,
74                &mut context.data_section,
75                register,
76                val,
77              ),
78              Operand::LiteralU64(val) => {
79                return Err("Cannot store a 64-bit integer in a 32-bit register".to_string());
80              }
81              Operand::LiteralBool(val) => {
82                load_immediate(&mut context.text_section, register, val as i32)
83              }
84              Operand::Identifier(var) => {
85                let var_register = context
86                  .register_map
87                  .get(&var)
88                  .ok_or_else(|| format!("Register {} not found", var))?
89                  .clone();
90
91                context
92                  .text_section
93                  .statements
94                  .push(Statement::Instruction(Instruction::Move(
95                    [
96                      InstructionArgument::Register(Register {
97                        name: register.clone(),
98                      }),
99                      InstructionArgument::Register(Register { name: var_register }),
100                    ]
101                    .into(),
102                  )));
103              }
104              // TODO: Handle float literals (should be in a different register)
105              Operand::LiteralF32(val) => todo!(),
106              Operand::LiteralF64(val) => todo!(),
107              Operand::Dereference(_) => todo!(),
108            },
109            Expr::BinaryOperation(bin_op) => match bin_op {
110              BinaryOperation::Arithmetic {
111                lhs, operator, rhs, ..
112              } => {
113                if is_register(&lhs) && is_register(&rhs) {
114                  let lhs = lhs.as_identifier()?;
115                  let rhs = rhs.as_identifier()?;
116
117                  let lhs_register = context
118                    .register_map
119                    .get(lhs)
120                    .ok_or_else(|| format!("Register {} not found", lhs))?
121                    .clone();
122                  let rhs_register = context
123                    .register_map
124                    .get(rhs)
125                    .ok_or_else(|| format!("Register {} not found", rhs))?
126                    .clone();
127
128                  context.text_section.statements.push(match operator {
129                    Operator::Add => create_instruction!(
130                      Instruction::Add,
131                      register,
132                      lhs_register,
133                      InstructionArgument::Register(Register { name: rhs_register })
134                    ),
135                    Operator::Sub => create_instruction!(
136                      Instruction::Sub,
137                      register,
138                      lhs_register,
139                      InstructionArgument::Register(Register { name: rhs_register })
140                    ),
141                    Operator::Mul => create_instruction!(
142                      Instruction::Mul,
143                      register,
144                      lhs_register,
145                      InstructionArgument::Register(Register { name: rhs_register })
146                    ),
147                    Operator::Div => create_instruction!(
148                      Instruction::Div,
149                      register,
150                      lhs_register,
151                      InstructionArgument::Register(Register { name: rhs_register })
152                    ),
153                  });
154                } else if is_register(&lhs) && is_immediate(&rhs) {
155                  let lhs = lhs.as_identifier()?;
156                  let lhs_register = context
157                    .register_map
158                    .get(lhs)
159                    .ok_or_else(|| format!("Register {} not found", lhs))?
160                    .clone();
161                  let rhs_value = rhs.as_immediate()?;
162
163                  context.text_section.statements.push(match operator {
164                    Operator::Add => create_instruction!(
165                      Instruction::Add,
166                      register,
167                      lhs_register,
168                      InstructionArgument::Immediate(rhs_value)
169                    ),
170                    Operator::Sub => create_instruction!(
171                      Instruction::Sub,
172                      register,
173                      lhs_register,
174                      InstructionArgument::Immediate(rhs_value)
175                    ),
176                    Operator::Mul => create_instruction!(
177                      Instruction::Mul,
178                      register,
179                      lhs_register,
180                      InstructionArgument::Immediate(rhs_value)
181                    ),
182                    Operator::Div => create_instruction!(
183                      Instruction::Div,
184                      register,
185                      lhs_register,
186                      InstructionArgument::Immediate(rhs_value)
187                    ),
188                  });
189                } else if is_immediate(&lhs) && is_immediate(&rhs) {
190                  let rhs = rhs.as_immediate()?;
191                  let lhs_value = lhs.as_immediate()?;
192
193                  // There is no instruction that can add two immediates, so we need to load one of them into a register first
194                  let lhs_register = new_register(&mut context.register_map);
195                  load_immediate(&mut context.text_section, lhs_register.clone(), lhs_value);
196
197                  context.text_section.statements.push(match operator {
198                    Operator::Add => create_instruction!(
199                      Instruction::Add,
200                      register,
201                      lhs_register,
202                      InstructionArgument::Immediate(rhs)
203                    ),
204                    Operator::Sub => create_instruction!(
205                      Instruction::Sub,
206                      register,
207                      lhs_register,
208                      InstructionArgument::Immediate(rhs)
209                    ),
210                    Operator::Mul => create_instruction!(
211                      Instruction::Mul,
212                      register,
213                      lhs_register,
214                      InstructionArgument::Immediate(rhs)
215                    ),
216                    Operator::Div => create_instruction!(
217                      Instruction::Div,
218                      register,
219                      lhs_register,
220                      InstructionArgument::Immediate(rhs)
221                    ),
222                    _ => todo!(),
223                  });
224                } else {
225                  Err(format!(
226                    "Invalid operands for arithmetic operation {} and {}",
227                    lhs, rhs
228                  ))?;
229                }
230              }
231              BinaryOperation::Conditional {
232                lhs,
233                condition,
234                rhs,
235                operation_type,
236              } => {
237                if operation_type != VarType::Bool {
238                  return Err("Conditional operations must be of type bool".to_string());
239                }
240
241                if is_register(&lhs) && is_register(&rhs) {
242                  let lhs = lhs.as_identifier()?;
243                  let rhs = rhs.as_identifier()?;
244
245                  let lhs_register = context
246                    .register_map
247                    .get(lhs)
248                    .ok_or_else(|| format!("Register {} not found", lhs))?
249                    .clone();
250                  let rhs_register = context
251                    .register_map
252                    .get(rhs)
253                    .ok_or_else(|| format!("Register {} not found", rhs))?
254                    .clone();
255                  let instruction = match condition {
256                    Condition::LessThan => Instruction::Slt,
257                    Condition::GreaterThan => Instruction::Sgt,
258                    Condition::LessThanOrEqual => Instruction::Sle,
259                    Condition::GreaterThanOrEqual => Instruction::Sge,
260                    Condition::Equal => Instruction::Seq,
261                    Condition::NotEqual => Instruction::Sne,
262                    Condition::And | Condition::Or => {
263                      return Err(
264                        "Cannot perform logical operations on immediate values".to_string(),
265                      );
266                    }
267                  };
268                  context.text_section.statements.push(create_instruction!(
269                    instruction,
270                    register,
271                    lhs_register,
272                    InstructionArgument::Register(Register { name: rhs_register })
273                  ));
274                } else if is_register(&lhs) && is_immediate(&rhs) {
275                  let lhs = lhs.as_identifier()?;
276                  let lhs_register = context
277                    .register_map
278                    .get(lhs)
279                    .ok_or_else(|| format!("Register {} not found", lhs))?
280                    .clone();
281                  let rhs_value = rhs.as_immediate()?;
282
283                  let instruction = match condition {
284                    Condition::LessThan => Instruction::Slt,
285                    Condition::GreaterThan => Instruction::Sgt,
286                    Condition::LessThanOrEqual => Instruction::Sle,
287                    Condition::GreaterThanOrEqual => Instruction::Sge,
288                    Condition::Equal => Instruction::Seq,
289                    Condition::NotEqual => Instruction::Sne,
290                    Condition::And | Condition::Or => {
291                      return Err(
292                        "Cannot perform logical operations on immediate values".to_string(),
293                      );
294                    }
295                  };
296
297                  context.text_section.statements.push(create_instruction!(
298                    instruction,
299                    register,
300                    lhs_register,
301                    InstructionArgument::Immediate(rhs_value)
302                  ));
303                } else if is_immediate(&lhs) && is_immediate(&rhs) {
304                  let lhs_value = lhs.as_immediate()?;
305                  let rhs_value = rhs.as_immediate()?;
306
307                  let lhs_register = new_register(&mut context.register_map);
308                  load_immediate(&mut context.text_section, lhs_register.clone(), lhs_value);
309
310                  let instruction = match condition {
311                    Condition::LessThan => Instruction::Slt,
312                    Condition::GreaterThan => Instruction::Sgt,
313                    Condition::LessThanOrEqual => Instruction::Sle,
314                    Condition::GreaterThanOrEqual => Instruction::Sge,
315                    Condition::Equal => Instruction::Seq,
316                    Condition::NotEqual => Instruction::Sne,
317                    Condition::And | Condition::Or => {
318                      return Err(
319                        "Cannot perform logical operations on immediate values".to_string(),
320                      );
321                    }
322                  };
323
324                  context.text_section.statements.push(create_instruction!(
325                    instruction,
326                    register,
327                    lhs_register,
328                    InstructionArgument::Immediate(rhs_value)
329                  ));
330                }
331              }
332            },
333            Expr::FunctionCall(function_call) => {
334              let function = context
335                .get_function(&function_call.name)
336                .ok_or_else(|| format!("Function {} not found", function_call.name))?;
337
338              if function.is_builtin {
339                match function.name.as_str() {
340                  "read_int" => {
341                    load_immediate(&mut context.text_section, "$v0".to_string(), 5);
342                  }
343                  "read_string" => {
344                    // $a0 = address of the buffer
345                    // $a1 = length of the buffer
346
347                    load_immediate(&mut context.text_section, "$v0".to_string(), 8);
348
349                    let size: i32 = if let Operand::LiteralU32(size) = &function_call.params[0] {
350                      *size as i32
351                    } else {
352                      return Err("Invalid argument for read_string".to_string());
353                    };
354
355                    context.buffer_counter += 1;
356                    let label = format!("__buffer_{label}", label = context.buffer_counter);
357                    context.data_section.variables.push(Variable {
358                      name: label.clone(),
359                      type_: Type::Space,
360                      value: Value::Bytes(size),
361                    });
362
363                    context.text_section.statements.append(
364                      &mut [
365                        Statement::Instruction(Instruction::La(
366                          [
367                            InstructionArgument::Register(Register {
368                              name: "$a0".to_string(),
369                            }),
370                            InstructionArgument::Label(label),
371                          ]
372                          .into(),
373                        )),
374                        Statement::Instruction(Instruction::Li(
375                          [
376                            InstructionArgument::Register(Register {
377                              name: "$a1".to_string(),
378                            }),
379                            InstructionArgument::Immediate(size),
380                          ]
381                          .into(),
382                        )),
383                      ]
384                      .into(),
385                    );
386                  }
387                  _ => Err(format!("Function {} not found", function.name))?,
388                };
389
390                context.text_section.statements.append(
391                  &mut [
392                    // Perform the syscall
393                    Statement::Instruction(Instruction::Syscall),
394                    // Move the result of the syscall to the register
395                    Statement::Instruction(Instruction::Move(
396                      [
397                        InstructionArgument::Register(Register {
398                          name: register.clone(),
399                        }),
400                        InstructionArgument::Register(Register {
401                          name: "$v0".to_string(),
402                        }),
403                      ]
404                      .into(),
405                    )),
406                  ]
407                  .into(),
408                );
409              } else {
410                function_call
411                  .params
412                  .iter()
413                  .enumerate()
414                  .for_each(|(i, param)| {
415                    let register = match param {
416                      Operand::Identifier(ident) => {
417                        find_or_create_reg(&mut context.register_map, ident.clone())
418                      }
419                      Operand::LiteralStr(str) => {
420                        let register = new_register(&mut context.register_map);
421                        load_string(
422                          &mut context.text_section,
423                          &mut context.data_section,
424                          register.clone(),
425                          str.clone(),
426                        );
427
428                        register
429                      }
430                      Operand::LiteralBool(value) => {
431                        load_immediate_to_new_register(context, *value as i32)
432                      }
433                      Operand::LiteralI8(value) => {
434                        load_immediate_to_new_register(context, *value as i32)
435                      }
436                      Operand::LiteralI16(value) => {
437                        load_immediate_to_new_register(context, *value as i32)
438                      }
439                      Operand::LiteralI32(value) => {
440                        load_immediate_to_new_register(context, *value as i32)
441                      }
442                      Operand::LiteralI64(value) => {
443                        load_immediate_to_new_register(context, *value as i32)
444                      }
445                      Operand::LiteralU8(value) => {
446                        load_immediate_to_new_register(context, *value as i32)
447                      }
448                      Operand::LiteralU16(value) => {
449                        load_immediate_to_new_register(context, *value as i32)
450                      }
451                      Operand::LiteralU32(value) => {
452                        load_immediate_to_new_register(context, *value as i32)
453                      }
454                      Operand::LiteralU64(value) => {
455                        load_immediate_to_new_register(context, *value as i32)
456                      }
457                      Operand::LiteralF32(_) => todo!(),
458                      Operand::LiteralF64(_) => todo!(),
459                      Operand::Dereference(_) => todo!(),
460                    };
461
462                    context.text_section.statements.push(Statement::Instruction(
463                      Instruction::Move(
464                        [
465                          InstructionArgument::Register(Register {
466                            name: format!("$a{}", i),
467                          }),
468                          InstructionArgument::Register(Register { name: register }),
469                        ]
470                        .into(),
471                      ),
472                    ));
473                  });
474              }
475            }
476          }
477        }
478        CompassStatement::ConditionalJump {
479          condition,
480          label,
481          location,
482        } => match condition {
483          Expr::BinaryOperation(op) => match op {
484            BinaryOperation::Conditional {
485              ref lhs,
486              ref condition,
487              ref rhs,
488              operation_type,
489            } => {
490              if operation_type != VarType::Bool {
491                return Err("Conditional operations must be of type bool".to_string());
492              }
493
494              if is_register(&lhs) && is_register(&rhs) {
495                let lhs = lhs.as_identifier()?;
496                let rhs = rhs.as_identifier()?;
497
498                let lhs_register = context
499                  .register_map
500                  .get(lhs)
501                  .ok_or_else(|| format!("Register {} not found", lhs))?
502                  .clone();
503                let rhs_register = context
504                  .register_map
505                  .get(rhs)
506                  .ok_or_else(|| format!("Register {} not found", rhs))?
507                  .clone();
508
509                match condition {
510                  Condition::And => {
511                    context.conditional_counter += 1;
512                    let new_label = format!("__and_{label}", label = context.conditional_counter);
513
514                    context.text_section.statements.append(
515                      &mut [
516                        Statement::Instruction(Instruction::Beqz(
517                          [
518                            InstructionArgument::Register(Register {
519                              name: lhs_register.clone(),
520                            }),
521                            InstructionArgument::Label(new_label.clone()),
522                          ]
523                          .into(),
524                        )),
525                        Statement::Instruction(Instruction::Beqz(
526                          [
527                            InstructionArgument::Register(Register {
528                              name: rhs_register.clone(),
529                            }),
530                            InstructionArgument::Label(new_label.clone()),
531                          ]
532                          .into(),
533                        )),
534                        Statement::Instruction(Instruction::J(
535                          [InstructionArgument::Label(label.clone())].into(),
536                        )),
537                        Statement::Label(new_label),
538                      ]
539                      .into(),
540                    );
541                  }
542                  Condition::Or => {
543                    context.text_section.statements.append(
544                      &mut [
545                        Statement::Instruction(Instruction::Bnez(
546                          [
547                            InstructionArgument::Register(Register {
548                              name: lhs_register.clone(),
549                            }),
550                            InstructionArgument::Label(label.clone()),
551                          ]
552                          .into(),
553                        )),
554                        Statement::Instruction(Instruction::Bnez(
555                          [
556                            InstructionArgument::Register(Register {
557                              name: rhs_register.clone(),
558                            }),
559                            InstructionArgument::Label(label.clone()),
560                          ]
561                          .into(),
562                        )),
563                      ]
564                      .into(),
565                    );
566                  }
567                  _ => {
568                    let condition_instruction = match condition {
569                      Condition::LessThan => Instruction::Blt,
570                      Condition::GreaterThan => Instruction::Bgt,
571                      Condition::LessThanOrEqual => Instruction::Ble,
572                      Condition::GreaterThanOrEqual => Instruction::Bge,
573                      Condition::Equal => Instruction::Beq,
574                      Condition::NotEqual => Instruction::Bne,
575                      _ => Err(format!(
576                        "Invalid binary operation for conditional jump {op:?}",
577                      ))?,
578                    };
579
580                    context.text_section.statements.push(create_instruction!(
581                      condition_instruction,
582                      lhs_register,
583                      rhs_register,
584                      InstructionArgument::Label(label)
585                    ));
586                  }
587                };
588              } else if is_register(&lhs) && is_immediate(&rhs) {
589                let lhs = lhs.as_identifier()?;
590                let lhs_register = context
591                  .register_map
592                  .get(lhs)
593                  .ok_or_else(|| format!("Register {} not found", lhs))?
594                  .clone();
595                let rhs_value = rhs.as_immediate()?;
596
597                match condition {
598                  Condition::And => {
599                    context.conditional_counter += 1;
600                    let new_label = format!("__and_{label}", label = context.conditional_counter);
601
602                    context.text_section.statements.append(
603                      &mut [
604                        Statement::Instruction(Instruction::Beqz(
605                          [
606                            InstructionArgument::Register(Register {
607                              name: lhs_register.clone(),
608                            }),
609                            InstructionArgument::Label(new_label.clone()),
610                          ]
611                          .into(),
612                        )),
613                        Statement::Instruction(Instruction::Beqz(
614                          [
615                            InstructionArgument::Immediate(rhs_value),
616                            InstructionArgument::Label(new_label.clone()),
617                          ]
618                          .into(),
619                        )),
620                        Statement::Instruction(Instruction::J(
621                          [InstructionArgument::Label(label.clone())].into(),
622                        )),
623                        Statement::Label(new_label),
624                      ]
625                      .into(),
626                    );
627                  }
628                  Condition::Or => {
629                    context.text_section.statements.append(
630                      &mut [
631                        Statement::Instruction(Instruction::Bnez(
632                          [
633                            InstructionArgument::Register(Register {
634                              name: lhs_register.clone(),
635                            }),
636                            InstructionArgument::Label(label.clone()),
637                          ]
638                          .into(),
639                        )),
640                        Statement::Instruction(Instruction::Bnez(
641                          [
642                            InstructionArgument::Immediate(rhs_value),
643                            InstructionArgument::Label(label.clone()),
644                          ]
645                          .into(),
646                        )),
647                      ]
648                      .into(),
649                    );
650                  }
651                  _ => {
652                    let condition_instruction = match condition {
653                      Condition::LessThan => Instruction::Blt,
654                      Condition::GreaterThan => Instruction::Bgt,
655                      Condition::LessThanOrEqual => Instruction::Ble,
656                      Condition::GreaterThanOrEqual => Instruction::Bge,
657                      Condition::Equal => Instruction::Beq,
658                      Condition::NotEqual => Instruction::Bne,
659                      _ => Err(format!(
660                        "Invalid binary operation for conditional jump {op:?}",
661                      ))?,
662                    };
663
664                    context.text_section.statements.push(Statement::Instruction(
665                      condition_instruction(
666                        [
667                          InstructionArgument::Register(Register { name: lhs_register }),
668                          InstructionArgument::Immediate(rhs_value),
669                          InstructionArgument::Label(label),
670                        ]
671                        .into(),
672                      ),
673                    ));
674                  }
675                };
676              } else if is_immediate(&lhs) && is_immediate(&rhs) {
677                let lhs_value = lhs.as_immediate()?;
678                let rhs_value = rhs.as_immediate()?;
679
680                let lhs_register = new_register(&mut context.register_map);
681                load_immediate(&mut context.text_section, lhs_register.clone(), lhs_value);
682
683                match condition {
684                  Condition::And => {
685                    context.conditional_counter += 1;
686                    let new_label = format!("__and_{label}", label = context.conditional_counter);
687
688                    context.text_section.statements.append(
689                      &mut [
690                        Statement::Instruction(Instruction::Beqz(
691                          [
692                            InstructionArgument::Register(Register {
693                              name: lhs_register.clone(),
694                            }),
695                            InstructionArgument::Label(new_label.clone()),
696                          ]
697                          .into(),
698                        )),
699                        Statement::Instruction(Instruction::Beqz(
700                          [
701                            InstructionArgument::Immediate(rhs_value),
702                            InstructionArgument::Label(new_label.clone()),
703                          ]
704                          .into(),
705                        )),
706                        Statement::Instruction(Instruction::J(
707                          [InstructionArgument::Label(label.clone())].into(),
708                        )),
709                        Statement::Label(new_label),
710                      ]
711                      .into(),
712                    );
713                  }
714                  Condition::Or => {
715                    context.text_section.statements.append(
716                      &mut [
717                        Statement::Instruction(Instruction::Bnez(
718                          [
719                            InstructionArgument::Register(Register {
720                              name: lhs_register.clone(),
721                            }),
722                            InstructionArgument::Label(label.clone()),
723                          ]
724                          .into(),
725                        )),
726                        Statement::Instruction(Instruction::Bnez(
727                          [
728                            InstructionArgument::Immediate(rhs_value),
729                            InstructionArgument::Label(label.clone()),
730                          ]
731                          .into(),
732                        )),
733                      ]
734                      .into(),
735                    );
736                  }
737                  _ => {
738                    let condition_instruction = match condition {
739                      Condition::LessThan => Instruction::Blt,
740                      Condition::GreaterThan => Instruction::Bgt,
741                      Condition::LessThanOrEqual => Instruction::Ble,
742                      Condition::GreaterThanOrEqual => Instruction::Bge,
743                      Condition::Equal => Instruction::Beq,
744                      Condition::NotEqual => Instruction::Bne,
745                      _ => Err(format!(
746                        "Invalid binary operation for conditional jump {op:?}",
747                      ))?,
748                    };
749
750                    context.text_section.statements.push(Statement::Instruction(
751                      condition_instruction(
752                        [
753                          InstructionArgument::Register(Register { name: lhs_register }),
754                          InstructionArgument::Immediate(rhs_value),
755                          InstructionArgument::Label(label),
756                        ]
757                        .into(),
758                      ),
759                    ));
760                  }
761                }
762              }
763            }
764            _ => Err(format!(
765              "Invalid binary operation for conditional jump {op:?}",
766            ))?,
767          },
768          Expr::Operand(op) => match op {
769            Operand::Identifier(ident) => {
770              let register = context
771                .register_map
772                .get(&ident)
773                .ok_or_else(|| format!("Register {} not found", ident))?
774                .clone();
775
776              context
777                .text_section
778                .statements
779                .push(Statement::Instruction(Instruction::Beqz(
780                  [
781                    InstructionArgument::Register(Register {
782                      name: register.clone(),
783                    }),
784                    InstructionArgument::Label(label),
785                  ]
786                  .into(),
787                )));
788            }
789            Operand::LiteralBool(value) => match value {
790              true => context
791                .text_section
792                .statements
793                .push(Statement::Instruction(Instruction::J(
794                  [InstructionArgument::Label(label)].into(),
795                ))),
796              false => (),
797            },
798            _ => Err(format!("Invalid operand for conditional jump {}", op))?,
799          },
800          Expr::FunctionCall(_) => unreachable!(),
801        },
802        CompassStatement::UnconditionalJump { label, location } => {
803          context
804            .text_section
805            .statements
806            .push(Statement::Instruction(Instruction::J(
807              [InstructionArgument::Label(label)].into(),
808            )));
809        }
810        CompassStatement::Label { name, location } => {
811          if name == "main" {
812            return Err("Cannot use 'main' as a label name".to_string());
813          }
814
815          context.text_section.statements.push(Statement::Label(name))
816        }
817        CompassStatement::FunctionDefinition(function) => {
818          let name = format!("__{name}", name = function.name.clone());
819
820          context.function_map.insert(name.clone(), function.clone());
821
822          // Function declarations should be added before the `main: flow`
823          let statements = vec![Statement::Label(name)];
824
825          // TODO: Insert variables into the scope
826
827          let mut save_statements = context.text_section.statements.clone();
828          context.text_section.statements = statements;
829          context.scope_level += 1;
830          self.generate(function.body, context)?;
831          context.scope_level -= 1;
832
833          context
834            .text_section
835            .statements
836            .push(Statement::Instruction(Instruction::Jr(
837              [InstructionArgument::Register(Register {
838                name: "$ra".to_string(),
839              })]
840              .into(),
841            )));
842
843          context.text_section.statements.append(&mut save_statements);
844        }
845        CompassStatement::Store { at, from, location } => match (&at, &from) {
846          (Operand::Dereference(at), Operand::Identifier(from)) => {
847            let at_register = context
848              .register_map
849              .get(at)
850              .ok_or_else(|| format!("Register {} not found", at))?
851              .clone();
852            let from_register = context
853              .register_map
854              .get(from)
855              .ok_or_else(|| format!("Register {} not found", from))?
856              .clone();
857
858            context
859              .text_section
860              .statements
861              .push(Statement::Instruction(Instruction::Sw(
862                [
863                  InstructionArgument::Register(Register {
864                    name: from_register,
865                  }),
866                  InstructionArgument::Register(Register { name: at_register }),
867                ]
868                .into(),
869              )));
870          }
871          _ => {
872            Err(format!(
873              "Invalid operands for store operation {} and {}",
874              at, from
875            ))?;
876          }
877        },
878        CompassStatement::Call(FunctionCall {
879          name,
880          params,
881          return_type,
882          location,
883        }) => {
884          let function = context
885            .get_function(&name)
886            .ok_or_else(|| format!("Function {} not found", name))?;
887
888          if function.is_builtin {
889            match function.name.as_str() {
890              "write_string" => {
891                // Perform the write_string syscall (v0 = 4, a0 = string address)
892                context
893                  .text_section
894                  .statements
895                  .push(Statement::Instruction(Instruction::Li(
896                    [
897                      InstructionArgument::Register(Register {
898                        name: "$v0".to_string(),
899                      }),
900                      InstructionArgument::Immediate(4),
901                    ]
902                    .into(),
903                  )));
904
905                let string_register = match &params[0] {
906                  Operand::Identifier(ident) => context
907                    .register_map
908                    .get(ident)
909                    .ok_or_else(|| format!("Register {} not found", ident))?
910                    .clone(),
911                  Operand::LiteralStr(str) => {
912                    let register = new_register(&mut context.register_map);
913                    load_string(
914                      &mut context.text_section,
915                      &mut context.data_section,
916                      register.clone(),
917                      str.clone(),
918                    );
919
920                    register
921                  }
922                  _ => todo!(),
923                };
924
925                context
926                  .text_section
927                  .statements
928                  .push(Statement::Instruction(Instruction::Move(
929                    [
930                      InstructionArgument::Register(Register {
931                        name: "$a0".to_string(),
932                      }),
933                      InstructionArgument::Register(Register {
934                        name: string_register,
935                      }),
936                    ]
937                    .into(),
938                  )));
939
940                context
941                  .text_section
942                  .statements
943                  .push(Statement::Instruction(Instruction::Syscall));
944              }
945              "write_int" => {
946                // Perform the write_int syscall (v0 = 1, a0 = integer value)
947                context
948                  .text_section
949                  .statements
950                  .push(Statement::Instruction(Instruction::Li(
951                    [
952                      InstructionArgument::Register(Register {
953                        name: "$v0".to_string(),
954                      }),
955                      InstructionArgument::Immediate(1),
956                    ]
957                    .into(),
958                  )));
959
960                let int_register = match &params[0] {
961                  Operand::Identifier(ident) => context
962                    .register_map
963                    .get(ident)
964                    .ok_or_else(|| format!("Register {} not found", ident))?
965                    .clone(),
966                  Operand::LiteralI8(val) => load_immediate_to_new_register(context, *val as i32),
967                  Operand::LiteralI16(val) => load_immediate_to_new_register(context, *val as i32),
968                  Operand::LiteralI32(val) => load_immediate_to_new_register(context, *val as i32),
969                  Operand::LiteralI64(val) => load_immediate_to_new_register(context, *val as i32),
970                  Operand::LiteralU8(val) => load_immediate_to_new_register(context, *val as i32),
971                  Operand::LiteralU16(val) => load_immediate_to_new_register(context, *val as i32),
972                  Operand::LiteralU32(val) => load_immediate_to_new_register(context, *val as i32),
973                  Operand::LiteralU64(val) => load_immediate_to_new_register(context, *val as i32),
974                  _ => todo!(),
975                };
976
977                context.text_section.statements.append(
978                  &mut [
979                    Statement::Instruction(Instruction::Move(
980                      [
981                        InstructionArgument::Register(Register {
982                          name: "$a0".to_string(),
983                        }),
984                        InstructionArgument::Register(Register { name: int_register }),
985                      ]
986                      .into(),
987                    )),
988                    Statement::Instruction(Instruction::Syscall),
989                  ]
990                  .into(),
991                );
992              }
993              _ => todo!(),
994            }
995          } else {
996            for (i, param) in params.iter().enumerate() {
997              let register = match param {
998                Operand::Identifier(ident) => {
999                  find_or_create_reg(&mut context.register_map, ident.clone())
1000                }
1001                Operand::LiteralStr(str) => {
1002                  let register = new_register(&mut context.register_map);
1003                  load_string(
1004                    &mut context.text_section,
1005                    &mut context.data_section,
1006                    register.clone(),
1007                    str.clone(),
1008                  );
1009
1010                  register
1011                }
1012                Operand::LiteralBool(value) => {
1013                  load_immediate_to_new_register(context, *value as i32)
1014                }
1015                Operand::LiteralI8(value) => load_immediate_to_new_register(context, *value as i32),
1016                Operand::LiteralI16(value) => {
1017                  load_immediate_to_new_register(context, *value as i32)
1018                }
1019                Operand::LiteralI32(value) => {
1020                  load_immediate_to_new_register(context, *value as i32)
1021                }
1022                Operand::LiteralI64(value) => {
1023                  load_immediate_to_new_register(context, *value as i32)
1024                }
1025                Operand::LiteralU8(value) => load_immediate_to_new_register(context, *value as i32),
1026                Operand::LiteralU16(value) => {
1027                  load_immediate_to_new_register(context, *value as i32)
1028                }
1029                Operand::LiteralU32(value) => {
1030                  load_immediate_to_new_register(context, *value as i32)
1031                }
1032                Operand::LiteralU64(value) => {
1033                  load_immediate_to_new_register(context, *value as i32)
1034                }
1035                Operand::LiteralF32(_) => todo!(),
1036                Operand::LiteralF64(_) => todo!(),
1037                Operand::Dereference(_) => todo!(),
1038              };
1039
1040              context
1041                .text_section
1042                .statements
1043                .push(Statement::Instruction(Instruction::Move(
1044                  [
1045                    InstructionArgument::Register(Register {
1046                      name: format!("$a{}", i),
1047                    }),
1048                    InstructionArgument::Register(Register { name: register }),
1049                  ]
1050                  .into(),
1051                )));
1052            }
1053
1054            context
1055              .text_section
1056              .statements
1057              .push(Statement::Instruction(Instruction::Jal(
1058                [InstructionArgument::Label(name)].into(),
1059              )));
1060          }
1061        }
1062        CompassStatement::NoOperation => {}
1063      }
1064    }
1065
1066    // Add halt
1067    context
1068      .text_section
1069      .statements
1070      .push(Statement::Instruction(Instruction::Halt));
1071
1072    let program = Program {
1073      data_section: context.data_section.clone(),
1074      text_section: context.text_section.clone(),
1075    };
1076
1077    Ok(program.to_string())
1078  }
1079}
1080
1081fn load_immediate_to_new_register(context: &mut Context, value: i32) -> String {
1082  let register = new_register(&mut context.register_map);
1083  load_immediate(&mut context.text_section, register.clone(), value);
1084  register
1085}
1086
1087fn find_or_create_reg(register_map: &mut HashMap<String, String>, name: String) -> String {
1088  if let Some(register) = register_map.get(&name) {
1089    register.clone()
1090  } else {
1091    let register = format!("$t{}", register_map.len());
1092    register_map.insert(name.clone(), register.clone());
1093    register
1094  }
1095}
1096
1097fn new_register(register_map: &mut HashMap<String, String>) -> String {
1098  let register = format!("$t{}", register_map.len());
1099  register_map.insert(register.clone(), register.clone());
1100  register
1101}
1102
1103fn is_register(value: &crate::ast::Operand) -> bool {
1104  match value {
1105    Operand::Identifier(_) => true,
1106    Operand::Dereference(_) => true,
1107    _ => false,
1108  }
1109}
1110
1111fn is_immediate(value: &crate::ast::Operand) -> bool {
1112  match value {
1113    Operand::Identifier(_) => false,
1114    Operand::Dereference(_) => false,
1115    _ => true,
1116  }
1117}
1118
1119fn load_immediate(text_section: &mut TextSection, register: String, value: i32) {
1120  text_section
1121    .statements
1122    .push(Statement::Instruction(Instruction::Li(
1123      [
1124        InstructionArgument::Register(Register { name: register }),
1125        InstructionArgument::Immediate(value),
1126      ]
1127      .into(),
1128    )));
1129}
1130
1131fn load_string(
1132  text_section: &mut TextSection,
1133  data_section: &mut DataSection,
1134  register: String,
1135  value: String,
1136) {
1137  // Search for an existing string label
1138  let mut label = String::new();
1139  let mut found = false;
1140  for (i, variable) in data_section.variables.iter().enumerate() {
1141    if let Value::String(existing_value) = &variable.value {
1142      if existing_value == &value {
1143        label = format!("str_{}", i);
1144        found = true;
1145        break;
1146      }
1147    }
1148  }
1149
1150  // If the string is not found, add it to the data section
1151  if !found {
1152    label = format!("str_{}", data_section.variables.len());
1153    data_section.variables.push(Variable {
1154      name: label.clone(),
1155      type_: Type::Asciiz,
1156      value: Value::String(value),
1157    });
1158  }
1159
1160  // Load the address of the string into the register
1161  text_section
1162    .statements
1163    .push(Statement::Instruction(Instruction::La(
1164      [
1165        InstructionArgument::Register(Register { name: register }),
1166        InstructionArgument::Label(label),
1167      ]
1168      .into(),
1169    )));
1170}