haloumi-picus 0.5.6

use std::{
    collections::{HashMap, HashSet},
    ops::Add,
};

use crate::pcl::{
    Module, ModuleRef,
    display::{Display, TextRepresentable, TextRepresentation},
    stmt::traits::{CallLike as _, MaybeCallLike as _},
    vars::VarKind,
};
use anyhow::{Result, anyhow};
use haloumi_ir::Prime;

#[derive(Debug)]
pub struct Program<K: VarKind> {
    prime: PrimeNumber,
    modules: Vec<Module<K>>,
}

impl<K: VarKind> Program<K> {
    pub fn modules(&self) -> &[Module<K>] {
        &self.modules
    }

    pub fn modules_mut(&mut self) -> &mut [Module<K>] {
        &mut self.modules
    }

    fn module_names(&self) -> HashSet<&str> {
        self.modules.iter().map(|m| m.name.as_str()).collect()
    }

    pub fn display(&self) -> Display<'_, K> {
        Display::new(self)
    }

    pub fn prime(&self) -> &Prime {
        self.prime.as_ref()
    }

    pub fn merge(&mut self, other: Program<K>) -> Result<()> {
        let collisions: HashSet<String> = self
            .module_names()
            .intersection(&other.module_names())
            .map(|s| (*s).to_owned())
            .collect();
        // Maps the old name to the new one
        let mut renames: HashMap<String, String> = Default::default();

        let renamed = other
            .modules
            .into_iter()
            .map(|m| -> Result<Module<K>> {
                if !collisions.contains(m.name.as_str()) {
                    return Ok(m);
                }

                let new_name = (0..)
                    .find_map(|i| {
                        let new_name = format!("{}{i}", *m.name);
                        if collisions.contains(new_name.as_str()) {
                            return None;
                        }
                        Some(new_name)
                    })
                    .ok_or_else(|| anyhow!("Failed to find a new name"))?;
                let mut m = m;
                renames.insert(m.name.to_string(), new_name.clone());
                *m.name = new_name;
                Ok(m)
            })
            // Collect the modules to make a barrier since the next step needs the full list of
            // renames
            .collect::<Result<Vec<_>>>()?;
        let renames = renames;
        self.modules.extend(renamed.into_iter().map(|m| {
            let mut m = m;

            m.stmts = m
                .stmts
                .into_iter()
                .map(|s| match s.as_call() {
                    Some(call) if renames.contains_key(call.callee()) => {
                        call.with_new_callee(renames[call.callee()].clone())
                    }
                    _ => s,
                })
                .collect();

            m
        }));

        Ok(())
    }

    pub fn new(prime: impl Into<Prime>, modules: Vec<ModuleRef<K>>) -> Self
    where
        K: VarKind + Clone,
    {
        Self {
            prime: PrimeNumber(prime.into()),
            modules: modules.into_iter().map(Into::into).collect(),
        }
    }
}

impl<K: VarKind> Add for Program<K> {
    type Output = Result<Program<K>>;

    fn add(self, rhs: Self) -> Self::Output {
        let mut lhs = self;
        lhs.merge(rhs)?;
        Ok(lhs)
    }
}

#[derive(Debug)]
struct PrimeNumber(Prime);

impl TextRepresentable for PrimeNumber {
    fn to_repr(&self) -> TextRepresentation<'_> {
        owned_list!("prime-number", &self.0).break_line()
    }

    fn width_hint(&self) -> usize {
        15 + self.0.width_hint()
    }
}

impl AsRef<Prime> for PrimeNumber {
    fn as_ref(&self) -> &Prime {
        &self.0
    }
}

type TR<'a> = TextRepresentation<'a>;

impl<K: VarKind> TextRepresentable for Program<K> {
    fn to_repr(&self) -> TextRepresentation<'_> {
        self.modules
            .iter()
            .fold(owned_list!(&self.prime), |acc, m| {
                acc + TR::br() + m.to_repr()
            })
            .with_punct("".into())
    }

    fn width_hint(&self) -> usize {
        self.prime.width_hint()
    }
}