repo_description: A portable programming language with a compact intermediate representation
repo_topics: c compiler compiler-design golang
code: use crate::{
      target::{Target, C},
      Identifier, StringLiteral,
  };
  use std::{
      collections::BTreeMap,
      fmt::{Debug, Display, Error, Formatter},
      fs::read_to_string,
      path::PathBuf,
  };
  
  #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
  pub enum AsmError {
      NonExistantExternFile(String),
      VariableNotDefined(Identifier),
      FunctionNotDefined(Identifier),
      NoEntryPoint,
  }
  
  impl Display for AsmError {
      fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
          match self {
              Self::NonExistantExternFile(filename) => {
                  write!(f, "could not find foreign file '{}'", filename)
              }
              Self::FunctionNotDefined(name) => write!(f, "function '{}' is not defined", name),
              Self::VariableNotDefined(name) => write!(f, "variable '{}' is not defined", name),
              Self::NoEntryPoint => write!(f, "no entry point defined"),
          }
      }
  }
  
  #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
  pub struct AsmType {
      ptr_level: i32,
      size: i32,
  }
  
  impl AsmType {
      pub fn new(size: i32) -> Self {
          Self { ptr_level: 0, size }
      }
  
      pub fn ch() -> Self {
          Self::new(1)
      }
      pub fn float() -> Self {
          Self::new(1)
      }
      pub fn void() -> Self {
          Self::new(0)
      }
  
      pub fn refer(&self) -> Self {
          let mut copy = *self;
          copy.ptr_level += 1;
          copy
      }
  
      pub fn deref(&self) -> Option<Self> {
          if self.ptr_level > 0 {
              let mut copy = *self;
              copy.ptr_level += 1;
              Some(copy)
          } else {
              None
          }
      }
  
      pub fn get_size(&self) -> i32 {
          if self.ptr_level > 0 {
              1
          } else {
              self.size
          }
      }
  }
  
  impl Debug for AsmType {
      fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
          for _ in 0..self.ptr_level {
              write!(f, "&")?;
          }
          write!(f, "{}", self.size)
      }
  }
  
  #[derive(Clone, Debug)]
  pub struct AsmProgram {
      externs: Vec<PathBuf>,
      funcs: Vec<AsmFunction>,
      memory_size: i32,
  }
  
  impl AsmProgram {
      const ENTRY_POINT: &'static str = "main";
  
      pub fn new(externs: Vec<PathBuf>, funcs: Vec<AsmFunction>, memory_size: i32) -> Self {
          Self {
              externs,
              funcs,
              memory_size,
          }
      }
  
      pub fn assemble(&self, target: &impl Target) -> Result<String, AsmError> {
          // Set up the output code
          let mut result = String::new();
  
          // Iterate over the external files to include
          for filename in &self.externs {
              // Find them in the current working directory
              if let Ok(contents) = read_to_string(filename.clone()) {
                  // Add the contents of the file to the result
                  result += &contents
              } else {
                  // If the file doesn't exist, throw an error
                  if let Ok(name) = filename.clone().into_os_string().into_string() {
                      return Err(AsmError::NonExistantExternFile(name));
                  } else {
                      return Err(AsmError::NonExistantExternFile(String::from("")));
                  }
              }
          }
  
          // Store the IDs of each function
          let mut func_ids = BTreeMap::new();
          // The number of cells to preemptively allocate on the stack before the program starts
          let mut global_scope_size = 0;
          for (id, func) in self.funcs.iter().enumerate() {
              // Store the function's ID
              func_ids.insert(func.name.clone(), id as i32);
              // Add the function header to the output code
              result += &target.fn_header(AsmFunction::get_assembled_name(id as i32));
          }
  
          // It is very important that the entry point is assembled last.
          // This is because of the way things are allocated on the stack.
          let mut entry_point = None;
          for func in &self.funcs {
              // Compile the function
              if !func.is_entry_point() {
                  result += &func.assemble(&func_ids, &mut global_scope_size, target)?;
              } else {
                  // Store the entry point for use later
                  // This has the side effect of ignoring multiple definitions
                  // of the `main` function, and just using the last one defined.
                  entry_point = Some(func);
              }
          }
  
          if let Some(func) = entry_point {
              if let Some(main_id) = func_ids.get(Self::ENTRY_POINT) {
                  // Assemble the entry point code
                  result += &func.assemble(&func_ids, &mut global_scope_size, target)?;
  
                  // Call the entry point
                  result += &target.begin_entry_point(global_scope_size, self.memory_size);
                  result += &target.call_fn(AsmFunction::get_assembled_name(*main_id));
                  result += &target.end_entry_point();
  
                  Ok(result)
              } else {
                  Err(AsmError::NoEntryPoint)
              }
          } else {
              Err(AsmError::NoEntryPoint)
          }
      }
  }
  
  #[derive(Clone, Debug)]
  pub struct AsmFunction {
      name: Identifier,
      args: Vec<(Identifier, AsmType)>,
      return_type: AsmType,
      body: Vec<AsmStatement>,
  }
  
  impl AsmFunction {
      pub fn new(
          name: Identifier,
          args: Vec<(Identifier, AsmType)>,
          return_type: AsmType,
          body: Vec<AsmStatement>,
      ) -> Self {
          Self {
              name,
              args,
              return_type,
              body,
          }
      }
  
      fn is_entry_point(&self) -> bool {
          self.name == AsmProgram::ENTRY_POINT
      }
  
      /// Use the function's ID to get the output code's name of the function.
      /// An ID is used to prevent invalid output code function names, or names
      /// that clash with standard library names such as "printf" or "malloc".
      fn get_assembled_name(id: i32) -> String {
          format!("fn{}", id)
      }
  
      fn assemble(
          &self,
          func_ids: &BTreeMap<String, i32>,
          global_scope_size: &mut i32,
          target: &impl Target,
      ) -> Result<String, AsmError> {
          let mut result = String::new();
          let mut arg_size = 0;
  
          // The local scope size starts at one. This is VERY important.
          // The reason the local scope size starts at one is to make room for
          // the virtual machine's base pointer on the stack before the stack
          // frame actually begins.
          let mut local_scope_size = 1;
  
          // Store the variables's addresses and types in the scope
          let mut vars = BTreeMap::new();
          for (arg_name, arg_type) in &self.args {
              // Add together the total size of all the arguments supplied to the function
              arg_size += arg_type.get_size();
  
              // Define each argument of the function
              result += &AsmStatement::Define(arg_name.clone(), *arg_type).assemble(
                  func_ids,
                  &mut vars,
                  global_scope_size,
                  &mut local_scope_size,
                  target,
              )?;
              result += &AsmStatement::Assign(*arg_type).assemble(
                  func_ids,
                  &mut vars,
                  global_scope_size,
                  &mut local_scope_size,
                  target,
              )?;
          }
  
          for stmt in &self.body {
              // Assemble each statement in the function body
              result += &stmt.assemble(
                  func_ids,
                  &mut vars,
                  global_scope_size,
                  &mut local_scope_size,
                  target,
              )?;
          }
  
          let start = target.establish_stack_frame(arg_size, local_scope_size);
          result += &target.end_stack_frame(self.return_type.get_size(), local_scope_size);
  
          // Write the function as output code
          if let Some(id) = func_ids.get(&self.name) {
              Ok(target.fn_definition(Self::get_assembled_name(*id), start + &result))
          } else {
              Err(AsmError::FunctionNotDefined(self.name.clone()))
          }
      }
  }
  
  #[derive(Clone, Debug)]
  pub enum AsmStatement {
      For(Vec<Self>, Vec<Self>, Vec<Self>, Vec<Self>),
      Define(Identifier, AsmType),
      Assign(AsmType),
      Expression(Vec<AsmExpression>),
  }
  
  impl AsmStatement {
      fn assemble(
          &self,
          func_ids: &BTreeMap<String, i32>,
          vars: &mut BTreeMap<String, (i32, AsmType)>,
          global_scope_size: &mut i32,
          local_scope_size: &mut i32,
          target: &impl Target,
      ) -> Result<String, AsmError> {
          Ok(match self {
              // Define a variable on the stack
              Self::Define(name, data_type) => {
                  let address = *local_scope_size;
                  // Add the variable's location and type to the scope
                  vars.insert(name.clone(), (address, *data_type));
                  // Increment the size of the program's variables
                  *local_scope_size += data_type.get_size();
                  // Push the address of the new variable onto the stack
                  target.push(address as f64) + &target.load_base_ptr() + &target.add()
              }
              // Pop an address off of the stack, pop an item of size `data_type`
              // off of the stack, and store the item at the address
              Self::Assign(data_type) => target.store(data_type.get_size()),
              Self::For(pre, cond, post, body) => {
                  let mut result = String::new();
                  // Run the code that preps the for loop
                  for stmt in pre {
                      result += &stmt.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  // Check the condition of the for loop
                  for expr in cond {
                      result += &expr.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  // Begin the loop body
                  result += &target.begin_while();
                  // Run the body of the loop
                  for stmt in body {
                      result += &stmt.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  // Run the code that procedes the body of the loop
                  for stmt in post {
                      result += &stmt.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  // Check the condition again
                  for expr in cond {
                      result += &expr.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  // End the loop body
                  result + &target.end_while()
              }
  
              Self::Expression(exprs) => {
                  let mut result = String::new();
                  for expr in exprs {
                      result += &expr.assemble(
                          func_ids,
                          vars,
                          global_scope_size,
                          local_scope_size,
                          target,
                      )?;
                  }
                  result
              }
          })
      }
  }
  
  #[derive(Clone, Debug)]
  pub enum AsmExpression {
      String(StringLiteral),
      Character(char),
      Float(f64),
      Void,
  
      ForeignCall(Identifier),
  
      Variable(Identifier),
      Call(Identifier),
      Refer(Identifier),
      Deref(i32),
  
      Alloc,
      Free,
  
      Divide,
      Multiply,
      Subtract,
      Add,
      Sign,
  }
  
  impl AsmExpression {
      fn assemble(
          &self,
          func_ids: &BTreeMap<String, i32>,
          vars: &mut BTreeMap<String, (i32, AsmType)>,
          global_scope_size: &mut i32,
          local_scope_size: &mut i32,
          target: &impl Target,
      ) -> Result<String, AsmError> {
          Ok(match self {
              Self::String(s) => {
                  // The address of the string is at the current first
                  // empty spot on the stack.
                  let address = *global_scope_size;
                  // The size of the string is the length of the characters,
                  // plus 1 for the zero terminated character.
                  let size = s.len() as i32 + 1;
  
                  // Push each character of the string onto the stack
                  let mut result = String::new();
                  for ch in s.chars() {
                      result += &target.push(ch as u8 as f64);
                  }
                  // Push the zero terminated character
                  result += &target.push(0.0);
  
                  // Store the characters at the address of the string,
                  // and push the address onto the stack.
                  result += &(target.push(address as f64)
                      + &target.store(size)
                      + &target.push(address as f64));
  
                  // Increment the amount of data stored on the stack
                  *global_scope_size += size;
                  result
              }
              // Push a character onto the stack
              Self::Character(ch) => target.push(*ch as u8 as f64),
              // Push a float onto the stack
              Self::Float(n) => target.push(*n),
              // Void expressions are a No-Op
              Self::Void => String::new(),
  
              // Load a variable onto the stack with a given type
              Self::Variable(name) => {
                  // Get the address of the variable on the stack
                  // and the type of the variable
                  if let Some((address, data_type)) = vars.get(name) {
                      // Push the address and load the data at that address
                      target.push(*address as f64)
                          + &target.load_base_ptr()
                          + &target.add()
                          + &target.load(data_type.get_size())
                  } else {
                      return Err(AsmError::VariableNotDefined(name.clone()));
                  }
              }
  
              // Call a function
              Self::Call(fn_name) => {
                  if let Some(fn_id) = func_ids.get(fn_name) {
                      target.call_fn(AsmFunction::get_assembled_name(*fn_id))
                  } else {
                      return Err(AsmError::FunctionNotDefined(fn_name.clone()));
                  }
              }
  
              // Call a foreign function
              Self::ForeignCall(fn_name) => target.call_foreign_fn(fn_name.clone()),
  
              // Allocate data on the heap
              Self::Alloc => target.allocate(),
              // Free data on the heap
              Self::Free => target.free(),
              // Get the address of a variable on the stack
              Self::Refer(name) => {
                  if let Some((addr, _)) = vars.get(name) {
                      target.push(*addr as f64) + &target.load_base_ptr() + &target.add()
                  } else {
                      return Err(AsmError::VariableNotDefined(name.clone()));
                  }
              }
              // Dereference an address
              Self::Deref(size) => target.load(*size),
  
              // Get the absolute value of a number on the stack
              Self::Sign => target.sign(),
              // Add two numbers on the stack
              Self::Add => target.add(),
              // Subtract two numbers on the stack
              Self::Subtract => target.subtract(),
              // Multiply two numbers on the stack
              Self::Multiply => target.multiply(),
              // Divide two numbers on the stack
              Self::Divide => target.divide(),
          })
      }
  }
  