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

//! Write operators like `+=`.

use crate::collections::symbol_map::Symbol;
use crate::eval::bc::compiler::expr::write_n_exprs;
use crate::eval::bc::instr_impl::InstrAddAssign;
use crate::eval::bc::instr_impl::InstrArrayIndex;
use crate::eval::bc::instr_impl::InstrArrayIndexSet;
use crate::eval::bc::instr_impl::InstrBitAnd;
use crate::eval::bc::instr_impl::InstrBitOrAssign;
use crate::eval::bc::instr_impl::InstrBitXor;
use crate::eval::bc::instr_impl::InstrDivide;
use crate::eval::bc::instr_impl::InstrFloorDivide;
use crate::eval::bc::instr_impl::InstrLeftShift;
use crate::eval::bc::instr_impl::InstrLoadModule;
use crate::eval::bc::instr_impl::InstrMultiply;
use crate::eval::bc::instr_impl::InstrObjectField;
use crate::eval::bc::instr_impl::InstrPercent;
use crate::eval::bc::instr_impl::InstrRightShift;
use crate::eval::bc::instr_impl::InstrSetObjectField;
use crate::eval::bc::instr_impl::InstrStoreModule;
use crate::eval::bc::instr_impl::InstrSub;
use crate::eval::bc::stack_ptr::BcSlotIn;
use crate::eval::bc::stack_ptr::BcSlotOut;
use crate::eval::bc::stack_ptr::BcSlotsN;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::stmt::AssignModifyLhs;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::syntax::ast::AssignOp;

impl AssignOp {
    fn write_bc(
        self,
        v0: BcSlotIn,
        v1: BcSlotIn,
        target: BcSlotOut,
        span: FrameSpan,
        bc: &mut BcWriter,
    ) {
        let arg = (v0, v1, target);
        match self {
            AssignOp::Add => bc.write_instr::<InstrAddAssign>(span, arg),
            AssignOp::Subtract => bc.write_instr::<InstrSub>(span, arg),
            AssignOp::Multiply => bc.write_instr::<InstrMultiply>(span, arg),
            AssignOp::Divide => bc.write_instr::<InstrDivide>(span, arg),
            AssignOp::FloorDivide => bc.write_instr::<InstrFloorDivide>(span, arg),
            AssignOp::Percent => bc.write_instr::<InstrPercent>(span, arg),
            AssignOp::BitAnd => bc.write_instr::<InstrBitAnd>(span, arg),
            AssignOp::BitOr => bc.write_instr::<InstrBitOrAssign>(span, arg),
            AssignOp::BitXor => bc.write_instr::<InstrBitXor>(span, arg),
            AssignOp::LeftShift => bc.write_instr::<InstrLeftShift>(span, arg),
            AssignOp::RightShift => bc.write_instr::<InstrRightShift>(span, arg),
        }
    }
}

impl AssignModifyLhs {
    /// After evaluation of `x[y] += ...`, variables `x` and `y` are definitely assigned.
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        match self {
            AssignModifyLhs::Dot(object, _field) => object.mark_definitely_assigned_after(bc),
            AssignModifyLhs::Array(array, index) => {
                array.mark_definitely_assigned_after(bc);
                index.mark_definitely_assigned_after(bc);
            }
            AssignModifyLhs::LocalCaptured(_) => {}
            AssignModifyLhs::Local(local) => bc.mark_definitely_assigned(local.node),
            AssignModifyLhs::Module(_) => {}
        }
    }

    pub(crate) fn write_bc(
        &self,
        span: FrameSpan,
        op: AssignOp,
        rhs: &IrSpanned<ExprCompiled>,
        bc: &mut BcWriter,
    ) {
        match *self {
            AssignModifyLhs::Dot(ref object, ref field) => {
                object.write_bc_cb(bc, |object, bc| {
                    bc.alloc_slots_c(|lhs_rhs: BcSlotsN<2>, bc| {
                        let field = Symbol::new(field.as_str());
                        bc.write_instr::<InstrObjectField>(
                            span,
                            (object, field.clone(), lhs_rhs.get::<0>().to_out()),
                        );
                        rhs.write_bc(lhs_rhs.get::<1>().to_out(), bc);
                        op.write_bc(
                            lhs_rhs.get::<0>().to_in(),
                            lhs_rhs.get::<1>().to_in(),
                            lhs_rhs.get::<1>().to_out(),
                            span,
                            bc,
                        );
                        bc.write_instr::<InstrSetObjectField>(
                            span,
                            (lhs_rhs.get::<1>().to_in(), object, field),
                        );
                    })
                });
            }
            AssignModifyLhs::Array(ref array, ref index) => {
                write_n_exprs([array, index], bc, |[array, index], bc| {
                    bc.alloc_slots_c(|slots: BcSlotsN<2>, bc| {
                        let temp_slot = slots.get::<0>();
                        let rhs_slot = slots.get::<1>();

                        bc.write_instr::<InstrArrayIndex>(span, (array, index, temp_slot.to_out()));
                        rhs.write_bc(rhs_slot.to_out(), bc);
                        op.write_bc(
                            temp_slot.to_in(),
                            rhs_slot.to_in(),
                            temp_slot.to_out(),
                            span,
                            bc,
                        );
                        bc.write_instr::<InstrArrayIndexSet>(
                            span,
                            (array, index, temp_slot.to_in()),
                        );
                    })
                });
            }
            AssignModifyLhs::Local(s) => bc.alloc_slots_c(|lhs_rhs: BcSlotsN<2>, bc| {
                let slot = s.node;
                bc.write_load_local(span, slot, lhs_rhs.get::<0>().to_out());
                rhs.write_bc(lhs_rhs.get::<1>().to_out(), bc);

                op.write_bc(
                    lhs_rhs.get::<0>().to_in(),
                    lhs_rhs.get::<1>().to_in(),
                    lhs_rhs.get::<1>().to_out(),
                    span,
                    bc,
                );
                bc.write_mov(span, lhs_rhs.get::<1>().to_in(), slot.to_bc_slot().to_out());
            }),
            AssignModifyLhs::LocalCaptured(s) => bc.alloc_slots_c(|lhs_rhs: BcSlotsN<2>, bc| {
                let slot = s.node;
                bc.write_load_local_captured(span, slot, lhs_rhs.get::<0>().to_out());
                rhs.write_bc(lhs_rhs.get::<1>().to_out(), bc);

                op.write_bc(
                    lhs_rhs.get::<0>().to_in(),
                    lhs_rhs.get::<1>().to_in(),
                    lhs_rhs.get::<1>().to_out(),
                    span,
                    bc,
                );
                bc.write_store_local_captured(span, lhs_rhs.get::<1>().to_in(), slot);
            }),
            AssignModifyLhs::Module(m) => bc.alloc_slots_c(|lhs_rhs: BcSlotsN<2>, bc| {
                let slot = m.node;
                bc.write_instr::<InstrLoadModule>(span, (slot, lhs_rhs.get::<0>().to_out()));
                rhs.write_bc(lhs_rhs.get::<1>().to_out(), bc);
                op.write_bc(
                    lhs_rhs.get::<0>().to_in(),
                    lhs_rhs.get::<1>().to_in(),
                    lhs_rhs.get::<1>().to_out(),
                    span,
                    bc,
                );
                bc.write_instr::<InstrStoreModule>(span, (lhs_rhs.get::<1>().to_in(), slot));
            }),
        }
    }
}


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

//! Compile function calls.

use either::Either;

use crate::collections::symbol_map::Symbol;
use crate::eval::bc::call::BcCallArgsFull;
use crate::eval::bc::call::BcCallArgsPos;
use crate::eval::bc::compiler::expr::write_expr_opt;
use crate::eval::bc::compiler::expr::write_exprs;
use crate::eval::bc::instr_impl::InstrCall;
use crate::eval::bc::instr_impl::InstrCallFrozen;
use crate::eval::bc::instr_impl::InstrCallFrozenDef;
use crate::eval::bc::instr_impl::InstrCallFrozenDefPos;
use crate::eval::bc::instr_impl::InstrCallFrozenNative;
use crate::eval::bc::instr_impl::InstrCallFrozenNativePos;
use crate::eval::bc::instr_impl::InstrCallFrozenPos;
use crate::eval::bc::instr_impl::InstrCallMaybeKnownMethod;
use crate::eval::bc::instr_impl::InstrCallMaybeKnownMethodPos;
use crate::eval::bc::instr_impl::InstrCallMethod;
use crate::eval::bc::instr_impl::InstrCallMethodPos;
use crate::eval::bc::instr_impl::InstrCallPos;
use crate::eval::bc::instr_impl::InstrLen;
use crate::eval::bc::instr_impl::InstrType;
use crate::eval::bc::native_function::BcNativeFunction;
use crate::eval::bc::stack_ptr::BcSlotOut;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::args::ArgsCompiledValue;
use crate::eval::compiler::call::CallCompiled;
use crate::eval::compiler::def::FrozenDef;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::values::function::NativeFunction;
use crate::values::types::known_methods::get_known_method;
use crate::values::FrozenValue;
use crate::values::FrozenValueTyped;

impl ArgsCompiledValue {
    /// After evaluation of function arguments like `foo(a, b=c[d], **e)`,
    /// variables `a`, `b`, `c`, `d`, and `e` are definitely assigned.
    fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        let ArgsCompiledValue {
            pos_named,
            names,
            args,
            kwargs,
        } = self;
        for n in pos_named {
            n.mark_definitely_assigned_after(bc);
        }
        let _ = names;
        if let Some(args) = args {
            args.mark_definitely_assigned_after(bc);
        }
        if let Some(kwargs) = kwargs {
            kwargs.mark_definitely_assigned_after(bc);
        }
    }

    fn write_bc(&self, bc: &mut BcWriter, k: impl FnOnce(BcCallArgsFull<Symbol>, &mut BcWriter)) {
        write_exprs(&self.pos_named, bc, |pos_named, bc| {
            write_expr_opt(&self.args, bc, |args, bc| {
                write_expr_opt(&self.kwargs, bc, |kwargs, bc| {
                    let args_full = BcCallArgsFull {
                        pos_named,
                        names: self.names.clone().into_boxed_slice(),
                        args,
                        kwargs,
                    };
                    k(args_full, bc)
                })
            })
        });
    }
}

impl CallCompiled {
    /// After evaluation of call like `a[b](c.d)`,
    /// variables `a`, `b`, and `c` are definitely assigned.
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        let CallCompiled { fun, args } = self;
        fun.mark_definitely_assigned_after(bc);
        args.mark_definitely_assigned_after(bc);
    }
}

impl IrSpanned<CallCompiled> {
    fn write_args(
        args: &ArgsCompiledValue,
        bc: &mut BcWriter,
        k: impl FnOnce(Either<BcCallArgsPos, BcCallArgsFull<Symbol>>, &mut BcWriter),
    ) {
        if let Some(pos) = args.pos_only() {
            write_exprs(pos, bc, |pos, bc| {
                let args = Either::Left(BcCallArgsPos { pos });
                k(args, bc)
            })
        } else {
            args.write_bc(bc, |args_full, bc| {
                let args = Either::Right(args_full);
                k(args, bc)
            })
        }
    }

    fn write_call_frozen(
        span: FrameSpan,
        fun: FrozenValue,
        args: &ArgsCompiledValue,
        target: BcSlotOut,
        bc: &mut BcWriter,
    ) {
        let file_span = bc.alloc_file_span(span);
        if let Some(fun) = FrozenValueTyped::<FrozenDef>::new(fun) {
            Self::write_args(args, bc, |args, bc| match args {
                Either::Left(npops) => {
                    bc.write_instr::<InstrCallFrozenDefPos>(span, (fun, npops, file_span, target))
                }
                Either::Right(args) => bc.write_instr::<InstrCallFrozenDef>(
                    span,
                    (fun, args.resolve(fun.as_ref()), file_span, target),
                ),
            })
        } else if let Some(fun) = FrozenValueTyped::<NativeFunction>::new(fun) {
            let fun = BcNativeFunction::new(fun);
            Self::write_args(args, bc, |args, bc| match args {
                Either::Left(npops) => {
                    bc.write_instr::<InstrCallFrozenNativePos>(
                        span,
                        (fun, npops, file_span, target),
                    );
                }
                Either::Right(args) => {
                    bc.write_instr::<InstrCallFrozenNative>(span, (fun, args, file_span, target));
                }
            })
        } else {
            Self::write_args(args, bc, |args, bc| match args {
                Either::Left(npops) => {
                    bc.write_instr::<InstrCallFrozenPos>(span, (fun, npops, file_span, target));
                }
                Either::Right(args) => {
                    bc.write_instr::<InstrCallFrozen>(span, (fun, args, file_span, target));
                }
            })
        }
    }

    fn write_call_method(
        target: BcSlotOut,
        span: FrameSpan,
        this: &IrSpanned<ExprCompiled>,
        symbol: &Symbol,
        args: &ArgsCompiledValue,
        bc: &mut BcWriter,
    ) {
        this.write_bc_cb(bc, |this, bc| {
            let file_span = bc.alloc_file_span(span);
            let known_method = get_known_method(symbol.as_str());
            if let Some(pos) = args.pos_only() {
                write_exprs(pos, bc, |pos, bc| {
                    if let Some(known_method) = known_method {
                        bc.write_instr::<InstrCallMaybeKnownMethodPos>(
                            span,
                            (
                                this,
                                symbol.clone(),
                                known_method,
                                BcCallArgsPos { pos },
                                file_span,
                                target,
                            ),
                        );
                    } else {
                        bc.write_instr::<InstrCallMethodPos>(
                            span,
                            (
                                this,
                                symbol.clone(),
                                BcCallArgsPos { pos },
                                file_span,
                                target,
                            ),
                        );
                    }
                });
            } else {
                args.write_bc(bc, |args, bc| {
                    if let Some(known_method) = known_method {
                        bc.write_instr::<InstrCallMaybeKnownMethod>(
                            span,
                            (this, symbol.clone(), known_method, args, file_span, target),
                        );
                    } else {
                        bc.write_instr::<InstrCallMethod>(
                            span,
                            (this, symbol.clone(), args, file_span, target),
                        );
                    }
                })
            }
        })
    }

    pub(crate) fn write_bc(&self, target: BcSlotOut, bc: &mut BcWriter) {
        if let Some(arg) = self.as_len() {
            return arg.write_bc_cb(bc, |arg, bc| {
                bc.write_instr::<InstrLen>(self.span, (arg, target));
            });
        } else if let Some(arg) = self.as_type() {
            return arg.write_bc_cb(bc, |arg, bc| {
                bc.write_instr::<InstrType>(self.span, (arg, target));
            });
        }

        let span = self.span;
        let file_span = bc.alloc_file_span(span);
        match self.method() {
            None => match self.fun.as_value() {
                Some(f) => Self::write_call_frozen(span, f, &self.args, target, bc),
                None => {
                    self.fun.write_bc_cb(bc, |fun, bc| {
                        Self::write_args(&self.args, bc, |args, bc| match args {
                            Either::Left(npops) => bc
                                .write_instr::<InstrCallPos>(span, (fun, npops, file_span, target)),
                            Either::Right(args) => {
                                bc.write_instr::<InstrCall>(span, (fun, args, file_span, target));
                            }
                        })
                    });
                }
            },
            Some((this, symbol, args)) => {
                Self::write_call_method(target, span, this, symbol, args, bc)
            }
        }
    }
}


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

//! Compile comprehensions.

use crate::eval::bc::compiler::expr::write_n_exprs;
use crate::eval::bc::compiler::if_compiler::write_if_then;
use crate::eval::bc::compiler::stmt::write_for;
use crate::eval::bc::instr_impl::InstrComprDictInsert;
use crate::eval::bc::instr_impl::InstrComprListAppend;
use crate::eval::bc::instr_impl::InstrDictNew;
use crate::eval::bc::instr_impl::InstrListNew;
use crate::eval::bc::stack_ptr::BcSlotOut;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::compr::ClauseCompiled;
use crate::eval::compiler::compr::ComprCompiled;
use crate::eval::compiler::expr::MaybeNot;
use crate::eval::runtime::frame_span::FrameSpan;

impl ClauseCompiled {
    fn write_bc(
        &self,
        bc: &mut BcWriter,
        rem: &[ClauseCompiled],
        term: impl FnOnce(&mut BcWriter),
    ) {
        write_for(&self.over, &self.var, self.over.span, bc, |bc| {
            for c in &self.ifs {
                write_if_then(c, MaybeNot::Not, |bc| bc.write_continue(c.span), bc);
            }

            match rem.split_last() {
                Some((first, rem)) => {
                    first.write_bc(bc, rem, term);
                }
                None => {
                    term(bc);
                }
            }
        })
    }
}

impl ComprCompiled {
    /// After evaluation of comprehension like `[(x, z) for x in y for z in w]`,
    /// we can mark `y` as definitely assigned.
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        let clauses = self.clauses();
        // We know that first loop argument is executed, and we don't know anything else.
        clauses
            .split_last()
            .0
            .over
            .mark_definitely_assigned_after(bc);
    }

    pub(crate) fn write_bc(&self, span: FrameSpan, target: BcSlotOut, bc: &mut BcWriter) {
        bc.alloc_slot(|temp, bc| {
            match self {
                ComprCompiled::List(ref expr, ref clauses) => {
                    bc.write_instr::<InstrListNew>(span, temp.to_out());
                    let (first, rem) = clauses.split_last();
                    first.write_bc(bc, rem, |bc| {
                        expr.write_bc_cb(bc, |expr_slot, bc| {
                            bc.write_instr::<InstrComprListAppend>(
                                expr.span,
                                (temp.to_in(), expr_slot),
                            )
                        });
                    });
                }
                ComprCompiled::Dict(k_v, clauses) => {
                    let (k, v) = &**k_v;
                    bc.write_instr::<InstrDictNew>(span, temp.to_out());
                    let (first, rem) = clauses.split_last();
                    first.write_bc(bc, rem, |bc| {
                        write_n_exprs([k, v], bc, |[k_slot, v_slot], bc| {
                            bc.write_instr::<InstrComprDictInsert>(
                                k.span,
                                (temp.to_in(), k_slot, v_slot),
                            );
                        });
                    });
                }
            };
            bc.write_mov(span, temp.to_in(), target);
        });
    }
}


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

//! Compile def.

use crate::eval::bc::instr_impl::InstrDef;
use crate::eval::bc::instr_impl::InstrDefData;
use crate::eval::bc::stack_ptr::BcSlotOut;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::def::DefCompiled;
use crate::eval::compiler::def::ParametersCompiled;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::slice_vec_ext::SliceExt;

impl DefCompiled {
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        // TODO(nga): argument default values and types can be used
        //   to mark variables definitely assigned.
        let _ = bc;
    }

    pub(crate) fn write_bc(&self, span: FrameSpan, target: BcSlotOut, bc: &mut BcWriter) {
        let DefCompiled {
            ref function_name,
            ref params,
            return_type,
            info,
        } = *self;
        let function_name = function_name.clone();

        let num_positional = params.num_positional;

        let how_many_slots_we_need = params.count_exprs();

        bc.alloc_slots(how_many_slots_we_need, |slots, bc| {
            let mut slots_i = slots.iter();
            let mut value_count = 0;
            let params = params.params.map(|p| {
                p.map(|p| {
                    p.map_expr(|e| {
                        e.write_bc(slots_i.next().unwrap().to_out(), bc);
                        value_count += 1;
                        value_count - 1
                    })
                })
            });

            let params = ParametersCompiled {
                params,
                num_positional,
            };
            let instr_def_data = InstrDefData {
                function_name,
                params,
                return_type,
                info,
            };

            assert!(slots_i.next().is_none());

            bc.write_instr::<InstrDef>(span, (slots.to_in(), instr_def_data, target));
        })
    }
}


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

//! Compile expressions.

use std::collections::HashSet;

use crate::collections::Hashed;
use crate::collections::SmallMap;
use crate::eval::bc::compiler::if_compiler::write_if_else;
use crate::eval::bc::instr_impl::*;
use crate::eval::bc::slow_arg::BcInstrSlowArg;
use crate::eval::bc::stack_ptr::BcSlot;
use crate::eval::bc::stack_ptr::BcSlotIn;
use crate::eval::bc::stack_ptr::BcSlotInRange;
use crate::eval::bc::stack_ptr::BcSlotOut;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::expr::Builtin1;
use crate::eval::compiler::expr::Builtin2;
use crate::eval::compiler::expr::CompareOp;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::expr::ExprLogicalBinOp;
use crate::eval::compiler::expr::MaybeNot;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::slice_vec_ext::SliceExt;
use crate::values::layout::value_not_special::FrozenValueNotSpecial;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;
use crate::values::ValueLike;

/// Try extract consecutive definitely initialized locals from expressions.
fn try_slot_range<'a>(
    exprs: impl IntoIterator<Item = &'a IrSpanned<ExprCompiled>>,
    bc: &BcWriter,
) -> Option<BcSlotInRange> {
    let mut range = BcSlotInRange::default();
    for expr in exprs {
        let local = expr.as_local_non_captured()?;
        let slot = bc.try_definitely_assigned(local)?;
        if !range.try_push(slot) {
            return None;
        }
    }
    Some(range)
}

/// Compile several expressions into consecutive registers.
pub(crate) fn write_exprs<'a>(
    exprs: impl IntoIterator<Item = &'a IrSpanned<ExprCompiled>>,
    bc: &mut BcWriter,
    k: impl FnOnce(BcSlotInRange, &mut BcWriter),
) {
    let exprs: Vec<_> = exprs.into_iter().collect();

    if let Some(slots) = try_slot_range(exprs.iter().copied(), bc) {
        k(slots, bc);
    } else {
        bc.alloc_slots_for_exprs(exprs, |slot, expr, bc| expr.write_bc(slot.to_out(), bc), k)
    }
}

pub(crate) fn write_expr_opt(
    expr: &Option<IrSpanned<ExprCompiled>>,
    bc: &mut BcWriter,
    k: impl FnOnce(Option<BcSlotIn>, &mut BcWriter),
) {
    if let Some(expr) = expr {
        expr.write_bc_cb(bc, |slot, bc| k(Some(slot), bc))
    } else {
        k(None, bc)
    }
}

pub(crate) fn write_n_exprs<const N: usize>(
    exprs: [&IrSpanned<ExprCompiled>; N],
    bc: &mut BcWriter,
    k: impl FnOnce([BcSlotIn; N], &mut BcWriter),
) {
    fn help<const N: usize>(
        mut filled: [BcSlotIn; N],
        rem_exprs: &[&IrSpanned<ExprCompiled>],
        bc: &mut BcWriter,
        k: impl FnOnce([BcSlotIn; N], &mut BcWriter),
    ) {
        match rem_exprs.split_first() {
            Some((first, rem)) => first.write_bc_cb(bc, |first, bc| {
                filled[N - rem.len() - 1] = first;
                help(filled, rem, bc, k)
            }),
            None => k(filled, bc),
        }
    }

    help([BcSlot(98765).to_in(); N], &exprs, bc, k)
}

impl ExprCompiled {
    /// Mark variables which are definitely assigned after execution of this expression.
    ///
    /// For example, when this expression if executed:
    ///
    /// ```ignore
    /// t if c else f
    /// ```
    ///
    /// `c` is definitely assigned (because if it is not, then execution fails),
    /// but we don't know about `t` or `f` because one of them was not executed.
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        match self {
            ExprCompiled::Value(_) => {}
            ExprCompiled::Local(local) => bc.mark_definitely_assigned(*local),
            ExprCompiled::LocalCaptured(_) => {}
            ExprCompiled::Module(_) => {}
            ExprCompiled::Tuple(xs) | ExprCompiled::List(xs) => {
                for x in xs {
                    x.mark_definitely_assigned_after(bc);
                }
            }
            ExprCompiled::Dict(xs) => {
                for (k, v) in xs {
                    k.mark_definitely_assigned_after(bc);
                    v.mark_definitely_assigned_after(bc);
                }
            }
            ExprCompiled::Compr(compr) => compr.mark_definitely_assigned_after(bc),
            ExprCompiled::If(c_t_f) => {
                let (c, _t, _f) = &**c_t_f;
                // Condition is executed unconditionally, so we use it to mark definitely assigned.
                // But we don't know which of the branches will be executed.
                c.mark_definitely_assigned_after(bc);
            }
            ExprCompiled::Slice(l_a_b_c) => {
                let (l, a, b, c) = &**l_a_b_c;
                l.mark_definitely_assigned_after(bc);
                if let Some(a) = a {
                    a.mark_definitely_assigned_after(bc);
                }
                if let Some(b) = b {
                    b.mark_definitely_assigned_after(bc);
                }
                if let Some(c) = c {
                    c.mark_definitely_assigned_after(bc);
                }
            }
            ExprCompiled::Builtin1(_op, expr) => {
                expr.mark_definitely_assigned_after(bc);
            }
            ExprCompiled::LogicalBinOp(_op, a_b) => {
                let (a, b) = &**a_b;
                // `a` is executed unconditionally, but `b` is not,
                // so we mark only `a` as definitely assigned.
                a.mark_definitely_assigned_after(bc);
                let _ = b;
            }
            ExprCompiled::Seq(a_b) => {
                let (a, b) = &**a_b;
                a.mark_definitely_assigned_after(bc);
                b.mark_definitely_assigned_after(bc);
            }
            ExprCompiled::Builtin2(_op, a_b) => {
                let (a, b) = &**a_b;
                a.mark_definitely_assigned_after(bc);
                b.mark_definitely_assigned_after(bc);
            }
            ExprCompiled::Index2(a_i0_i1) => {
                let (a, i0, i1) = &**a_i0_i1;
                a.mark_definitely_assigned_after(bc);
                i0.mark_definitely_assigned_after(bc);
                i1.mark_definitely_assigned_after(bc);
            }
            ExprCompiled::Call(c) => c.mark_definitely_assigned_after(bc),
            ExprCompiled::Def(d) => d.mark_definitely_assigned_after(bc),
        }
    }
}

impl IrSpanned<ExprCompiled> {
    fn try_dict_of_consts(
        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
    ) -> Option<SmallMap<FrozenValue, FrozenValue>> {
        let mut res = SmallMap::new();
        for (k, v) in xs {
            let k = k.as_value()?.get_hashed().ok()?;
            let v = v.as_value()?;
            let prev = res.insert_hashed(k, v);
            if prev.is_some() {
                // If there are duplicates, so don't take the fast-literal
                // path and go down the slow runtime path (which will raise the error).
                // We have a lint that will likely fire on this issue (and others).
                return None;
            }
        }
        Some(res)
    }

    fn try_dict_const_keys(
        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
    ) -> Option<Box<[Hashed<FrozenValue>]>> {
        let mut keys = Vec::new();
        let mut keys_unique = HashSet::new();
        for (k, _) in xs {
            let k = k.as_value()?.get_hashed().ok()?;
            keys.push(k);
            let inserted = keys_unique.insert(k);
            if !inserted {
                // Otherwise fail at runtime
                return None;
            }
        }
        Some(keys.into_boxed_slice())
    }

    fn write_dict(
        span: FrameSpan,
        xs: &[(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)],
        target: BcSlotOut,
        bc: &mut BcWriter,
    ) {
        if xs.is_empty() {
            bc.write_instr::<InstrDictNew>(span, target);
        } else if let Some(d) = Self::try_dict_of_consts(xs) {
            bc.write_instr::<InstrDictOfConsts>(span, (d, target));
        } else if let Some(keys) = Self::try_dict_const_keys(xs) {
            assert_eq!(keys.len(), xs.len());
            write_exprs(xs.iter().map(|(_, v)| v), bc, |values, bc| {
                assert_eq!(values.len() as usize, keys.len());
                bc.write_instr::<InstrDictConstKeys>(span, (keys, values.to_range_from(), target));
            });
        } else {
            let key_spans = xs.map(|(k, _v)| k.span);
            write_exprs(xs.iter().flat_map(|(k, v)| [k, v]), bc, |kvs, bc| {
                bc.write_instr_explicit::<InstrDictNPop>(
                    BcInstrSlowArg {
                        span,
                        spans: key_spans,
                    },
                    (kvs, target),
                );
            });
        }
    }

    fn write_not(expr: &IrSpanned<ExprCompiled>, target: BcSlotOut, bc: &mut BcWriter) {
        expr.write_bc_cb(bc, |slot, bc| {
            bc.write_instr::<InstrNot>(expr.span, (slot, target));
        });
    }

    fn write_equals_const(
        span: FrameSpan,
        a: &IrSpanned<ExprCompiled>,
        b: FrozenValue,
        target: BcSlotOut,
        bc: &mut BcWriter,
    ) {
        a.write_bc_cb(bc, |a, bc| {
            if let Some(b) = b.to_value().unpack_int_value() {
                bc.write_instr::<InstrEqInt>(span, (a, b, target));
            } else if b.eq_is_ptr_eq() {
                bc.write_instr::<InstrEqPtr>(span, (a, b, target));
            } else if let Some(b) = FrozenStringValue::new(b) {
                bc.write_instr::<InstrEqStr>(span, (a, b, target));
            } else if let Some(b) = FrozenValueNotSpecial::new(b) {
                bc.write_instr::<InstrEqConst>(span, (a, b, target));
            } else {
                unreachable!("FrozenValue must be either i32, str or not-special");
            }
        });
    }

    fn write_equals(
        span: FrameSpan,
        a: &IrSpanned<ExprCompiled>,
        b: &IrSpanned<ExprCompiled>,
        target: BcSlotOut,
        bc: &mut BcWriter,
    ) {
        if let Some(a) = a.as_value() {
            Self::write_equals_const(span, b, a, target, bc);
        } else if let Some(b) = b.as_value() {
            Self::write_equals_const(span, a, b, target, bc);
        } else {
            write_n_exprs([a, b], bc, |[a, b], bc| {
                bc.write_instr::<InstrEq>(span, (a, b, target));
            });
        }
    }

    pub(crate) fn write_bc(&self, target: BcSlotOut, bc: &mut BcWriter) {
        let span = self.span;
        match &self.node {
            ExprCompiled::Value(v) => {
                bc.write_const(span, *v, target);
            }
            ExprCompiled::Local(slot) => {
                bc.write_load_local(span, *slot, target);
            }
            ExprCompiled::LocalCaptured(slot) => {
                bc.write_load_local_captured(span, *slot, target);
            }
            ExprCompiled::Module(slot) => {
                bc.write_instr::<InstrLoadModule>(span, (*slot, target));
            }
            ExprCompiled::Tuple(xs) => {
                write_exprs(xs, bc, |xs, bc| {
                    bc.write_instr::<InstrTupleNPop>(span, (xs, target));
                });
            }
            ExprCompiled::List(ref xs) => {
                if xs.is_empty() {
                    bc.write_instr::<InstrListNew>(span, target);
                } else if xs.iter().all(|x| x.as_value().is_some()) {
                    let content = xs.map(|v| v.as_value().unwrap()).into_boxed_slice();
                    bc.write_instr::<InstrListOfConsts>(span, (content, target));
                } else {
                    write_exprs(xs, bc, |xs, bc| {
                        bc.write_instr::<InstrListNPop>(span, (xs, target));
                    });
                }
            }
            ExprCompiled::Dict(ref xs) => Self::write_dict(span, xs, target, bc),
            ExprCompiled::Compr(ref compr) => compr.write_bc(span, target, bc),
            ExprCompiled::Slice(l_start_stop_step) => {
                let (l, start, stop, step) = &**l_start_stop_step;
                l.write_bc_cb(bc, |l, bc| {
                    write_expr_opt(start, bc, |start, bc| {
                        write_expr_opt(stop, bc, |stop, bc| {
                            write_expr_opt(step, bc, |step, bc| {
                                bc.write_instr::<InstrSlice>(span, (l, start, stop, step, target))
                            })
                        })
                    })
                });
            }
            ExprCompiled::Builtin1(Builtin1::Not, expr) => Self::write_not(expr, target, bc),
            ExprCompiled::Builtin1(op, expr) => {
                expr.write_bc_cb(bc, |expr, bc| {
                    let arg = (expr, target);
                    match op {
                        Builtin1::Not => unreachable!("handled above"),
                        Builtin1::Minus => bc.write_instr::<InstrMinus>(span, arg),
                        Builtin1::Plus => bc.write_instr::<InstrPlus>(span, arg),
                        Builtin1::BitNot => bc.write_instr::<InstrBitNot>(span, arg),
                        Builtin1::TypeIs(t) => {
                            bc.write_instr::<InstrTypeIs>(span, (expr, *t, target))
                        }
                        Builtin1::PercentSOne(before, after) => bc
                            .write_instr::<InstrPercentSOne>(span, (*before, expr, *after, target)),
                        Builtin1::FormatOne(before, after) => {
                            bc.write_instr::<InstrFormatOne>(span, (*before, expr, *after, target))
                        }
                        Builtin1::Dot(field) => {
                            bc.write_instr::<InstrObjectField>(span, (expr, field.clone(), target))
                        }
                    }
                });
            }
            ExprCompiled::If(cond_t_f) => {
                let (cond, t, f) = &**cond_t_f;
                write_if_else(
                    cond,
                    |bc| t.write_bc(target, bc),
                    |bc| f.write_bc(target, bc),
                    bc,
                );
            }
            ExprCompiled::LogicalBinOp(op, l_r) => {
                let (l, r) = &**l_r;
                l.write_bc_cb(bc, |l_slot, bc| {
                    let maybe_not = match op {
                        ExprLogicalBinOp::And => MaybeNot::Id,
                        ExprLogicalBinOp::Or => MaybeNot::Not,
                    };
                    bc.write_if_else(
                        l_slot,
                        maybe_not,
                        l.span,
                        |bc| r.write_bc(target, bc),
                        |bc| bc.write_mov(span, l_slot, target),
                    );
                });
            }
            ExprCompiled::Seq(l_r) => {
                let (l, r) = &**l_r;
                l.write_bc_for_effect(bc);
                r.write_bc(target, bc);
            }
            ExprCompiled::Builtin2(Builtin2::Equals, l_r) => {
                let (l, r) = &**l_r;
                Self::write_equals(span, l, r, target, bc)
            }
            ExprCompiled::Builtin2(op, l_r) => {
                let (l, r) = &**l_r;
                write_n_exprs([l, r], bc, |[l, r], bc| {
                    let arg = (l, r, target);
                    match op {
                        Builtin2::Equals => unreachable!("handled above"),
                        Builtin2::Compare(CompareOp::Less) => {
                            bc.write_instr::<InstrLess>(span, arg)
                        }
                        Builtin2::Compare(CompareOp::Greater) => {
                            bc.write_instr::<InstrGreater>(span, arg)
                        }
                        Builtin2::Compare(CompareOp::LessOrEqual) => {
                            bc.write_instr::<InstrLessOrEqual>(span, arg)
                        }
                        Builtin2::Compare(CompareOp::GreaterOrEqual) => {
                            bc.write_instr::<InstrGreaterOrEqual>(span, arg)
                        }
                        Builtin2::In => bc.write_instr::<InstrIn>(span, arg),
                        Builtin2::Sub => bc.write_instr::<InstrSub>(span, arg),
                        Builtin2::Add => bc.write_instr::<InstrAdd>(span, arg),
                        Builtin2::Multiply => bc.write_instr::<InstrMultiply>(span, arg),
                        Builtin2::Divide => bc.write_instr::<InstrDivide>(span, arg),
                        Builtin2::FloorDivide => bc.write_instr::<InstrFloorDivide>(span, arg),
                        Builtin2::Percent => bc.write_instr::<InstrPercent>(span, arg),
                        Builtin2::BitAnd => bc.write_instr::<InstrBitAnd>(span, arg),
                        Builtin2::BitOr => bc.write_instr::<InstrBitOr>(span, arg),
                        Builtin2::BitXor => bc.write_instr::<InstrBitXor>(span, arg),
                        Builtin2::LeftShift => bc.write_instr::<InstrLeftShift>(span, arg),
                        Builtin2::RightShift => bc.write_instr::<InstrRightShift>(span, arg),
                        Builtin2::ArrayIndex => bc.write_instr::<InstrArrayIndex>(span, arg),
                    }
                });
            }
            ExprCompiled::Index2(a_i0_i1) => {
                let (a, i0, i1) = &**a_i0_i1;
                write_n_exprs([a, i0, i1], bc, |[a, i0, i1], bc| {
                    bc.write_instr::<InstrArrayIndex2>(span, (a, i0, i1, target))
                });
            }
            ExprCompiled::Call(ref call) => call.write_bc(target, bc),
            ExprCompiled::Def(ref def) => def.write_bc(span, target, bc),
        }
    }

    /// Allocate temporary slot, write expression into it,
    /// and then consume the slot with the callback.
    pub(crate) fn write_bc_cb<R>(
        &self,
        bc: &mut BcWriter,
        k: impl FnOnce(BcSlotIn, &mut BcWriter) -> R,
    ) -> R {
        if let Some(local) = self.as_local_non_captured() {
            // Local is known to be definitely assigned, so there's no need
            // to "load" it just to trigger check that it is assigned.
            if let Some(slot) = bc.try_definitely_assigned(local) {
                return k(slot, bc);
            }
        }

        bc.alloc_slot(|slot, bc| {
            self.write_bc(slot.to_out(), bc);
            k(slot.to_in(), bc)
        })
    }

    pub(crate) fn write_bc_for_effect(&self, bc: &mut BcWriter) {
        self.write_bc_cb(bc, |slot, _bc| {
            let _ = slot;
        });
    }
}


/*
 * 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::eval::bc::instrs::PatchAddr;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::expr::Builtin1;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::expr::ExprLogicalBinOp;
use crate::eval::compiler::expr::MaybeNot;
use crate::eval::compiler::span::IrSpanned;

/// Common code for compiling if statements and if expressions.
pub(crate) fn write_if_else(
    c: &IrSpanned<ExprCompiled>,
    t: impl FnOnce(&mut BcWriter),
    f: impl FnOnce(&mut BcWriter),
    bc: &mut BcWriter,
) {
    write_if_else_impl(c, MaybeNot::Id, t, Some(f), bc);
}

/// Common code for compiling if statements and if conditions in comprehensions.
pub(crate) fn write_if_then(
    c: &IrSpanned<ExprCompiled>,
    maybe_not: MaybeNot,
    t: impl FnOnce(&mut BcWriter),
    bc: &mut BcWriter,
) {
    // Deal with some typechecker issue.
    fn wr<T, F>(c: &IrSpanned<ExprCompiled>, maybe_not: MaybeNot, t: T, _f: F, bc: &mut BcWriter)
    where
        T: FnOnce(&mut BcWriter),
        F: FnOnce(&mut BcWriter),
    {
        write_if_else_impl::<T, F>(c, maybe_not, t, None, bc);
    }

    wr(c, maybe_not, t, |_| unreachable!(), bc);
}

/// Common code for writing if-then or if-then-else expression or statement.
fn write_if_else_impl<T, F>(
    cond: &IrSpanned<ExprCompiled>,
    maybe_not: MaybeNot,
    t: T,
    f: Option<F>,
    bc: &mut BcWriter,
) where
    T: FnOnce(&mut BcWriter),
    F: FnOnce(&mut BcWriter),
{
    let mut then_addrs = Vec::new();
    let mut else_addrs = Vec::new();

    write_cond(cond, maybe_not, &mut then_addrs, &mut else_addrs, bc);

    let definitely_assigned = bc.save_definitely_assigned();

    bc.patch_addrs(then_addrs);
    t(bc);
    if let Some(f) = f {
        let end_addr = bc.write_br(cond.span);

        bc.restore_definitely_assigned(definitely_assigned.clone());

        bc.patch_addrs(else_addrs);
        f(bc);

        bc.patch_addr(end_addr);
    } else {
        bc.patch_addrs(else_addrs);
    }

    bc.restore_definitely_assigned(definitely_assigned);
}

/// Write boolean binary condition.
///
/// The condition is: `maybe_not(x bin_op y)`.
///
/// See `write_cond` for semantics of `t`, `f` parameters.
fn write_cond_bin_op(
    x: &IrSpanned<ExprCompiled>,
    y: &IrSpanned<ExprCompiled>,
    bin_op: ExprLogicalBinOp,
    maybe_not: MaybeNot,
    t: &mut Vec<PatchAddr>,
    f: &mut Vec<PatchAddr>,
    bc: &mut BcWriter,
) {
    if (bin_op == ExprLogicalBinOp::And) == (maybe_not == MaybeNot::Id) {
        // This branch handles either of expressions:
        // expression   | bin_op | maybe_not
        // --------------+-------+----------
        // x and y      | and    | id
        // not (x or y) | or     | not

        let mut x_skip = Vec::new();
        write_cond(x, maybe_not, &mut x_skip, f, bc);
        bc.patch_addrs(x_skip);

        write_cond(y, maybe_not, t, f, bc);
    } else {
        // This branch handles either of expressions:
        // expression    | bin_op | maybe_not
        // --------------+-----+--+----------
        // x or y        | or     | id
        // not (x and y) | and    | not

        let mut x_skip = Vec::new();
        write_cond(x, maybe_not.negate(), &mut x_skip, t, bc);
        bc.patch_addrs(x_skip);

        write_cond(y, maybe_not, t, f, bc);
    }
}

/// Write if condition bytecode.
///
/// The condition is `maybe_not(cond)`.
///
/// This function assumes there are two address:
/// * address of then block
/// * address of else block
///
/// Generated code will:
/// * jump to else address if condition is false
/// * jump to then address **or** fall through if condition is true
///
/// This function will populate `t` and `f` with addresses of instructions
/// which jump to then or else block respectively. Caller needs to patch these.
fn write_cond(
    cond: &IrSpanned<ExprCompiled>,
    maybe_not: MaybeNot,
    t: &mut Vec<PatchAddr>,
    f: &mut Vec<PatchAddr>,
    bc: &mut BcWriter,
) {
    match &cond.node {
        ExprCompiled::Builtin1(Builtin1::Not, cond) => {
            write_cond(cond, maybe_not.negate(), t, f, bc);
        }
        ExprCompiled::LogicalBinOp(op, x_y) => {
            let (x, y) = &**x_y;
            write_cond_bin_op(x, y, *op, maybe_not, t, f, bc);
        }
        _ => {
            cond.write_bc_cb(bc, |cond_slot, bc| {
                let addr = match maybe_not {
                    MaybeNot::Id => bc.write_if_not_br(cond_slot, cond.span),
                    MaybeNot::Not => bc.write_if_br(cond_slot, cond.span),
                };
                f.push(addr);
            });
        }
    }
}


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

//! Compile module or function to bytecode.

pub(crate) mod assign;
pub(crate) mod assign_modify;
pub(crate) mod call;
pub(crate) mod compr;
pub(crate) mod def;
pub(crate) mod expr;
pub(crate) mod if_compiler;
pub(crate) mod stmt;


/*
 * 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::eval::bc::bytecode::Bc;
use crate::eval::bc::compiler::if_compiler::write_if_else;
use crate::eval::bc::compiler::if_compiler::write_if_then;
use crate::eval::bc::instr_impl::InstrCheckType;
use crate::eval::bc::instr_impl::InstrPossibleGc;
use crate::eval::bc::instr_impl::InstrReturn;
use crate::eval::bc::instr_impl::InstrReturnCheckType;
use crate::eval::bc::instr_impl::InstrReturnConst;
use crate::eval::bc::stack_ptr::BcSlotIn;
use crate::eval::bc::writer::BcWriter;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::expr::MaybeNot;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::stmt::AssignCompiledValue;
use crate::eval::compiler::stmt::StmtCompileContext;
use crate::eval::compiler::stmt::StmtCompiled;
use crate::eval::compiler::stmt::StmtsCompiled;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::values::typing::TypeCompiled;
use crate::values::FrozenHeap;
use crate::values::FrozenRef;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;

pub(crate) fn write_for(
    over: &IrSpanned<ExprCompiled>,
    var: &IrSpanned<AssignCompiledValue>,
    span: FrameSpan,
    bc: &mut BcWriter,
    body: impl FnOnce(&mut BcWriter),
) {
    let definitely_assigned = bc.save_definitely_assigned();

    over.write_bc_cb(bc, |over, bc| {
        if let Some(var) = var.as_local_non_captured() {
            // Typical case: `for x in ...: ...`,
            // compile loop assignment directly to a local variable.
            bc.write_for(over, var.to_bc_slot().to_out(), span, |bc| {
                bc.mark_definitely_assigned(var);
                body(bc);
            })
        } else {
            // General case, e. g. `for (x, y[0]) in ...: ...`,
            // compile loop assignment to a temporary variable,
            // and reassign it in the loop body.
            bc.alloc_slot(|var_slot, bc| {
                bc.write_for(over, var_slot.to_out(), span, |bc| {
                    var.write_bc(var_slot.to_in(), bc);
                    var.mark_definitely_assigned_after(bc);
                    body(bc);
                })
            })
        }
    });

    bc.restore_definitely_assigned(definitely_assigned);
}

impl StmtsCompiled {
    pub(crate) fn write_bc(&self, compiler: &StmtCompileContext, bc: &mut BcWriter) {
        for stmt in self.stmts() {
            stmt.write_bc(compiler, bc);
        }
    }
}

impl StmtCompiled {
    /// Mark local variables are definitely assigned after this statement executed.
    pub(crate) fn mark_definitely_assigned_after(&self, bc: &mut BcWriter) {
        match self {
            StmtCompiled::PossibleGc => {}
            StmtCompiled::Return(e) => {
                // `e` is definitely assigned after `return` statement,
                // but no code is executed after `return`, so marking would be useless.
                let _ = e;
            }
            StmtCompiled::Expr(e) => e.mark_definitely_assigned_after(bc),
            StmtCompiled::Assign(lhs, ty, rhs) => {
                lhs.mark_definitely_assigned_after(bc);
                rhs.mark_definitely_assigned_after(bc);
                // We might have evaluate types turned off
                let _ = ty;
            }
            StmtCompiled::AssignModify(lhs, _op, rhs) => {
                rhs.mark_definitely_assigned_after(bc);
                lhs.mark_definitely_assigned_after(bc);
            }
            StmtCompiled::If(cond_t_f) => {
                let (cond, t, f) = &**cond_t_f;
                cond.mark_definitely_assigned_after(bc);
                // We could merge `t` and `f` definitely assigned, e. g.
                // ```
                // if cond:
                //   x = 1
                // else:
                //   x = 2
                // ```
                // we could mark `x` as definitely assigned.
                let _ = (t, f);
            }
            StmtCompiled::For(var_over_body) => {
                let (_var, over, _body) = &**var_over_body;
                over.mark_definitely_assigned_after(bc);
            }
            StmtCompiled::Break => {}
            StmtCompiled::Continue => {}
        }
    }

    /// If statement is `return x`, return `x`.
    pub(crate) fn as_return(&self) -> Option<&IrSpanned<ExprCompiled>> {
        match self {
            StmtCompiled::Return(e) => Some(e),
            _ => None,
        }
    }
}

impl IrSpanned<StmtCompiled> {
    fn write_bc(&self, compiler: &StmtCompileContext, bc: &mut BcWriter) {
        bc.mark_before_stmt(self.span);
        self.write_bc_inner(compiler, bc);
        self.mark_definitely_assigned_after(bc);
    }

    fn write_if_then(
        compiler: &StmtCompileContext,
        bc: &mut BcWriter,
        c: &IrSpanned<ExprCompiled>,
        maybe_not: MaybeNot,
        t: &dyn Fn(&StmtCompileContext, &mut BcWriter),
    ) {
        write_if_then(
            c,
            maybe_not,
            |bc| {
                t(compiler, bc);
            },
            bc,
        );
    }

    fn write_if_else(
        c: &IrSpanned<ExprCompiled>,
        t: &StmtsCompiled,
        f: &StmtsCompiled,
        compiler: &StmtCompileContext,
        bc: &mut BcWriter,
    ) {
        assert!(!t.is_empty() || !f.is_empty());
        if f.is_empty() {
            Self::write_if_then(compiler, bc, c, MaybeNot::Id, &|compiler, bc| {
                t.write_bc(compiler, bc)
            });
        } else if t.is_empty() {
            Self::write_if_then(compiler, bc, c, MaybeNot::Not, &|compiler, bc| {
                f.write_bc(compiler, bc)
            });
        } else {
            write_if_else(
                c,
                |bc| t.write_bc(compiler, bc),
                |bc| f.write_bc(compiler, bc),
                bc,
            );
        }
    }

    fn write_return(
        span: FrameSpan,
        expr: &IrSpanned<ExprCompiled>,
        compiler: &StmtCompileContext,
        bc: &mut BcWriter,
    ) {
        bc.write_iter_stop(span);
        if compiler.has_return_type {
            expr.write_bc_cb(bc, |slot, bc| {
                bc.write_instr::<InstrReturnCheckType>(span, slot);
            });
        } else {
            if let Some(value) = expr.as_value() {
                bc.write_instr::<InstrReturnConst>(span, value);
            } else {
                expr.write_bc_cb(bc, |slot, bc| {
                    bc.write_instr::<InstrReturn>(span, slot);
                });
            }
        }
    }

    fn write_bc_inner(&self, compiler: &StmtCompileContext, bc: &mut BcWriter) {
        let span = self.span;
        match &self.node {
            StmtCompiled::PossibleGc => bc.write_instr::<InstrPossibleGc>(span, ()),
            StmtCompiled::Return(expr) => Self::write_return(span, expr, compiler, bc),
            StmtCompiled::Expr(expr) => {
                expr.write_bc_for_effect(bc);
            }
            StmtCompiled::Assign(lhs, ty, rhs) => {
                fn check_type(
                    ty: &Option<IrSpanned<TypeCompiled<FrozenValue>>>,
                    slot_expr: BcSlotIn,
                    bc: &mut BcWriter,
                ) {
                    if let Some(ty) = ty {
                        bc.write_instr::<InstrCheckType>(ty.span, (slot_expr, ty.node))
                    }
                }

                if let Some(local) = lhs.as_local_non_captured() {
                    // Write expression directly to local slot.
                    rhs.write_bc(local.to_bc_slot().to_out(), bc);
                    check_type(ty, local.to_bc_slot().to_in(), bc);
                } else {
                    rhs.write_bc_cb(bc, |slot, bc| {
                        check_type(ty, slot, bc);
                        lhs.write_bc(slot, bc);
                    });
                }
            }
            StmtCompiled::AssignModify(lhs, op, rhs) => {
                lhs.write_bc(span, *op, rhs, bc);
            }
            StmtCompiled::If(c_t_f) => {
                let (c, t, f) = &**c_t_f;
                Self::write_if_else(c, t, f, compiler, bc);
            }
            StmtCompiled::For(assign_over_body) => {
                let (assign, over, body) = &**assign_over_body;
                write_for(over, assign, span, bc, |bc| body.write_bc(compiler, bc));
            }
            StmtCompiled::Break => {
                bc.write_break(span);
            }
            StmtCompiled::Continue => {
                bc.write_continue(span);
            }
        }
    }
}

impl StmtsCompiled {
    pub(crate) fn as_bc(
        &self,
        compiler: &StmtCompileContext,
        local_names: FrozenRef<'static, [FrozenStringValue]>,
        param_count: u32,
        heap: &FrozenHeap,
    ) -> Bc {
        let mut bc = BcWriter::new(local_names, param_count, heap);
        self.write_bc(compiler, &mut bc);

        // Small optimization: if the last statement is return,
        // we do not need to write another return.
        if !matches!(self.last().map(|s| &s.node), Some(StmtCompiled::Return(..))) {
            let span = self.last().map(|s| s.span.end_span()).unwrap_or_default();
            if compiler.has_return_type {
                bc.alloc_slot(|slot, bc| {
                    bc.write_const(span, FrozenValue::new_none(), slot.to_out());
                    bc.write_instr::<InstrReturnCheckType>(span, slot.to_in());
                });
            } else {
                bc.write_instr::<InstrReturnConst>(span, FrozenValue::new_none());
            }
        }

        bc.finish()
    }
}


/*
 * 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::coerce::coerce;
use crate::collections::symbol_map::Symbol;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::opt_ctx::OptCtx;
use crate::eval::compiler::scope::payload::CstArgument;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::Compiler;
use crate::eval::runtime::arguments::ArgNames;
use crate::eval::runtime::arguments::ArgumentsFull;
use crate::eval::Arguments;
use crate::slice_vec_ext::SliceExt;
use crate::syntax::ast::ArgumentP;
use crate::values::FrozenStringValue;
use crate::values::FrozenValue;
use crate::values::Value;

#[derive(Default, Clone, Debug, VisitSpanMut)]
pub(crate) struct ArgsCompiledValue {
    pub(crate) pos_named: Vec<IrSpanned<ExprCompiled>>,
    /// Named arguments compiled.
    ///
    /// Note names are guaranteed to be unique here because names are validated in AST:
    /// named arguments in [`Expr::Call`] are unique.
    pub(crate) names: Vec<(Symbol, FrozenStringValue)>,
    pub(crate) args: Option<IrSpanned<ExprCompiled>>,
    pub(crate) kwargs: Option<IrSpanned<ExprCompiled>>,
}

impl ArgsCompiledValue {
    /// Check if arguments is one positional argument.
    pub(crate) fn one_pos(&self) -> Option<&IrSpanned<ExprCompiled>> {
        let ArgsCompiledValue {
            pos_named,
            names,
            args,
            kwargs,
        } = self;
        match (pos_named.as_slice(), names.as_slice(), args, kwargs) {
            ([pos], [], None, None) => Some(pos),
            _ => None,
        }
    }

    pub(crate) fn pos_only(&self) -> Option<&[IrSpanned<ExprCompiled>]> {
        if self.names.is_empty() && self.args.is_none() && self.kwargs.is_none() {
            Some(&self.pos_named)
        } else {
            None
        }
    }

    fn split_pos_names(&self) -> (&[IrSpanned<ExprCompiled>], &[IrSpanned<ExprCompiled>]) {
        self.pos_named
            .as_slice()
            .split_at(self.pos_named.len() - self.names.len())
    }

    /// Invoke a callback if all arguments are frozen values.
    pub(crate) fn all_values<'v, R>(
        &self,
        handler: impl FnOnce(&Arguments<'v, '_>) -> R,
    ) -> Option<R> {
        self.all_values_generic(|e| e.as_value().map(FrozenValue::to_value), handler)
    }

    /// Invoke a callback if all arguments are frozen values.
    pub(crate) fn all_values_generic<'v, R>(
        &self,
        expr_to_value: impl Fn(&ExprCompiled) -> Option<Value<'v>>,
        handler: impl FnOnce(&Arguments<'v, '_>) -> R,
    ) -> Option<R> {
        let (pos, named) = self.split_pos_names();
        let pos = pos.try_map(|e| expr_to_value(e).ok_or(())).ok()?;
        let named = named.try_map(|e| expr_to_value(e).ok_or(())).ok()?;
        let args = self
            .args
            .as_ref()
            .map(|args| expr_to_value(args).ok_or(()))
            .transpose()
            .ok()?;
        let kwargs = self
            .kwargs
            .as_ref()
            .map(|kwargs| expr_to_value(kwargs).ok_or(()))
            .transpose()
            .ok()?;
        Some(handler(&Arguments(ArgumentsFull {
            pos: &pos,
            named: &named,
            names: ArgNames::new(coerce(&self.names)),
            args,
            kwargs,
        })))
    }

    /// Expressions of all arguments: positional, named, star-args, star-star-args.
    pub(crate) fn arg_exprs(&self) -> impl Iterator<Item = &IrSpanned<ExprCompiled>> {
        self.pos_named
            .iter()
            .chain(self.args.iter())
            .chain(self.kwargs.iter())
    }

    pub(crate) fn map_exprs<E>(
        &self,
        mut f: impl FnMut(&IrSpanned<ExprCompiled>) -> Result<IrSpanned<ExprCompiled>, E>,
    ) -> Result<ArgsCompiledValue, E> {
        let ArgsCompiledValue {
            pos_named,
            names,
            args,
            kwargs,
        } = self;
        Ok(ArgsCompiledValue {
            pos_named: pos_named.try_map(&mut f)?,
            names: names.clone(),
            args: args.as_ref().map(&mut f).transpose()?,
            kwargs: kwargs.as_ref().map(&mut f).transpose()?,
        })
    }

    pub(crate) fn optimize(&self, ctx: &mut OptCtx) -> ArgsCompiledValue {
        enum Never {}
        self.map_exprs(|e| Ok(e.optimize(ctx)))
            .unwrap_or_else(|e: Never| match e {})
    }
}

impl Compiler<'_, '_, '_> {
    pub(crate) fn args(&mut self, args: &[CstArgument]) -> ArgsCompiledValue {
        let mut res = ArgsCompiledValue::default();
        for x in args {
            match &x.node {
                ArgumentP::Positional(x) => res.pos_named.push(self.expr(x)),
                ArgumentP::Named(name, value) => {
                    let fv = self
                        .eval
                        .module_env
                        .frozen_heap()
                        .alloc_str(name.node.as_str());
                    res.names.push((Symbol::new(&name.node), fv));
                    res.pos_named.push(self.expr(value));
                }
                ArgumentP::Args(x) => res.args = Some(self.expr(x)),
                ArgumentP::KwArgs(x) => res.kwargs = Some(self.expr(x)),
            }
        }
        res
    }
}


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

//! Compile function calls.

use std::cell::Cell;

use starlark_derive::VisitSpanMut;

use crate::collections::symbol_map::Symbol;
use crate::eval::compiler::args::ArgsCompiledValue;
use crate::eval::compiler::def_inline::local_as_value::local_as_value;
use crate::eval::compiler::def_inline::InlineDefBody;
use crate::eval::compiler::def_inline::InlineDefCallSite;
use crate::eval::compiler::expr::Builtin1;
use crate::eval::compiler::expr::ExprCompiled;
use crate::eval::compiler::opt_ctx::OptCtx;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::runtime::frame_span::FrameSpan;
use crate::eval::runtime::inlined_frame::InlinedFrameAlloc;
use crate::eval::runtime::visit_span::VisitSpanMut;
use crate::slice_vec_ext::VecExt;
use crate::values::enumeration::FrozenEnumType;
use crate::values::string::dot_format::parse_format_one;
use crate::values::FrozenStringValue;
use crate::values::Value;

#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) struct CallCompiled {
    pub(crate) fun: IrSpanned<ExprCompiled>,
    pub(crate) args: ArgsCompiledValue,
}

impl CallCompiled {
    pub(crate) fn new_method(
        span: FrameSpan,
        this: IrSpanned<ExprCompiled>,
        field: &Symbol,
        getattr_span: FrameSpan,
        args: ArgsCompiledValue,
        ctx: &mut OptCtx,
    ) -> ExprCompiled {
        if let Some(this) = this.as_value() {
            if let Some(v) = ExprCompiled::compile_time_getattr(this, field, ctx) {
                let v = ExprCompiled::Value(v);
                let v = IrSpanned {
                    span: getattr_span,
                    node: v,
                };
                return CallCompiled::call(span, v, args, ctx);
            }
        }

        ExprCompiled::Call(Box::new(IrSpanned {
            span,
            node: CallCompiled {
                fun: IrSpanned {
                    span: getattr_span,
                    node: ExprCompiled::dot(this, field, ctx),
                },
                args,
            },
        }))
    }

    /// If this call expression is `len(x)`, return `x`.
    pub(crate) fn as_len(&self) -> Option<&IrSpanned<ExprCompiled>> {
        if !self.fun.is_fn_len() {
            return None;
        }
        self.args.one_pos()
    }

    /// If this call expression is `type(x)`, return `x`.
    pub(crate) fn as_type(&self) -> Option<&IrSpanned<ExprCompiled>> {
        if !self.fun.is_fn_type() {
            return None;
        }
        self.args.one_pos()
    }

    /// This call is infallible and has no side effects.
    pub(crate) fn is_pure_infallible(&self) -> bool {
        match self.as_type() {
            Some(arg) => arg.is_pure_infallible(),
            None => false,
        }
    }

    /// This call is a method call.
    pub(crate) fn method(&self) -> Option<(&IrSpanned<ExprCompiled>, &Symbol, &ArgsCompiledValue)> {
        match &self.fun.node {
            ExprCompiled::Builtin1(Builtin1::Dot(name), expr) => Some((expr, name, &self.args)),
            _ => None,
        }
    }

    /// Try to inline a function like `lambda x: type(x) == "y"`.
    fn try_type_is(fun: &ExprCompiled, args: &ArgsCompiledValue) -> Option<ExprCompiled> {
        let fun = fun.as_frozen_def()?;
        let pos = args.one_pos()?;
        if let Some(InlineDefBody::ReturnTypeIs(t)) = &fun.def_info.inline_def_body {
            Some(ExprCompiled::type_is(pos.clone(), *t))
        } else {
            None
        }
    }

    /// Inline calls to functions which are safe to inline.
    fn try_inline(
        span: FrameSpan,
        fun: &ExprCompiled,
        args: &ArgsCompiledValue,
        ctx: &mut OptCtx,
    ) -> Option<IrSpanned<ExprCompiled>> {
        let fun = fun.as_frozen_def()?;

        if fun.parameters.has_args_or_kwargs() {
            // Functions with `*args` or `**kwargs` are not marked safe to inline,
            // but it is safer to also handle it explicitly here.
            return None;
        }

        let expr = if let Some(InlineDefBody::ReturnSafeToInlineExpr(expr)) =
            &fun.def_info.inline_def_body
        {
            expr
        } else {
            return None;
        };

        let param_count = ctx.param_count;
        let expr_to_value = |expr: &ExprCompiled| -> Option<Value> {
            match expr {
                ExprCompiled::Value(v) => Some(v.to_value()),
                ExprCompiled::Local(local) if local.0 < param_count => {
                    // Definitely assigned local variable.
                    //
                    // Consider this example:
                    // ```
                    // def foo(x): bar() + x
                    // ```
                    // We can inline calls like `foo(x)` when `x` is definitely assigned,
                    // but if `x` is not we cannot do that, because
                    // we should emit `x` is not assigned error before call to `bar()`
                    // which may fail. We can implement inlining of variables which
                    // may be not assigned, or inlining of any expression arguments,
                    // but more work is needed for that.
                    Some(local_as_value(*local)?.to_value())
                }
                _ => None,
            }
        };

        args.all_values_generic(expr_to_value, |arguments| {
            let slots = vec![Cell::new(None); fun.parameters.len()];
            fun.parameters
                .collect(arguments.frozen_to_v(), &slots, ctx.heap())
                .ok()?;

            let slots = slots
                .into_try_map(|value| {
                    // Value must be set, but better ignore optimization here than panic.
                    let value = value.get().ok_or(())?;
                    // Everything should be frozen here, but if not,
                    // it is safer to abandon optimization.
                    value.unpack_frozen().ok_or(())
                })
                .ok()?;

            let mut expr = IrSpanned {
                span,
                node: expr.node.clone(),
            };
            let mut span_alloc = InlinedFrameAlloc::new(ctx.frozen_heap());
            expr.visit_spans(&mut |expr_span: &mut FrameSpan| {
                expr_span
                    .inlined_frames
                    .inline_into(span, fun.to_frozen_value(), &mut span_alloc);
            });
            InlineDefCallSite { ctx, slots: &slots }.inline(&expr).ok()
        })?
    }

    fn try_spec_exec(
        span: FrameSpan,
        fun: &ExprCompiled,
        args: &ArgsCompiledValue,
        ctx: &mut OptCtx,
    ) -> Option<ExprCompiled> {
        let fun = fun.as_value()?;

        if !fun.speculative_exec_safe() {
            return None;
        }

        let eval = ctx.eval()?;

        // Only if all call arguments are frozen values.
        args.all_values(|arguments| {
            let v = fun.to_value().invoke(arguments.frozen_to_v(), eval).ok()?;
            ExprCompiled::try_value(span, v, eval.module_env.frozen_heap())
        })?
    }

    // Optimize `MyEnum(arg)`.
    fn try_enum_value(
        fun: &IrSpanned<ExprCompiled>,
        args: &ArgsCompiledValue,
    ) -> Option<ExprCompiled> {
        let fun = fun.as_value()?.downcast_frozen_ref::<FrozenEnumType>()?;
        let arg = args.one_pos()?.as_value()?;
        Some(ExprCompiled::Value(
            fun.value.construct(arg.to_value()).ok()?,
        ))
    }

    // Optimize `"aaa{}bbb".format(arg)`.
    fn try_format(
        fun: &IrSpanned<ExprCompiled>,
        args: &ArgsCompiledValue,
        ctx: &mut OptCtx,
    ) -> Option<ExprCompiled> {
        let fun = fun.as_frozen_bound_method()?;
        let format = FrozenStringValue::new(fun.this)?;
        if fun.method.name != "format" {
            return None;
        }
        let arg = args.one_pos()?;

        let (before, after) = parse_format_one(&format)?;

        let before = ctx.frozen_heap().alloc_str(&before);
        let after = ctx.frozen_heap().alloc_str(&after);
        Some(ExprCompiled::format_one(before, arg.clone(), after, ctx))
    }

    pub(crate) fn call(
        span: FrameSpan,
        fun: IrSpanned<ExprCompiled>,
        args: ArgsCompiledValue,
        ctx: &mut OptCtx,
    ) -> ExprCompiled {
        if let Some(type_is) = CallCompiled::try_type_is(&fun, &args) {
            return type_is;
        }

        if let Some(inline) = CallCompiled::try_inline(span, &fun, &args, ctx) {
            return inline.node;
        }

        if fun.is_fn_len() {
            if let Some(arg) = args.one_pos() {
                return ExprCompiled::len(span, arg.clone());
            }
        }

        if fun.is_fn_type() {
            if let Some(arg) = args.one_pos() {
                return ExprCompiled::typ(span, arg.clone());
            }
        }

        if let Some(r) = CallCompiled::try_enum_value(&fun, &args) {
            return r;
        }

        if let Some(r) = CallCompiled::try_spec_exec(span, &fun, &args, ctx) {
            return r;
        }

        if let Some(r) = CallCompiled::try_format(&fun, &args, ctx) {
            return r;
        }

        if let ExprCompiled::Builtin1(Builtin1::Dot(field), this) = &fun.node {
            return CallCompiled::new_method(span, (**this).clone(), field, fun.span, args, ctx);
        }

        ExprCompiled::Call(Box::new(IrSpanned {
            span,
            node: CallCompiled { fun, args },
        }))
    }
}

impl IrSpanned<CallCompiled> {
    pub(crate) fn optimize(&self, ctx: &mut OptCtx) -> ExprCompiled {
        let CallCompiled { fun: expr, args } = &self.node;
        let expr = expr.optimize(ctx);
        let args = args.optimize(ctx);
        CallCompiled::call(self.span, expr, args, ctx)
    }
}


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

//! List/dict/set comprenension evaluation.

use starlark_derive::VisitSpanMut;

use crate::eval::compiler::expr::ExprCompiled;
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::CstExpr;
use crate::eval::compiler::scope::payload::CstPayload;
use crate::eval::compiler::span::IrSpanned;
use crate::eval::compiler::stmt::AssignCompiledValue;
use crate::eval::compiler::Compiler;
use crate::slice_vec_ext::SliceExt;
use crate::syntax::ast::ClauseP;
use crate::syntax::ast::ForClauseP;

impl Compiler<'_, '_, '_> {
    pub fn list_comprehension(
        &mut self,
        x: &CstExpr,
        for_: &ForClauseP<CstPayload>,
        clauses: &[ClauseP<CstPayload>],
    ) -> ExprCompiled {
        let clauses = self.compile_clauses(for_, clauses);
        let x = self.expr(x);
        ExprCompiled::compr(ComprCompiled::List(Box::new(x), clauses))
    }

    pub fn dict_comprehension(
        &mut self,
        k: &CstExpr,
        v: &CstExpr,
        for_: &ForClauseP<CstPayload>,
        clauses: &[ClauseP<CstPayload>],
    ) -> ExprCompiled {
        let clauses = self.compile_clauses(for_, clauses);
        let k = self.expr(k);
        let v = self.expr(v);
        ExprCompiled::compr(ComprCompiled::Dict(Box::new((k, v)), clauses))
    }

    /// Peel the final if's from clauses, and return them (in the order they started), plus the next for you get to
    fn compile_ifs(
        &mut self,
        clauses: &mut Vec<ClauseP<CstPayload>>,
    ) -> (Option<ForClauseP<CstPayload>>, Vec<IrSpanned<ExprCompiled>>) {
        let mut ifs = Vec::new();
        while let Some(x) = clauses.pop() {
            match x {
                ClauseP::For(f) => {
                    ifs.reverse();
                    return (Some(f), ifs);
                }
                ClauseP::If(x) => {
                    let x = self.expr_truth(&x);
                    if let ExprCompiledBool::Const(true) = &x.node {
                        // If the condition is always true, skip the clause.
                        continue;
                    }
                    // TODO(nga): if condition is false, do something better.
                    ifs.push(x.into_expr());
                }
            }
        }
        ifs.reverse();
        (None, ifs)
    }

    fn compile_clauses(
        &mut self,
        for_: &ForClauseP<CstPayload>,
        clauses: &[ClauseP<CstPayload>],
    ) -> ClausesCompiled {
        // The first for.over is scoped before we enter the list comp
        let over = self.expr(&list_to_tuple(&for_.over));

        // TODO(nga): unnecessary clone.
        let mut clauses = clauses.to_vec();

        // Now we want to group them into a `for`, followed by any number of `if`.
        // The evaluator wants to use pop to consume them, so reverse the order.
        let mut res = Vec::new();
        loop {
            let (next_for, ifs) = self.compile_ifs(&mut clauses);
            match next_for {
                None => {
                    let last = ClauseCompiled {
                        var: self.assign(&for_.var),
                        over,
                        ifs,
                    };
                    return ClausesCompiled::new(res, last);
                }
                Some(f) => {
                    res.push(ClauseCompiled {
                        over: self.expr(&f.over),
                        var: self.assign(&f.var),
                        ifs,
                    });
                }
            }
        }
    }
}

#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) enum ComprCompiled {
    List(Box<IrSpanned<ExprCompiled>>, ClausesCompiled),
    Dict(
        Box<(IrSpanned<ExprCompiled>, IrSpanned<ExprCompiled>)>,
        ClausesCompiled,
    ),
}

impl ComprCompiled {
    pub(crate) fn clauses(&self) -> &ClausesCompiled {
        match self {
            ComprCompiled::List(_, clauses) => clauses,
            ComprCompiled::Dict(_, clauses) => clauses,
        }
    }

    pub(crate) fn optimize(&self, ctx: &mut OptCtx) -> ExprCompiled {
        match self {
            ComprCompiled::List(ref x, ref clauses) => {
                let clauses = clauses.optimize(ctx);
                ExprCompiled::compr(ComprCompiled::List(Box::new(x.optimize(ctx)), clauses))
            }
            ComprCompiled::Dict(k_v, ref clauses) => {
                let (k, v) = &**k_v;
                let clauses = clauses.optimize(ctx);
                ExprCompiled::compr(ComprCompiled::Dict(
                    Box::new((k.optimize(ctx), v.optimize(ctx))),
                    clauses.optimize(ctx),
                ))
            }
        }
    }
}

#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) struct ClauseCompiled {
    pub(crate) var: IrSpanned<AssignCompiledValue>,
    pub(crate) over: IrSpanned<ExprCompiled>,
    pub(crate) ifs: Vec<IrSpanned<ExprCompiled>>,
}

impl ClauseCompiled {
    fn optimize(&self, ctx: &mut OptCtx) -> ClauseCompiled {
        let ClauseCompiled {
            ref var,
            ref over,
            ref ifs,
        } = *self;
        ClauseCompiled {
            var: var.optimize(ctx),
            over: over.optimize(ctx),
            ifs: ifs
                .iter()
                .filter_map(|e| {
                    let e = e.optimize(ctx);
                    let e = ExprCompiledBool::new(e);
                    match &e.node {
                        ExprCompiledBool::Const(true) => None,
                        _ => Some(e.into_expr()),
                    }
                })
                .collect(),
        }
    }
}

/// All clauses in a comprehension. Never empty.
#[derive(Clone, Debug, VisitSpanMut)]
pub(crate) struct ClausesCompiled {
    /// Not empty.
    ///
    /// Clauses are in reverse order, i. e. the first executed clause is the last in the list.
    clauses: Vec<ClauseCompiled>,
}

impl ClausesCompiled {
    fn new(clauses: Vec<ClauseCompiled>, last: ClauseCompiled) -> ClausesCompiled {
        let mut clauses = clauses;
        clauses.push(last);
        ClausesCompiled { clauses }
    }

    /// Clauses are definitely no-op, i. e. zero iterations, and no side effects of iteration.
    pub(crate) fn is_nop(&self) -> bool {
        // NOTE(nga): if the first loop argument is empty collection, clauses are definitely no-op.
        //   But this is not true for the rest of loops: if inner loop collection is empty,
        //   clauses produce no iterations, but the outer loop may still has side effects.
        //   There are missing optimizations here:
        //   * we could separate effects and emit empty list/dict.
        //   * or at least do not generate comprehension terminator.
        self.split_last().0.over.is_iterable_empty()
    }

    /// Last clause is the one which is executed first.
    pub(crate) fn split_last(&self) -> (&ClauseCompiled, &[ClauseCompiled]) {
        self.clauses.split_last().unwrap()
    }

    fn optimize(&self, ctx: &mut OptCtx) -> ClausesCompiled {
        ClausesCompiled {
            clauses: self.clauses.map(|c| c.optimize(ctx)),
        }
    }
}


/*
 * 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 dupe::Dupe;
use once_cell::sync::Lazy;

use crate::environment::Globals;
use crate::values::FrozenValue;

#[derive(Copy, Clone, Dupe, Debug)]
pub(crate) struct BuiltinFn(pub(crate) FrozenValue);

impl PartialEq<FrozenValue> for BuiltinFn {
    fn eq(&self, other: &FrozenValue) -> bool {
        // Pointer equality works because `#[starlark_module]` proc macro
        // generates a singleton which allocates the function only once
        // even if builder function is called multiple times.
        self.0.to_value().ptr_eq(other.to_value())
    }
}

impl PartialEq<BuiltinFn> for FrozenValue {
    fn eq(&self, other: &BuiltinFn) -> bool {
        other == self
    }
}

pub(crate) struct Constants {
    pub(crate) fn_len: BuiltinFn,
    pub(crate) fn_type: BuiltinFn,
    pub(crate) fn_list: BuiltinFn,
    pub(crate) fn_dict: BuiltinFn,
}

impl Constants {
    pub fn get() -> &'static Constants {
        static RES: Lazy<Constants> = Lazy::new(|| {
            let g = Globals::standard();
            Constants {
                fn_len: BuiltinFn(g.get_frozen("len").unwrap()),
                fn_type: BuiltinFn(g.get_frozen("type").unwrap()),
                fn_list: BuiltinFn(g.get_frozen("list").unwrap()),
                fn_dict: BuiltinFn(g.get_frozen("dict").unwrap()),
            }
        });
        Lazy::force(&RES)
    }
}

#[cfg(test)]
mod tests {
    use crate::environment::Globals;
    use crate::eval::compiler::constants::Constants;

    #[test]
    fn test_constants() {
        assert_eq!(
            Globals::standard().get_frozen("len").unwrap(),
            Constants::get().fn_len
        );
        assert_eq!(
            Globals::extended().get_frozen("len").unwrap(),
            Constants::get().fn_len
        );
    }
}