❓❓❓❓❓❓❓ (WIP, no judging)

extractor.rs

use crate::{
	expressions::object_literal::ObjectLiteralMethodBase, statements::class::ClassConstructorBase,
	statements::class::ClassFunctionBase, statements::StatementFunctionBase, ArrowFunctionBase,
	ExpressionFunctionBase,
};
use crate::{FunctionBase, FunctionBased, FunctionId, Visitable};
use std::collections::HashMap;
use std::marker::PhantomData;

// TODO generic function_id
pub trait GetFunction<T: FunctionBased> {
	fn get_function_ref(&self, id: FunctionId) -> Option<&FunctionBase<T>>;

	fn get_function(&mut self, id: FunctionId) -> FunctionBase<T>;

	fn insert_function(&mut self, func: FunctionBase<T>);

	fn new_function(&mut self, func: FunctionBase<T>) -> ExtractedFunction<T> {
		let id = func.get_function_id();
		self.insert_function(func);
		ExtractedFunction(id, PhantomData::default())
	}
}

macro_rules! make_extracted_functions {
	{$($ty:ident),*} => {

		#[derive(Default, Debug)]
		#[allow(non_snake_case)]
		pub struct ExtractedFunctions {
			$(
				$ty: HashMap<FunctionId, FunctionBase<$ty>>,
			)*
		}

		$(
			impl GetFunction<$ty> for ExtractedFunctions {
				fn get_function_ref(&self, id: FunctionId) -> Option<&FunctionBase<$ty>> {
					self.$ty.get(&id)
				}

				fn get_function(&mut self, id: FunctionId) -> FunctionBase<$ty> {
					self.$ty.remove(&id).unwrap()
				}

				fn insert_function(&mut self, func: FunctionBase<$ty>) {
					let id = func.get_function_id();
					self.$ty.insert(id, func);
				}
			}
		)*
	}
}

make_extracted_functions! {
	ArrowFunctionBase,
	ExpressionFunctionBase,
	StatementFunctionBase,
	ObjectLiteralMethodBase,
	ClassConstructorBase,
	ClassFunctionBase
}

/// New type purely for [Visitable] implementation
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct ExtractedFunction<T: FunctionBased>(pub FunctionId, PhantomData<T>);

impl<T: FunctionBased> Visitable for ExtractedFunction<T>
where
	ExtractedFunctions: GetFunction<T>,
	FunctionBase<T>: Visitable
{
	fn visit<TData>(
		&self,
		visitors: &mut (impl crate::VisitorReceiver<TData> + ?Sized),
		data: &mut TData,
		settings: &crate::VisitSettings,
		functions: &mut ExtractedFunctions,
		chain: &mut temporary_annex::Annex<crate::Chain>,
	) {
		let function: FunctionBase<T> = functions.get_function(self.0);
		function.visit(visitors, data, settings, functions, chain);
		functions.insert_function(function)
	}

	fn visit_mut<TData>(
		&mut self,
		visitors: &mut (impl crate::VisitorMutReceiver<TData> + ?Sized),
		data: &mut TData,
		settings: &crate::VisitSettings,
		functions: &mut ExtractedFunctions,
		chain: &mut temporary_annex::Annex<crate::Chain>,
	) {
		let mut function: FunctionBase<T> = functions.get_function(self.0);
		function.visit_mut(visitors, data, settings, functions, chain);
		functions.insert_function(function)
	}
}