use std::{fmt, ops};

use clippy_config::Conf;
use clippy_utils::diagnostics::span_lint_and_then;
use clippy_utils::source::{HasSession, SpanRangeExt};
use clippy_utils::{fn_has_unsatisfiable_preds, is_entrypoint_fn, is_in_test};
use rustc_errors::Diag;
use rustc_hir::def_id::LocalDefId;
use rustc_hir::intravisit::FnKind;
use rustc_hir::{Body, FnDecl};
use rustc_lexer::is_ident;
use rustc_lint::{LateContext, LateLintPass};
use rustc_session::impl_lint_pass;
use rustc_span::{Span, SyntaxContext};

declare_clippy_lint! {
    /// ### What it does
    /// Checks for functions that use a lot of stack space.
    ///
    /// This often happens when constructing a large type, such as an array with a lot of elements,
    /// or constructing *many* smaller-but-still-large structs, or copying around a lot of large types.
    ///
    /// This lint is a more general version of [`large_stack_arrays`](https://rust-lang.github.io/rust-clippy/master/#large_stack_arrays)
    /// that is intended to look at functions as a whole instead of only individual array expressions inside of a function.
    ///
    /// ### Why is this bad?
    /// The stack region of memory is very limited in size (usually *much* smaller than the heap) and attempting to
    /// use too much will result in a stack overflow and crash the program.
    /// To avoid this, you should consider allocating large types on the heap instead (e.g. by boxing them).
    ///
    /// Keep in mind that the code path to construction of large types does not even need to be reachable;
    /// it purely needs to *exist* inside of the function to contribute to the stack size.
    /// For example, this causes a stack overflow even though the branch is unreachable:
    /// ```rust,ignore
    /// fn main() {
    ///     if false {
    ///         let x = [0u8; 10000000]; // 10 MB stack array
    ///         black_box(&x);
    ///     }
    /// }
    /// ```
    ///
    /// ### Known issues
    /// False positives. The stack size that clippy sees is an estimated value and can be vastly different
    /// from the actual stack usage after optimizations passes have run (especially true in release mode).
    /// Modern compilers are very smart and are able to optimize away a lot of unnecessary stack allocations.
    /// In debug mode however, it is usually more accurate.
    ///
    /// This lint works by summing up the size of all variables that the user typed, variables that were
    /// implicitly introduced by the compiler for temporaries, function arguments and the return value,
    /// and comparing them against a (configurable, but high-by-default).
    ///
    /// ### Example
    /// This function creates four 500 KB arrays on the stack. Quite big but just small enough to not trigger `large_stack_arrays`.
    /// However, looking at the function as a whole, it's clear that this uses a lot of stack space.
    /// ```no_run
    /// struct QuiteLargeType([u8; 500_000]);
    /// fn foo() {
    ///     // ... some function that uses a lot of stack space ...
    ///     let _x1 = QuiteLargeType([0; 500_000]);
    ///     let _x2 = QuiteLargeType([0; 500_000]);
    ///     let _x3 = QuiteLargeType([0; 500_000]);
    ///     let _x4 = QuiteLargeType([0; 500_000]);
    /// }
    /// ```
    ///
    /// Instead of doing this, allocate the arrays on the heap.
    /// This currently requires going through a `Vec` first and then converting it to a `Box`:
    /// ```no_run
    /// struct NotSoLargeType(Box<[u8]>);
    ///
    /// fn foo() {
    ///     let _x1 = NotSoLargeType(vec![0; 500_000].into_boxed_slice());
    /// //                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  Now heap allocated.
    /// //                                                                The size of `NotSoLargeType` is 16 bytes.
    /// //  ...
    /// }
    /// ```
    #[clippy::version = "1.72.0"]
    pub LARGE_STACK_FRAMES,
    nursery,
    "checks for functions that allocate a lot of stack space"
}

pub struct LargeStackFrames {
    maximum_allowed_size: u64,
    allow_large_stack_frames_in_tests: bool,
}

impl LargeStackFrames {
    pub fn new(conf: &'static Conf) -> Self {
        Self {
            maximum_allowed_size: conf.stack_size_threshold,
            allow_large_stack_frames_in_tests: conf.allow_large_stack_frames_in_tests,
        }
    }
}

impl_lint_pass!(LargeStackFrames => [LARGE_STACK_FRAMES]);

#[derive(Copy, Clone)]
enum Space {
    Used(u64),
    Overflow,
}

impl Space {
    pub fn exceeds_limit(self, limit: u64) -> bool {
        match self {
            Self::Used(used) => used > limit,
            Self::Overflow => true,
        }
    }
}

impl fmt::Display for Space {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Space::Used(1) => write!(f, "1 byte"),
            Space::Used(n) => write!(f, "{n} bytes"),
            Space::Overflow => write!(f, "over 2⁶⁴-1 bytes"),
        }
    }
}

impl ops::Add<u64> for Space {
    type Output = Self;
    fn add(self, rhs: u64) -> Self {
        match self {
            Self::Used(lhs) => match lhs.checked_add(rhs) {
                Some(sum) => Self::Used(sum),
                None => Self::Overflow,
            },
            Self::Overflow => self,
        }
    }
}

impl<'tcx> LateLintPass<'tcx> for LargeStackFrames {
    fn check_fn(
        &mut self,
        cx: &LateContext<'tcx>,
        fn_kind: FnKind<'tcx>,
        _: &'tcx FnDecl<'tcx>,
        _: &'tcx Body<'tcx>,
        entire_fn_span: Span,
        local_def_id: LocalDefId,
    ) {
        let def_id = local_def_id.to_def_id();
        // Building MIR for `fn`s with unsatisfiable preds results in ICE.
        if fn_has_unsatisfiable_preds(cx, def_id) {
            return;
        }

        let mir = cx.tcx.optimized_mir(def_id);
        let typing_env = mir.typing_env(cx.tcx);

        let sizes_of_locals = mir
            .local_decls
            .iter()
            .filter_map(|local| {
                let layout = cx.tcx.layout_of(typing_env.as_query_input(local.ty)).ok()?;
                Some((local, layout.size.bytes()))
            })
            .collect::<Vec<_>>();

        let frame_size = sizes_of_locals
            .iter()
            .fold(Space::Used(0), |sum, (_, size)| sum + *size);

        let limit = self.maximum_allowed_size;
        if frame_size.exceeds_limit(limit) {
            // Point at just the function name if possible, because lints that span
            // the entire body and don't have to are less legible.
            let (fn_span, fn_name) = match fn_kind {
                FnKind::ItemFn(ident, _, _) => (ident.span, format!("function `{}`", ident.name)),
                FnKind::Method(ident, _) => (ident.span, format!("method `{}`", ident.name)),
                FnKind::Closure => (entire_fn_span, "closure".to_string()),
            };

            // Don't lint inside tests if configured to not do so.
            if self.allow_large_stack_frames_in_tests && is_in_test(cx.tcx, cx.tcx.local_def_id_to_hir_id(local_def_id))
            {
                return;
            }

            let explain_lint = |diag: &mut Diag<'_, ()>, ctxt: SyntaxContext| {
                // Point out the largest individual contribution to this size, because
                // it is the most likely to be unintentionally large.
                if let Some((local, size)) = sizes_of_locals.iter().max_by_key(|&(_, size)| size)
                    && let local_span = local.source_info.span
                    && local_span.ctxt() == ctxt
                {
                    let size = Space::Used(*size); // pluralizes for us
                    let ty = local.ty;

                    // TODO: Is there a cleaner, robust way to ask this question?
                    // The obvious `LocalDecl::is_user_variable()` panics on "unwrapping cross-crate data",
                    // and that doesn't get us the true name in scope rather than the span text either.
                    if let Some(name) = local_span.get_source_text(cx)
                        && is_ident(&name)
                    {
                        // If the local is an ordinary named variable,
                        // print its name rather than relying solely on the span.
                        diag.span_label(
                            local_span,
                            format!("`{name}` is the largest part, at {size} for type `{ty}`"),
                        );
                    } else {
                        diag.span_label(
                            local_span,
                            format!("this is the largest part, at {size} for type `{ty}`"),
                        );
                    }
                }

                // Explain why we are linting this and not other functions.
                diag.note(format!(
                    "{frame_size} is larger than Clippy's configured `stack-size-threshold` of {limit}"
                ));

                // Explain why the user should care, briefly.
                diag.note_once(
                    "allocating large amounts of stack space can overflow the stack \
                        and cause the program to abort",
                );
            };

            if fn_span.from_expansion() {
                // Don't lint on the main function generated by `--test` target
                if cx.sess().is_test_crate() && is_entrypoint_fn(cx, local_def_id.to_def_id()) {
                    return;
                }

                let is_from_external_macro = fn_span.in_external_macro(cx.sess().source_map());
                span_lint_and_then(
                    cx,
                    LARGE_STACK_FRAMES,
                    fn_span.source_callsite(),
                    format!(
                        "{} generated by this macro may allocate a lot of stack space",
                        if is_from_external_macro {
                            cx.tcx.def_descr(local_def_id.into())
                        } else {
                            fn_name.as_str()
                        }
                    ),
                    |diag| {
                        if is_from_external_macro {
                            return;
                        }

                        diag.span_label(
                            fn_span,
                            format!(
                                "this {} has a stack frame size of {frame_size}",
                                cx.tcx.def_descr(local_def_id.into())
                            ),
                        );

                        explain_lint(diag, fn_span.ctxt());
                    },
                );
                return;
            }

            span_lint_and_then(
                cx,
                LARGE_STACK_FRAMES,
                fn_span,
                format!("this function may allocate {frame_size} on the stack"),
                |diag| {
                    explain_lint(diag, SyntaxContext::root());
                },
            );
        }
    }
}