// Copyright (c) 2016-2019 Fabian Schuiki

//! A parameter environment generated by an instantiation.

use crate::{
    ast_map::AstNode,
    crate_prelude::*,
    hir::{HirNode, NamedParam, PosParam},
    ty::Type,
    value::Value,
};

/// A parameter environment.
///
/// This is merely an handle that is cheap to copy and pass around. Use the
/// [`Context`] to resolve this to the actual [`ParamEnvData`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ParamEnv(pub(crate) u32);

/// A node id with corresponding parameter environment.
pub type NodeEnvId = (NodeId, ParamEnv);

/// A parameter environment.
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
pub struct ParamEnvData<'t> {
    values: Vec<(NodeId, ParamEnvBinding<Value<'t>>)>,
    types: Vec<(NodeId, ParamEnvBinding<Type<'t>>)>,
}

impl<'t> ParamEnvData<'t> {
    /// Find the value assigned to a node.
    pub fn find_value(&self, node_id: NodeId) -> Option<ParamEnvBinding<Value<'t>>> {
        self.values
            .iter()
            .find(|&&(id, _)| id == node_id)
            .map(|&(_, id)| id)
    }

    /// Find the type assigned to a node.
    pub fn find_type(&self, node_id: NodeId) -> Option<ParamEnvBinding<Type<'t>>> {
        self.types
            .iter()
            .find(|&&(id, _)| id == node_id)
            .map(|&(_, id)| id)
    }

    /// Assign a value to a node.
    pub fn set_value(&mut self, node_id: NodeId, value: Value<'t>) {
        self.values.retain(|&(n, _)| n != node_id);
        self.values.push((node_id, ParamEnvBinding::Direct(value)));
    }
}

/// A binding in a parameter environment.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum ParamEnvBinding<T> {
    /// A direct binding, directly assigning a type or value to a node.
    Direct(T),
    /// An indirect binding, pointing at another node's type or value.
    Indirect(NodeEnvId),
}

/// A location that implies a parameter environment.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ParamEnvSource<'hir> {
    ModuleInst {
        module: NodeId,
        inst: NodeId,
        env: ParamEnv,
        pos: &'hir [PosParam],
        named: &'hir [NamedParam],
    },
}

pub(crate) fn compute<'gcx>(
    cx: &impl Context<'gcx>,
    src: ParamEnvSource<'gcx>,
) -> Result<ParamEnv> {
    match src {
        ParamEnvSource::ModuleInst {
            module,
            inst,
            env,
            pos,
            named,
        } => {
            let module = match cx.hir_of(module)? {
                HirNode::Module(m) => m,
                _ => panic!("expected module"),
            };

            // Collect a list of module parameters.
            let module_params: Vec<_> = module
                .params
                .iter()
                .cloned()
                .chain(module.block.params.iter().cloned())
                .collect();

            // Associate the positional and named assignments with the actual
            // parameters of the module.
            let param_iter = pos
                .iter()
                .enumerate()
                .map(
                    |(index, &(span, assign_id))| match module_params.get(index) {
                        Some(&param_id) => Ok((param_id, (assign_id, env))),
                        None => {
                            cx.emit(
                                DiagBuilder2::error(format!(
                                    "{} only has {} parameter(s)",
                                    module.desc_full(),
                                    module_params.len()
                                ))
                                .span(span),
                            );
                            Err(())
                        }
                    },
                )
                .chain(named.iter().map(|&(_span, name, assign_id)| {
                    let names: Vec<_> = module_params
                        .iter()
                        .flat_map(|&id| match cx.ast_of(id) {
                            Ok(AstNode::TypeParam(_, p)) => Some((p.name.name, id)),
                            Ok(AstNode::ValueParam(_, p)) => Some((p.name.name, id)),
                            Ok(_) => unreachable!(),
                            Err(()) => None,
                        })
                        .collect();
                    match names
                        .iter()
                        .find(|&(param_name, _)| *param_name == name.value)
                    {
                        Some(&(_, param_id)) => Ok((param_id, (assign_id, env))),
                        None => {
                            cx.emit(
                                DiagBuilder2::error(format!(
                                    "no parameter `{}` in {}",
                                    name,
                                    module.desc_full(),
                                ))
                                .span(name.span)
                                .add_note(format!(
                                    "declared parameters are {}",
                                    names
                                        .iter()
                                        .map(|&(n, _)| format!("`{}`", n))
                                        .collect::<Vec<_>>()
                                        .join(", ")
                                )),
                            );
                            Err(())
                        }
                    }
                }));
            let param_iter = param_iter
                .collect::<Vec<_>>()
                .into_iter()
                .collect::<Result<Vec<_>>>()?
                .into_iter();

            // Split up type and value parameters.
            let mut types = vec![];
            let mut values = vec![];
            for (param_id, assign_id) in param_iter {
                let assign_id = match assign_id {
                    (Some(i), n) => (i, n),
                    _ => continue,
                };
                match cx.ast_of(param_id)? {
                    AstNode::TypeParam(..) => {
                        cx.set_lowering_hint(assign_id.0, hir::Hint::Type);
                        types.push((param_id, ParamEnvBinding::Indirect(assign_id)))
                    }
                    AstNode::ValueParam(..) => {
                        cx.set_lowering_hint(assign_id.0, hir::Hint::Expr);
                        values.push((param_id, ParamEnvBinding::Indirect(assign_id)))
                    }
                    _ => unreachable!(),
                }
            }

            let env = cx.intern_param_env(ParamEnvData { types, values });
            cx.add_param_env_context(env, inst);
            Ok(env)
        }
    }
}