solveq/src/normal_form.rs

use crate::language::{EGraph, EquationLanguage, Rational, RATIONAL_ONE, RATIONAL_ZERO};
use std::collections::HashMap;
use egg::Id;

#[derive(Debug,Clone)]
pub enum SpecialTerm {
	Constant(Rational),
	PowerOfX(usize),
	Monomial(usize, Rational),
	MonicNonconstPoly(Vec<Rational>),
	Factorization(Rational, Vec<Vec<Rational>>),
	Other,
}

fn search_for<F, T>(egraph: &EGraph, f: F) -> HashMap<Id, T>
where
	F: Fn(Id, &EquationLanguage, &HashMap<Id, T>) -> Option<T> {

	let mut result: HashMap<Id, T> = HashMap::new();
	let mut modifications: usize = 1;

	while modifications > 0 {
		modifications = 0;

		for cls in egraph.classes() {
			let id = cls.id;
			if result.contains_key(&id) {
				continue;
			}

			for node in &cls.nodes {
				if let Some(x) = f(id, node, &result) {
					result.insert(id, x);
					modifications += 1;
				}
			}
		}

		println!("{} modifications!", modifications);
	}

	result
}

pub fn analyze3(egraph: &EGraph, eclass: Id) {
	let constants = search_for(egraph, |id, _, _|
		egraph[id].data.as_ref().map(|c|c.clone())
	);

	println!("{:?}", constants);

	let powers_of_x = search_for(egraph, |_, node, matches| match *node {
		EquationLanguage::Unknown => Some(1),
		EquationLanguage::Mul([a,b]) => {
			if !matches.contains_key(&a) || !matches.contains_key(&b) {
				return None;
			}

			let (dega, degb) = (matches[&a], matches[&b]);
			Some(dega + degb)
		},
		_ => None,
	});

	println!("{:?}", powers_of_x);

	let monomials = search_for(egraph, |id, node, _| {
		if let Some(deg) = powers_of_x.get(&id) {
			return Some((*deg, RATIONAL_ONE.clone()));
		}

		if let Some(c) = constants.get(&id) {
			return Some((0, c.clone()));
		}

		match *node {
			EquationLanguage::Mul([a,b]) => {
				if !constants.contains_key(&a) || !powers_of_x.contains_key(&b) {
					return None;
				}

				let (coeff, deg) = (&constants[&a], powers_of_x[&b]);
				Some((deg, coeff.clone()))
			},
			_ => None,
		}
	});

	println!("{:?}", monomials);

	let monic_polynomials = search_for(egraph, |id, node, matches| {
		if let Some(deg) = powers_of_x.get(&id) {
			let mut poly: Vec<Rational> = Vec::new();
			poly.resize(*deg, RATIONAL_ZERO);
			poly.push(RATIONAL_ONE.clone());
			Some(poly)
		} else {
			match *node {
				EquationLanguage::Add([a,b]) => {
					if !matches.contains_key(&a) || !monomials.contains_key(&b) {
						return None;
					}

					let (leading, (deg, coeff)) = (&matches[&a], &monomials[&b]);
					if leading.len() <= *deg || leading[*deg] != RATIONAL_ZERO {
						return None;
					}

					let mut poly = leading.clone();
					poly[*deg] = coeff.clone();
					Some(poly)
				},
				_ => None,
			}
		}
	});

	for p in &monic_polynomials {
		println!("{:?}", p);
	}

	let factorizations: HashMap<Id, (Rational, Vec<Vec<Rational>>)> = search_for(egraph, |id, node, matches| {
		if let Some(c) = constants.get(&id) {
			return Some((c.clone(), vec![]));
		}

		if let Some(poly) = monic_polynomials.get(&id) {
			return Some((RATIONAL_ONE.clone(), vec![poly.clone()]));
		}

		match *node {
			EquationLanguage::Mul([a,b]) => {
				if !matches.contains_key(&a) || !monic_polynomials.contains_key(&b) {
					return None;
				}

				let ((factor, polys), newpoly) = (&matches[&a], &monic_polynomials[&b]);

				let mut combined: Vec<Vec<Rational>> = polys.clone();
				combined.push(newpoly.clone());
				Some((factor.clone(), combined))
			},
			_ => None,
		}
	});

	/*
	for p in &factorizations {
		println!("{:?}", p);
}
	*/

	println!("{:?}", factorizations[&eclass]);
}

pub fn analyze2(egraph: &EGraph) -> HashMap<Id, SpecialTerm> {
	let mut types: HashMap<Id, SpecialTerm> = HashMap::new();

	let mut modifications: usize = 1;

	while modifications > 0 {
		modifications = 0;

		for cls in egraph.classes() {
			let id = cls.id;
			if types.contains_key(&id) {
				continue;
			}

			if let Some(c) = &egraph[id].data {
				types.insert(id, SpecialTerm::Constant(c.clone()));
				modifications += 1;
				continue;
			}

			for node in &cls.nodes {
				match *node {
					EquationLanguage::Unknown => {
						types.insert(id, SpecialTerm::PowerOfX(1));
						modifications += 1;
					},
					EquationLanguage::Mul([a,b]) => {
						// as we don't know a and b yet, defer to future iteration
						if !types.contains_key(&a) || !types.contains_key(&b) {
							continue;
						}

						match (&types[&a], &types[&b]) {
							(SpecialTerm::PowerOfX(dega), SpecialTerm::PowerOfX(degb)) => {
								types.insert(id, SpecialTerm::PowerOfX(*dega + *degb));
								modifications += 1;
							},
							(SpecialTerm::Constant(coeff), SpecialTerm::PowerOfX(deg)) => {
								types.insert(id, SpecialTerm::Monomial(*deg, coeff.clone()));
								modifications += 1;
							},
							_ => { },
						}
					},
					EquationLanguage::Add([a,b]) => {
						// as we don't know a and b yet, defer to future iteration
						if !types.contains_key(&a) || !types.contains_key(&b) {
							continue;
						}

						match (&types[&a], &types[&b]) {
							(SpecialTerm::MonicNonconstPoly(poly), SpecialTerm::Monomial(deg, coeff)) => {
								if poly.len() <= *deg || poly[*deg] != RATIONAL_ZERO {
									continue;
								}

								let mut poly = poly.clone();
								poly[*deg] = coeff.clone();
								types.insert(id, SpecialTerm::MonicNonconstPoly(poly));
								modifications += 1;
							},
							_ => { },
						}
					},
					_ => {}
				}
			}
		}

		println!("{} modifications!", modifications);
	}

	types
}