/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! This cache maps from spreads (indexed by spread ids, which are generated by eval_destructor
//! and Statement.ml) to one of the lower bounds of their operands in each index and the reasons
//! that reach that spread operand.
//!
//! cache(spread_id)(resolve_idx) gets you a slice and reasons for a spread index. If any
//! two indices have more than one lower bound, then we will get at least quadratic results
//! by computing the spreads. Instead of going through with the computation, we emit an error.
//!
//! We need to track the slice explicitly in addition to the list of reasons because
//! only count an index as having multiple lower bounds if those lower bounds are actually
//! different. This is important: if an early index receives multiple lower bounds, it will
//! kick off a resolution job for each. Both of those resolution jobs share a spread id,
//! so it's important that we don't double count a later operand being resolved as multiple
//! lower bounds when it's just resolved twice because of the separate resolution jobs.

use std::cell::RefCell;

use dupe::Dupe;
use flow_aloc::ALoc;
use flow_common::reason::VirtualReason;
use flow_data_structure_wrapper::ord_map::FlowOrdMap;
use flow_typing_errors::intermediate_error_types::ExponentialSpreadReasonGroup;
use flow_typing_type::type_::object::Slice as ObjectSlice;
use vec1::Vec1;

pub type ReasonState = ExponentialSpreadReasonGroup<ALoc>;

#[derive(Debug, Clone)]
pub enum CacheState {
    Ok {
        resolve_idx: Option<i32>,
        slices: FlowOrdMap<i32, (ObjectSlice, ReasonState)>,
    },
    Error(ReasonState, ReasonState),
}

#[derive(Debug, Clone, Default, Dupe)]
pub struct SpreadCache(FlowOrdMap<i32, CacheState>);

impl SpreadCache {
    pub fn new() -> Self {
        thread_local! {
            static CACHED: FlowOrdMap<i32, CacheState> = FlowOrdMap::new();
        }
        SpreadCache(CACHED.with(|c| c.clone()))
    }
}

pub fn add_lower_bound(
    cache: &RefCell<SpreadCache>,
    spread_id: i32,
    resolve_idx: i32,
    r: VirtualReason<ALoc>,
    objtypes: &Vec1<ObjectSlice>,
) {
    let state = cache
        .borrow_mut()
        .0
        .remove(&spread_id)
        .unwrap_or_else(|| CacheState::Ok {
            resolve_idx: None,
            slices: FlowOrdMap::new(),
        });
    let state_prime = match state {
        CacheState::Error(g1, g2) => CacheState::Error(g1, g2),
        CacheState::Ok {
            resolve_idx: idx_option,
            slices: mut map,
        } => {
            let (new_entry, has_multiple_lower_bounds) = {
                match map.remove(&resolve_idx) {
                    None => {
                        let entry = (
                            objtypes.first().clone(),
                            ExponentialSpreadReasonGroup {
                                first_reason: r.dupe(),
                                second_reason: None,
                            },
                        );
                        let has_multi = objtypes.len() != 1;
                        (entry, has_multi)
                    }
                    Some((slice, reason_state)) => {
                        let has_new_lower_bound = objtypes.iter().any(|s| &slice != s);
                        let reason_state = ExponentialSpreadReasonGroup {
                            first_reason: reason_state.first_reason.dupe(),
                            second_reason: Some(r.dupe()),
                        };
                        ((slice, reason_state), has_new_lower_bound)
                    }
                }
            };
            map.insert(resolve_idx, new_entry);
            // If we have multiple lower bounds and this isn't the only idx with multiple
            // lower bounds then we enter our error state
            if has_multiple_lower_bounds {
                match idx_option {
                    Some(idx) if idx != resolve_idx => {
                        let group1 = map.get(&resolve_idx).unwrap().1.dupe();
                        let group2 = map.get(&idx).unwrap().1.dupe();
                        CacheState::Error(group1, group2)
                    }
                    _ => CacheState::Ok {
                        resolve_idx: Some(resolve_idx),
                        slices: map,
                    },
                }
            } else {
                CacheState::Ok {
                    resolve_idx: idx_option,
                    slices: map,
                }
            }
        }
    };
    cache.borrow_mut().0.insert(spread_id, state_prime);
}

pub fn get_error_groups(
    cache: &RefCell<SpreadCache>,
    spread_id: i32,
) -> (ReasonState, ReasonState) {
    match cache.borrow().0.get(&spread_id) {
        Some(CacheState::Error(group1, group2)) => (group1.dupe(), group2.dupe()),
        _ => panic!(
            "Invariant violation: make sure can_spread is false before calling get_error_groups"
        ),
    }
}

pub fn can_spread(cache: &RefCell<SpreadCache>, spread_id: i32) -> bool {
    match cache.borrow().0.get(&spread_id) {
        Some(CacheState::Error(_, _)) => false,
        _ => true,
    }
}
