/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

pub(crate) mod payload;
mod tests;

use std::cmp;
use std::collections::HashMap;
use std::iter;
use std::marker::PhantomData;
use std::mem;
use std::slice;

use dupe::Dupe;
use starlark_derive::VisitSpanMut;
use starlark_map::small_map;
use starlark_map::small_map::SmallMap;

use crate::codemap::CodeMap;
use crate::codemap::Span;
use crate::environment::names::MutableNames;
use crate::environment::slots::ModuleSlotId;
use crate::environment::Globals;
use crate::environment::Module;
use crate::errors::did_you_mean::did_you_mean;
use crate::eval::compiler::def::CopySlotFromParent;
use crate::eval::compiler::scope::payload::CstAssign;
use crate::eval::compiler::scope::payload::CstAssignIdent;
use crate::eval::compiler::scope::payload::CstExpr;
use crate::eval::compiler::scope::payload::CstIdent;
use crate::eval::compiler::scope::payload::CstParameter;
use crate::eval::compiler::scope::payload::CstPayload;
use crate::eval::compiler::scope::payload::CstStmt;
use crate::eval::compiler::scope::payload::CstTypeExpr;
use crate::eval::compiler::EvalException;
use crate::eval::runtime::slots::LocalSlotIdCapturedOrNot;
use crate::syntax::ast::Assign;
use crate::syntax::ast::AssignIdent;
use crate::syntax::ast::AstAssignIdentP;
use crate::syntax::ast::AstStmt;
use crate::syntax::ast::ClauseP;
use crate::syntax::ast::DefP;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::ForClauseP;
use crate::syntax::ast::LambdaP;
use crate::syntax::ast::Stmt;
use crate::syntax::ast::StmtP;
use crate::syntax::ast::Visibility;
use crate::syntax::uniplate::VisitMut;
use crate::syntax::Dialect;
use crate::typing::Interface;
use crate::values::FrozenHeap;
use crate::values::FrozenRef;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;

#[derive(Debug, thiserror::Error)]
enum ScopeError {
    #[error("Variable `{0}` not found")]
    VariableNotFound(String),
    #[error("Variable `{0}` not found, did you mean `{1}`?")]
    VariableNotFoundDidYouMean(String, String),
    #[error("Identifiers in type expressions can only refer globals or builtins: `{0}`")]
    TypeExpressionGlobalOrBuiltin(String),
}

/// All scopes and bindings in a module.
struct ModuleScopeBuilder<'a> {
    scope_data: ModuleScopeData<'a>,
    module: &'a MutableNames,
    frozen_heap: &'a FrozenHeap,
    module_bindings: SmallMap<FrozenStringValue, BindingId>,
    // The first scope is a module-level scope (including comprehensions in module scope).
    // The rest are scopes for functions (which include their comprehensions).
    locals: Vec<ScopeId>,
    unscopes: Vec<Unscope>,
    codemap: FrozenRef<'static, CodeMap>,
    globals: FrozenRef<'static, Globals>,
    errors: Vec<EvalException>,
    top_level_stmt_count: usize,
    /// Index of last statement that defines a type.
    last_stmt_defining_type: Option<TopLevelStmtIndex>,
}

pub(crate) struct ModuleScopes<'f> {
    pub(crate) scope_data: ModuleScopeData<'f>,
    pub(crate) module_slot_count: u32,
    pub(crate) cst: CstStmt,
    pub(crate) module_bindings: SmallMap<FrozenStringValue, BindingId>,
    /// Number of top-level statements in the module.
    pub(crate) top_level_stmt_count: usize,
    /// Index of last statement that defines a type.
    pub(crate) last_stmt_defining_type: Option<TopLevelStmtIndex>,
}

struct UnscopeBinding {
    /// Variable mappings in local scope are overwritten by comprehension variables.
    ///
    /// When we enter comprehension, we replace local scope variable slots with comprehension
    /// scope slots. This field stores the original slot in the local scope,
    /// or `None` if there was no mapping for the variable.
    ///
    /// When we pop the comprehension scope, we restore the mapping from this value.
    undo: Option<(LocalSlotIdCapturedOrNot, BindingId)>,
}

#[derive(Default)]
struct Unscope(SmallMap<FrozenStringValue, UnscopeBinding>);

#[derive(Default, Debug)]
pub(crate) struct ScopeNames<'f> {
    /// `Some` when scope is initialized.
    /// For module scope, the value is zero.
    pub param_count: Option<u32>,
    /// Slots this scope uses, including for parameters and `parent`.
    /// Indexed by [`LocalSlotId`], values are variable names.
    pub used: Vec<FrozenStringValue>,
    /// The names that are in this scope
    pub mp: SmallMap<FrozenStringValue, (LocalSlotIdCapturedOrNot, BindingId)>,
    /// Slots to copy from the parent.
    /// Module-level identifiers are not copied over, to avoid excess copying.
    pub parent: Vec<CopySlotFromParent>,
    /// We store frozen strings.
    _heap: PhantomData<&'f ()>,
}

impl<'f> ScopeNames<'f> {
    fn set_param_count(&mut self, param_count: u32) {
        assert!(self.param_count.is_none());
        self.param_count = Some(param_count);
    }

    pub(crate) fn param_count(&self) -> u32 {
        self.param_count
            .expect("param_count must be set during analysis")
    }

    fn copy_parent(
        &mut self,
        parent_slot: LocalSlotIdCapturedOrNot,
        binding_id: BindingId,
        name: FrozenStringValue,
    ) -> LocalSlotIdCapturedOrNot {
        assert!(self.get_name(name).is_none()); // Or we'll be overwriting our variable
        let res = self.add_name(name, binding_id);
        self.parent.push(CopySlotFromParent {
            parent: parent_slot,
            child: res,
        });
        res
    }

    fn next_slot(&mut self, name: FrozenStringValue) -> LocalSlotIdCapturedOrNot {
        let res = LocalSlotIdCapturedOrNot(self.used.len().try_into().unwrap());
        self.used.push(name);
        res
    }

    fn add_name(
        &mut self,
        name: FrozenStringValue,
        binding_id: BindingId,
    ) -> LocalSlotIdCapturedOrNot {
        let slot = self.next_slot(name);
        let old = self.mp.insert_hashed(name.get_hashed(), (slot, binding_id));
        assert!(old.is_none());
        slot
    }

    fn add_scoped(
        &mut self,
        name: FrozenStringValue,
        binding_id: BindingId,
        unscope: &mut Unscope,
    ) -> LocalSlotIdCapturedOrNot {
        let slot = self.next_slot(name);
        let undo = match self.mp.get_mut_hashed(name.get_hashed().as_ref()) {
            Some(v) => {
                let old = *v;
                *v = (slot, binding_id);
                Some(old)
            }
            None => {
                self.mp.insert_hashed(name.get_hashed(), (slot, binding_id));
                None
            }
        };
        assert!(
            unscope
                .0
                .insert_hashed(name.get_hashed(), UnscopeBinding { undo })
                .is_none()
        );
        slot
    }

    fn unscope(&mut self, unscope: Unscope) {
        for (name, UnscopeBinding { undo }) in unscope.0 {
            match undo {
                None => {
                    self.mp.remove(&name);
                }
                Some(v) => *self.mp.get_mut(&name).unwrap() = v,
            }
        }
    }

    fn get_name(&self, name: FrozenStringValue) -> Option<(LocalSlotIdCapturedOrNot, BindingId)> {
        self.mp.get_hashed(name.get_hashed().as_ref()).copied()
    }
}

#[derive(Copy, Clone, Dupe, Debug)]
pub(crate) enum Slot {
    /// Top-level module scope.
    Module(ModuleSlotId),
    /// Local scope, always mutable.
    Local(LocalSlotIdCapturedOrNot),
}

#[derive(Clone, Copy, Dupe)]
enum ResolveIdentScope {
    /// Resolving normal identifier.
    Any,
    /// Resolving identifier in type expression.
    GlobalForTypeExpression,
}

impl<'f> ModuleScopeBuilder<'f> {
    fn top_scope_id(&self) -> ScopeId {
        *self.locals.last().unwrap()
    }

    fn scope_at_level(&self, level: usize) -> &ScopeNames<'f> {
        let scope_id = self.locals[level];
        self.scope_data.get_scope(scope_id)
    }

    fn scope_at_level_mut(&mut self, level: usize) -> &mut ScopeNames<'f> {
        let scope_id = self.locals[level];
        self.scope_data.mut_scope(scope_id)
    }

    /// Resolve symbols in a module.
    ///
    /// Checks all the symbols are resolved to locals/globals/captured/etc.
    /// Do not check types yet. But validate type expressions.
    ///
    /// This function does not fail, errors are stored in the `errors` field.
    fn enter_module(
        module: &'f MutableNames,
        frozen_heap: &'f FrozenHeap,
        loads: &HashMap<String, Interface>,
        stmt: AstStmt,
        globals: FrozenRef<'static, Globals>,
        codemap: FrozenRef<'static, CodeMap>,
        dialect: &Dialect,
    ) -> (CstStmt, ModuleScopeBuilder<'f>) {
        let mut scope_data = ModuleScopeData::new();
        let scope_id = scope_data.new_scope().0;
        let mut cst = CstStmt::from_ast(stmt, &mut scope_data, loads);

        let top_level_stmts: &mut [CstStmt] = match &mut cst.node {
            StmtP::Statements(stmts) => {
                // TODO(nga): single-line top-level statements like `a(); b()`
                //   will be treated as multiple top-level statements.
                stmts
            }
            _ => slice::from_mut(&mut cst),
        };

        // Not really important, sanity check
        assert_eq!(scope_id, ScopeId::module());

        scope_data.mut_scope(scope_id).set_param_count(0);

        let mut locals: SmallMap<FrozenStringValue, _> = SmallMap::new();

        let existing_module_names_and_visibilites = module.all_names_and_visibilities();
        for (name, vis) in existing_module_names_and_visibilites.iter() {
            let (binding_id, _binding) = scope_data.new_binding(
                *name,
                BindingSource::FromModule,
                *vis,
                AssignCount::AtMostOnce,
            );
            locals.insert_hashed(name.get_hashed(), binding_id);
        }

        for (i, stmt) in top_level_stmts.iter_mut().enumerate() {
            Stmt::collect_defines(
                stmt,
                TopLevelStmtIndex(i),
                InLoop::No,
                &mut scope_data,
                frozen_heap,
                &mut locals,
                dialect,
            );
        }

        let mut module_bindings = SmallMap::new();
        for (x, binding_id) in locals {
            let binding = scope_data.mut_binding(binding_id);
            let slot = module.add_name_visibility(x, binding.vis);
            let old_slot = mem::replace(&mut binding.slot, Some(Slot::Module(slot)));
            assert!(old_slot.is_none());
            let old_binding = module_bindings.insert_hashed(x.get_hashed(), binding_id);
            assert!(old_binding.is_none());
        }

        // Here we traverse the AST second time to collect scopes of defs
        for (i, stmt) in top_level_stmts.iter_mut().enumerate() {
            ModuleScopeBuilder::collect_defines_recursively(
                &mut scope_data,
                stmt,
                TopLevelStmtIndex(i),
                frozen_heap,
                dialect,
            );
        }
        let mut scope = ModuleScopeBuilder {
            scope_data,
            frozen_heap,
            module,
            module_bindings,
            locals: vec![scope_id],
            unscopes: Vec::new(),
            codemap,
            globals,
            errors: Vec::new(),
            top_level_stmt_count: top_level_stmts.len(),
            last_stmt_defining_type: None,
        };
        for (i, stmt) in top_level_stmts.iter_mut().enumerate() {
            scope.resolve_idents(stmt, TopLevelStmtIndex(i));
        }
        (cst, scope)
    }
}

impl<'f> ModuleScopeBuilder<'f> {
    // Number of module slots I need, a struct holding all scopes, and module bindings.
    fn exit_module(
        mut self,
    ) -> (
        u32,
        ModuleScopeData<'f>,
        SmallMap<FrozenStringValue, BindingId>,
    ) {
        assert!(self.locals.len() == 1);
        assert!(self.unscopes.is_empty());
        let scope_id = self.locals.pop().unwrap();
        assert!(scope_id == ScopeId::module());
        let scope = self.scope_data.get_scope(scope_id);
        assert!(scope.parent.is_empty());
        (
            self.module.slot_count(),
            self.scope_data,
            self.module_bindings,
        )
    }
}

impl<'f> ModuleScopes<'f> {
    pub(crate) fn check_module_err(
        module: &'f MutableNames,
        frozen_heap: &'f FrozenHeap,
        loads: &HashMap<String, Interface>,
        stmt: AstStmt,
        globals: FrozenRef<'static, Globals>,
        codemap: FrozenRef<'static, CodeMap>,
        dialect: &Dialect,
    ) -> anyhow::Result<ModuleScopes<'f>> {
        let (errors, scopes) =
            ModuleScopes::check_module(module, frozen_heap, loads, stmt, globals, codemap, dialect);
        if let Some(error) = errors.into_iter().next() {
            return Err(error.into_anyhow());
        }
        Ok(scopes)
    }

    pub(crate) fn check_module(
        module: &'f MutableNames,
        frozen_heap: &'f FrozenHeap,
        loads: &HashMap<String, Interface>,
        stmt: AstStmt,
        globals: FrozenRef<'static, Globals>,
        codemap: FrozenRef<'static, CodeMap>,
        dialect: &Dialect,
    ) -> (Vec<EvalException>, ModuleScopes<'f>) {
        let (stmt, mut scope) = ModuleScopeBuilder::enter_module(
            module,
            frozen_heap,
            loads,
            stmt,
            globals,
            codemap,
            dialect,
        );
        let top_level_stmt_count = scope.top_level_stmt_count;
        let last_stmt_defining_type = scope.last_stmt_defining_type;
        let errors = mem::take(&mut scope.errors);
        let (module_slot_count, scope_data, module_bindings) = scope.exit_module();
        (
            errors,
            ModuleScopes {
                cst: stmt,
                scope_data,
                module_bindings,
                module_slot_count,
                top_level_stmt_count,
                last_stmt_defining_type,
            },
        )
    }
}

impl<'f> ModuleScopeBuilder<'f> {
    fn collect_defines_in_def(
        scope_data: &mut ModuleScopeData,
        scope_id: ScopeId,
        params: &mut [CstParameter],
        body: Option<&mut CstStmt>,
        top_level_stmt_index: TopLevelStmtIndex,
        frozen_heap: &FrozenHeap,
        dialect: &Dialect,
    ) {
        let params: Vec<&mut AstAssignIdentP<_>> = params
            .iter_mut()
            .filter_map(|p| p.node.split_mut().0)
            .collect::<Vec<_>>();
        scope_data
            .mut_scope(scope_id)
            .set_param_count(params.len().try_into().unwrap());
        let mut locals: SmallMap<FrozenStringValue, _> = SmallMap::new();
        for p in params {
            let name = frozen_heap.alloc_str_intern(&p.0);
            // Subtle invariant: the slots for the params must be ordered and at the
            // beginning
            let binding_id = scope_data
                .new_binding(
                    name,
                    BindingSource::Source(p.span, top_level_stmt_index),
                    Visibility::Public,
                    AssignCount::AtMostOnce,
                )
                .0;
            p.1 = Some(binding_id);
            let old_local = locals.insert_hashed(name.get_hashed(), binding_id);
            assert!(old_local.is_none());
        }
        if let Some(code) = body {
            Stmt::collect_defines(
                code,
                top_level_stmt_index,
                InLoop::No,
                scope_data,
                frozen_heap,
                &mut locals,
                dialect,
            );
        }
        for (name, binding_id) in locals.into_iter() {
            let slot = scope_data.mut_scope(scope_id).add_name(name, binding_id);
            let binding = scope_data.mut_binding(binding_id);
            let old_slot = mem::replace(&mut binding.slot, Some(Slot::Local(slot)));
            assert!(old_slot.is_none());
        }
    }

    fn collect_defines_recursively(
        scope_data: &mut ModuleScopeData,
        code: &mut CstStmt,
        top_level_stmt_index: TopLevelStmtIndex,
        frozen_heap: &FrozenHeap,
        dialect: &Dialect,
    ) {
        if let StmtP::Def(DefP {
            name: _,
            params,
            return_type: _,
            body,
            payload: scope_id,
        }) = &mut code.node
        {
            // Here we traverse the AST twice: once for this def scope,
            // second time below for nested defs.
            Self::collect_defines_in_def(
                scope_data,
                *scope_id,
                params,
                Some(body),
                top_level_stmt_index,
                frozen_heap,
                dialect,
            );
        }

        code.visit_children_mut(&mut |visit| match visit {
            VisitMut::Expr(e) => Self::collect_defines_recursively_in_expr(
                scope_data,
                e,
                top_level_stmt_index,
                frozen_heap,
                dialect,
            ),
            VisitMut::Stmt(s) => Self::collect_defines_recursively(
                scope_data,
                s,
                top_level_stmt_index,
                frozen_heap,
                dialect,
            ),
        });
    }

    fn collect_defines_recursively_in_expr(
        scope_data: &mut ModuleScopeData,
        code: &mut CstExpr,
        top_level_stmt_index: TopLevelStmtIndex,
        frozen_heap: &FrozenHeap,
        dialect: &Dialect,
    ) {
        if let ExprP::Lambda(LambdaP {
            params,
            body: _,
            payload: scope_id,
        }) = &mut code.node
        {
            Self::collect_defines_in_def(
                scope_data,
                *scope_id,
                params,
                None,
                top_level_stmt_index,
                frozen_heap,
                dialect,
            );
        }

        code.visit_expr_mut(|e| {
            Self::collect_defines_recursively_in_expr(
                scope_data,
                e,
                top_level_stmt_index,
                frozen_heap,
                dialect,
            )
        });
    }

    fn resolve_idents(&mut self, code: &mut CstStmt, top_level_stmt_index: TopLevelStmtIndex) {
        match &mut code.node {
            StmtP::Def(DefP {
                name: _,
                params,
                return_type,
                body,
                payload: scope_id,
            }) => self.resolve_idents_in_def(
                *scope_id,
                params,
                return_type.as_mut().map(|r| &mut **r),
                Some(body),
                None,
                top_level_stmt_index,
            ),
            StmtP::Assign(lhs, ty_rhs) => {
                let (ty, rhs) = &mut **ty_rhs;
                self.resolve_idents_in_assign(lhs, top_level_stmt_index);
                if let Some(ty) = ty {
                    self.resolve_idents_in_type_expr(ty, top_level_stmt_index);
                }
                self.resolve_idents_in_expr(rhs, top_level_stmt_index);
            }
            _ => code.visit_children_mut(|visit| match visit {
                VisitMut::Stmt(stmt) => self.resolve_idents(stmt, top_level_stmt_index),
                VisitMut::Expr(expr) => self.resolve_idents_in_expr(expr, top_level_stmt_index),
            }),
        }
    }

    fn resolve_idents_in_assign(
        &mut self,
        assign: &mut CstAssign,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        assign.visit_expr_mut(|expr| self.resolve_idents_in_expr(expr, top_level_stmt_index));
    }

    fn resolve_idents_in_def(
        &mut self,
        scope_id: ScopeId,
        params: &mut [CstParameter],
        ret: Option<&mut CstTypeExpr>,
        body_stmt: Option<&mut CstStmt>,
        body_expr: Option<&mut CstExpr>,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        for param in params {
            let (_, ty, def) = param.split_mut();
            if let Some(ty) = ty {
                self.resolve_idents_in_type_expr(ty, top_level_stmt_index);
            }
            if let Some(def) = def {
                self.resolve_idents_in_expr(def, top_level_stmt_index);
            }
        }
        if let Some(ret) = ret {
            self.resolve_idents_in_type_expr(ret, top_level_stmt_index);
        }

        self.enter_def(scope_id);
        if let Some(body_stmt) = body_stmt {
            self.resolve_idents(body_stmt, top_level_stmt_index);
        }
        if let Some(body_expr) = body_expr {
            self.resolve_idents_in_expr(body_expr, top_level_stmt_index);
        }
        self.exit_def();
    }

    fn resolve_idents_in_expr_impl(
        &mut self,
        scope: ResolveIdentScope,
        expr: &mut CstExpr,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        match &mut expr.node {
            ExprP::Identifier(ident) => self.resolve_ident(scope, ident),
            ExprP::Lambda(LambdaP {
                params,
                body,
                payload: scope_id,
            }) => self.resolve_idents_in_def(
                *scope_id,
                params,
                None,
                None,
                Some(body),
                top_level_stmt_index,
            ),
            ExprP::ListComprehension(expr, first_for, clauses) => {
                self.resolve_idents_in_compr(&mut [expr], first_for, clauses, top_level_stmt_index)
            }
            ExprP::DictComprehension(k_v, first_for, clauses) => {
                let (k, v) = &mut **k_v;
                self.resolve_idents_in_compr(&mut [k, v], first_for, clauses, top_level_stmt_index)
            }
            _ => expr.visit_expr_mut(|expr| {
                self.resolve_idents_in_expr_impl(scope, expr, top_level_stmt_index)
            }),
        }
    }

    fn resolve_idents_in_expr(
        &mut self,
        expr: &mut CstExpr,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        self.resolve_idents_in_expr_impl(ResolveIdentScope::Any, expr, top_level_stmt_index);
    }

    fn resolve_idents_in_type_expr(
        &mut self,
        expr: &mut CstTypeExpr,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        self.resolve_idents_in_expr_impl(
            ResolveIdentScope::GlobalForTypeExpression,
            &mut expr.node.expr,
            top_level_stmt_index,
        );
    }

    fn current_scope_all_visible_names_for_did_you_mean(&self) -> Vec<FrozenStringValue> {
        // It is OK to return non-unique identifiers
        let mut r: Vec<FrozenStringValue> = Vec::new();
        for &scope_id in self.locals.iter().rev() {
            let scope = self.scope_data.get_scope(scope_id);
            r.extend(scope.mp.keys().copied());
        }
        r.extend(self.module_bindings.keys().copied());
        r.extend(self.globals.names());
        r
    }

    fn variable_not_found_err(&self, ident: &CstIdent) -> EvalException {
        let variants = self.current_scope_all_visible_names_for_did_you_mean();
        let better = did_you_mean(ident.node.0.as_str(), variants.iter().map(|s| s.as_str()));
        EvalException::new(
            match better {
                Some(better) => {
                    ScopeError::VariableNotFoundDidYouMean(ident.node.0.clone(), better.to_owned())
                }
                None => ScopeError::VariableNotFound(ident.node.0.clone()),
            }
            .into(),
            ident.span,
            &self.codemap,
        )
    }

    fn update_latest_binding_used_in_type_expr(&mut self, binding_id: BindingId) {
        let binding = self.scope_data.get_binding(binding_id);
        let top_level_stmt_index = match binding.source {
            BindingSource::Source(_, top_level_stmt_index) => top_level_stmt_index,
            BindingSource::FromModule => return,
        };
        match &mut self.last_stmt_defining_type {
            None => self.last_stmt_defining_type = Some(top_level_stmt_index),
            Some(last_stmt) => *last_stmt = cmp::max(*last_stmt, top_level_stmt_index),
        }
    }

    fn resolve_ident(&mut self, scope: ResolveIdentScope, ident: &mut CstIdent) {
        assert!(ident.node.1.is_none());
        let resolved = match self.get_name(self.frozen_heap.alloc_str_intern(&ident.node.0)) {
            None => {
                // Must be a global, since we know all variables
                match self.globals.get_frozen(&ident.node.0) {
                    None => {
                        self.errors.push(self.variable_not_found_err(ident));
                        return;
                    }
                    Some(v) => ResolvedIdent::Global(v),
                }
            }
            Some((slot, binding_id)) => ResolvedIdent::Slot(slot, binding_id),
        };
        match scope {
            ResolveIdentScope::Any => {}
            ResolveIdentScope::GlobalForTypeExpression => match resolved {
                ResolvedIdent::Slot(Slot::Local(_), _) => {
                    self.errors.push(EvalException::new(
                        ScopeError::TypeExpressionGlobalOrBuiltin(ident.node.0.clone()).into(),
                        ident.span,
                        &self.codemap,
                    ));
                    return;
                }
                ResolvedIdent::Slot(Slot::Module(_), binding_id) => {
                    self.update_latest_binding_used_in_type_expr(binding_id);
                }
                ResolvedIdent::Global(_) => {}
            },
        }
        ident.node.1 = Some(resolved);
    }

    fn resolve_idents_in_compr(
        &mut self,
        exprs: &mut [&mut CstExpr],
        first_for: &mut ForClauseP<CstPayload>,
        clauses: &mut [ClauseP<CstPayload>],
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        // First for is resolved in outer scope
        self.resolve_idents_in_for_clause(first_for, top_level_stmt_index);

        self.enter_compr();

        // Add identifiers to compr scope

        self.add_compr(
            iter::once(&mut first_for.var).chain(clauses.iter_mut().filter_map(
                |clause| match clause {
                    ClauseP::For(for_clause) => Some(&mut for_clause.var),
                    ClauseP::If(..) => None,
                },
            )),
            top_level_stmt_index,
        );

        // Now resolve idents in compr scope

        for clause in clauses.iter_mut() {
            match clause {
                ClauseP::For(for_clause) => {
                    self.resolve_idents_in_for_clause(for_clause, top_level_stmt_index)
                }
                ClauseP::If(cond) => self.resolve_idents_in_expr(cond, top_level_stmt_index),
            }
        }

        // Finally, resolve the item expression

        for expr in exprs {
            self.resolve_idents_in_expr(expr, top_level_stmt_index);
        }

        self.exit_compr();
    }

    fn resolve_idents_in_for_clause(
        &mut self,
        for_clause: &mut ForClauseP<CstPayload>,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        self.resolve_idents_in_expr(&mut for_clause.over, top_level_stmt_index);
        self.resolve_idents_in_assign(&mut for_clause.var, top_level_stmt_index);
    }

    pub fn enter_def(&mut self, scope_id: ScopeId) {
        assert!(scope_id != ScopeId::module());
        self.locals.push(scope_id);
    }

    // Which slots to grab from the current scope to the parent scope, size of your
    // self scope Future state: Should return the slots to use from the parent
    // scope
    pub fn exit_def(&mut self) -> &mut ScopeNames<'f> {
        let scope_id = self.locals.pop().unwrap();
        self.scope_data.mut_scope(scope_id)
    }

    fn enter_compr(&mut self) {
        self.unscopes.push(Unscope::default());
    }

    fn add_compr<'x>(
        &mut self,
        var: impl IntoIterator<Item = &'x mut CstAssign>,
        top_level_stmt_index: TopLevelStmtIndex,
    ) {
        let scope_id = self.top_scope_id();
        let mut locals = SmallMap::new();
        for var in var {
            Assign::collect_defines_lvalue(
                var,
                top_level_stmt_index,
                InLoop::Yes,
                &mut self.scope_data,
                self.frozen_heap,
                &mut locals,
            );
        }
        for (name, binding_id) in locals.into_iter() {
            let slot = self.scope_data.mut_scope(scope_id).add_scoped(
                name,
                binding_id,
                self.unscopes.last_mut().unwrap(),
            );
            let binding = self.scope_data.mut_binding(binding_id);
            assert!(mem::replace(&mut binding.slot, Some(Slot::Local(slot))).is_none());
        }
    }

    fn exit_compr(&mut self) {
        self.scope_data
            .mut_scope(self.top_scope_id())
            .unscope(self.unscopes.pop().unwrap());
    }

    fn get_name(&mut self, name: FrozenStringValue) -> Option<(Slot, BindingId)> {
        // look upwards to find the first place the variable occurs
        // then copy that variable downwards
        for i in (0..self.locals.len()).rev() {
            if let Some((mut v, binding_id)) = self.scope_at_level(i).get_name(name) {
                if i + 1 != self.locals.len() {
                    self.scope_data.mut_binding(binding_id).captured = Captured::Yes;
                }
                for j in (i + 1)..self.locals.len() {
                    v = self.scope_at_level_mut(j).copy_parent(v, binding_id, name);
                }
                return Some((Slot::Local(v), binding_id));
            }
        }
        let binding_id = self
            .module_bindings
            .get_hashed(name.get_hashed().as_ref())
            .copied();
        match binding_id {
            Some(binding_id) => {
                let binding = self.scope_data.mut_binding(binding_id);
                if self.locals.len() > 1 {
                    binding.captured = Captured::Yes;
                }
                Some((binding.slot.unwrap(), binding_id))
            }
            None => None,
        }
    }
}

/// While performing analysis.
#[derive(Copy, Clone, Dupe)]
enum InLoop {
    /// Current statement has an enclosing loop in the current scope.
    Yes,
    /// Current statement has no enclosing loop in the current scope.
    No,
}

impl Stmt {
    // Collect all the variables that are defined in this scope
    fn collect_defines<'a>(
        stmt: &'a mut CstStmt,
        top_level_stmt_index: TopLevelStmtIndex,
        in_loop: InLoop,
        scope_data: &mut ModuleScopeData,
        frozen_heap: &FrozenHeap,
        result: &mut SmallMap<FrozenStringValue, BindingId>,
        dialect: &Dialect,
    ) {
        match &mut stmt.node {
            StmtP::Assign(dest, _) | StmtP::AssignModify(dest, _, _) => {
                Assign::collect_defines_lvalue(
                    dest,
                    top_level_stmt_index,
                    in_loop,
                    scope_data,
                    frozen_heap,
                    result,
                );
            }
            StmtP::For(dest, over_body) => {
                let (_over, body) = &mut **over_body;
                Assign::collect_defines_lvalue(
                    dest,
                    top_level_stmt_index,
                    InLoop::Yes,
                    scope_data,
                    frozen_heap,
                    result,
                );
                StmtP::collect_defines(
                    body,
                    top_level_stmt_index,
                    InLoop::Yes,
                    scope_data,
                    frozen_heap,
                    result,
                    dialect,
                );
            }
            StmtP::Def(DefP { name, .. }) => AssignIdent::collect_assign_ident(
                name,
                top_level_stmt_index,
                in_loop,
                Visibility::Public,
                scope_data,
                frozen_heap,
                result,
            ),
            StmtP::Load(load) => {
                let vis = dialect.load_visibility();
                for (name, _) in &mut load.args {
                    let mut vis = vis;
                    if Module::default_visibility(&name.0) == Visibility::Private {
                        vis = Visibility::Private;
                    }
                    AssignIdent::collect_assign_ident(
                        name,
                        top_level_stmt_index,
                        in_loop,
                        vis,
                        scope_data,
                        frozen_heap,
                        result,
                    );
                }
            }
            stmt => stmt.visit_stmt_mut(|x| {
                Stmt::collect_defines(
                    x,
                    top_level_stmt_index,
                    in_loop,
                    scope_data,
                    frozen_heap,
                    result,
                    dialect,
                )
            }),
        }
    }
}

impl AssignIdent {
    fn collect_assign_ident<'a>(
        assign: &'a mut CstAssignIdent,
        top_level_stmt_index: TopLevelStmtIndex,
        in_loop: InLoop,
        vis: Visibility,
        scope_data: &mut ModuleScopeData,
        frozen_heap: &FrozenHeap,
        result: &mut SmallMap<FrozenStringValue, BindingId>,
    ) {
        // Helper function to untangle lifetimes: we read and modify `assign` fields.
        fn assign_ident_impl<'b>(
            name: FrozenStringValue,
            span: Span,
            top_level_stmt_index: TopLevelStmtIndex,
            binding: &'b mut Option<BindingId>,
            in_loop: InLoop,
            mut vis: Visibility,
            scope_data: &mut ModuleScopeData,
            result: &mut SmallMap<FrozenStringValue, BindingId>,
        ) {
            assert!(
                binding.is_none(),
                "binding can be assigned only once: `{}`",
                name.as_str()
            );
            if vis == Visibility::Public {
                vis = Module::default_visibility(&name);
            }
            match result.entry_hashed(name.get_hashed()) {
                small_map::Entry::Occupied(e) => {
                    let prev_binding_id = *e.get();
                    let prev_binding = scope_data.mut_binding(prev_binding_id);
                    // If we are in the map as Public and Private, then Public wins.
                    // Everything but Load is definitely Public.
                    // So only insert if it wasn't already there.
                    if vis == Visibility::Public {
                        prev_binding.vis = Visibility::Public;
                    }
                    prev_binding.assign_count = AssignCount::Any;
                    *binding = Some(prev_binding_id);
                }
                small_map::Entry::Vacant(e) => {
                    let assign_count = match in_loop {
                        InLoop::Yes => AssignCount::Any,
                        InLoop::No => AssignCount::AtMostOnce,
                    };
                    let (new_binding_id, _) = scope_data.new_binding(
                        name,
                        BindingSource::Source(span, top_level_stmt_index),
                        vis,
                        assign_count,
                    );
                    e.insert(new_binding_id);
                    *binding = Some(new_binding_id);
                }
            };
        }
        assign_ident_impl(
            frozen_heap.alloc_str_intern(&assign.node.0),
            assign.span,
            top_level_stmt_index,
            &mut assign.node.1,
            in_loop,
            vis,
            scope_data,
            result,
        );
    }
}

impl Assign {
    // Collect variables defined in an expression on the LHS of an assignment (or
    // for variable etc)
    fn collect_defines_lvalue<'a>(
        expr: &'a mut CstAssign,
        top_level_stmt_index: TopLevelStmtIndex,
        in_loop: InLoop,
        scope_data: &mut ModuleScopeData,
        frozen_heap: &FrozenHeap,
        result: &mut SmallMap<FrozenStringValue, BindingId>,
    ) {
        expr.node.visit_lvalue_mut(|x| {
            AssignIdent::collect_assign_ident(
                x,
                top_level_stmt_index,
                in_loop,
                Visibility::Public,
                scope_data,
                frozen_heap,
                result,
            )
        });
    }
}

/// Storage of objects referenced by AST.
#[derive(Default)]
pub(crate) struct ModuleScopeData<'f> {
    /// Bindings by id.
    bindings: Vec<Binding<'f>>,
    /// Scopes by id.
    scopes: Vec<ScopeNames<'f>>,
}

#[derive(Debug, Eq, PartialEq)]
pub(crate) enum AssignCount {
    /// Variable is assigned at most once during the execution of the scope.
    AtMostOnce,
    /// Variable may be assigned more than once during execution of the scope.
    Any,
}

/// Was a binding captured by nested def or lambda scopes?
#[derive(Debug, Copy, Clone, Dupe, Eq, PartialEq, VisitSpanMut)]
pub(crate) enum Captured {
    Yes,
    No,
}

#[derive(Debug)]
pub(crate) enum BindingSource {
    /// Variable is defined in the source of the module.
    Source(Span, TopLevelStmtIndex),
    /// Variable came from `Module`, not defined in the source file.
    FromModule,
}

/// Binding defines a place for a variable.
///
/// For example, in code `x = 1; x = 2`, there's one binding for name `x`.
///
/// In code `x = 1; def f(): x = 2`, there are two bindings for name `x`.
#[derive(Debug)]
pub(crate) struct Binding<'f> {
    pub(crate) name: FrozenStringValue,
    pub(crate) source: BindingSource,
    pub(crate) vis: Visibility,
    /// `slot` is `None` when it is not initialized yet.
    /// When analysis is completed, `slot` is always `Some`.
    pub(crate) slot: Option<Slot>,
    pub(crate) assign_count: AssignCount,
    // Whether a variable defined in a scope gets captured in nested def or lambda scope.
    // (Comprehension scopes do not count, because they are considered
    // local by the runtime and do not allocate a frame).
    pub(crate) captured: Captured,
    _marker: PhantomData<&'f ()>,
}

impl<'f> Binding<'f> {
    fn new(
        name: FrozenStringValue,
        source: BindingSource,
        vis: Visibility,
        assign_count: AssignCount,
    ) -> Binding<'f> {
        Binding {
            name,
            source,
            vis,
            slot: None,
            assign_count,
            captured: Captured::No,
            _marker: PhantomData,
        }
    }
}

/// If of a binding within current module.
#[derive(Copy, Clone, Dupe, Debug, Hash, PartialEq, Eq)]
pub(crate) struct BindingId(usize);

/// Id of a scope within current module.
#[derive(Copy, Clone, Dupe, Debug, Eq, PartialEq)]
pub(crate) struct ScopeId(usize);

impl ScopeId {
    pub(crate) fn module() -> ScopeId {
        ScopeId(0)
    }
}

/// Index of top-level statement.
#[derive(Copy, Clone, Dupe, Debug, PartialEq, Eq, Ord, PartialOrd)]
pub(crate) struct TopLevelStmtIndex(pub(crate) usize);

impl<'f> ModuleScopeData<'f> {
    pub(crate) fn new() -> ModuleScopeData<'f> {
        ModuleScopeData::default()
    }

    pub(crate) fn get_binding(&self, BindingId(id): BindingId) -> &Binding<'f> {
        &self.bindings[id]
    }

    fn mut_binding(&mut self, BindingId(id): BindingId) -> &mut Binding<'f> {
        &mut self.bindings[id]
    }

    fn new_binding(
        &mut self,
        name: FrozenStringValue,
        source: BindingSource,
        vis: Visibility,
        assigned_count: AssignCount,
    ) -> (BindingId, &mut Binding<'f>) {
        let binding_id = BindingId(self.bindings.len());
        self.bindings
            .push(Binding::new(name, source, vis, assigned_count));
        (binding_id, self.bindings.last_mut().unwrap())
    }

    pub(crate) fn get_scope(&self, ScopeId(id): ScopeId) -> &ScopeNames<'f> {
        &self.scopes[id]
    }

    pub(crate) fn mut_scope(&mut self, ScopeId(id): ScopeId) -> &mut ScopeNames<'f> {
        &mut self.scopes[id]
    }

    pub(crate) fn new_scope(&mut self) -> (ScopeId, &mut ScopeNames<'f>) {
        let scope_id = ScopeId(self.scopes.len());
        self.scopes.push(ScopeNames::default());
        (scope_id, self.scopes.last_mut().unwrap())
    }

    /// Get resolved slot for assigning identifier.
    pub(crate) fn get_assign_ident_slot(&self, ident: &CstAssignIdent) -> (Slot, Captured) {
        let binding_id = ident.1.expect("binding not assigned for ident");
        let binding = self.get_binding(binding_id);
        let slot = binding.slot;
        (
            slot.expect("binding slot is not initialized"),
            binding.captured,
        )
    }
}

#[derive(Debug, Clone, Dupe)]
pub(crate) enum ResolvedIdent {
    Slot(Slot, BindingId),
    Global(FrozenValue),
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! AST payload for type checking.

// We use CST as acronym for compiler-specific AST.

use std::collections::HashMap;

use dupe::OptionDupedExt;

use crate::codemap::CodeMap;
use crate::eval::compiler::scope::BindingId;
use crate::eval::compiler::scope::ModuleScopeData;
use crate::eval::compiler::scope::ResolvedIdent;
use crate::eval::compiler::scope::ScopeId;
use crate::syntax::ast::AstArgumentP;
use crate::syntax::ast::AstAssignIdentP;
use crate::syntax::ast::AstAssignP;
use crate::syntax::ast::AstExprP;
use crate::syntax::ast::AstIdentP;
use crate::syntax::ast::AstNoPayload;
use crate::syntax::ast::AstParameterP;
use crate::syntax::ast::AstPayload;
use crate::syntax::ast::AstStmt;
use crate::syntax::ast::AstStmtP;
use crate::syntax::ast::AstTypeExprP;
use crate::syntax::payload_map::AstPayloadFunction;
use crate::typing::error::InternalError;
use crate::typing::Interface;
use crate::values::typing::TypeCompiled;
use crate::values::FrozenValue;

/// Compiler-specific AST payload.
#[derive(Debug, Clone)]
pub(crate) struct CstPayload;
impl AstPayload for CstPayload {
    type LoadPayload = Interface;
    /// Information about how identifier binding is resolved.
    ///
    /// This is `None` when CST is created.
    /// All payload objects are filled with binding ids for all assign idents
    /// during analysis.
    ///
    /// When compilation starts, all payloads are `Some`.
    type IdentPayload = Option<ResolvedIdent>;
    /// Binding for an identifier in assignment position.
    ///
    /// This is `None` when CST is created.
    /// All payload objects are filled with binding ids for all assign idents
    /// during analysis.
    ///
    /// When compilation starts, all payloads are `Some`.
    type IdentAssignPayload = Option<BindingId>;
    type DefPayload = ScopeId;
    /// Populated before evaluation of top level statements.
    type TypeExprPayload = Option<TypeCompiled<FrozenValue>>;
}

struct CompilerAstMap<'a, 'f> {
    scope_data: &'a mut ModuleScopeData<'f>,
    loads: &'a HashMap<String, Interface>,
}

impl AstPayloadFunction<AstNoPayload, CstPayload> for CompilerAstMap<'_, '_> {
    fn map_load(&mut self, import_path: &str, (): ()) -> Interface {
        self.loads
            .get(import_path)
            .duped()
            .unwrap_or_else(Interface::empty)
    }

    fn map_ident(&mut self, (): ()) -> Option<ResolvedIdent> {
        None
    }

    fn map_ident_assign(&mut self, (): ()) -> Option<BindingId> {
        None
    }

    fn map_def(&mut self, (): ()) -> ScopeId {
        self.scope_data.new_scope().0
    }

    fn map_type_expr(&mut self, (): ()) -> Option<TypeCompiled<FrozenValue>> {
        None
    }
}

impl CstStmt {
    pub(crate) fn from_ast(
        stmt: AstStmt,
        scope_data: &mut ModuleScopeData,
        loads: &HashMap<String, Interface>,
    ) -> CstStmt {
        stmt.into_map_payload(&mut CompilerAstMap { scope_data, loads })
    }
}

impl CstAssignIdent {
    pub(crate) fn resolved_binding_id(
        &self,
        codemap: &CodeMap,
    ) -> Result<BindingId, InternalError> {
        match self.1 {
            Some(binding_id) => Ok(binding_id),
            None => Err(InternalError::msg(
                "Binding id is not filled",
                self.span,
                codemap,
            )),
        }
    }
}

pub(crate) type CstExpr = AstExprP<CstPayload>;
pub(crate) type CstTypeExpr = AstTypeExprP<CstPayload>;
pub(crate) type CstAssign = AstAssignP<CstPayload>;
pub(crate) type CstAssignIdent = AstAssignIdentP<CstPayload>;
pub(crate) type CstIdent = AstIdentP<CstPayload>;
pub(crate) type CstArgument = AstArgumentP<CstPayload>;
pub(crate) type CstParameter = AstParameterP<CstPayload>;
pub(crate) type CstStmt = AstStmtP<CstPayload>;


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#![cfg(test)]

use std::collections::HashMap;
use std::fmt::Write;

use dupe::Dupe;

use crate::environment::names::MutableNames;
use crate::environment::Globals;
use crate::eval::compiler::scope::payload::CstAssign;
use crate::eval::compiler::scope::payload::CstAssignIdent;
use crate::eval::compiler::scope::payload::CstExpr;
use crate::eval::compiler::scope::payload::CstStmt;
use crate::eval::compiler::scope::AssignCount;
use crate::eval::compiler::scope::Captured;
use crate::eval::compiler::scope::ModuleScopes;
use crate::eval::compiler::scope::ResolvedIdent;
use crate::eval::compiler::scope::Slot;
use crate::syntax::ast::DefP;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::StmtP;
use crate::syntax::uniplate::Visit;
use crate::syntax::AstModule;
use crate::syntax::Dialect;
use crate::values::FrozenHeap;
use crate::values::FrozenRef;

fn test_with_module(program: &str, expected: &str, module: &MutableNames) {
    let ast = AstModule::parse("t.star", program.to_owned(), &Dialect::Extended).unwrap();
    let frozen_heap = FrozenHeap::new();
    let codemap = frozen_heap.alloc_any_display_from_debug(ast.codemap.dupe());
    let ModuleScopes {
        cst, scope_data, ..
    } = ModuleScopes::check_module_err(
        module,
        &frozen_heap,
        &HashMap::new(),
        ast.statement,
        FrozenRef::new(Globals::empty()),
        codemap,
        &Dialect::Extended,
    )
    .unwrap();
    let mut r = String::new();
    for (i, binding) in scope_data.bindings.iter().enumerate() {
        if i != 0 {
            r.push(' ');
        }
        let slot = match binding.slot.unwrap() {
            Slot::Module(slot) => format!("m={}", slot.0),
            Slot::Local(slot) => format!("l={}", slot.0),
        };
        let assign_count = match binding.assign_count {
            AssignCount::AtMostOnce => "",
            AssignCount::Any => "+",
        };
        let captured = match binding.captured {
            Captured::Yes => "&",
            Captured::No => "",
        };
        write!(r, "{}:{}{}{}", i, slot, assign_count, captured).unwrap();
    }

    write!(r, " |").unwrap();

    struct Visitor<'a> {
        r: &'a mut String,
    }

    impl Visitor<'_> {
        fn visit_expr(&mut self, expr: &CstExpr) {
            if let ExprP::Identifier(ident) = &expr.node {
                let resolved = match ident.node.1.as_ref().unwrap() {
                    ResolvedIdent::Slot(_slot, binding_id) => binding_id.0.to_string(),
                    ResolvedIdent::Global(_) => "G".to_owned(),
                };
                write!(&mut self.r, " {}:{}", ident.node.0, resolved).unwrap();
            }

            expr.visit_expr(|expr| self.visit_expr(expr));
        }

        fn visit_exprs<'a>(&mut self, exprs: impl IntoIterator<Item = &'a CstExpr>) {
            for expr in exprs {
                self.visit_expr(expr);
            }
        }

        fn visit_lvalue(&mut self, ident: &CstAssignIdent) {
            write!(&mut self.r, " {}:{}", ident.0, ident.1.unwrap().0).unwrap();
        }

        fn visit_stmt_children(&mut self, stmt: &CstStmt) {
            stmt.visit_children(|visit| match visit {
                Visit::Stmt(stmt) => self.visit_stmt(stmt),
                Visit::Expr(expr) => self.visit_expr(expr),
            });
        }

        fn visit_assign(&mut self, assign: &CstAssign) {
            assign.visit_lvalue(|ident| self.visit_lvalue(ident));
        }

        fn visit_stmt(&mut self, stmt: &CstStmt) {
            match &stmt.node {
                StmtP::Assign(lhs, _rhs) => self.visit_assign(lhs),
                StmtP::Def(DefP { name, params, .. }) => {
                    self.visit_lvalue(name);
                    for param in params {
                        let (name, def, typ) = param.split();
                        if let Some(name) = name {
                            self.visit_lvalue(name);
                        }
                        self.visit_exprs(def.map(|def| &def.node.expr));
                        self.visit_exprs(typ);
                    }
                }
                StmtP::For(assign, ..) => self.visit_assign(assign),
                _ => {}
            }

            self.visit_stmt_children(stmt);
        }
    }

    Visitor { r: &mut r }.visit_stmt(&cst);

    assert_eq!(expected, &r);
}

fn t(program: &str, expected: &str) {
    let module = MutableNames::new();
    test_with_module(program, expected, &module);
}

// Expected test output (second parameter to `t` function) is:
// * list of bindings in format like `1:l=2` means binding id = 1, local slot 2
// * list of variables with references to binding ids

#[test]
fn basic() {
    t("x = 1; y = 2", "0:m=0 1:m=1 | x:0 y:1");
}

#[test]
fn module_reassignment() {
    t("x = 1; x = 2", "0:m=0+ | x:0 x:0");
}

#[test]
fn reassignment_in_loop() {
    t("for x in []: y = x", "0:m=0+ 1:m=1+ | x:0 y:1 x:0");
}

#[test]
fn def_capture() {
    t("x = 1\ndef f(): x", "0:m=0& 1:m=1 | x:0 f:1 x:0");
}

#[test]
fn def_shadow() {
    t("x = 1\ndef f(): x = 2", "0:m=0 1:m=1 2:l=0 | x:0 f:1 x:2");
}

#[test]
fn def_param_bindings() {
    t("def f(x): return x", "0:m=0 1:l=0 | f:0 x:1 x:1");
}

#[test]
fn nested_def_capture() {
    t(
        "\
def f():
    x = 1
    def g(): return x",
        "0:m=0 1:l=0& 2:l=1 | f:0 x:1 g:2 x:1",
    )
}

#[test]
fn existing_module_with_names() {
    let frozen_heap = FrozenHeap::new();
    let module = MutableNames::new();
    module.add_name(frozen_heap.alloc_str_intern("x"));
    module.add_name(frozen_heap.alloc_str_intern("y"));
    test_with_module("x = y", "0:m=0+ 1:m=1 | x:0 y:1", &module);
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::fmt::Debug;
use std::iter;
use std::marker::PhantomData;

use dupe::Clone_;
use dupe::Dupe_;
use either::Either;
use thiserror::Error;

use crate::cast::transmute;
use crate::coerce::coerce;
use crate::coerce::Coerce;
use crate::collections::symbol_map::Symbol;
use crate::collections::Hashed;
use crate::collections::SmallMap;
use crate::collections::StarlarkHashValue;
use crate::eval::runtime::params::ParametersSpec;
use crate::hint::unlikely;
use crate::values::dict::Dict;
use crate::values::dict::DictRef;
use crate::values::iter::StarlarkIterator;
use crate::values::Heap;
use crate::values::StringValue;
use crate::values::UnpackValue;
use crate::values::Value;
use crate::values::ValueError;
use crate::values::ValueLike;

#[derive(Debug, Clone, Error)]
pub(crate) enum FunctionError {
    #[error("Missing parameter `{name}` for call to {function}")]
    MissingParameter { name: String, function: String },
    #[error("Found {count} extra positional argument(s) for call to {function}")]
    ExtraPositionalArg { count: usize, function: String },
    #[error("Found `{}` extra named parameter(s) for call to {function}", .names.join("` `"))]
    ExtraNamedArg {
        names: Vec<String>,
        function: String,
    },
    #[error("Argument `{name}` occurs more than once")]
    RepeatedArg { name: String },
    #[error("The argument provided for *args is not an identifier")]
    ArgsValueIsNotString,
    #[error("The argument provided for *args is not iterable")]
    ArgsArrayIsNotIterable,
    #[error("The argument provided for **kwargs is not a dictionary")]
    KwArgsIsNotDict,
    #[error("Wrong number of positional arguments, expected {}, got {got}",
        if min == max {min.to_string()} else {format!("between {} and {}", min, max)})]
    WrongNumberOfArgs { min: usize, max: usize, got: usize },
}

/// An object accompanying argument name for faster argument resolution.
pub(crate) trait ArgSymbol: Debug + Coerce<Self> + 'static {
    fn get_index_from_param_spec<'v, V: ValueLike<'v>>(
        &self,
        ps: &ParametersSpec<V>,
    ) -> Option<usize>;

    fn small_hash(&self) -> StarlarkHashValue;
}

impl ArgSymbol for Symbol {
    fn get_index_from_param_spec<'v, V: ValueLike<'v>>(
        &self,
        ps: &ParametersSpec<V>,
    ) -> Option<usize> {
        ps.names.get(self).map(|i| *i as usize)
    }

    fn small_hash(&self) -> StarlarkHashValue {
        self.small_hash()
    }
}

/// `Symbol` resolved to function parameter index.
#[derive(Debug)]
pub(crate) struct ResolvedArgName {
    /// Hash of the argument name.
    pub(crate) hash: StarlarkHashValue,
    /// Parameter index or `None` if the argument should go to kwargs.
    pub(crate) param_index: Option<u32>,
}

impl ArgSymbol for ResolvedArgName {
    fn get_index_from_param_spec<'v, V: ValueLike<'v>>(
        &self,
        _ps: &ParametersSpec<V>,
    ) -> Option<usize> {
        self.param_index.map(|i| i as usize)
    }

    fn small_hash(&self) -> StarlarkHashValue {
        self.hash
    }
}

unsafe impl Coerce<ResolvedArgName> for ResolvedArgName {}

#[derive(Debug, Clone_, Dupe_)]
pub(crate) struct ArgNames<'a, 'v, S: ArgSymbol> {
    /// Names are not guaranteed to be unique here.
    names: &'a [(S, StringValue<'v>)],
}

impl<'a, 'v, S: ArgSymbol> Default for ArgNames<'a, 'v, S> {
    fn default() -> Self {
        ArgNames { names: &[] }
    }
}

impl<'a, 'v, S: ArgSymbol> Copy for ArgNames<'a, 'v, S> {}

impl<'a, 'v, S: ArgSymbol> ArgNames<'a, 'v, S> {
    /// Names are allowed to be not-unique.
    /// String in `Symbol` must be equal to the `StringValue`,
    /// it is caller responsibility to ensure that.
    pub(crate) fn new(names: &'a [(S, StringValue<'v>)]) -> ArgNames<'a, 'v, S> {
        ArgNames { names }
    }

    pub(crate) fn names(&self) -> &'a [(S, StringValue<'v>)] {
        self.names
    }

    pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = &'a (S, StringValue<'v>)> {
        self.names.iter()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.names.is_empty()
    }

    pub(crate) fn len(&self) -> usize {
        self.names.len()
    }
}

/// Either full arguments, or short arguments for positional-only calls.
pub(crate) trait ArgumentsImpl<'v, 'a>: Debug {
    type ArgSymbol: ArgSymbol;
    fn pos(&self) -> &[Value<'v>];
    fn named(&self) -> &[Value<'v>];
    fn names(&self) -> ArgNames<'a, 'v, Self::ArgSymbol>;
    fn args(&self) -> Option<Value<'v>>;
    fn kwargs(&self) -> Option<Value<'v>>;
}

/// Arguments object is passed from the starlark interpreter to function implementation
/// when evaluation function or method calls.
#[derive(Clone_, Dupe_, Debug)]
pub(crate) struct ArgumentsFull<'v, 'a, S: ArgSymbol> {
    /// Positional arguments.
    pub(crate) pos: &'a [Value<'v>],
    /// Named arguments.
    pub(crate) named: &'a [Value<'v>],
    /// Names of named arguments.
    ///
    /// `named` length must be equal to `names` length.
    pub(crate) names: ArgNames<'a, 'v, S>,
    /// `*args` argument.
    pub(crate) args: Option<Value<'v>>,
    /// `**kwargs` argument.
    pub(crate) kwargs: Option<Value<'v>>,
}

impl<'v, 'a, S: ArgSymbol> Default for ArgumentsFull<'v, 'a, S> {
    fn default() -> Self {
        ArgumentsFull {
            pos: &[],
            named: &[],
            names: ArgNames::default(),
            args: None,
            kwargs: None,
        }
    }
}

impl<'v, 'a, S: ArgSymbol> ArgumentsImpl<'v, 'a> for ArgumentsFull<'v, 'a, S> {
    type ArgSymbol = S;

    #[inline]
    fn pos(&self) -> &[Value<'v>] {
        self.pos
    }

    #[inline]
    fn named(&self) -> &[Value<'v>] {
        self.named
    }

    #[inline]
    fn names(&self) -> ArgNames<'a, 'v, S> {
        self.names
    }

    #[inline]
    fn args(&self) -> Option<Value<'v>> {
        self.args
    }

    #[inline]
    fn kwargs(&self) -> Option<Value<'v>> {
        self.kwargs
    }
}

/// Positional-only arguments, smaller and faster than `ArgumentsFull`.
#[derive(Debug)]
pub(crate) struct ArgumentsPos<'v, 'a, S: ArgSymbol> {
    pub(crate) pos: &'a [Value<'v>],
    pub(crate) names: PhantomData<&'static S>,
}

impl<'a, 'v, S: ArgSymbol> ArgumentsImpl<'v, 'a> for ArgumentsPos<'v, 'a, S> {
    type ArgSymbol = S;

    #[inline]
    fn pos(&self) -> &[Value<'v>] {
        self.pos
    }

    #[inline]
    fn named(&self) -> &[Value<'v>] {
        &[]
    }

    #[inline]
    fn names(&self) -> ArgNames<'a, 'v, S> {
        ArgNames::default()
    }

    #[inline]
    fn args(&self) -> Option<Value<'v>> {
        None
    }

    #[inline]
    fn kwargs(&self) -> Option<Value<'v>> {
        None
    }
}

/// Arguments object is passed from the starlark interpreter to function implementation
/// when evaluation function or method calls.
#[derive(Default, Clone, Dupe_)]
pub struct Arguments<'v, 'a>(pub(crate) ArgumentsFull<'v, 'a, Symbol>);

impl<'v, 'a> Arguments<'v, 'a> {
    /// Unwrap all named arguments (both explicit and in `**kwargs`) into a map.
    ///
    /// This operation fails if named argument names are not unique.
    pub fn names_map(&self) -> anyhow::Result<SmallMap<StringValue<'v>, Value<'v>>> {
        match self.unpack_kwargs()? {
            None => {
                let mut result = SmallMap::with_capacity(self.0.names.len());
                for (k, v) in self.0.names.iter().zip(self.0.named) {
                    let old =
                        result.insert_hashed(Hashed::new_unchecked(k.0.small_hash(), k.1), *v);
                    if unlikely(old.is_some()) {
                        return Err(FunctionError::RepeatedArg {
                            name: k.1.as_str().to_owned(),
                        }
                        .into());
                    }
                }
                Ok(result)
            }
            Some(kwargs) => {
                if self.0.names.is_empty() {
                    match kwargs.downcast_ref_key_string() {
                        Some(kwargs) => Ok(kwargs.clone()),
                        None => Err(FunctionError::ArgsValueIsNotString.into()),
                    }
                } else {
                    // We have to insert the names before the kwargs since the iteration order is observable
                    let mut result = SmallMap::with_capacity(self.0.names.len() + kwargs.len());
                    for (k, v) in self.0.names.iter().zip(self.0.named) {
                        let old =
                            result.insert_hashed(Hashed::new_unchecked(k.0.small_hash(), k.1), *v);
                        if unlikely(old.is_some()) {
                            return Err(FunctionError::RepeatedArg {
                                name: k.1.as_str().to_owned(),
                            }
                            .into());
                        }
                    }
                    for (k, v) in kwargs.iter_hashed() {
                        let s = Arguments::unpack_kwargs_key_as_value(*k.key())?;
                        let k = Hashed::new_unchecked(k.hash(), s);
                        let old = result.insert_hashed(k, v);
                        if unlikely(old.is_some()) {
                            return Err(FunctionError::RepeatedArg {
                                name: s.as_str().to_owned(),
                            }
                            .into());
                        }
                    }
                    Ok(result)
                }
            }
        }
    }

    /// The number of arguments, where those inside a args/kwargs are counted as multiple arguments.
    ///
    /// This operation fails if the `kwargs` is not a dictionary, or `args` does not support `len`.
    pub fn len(&self) -> anyhow::Result<usize> {
        let args = match self.0.args {
            None => 0,
            Some(v) => v.length()? as usize,
        };
        let kwargs = match self.unpack_kwargs()? {
            None => 0,
            Some(v) => v.len(),
        };
        Ok(self.0.pos.len() + self.0.named.len() + args + kwargs)
    }

    /// Unwrap all named arguments (both explicit and in `**kwargs`) into a dictionary.
    ///
    /// This operation fails if named argument names are not unique.
    pub fn names(&self) -> anyhow::Result<Dict<'v>> {
        Ok(Dict::new(coerce(self.names_map()?)))
    }

    /// Unpack all positional parameters into an iterator.
    pub fn positions<'b>(
        &'b self,
        heap: &'v Heap,
    ) -> anyhow::Result<impl Iterator<Item = Value<'v>> + 'b> {
        let tail = match self.0.args {
            None => Either::Left(iter::empty()),
            Some(args) => Either::Right(args.iterate(heap)?),
        };
        Ok(self.0.pos.iter().copied().chain(tail))
    }

    /// Examine the `kwargs` field, converting it to a [`Dict`] or failing.
    /// Note that even if this operation succeeds, the keys in the kwargs
    /// will _not_ have been validated to be strings (as they must be).
    /// The arguments may also overlap with named, which would be an error.
    #[inline(always)]
    pub fn unpack_kwargs(&self) -> anyhow::Result<Option<DictRef<'v>>> {
        match self.0.kwargs {
            None => Ok(None),
            Some(kwargs) => match DictRef::from_value(kwargs) {
                None => Err(FunctionError::KwArgsIsNotDict.into()),
                Some(x) => Ok(Some(x)),
            },
        }
    }

    /// Confirm that a key in the `kwargs` field is indeed a string, or [`Err`].
    #[inline(always)]
    pub(crate) fn unpack_kwargs_key_as_value(k: Value<'v>) -> anyhow::Result<StringValue<'v>> {
        match StringValue::new(k) {
            None => Err(FunctionError::ArgsValueIsNotString.into()),
            Some(k) => Ok(k),
        }
    }

    /// Confirm that a key in the `kwargs` field is indeed a string, or [`Err`].
    #[inline(always)]
    pub fn unpack_kwargs_key(k: Value<'v>) -> anyhow::Result<&'v str> {
        Arguments::unpack_kwargs_key_as_value(k).map(|k| k.as_str())
    }

    /// Produce [`Err`] if there are any positional arguments.
    #[inline(always)]
    pub fn no_positional_args(&self, heap: &'v Heap) -> anyhow::Result<()> {
        let [] = self.positional(heap)?;
        Ok(())
    }

    /// Produce [`Err`] if there are any named (i.e. non-positional) arguments.
    #[inline(always)]
    pub fn no_named_args(&self) -> anyhow::Result<()> {
        #[cold]
        #[inline(never)]
        fn bad(x: &Arguments) -> anyhow::Result<()> {
            // We might have a empty kwargs dictionary, but probably have an error
            let mut extra = Vec::new();
            extra.extend(x.0.names.iter().map(|x| x.0.as_str().to_owned()));
            if let Some(kwargs) = x.unpack_kwargs()? {
                for k in kwargs.keys() {
                    extra.push(Arguments::unpack_kwargs_key(k)?.to_owned());
                }
            }
            if extra.is_empty() {
                Ok(())
            } else {
                // Would be nice to give a better name here, but it's in the call stack, so no big deal
                Err(FunctionError::ExtraNamedArg {
                    names: extra,
                    function: "function".to_owned(),
                }
                .into())
            }
        }

        if self.0.named.is_empty() && self.0.kwargs.is_none() {
            Ok(())
        } else {
            bad(self)
        }
    }

    /// Collect exactly `N` positional arguments from the [`Arguments`], failing if there are too many/few
    /// arguments. Ignores named arguments.
    #[inline(always)]
    pub fn positional<const N: usize>(&self, heap: &'v Heap) -> anyhow::Result<[Value<'v>; N]> {
        #[cold]
        #[inline(never)]
        fn rare<'v, const N: usize>(
            x: &Arguments<'v, '_>,
            heap: &'v Heap,
        ) -> anyhow::Result<[Value<'v>; N]> {
            // Very sad that we allocate into a vector, but I expect calling into a small positional argument
            // with a *args is very rare.
            let xs =
                x.0.pos
                    .iter()
                    .copied()
                    .chain(x.0.args.unwrap().iterate(heap)?)
                    .collect::<Vec<_>>();
            xs.as_slice().try_into().map_err(|_| {
                FunctionError::WrongNumberOfArgs {
                    min: N,
                    max: N,
                    got: x.0.pos.len(),
                }
                .into()
            })
        }

        if self.0.args.is_none() {
            self.0.pos.try_into().map_err(|_| {
                FunctionError::WrongNumberOfArgs {
                    min: N,
                    max: N,
                    got: self.0.pos.len(),
                }
                .into()
            })
        } else {
            rare(self, heap)
        }
    }

    /// Collect exactly `REQUIRED` positional arguments, plus at most `OPTIONAL` positional arguments
    /// from the [`Arguments`], failing if there are too many/few arguments. Ignores named arguments.
    /// The `OPTIONAL` array will never have a [`Some`] after a [`None`].
    #[inline(always)]
    pub fn optional<const REQUIRED: usize, const OPTIONAL: usize>(
        &self,
        heap: &'v Heap,
    ) -> anyhow::Result<([Value<'v>; REQUIRED], [Option<Value<'v>>; OPTIONAL])> {
        #[cold]
        #[inline(never)]
        fn rare<'v, const REQUIRED: usize, const OPTIONAL: usize>(
            x: &Arguments<'v, '_>,
            heap: &'v Heap,
        ) -> anyhow::Result<([Value<'v>; REQUIRED], [Option<Value<'v>>; OPTIONAL])> {
            // Very sad that we allocate into a vector, but I expect calling into a small positional argument
            // with a *args is very rare.
            let args = match x.0.args {
                None => StarlarkIterator::empty(heap),
                Some(args) => args.iterate(heap)?,
            };
            let xs = x.0.pos.iter().copied().chain(args).collect::<Vec<_>>();
            if xs.len() >= REQUIRED && xs.len() <= REQUIRED + OPTIONAL {
                let required = xs[0..REQUIRED].try_into().unwrap();
                let mut optional = [None; OPTIONAL];
                for (a, b) in optional.iter_mut().zip(&xs[REQUIRED..]) {
                    *a = Some(*b);
                }
                Ok((required, optional))
            } else {
                Err(FunctionError::WrongNumberOfArgs {
                    min: REQUIRED,
                    max: REQUIRED + OPTIONAL,
                    got: xs.len(),
                }
                .into())
            }
        }

        if self.0.args.is_none()
            && self.0.pos.len() >= REQUIRED
            && self.0.pos.len() <= REQUIRED + OPTIONAL
        {
            let required = self.0.pos[0..REQUIRED].try_into().unwrap();
            let mut optional = [None; OPTIONAL];
            for (a, b) in optional.iter_mut().zip(&self.0.pos[REQUIRED..]) {
                *a = Some(*b);
            }
            Ok((required, optional))
        } else {
            rare(self, heap)
        }
    }

    /// Collect 1 positional arguments from the [`Arguments`], failing if there are too many/few
    /// arguments. Ignores named arguments.
    #[inline(always)]
    pub fn positional1(&self, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        // Could be implemented more directly, let's see if profiling shows it up
        let [x] = self.positional(heap)?;
        Ok(x)
    }

    /// Collect up to 1 optional arguments from the [`Arguments`], failing if there are too many
    /// arguments. Ignores named arguments.
    #[inline(always)]
    pub fn optional1(&self, heap: &'v Heap) -> anyhow::Result<Option<Value<'v>>> {
        // Could be implemented more directly, let's see if profiling shows it up
        let ([], [x]) = self.optional(heap)?;
        Ok(x)
    }
}

impl Arguments<'_, '_> {
    /// Utility for checking a `this` parameter matches what you expect.
    pub fn check_this<'v, T: UnpackValue<'v>>(this: Value<'v>) -> anyhow::Result<T> {
        T::unpack_named_param(this, "this")
    }

    /// Utility for checking a required parameter matches what you expect.
    pub fn check_required<'v, T: UnpackValue<'v>>(
        name: &str,
        x: Option<Value<'v>>,
    ) -> anyhow::Result<T> {
        let x = x.ok_or_else(|| ValueError::MissingRequired(name.to_owned()))?;
        T::unpack_named_param(x, name)
    }

    /// Utility for checking an optional parameter matches what you expect.
    pub fn check_optional<'v, T: UnpackValue<'v>>(
        name: &str,
        x: Option<Value<'v>>,
    ) -> anyhow::Result<Option<T>> {
        match x {
            None => Ok(None),
            Some(x) => Ok(Some(T::unpack_value(x).ok_or_else::<anyhow::Error, _>(
                || {
                    ValueError::IncorrectParameterTypeNamedWithExpected(
                        name.to_owned(),
                        T::expected(),
                        x.get_type().to_owned(),
                    )
                    .into()
                },
            )?)),
        }
    }
}

impl<'a> Arguments<'static, 'a> {
    /// Convert `Arguments` with `FrozenValue` (because no other values can have `'v` lifetime)
    /// to arbitrary `'v` lifetime.
    pub(crate) fn frozen_to_v<'v>(&self) -> &Arguments<'v, 'a> {
        unsafe { transmute!(&Arguments, &Arguments, self) }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::const_frozen_string;
    use crate::values::StringValueLike;

    #[test]
    fn test_parameter_unpack() {
        let heap = Heap::new();
        fn f<'v, F: Fn(&Arguments<'v, '_>), const N: usize>(heap: &'v Heap, op: F) {
            for i in 0..=N {
                let mut p = Arguments::default();
                let pos = (0..i)
                    .map(|x| Value::testing_new_int(x as i32))
                    .collect::<Vec<_>>();
                let args = (i..N)
                    .map(|x| Value::testing_new_int(x as i32))
                    .collect::<Vec<_>>();
                let empty_args = args.is_empty();
                p.0.pos = &pos;
                p.0.args = Some(heap.alloc(args));
                op(&p);
                if empty_args {
                    p.0.args = None;
                    op(&p);
                }
                assert_eq!(p.len().unwrap(), N);
            }
        }

        f::<_, 0>(&heap, |p| {
            assert_eq!(&p.positional::<0>(&heap).unwrap(), &[]);
            assert!(&p.positional::<1>(&heap).is_err());
            assert!(&p.positional::<2>(&heap).is_err());
            assert_eq!(&p.optional::<0, 1>(&heap).unwrap(), &([], [None]));
            assert!(&p.optional::<1, 1>(&heap).is_err());
            assert_eq!(&p.optional::<0, 2>(&heap).unwrap(), &([], [None, None]));
        });
        f::<_, 1>(&heap, |p| {
            assert!(&p.positional::<0>(&heap).is_err());
            assert_eq!(
                &p.positional::<1>(&heap).unwrap(),
                &[Value::testing_new_int(0)]
            );
            assert!(&p.positional::<2>(&heap).is_err());
            assert_eq!(
                &p.optional::<0, 1>(&heap).unwrap(),
                &([], [Some(Value::testing_new_int(0))])
            );
            assert_eq!(
                &p.optional::<1, 1>(&heap).unwrap(),
                &([Value::testing_new_int(0)], [None])
            );
            assert_eq!(
                &p.optional::<0, 2>(&heap).unwrap(),
                &([], [Some(Value::testing_new_int(0)), None])
            );
        });
        f::<_, 2>(&heap, |p| {
            assert!(&p.positional::<0>(&heap).is_err());
            assert!(&p.positional::<1>(&heap).is_err());
            assert_eq!(
                &p.positional::<2>(&heap).unwrap(),
                &[Value::testing_new_int(0), Value::testing_new_int(1)]
            );
            assert!(p.optional::<0, 1>(&heap).is_err());
            assert_eq!(
                &p.optional::<1, 1>(&heap).unwrap(),
                &(
                    [Value::testing_new_int(0)],
                    [Some(Value::testing_new_int(1))]
                )
            );
            assert_eq!(
                &p.optional::<0, 2>(&heap).unwrap(),
                &(
                    [],
                    [
                        Some(Value::testing_new_int(0)),
                        Some(Value::testing_new_int(1))
                    ]
                )
            );
        });
        f::<_, 3>(&heap, |p| {
            assert!(&p.positional::<0>(&heap).is_err());
            assert!(&p.positional::<1>(&heap).is_err());
            assert!(&p.positional::<2>(&heap).is_err());
            assert!(p.optional::<0, 1>(&heap).is_err());
            assert!(p.optional::<1, 1>(&heap).is_err());
            assert!(p.optional::<0, 2>(&heap).is_err());
        });
    }

    #[test]
    fn test_parameter_no_named() {
        let heap = Heap::new();
        let mut p = Arguments::default();
        assert!(p.no_named_args().is_ok());
        assert_eq!(p.len().unwrap(), 0);

        // Test lots of forms of kwargs work properly
        p.0.kwargs = Some(Value::new_none());
        assert!(p.no_named_args().is_err());
        p.0.kwargs = Some(heap.alloc(Dict::default()));
        assert!(p.no_named_args().is_ok());
        assert_eq!(p.len().unwrap(), 0);
        let mut sm = SmallMap::new();
        sm.insert_hashed(heap.alloc_str("test").get_hashed(), Value::new_none());
        p.0.kwargs = Some(heap.alloc(Dict::new(coerce(sm))));
        assert!(p.no_named_args().is_err());
        assert_eq!(p.len().unwrap(), 1);

        // Test named arguments work properly
        p.0.kwargs = None;
        let named = [Value::new_none()];
        p.0.named = &named;
        let names = [(Symbol::new("test"), heap.alloc_str("test"))];
        p.0.names = ArgNames::new(&names);
        assert!(p.no_named_args().is_err());
        assert_eq!(p.len().unwrap(), 1);
    }

    #[test]
    fn test_names_map_repeated_name_in_arg_names() {
        let named = vec![Value::testing_new_int(10), Value::new_bool(true)];
        let names = vec![
            (
                Symbol::new("a"),
                const_frozen_string!("a").to_string_value(),
            ),
            (
                Symbol::new("a"),
                const_frozen_string!("a").to_string_value(),
            ),
        ];
        let error = Arguments(ArgumentsFull {
            pos: &[],
            named: &named,
            names: ArgNames::new(&names),
            args: None,
            kwargs: None,
        })
        .names_map()
        .unwrap_err();
        assert!(error.to_string().contains("occurs more than once"));
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Configuration of `BeforeStmt` instrumentation of bytecode.

use crate::codemap::FileSpanRef;
use crate::eval::Evaluator;

/// Configuration of `BeforeStmt` instrumentation of bytecode.
#[derive(Default)]
pub(crate) struct BeforeStmt<'a> {
    /// Functions to run before each statement.
    pub(crate) before_stmt: Vec<BeforeStmtFunc<'a>>,
    /// Explicitly request generation of `BeforeStmt` instructions
    /// even if no `before_stmt` functions are registered.
    /// This is needed when compiling dependencies of a file to be profiled.
    pub(crate) instrument: bool,
}

/// This is used by DAP, and it is not public API.
// TODO(cjhopman): pull DAP into the crate, and hide this function.
#[doc(hidden)]
pub enum BeforeStmtFunc<'a> {
    Fn(&'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a>)),
    Dyn(Box<dyn BeforeStmtFuncDyn<'a>>),
}

impl<'a> BeforeStmtFunc<'a> {
    pub(crate) fn call<'v>(&mut self, span: FileSpanRef, eval: &mut Evaluator<'v, 'a>) {
        match self {
            BeforeStmtFunc::Fn(f) => f(span, eval),
            BeforeStmtFunc::Dyn(d) => d.call(span, eval),
        }
    }
}

/// This is used by DAP, and it is not public API.
// TODO(cjhopman): pull DAP into the crate, and hide this function.
#[doc(hidden)]
pub trait BeforeStmtFuncDyn<'a> {
    /// This is used by DAP, and it is not public API.
    // TODO(cjhopman): pull DAP into the crate, and hide this function.
    #[doc(hidden)]
    fn call<'v>(&mut self, span: FileSpanRef, eval: &mut Evaluator<'v, 'a>);
}

impl<'a> BeforeStmt<'a> {
    pub(crate) fn enabled(&self) -> bool {
        self.instrument || !self.before_stmt.is_empty()
    }
}

impl<'a> From<&'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a>)> for BeforeStmtFunc<'a> {
    fn from(value: &'a dyn for<'v1> Fn(FileSpanRef, &mut Evaluator<'v1, 'a>)) -> Self {
        Self::Fn(value)
    }
}

impl<'a> From<Box<dyn BeforeStmtFuncDyn<'a>>> for BeforeStmtFunc<'a> {
    fn from(value: Box<dyn BeforeStmtFuncDyn<'a>>) -> Self {
        Self::Dyn(value)
    }
}


/*
 * Copyright 2019 The Starlark in Rust Authors.
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Starlark call stack.

// FIXME: I think we should rewrite the CallStack stuff entirely:
// * Don't keep a call stack, just a call stack depth.
// * When people did StackGuard.inc just do CallStack.inc, we need less info
// once it's an int so can reuse.
// * When an exception happens, decorate it with the call stack on the way back
//   up, in eval_call.

use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;

use dupe::Dupe;

use crate::codemap::FileSpan;
use crate::errors::Frame;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::eval::runtime::inlined_frame::InlinedFrames;
use crate::hint::unlikely;
use crate::slice_vec_ext::SliceExt;
use crate::values::FrozenRef;
use crate::values::Trace;
use crate::values::Tracer;
use crate::values::Value;

// A value akin to Frame, but can be created cheaply, since it doesn't resolve
// anything in advance.
// The downside is it has a lifetime on 'v and keeps alive the whole CodeMap.
#[derive(Clone, Copy, Dupe)]
struct CheapFrame<'v> {
    function: Value<'v>,
    span: Option<FrozenRef<'static, FrameSpan>>,
}

impl CheapFrame<'_> {
    fn location(&self) -> Option<FileSpan> {
        self.span.map(|span| span.span.to_file_span())
    }

    fn extend_frames(&self, frames: &mut Vec<Frame>) {
        if let Some(span) = self.span {
            span.inlined_frames.extend_frames(frames);
        }
        frames.push(self.to_frame());
    }

    fn to_frame(&self) -> Frame {
        Frame {
            name: self.function.name_for_call_stack(),
            location: self.location(),
        }
    }
}

impl Debug for CheapFrame<'_> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut x = f.debug_struct("Frame");
        x.field("function", &self.function);
        x.field("span", &self.span);
        x.finish()
    }
}

#[derive(Debug, thiserror::Error)]
enum CallStackError {
    #[error("Requested {0}-th top frame, but stack size is {1} (internal error)")]
    StackIsTooShallowForNthTopFrame(usize, usize),
    #[error("Starlark call stack overflow")]
    Overflow,
}

/// Starlark call stack.
#[derive(Debug)]
pub(crate) struct CheapCallStack<'v> {
    count: usize,
    stack: [CheapFrame<'v>; MAX_CALLSTACK_RECURSION],
}

impl<'v> Default for CheapCallStack<'v> {
    fn default() -> Self {
        Self {
            count: 0,
            stack: [CheapFrame {
                function: Value::new_none(),
                span: None,
            }; MAX_CALLSTACK_RECURSION],
        }
    }
}

// Currently, each frame typically allocates about 1K of native stack size (see `test_frame_size`),
// but it is a bit more complicated:
// * each for loop in a frame allocates more native stack
// * inlined functions do not allocate native stack
// Practically max call stack depends on native stack size,
// and depending on environment, it may be configured differently, for example:
// * macOS default stack size is 512KB
// * Linux default stack size is 8MB
// * [tokio default stack size is 2MB][1]
// [1] https://docs.rs/tokio/0.2.1/tokio/runtime/struct.Builder.html#method.thread_stack_size
// TODO(nga): count loops in call stack size.
// TODO(nga): make it configurable.
const MAX_CALLSTACK_RECURSION: usize = 50;

unsafe impl<'v> Trace<'v> for CheapCallStack<'v> {
    fn trace(&mut self, tracer: &Tracer<'v>) {
        let (used, unused) = self.stack.split_at_mut(self.count);
        for x in used {
            x.function.trace(tracer);
        }
        // Not required, but since we are choosing not to walk those above
        // the current stack depth, it's good practice to blank those values out
        for x in unused {
            x.function = Value::new_none();
            x.span = None;
        }
    }
}

impl<'v> CheapCallStack<'v> {
    /// Push an element to the stack. It is important the each `push` is paired
    /// with a `pop`.
    pub(crate) fn push(
        &mut self,
        function: Value<'v>,
        span: Option<FrozenRef<'static, FrameSpan>>,
    ) -> anyhow::Result<()> {
        if unlikely(self.count >= MAX_CALLSTACK_RECURSION) {
            return Err(CallStackError::Overflow.into());
        }
        self.stack[self.count] = CheapFrame { function, span };
        self.count += 1;
        Ok(())
    }

    /// Remove the top element from the stack. Called after `push`.
    pub(crate) fn pop(&mut self) {
        debug_assert!(self.count >= 1);
        // We could clear the elements, but don't need to bother
        self.count -= 1;
    }

    /// Current size (in frames) of the stack.
    pub(crate) fn count(&self) -> usize {
        self.count
    }

    /// The frame at the top of the stack. May be `None` if
    /// either there the stack is empty, or the top of the stack lacks location
    /// information (e.g. called from Rust).
    pub(crate) fn top_frame(&self) -> Option<Frame> {
        Some(self.stack.last().as_ref()?.to_frame())
    }

    /// The location at the top of the stack. May be `None` if
    /// either there the stack is empty, or the top of the stack lacks location
    /// information (e.g. called from Rust).
    pub(crate) fn top_location(&self) -> Option<FileSpan> {
        if self.count == 0 {
            None
        } else {
            self.stack[self.count - 1].location()
        }
    }

    /// `n`-th element from the top of the stack.
    pub(crate) fn top_nth_function(&self, n: usize) -> anyhow::Result<Value<'v>> {
        let index = self
            .count
            .checked_sub(1)
            .and_then(|x| x.checked_sub(n))
            .ok_or(CallStackError::StackIsTooShallowForNthTopFrame(
                n, self.count,
            ))?;
        Ok(self.stack[index].function)
    }

    pub(crate) fn to_diagnostic_frames(&self, inlined_frames: InlinedFrames) -> CallStack {
        // The first entry is just the entire module, so skip it
        let mut frames = Vec::new();
        for frame in &self.stack[1..self.count] {
            frame.extend_frames(&mut frames);
        }
        inlined_frames.extend_frames(&mut frames);
        CallStack { frames }
    }

    /// List the entries on the stack as values
    pub(crate) fn to_function_values(&self) -> Vec<Value<'v>> {
        self.stack[1..self.count].map(|x| x.function)
    }
}

/// Owned call stack.
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
pub struct CallStack {
    frames: Vec<Frame>,
}

impl CallStack {
    /// Is the call stack empty?
    pub fn is_empty(&self) -> bool {
        self.frames.is_empty()
    }

    /// Take the contained frames.
    pub fn into_frames(self) -> Vec<Frame> {
        self.frames
    }
}

impl Display for CallStack {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if !self.frames.is_empty() {
            // Match Python output.
            writeln!(f, "Traceback (most recent call last):")?;
            // TODO(nga): use real module name.
            let mut prev = "<module>";
            for x in &self.frames {
                x.write_two_lines("  ", prev, f)?;
                prev = &x.name;
            }
        }
        Ok(())
    }
}