DSA/tooling/dscs_compiler/src/backend/ir_lowering_codegen.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use clap::builder::Resettable::Value;
use dsa_tasks::ProcessAgent;
use crate::backend::scope::Allocator;
use crate::model::{Block, Call, CompilerError, Declaration, Expression, Program, Statement, TypeId, Variable};

use crate::backend::intermediate_rep as IR;
use crate::backend::intermediate_rep::{BasicBlock, Label, Terminator, VirtualReg};

pub struct IrLoweringCodeGenerator {
	ast: Program,

	functions: Vec<IR::Function>
}

impl IrLoweringCodeGenerator {
	pub fn new(ast: Program) -> Self {
		Self { ast, functions: vec![] }
	}

	pub fn run(&mut self, ctx: &ProcessAgent) -> Result<IR::IRProgram, CompilerError> {
		for block in self.ast.clone().declarations {
			match block {
				Declaration::Variable { ..} => todo!(),
				Declaration::Dependency(_) => todo!(),
				Declaration::Function { name, ..} => todo!(),
				Declaration::Struct { .. } => todo!(),
			}
		}

		for block in self.ast.clone().declarations {
			self.generate_block(&block)?;
		}

		self.output(ctx)
	}

	pub fn output(&mut self, ctx: &ProcessAgent) -> Result<IR::IRProgram, CompilerError> {
		todo!()
	}

	pub fn generate_block(&mut self, decl: &Declaration) -> Result<(), CompilerError> {
		match decl {
			Declaration::Function { name, params, body, return_type } => {
				let function = self.generate_function(name, params, body, return_type)?;
				self.functions.push(function);
			},
			Declaration::Struct { .. } => todo!("no codegen required for structs yet"),
			Declaration::Dependency(_) => todo!("no codegen required for dependencies yet"),
			Declaration::Variable { .. } => todo!("codegen for statics and constants"),
		}

		Ok(())
	}

	pub fn generate_function(
		&mut self, function_name: &str,
		params: &[Variable],
		body: &Block,
		return_type: &TypeId
	) -> Result<IR::Function, CompilerError> {
		let mut context = FunctionContext::new(Label::new(function_name.to_string()));

		let params = (0..params.len())
			.map(|_| context.new_vreg()).collect();

		self.collapse_block(body, &mut context)?;

		let blocks = context.terminate(IR::Terminator::ReturnVoid);

		Ok(IR::Function {
			name: function_name.to_string(), params, blocks,
		})
	}

	fn collapse_block(
		&self,
		block: &Block,
		ctx: &mut FunctionContext,
	) -> Result<(), CompilerError> {
		for statement in block {
			self.collapse_statement(statement, ctx)?;
		}
		Ok(())
	}

	fn collapse_statement(
		&self,
		stmt: &Statement,
		ctx: &mut FunctionContext,
	) -> Result<(), CompilerError> {
		match stmt {
			Statement::Block(block) => {
				self.collapse_block(block, ctx)
			}

			Statement::If { condition, then_stmt, else_stmt } => {
				/* if statement expansion. produces the following layout
					reg_cond = ... expansion of {condition} ...
					if reg_cond goto then_{id} else goto else_{id}
				then_{id}:
					... expansion of {then_stmt} ...
					goto end_{id}
				else_{id}:
					... expansion of {else_stmt} ...
					goto end_{id}
				end_{id}:
					... not our responsibility ...
				*/
				let id = Label::get_new_id();
				let then_label = Label::new_with_id("then", id);
				let else_label = Label::new_with_id("else", id);
				let end_label = Label::new_with_id("endif", id);

				let eval_result = self.collapse_expression(condition, ctx)?;

				ctx.next_block(
					Terminator::ConditionalJump {
						if_flag: eval_result,
						then_goto: then_label.clone(),
						else_goto: else_label.clone(),
					},
					then_label.clone()
				);

				self.collapse_block(then_stmt, ctx)?;

				ctx.next_block(
					Terminator::Jump(end_label.clone()),
					else_label.clone()
				);

				self.collapse_block(else_stmt, ctx)?;

				ctx.next_block(
					Terminator::Jump(end_label.clone()),
					end_label.clone()
				);

				Ok(())
			}

			Statement::Defer(call) => {
				// defers stmt. all currently in scope defers are called in FILO order on return
				ctx.defer(call.clone());
				Ok(())
			},

			Statement::Return(expr) => {
				// return stmt. calls relevant defer methods and expands the body of the return statement.

				let result = if let Some(value) = expr {
					self.collapse_expression(value, ctx)?
				} else { VirtualReg::void() };

				// run cleanup logic (defers)
				for call in ctx.get_deferred_calls() {
					self.collapse_call(&call, ctx)?;
				}

				ctx.next_block(Terminator::Return(result), Label::new("unreachable"));

				Ok(())
			}

			Statement::Expression { expr } => {
				// generate the code for the expression but discard the result
				// let the optimiser remove it if it provably has no side effects.

				let _ = self.collapse_expression(expr, ctx)?;
				Ok(())
			}

			Statement::Continue => {
				// continue stmt. goes to the 'start' block of the loop

				let Some(scope) = ctx.current_scope() else {
					return Err(CompilerError::generic("Invalid 'Continue' statement (not inside any loop)"));
				};

				ctx.next_block(Terminator::Jump(scope.start), Label::new("unreachable"));
				Ok(())
			},

			Statement::Break => {
				// break stmt. jumps to the 'end' block of the loop

				let Some(scope) = ctx.current_scope() else {
					return Err(CompilerError::generic("Invalid 'Continue' statement (not inside any loop)"));
				};

				ctx.next_block(Terminator::Jump(scope.end), Label::new("unreachable"));
				Ok(())
			},

			Statement::While { condition, body } => {
				/* while loop. produces the following layout
					goto while_check_{id}
				while_check_{id}:
					if {condition} goto while_body_{id} else goto while_end_{id}
				while_body_{id}:
					... expansion of {body} ...
					goto while_check_{id}
				while_end_{id}:
					... not our responsibility ...
				*/

				let id = Label::get_new_id();
				let check_label = Label::new_with_id("while_check", id);
				let body_label = Label::new_with_id("while_body", id);
				let end_label =	Label::new_with_id("while_end", id);

				ctx.next_block(Terminator::Jump(check_label.clone()), check_label.clone());

				let result = self.collapse_expression(condition, ctx)?;

				ctx.next_block(
					Terminator::ConditionalJump {
						if_flag: result,
						then_goto: body_label.clone(),
						else_goto: end_label.clone()
					},
					body_label.clone(),
				);

				self.collapse_block(body, ctx)?;

				ctx.next_block(
					Terminator::Jump(check_label.clone()),
					end_label.clone()
				);

				Ok(())
			}

			Statement::Loop(block) => {
				/* indefinite loop. produces the following layout
					goto start_loop_{id}
				start_loop_{id}:
					goto start_loop_{id}
				while_end_{id}:
					... not our responsibility ...
				*/

				let id = Label::get_new_id();
				let start_label = Label::new_with_id("start_loop", id);
				let end_label = Label::new_with_id("end_loop", id);

				ctx.next_block(Terminator::Jump(start_label.clone()), start_label.clone());

				ctx.enter_scope(Scope::new(start_label.clone(), end_label.clone()));

				self.collapse_block(block, ctx)?;

				ctx.exit_scope();

				ctx.next_block(Terminator::Jump(start_label.clone()), end_label);

				Ok(())
			},
			_ => todo!()
		}
	}

	fn collapse_call(
		&self,
		call: &Call,
		ctx: &mut FunctionContext
	) -> Result<VirtualReg, CompilerError> {
		let reg = ctx.new_vreg();

		let args = call.args.iter()
			.map(|arg| self.collapse_expression(&arg, ctx))
			.collect::<Result<Vec<VirtualReg>, _>>()?;

		ctx.push_stmt(IR::Statement::Call {
			return_value: reg,
			func: call.name.to_string(),
			args
		});

		Ok(reg)
	}

	fn collapse_expression(
		&self,
		expr: &Expression,
		ctx: &mut FunctionContext,
	) -> Result<VirtualReg, CompilerError> {
		todo!()
	}
}

struct BasicBlockBuilder {
	label: Label,
	statements: Vec<IR::Statement>
}

impl BasicBlockBuilder {
	fn new(label: Label) -> Self {
		Self {
			label,
			statements: vec![]
		}
	}

	pub fn push(&mut self, statement: IR::Statement) {
		self.statements.push(statement);
	}

	pub fn next_block(&mut self, terminator: Terminator, label: Label) -> BasicBlock {
		BasicBlock {
			label: std::mem::replace(&mut self.label, label),
			statements: std::mem::take(&mut self.statements),
			terminator,
		}
	}

	pub fn terminate_block(self, terminator: Terminator) -> BasicBlock {
		BasicBlock {
			label: self.label,
			statements: self.statements,
			terminator,
		}
	}
}

pub struct Scope {
	start: Label,
	end: Label,

	defers: Vec<Call>,
	locals: Vec<(VirtualReg, Variable)>,
}

impl Scope {
	pub fn new(start: Label, end: Label) -> Self {
		Self { start, end, defers: vec![], locals: vec![] }
	}
}

struct FunctionContext {
	entrypoint: Label,
	scopes: Vec<Scope>,
	blocks: Vec<BasicBlock>,

	builder: BasicBlockBuilder,

	v_register_idx: usize,
}

impl FunctionContext {
	pub fn new(entrypoint: Label) -> Self {
		Self {
			entrypoint: entrypoint.clone(),
			blocks: Vec::new(),
			scopes: vec![Scope::new(entrypoint.clone(), Label::new("unreachable"))],
			builder: BasicBlockBuilder::new(entrypoint),
			v_register_idx: 0,
		}
	}

	pub fn new_vreg(&mut self) -> VirtualReg {
		self.v_register_idx += 1;
		VirtualReg(self.v_register_idx as isize - 1)
	}

	pub fn defer(&mut self, call: Call) {
		self.scopes.last_mut().unwrap().defers.push(call);
	}

	pub fn new_local(&mut self, variable: Variable) -> VirtualReg {
		let vreg = self.new_vreg();
		self.scopes.last_mut().unwrap().locals.push((vreg, variable));
		vreg
	}

	pub fn enter_scope(&mut self, scope: Scope) {
		self.scopes.push(scope);
	}

	pub fn current_scope(&mut self) -> Option<Scope> {
		self.scopes.pop()
	}

	pub fn get_deferred_calls(&self) -> Vec<Call> {
		self.scopes
			.iter()
			.rev()
			.flat_map(|s|
				s.defers
				.iter()
				.cloned()
				.rev()
			).collect()
	}

	pub fn find_local(&self, var: &Variable) -> Option<VirtualReg> {
		for scope in self.scopes.iter().rev() {
			for local in &scope.locals {
				if local.1 == *var {
					return Some(local.0)
				}
			}
		}

		None
	}

	pub fn exit_scope(&mut self) {
		self.scopes.pop().expect("(BUG) Cannot exit outermost scope");
	}

	pub fn next_block(&mut self, terminator: Terminator, label: Label) {
		self.blocks.push(self.builder.next_block(terminator, label));
	}

	pub fn push_stmt(&mut self, stmt: IR::Statement) {
		self.builder.push(stmt);
	}

	pub fn terminate(mut self, terminator: Terminator) -> Vec<IR::BasicBlock> {
		self.blocks.push(self.builder.terminate_block(terminator));
		 self.blocks
	}

	pub fn take(&mut self) -> Vec<IR::BasicBlock> {
		std::mem::take(&mut self.blocks)
	}
}