mod fulfill;
pub mod lib;
mod search_graph;
mod stack;

use self::fulfill::{Fulfill, RecursiveInferenceTable, RecursiveSolver};
use self::lib::{Guidance, Minimums, Solution, UCanonicalGoal};
use self::search_graph::{DepthFirstNumber, SearchGraph};
use self::stack::{Stack, StackDepth};
use crate::clauses::program_clauses_for_goal;
use crate::infer::{InferenceTable, ParameterEnaVariableExt};
use crate::solve::truncate;
use crate::{coinductive_goal::IsCoinductive, RustIrDatabase};
use chalk_ir::fold::Fold;
use chalk_ir::interner::{HasInterner, Interner};
use chalk_ir::visit::Visit;
use chalk_ir::zip::Zip;
use chalk_ir::{debug, debug_heading, info, info_heading};
use chalk_ir::{
    Binders, Canonical, ClausePriority, ConstrainedSubst, Constraint, DomainGoal, Environment,
    Fallible, Floundered, GenericArg, Goal, GoalData, InEnvironment, NoSolution, ProgramClause,
    ProgramClauseData, ProgramClauseImplication, Substitution, UCanonical, UniverseMap,
    VariableKinds,
};
use rustc_hash::FxHashMap;
use std::fmt::Debug;

pub(crate) struct RecursiveContext<I: Interner> {
    stack: Stack,

    /// The "search graph" stores "in-progress results" that are still being
    /// solved.
    search_graph: SearchGraph<I>,

    /// The "cache" stores results for goals that we have completely solved.
    /// Things are added to the cache when we have completely processed their
    /// result.
    cache: FxHashMap<UCanonicalGoal<I>, Fallible<Solution<I>>>,

    caching_enabled: bool,
}

pub(crate) struct RecursiveInferenceTableImpl<I: Interner> {
    pub(crate) infer: InferenceTable<I>,
}

impl<I: Interner> RecursiveInferenceTable<I> for RecursiveInferenceTableImpl<I> {
    fn instantiate_binders_universally<'a, T>(
        &mut self,
        interner: &'a I,
        arg: &'a Binders<T>,
    ) -> T::Result
    where
        T: Fold<I> + HasInterner<Interner = I>,
    {
        self.infer.instantiate_binders_universally(interner, arg)
    }

    fn instantiate_binders_existentially<'a, T>(
        &mut self,
        interner: &'a I,
        arg: &'a Binders<T>,
    ) -> T::Result
    where
        T: Fold<I> + HasInterner<Interner = I>,
    {
        self.infer.instantiate_binders_existentially(interner, arg)
    }

    fn canonicalize<T>(
        &mut self,
        interner: &I,
        value: &T,
    ) -> (Canonical<T::Result>, Vec<GenericArg<I>>)
    where
        T: Fold<I>,
        T::Result: HasInterner<Interner = I>,
    {
        let res = self.infer.canonicalize(interner, value);
        let free_vars = res
            .free_vars
            .into_iter()
            .map(|free_var| free_var.to_generic_arg(interner))
            .collect();
        (res.quantified, free_vars)
    }

    fn u_canonicalize<T>(
        &mut self,
        interner: &I,
        value0: &Canonical<T>,
    ) -> (UCanonical<T::Result>, UniverseMap)
    where
        T: HasInterner<Interner = I> + Fold<I> + Visit<I>,
        T::Result: HasInterner<Interner = I>,
    {
        let res = self.infer.u_canonicalize(interner, value0);
        (res.quantified, res.universes)
    }

    fn unify<T>(
        &mut self,
        interner: &I,
        environment: &Environment<I>,
        a: &T,
        b: &T,
    ) -> Fallible<(
        Vec<InEnvironment<DomainGoal<I>>>,
        Vec<InEnvironment<Constraint<I>>>,
    )>
    where
        T: ?Sized + Zip<I>,
    {
        let res = self.infer.unify(interner, environment, a, b)?;
        Ok((res.goals, res.constraints))
    }

    fn instantiate_canonical<T>(&mut self, interner: &I, bound: &Canonical<T>) -> T::Result
    where
        T: HasInterner<Interner = I> + Fold<I> + Debug,
    {
        self.infer.instantiate_canonical(interner, bound)
    }

    fn invert_then_canonicalize<T>(
        &mut self,
        interner: &I,
        value: &T,
    ) -> Option<Canonical<T::Result>>
    where
        T: Fold<I, Result = T> + HasInterner<Interner = I>,
    {
        self.infer.invert_then_canonicalize(interner, value)
    }

    fn needs_truncation(&mut self, interner: &I, max_size: usize, value: impl Visit<I>) -> bool {
        truncate::needs_truncation(interner, &mut self.infer, max_size, value)
    }
}

/// A Solver is the basic context in which you can propose goals for a given
/// program. **All questions posed to the solver are in canonical, closed form,
/// so that each question is answered with effectively a "clean slate"**. This
/// allows for better caching, and simplifies management of the inference
/// context.
pub(crate) struct Solver<'me, I: Interner> {
    program: &'me dyn RustIrDatabase<I>,
    context: &'me mut RecursiveContext<I>,
}

/// An extension trait for merging `Result`s
trait MergeWith<T> {
    fn merge_with<F>(self, other: Self, f: F) -> Self
    where
        F: FnOnce(T, T) -> T;
}

impl<T> MergeWith<T> for Fallible<T> {
    fn merge_with<F>(self: Fallible<T>, other: Fallible<T>, f: F) -> Fallible<T>
    where
        F: FnOnce(T, T) -> T,
    {
        match (self, other) {
            (Err(_), Ok(v)) | (Ok(v), Err(_)) => Ok(v),
            (Ok(v1), Ok(v2)) => Ok(f(v1, v2)),
            (Err(_), Err(e)) => Err(e),
        }
    }
}

impl<I: Interner> RecursiveContext<I> {
    pub(crate) fn new(overflow_depth: usize, caching_enabled: bool) -> Self {
        RecursiveContext {
            stack: Stack::new(overflow_depth),
            search_graph: SearchGraph::new(),
            cache: FxHashMap::default(),
            caching_enabled,
        }
    }

    pub(crate) fn solver<'me>(
        &'me mut self,
        program: &'me dyn RustIrDatabase<I>,
    ) -> Solver<'me, I> {
        Solver {
            program,
            context: self,
        }
    }
}

impl<'me, I: Interner> Solver<'me, I> {
    pub(crate) fn new_inference_table<
        T: Fold<I, I, Result = T> + HasInterner<Interner = I> + Clone,
    >(
        &self,
        ucanonical_goal: &UCanonical<InEnvironment<T>>,
    ) -> (
        RecursiveInferenceTableImpl<I>,
        Substitution<I>,
        InEnvironment<T::Result>,
    ) {
        let (infer, subst, canonical_goal) = InferenceTable::from_canonical(
            self.program.interner(),
            ucanonical_goal.universes,
            &ucanonical_goal.canonical,
        );
        let infer = crate::recursive::RecursiveInferenceTableImpl { infer };
        (infer, subst, canonical_goal)
    }

    /// Solves a canonical goal. The substitution returned in the
    /// solution will be for the fully decomposed goal. For example, given the
    /// program
    ///
    /// ```ignore
    /// struct u8 { }
    /// struct SomeType<T> { }
    /// trait Foo<T> { }
    /// impl<U> Foo<u8> for SomeType<U> { }
    /// ```
    ///
    /// and the goal `exists<V> { forall<U> { SomeType<U>: Foo<V> }
    /// }`, `into_peeled_goal` can be used to create a canonical goal
    /// `SomeType<!1>: Foo<?0>`. This function will then return a
    /// solution with the substitution `?0 := u8`.
    pub(crate) fn solve_root_goal(
        &mut self,
        canonical_goal: &UCanonicalGoal<I>,
    ) -> Fallible<Solution<I>> {
        debug!("solve_root_goal(canonical_goal={:?})", canonical_goal);
        assert!(self.context.stack.is_empty());
        let minimums = &mut Minimums::new();
        self.solve_goal(canonical_goal.clone(), minimums)
    }

    fn solve_new_subgoal(
        &mut self,
        canonical_goal: UCanonicalGoal<I>,
        depth: StackDepth,
        dfn: DepthFirstNumber,
    ) -> Minimums {
        debug_heading!(
            "solve_new_subgoal(canonical_goal={:?}, depth={:?}, dfn={:?})",
            canonical_goal,
            depth,
            dfn,
        );

        // We start with `answer = None` and try to solve the goal. At the end of the iteration,
        // `answer` will be updated with the result of the solving process. If we detect a cycle
        // during the solving process, we cache `answer` and try to solve the goal again. We repeat
        // until we reach a fixed point for `answer`.
        // Considering the partial order:
        // - None < Some(Unique) < Some(Ambiguous)
        // - None < Some(CannotProve)
        // the function which maps the loop iteration to `answer` is a nondecreasing function
        // so this function will eventually be constant and the loop terminates.
        let minimums = &mut Minimums::new();
        loop {
            let UCanonical {
                universes,
                canonical:
                    Canonical {
                        binders,
                        value: InEnvironment { environment, goal },
                    },
            } = canonical_goal.clone();

            let (current_answer, current_prio) = match goal.data(self.program.interner()) {
                GoalData::DomainGoal(domain_goal) => {
                    let canonical_goal = UCanonical {
                        universes,
                        canonical: Canonical {
                            binders,
                            value: InEnvironment {
                                environment,
                                goal: domain_goal.clone(),
                            },
                        },
                    };

                    // "Domain" goals (i.e., leaf goals that are Rust-specific) are
                    // always solved via some form of implication. We can either
                    // apply assumptions from our environment (i.e. where clauses),
                    // or from the lowered program, which includes fallback
                    // clauses. We try each approach in turn:

                    let InEnvironment { environment, goal } = &canonical_goal.canonical.value;

                    let (prog_solution, prog_prio) = {
                        debug_heading!("prog_clauses");

                        let prog_clauses = self.program_clauses_for_goal(environment, &goal);
                        match prog_clauses {
                            Ok(clauses) => {
                                self.solve_from_clauses(&canonical_goal, clauses, minimums)
                            }
                            Err(Floundered) => {
                                (Ok(Solution::Ambig(Guidance::Unknown)), ClausePriority::High)
                            }
                        }
                    };
                    debug!("prog_solution={:?}", prog_solution);

                    (prog_solution, prog_prio)
                }

                _ => {
                    let canonical_goal = UCanonical {
                        universes,
                        canonical: Canonical {
                            binders,
                            value: InEnvironment { environment, goal },
                        },
                    };

                    self.solve_via_simplification(&canonical_goal, minimums)
                }
            };

            debug!(
                "solve_new_subgoal: loop iteration result = {:?} with minimums {:?}",
                current_answer, minimums
            );

            if !self.context.stack[depth].read_and_reset_cycle_flag() {
                // None of our subgoals depended on us directly.
                // We can return.
                self.context.search_graph[dfn].solution = current_answer;
                self.context.search_graph[dfn].solution_priority = current_prio;
                return *minimums;
            }

            let old_answer = &self.context.search_graph[dfn].solution;
            let old_prio = self.context.search_graph[dfn].solution_priority;

            let (current_answer, current_prio) = combine_with_priorities_for_goal(
                self.program.interner(),
                &canonical_goal.canonical.value.goal,
                old_answer.clone(),
                old_prio,
                current_answer,
                current_prio,
            );

            // Some of our subgoals depended on us. We need to re-run
            // with the current answer.
            if self.context.search_graph[dfn].solution == current_answer {
                // Reached a fixed point.
                return *minimums;
            }

            let current_answer_is_ambig = match &current_answer {
                Ok(s) => s.is_ambig(),
                Err(_) => false,
            };

            self.context.search_graph[dfn].solution = current_answer;
            self.context.search_graph[dfn].solution_priority = current_prio;

            // Subtle: if our current answer is ambiguous, we can just stop, and
            // in fact we *must* -- otherwise, we sometimes fail to reach a
            // fixed point. See `multiple_ambiguous_cycles` for more.
            if current_answer_is_ambig {
                return *minimums;
            }

            // Otherwise: rollback the search tree and try again.
            self.context.search_graph.rollback_to(dfn + 1);
        }
    }

    fn solve_via_simplification(
        &mut self,
        canonical_goal: &UCanonicalGoal<I>,
        minimums: &mut Minimums,
    ) -> (Fallible<Solution<I>>, ClausePriority) {
        debug_heading!("solve_via_simplification({:?})", canonical_goal);
        let (infer, subst, goal) = self.new_inference_table(canonical_goal);
        match Fulfill::new_with_simplification(self, infer, subst, goal) {
            Ok(fulfill) => (fulfill.solve(minimums), ClausePriority::High),
            Err(e) => (Err(e), ClausePriority::High),
        }
    }

    /// See whether we can solve a goal by implication on any of the given
    /// clauses. If multiple such solutions are possible, we attempt to combine
    /// them.
    fn solve_from_clauses<C>(
        &mut self,
        canonical_goal: &UCanonical<InEnvironment<DomainGoal<I>>>,
        clauses: C,
        minimums: &mut Minimums,
    ) -> (Fallible<Solution<I>>, ClausePriority)
    where
        C: IntoIterator<Item = ProgramClause<I>>,
    {
        let mut cur_solution = None;
        for program_clause in clauses {
            debug_heading!("clause={:?}", program_clause);

            // If we have a completely ambiguous answer, it's not going to get better, so stop
            if cur_solution == Some((Solution::Ambig(Guidance::Unknown), ClausePriority::High)) {
                return (Ok(Solution::Ambig(Guidance::Unknown)), ClausePriority::High);
            }

            let res = match program_clause.data(self.program.interner()) {
                ProgramClauseData::Implies(implication) => self.solve_via_implication(
                    canonical_goal,
                    &Binders::new(
                        VariableKinds::from(self.program.interner(), vec![]),
                        implication.clone(),
                    ),
                    minimums,
                ),
                ProgramClauseData::ForAll(implication) => {
                    self.solve_via_implication(canonical_goal, implication, minimums)
                }
            };
            if let (Ok(solution), priority) = res {
                debug!("ok: solution={:?} prio={:?}", solution, priority);
                cur_solution = Some(match cur_solution {
                    None => (solution, priority),
                    Some((cur, cur_priority)) => combine_with_priorities(
                        self.program.interner(),
                        &canonical_goal.canonical.value.goal,
                        cur,
                        cur_priority,
                        solution,
                        priority,
                    ),
                });
            } else {
                debug!("error");
            }
        }
        cur_solution.map_or((Err(NoSolution), ClausePriority::High), |(s, p)| (Ok(s), p))
    }

    /// Modus ponens! That is: try to apply an implication by proving its premises.
    fn solve_via_implication(
        &mut self,
        canonical_goal: &UCanonical<InEnvironment<DomainGoal<I>>>,
        clause: &Binders<ProgramClauseImplication<I>>,
        minimums: &mut Minimums,
    ) -> (Fallible<Solution<I>>, ClausePriority) {
        info_heading!(
            "solve_via_implication(\
             \n    canonical_goal={:?},\
             \n    clause={:?})",
            canonical_goal,
            clause
        );

        let (infer, subst, goal) = self.new_inference_table(canonical_goal);
        match Fulfill::new_with_clause(self, infer, subst, goal, clause) {
            Ok(fulfill) => (fulfill.solve(minimums), clause.skip_binders().priority),
            Err(e) => (Err(e), ClausePriority::High),
        }
    }

    fn program_clauses_for_goal(
        &self,
        environment: &Environment<I>,
        goal: &DomainGoal<I>,
    ) -> Result<Vec<ProgramClause<I>>, Floundered> {
        program_clauses_for_goal(self.program, environment, goal)
    }
}

impl<'me, I: Interner> RecursiveSolver<I> for Solver<'me, I> {
    /// Attempt to solve a goal that has been fully broken down into leaf form
    /// and canonicalized. This is where the action really happens, and is the
    /// place where we would perform caching in rustc (and may eventually do in Chalk).
    fn solve_goal(
        &mut self,
        goal: UCanonicalGoal<I>,
        minimums: &mut Minimums,
    ) -> Fallible<Solution<I>> {
        info_heading!("solve_goal({:?})", goal);

        // First check the cache.
        if let Some(value) = self.context.cache.get(&goal) {
            debug!("solve_reduced_goal: cache hit, value={:?}", value);
            return value.clone();
        }

        // Next, check if the goal is in the search tree already.
        if let Some(dfn) = self.context.search_graph.lookup(&goal) {
            // Check if this table is still on the stack.
            if let Some(depth) = self.context.search_graph[dfn].stack_depth {
                // Is this a coinductive goal? If so, that is success,
                // so we can return normally. Note that this return is
                // not tabled.
                //
                // XXX how does caching with coinduction work?
                if self.context.stack.coinductive_cycle_from(depth) {
                    let value = ConstrainedSubst {
                        subst: goal.trivial_substitution(self.program.interner()),
                        constraints: vec![],
                    };
                    debug!("applying coinductive semantics");
                    return Ok(Solution::Unique(Canonical {
                        value,
                        binders: goal.canonical.binders,
                    }));
                }

                self.context.stack[depth].flag_cycle();
            }

            minimums.update_from(self.context.search_graph[dfn].links);

            // Return the solution from the table.
            let previous_solution = self.context.search_graph[dfn].solution.clone();
            let previous_solution_priority = self.context.search_graph[dfn].solution_priority;
            info!(
                "solve_goal: cycle detected, previous solution {:?} with prio {:?}",
                previous_solution, previous_solution_priority
            );
            previous_solution
        } else {
            // Otherwise, push the goal onto the stack and create a table.
            // The initial result for this table is error.
            let coinductive_goal = goal.is_coinductive(self.program);
            let depth = self.context.stack.push(coinductive_goal);
            let dfn = self.context.search_graph.insert(&goal, depth);
            let subgoal_minimums = self.solve_new_subgoal(goal, depth, dfn);
            self.context.search_graph[dfn].links = subgoal_minimums;
            self.context.search_graph[dfn].stack_depth = None;
            self.context.stack.pop(depth);
            minimums.update_from(subgoal_minimums);

            // Read final result from table.
            let result = self.context.search_graph[dfn].solution.clone();
            let priority = self.context.search_graph[dfn].solution_priority;

            // If processing this subgoal did not involve anything
            // outside of its subtree, then we can promote it to the
            // cache now. This is a sort of hack to alleviate the
            // worst of the repeated work that we do during tabling.
            if subgoal_minimums.positive >= dfn {
                if self.context.caching_enabled {
                    self.context
                        .search_graph
                        .move_to_cache(dfn, &mut self.context.cache);
                    debug!("solve_reduced_goal: SCC head encountered, moving to cache");
                } else {
                    debug!(
                        "solve_reduced_goal: SCC head encountered, rolling back as caching disabled"
                    );
                    self.context.search_graph.rollback_to(dfn);
                }
            }

            info!("solve_goal: solution = {:?} prio {:?}", result, priority);
            result
        }
    }

    fn interner(&self) -> &I {
        &self.program.interner()
    }
}

fn calculate_inputs<I: Interner>(
    interner: &I,
    domain_goal: &DomainGoal<I>,
    solution: &Solution<I>,
) -> Vec<GenericArg<I>> {
    if let Some(subst) = solution.constrained_subst() {
        let subst_goal = subst.value.subst.apply(&domain_goal, interner);
        subst_goal.inputs(interner)
    } else {
        domain_goal.inputs(interner)
    }
}

fn combine_with_priorities_for_goal<I: Interner>(
    interner: &I,
    goal: &Goal<I>,
    a: Fallible<Solution<I>>,
    prio_a: ClausePriority,
    b: Fallible<Solution<I>>,
    prio_b: ClausePriority,
) -> (Fallible<Solution<I>>, ClausePriority) {
    let domain_goal = match goal.data(interner) {
        GoalData::DomainGoal(domain_goal) => domain_goal,
        _ => {
            // non-domain goals currently have no priorities, so we always take the new solution here
            return (b, prio_b);
        }
    };
    match (a, b) {
        (Ok(a), Ok(b)) => {
            let (solution, prio) =
                combine_with_priorities(interner, domain_goal, a, prio_a, b, prio_b);
            (Ok(solution), prio)
        }
        (Ok(solution), Err(_)) => (Ok(solution), prio_a),
        (Err(_), Ok(solution)) => (Ok(solution), prio_b),
        (Err(_), Err(e)) => (Err(e), prio_b),
    }
}

fn combine_with_priorities<I: Interner>(
    interner: &I,
    domain_goal: &DomainGoal<I>,
    a: Solution<I>,
    prio_a: ClausePriority,
    b: Solution<I>,
    prio_b: ClausePriority,
) -> (Solution<I>, ClausePriority) {
    match (prio_a, prio_b, a, b) {
        (ClausePriority::High, ClausePriority::Low, higher, lower)
        | (ClausePriority::Low, ClausePriority::High, lower, higher) => {
            // if we have a high-priority solution and a low-priority solution,
            // the high-priority solution overrides *if* they are both for the
            // same inputs -- we don't want a more specific high-priority
            // solution overriding a general low-priority one. Currently inputs
            // only matter for projections; in a goal like `AliasEq(<?0 as
            // Trait>::Type = ?1)`, ?0 is the input.
            let inputs_higher = calculate_inputs(interner, domain_goal, &higher);
            let inputs_lower = calculate_inputs(interner, domain_goal, &lower);
            if inputs_higher == inputs_lower {
                debug!(
                    "preferring solution: {:?} over {:?} because of higher prio",
                    higher, lower
                );
                (higher, ClausePriority::High)
            } else {
                (higher.combine(lower, interner), ClausePriority::High)
            }
        }
        (_, _, a, b) => (a.combine(b, interner), prio_a),
    }
}
