/*
 * 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 crate::environment::FrozenModuleData;
use crate::eval::compiler::stmt::OptimizeOnFreezeContext;
use crate::eval::Evaluator;
use crate::values::FrozenHeap;
use crate::values::Heap;

pub(crate) trait OptCtxEval<'v, 'a> {
    fn heap(&self) -> &'v Heap;
    fn frozen_heap(&self) -> &FrozenHeap;
    fn eval(&mut self) -> Option<&mut Evaluator<'v, 'a>>;
    fn frozen_module(&self) -> Option<&FrozenModuleData>;
}

impl<'v, 'a> OptCtxEval<'v, 'a> for OptimizeOnFreezeContext<'v, 'a> {
    fn heap(&self) -> &'v Heap {
        self.heap
    }

    fn frozen_heap(&self) -> &FrozenHeap {
        self.frozen_heap
    }

    fn eval(&mut self) -> Option<&mut Evaluator<'v, 'a>> {
        None
    }

    fn frozen_module(&self) -> Option<&FrozenModuleData> {
        Some(self.module)
    }
}

impl<'v, 'a> OptCtxEval<'v, 'a> for Evaluator<'v, 'a> {
    fn heap(&self) -> &'v Heap {
        self.heap()
    }

    fn frozen_heap(&self) -> &FrozenHeap {
        self.frozen_heap()
    }

    fn eval(&mut self) -> Option<&mut Evaluator<'v, 'a>> {
        Some(self)
    }

    fn frozen_module(&self) -> Option<&FrozenModuleData> {
        None
    }
}

/// Optimization context.
///
/// We perform optimization
/// * during compilation of AST to IR, and
/// * when freezing the heap.
pub(crate) struct OptCtx<'v, 'a, 'e> {
    pub(crate) eval: &'e mut dyn OptCtxEval<'v, 'a>,
    /// Current function parameter slot count. Zero when compiling module.
    pub(crate) param_count: u32,
}

impl<'v, 'a, 'e> OptCtx<'v, 'a, 'e> {
    pub(crate) fn new(
        eval: &'e mut dyn OptCtxEval<'v, 'a>,
        param_count: u32,
    ) -> OptCtx<'v, 'a, 'e> {
        OptCtx { eval, param_count }
    }

    pub(crate) fn heap(&self) -> &'v Heap {
        self.eval.heap()
    }

    pub(crate) fn frozen_heap(&self) -> &FrozenHeap {
        self.eval.frozen_heap()
    }

    pub(crate) fn eval(&mut self) -> Option<&mut Evaluator<'v, 'a>> {
        self.eval.eval()
    }

    pub(crate) fn frozen_module(&self) -> Option<&FrozenModuleData> {
        self.eval.frozen_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.
 */

//! Small `Vec`.

use std::fmt::Debug;
use std::mem;
use std::slice;

/// A small vector.
#[derive(Clone)]
pub(crate) enum SmallVec1<T> {
    Empty,
    One(T),
    Many(Vec<T>),
}

impl<T: Debug> Debug for SmallVec1<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_list().entries(self.as_slice()).finish()
    }
}

impl<T> SmallVec1<T> {
    pub(crate) fn as_slice(&self) -> &[T] {
        match self {
            SmallVec1::Empty => &[],
            SmallVec1::One(x) => slice::from_ref(x),
            SmallVec1::Many(xs) => xs,
        }
    }

    pub(crate) fn extend(&mut self, that: SmallVec1<T>) {
        *self = match (mem::replace(self, SmallVec1::Empty), that) {
            (SmallVec1::Empty, right) => right,
            (left, SmallVec1::Empty) => left,
            (SmallVec1::One(left), SmallVec1::One(right)) => SmallVec1::Many(vec![left, right]),
            (SmallVec1::One(left), SmallVec1::Many(mut right)) => {
                right.insert(0, left);
                SmallVec1::Many(right)
            }
            (SmallVec1::Many(mut left), SmallVec1::One(right)) => {
                left.push(right);
                SmallVec1::Many(left)
            }
            (SmallVec1::Many(mut left), SmallVec1::Many(right)) => {
                left.extend(right);
                SmallVec1::Many(left)
            }
        }
    }
}


/*
 * 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::ops::Deref;

use crate::eval::runtime::frame_span::FrameSpan;

/// Similar to `Spanned<T>` but with file span.
///
/// For intermediate representation.
#[derive(Clone, Copy, Debug)]
pub(crate) struct IrSpanned<T> {
    pub(crate) span: FrameSpan,
    pub(crate) node: T,
}

impl<T> IrSpanned<T> {
    pub fn map<U>(&self, f: impl FnOnce(&T) -> U) -> IrSpanned<U> {
        IrSpanned {
            node: f(&self.node),
            span: self.span,
        }
    }
}

impl<T> Deref for IrSpanned<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.node
    }
}


/*
 * Copyright 2018 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.
 */

//! Evaluation environment, provide converters from Ast* element to value.
//!
//! # Starlark and BUILD dialect
//!
//! All evaluation function can evaluate the full Starlark language (i.e.
//! Bazel's .bzl files) or the BUILD file dialect (i.e. used to interpret
//! Bazel's BUILD file). The BUILD dialect does not allow `def` statements.

use std::cmp;

use starlark_derive::VisitSpanMut;
use thiserror::Error;

use crate::codemap::Span;
use crate::codemap::Spanned;
use crate::environment::slots::ModuleSlotId;
use crate::environment::FrozenModuleData;
use crate::eval::compiler::expr::Builtin1;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::expr::ExprLogicalBinOp;
use crate::eval::compiler::expr_bool::ExprCompiledBool;
use crate::eval::compiler::known::list_to_tuple;
use crate::eval::compiler::opt_ctx::OptCtx;
use crate::eval::compiler::scope::payload::CstAssign;
use crate::eval::compiler::scope::payload::CstExpr;
use crate::eval::compiler::scope::payload::CstStmt;
use crate::eval::compiler::scope::Captured;
use crate::eval::compiler::scope::Slot;
use crate::eval::compiler::small_vec_1::SmallVec1;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::Compiler;
use crate::eval::runtime::evaluator::Evaluator;
use crate::eval::runtime::evaluator::GC_THRESHOLD;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::eval::runtime::frozen_file_span::FrozenFileSpan;
use crate::eval::runtime::slots::LocalCapturedSlotId;
use crate::eval::runtime::slots::LocalSlotId;
use crate::slice_vec_ext::SliceExt;
use crate::syntax::ast::AssignOp;
use crate::syntax::ast::AssignP;
use crate::syntax::ast::DefP;
use crate::syntax::ast::StmtP;
use crate::values::dict::Dict;
use crate::values::dict::DictMut;
use crate::values::dict::DictRef;
use crate::values::types::list::value::ListData;
use crate::values::typing::TypeCompiled;
use crate::values::FrozenHeap;
use crate::values::FrozenValue;
use crate::values::Heap;
use crate::values::Value;
use crate::values::ValueError;

#[derive(Clone, Debug)]
pub(crate) enum AssignModifyLhs {
    Dot(IrSpanned<ExprCompiled>, String),
    Array(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>),
    Local(IrSpanned<LocalSlotId>),
    LocalCaptured(IrSpanned<LocalCapturedSlotId>),
    Module(IrSpanned<ModuleSlotId>),
}

#[derive(Clone, Debug)]
pub(crate) enum StmtCompiled {
    PossibleGc,
    Return(IrSpanned<ExprCompiled>),
    Expr(IrSpanned<ExprCompiled>),
    Assign(
        IrSpanned<AssignCompiledValue>,
        Option<IrSpanned<TypeCompiled<FrozenValue>>>,
        IrSpanned<ExprCompiled>,
    ),
    AssignModify(AssignModifyLhs, AssignOp, IrSpanned<ExprCompiled>),
    If(Box<(IrSpanned<ExprCompiled>, StmtsCompiled, StmtsCompiled)>),
    For(
        Box<(
            IrSpanned<AssignCompiledValue>,
            IrSpanned<ExprCompiled>,
            StmtsCompiled,
        )>,
    ),
    Break,
    Continue,
}

#[derive(Debug, Default)]
pub(crate) struct StmtCompileContext {
    /// Current function has return type.
    pub(crate) has_return_type: bool,
}

pub(crate) struct OptimizeOnFreezeContext<'v, 'a> {
    pub(crate) module: &'a FrozenModuleData,
    /// Nothing useful should be left in the heap after the freeze,
    /// but having a heap is useful to allocate objects temporarily
    /// (when invoking operations which require heap).
    pub(crate) heap: &'v Heap,
    pub(crate) frozen_heap: &'a FrozenHeap,
}

impl AssignModifyLhs {
    fn optimize(&self, ctx: &mut OptCtx) -> AssignModifyLhs {
        match self {
            AssignModifyLhs::Dot(expr, name) => {
                AssignModifyLhs::Dot(expr.optimize(ctx), name.clone())
            }
            AssignModifyLhs::Array(expr, index) => {
                AssignModifyLhs::Array(expr.optimize(ctx), index.optimize(ctx))
            }
            l @ (AssignModifyLhs::Local(..)
            | AssignModifyLhs::LocalCaptured(..)
            | AssignModifyLhs::Module(..)) => l.clone(),
        }
    }
}

impl IrSpanned<StmtCompiled> {
    fn optimize(&self, ctx: &mut OptCtx) -> StmtsCompiled {
        let span = self.span;
        match &self.node {
            StmtCompiled::Return(e) => StmtsCompiled::one(IrSpanned {
                span,
                node: StmtCompiled::Return(e.optimize(ctx)),
            }),
            StmtCompiled::Expr(expr) => {
                let expr = expr.optimize(ctx);
                StmtsCompiled::expr(expr)
            }
            StmtCompiled::Assign(lhs, ty, rhs) => {
                let lhs = lhs.optimize(ctx);
                let rhs = rhs.optimize(ctx);
                StmtsCompiled::one(IrSpanned {
                    span,
                    node: StmtCompiled::Assign(lhs, *ty, rhs),
                })
            }
            StmtCompiled::If(cond_t_f) => {
                let (cond, t, f) = &**cond_t_f;
                let cond = cond.optimize(ctx);
                let t = t.optimize(ctx);
                let f = f.optimize(ctx);
                StmtsCompiled::if_stmt(span, cond, t, f)
            }
            StmtCompiled::For(var_over_body) => {
                let (var, over, body) = &**var_over_body;
                let var = var.optimize(ctx);
                let over = over.optimize(ctx);
                let body = body.optimize(ctx);
                StmtsCompiled::for_stmt(span, var, over, body)
            }
            s @ (StmtCompiled::PossibleGc | StmtCompiled::Break | StmtCompiled::Continue) => {
                StmtsCompiled::one(IrSpanned {
                    span,
                    node: s.clone(),
                })
            }
            StmtCompiled::AssignModify(lhs, op, rhs) => StmtsCompiled::one(IrSpanned {
                span,
                node: StmtCompiled::AssignModify(lhs.optimize(ctx), *op, rhs.optimize(ctx)),
            }),
        }
    }
}

#[derive(Clone, Debug)]
pub(crate) struct StmtsCompiled(SmallVec1<IrSpanned<StmtCompiled>>);

impl StmtsCompiled {
    pub(crate) fn empty() -> StmtsCompiled {
        StmtsCompiled(SmallVec1::Empty)
    }

    pub(crate) fn one(stmt: IrSpanned<StmtCompiled>) -> StmtsCompiled {
        StmtsCompiled(SmallVec1::One(stmt))
    }

    pub(crate) fn is_empty(&self) -> bool {
        match &self.0 {
            SmallVec1::Empty => true,
            SmallVec1::One(_) => false,
            SmallVec1::Many(stmts) => {
                debug_assert!(stmts.len() > 1);
                false
            }
        }
    }

    pub(crate) fn stmts(&self) -> &[IrSpanned<StmtCompiled>] {
        self.0.as_slice()
    }

    /// Last statement in this block is `break`, `continue` or `return`.
    fn is_terminal(&self) -> bool {
        if let Some(stmt) = self.last() {
            match &stmt.node {
                StmtCompiled::Break | StmtCompiled::Continue | StmtCompiled::Return(..) => true,
                _ => false,
            }
        } else {
            false
        }
    }

    pub(crate) fn extend(&mut self, right: StmtsCompiled) {
        // Do not add any code after `break`, `continue` or `return`.
        if self.is_terminal() {
            return;
        }
        self.0.extend(right.0);
    }

    pub(crate) fn optimize(&self, ctx: &mut OptCtx) -> StmtsCompiled {
        let mut stmts = StmtsCompiled::empty();
        match &self.0 {
            SmallVec1::Empty => {}
            SmallVec1::One(s) => stmts.extend(s.optimize(ctx)),
            SmallVec1::Many(ss) => {
                for s in ss {
                    if stmts.is_terminal() {
                        break;
                    }
                    stmts.extend(s.optimize(ctx));
                }
            }
        }
        stmts
    }

    pub(crate) fn first(&self) -> Option<&IrSpanned<StmtCompiled>> {
        match &self.0 {
            SmallVec1::Empty => None,
            SmallVec1::One(s) => Some(s),
            SmallVec1::Many(ss) => ss.first(),
        }
    }

    pub(crate) fn last(&self) -> Option<&IrSpanned<StmtCompiled>> {
        match &self.0 {
            SmallVec1::Empty => None,
            SmallVec1::One(s) => Some(s),
            SmallVec1::Many(ss) => ss.last(),
        }
    }

    fn expr(expr: IrSpanned<ExprCompiled>) -> StmtsCompiled {
        let span = expr.span;
        match expr.node {
            expr if expr.is_pure_infallible() => StmtsCompiled::empty(),
            ExprCompiled::List(xs) | ExprCompiled::Tuple(xs) => {
                let mut stmts = StmtsCompiled::empty();
                for x in xs {
                    stmts.extend(Self::expr(x));
                }
                stmts
            }
            // Unwrap infallible expressions.
            ExprCompiled::Builtin1(Builtin1::Not | Builtin1::TypeIs(_), x) => Self::expr(*x),
            // "And" and "or" for effect are equivalent to `if`.
            ExprCompiled::LogicalBinOp(ExprLogicalBinOp::And, x_y) => {
                let (x, y) = *x_y;
                Self::if_stmt(expr.span, x, Self::expr(y), StmtsCompiled::empty())
            }
            ExprCompiled::LogicalBinOp(ExprLogicalBinOp::Or, x_y) => {
                let (x, y) = *x_y;
                Self::if_stmt(expr.span, x, StmtsCompiled::empty(), Self::expr(y))
            }
            expr => {
                if let Some(t) = expr.as_type() {
                    StmtsCompiled::expr(t.clone())
                } else {
                    StmtsCompiled::one(IrSpanned {
                        span,
                        node: StmtCompiled::Expr(IrSpanned { span, node: expr }),
                    })
                }
            }
        }
    }

    fn if_stmt(
        span: FrameSpan,
        cond: IrSpanned<ExprCompiled>,
        t: StmtsCompiled,
        f: StmtsCompiled,
    ) -> StmtsCompiled {
        let cond = ExprCompiledBool::new(cond);
        match cond.node {
            ExprCompiledBool::Const(true) => t,
            ExprCompiledBool::Const(false) => f,
            ExprCompiledBool::Expr(cond) => match cond {
                ExprCompiled::Builtin1(Builtin1::Not, cond) => Self::if_stmt(span, *cond, f, t),
                ExprCompiled::Seq(x_cond) => {
                    let (x, cond) = *x_cond;
                    let mut stmt = StmtsCompiled::empty();
                    stmt.extend(Self::expr(x));
                    stmt.extend(Self::if_stmt(span, cond, t, f));
                    stmt
                }
                cond => {
                    let cond = IrSpanned { span, node: cond };
                    if t.is_empty() && f.is_empty() {
                        Self::expr(cond)
                    } else {
                        StmtsCompiled::one(IrSpanned {
                            span,
                            node: StmtCompiled::If(Box::new((cond, t, f))),
                        })
                    }
                }
            },
        }
    }

    fn for_stmt(
        span: FrameSpan,
        var: IrSpanned<AssignCompiledValue>,
        over: IrSpanned<ExprCompiled>,
        body: StmtsCompiled,
    ) -> StmtsCompiled {
        if over.is_iterable_empty() {
            return StmtsCompiled::empty();
        }
        StmtsCompiled::one(IrSpanned {
            span,
            node: StmtCompiled::For(Box::new((var, over, body))),
        })
    }
}

#[derive(Debug, Error)]
pub(crate) enum AssignError {
    // Incorrect number of value to unpack (expected, got)
    #[error("Unpacked {1} values but expected {0}")]
    IncorrectNumberOfValueToUnpack(i32, i32),
}

#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) enum AssignCompiledValue {
    Dot(IrSpanned<ExprCompiled>, String),
    Index(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>),
    Tuple(Vec<IrSpanned<AssignCompiledValue>>),
    Local(LocalSlotId),
    LocalCaptured(LocalCapturedSlotId),
    Module(ModuleSlotId, String),
}

impl AssignCompiledValue {
    /// Assignment to a local non-captured variable.
    pub(crate) fn as_local_non_captured(&self) -> Option<LocalSlotId> {
        match self {
            AssignCompiledValue::Local(id) => Some(*id),
            _ => None,
        }
    }
}

impl IrSpanned<AssignCompiledValue> {
    pub(crate) fn optimize(&self, ctx: &mut OptCtx) -> IrSpanned<AssignCompiledValue> {
        let span = self.span;
        let assign = match self.node {
            AssignCompiledValue::Dot(ref object, ref field) => {
                let object = object.optimize(ctx);
                let field = field.clone();
                AssignCompiledValue::Dot(object, field)
            }
            AssignCompiledValue::Index(ref array, ref index) => {
                let array = array.optimize(ctx);
                let index = index.optimize(ctx);
                AssignCompiledValue::Index(array, index)
            }
            AssignCompiledValue::Tuple(ref xs) => {
                let xs = xs.map(|x| x.optimize(ctx));
                AssignCompiledValue::Tuple(xs)
            }
            ref e @ (AssignCompiledValue::Local(..)
            | AssignCompiledValue::LocalCaptured(..)
            | AssignCompiledValue::Module(..)) => e.clone(),
        };
        IrSpanned { node: assign, span }
    }
}

impl Compiler<'_, '_, '_> {
    pub fn assign(&mut self, expr: &CstAssign) -> IrSpanned<AssignCompiledValue> {
        let span = FrameSpan::new(FrozenFileSpan::new(self.codemap, expr.span));
        let assign = match &expr.node {
            AssignP::Dot(e, s) => {
                let e = self.expr(e);
                let s = &s.node;
                AssignCompiledValue::Dot(e, s.to_owned())
            }
            AssignP::Index(e_idx) => {
                let (e, idx) = &**e_idx;
                let e = self.expr(e);
                let idx = self.expr(idx);
                AssignCompiledValue::Index(e, idx)
            }
            AssignP::Tuple(v) => {
                let v = v.map(|x| self.assign(x));
                AssignCompiledValue::Tuple(v)
            }
            AssignP::Identifier(ident) => {
                let name = ident.node.0.as_str();
                let binding_id = ident
                    .node
                    .1
                    .unwrap_or_else(|| panic!("unresolved binding: `{}`", name));
                let binding = self.scope_data.get_binding(binding_id);
                let slot = binding
                    .slot
                    .unwrap_or_else(|| panic!("unresolved binding: `{}`", name));
                match (slot, binding.captured) {
                    (Slot::Local(slot), Captured::No) => {
                        AssignCompiledValue::Local(LocalSlotId(slot.0))
                    }
                    (Slot::Local(slot), Captured::Yes) => {
                        AssignCompiledValue::LocalCaptured(LocalCapturedSlotId(slot.0))
                    }
                    (Slot::Module(slot), _) => AssignCompiledValue::Module(slot, name.to_owned()),
                }
            }
        };
        IrSpanned { node: assign, span }
    }

    fn assign_modify(
        &mut self,
        span_stmt: Span,
        lhs: &CstAssign,
        rhs: IrSpanned<ExprCompiled>,
        op: AssignOp,
    ) -> StmtsCompiled {
        let span_stmt = FrameSpan::new(FrozenFileSpan::new(self.codemap, span_stmt));
        let span_lhs = FrameSpan::new(FrozenFileSpan::new(self.codemap, lhs.span));
        match &lhs.node {
            AssignP::Dot(e, s) => {
                let e = self.expr(e);
                StmtsCompiled::one(IrSpanned {
                    span: span_stmt,
                    node: StmtCompiled::AssignModify(
                        AssignModifyLhs::Dot(e, s.node.clone()),
                        op,
                        rhs,
                    ),
                })
            }
            AssignP::Index(e_idx) => {
                let (e, idx) = &**e_idx;
                let e = self.expr(e);
                let idx = self.expr(idx);
                StmtsCompiled::one(IrSpanned {
                    span: span_stmt,
                    node: StmtCompiled::AssignModify(AssignModifyLhs::Array(e, idx), op, rhs),
                })
            }
            AssignP::Identifier(ident) => {
                let (slot, captured) = self.scope_data.get_assign_ident_slot(&ident);
                match (slot, captured) {
                    (Slot::Local(slot), Captured::No) => {
                        let lhs = IrSpanned {
                            node: LocalSlotId(slot.0),
                            span: span_lhs,
                        };
                        StmtsCompiled::one(IrSpanned {
                            span: span_stmt,
                            node: StmtCompiled::AssignModify(AssignModifyLhs::Local(lhs), op, rhs),
                        })
                    }
                    (Slot::Local(slot), Captured::Yes) => {
                        let lhs = IrSpanned {
                            node: LocalCapturedSlotId(slot.0),
                            span: span_lhs,
                        };
                        StmtsCompiled::one(IrSpanned {
                            span: span_stmt,
                            node: StmtCompiled::AssignModify(
                                AssignModifyLhs::LocalCaptured(lhs),
                                op,
                                rhs,
                            ),
                        })
                    }
                    (Slot::Module(slot), _) => {
                        let lhs = IrSpanned {
                            node: slot,
                            span: span_lhs,
                        };
                        StmtsCompiled::one(IrSpanned {
                            span: span_stmt,
                            node: StmtCompiled::AssignModify(AssignModifyLhs::Module(lhs), op, rhs),
                        })
                    }
                }
            }
            AssignP::Tuple(_) => {
                unreachable!("Assign modify validates that the LHS is never a tuple")
            }
        }
    }
}

// There are two requirements to perform a GC:
//
// 1. We can't be profiling, since profiling relies on the redundant heap
//    entries. When profiling we set disable_gc.
// 2. We must be able to access all roots.
//
// We track as many roots as possible, and eventually aim to track them all, but
// for the moment we're only sure we have all roots when we are in the module
// evaluation eval. There are three roots we don't yet know about:
//
// 1. When evaluating an expression which has multiple subexpressions, e.g. List
//    we can't GC during that, as we can't see the root of the first list.
// 2. When evaluating inside a native function, especially if that native
//    function calls back to a non-native function, e.g. sort with a comparison
//    function.
// 3. When iterating we freeze the iteration variable, which means it
//    can't be moved by a GC. A special type of root.
//
// The first issue can be solved by moving to a bytecode interpreter and
// evaluation stack. The second issue can be solved by disabling GC while in
// such functions (it's probably rare). The third issue could be solved by
// making the freeze for iteration a separate flag to the RefCell, at the cost
// of an extra word in ValueMem. Or we could disable GC while iterating.
//
// For the moment we only GC when executing a statement at the root of the
// module, which we know is safe with respect to all three conditions.
//
// We also require that `extra_v` is None, since otherwise the user might have
// additional values stashed somewhere.
pub(crate) fn possible_gc(eval: &mut Evaluator) {
    if !eval.disable_gc && eval.heap().allocated_bytes() >= eval.next_gc_level {
        // When we are at a module scope (as checked above) the eval contains
        // references to all values, so walking covers everything and the unsafe
        // is satisfied.
        unsafe { eval.garbage_collect() }
        eval.next_gc_level = cmp::max(eval.heap().allocated_bytes() * 2, GC_THRESHOLD);
    }
}

/// Implement lhs |= rhs, which is special in Starlark, because dicts are mutated,
/// while all other types are not.
pub(crate) fn bit_or_assign<'v>(
    lhs: Value<'v>,
    rhs: Value<'v>,
    heap: &'v Heap,
) -> anyhow::Result<Value<'v>> {
    // The Starlark spec says dict |= mutates, while nothing else does.
    // When mutating, be careful if they alias, so we don't have `lhs`
    // mutably borrowed when we iterate over `rhs`, as they might alias.

    let lhs_aref = lhs.get_ref();
    let lhs_ty = lhs_aref.static_type_of_value();

    if Dict::is_dict_type(lhs_ty) {
        let mut dict = DictMut::from_value(lhs)?;
        if lhs.ptr_eq(rhs) {
            // Nothing to do as union is idempotent
        } else {
            let rhs = DictRef::from_value(rhs).map_or_else(
                || ValueError::unsupported_owned(lhs_aref.get_type(), "|=", Some(rhs.get_type())),
                Ok,
            )?;
            for (k, v) in rhs.iter_hashed() {
                dict.insert_hashed(k, v);
            }
        }
        Ok(lhs)
    } else {
        lhs_aref.bit_or(rhs, heap)
    }
}

/// Implement lhs += rhs, which is special in Starlark, because lists are mutated,
/// while all other types are not.
pub(crate) fn add_assign<'v>(
    lhs: Value<'v>,
    rhs: Value<'v>,
    heap: &'v Heap,
) -> anyhow::Result<Value<'v>> {
    // Checking whether a value is an integer or a string is cheap (no virtual call),
    // and `Value::add` has optimizations for these types, so check them first
    // and delegate to `Value::add`.
    if lhs.unpack_inline_int().is_some() || lhs.is_str() {
        return lhs.add(rhs, heap);
    }

    // The Starlark spec says list += mutates, while nothing else does.
    // When mutating, be careful if they alias, so we don't have `lhs`
    // mutably borrowed when we iterate over `rhs`, as they might alias.

    // In practice, select is the only thing that implements radd.
    // If the users does x += select(...) we don't want an error,
    // we really want to x = x + select, so check radd first.
    let lhs_aref = lhs.get_ref();
    let lhs_ty = lhs_aref.static_type_of_value();

    if ListData::is_list_type(lhs_ty) {
        if let Some(v) = rhs.get_ref().radd(lhs, heap) {
            v
        } else {
            let list = ListData::from_value_mut(lhs)?;
            if lhs.ptr_eq(rhs) {
                list.double(heap);
            } else {
                // TODO: if RHS is list, consider calling `List::extend_from_slice`.
                list.extend(rhs.iterate(heap)?, heap);
            }
            Ok(lhs)
        }
    } else {
        lhs.add(rhs, heap)
    }
}

impl Compiler<'_, '_, '_> {
    pub(crate) fn compile_context(&self, has_return_type: bool) -> StmtCompileContext {
        StmtCompileContext { has_return_type }
    }

    pub(crate) fn stmt(&mut self, stmt: &CstStmt, allow_gc: bool) -> StmtsCompiled {
        let span = FrameSpan::new(FrozenFileSpan::new(self.codemap, stmt.span));
        let is_statements = matches!(&stmt.node, StmtP::Statements(_));
        let res = self.stmt_direct(stmt, allow_gc);
        // No point inserting a GC point around statements, since they will contain inner statements we can do
        if allow_gc && !is_statements {
            // We could do this more efficiently by fusing the possible_gc
            // into the inner closure, but no real need - we insert allow_gc fairly rarely
            let mut with_gc = StmtsCompiled::one(IrSpanned {
                span,
                node: StmtCompiled::PossibleGc,
            });
            with_gc.extend(res);
            with_gc
        } else {
            res
        }
    }

    pub(crate) fn module_top_level_stmt(&mut self, stmt: &CstStmt) -> StmtsCompiled {
        match &stmt.node {
            StmtP::Statements(..) => {
                unreachable!("top level statement lists are handled by outer loop")
            }
            StmtP::Expression(expr) => {
                let stmt = Spanned {
                    span: expr.span,
                    // When top level statement is an expression, compile it as return.
                    // This is used to obtain the result of evaluation
                    // of the last statement-expression in module.
                    // TODO(nga): unnecessary clone.
                    node: StmtP::Return(Some(expr.clone())),
                };
                self.stmt(&stmt, true)
            }
            _ => self.stmt(stmt, true),
        }
    }

    fn stmt_if(
        &mut self,
        span: FrameSpan,
        cond: &CstExpr,
        then_block: &CstStmt,
        allow_gc: bool,
    ) -> StmtsCompiled {
        let cond = self.expr(cond);
        let then_block = self.stmt(then_block, allow_gc);
        StmtsCompiled::if_stmt(span, cond, then_block, StmtsCompiled::empty())
    }

    fn stmt_if_else(
        &mut self,
        span: FrameSpan,
        cond: &CstExpr,
        then_block: &CstStmt,
        else_block: &CstStmt,
        allow_gc: bool,
    ) -> StmtsCompiled {
        let cond = self.expr(cond);
        let then_block = self.stmt(then_block, allow_gc);
        let else_block = self.stmt(else_block, allow_gc);
        StmtsCompiled::if_stmt(span, cond, then_block, else_block)
    }

    fn stmt_expr(&mut self, expr: &CstExpr) -> StmtsCompiled {
        let expr = self.expr(expr);
        StmtsCompiled::expr(expr)
    }

    fn stmt_direct(&mut self, stmt: &CstStmt, allow_gc: bool) -> StmtsCompiled {
        let span = FrameSpan::new(FrozenFileSpan::new(self.codemap, stmt.span));
        match &stmt.node {
            StmtP::Def(def) => {
                let signature_span = def.signature_span();
                let signature_span = FrozenFileSpan::new(self.codemap, signature_span);
                let DefP {
                    name,
                    params,
                    return_type,
                    body,
                    payload: scope_id,
                } = def;
                let rhs = IrSpanned {
                    node: self.function(
                        &name.0,
                        signature_span,
                        *scope_id,
                        params,
                        return_type.as_deref(),
                        body,
                    ),
                    span,
                };
                let lhs = self.assign(&Spanned {
                    span: name.span,
                    node: AssignP::Identifier(name.clone()),
                });
                StmtsCompiled::one(IrSpanned {
                    span,
                    node: StmtCompiled::Assign(lhs, None, rhs),
                })
            }
            StmtP::For(var, over_body) => {
                let (over, body) = &**over_body;
                let over = list_to_tuple(over);
                let var = self.assign(var);
                let over = self.expr(&over);
                let st = self.stmt(body, false);
                StmtsCompiled::for_stmt(span, var, over, st)
            }
            StmtP::Return(None) => StmtsCompiled::one(IrSpanned {
                node: StmtCompiled::Return(IrSpanned {
                    span,
                    node: ExprCompiled::Value(FrozenValue::new_none()),
                }),
                span,
            }),
            StmtP::Return(Some(e)) => StmtsCompiled::one(IrSpanned {
                node: StmtCompiled::Return(self.expr(e)),
                span,
            }),
            StmtP::If(cond, then_block) => self.stmt_if(span, cond, then_block, allow_gc),
            StmtP::IfElse(cond, then_block_else_block) => {
                let (then_block, else_block) = &**then_block_else_block;
                self.stmt_if_else(span, cond, then_block, else_block, allow_gc)
            }
            StmtP::Statements(stmts) => {
                let mut r = StmtsCompiled::empty();
                for stmt in stmts {
                    if r.is_terminal() {
                        break;
                    }
                    r.extend(self.stmt(stmt, allow_gc));
                }
                r
            }
            StmtP::Expression(e) => self.stmt_expr(e),
            StmtP::Assign(lhs, ty_rhs) => {
                let (ty, rhs) = &**ty_rhs;
                let rhs = self.expr(rhs);
                let ty = self.expr_for_type(ty.as_ref());
                let lhs = self.assign(lhs);
                StmtsCompiled::one(IrSpanned {
                    span,
                    node: StmtCompiled::Assign(lhs, ty, rhs),
                })
            }
            StmtP::AssignModify(lhs, op, rhs) => {
                let rhs = self.expr(rhs);
                self.assign_modify(span.span.span(), lhs, rhs, *op)
            }
            StmtP::Load(..) => unreachable!(),
            StmtP::Pass => StmtsCompiled::empty(),
            StmtP::Break => StmtsCompiled::one(IrSpanned {
                span,
                node: StmtCompiled::Break,
            }),
            StmtP::Continue => StmtsCompiled::one(IrSpanned {
                span,
                node: StmtCompiled::Continue,
            }),
        }
    }
}


/*
 * 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 crate::codemap::Span;
use crate::codemap::Spanned;
use crate::eval::compiler::constants::Constants;
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::scope::ResolvedIdent;
use crate::eval::compiler::scope::Slot;
use crate::eval::compiler::scope::TopLevelStmtIndex;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::Compiler;
use crate::eval::compiler::EvalException;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::eval::runtime::frozen_file_span::FrozenFileSpan;
use crate::slice_vec_ext::VecExt;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::StmtP;
use crate::syntax::type_expr::TypeExprUnpackP;
use crate::syntax::uniplate::VisitMut;
use crate::values::typing::TypeCompiled;
use crate::values::FrozenValue;
use crate::values::Value;

#[derive(Debug, thiserror::Error)]
enum TypesError {
    #[error("Type already initialized (internal error)")]
    TypeAlreadySet,
    #[error("Identifier is not resolved (internal error)")]
    UnresolvedIdentifier,
    #[error("Identifier is resolve as local variable (internal error)")]
    LocalIdentifier,
    #[error("Module variable is not set: `{0}`")]
    ModuleVariableNotSet(String),
    #[error("Type payload not set (internal error)")]
    TypePayloadNotSet,
    #[error("[] can only be applied to list function in type expression")]
    TypeIndexOnNonList,
    #[error("[,] can only be applied to dict function in type expression")]
    TypeIndexOnNonDict,
}

impl<'v> Compiler<'v, '_, '_> {
    /// Compile expression when it is expected to be interpreted as type.
    pub(crate) fn expr_for_type(
        &mut self,
        expr: Option<&CstTypeExpr>,
    ) -> Option<IrSpanned<TypeCompiled<FrozenValue>>> {
        if !self.check_types {
            return None;
        }
        let expr = expr?;
        let span = FrameSpan::new(FrozenFileSpan::new(self.codemap, expr.span));
        let Some(type_value) = expr.payload else {
            // This is unreachable. But unfortunately we do not return error here.
            // Still make an error in panic to produce nice panic message.
            panic!("{:?}", EvalException::new(
                TypesError::TypePayloadNotSet.into(),
                expr.span,
                &self.codemap));
        };
        if type_value.type_is_wildcard() {
            return None;
        }
        Some(IrSpanned {
            span,
            node: type_value,
        })
    }

    /// We evaluated type expression to `Value`, now convert it to `FrozenValue`.
    // TODO(nga): this step is not really necessary, we should just create `TypeCompiled` directly.
    fn alloc_value_for_type(
        &mut self,
        value: Value<'v>,
        span: Span,
    ) -> Result<TypeCompiled<Value<'v>>, EvalException> {
        let ty = TypeCompiled::new(value, self.eval.heap());
        ty.map_err(|e| EvalException::new(e, span, &self.codemap))
    }

    fn eval_ident_in_type_expr(&mut self, ident: &CstIdent) -> Result<Value<'v>, EvalException> {
        let Some(ident_payload) = &ident.node.1 else {
            return Err(EvalException::new(
                TypesError::UnresolvedIdentifier.into(),
                ident.span,
                &self.codemap,
            ));
        };
        match ident_payload {
            ResolvedIdent::Slot(Slot::Local(..), _) => Err(EvalException::new(
                TypesError::LocalIdentifier.into(),
                ident.span,
                &self.codemap,
            )),
            ResolvedIdent::Slot(Slot::Module(module_slot_id), _) => {
                match self.eval.module_env.slots().get_slot(*module_slot_id) {
                    Some(v) => Ok(v),
                    None => Err(EvalException::new(
                        TypesError::ModuleVariableNotSet(ident.node.0.clone()).into(),
                        ident.span,
                        &self.codemap,
                    )),
                }
            }
            ResolvedIdent::Global(v) => Ok(v.to_value()),
        }
    }

    /// We may use non-frozen values as types, so we don't reuse `expr_ident` function
    /// which is used in normal compilation.
    fn eval_path_as_type(
        &mut self,
        first: &CstIdent,
        rem: &[Spanned<&str>],
    ) -> Result<TypeCompiled<Value<'v>>, EvalException> {
        let mut value = self.eval_ident_in_type_expr(first)?;
        for step in rem {
            value = value
                .get_attr_error(step.node, self.eval.heap())
                .map_err(|e| EvalException::new(e, step.span, &self.codemap))?;
        }
        let mut span = first.span;
        if let Some(last) = rem.last() {
            span = span.merge(last.span);
        }
        self.alloc_value_for_type(value, span)
    }

    fn eval_expr_as_type(
        &mut self,
        expr: Spanned<TypeExprUnpackP<CstPayload>>,
    ) -> Result<TypeCompiled<Value<'v>>, EvalException> {
        match expr.node {
            TypeExprUnpackP::Path(ident, rem) => self.eval_path_as_type(ident, &rem),
            TypeExprUnpackP::Index(a, i) => {
                let a = self.eval_ident_in_type_expr(a)?;
                if !a.ptr_eq(Constants::get().fn_list.0.to_value()) {
                    return Err(EvalException::new(
                        TypesError::TypeIndexOnNonList.into(),
                        expr.span,
                        &self.codemap,
                    ));
                }
                let i = self.eval_expr_as_type(*i)?;
                let t = a
                    .get_ref()
                    .at(i.to_inner(), self.eval.heap())
                    .map_err(|e| EvalException::new(e, expr.span, &self.codemap))?;
                Ok(TypeCompiled::new(t, self.eval.heap())
                    .map_err(|e| EvalException::new(e, expr.span, &self.codemap))?)
            }
            TypeExprUnpackP::Index2(a, i0, i1) => {
                let a = self.eval_ident_in_type_expr(a)?;
                if !a.ptr_eq(Constants::get().fn_dict.0.to_value()) {
                    return Err(EvalException::new(
                        TypesError::TypeIndexOnNonDict.into(),
                        expr.span,
                        &self.codemap,
                    ));
                }
                let i0 = self.eval_expr_as_type(*i0)?;
                let i1 = self.eval_expr_as_type(*i1)?;
                let t = a
                    .get_ref()
                    .at2(i0.to_inner(), i1.to_inner(), self.eval.heap())
                    .map_err(|e| EvalException::new(e, expr.span, &self.codemap))?;
                Ok(TypeCompiled::new(t, self.eval.heap())
                    .map_err(|e| EvalException::new(e, expr.span, &self.codemap))?)
            }
            TypeExprUnpackP::Any(xs) => {
                let xs = xs.into_try_map(|x| self.eval_expr_as_type(x))?;
                Ok(TypeCompiled::type_any_of(xs, self.eval.heap()))
            }
            TypeExprUnpackP::ListOf(x) => {
                let x = self.eval_expr_as_type(*x)?;
                Ok(TypeCompiled::type_list_of(x, self.eval.heap()))
            }
            TypeExprUnpackP::DictOf(k, v) => {
                let k = self.eval_expr_as_type(*k)?;
                let v = self.eval_expr_as_type(*v)?;
                Ok(TypeCompiled::type_dict_of(k, v, self.eval.heap()))
            }
            TypeExprUnpackP::Tuple(xs) => {
                let xs = xs.into_try_map(|x| self.eval_expr_as_type(x))?;
                Ok(TypeCompiled::type_tuple_of(xs, self.eval.heap()))
            }
            TypeExprUnpackP::Literal(s) => Ok(TypeCompiled::from_str(s.node, self.eval.heap())),
        }
    }

    fn populate_types_in_type_expr(
        &mut self,
        type_expr: &mut CstTypeExpr,
    ) -> Result<(), EvalException> {
        if type_expr.payload.is_some() {
            return Err(EvalException::new(
                TypesError::TypeAlreadySet.into(),
                type_expr.span,
                &self.codemap,
            ));
        }
        // This should not fail because we validated it at parse time.
        let unpack = TypeExprUnpackP::unpack(&type_expr.expr, &self.codemap)?;
        let type_value = self.eval_expr_as_type(unpack)?;
        let type_value = type_value.to_frozen(self.eval.frozen_heap());
        type_expr.payload = Some(type_value);
        Ok(())
    }

    fn populate_types_in_params(
        &mut self,
        params: &mut [CstParameter],
    ) -> Result<(), EvalException> {
        for param in params {
            let (_, ty, _) = param.split_mut();
            if let Some(ty) = ty {
                self.populate_types_in_type_expr(ty)?;
            }
        }
        Ok(())
    }

    fn populate_types_in_expr(&mut self, expr: &mut CstExpr) -> Result<(), EvalException> {
        match &mut expr.node {
            ExprP::Lambda(lambda) => {
                self.populate_types_in_params(&mut lambda.params)?;
            }
            _ => {}
        }
        expr.visit_expr_err_mut(|expr| self.populate_types_in_expr(expr))
    }

    fn populate_types_in_stmt(&mut self, stmt: &mut CstStmt) -> Result<(), EvalException> {
        match &mut stmt.node {
            StmtP::Assign(_lhs, ty_def) => {
                let (ty, _) = &mut **ty_def;
                if let Some(ty) = ty {
                    self.populate_types_in_type_expr(ty)?;
                }
            }
            StmtP::Def(def) => {
                self.populate_types_in_params(&mut def.params)?;
                if let Some(ret) = &mut def.return_type {
                    self.populate_types_in_type_expr(ret)?;
                }
            }
            _ => {}
        }
        stmt.visit_children_err_mut(|visit| match visit {
            VisitMut::Stmt(stmt) => self.populate_types_in_stmt(stmt),
            VisitMut::Expr(expr) => self.populate_types_in_expr(expr),
        })
    }

    pub(crate) fn populate_types_in_stmts(
        &mut self,
        stmts: &mut [CstStmt],
        up_to: TopLevelStmtIndex,
    ) -> Result<(), EvalException> {
        if self.last_stmt_with_populated_types >= up_to {
            return Ok(());
        }
        for stmt in &mut stmts[self.last_stmt_with_populated_types.0..up_to.0] {
            self.populate_types_in_stmt(stmt)?;
        }
        self.last_stmt_with_populated_types = up_to;
        Ok(())
    }
}


/*
 * 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 starlark_derive::VisitSpanMut;

use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::span::IrSpanned;
use crate::values::typing::TypeCompiled;

/// IR expression in type position.
#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) struct TypeExprCompiled {
    expr: IrSpanned<ExprCompiled>,
}

impl TypeExprCompiled {
    pub(crate) fn new_expr(expr: IrSpanned<ExprCompiled>) -> TypeExprCompiled {
        TypeExprCompiled { expr }
    }

    /// True if type is known to match any type.
    pub(crate) fn is_wildcard(&self) -> bool {
        if let Some(value) = self.expr.as_value() {
            TypeCompiled::is_wildcard_value(value.to_value())
        } else {
            false
        }
    }
}


/*
 * 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.
 */

//! Local slot id as value.
//!
//! To be able to propagate the local slot number through parameter binding machinery.

use std::array;

use allocative::Allocative;
use once_cell::sync::Lazy;
use starlark_derive::starlark_value;
use starlark_derive::NoSerialize;

use crate as starlark;
use crate::eval::runtime::slots::LocalSlotId;
use crate::starlark_simple_value;
use crate::values::FrozenHeap;
use crate::values::FrozenHeapRef;
use crate::values::FrozenValueTyped;
use crate::values::ProvidesStaticType;
use crate::values::StarlarkValue;

/// Local slot id as `FrozenValue`. This object only using during compilation
/// and never appears in the executed program.
#[derive(
    derive_more::Display,
    Debug,
    ProvidesStaticType,
    NoSerialize,
    Allocative
)]
#[display(fmt = "{:?}", self)]
pub(crate) struct LocalAsValue {
    #[allocative(skip)]
    pub(crate) local: LocalSlotId,
}

#[starlark_value(type = "LocalAsValue")]
impl<'v> StarlarkValue<'v> for LocalAsValue {}

starlark_simple_value!(LocalAsValue);

/// Create a value which represents a reference to local slot id during optimization.
pub(crate) fn local_as_value(
    local: LocalSlotId,
) -> Option<FrozenValueTyped<'static, LocalAsValue>> {
    // 100 is practically enough.
    static LOCALS: Lazy<(
        FrozenHeapRef,
        [FrozenValueTyped<'static, LocalAsValue>; 100],
    )> = Lazy::new(|| {
        let heap = FrozenHeap::new();
        let locals = array::from_fn(|i| {
            heap.alloc_simple_typed(LocalAsValue {
                local: LocalSlotId(i as u32),
            })
        });
        (heap.into_ref(), locals)
    });
    LOCALS.1.get(local.0 as usize).copied()
}


/*
 * 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.
 */

//! Inline functions.

pub(crate) mod local_as_value;

use crate::eval::compiler::args::ArgsCompiledValue;
use crate::eval::compiler::call::CallCompiled;
use crate::eval::compiler::def::ParametersCompiled;
use crate::eval::compiler::def_inline::local_as_value::LocalAsValue;
use crate::eval::compiler::expr::Builtin1;
use crate::eval::compiler::expr::Builtin2;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::expr::ExprLogicalBinOp;
use crate::eval::compiler::opt_ctx::OptCtx;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::stmt::StmtCompiled;
use crate::eval::compiler::stmt::StmtsCompiled;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::eval::runtime::slots::LocalSlotId;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;
use crate::values::FrozenValueTyped;

/// Function body suitable for inlining.
#[derive(Debug)]
pub(crate) enum InlineDefBody {
    /// Function body is `return type(x) == "y"`
    ReturnTypeIs(FrozenStringValue),
    /// Any expression which can be safely inlined.
    ///
    /// See the function where this enum variant is computed for the definition
    /// of safe to inline expression.
    ReturnSafeToInlineExpr(IrSpanned<ExprCompiled>),
}

/// If a statement is `return type(x) == "y"` where `x` is a first slot.
fn is_return_type_is(stmt: &StmtsCompiled) -> Option<FrozenStringValue> {
    let (x, t) = stmt.first()?.as_return()?.as_type_is()?;
    match &x.node {
        // Slot 0 is a slot for the first function parameter.
        ExprCompiled::Local(LocalSlotId(0)) => Some(t),
        _ => None,
    }
}

struct IsSafeToInlineExpr {
    /// Function parameter count.
    param_count: u32,
    /// How many expressions we visited already.
    counter: u32,
}

impl IsSafeToInlineExpr {
    fn new(param_count: u32) -> IsSafeToInlineExpr {
        Self {
            param_count,
            counter: 0,
        }
    }

    fn is_safe_to_inline_opt_expr(&mut self, expr: &Option<IrSpanned<ExprCompiled>>) -> bool {
        if let Some(expr) = expr {
            self.is_safe_to_inline_expr(expr)
        } else {
            true
        }
    }

    /// Expression which is has no access to locals or globals.
    fn is_safe_to_inline_expr(&mut self, expr: &ExprCompiled) -> bool {
        // Do not inline too large functions.
        if self.counter > 100 {
            return false;
        }
        self.counter += 1;
        match expr {
            ExprCompiled::Value(..) => true,
            ExprCompiled::LocalCaptured(..) | ExprCompiled::Module(..) | ExprCompiled::Def(..) => {
                false
            }
            ExprCompiled::Local(l) => {
                // `l >= param_count` should be unreachable, but it is safer this way.
                l.0 < self.param_count
            }
            ExprCompiled::Call(call) => {
                self.is_safe_to_inline_expr(&call.fun)
                    && call
                        .args
                        .arg_exprs()
                        .all(|e| self.is_safe_to_inline_expr(e))
            }
            ExprCompiled::Compr(..) => {
                // TODO: some comprehensions are safe to inline.
                false
            }
            ExprCompiled::Slice(a_b_c_d) => {
                let (a, b, c, d) = &**a_b_c_d;
                self.is_safe_to_inline_expr(a)
                    && self.is_safe_to_inline_opt_expr(b)
                    && self.is_safe_to_inline_opt_expr(c)
                    && self.is_safe_to_inline_opt_expr(d)
            }
            ExprCompiled::Builtin2(bin_op, a_b) => {
                let (a, b) = &**a_b;
                let _: &Builtin2 = bin_op;
                self.is_safe_to_inline_expr(a) && self.is_safe_to_inline_expr(b)
            }
            ExprCompiled::Index2(a_i0_i1) => {
                let (a, i0, i1) = &**a_i0_i1;
                self.is_safe_to_inline_expr(a)
                    && self.is_safe_to_inline_expr(i0)
                    && self.is_safe_to_inline_expr(i1)
            }
            ExprCompiled::Builtin1(un_op, arg) => {
                let _: &Builtin1 = un_op;
                self.is_safe_to_inline_expr(arg)
            }
            ExprCompiled::Tuple(xs) | ExprCompiled::List(xs) => {
                xs.iter().all(|x| self.is_safe_to_inline_expr(x))
            }
            ExprCompiled::Dict(xs) => xs
                .iter()
                .all(|(x, y)| self.is_safe_to_inline_expr(x) && self.is_safe_to_inline_expr(y)),
            ExprCompiled::If(c_t_f) => {
                let (c, t, f) = &**c_t_f;
                self.is_safe_to_inline_expr(c)
                    && self.is_safe_to_inline_expr(t)
                    && self.is_safe_to_inline_expr(f)
            }
            ExprCompiled::LogicalBinOp(op, x_y) => {
                let (x, y) = &**x_y;
                let _: &ExprLogicalBinOp = op;
                self.is_safe_to_inline_expr(x) && self.is_safe_to_inline_expr(y)
            }
            ExprCompiled::Seq(x_y) => {
                let (x, y) = &**x_y;
                self.is_safe_to_inline_expr(x) && self.is_safe_to_inline_expr(y)
            }
        }
    }
}

/// Function body is a `return` safe to inline expression (as defined above).
fn is_return_safe_to_inline_expr(
    stmts: &StmtsCompiled,
    param_count: u32,
) -> Option<IrSpanned<ExprCompiled>> {
    match stmts.first() {
        None => {
            // Empty function is equivalent to `return None`.
            Some(IrSpanned {
                span: FrameSpan::default(),
                node: ExprCompiled::Value(FrozenValue::new_none()),
            })
        }
        Some(stmt) => match &stmt.node {
            StmtCompiled::Return(expr)
                if IsSafeToInlineExpr::new(param_count).is_safe_to_inline_expr(expr) =>
            {
                Some(expr.clone())
            }
            _ => None,
        },
    }
}

pub(crate) fn inline_def_body(
    params: &ParametersCompiled<IrSpanned<ExprCompiled>>,
    body: &StmtsCompiled,
) -> Option<InlineDefBody> {
    if params.params.len() == 1 && params.params[0].accepts_positional() {
        if let Some(t) = is_return_type_is(body) {
            return Some(InlineDefBody::ReturnTypeIs(t));
        }
    }
    if !params.has_args_or_kwargs() {
        // It is possible to sometimes inline functions with `*args` or `**kwargs`,
        // but let's postpone that for now.
        let param_count = params.count_param_variables();
        if let Some(expr) = is_return_safe_to_inline_expr(body, param_count) {
            return Some(InlineDefBody::ReturnSafeToInlineExpr(expr));
        }
    }
    None
}

pub(crate) struct CannotInline;

/// Utility to inline function body at call site.
pub(crate) struct InlineDefCallSite<'s, 'v, 'a, 'e> {
    pub(crate) ctx: &'s mut OptCtx<'v, 'a, 'e>,
    // Values in the slots are either real frozen values
    // or `LocalAsValue` which are the parameters to be substituted with caller locals.
    pub(crate) slots: &'s [FrozenValue],
}

impl<'s, 'v, 'a, 'e> InlineDefCallSite<'s, 'v, 'a, 'e> {
    fn inline_opt(
        &mut self,
        expr: Option<&IrSpanned<ExprCompiled>>,
    ) -> Result<Option<IrSpanned<ExprCompiled>>, CannotInline> {
        match expr {
            None => Ok(None),
            Some(expr) => Ok(Some(self.inline(expr)?)),
        }
    }

    fn inline_args(&mut self, args: &ArgsCompiledValue) -> Result<ArgsCompiledValue, CannotInline> {
        args.map_exprs(|expr| self.inline(expr))
    }

    fn inline_call(
        &mut self,
        call: &IrSpanned<CallCompiled>,
    ) -> Result<IrSpanned<ExprCompiled>, CannotInline> {
        let span = call.span;
        let CallCompiled { fun, args } = &call.node;
        let fun = self.inline(fun)?;
        let args = self.inline_args(args)?;
        Ok(IrSpanned {
            span,
            node: CallCompiled::call(span, fun, args, self.ctx),
        })
    }

    pub(crate) fn inline(
        &mut self,
        expr: &IrSpanned<ExprCompiled>,
    ) -> Result<IrSpanned<ExprCompiled>, CannotInline> {
        let span = expr.span;
        Ok(match &expr.node {
            e @ ExprCompiled::Value(..) => IrSpanned {
                span,
                node: e.clone(),
            },
            ExprCompiled::Local(local) => {
                let value = self.slots[local.0 as usize];
                let expr = if let Some(local) = FrozenValueTyped::<LocalAsValue>::new(value) {
                    ExprCompiled::Local(local.local)
                } else {
                    ExprCompiled::Value(value)
                };
                IrSpanned { span, node: expr }
            }
            ExprCompiled::If(c_t_f) => {
                let (c, t, f) = &**c_t_f;
                let c = self.inline(c)?;
                let t = self.inline(t)?;
                let f = self.inline(f)?;
                ExprCompiled::if_expr(c, t, f)
            }
            ExprCompiled::LogicalBinOp(op, l_r) => {
                let (l, r) = &**l_r;
                let l = self.inline(l)?;
                let r = self.inline(r)?;
                ExprCompiled::logical_bin_op(*op, l, r)
            }
            ExprCompiled::List(xs) => {
                let xs = xs
                    .iter()
                    .map(|x| self.inline(x))
                    .collect::<Result<Vec<_>, CannotInline>>()?;
                IrSpanned {
                    span,
                    node: ExprCompiled::List(xs),
                }
            }
            ExprCompiled::Tuple(xs) => {
                let xs = xs
                    .iter()
                    .map(|x| self.inline(x))
                    .collect::<Result<Vec<_>, CannotInline>>()?;
                IrSpanned {
                    span,
                    node: ExprCompiled::tuple(xs, self.ctx.frozen_heap()),
                }
            }
            ExprCompiled::Dict(xs) => {
                let xs = xs
                    .iter()
                    .map(|(x, y)| Ok((self.inline(x)?, self.inline(y)?)))
                    .collect::<Result<Vec<_>, CannotInline>>()?;
                IrSpanned {
                    span,
                    node: ExprCompiled::Dict(xs),
                }
            }
            ExprCompiled::Builtin2(op, l_r) => {
                let (l, r) = &**l_r;
                let l = self.inline(l)?;
                let r = self.inline(r)?;
                IrSpanned {
                    span,
                    node: ExprCompiled::bin_op(*op, l, r, self.ctx),
                }
            }
            ExprCompiled::Index2(a_i0_i1) => {
                let (a, i0, i1) = &**a_i0_i1;
                let a = self.inline(a)?;
                let i0 = self.inline(i0)?;
                let i1 = self.inline(i1)?;
                IrSpanned {
                    span,
                    node: ExprCompiled::Index2(Box::new((a, i0, i1))),
                }
            }
            ExprCompiled::Builtin1(op, x) => {
                let x = self.inline(x)?;
                IrSpanned {
                    span,
                    node: ExprCompiled::un_op(span, op, x, self.ctx),
                }
            }
            ExprCompiled::Slice(l_a_b_c) => {
                let (l, a, b, c) = &**l_a_b_c;
                let l = self.inline(l)?;
                let a = self.inline_opt(a.as_ref())?;
                let b = self.inline_opt(b.as_ref())?;
                let c = self.inline_opt(c.as_ref())?;
                IrSpanned {
                    span,
                    node: ExprCompiled::Slice(Box::new((l, a, b, c))),
                }
            }
            ExprCompiled::Seq(a_b) => {
                let (a, b) = &**a_b;
                let a = self.inline(a)?;
                let b = self.inline(b)?;
                ExprCompiled::seq(a, b)
            }
            ExprCompiled::Call(call) => return self.inline_call(call),
            // These should be unreachable, but it is safer
            // to do unnecessary work in compiler than crash.
            ExprCompiled::LocalCaptured(..)
            | ExprCompiled::Module(..)
            | ExprCompiled::Compr(..)
            | ExprCompiled::Def(..) => return Err(CannotInline),
        })
    }
}