N2T/jack_compiler/src/code_writer.rs

use crate::symbol_table::get_empty_symbol_table;
use crate::symbol_table::SymbolTable;
use crate::tokenizer::Keyword;
use crate::tokenizer::Symbol;
use crate::tokenizer::Token;
use std::fs;

#[derive(Debug)]
pub struct Writer {
	pub table: SymbolTable,
	pub code: String,
	pub class_name: String,
	pub subroutine_name: String,
	pub in_constructor: bool,
	pub in_method: bool,
	if_label_count: usize,
	while_label_count: usize,
}

pub fn get_code_writer() -> Writer {
	Writer {
		table: get_empty_symbol_table(),
		code: String::new(),
		class_name: String::new(),
		subroutine_name: String::new(),
		in_constructor: false,
		in_method: false,
		if_label_count: 0,
		while_label_count: 0,
	}
}

fn segment_to_string(segment: Keyword) -> String {
	use crate::tokenizer::Keyword::*;
	match segment {
		Constant => String::from("constant"),
		Argument => String::from("argument"),
		Local => String::from("local"),
		Temp => String::from("temp"),
		Field => String::from("this"),
		Static => String::from("static"),
		That => String::from("that"),
		Pointer => String::from("pointer"),
		_ => panic!("Unexpected segment {:?}", segment),
	}
}

impl Writer {
	pub fn write_to_file(&self, file: &String) {
		fs::write(file, self.code.to_string()).expect("Unable to write file");
	}

	pub fn start_subroutine(&mut self) {
		self.table.remove_subroutine_symbols();
		self.in_constructor = false;
		self.in_method = false;
		self.if_label_count = 0;
		self.while_label_count = 0;
	}

	pub fn start_class(&mut self) {
		self.table = get_empty_symbol_table();
	}

	pub fn define_symbol(&mut self, name: String, symbol_type: Token, kind: Keyword) {
		self.table.define_symbol(name, symbol_type, kind);
	}

	pub fn write_push(&mut self, segment: Keyword, index: usize) {
		let segment = segment_to_string(segment);
		let s = format!("push {} {}\n", segment, index);
		self.code.push_str(&s);
	}

	pub fn write_pop(&mut self, segment: Keyword, index: usize) {
		let segment = segment_to_string(segment);
		let s = format!("pop {} {}\n", segment, index);
		self.code.push_str(&s);
	}

	pub fn write_arithmetic(&mut self, op: Symbol) {
		use crate::tokenizer::Symbol::*;
		let s = match op {
			Plus => "add\n",
			Minus => "sub\n",
			Mul => "call Math.multiply 2\n",
			Div => "call Math.divide 2\n",
			ExclusiveAnd => "and\n",
			ExclusiveOr => "or\n",
			Smaller => "lt\n",
			Greater => "gt\n",
			Equal => "eq\n",
			UnaryMinus => "neg\n",
			Not => "not\n",
			_ => panic!("Unsupported operator {:?}.", op),
		};
		self.code.push_str(&s);
	}

	pub fn write_function(&mut self) {
		let n_locals = self.table.get_count(Keyword::Local);
		let s = format!(
			"function {}.{} {}\n",
			self.class_name, self.subroutine_name, n_locals
		);
		self.code.push_str(&s);

		if self.in_constructor {
			// Allocate class memory and initialize This.
			let n_fields = self.table.get_count(Keyword::Field);
			self.write_push(Keyword::Constant, n_fields);
			self.write_call(&"Memory".to_string(), &"alloc".to_string(), 1);
			self.write_pop(Keyword::Pointer, 0);
		} else if self.in_method {
			self.write_push(Keyword::Argument, 0);
			self.write_pop(Keyword::Pointer, 0);
		}
	}

	pub fn write_call(&mut self, class_name: &String, subroutine_name: &String, n_args: usize) {
		let s = format!("call {}.{} {}\n", class_name, subroutine_name, n_args);
		self.code.push_str(&s);
	}

	pub fn write_return(&mut self) {
		self.code.push_str("return\n");
	}

	pub fn write_var_assignment(&mut self, var_name: &String) {
		let index = self.table.index_of(var_name);
		let symbol_type = self.table.kind_of(var_name);
		// if symbol_type == Keyword::Static {
		// 	println!("{:?}", self.table);
		// 	panic!("assignment to static not supported, yet");
		// }
		self.write_pop(symbol_type, index);
	}

	pub fn write_var_read(&mut self, var_name: &String) {
		let index = self.table.index_of(var_name);
		let symbol_type = self.table.kind_of(var_name);
		// if symbol_type == Keyword::Static {
		// 	panic!("read from static not supported, yet");
		// }
		self.write_push(symbol_type, index);
	}

	pub fn write_label(&mut self, label_name: &String) {
		let s = format!("label {}\n", label_name);
		self.code.push_str(&s);
	}

	pub fn write_if_goto(&mut self, label_name: &String) {
		let s = format!("if-goto {}\n", label_name);
		self.code.push_str(&s);
	}

	pub fn write_goto(&mut self, label_name: &String) {
		let s = format!("goto {}\n", label_name);
		self.code.push_str(&s);
	}

	pub fn get_if_labels(&mut self) -> (String, String, String) {
		let l1 = format!("IF_FALSE{}", self.if_label_count);
		let l2 = format!("IF_TRUE{}", self.if_label_count);
		let l3 = format!("IF_END{}", self.if_label_count);
		self.if_label_count += 1;
		return (l1, l2, l3);
	}

	pub fn get_while_labels(&mut self) -> (String, String, String) {
		let l1 = format!("WHILE_EXP{}", self.while_label_count);
		let l2 = format!("WHILE_START{}", self.while_label_count);
		let l3 = format!("WHILE_END{}", self.while_label_count);
		self.while_label_count += 1;
		return (l1, l2, l3);
	}
}