leaf/assembly/src/functions/ir.rs

use crate::{
	assembly::Ctx,
	functions::{Function, FunctionBody},
	types::{Type, derivations::*},
	values::{Value, ValueFlags},
};
use derive_more::{Debug, Display};
use std::{borrow::Cow, cell::UnsafeCell, hash::Hash, ops::Deref, sync::OnceLock};

// Maybe unsafe but honestly it's extremely unlikely that this will go wrong and it won't cause any issues.
struct Id(UnsafeCell<u32>);
unsafe impl Send for Id {}
unsafe impl Sync for Id {}

#[derive(Debug, Display)]
#[debug("{variant:?}")]
#[display("%{}", unsafe { *id.0.get() })]
pub struct Instruction<'l> {
	id: Id,
	pub parent_block: &'l Block<'l>,
	pub variant: InstructionVariant<'l>,
}

impl Eq for Instruction<'_> {}

impl PartialEq for Instruction<'_> {
	#[inline]
	fn eq(&self, _: &Self) -> bool {
		false
	}
}

impl Hash for Instruction<'_> {
	#[inline]
	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
		(self as *const Self).hash(state);
	}
}

impl<'l> Deref for Instruction<'l> {
	type Target = InstructionVariant<'l>;

	#[inline]
	fn deref(&self) -> &Self::Target {
		&self.variant
	}
}

impl<'l> Instruction<'l> {
	#[inline]
	pub fn ctx(&self) -> Ctx<'l> {
		self.parent_block.func.declaring_assembly.ctx()
	}

	#[inline]
	pub fn id(&self) -> u32 {
		unsafe { *self.id.0.get() }
	}

	pub fn value_flags(&self) -> ValueFlags {
		match self.variant {
			_ => todo!(),
		}
	}

	pub fn value_ty(&'l self) -> Type<'l> {
		match self.variant {
			InstructionVariant::Return(_) => Type::Void,
			InstructionVariant::Store(_, _) => Type::Void,
			InstructionVariant::StackAlloc(ty) => ty.make_ptr(true).into(),
			InstructionVariant::Load(value) => match value.ty() {
				Type::Ptr(PtrT { base, .. }) => *base,
				_ => unreachable!(),
			},
			InstructionVariant::IAdd(a, _) => a.ty(),
			InstructionVariant::FAdd(a, _) => a.ty(),
			InstructionVariant::ICmp(_, _, _) => Type::Bool,
			_ => todo!("{self:?}"),
		}
	}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Cmp {
	Eq,
	Lt,
	Gt,
	Le,
	Ge,
}

#[derive(PartialEq, Eq)]
pub enum InstructionVariant<'l> {
	StackAlloc(Type<'l>),
	GCAlloc(Type<'l>),

	Load(Value<'l>),
	Store(Value<'l>, Value<'l>),

	IAdd(Value<'l>, Value<'l>),
	FAdd(Value<'l>, Value<'l>),

	ICmp(Value<'l>, Value<'l>, Cmp),
	FCmp(Value<'l>, Value<'l>, Cmp),

	Call(&'l Function<'l>, Vec<Value<'l>>),
	Jump(&'l Block<'l>),
	Branch {
		cond: Value<'l>,
		true_case: &'l Block<'l>,
		false_case: &'l Block<'l>,
	},
	Return(Option<Value<'l>>),
}

impl InstructionVariant<'_> {
	pub fn is_block_termination(&self) -> bool {
		match self {
			Self::Return(_) => true,
			Self::Jump(_) => true,
			Self::Branch { .. } => true,
			_ => false,
		}
	}
}

impl std::fmt::Debug for InstructionVariant<'_> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			Self::StackAlloc(ty) => write!(f, "stackalloc {ty}"),
			Self::GCAlloc(ty) => write!(f, "gcalloc {ty}"),
			Self::Load(v) => write!(f, "load {v}"),
			Self::Store(t, v) => write!(f, "store {t}, {v}"),
			Self::IAdd(a, b) => write!(f, "iadd {a}, {b}"),
			Self::FAdd(a, b) => write!(f, "fadd {a}, {b}"),
			Self::ICmp(a, b, c) => write!(f, "icmp {c:?} {a}, {b}"),
			Self::FCmp(a, b, c) => write!(f, "fcmp {c:?} {a}, {b}"),
			Self::Call(func, args) => {
				write!(f, "call {:#?}(", *func as *const Function)?;
				let mut separator = "";
				for arg in args {
					write!(f, "{separator}{arg}")?;
					separator = ", ";
				}
				write!(f, ")")
			}
			Self::Branch {
				cond,
				true_case,
				false_case,
			} => write!(f, "br {cond} #{}, #{}", true_case.id, false_case.id),
			Self::Jump(b) => write!(f, "jump #{}", b.id),
			Self::Return(None) => write!(f, "return"),
			Self::Return(Some(v)) => write!(f, "return {v}"),
		}
	}
}

pub struct Block<'l> {
	pub id: u32,
	pub func: &'l Function<'l>,
	instructions: OnceLock<Vec<&'l Instruction<'l>>>,
}

impl<'l> Block<'l> {
	#[inline]
	pub fn instructions(&self) -> &[&'l Instruction<'l>] {
		match self.instructions.get() {
			None => &[],
			Some(v) => v,
		}
	}
}

impl Eq for Block<'_> {}

impl PartialEq for Block<'_> {
	fn eq(&self, other: &Self) -> bool {
		std::ptr::eq(self, other)
	}
}

pub struct BlockBuilder<'l> {
	block: &'l Block<'l>,
	instructions: Vec<&'l Instruction<'l>>,
}

pub type BlockBuilderError<'l> = Cow<'l, str>;
pub type BlockBuilderResult<'l, T> = Result<T, BlockBuilderError<'l>>;

impl<'l> BlockBuilder<'l> {
	pub fn stack_alloc(&mut self, ty: Type<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		let inst = self.push_instruction(InstructionVariant::StackAlloc(ty))?;
		Ok(inst.into())
	}

	pub fn store(
		&mut self,
		target: Value<'l>,
		value: Value<'l>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		let value_ty = value.ty();
		let target_ty = target.ty();
		match target_ty {
			Type::Ptr(PtrT {
				base,
				mutable: true,
			}) if *base == value_ty => {}

			Type::Ref(RefT {
				base,
				mutable: true,
			}) if *base == value_ty => {}

			_ => {
				return Err(format!(
					"Cannot store value of type `{value_ty}` into target of type `{target_ty}.`"
				)
				.into());
			}
		};
		let inst = self.push_instruction(InstructionVariant::Store(target, value))?;
		Ok(inst.into())
	}

	pub fn load(&mut self, value: Value<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		let value_ty = value.ty();
		match value_ty {
			Type::Ptr(PtrT { .. }) => {}
			Type::Ref(RefT { .. }) => {}
			_ => {
				return Err(format!("Cannot load value of type `{}`.`", value_ty).into());
			}
		}
		let inst = self.push_instruction(InstructionVariant::Load(value))?;
		Ok(inst.into())
	}

	pub fn add(&mut self, a: Value<'l>, b: Value<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		let [a_ty, b_ty] = [a.ty(), b.ty()];

		match (a_ty, b_ty) {
			(Type::Int(a_ty), Type::Int(b_ty)) if a_ty == b_ty => {
				let inst = self.push_instruction(InstructionVariant::IAdd(a, b))?;
				Ok(inst.into())
			}
			(Type::Float(a_ty), Type::Float(b_ty)) if a_ty == b_ty => {
				let inst = self.push_instruction(InstructionVariant::FAdd(a, b))?;
				Ok(inst.into())
			}
			_ => Err(format!("Cannot add values of type `{a_ty}` and `b_ty`.").into()),
		}
	}

	pub fn cmp(
		&mut self,
		a: Value<'l>,
		b: Value<'l>,
		cmp: Cmp,
	) -> BlockBuilderResult<'l, Value<'l>> {
		let [a_ty, b_ty] = [a.ty(), b.ty()];

		match (a_ty, b_ty) {
			(Type::Int(a_ty), Type::Int(b_ty)) if a_ty == b_ty => {
				let inst = self.push_instruction(InstructionVariant::ICmp(a, b, cmp))?;
				Ok(inst.into())
			}
			(Type::Float(a_ty), Type::Float(b_ty)) if a_ty == b_ty => {
				let inst = self.push_instruction(InstructionVariant::FCmp(a, b, cmp))?;
				Ok(inst.into())
			}
			_ => Err(format!("Cannot compare values of type `{a_ty}` and `b_ty`.").into()),
		}
	}

	pub fn jump(&mut self, block: &'l Block<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		if !std::ptr::eq(block.func, self.block.func) {
			return Err("Block does not belong to this function.".into());
		}
		let inst = self.push_instruction(InstructionVariant::Jump(block))?;
		Ok(inst.into())
	}

	pub fn branch(
		&mut self,
		cond: Value<'l>,
		true_case: &'l Block<'l>,
		false_case: &'l Block<'l>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		if !std::ptr::eq(true_case.func, self.block.func) {
			return Err("Block does not belong to this function.".into());
		}
		if !std::ptr::eq(false_case.func, self.block.func) {
			return Err("Block does not belong to this function.".into());
		}
		if !matches!(cond.ty(), Type::Bool) {
			return Err(format!("Expected value of type `bool`, found `{}`.", cond.ty()).into());
		}
		let inst = self.push_instruction(InstructionVariant::Branch {
			cond,
			true_case,
			false_case,
		})?;
		Ok(inst.into())
	}

	pub fn call(
		&mut self,
		func: &'l Function<'l>,
		args: Vec<Value<'l>>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		let par_t = &*func.ty.par_t;
		if par_t.len() != args.len() {
			return Err(format!(
				"Invalid parameter count. Expected {}, found {}.",
				par_t.len(),
				args.len()
			)
			.into());
		}

		if let Some(i) = par_t.iter().zip(&args).position(|(a, b)| *a != b.ty()) {
			return Err(format!(
				"Invalid parameter at position {i}. Expected type `{}`, found `{}`.",
				&par_t[i],
				args[i].ty()
			)
			.into());
		}

		let inst = self.push_instruction(InstructionVariant::Call(func, args))?;
		Ok(inst.into())
	}

	pub fn ret(&mut self, value: Option<Value<'l>>) -> BlockBuilderResult<'l, Value<'l>> {
		let ret_t = self.block.func.ty.ret_t;
		let value_ty = match value {
			Some(v) => v.ty(),
			None => Type::Void,
		};
		if value_ty != ret_t {
			return Err(format!(
				"Cannot return value of type `{value_ty}`. Expected type `{ret_t}`."
			)
			.into());
		}
		let inst = self.push_instruction(InstructionVariant::Return(value))?;
		Ok(inst.into())
	}

	#[inline]
	pub fn has_termination(&self) -> bool {
		match self.instructions.as_slice() {
			[.., i] => i.is_block_termination(),
			_ => false,
		}
	}

	pub fn build(self) -> BlockBuilderResult<'l, &'l Block<'l>> {
		if !self.has_termination() {
			return Err(format!("Block #{} has not termination.", self.block.id).into());
		}
		self.block.instructions.set(self.instructions).unwrap();
		Ok(self.block)
	}

	fn push_instruction(
		&mut self,
		variant: InstructionVariant<'l>,
	) -> BlockBuilderResult<'l, &'l Instruction<'l>> {
		if self.has_termination() {
			return Err(format!("Block #{} has already terminated", self.block.id).into());
		}
		let instruction = &*self.block.func.ctx().alloc.alloc(Instruction {
			id: Id(UnsafeCell::new(u32::MAX)),
			parent_block: self.block,
			variant,
		});
		self.instructions.push(instruction);
		Ok(instruction)
	}
}

pub struct FunctionBodyBuilder<'l> {
	current_block: usize,
	func: &'l Function<'l>,
	blocks: Vec<BlockBuilder<'l>>,
}

impl<'l> FunctionBodyBuilder<'l> {
	pub fn new(func: &'l Function<'l>) -> Self {
		let mut builder = Self {
			func,
			blocks: vec![],
			current_block: 0,
		};
		builder.create_block();
		builder
	}

	pub fn current_block(&self) -> &'l Block<'l> {
		self.blocks[self.current_block].block
	}

	pub fn set_current_block(&mut self, block: &'l Block<'l>) -> Option<&'l Block<'l>> {
		match std::ptr::eq(block.func, self.func) {
			false => None,
			true => {
				let current = self.current_block();
				self.current_block = block.id as usize;
				Some(current)
			}
		}
	}

	pub fn create_block(&mut self) -> &'l Block<'l> {
		let block = &*self.func.ctx().alloc.alloc(Block {
			id: self.blocks.len() as u32,
			func: self.func,
			instructions: OnceLock::new(),
		});
		let builder = BlockBuilder {
			block,
			instructions: vec![],
		};
		self.blocks.push(builder);
		block
	}

	pub fn build(self) -> Result<&'l FunctionBody<'l>, Cow<'l, str>> {
		let mut next_id = 0..;
		let mut blocks = Vec::with_capacity(self.blocks.len());
		for block in self.blocks {
			let block = block.build()?;
			for inst in block.instructions() {
				unsafe {
					let ptr = inst.id.0.get();
					std::ptr::write(ptr, next_id.next().unwrap());
				}
			}
			blocks.push(block);
		}
		if self.func.body.set(FunctionBody { blocks }).is_err() {
			return Err("Function body already exists.".into());
		}
		Ok(self.func.body.get().unwrap())
	}

	pub fn stack_alloc(&mut self, ty: Type<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().stack_alloc(ty)
	}

	pub fn store(
		&mut self,
		target: Value<'l>,
		value: Value<'l>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().store(target, value)
	}

	pub fn load(&mut self, value: Value<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().load(value)
	}

	pub fn add(&mut self, a: Value<'l>, b: Value<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().add(a, b)
	}

	pub fn cmp(
		&mut self,
		a: Value<'l>,
		b: Value<'l>,
		cmp: Cmp,
	) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().cmp(a, b, cmp)
	}

	pub fn jump(&mut self, block: &'l Block<'l>) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().jump(block)
	}

	pub fn branch(
		&mut self,
		cond: Value<'l>,
		true_case: &'l Block<'l>,
		false_case: &'l Block<'l>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().branch(cond, true_case, false_case)
	}

	pub fn call(
		&mut self,
		func: &'l Function<'l>,
		args: Vec<Value<'l>>,
	) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().call(func, args)
	}

	pub fn ret(&mut self, value: Option<Value<'l>>) -> BlockBuilderResult<'l, Value<'l>> {
		self.current_builder().ret(value)
	}

	fn current_builder(&mut self) -> &mut BlockBuilder<'l> {
		&mut self.blocks[self.current_block]
	}
}