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

use either::Either;
use itertools::Itertools;
use starlark_derive::starlark_module;

use crate as starlark;
use crate::assert::Assert;
use crate::environment::GlobalsBuilder;
use crate::values::dict::DictOf;
use crate::values::list::ListOf;
use crate::values::structs::StructOf;
use crate::values::Value;
use crate::values::ValueOf;

// TODO(nmj): Figure out default values here. ValueOf<i32> = 5 should work.
#[starlark_module]
fn validate_module(builder: &mut GlobalsBuilder) {
    fn with_int<'v>(v: ValueOf<'v, i32>) -> anyhow::Result<(Value<'v>, String)> {
        Ok((*v, format!("{}", v.typed)))
    }
    fn with_int_list<'v>(v: ListOf<'v, i32>) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v.to_vec().iter().join(", ");
        Ok((*v, repr))
    }
    fn with_list_list<'v>(v: ListOf<'v, ListOf<'v, i32>>) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_vec()
            .iter()
            .map(|l| l.to_vec().iter().join(", "))
            .join(" + ");
        Ok((*v, repr))
    }
    fn with_dict_list<'v>(
        v: ListOf<'v, DictOf<'v, i32, i32>>,
    ) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_vec()
            .iter()
            .map(|l| {
                l.to_dict()
                    .iter()
                    .map(|(k, v)| format!("{}: {}", k, v))
                    .join(", ")
            })
            .join(" + ");
        Ok((*v, repr))
    }
    fn with_int_dict<'v>(v: DictOf<'v, i32, i32>) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_dict()
            .iter()
            .map(|(k, v)| format!("{}: {}", k, v))
            .join(" + ");
        Ok((*v, repr))
    }
    fn with_list_dict<'v>(
        v: DictOf<'v, i32, ListOf<'v, i32>>,
    ) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_dict()
            .iter()
            .map(|(k, v)| format!("{}: {}", k, v.to_vec().iter().join(", ")))
            .join(" + ");
        Ok((*v, repr))
    }
    fn with_dict_dict<'v>(
        v: DictOf<'v, i32, DictOf<'v, i32, i32>>,
    ) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_dict()
            .iter()
            .map(|(k, v)| {
                let inner_repr = v
                    .to_dict()
                    .iter()
                    .map(|(k2, v2)| format!("{}:{}", k2, v2))
                    .join(", ");
                format!("{}: {}", k, inner_repr)
            })
            .join(" + ");
        Ok((*v, repr))
    }
    fn with_struct_int<'v>(v: StructOf<'v, i32>) -> anyhow::Result<(Value<'v>, String)> {
        let repr = v
            .to_map()
            .iter()
            .map(|(k, v)| format!("{}={}", k, v))
            .join(" + ");
        Ok((v.to_value(), repr))
    }
    fn with_either(v: Either<i32, Either<String, ListOf<i32>>>) -> anyhow::Result<String> {
        match v {
            Either::Left(i) => Ok(i.to_string()),
            Either::Right(nested) => match nested {
                Either::Left(s) => Ok(s),
                Either::Right(l) => Ok(l.to_repr()),
            },
        }
    }
}

// The standard error these raise on incorrect types
const BAD: &str = "Type of parameter";

#[test]
fn test_value_of() {
    let mut a = Assert::new();
    a.globals_add(validate_module);
    a.eq("(1, '1')", "with_int(1)");
    a.fail("with_int(noop(None))", BAD);
}

#[test]
fn test_list_of() {
    let mut a = Assert::new();
    a.globals_add(validate_module);
    a.eq("([1, 2, 3], '1, 2, 3')", "with_int_list([1, 2, 3])");
    a.fail("with_int_list(noop(1))", BAD);
    a.fail("with_int_list(noop([1, 'foo']))", BAD);
    a.fail("with_int_list(noop([[]]))", BAD);

    a.eq(
        "([[1, 2], [3]], '1, 2 + 3')",
        "with_list_list([[1, 2], [3]])",
    );

    let expected = r#"([{1: 2, 3: 4}, {5: 6}], "1: 2, 3: 4 + 5: 6")"#;
    let test = r#"with_dict_list([{1: 2, 3: 4}, {5: 6}])"#;
    a.eq(expected, test);
}

#[test]
fn test_dict_of() {
    let mut a = Assert::new();
    a.globals_add(validate_module);
    a.eq("({1: 2}, '1: 2')", "with_int_dict({1: 2})");
    a.fail(r#"with_int_dict(noop(1))"#, BAD);
    a.fail(r#"with_int_dict(noop({1: "str"}))"#, BAD);
    a.fail(r#"with_int_dict(noop({1: {}}))"#, BAD);

    let expected = r#"({1: [2, 3], 4: [5]}, "1: 2, 3 + 4: 5")"#;
    let test = r#"with_list_dict({1: [2, 3], 4: [5]})"#;
    a.eq(expected, test);

    let expected = r#"({1: {2: 3, 4: 5}, 6: {7: 8}}, "1: 2:3, 4:5 + 6: 7:8")"#;
    let test = r#"with_dict_dict({1: {2: 3, 4: 5}, 6: {7: 8}})"#;
    a.eq(expected, test);
}

#[test]
fn test_struct_of() {
    let mut a = Assert::new();
    a.globals_add(validate_module);
    a.eq("(struct(a=1), '\"a\"=1')", "with_struct_int(struct(a=1))");
    a.fail("with_struct_int(struct(a=True))", BAD);
}

#[test]
fn test_either_of() {
    let mut a = Assert::new();
    a.globals_add(validate_module);
    a.eq("'2'", "with_either(2)");
    a.eq("'[2, 3]'", "with_either([2,3])");
    a.eq("'s'", "with_either('s')");
    a.fail("with_either(noop(None))", BAD);
    a.fail("with_either(noop({}))", BAD);
}


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

use starlark_derive::Trace;

use crate as starlark;
use crate::values::Value;

#[allow(dead_code)] // Just check it compiles.
#[derive(Trace)]
enum TestEnumForTrace<'v> {
    Foo(String),
    Bar(i32, Value<'v>),
}


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

mod enums;
mod statics;


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

use starlark_derive::Trace;

use crate as starlark;
use crate::values::Value;

#[allow(dead_code)] // Just check it compiles.
#[derive(Trace)]
struct TraceWithStatic<'v> {
    actual_value: Value<'v>,
    // This field doesn't have a Trace trait, but should be ignored
    // because it looks like it is static
    ignored_because_static: StaticType<'static, std::string::String>,
}

#[allow(dead_code)]
struct StaticType<'a, T> {
    inner: &'a T,
}


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

use crate::docs::Doc;
use crate::docs::DocItem;
use crate::docs::MarkdownFlavor;
use crate::docs::RenderMarkdown;
use crate::tests::golden_test_template::golden_test_template;

pub(crate) fn docs_golden_test(test_name: &str, doc: DocItem) -> String {
    let output = Doc::named_item("name".to_owned(), doc).render_markdown(MarkdownFlavor::DocFile);

    golden_test_template(
        &format!("src/tests/docs/golden/{test_name}.golden.md"),
        &output,
    );

    output
}


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

use allocative::Allocative;
use derive_more::Display;
use serde::Serialize;
use starlark_derive::starlark_module;
use starlark_derive::starlark_value;
use starlark_derive::NoSerialize;

use crate as starlark;
use crate::any::ProvidesStaticType;
use crate::assert;
use crate::docs::Doc;
use crate::docs::DocItem;
use crate::docs::MarkdownFlavor;
use crate::docs::RenderMarkdown;
use crate::environment::GlobalsBuilder;
use crate::environment::Methods;
use crate::environment::MethodsBuilder;
use crate::environment::MethodsStatic;
use crate::tests::docs::golden::docs_golden_test;
use crate::values::none::NoneType;
use crate::values::Heap;
use crate::values::StarlarkValue;
use crate::values::Value;

mod golden;
mod rustdocs;

const STARLARK_CODE: &str = r#"
"""
This is the summary of the module's docs

Some extra details can go here,
    and indentation is kept as expected
"""

def f1(a, b: "string", c:"int" = 5, *, d:"string" = "some string", **kwargs) -> ["string"]:
    """
    Summary line goes here

    Args:
        a: The docs for a
        b: The docs for b
        c: The docs for c, but these
           go onto two lines
        **kwargs: Docs for the keyword args

    Returns:
        A string repr of the args
    """
    return [str((a, b, c, d, repr(kwargs)))]

def f2(a, *args: ["string"]):
    """
    This is a function with *args, and no return type

    Args:
        *args: Only doc this arg
    """
    return None

def f3(a: "string") -> "string":
    return a

def f4(a: "string") -> "string":
    """ This is a docstring with no 'Args:' section """
    return a

# Not public, so shouldn't show up
def _do_not_export():
    pass
"#;

#[derive(
    Debug,
    derive_more::Display,
    ProvidesStaticType,
    Allocative,
    NoSerialize
)]
#[display(fmt = "magic")]
struct Magic;

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

/// These are where the module docs go
#[starlark_module]
fn module(builder: &mut GlobalsBuilder) {
    const MAGIC: i32 = 42;

    /// Docs for func1
    ///
    /// # Arguments
    ///     * `foo`: Docs for foo
    ///
    /// # Returns
    /// The string 'func1'
    fn func1(foo: String) -> anyhow::Result<String> {
        let _ignore = foo;
        Ok("func1".to_owned())
    }

    fn func2() -> anyhow::Result<String> {
        Ok("func2".to_owned())
    }

    /// A function with only positional arguments.
    ///
    /// And a slightly longer description. With some example code:
    ///
    /// ```python
    /// func3(1)
    /// ```
    ///
    /// And some assertions:
    ///
    /// ```rust
    /// # starlark::assert::all_true(r#"
    /// 1 == 1
    /// # "#);
    /// ```
    #[starlark(as_type = Magic)]
    fn func3(
        #[starlark(require = pos)] a1: i32,
        #[starlark(require = pos)] a2: Option<i32>,
        #[starlark(require = pos, default = 1)] step: i32,
    ) -> anyhow::Result<String> {
        let _ = (a1, a2, step);
        Ok("func3".to_owned())
    }

    fn with_defaults<'v>(
        #[starlark(default=Vec::new())] explicit_default: Vec<String>,
        hidden_default: Option<Vec<String>>,
        #[starlark(default = "my_default")] string_default: &str,
    ) -> anyhow::Result<NoneType> {
        let _unused = (explicit_default, hidden_default, string_default);
        Ok(NoneType)
    }

    fn pos_either_named(
        #[starlark(require = pos)] a: i32,
        b: i32,
        #[starlark(require = named)] c: i32,
    ) -> anyhow::Result<NoneType> {
        let _unused = (a, b, c);
        Ok(NoneType)
    }

    fn notypes<'v>(a: Value<'v>) -> anyhow::Result<Value<'v>> {
        Ok(a)
    }
}

#[derive(ProvidesStaticType, Debug, Display, Allocative, Serialize)]
#[display(format = "obj")]
struct Obj;

#[starlark_value(type = "obj")]
impl<'v> StarlarkValue<'v> for Obj {
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(object)
    }
}

/// These are where the module docs go
#[starlark_module]
fn object(builder: &mut MethodsBuilder) {
    /// Docs for attr1
    #[starlark(attribute)]
    fn attr1<'v>(this: Value<'v>) -> anyhow::Result<String> {
        Ok("attr1".to_owned())
    }

    #[starlark(attribute)]
    fn attr2<'v>(this: Value<'v>) -> anyhow::Result<String> {
        Ok("attr2".to_owned())
    }

    /// Docs for func1
    ///
    /// # Arguments
    ///     * `foo`: Docs for foo
    ///
    /// # Returns
    /// The string 'func1'
    fn func1<'v>(this: Value<'v>, foo: String) -> anyhow::Result<String> {
        let _ignore = (this, foo);
        Ok("func1".to_owned())
    }

    fn func2<'v>(this: Value<'v>) -> anyhow::Result<String> {
        let _ = this;
        Ok("func2".to_owned())
    }

    /// Needs to be escaped when rendered in markdown.
    fn __exported__<'v>(this: Value<'v>) -> anyhow::Result<NoneType> {
        let _ = this;
        Ok(NoneType)
    }
}

#[test]
fn golden_docs_starlark() {
    let res = docs_golden_test(
        "starlark",
        DocItem::Module(assert::pass_module(STARLARK_CODE).documentation()),
    );
    assert!(!res.contains("_do_not_export"));
}

#[test]
fn golden_docs_module() {
    let res = docs_golden_test(
        "module",
        DocItem::Module(GlobalsBuilder::new().with(module).build().documentation()),
    );
    assert!(!res.contains("starlark::assert::all_true"));
    assert!(res.contains(r#"string_default: str.type = "my_default"#));
}

#[test]
fn golden_docs_object() {
    let res = docs_golden_test("object", Obj.documentation().unwrap());
    assert!(res.contains(r#"name.\_\_exported\_\_"#));
}

#[test]
fn inner_object_functions_have_docs() {
    let heap = Heap::new();
    let obj = heap.alloc_simple(Obj);
    let item = obj
        .get_attr("func1", &heap)
        .unwrap()
        .unwrap()
        .documentation()
        .unwrap();
    let res = Doc::named_item("func1".to_owned(), item).render_markdown(MarkdownFlavor::DocFile);
    assert!(res.contains("Docs for func1"));
}

#[test]
fn inner_module_functions_have_docs() {
    let item = GlobalsBuilder::new()
        .with(module)
        .build()
        .get("func1")
        .unwrap()
        .documentation()
        .unwrap();
    let res = Doc::named_item("func1".to_owned(), item).render_markdown(MarkdownFlavor::DocFile);
    assert!(res.contains("Docs for func1"))
}


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

use std::collections::HashMap;

use starlark_derive::starlark_module;
use starlark_map::small_map::SmallMap;

use crate as starlark;
use crate::assert;
use crate::docs::DocItem;
use crate::docs::DocMember;
use crate::environment::GlobalsBuilder;
use crate::eval::Arguments;
use crate::eval::Evaluator;
use crate::typing::Ty;
use crate::values::none::NoneType;
use crate::values::type_repr::StarlarkTypeRepr;
use crate::values::Heap;
use crate::values::StringValue;
use crate::values::Value;
use crate::values::ValueOfUnchecked;

struct InputTypeRepr;
struct OutputTypeRepr;

impl StarlarkTypeRepr for InputTypeRepr {
    fn starlark_type_repr() -> Ty {
        Ty::name("input")
    }
}
impl StarlarkTypeRepr for OutputTypeRepr {
    fn starlark_type_repr() -> Ty {
        Ty::name("output")
    }
}

#[starlark_module]
#[allow(unused_variables)] // Since this is for a test
fn globals(builder: &mut GlobalsBuilder) {
    fn simple(
        arg_int: i32,
        arg_bool: bool,
        arg_vec: Vec<&str>,
        arg_dict: SmallMap<String, (bool, i32)>,
    ) -> anyhow::Result<NoneType> {
        unimplemented!()
    }

    fn default_arg<'v>(
        arg1: Option<Value<'v>>,
        #[starlark(default = NoneType)] arg2: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Vec<String>> {
        unimplemented!()
    }

    fn args_kwargs<'v>(
        #[starlark(args)] args: Vec<Value<'v>>,
        #[starlark(kwargs)] kwargs: Value<'v>,
    ) -> anyhow::Result<NoneType> {
        unimplemented!()
    }

    fn custom_types<'v>(
        arg1: StringValue<'v>,
        arg2: ValueOfUnchecked<'v, InputTypeRepr>,
        heap: &'v Heap,
    ) -> anyhow::Result<ValueOfUnchecked<'v, OutputTypeRepr>> {
        unimplemented!()
    }

    fn pos_named(arg1: i32, #[starlark(require = named)] arg2: i32) -> anyhow::Result<i32> {
        unimplemented!()
    }

    fn with_arguments(args: &Arguments) -> anyhow::Result<i32> {
        unimplemented!()
    }
}

/// Test that a Rust starlark_module produces the right documentation.
#[test]
fn test_rustdoc() {
    let got = GlobalsBuilder::new().with(globals).build();
    let expected = assert::pass_module(
        r#"
def args_kwargs(*args, **kwargs: "") -> None: pass
def custom_types(arg1: str.type, arg2: "input") -> "output": pass
def default_arg(arg1 = "_", arg2: "" = None) -> [str.type]: pass
def pos_named(arg1: int.type, *, arg2: int.type) -> int.type: pass
def simple(arg_int: int.type, arg_bool: bool.type, arg_vec: [str.type], arg_dict: {str.type: (bool.type, int.type)}) -> None: pass
def with_arguments(*args, **kwargs) -> int.type: pass
"#,
    );

    fn unpack(x: DocItem) -> HashMap<String, DocItem> {
        match x {
            DocItem::Module(obj) => obj
                .members
                .into_iter()
                .filter_map(|(name, member)| match member {
                    DocMember::Property(_) => None,
                    DocMember::Function(f) => Some((name, DocItem::Function(f))),
                })
                .collect(),
            _ => HashMap::new(),
        }
    }

    fn cleanup_types(x: &str) -> String {
        x.replace("\\\"int\\\"", "int.type")
            .replace("\\\"bool\\\"", "bool.type")
            .replace("\\\"string\\\"", "str.type")
            .replace("Some(DocType { raw_type: Any })", "None")
            .replace("\\\"_\\\"", "_")
    }

    let expected = expected.documentation().members;
    let got = unpack(DocItem::Module(got.documentation()));
    assert_eq!(expected.len(), got.len());
    for (name, expected1) in expected.iter() {
        let got1 = got.get(name).unwrap();
        assert_eq!(
            cleanup_types(&format!("{:?}", expected1)),
            cleanup_types(&format!("{:?}", got1)),
            "Function {}",
            name
        );
    }
}


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

use crate::assert::Assert;
use crate::tests::bc::golden::bc_golden_test;

#[test]
fn test_fold_list_add() {
    bc_golden_test("constant_folding_list_add", "def test(): return [1] + [2]");
}

#[test]
fn test_fold_list_add_too_large() {
    let mut a = Assert::new();
    a.module("a0.bzl", "def a0(): return [1]");
    for i in 1..100 {
        // `a_17()` is inlined into `a_18()`.
        // If inlining is not limited, this test will run for very long time
        // and eat up all the memory.
        let i_1 = i - 1;
        a.module(
            &format!("a{i}.bzl"),
            &format!(
                "load('a{i_1}.bzl', 'a{i_1}')\n\
                def a{i}(): return a{i_1}() + a{i_1}() + a{i_1}() + a{i_1}() + a{i_1}()"
            ),
        );
    }
}


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

//! Test function bodies inlined.

use crate::assert::Assert;
use crate::eval::bc::opcode::BcOpcode;
use crate::eval::compiler::def::FrozenDef;
use crate::tests::bc::golden::bc_golden_test;
use crate::values::ValueLike;

#[test]
fn test_def_const_inlined() {
    bc_golden_test(
        "def_inline_const_inlined",
        r#"
def trivial():
    return 10

def test():
    return trivial()
"#,
    )
}

#[test]
fn test_def_list_inlined() {
    bc_golden_test(
        "def_inline_list_inlined",
        r#"
def test():
    return returns_list()

# Also test function is inlined if it is defined after the caller.
def returns_list():
    return [10, True]
"#,
    )
}

#[test]
fn test_dict_inlined() {
    bc_golden_test(
        "def_inline_dict_inlined",
        r#"
def returns_dict():
    # This should fail at runtime.
    return {[]: 10}

def test():
    return returns_dict()
"#,
    );
}

#[test]
fn test_do_not_inline_functions_with_return_type() {
    bc_golden_test(
        "def_inline_return_type_inlined",
        r#"
def smth() -> "string":
    return "10"

def test():
    # This call should not be inlined.
    return smth()
"#,
    );
}

#[test]
fn test_dict_inlined_call_stack() {
    let mut a = Assert::new();
    a.module("f.bzl", "def f(): return {[]: 10}");
    // For now inlining doesn't work within one module, so do different modules.
    let m_g = a.module("g.bzl", "load('f.bzl', 'f')\ndef g(): return f()");
    let m_h = a.module("h.bzl", "load('g.bzl', 'g')\ndef h(): return g()");

    // Check `f` is inlined into `g` and `h`.
    for (m, f) in [(m_g, "g"), (m_h, "h")] {
        let f = m.get(f).unwrap();
        let f = f.value().downcast_ref::<FrozenDef>().unwrap();
        assert_eq!(
            BcOpcode::ListNew,
            f.bc().instrs.opcodes().as_slice()[0],
            "in `{}`",
            f,
        );
    }

    let error = a.fail(
        r"
load('h.bzl', 'h')
h()
",
        "",
    );

    assert_eq!(
        r"
Traceback (most recent call last):
  * assert.bzl:3, in <module>
      h()
  * h.bzl.bzl:2, in h
      def h(): return g()
  * g.bzl.bzl:2, in g
      def g(): return f()
error: Value of type `list` is not hashable
 --> f.bzl.bzl:1:18
  |
1 | def f(): return {[]: 10}
  |                  ^^
  |
",
        &format!("\n{:#}", error)
    );
}

#[test]
fn test_do_not_inline_too_large_functions() {
    let mut a = Assert::new();
    a.module("a0.bzl", "def a0(): return noop()");
    for i in 1..100 {
        // `a_17()` is inlined into `a_18()`.
        // If inlining is not limited, this test will run for very long time
        // and eat up all the memory.
        let i_1 = i - 1;
        a.module(
            &format!("a{i}.bzl"),
            &format!(
                "load('a{i_1}.bzl', 'a{i_1}')\n\
                def a{i}(): return a{i_1}() or a{i_1}()"
            ),
        );
    }
}

#[test]
fn test_calls_with_const_args_inlined() {
    bc_golden_test(
        "def_inline_const_args_inlined",
        r#"
def foo(x, y):
    return noop(y, x)

def test():
    # This call should be inlined.
    return foo(10, True)
"#,
    );
}

#[test]
fn test_calls_with_locals_inlined() {
    bc_golden_test(
        "def_inline_locals_inlined",
        r#"
def foo(x, y):
    return noop(y, x)

def test(x, y):
    # This call should be inlined.
    return foo(y, x)
"#,
    );
}


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

//! Permutations tests for if condition evaluation.

use std::cell::Cell;
use std::fmt;
use std::fmt::Display;

use derive_more::Display;
use dupe::Dupe;
use rand::rngs::SmallRng;
use rand::Rng;
use rand::SeedableRng;
use starlark_derive::starlark_module;

use crate as starlark;
use crate::any::ProvidesStaticType;
use crate::environment::GlobalsBuilder;
use crate::environment::Module;
use crate::eval::Evaluator;
use crate::syntax::AstModule;
use crate::syntax::Dialect;

/// Count side effects. For example, in expression like:
///
/// ```text
/// false() and true()
/// ```
///
/// `true()` should not be evaluated. After evaluation, the counter should be 1.
#[derive(Debug, ProvidesStaticType, Default, PartialEq)]
struct CountCalls {
    calls: Cell<usize>,
}

#[starlark_module]
fn bool_fns(globals: &mut GlobalsBuilder) {
    /// Return `true` and record side effect.
    fn r#true(eval: &mut Evaluator) -> anyhow::Result<bool> {
        let calls = eval
            .extra
            .as_ref()
            .unwrap()
            .downcast_ref::<CountCalls>()
            .unwrap();
        calls.calls.set(calls.calls.get() + 1);
        Ok(true)
    }

    fn r#false(eval: &mut Evaluator) -> anyhow::Result<bool> {
        let calls = eval
            .extra
            .as_ref()
            .unwrap()
            .downcast_ref::<CountCalls>()
            .unwrap();
        calls.calls.set(calls.calls.get() + 1);
        Ok(false)
    }
}

#[derive(Display, Debug, Copy, Clone, Dupe)]
enum TestBinOp {
    #[display(fmt = "and")]
    And,
    #[display(fmt = "or")]
    Or,
}

impl TestBinOp {
    fn eval(self, x: bool, y: impl FnOnce() -> bool) -> bool {
        match self {
            TestBinOp::And => x && y(),
            TestBinOp::Or => x || y(),
        }
    }
}

#[derive(Clone, Debug)]
enum TestExpr {
    /// `True` or `False`.
    Const(bool),
    /// `true()` or `false()`.
    Count(bool),
    /// Binary operation.
    BinOp(TestBinOp, Box<(TestExpr, TestExpr)>),
    /// `not` operation.
    Not(Box<TestExpr>),
}

impl TestExpr {
    /// Evaluate the expression the same way Starlark would evaluate it.
    fn eval(&self, count: &CountCalls) -> bool {
        match self {
            TestExpr::Const(x) => *x,
            TestExpr::Count(x) => {
                // Record side effect.
                count.calls.set(count.calls.get() + 1);
                *x
            }
            TestExpr::BinOp(op, x_y) => {
                let (x, y) = &**x_y;
                op.eval(x.eval(count), || y.eval(count))
            }
            TestExpr::Not(x) => !x.eval(count),
        }
    }
}

impl Display for TestExpr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            TestExpr::Const(x) => match x {
                true => write!(f, "True"),
                false => write!(f, "False"),
            },
            TestExpr::Count(x) => match x {
                true => write!(f, "true()"),
                false => write!(f, "false()"),
            },
            TestExpr::BinOp(op, x_y) => {
                let (x, y) = &**x_y;
                write!(f, "({} {} {})", x, op, y)
            }
            TestExpr::Not(x) => write!(f, "(not {})", x),
        }
    }
}

/// Evaluate the program.
///
/// * Return the result which is expected to be `bool`.
/// * Count side effects.
fn eval_program(program: &str) -> (bool, CountCalls) {
    let module = Module::new();
    let ast = AstModule::parse("t.star", program.to_owned(), &Dialect::Extended).unwrap();

    let mut globals = GlobalsBuilder::standard();
    bool_fns(&mut globals);
    let globals = globals.build();

    let counts = CountCalls::default();
    let r = {
        let mut eval = Evaluator::new(&module);
        eval.extra = Some(&counts);
        let r = eval.eval_module(ast, &globals).unwrap();
        r.unpack_bool().unwrap()
    };
    (r, counts)
}

fn eval_if_else_with_starlark(expr: &TestExpr) -> (bool, CountCalls) {
    let program = format!(
        r#"
if {}:
    r = True
else:
    r = False
r
"#,
        expr
    );
    eval_program(&program)
}

fn eval_if_with_starlark(expr: &TestExpr) -> (bool, CountCalls) {
    let program = format!(
        r#"
r = False
if {}:
    r = True
r
"#,
        expr
    );
    eval_program(&program)
}

fn eval_expr_result(expr: &TestExpr) -> (bool, CountCalls) {
    eval_program(&expr.to_string())
}

fn eval_manually(expr: &TestExpr) -> (bool, CountCalls) {
    let counts = CountCalls::default();
    let r = expr.eval(&counts);
    (r, counts)
}

fn test_if_else(expr: &TestExpr) {
    let expected = eval_manually(expr);
    let actual = eval_if_else_with_starlark(expr);
    assert_eq!(expected, actual, "expression: {}", expr);
}

fn test_if(expr: &TestExpr) {
    let expected = eval_manually(expr);
    let actual = eval_if_with_starlark(expr);
    assert_eq!(expected, actual, "expression: {}", expr);
}

fn test_expr_result(expr: &TestExpr) {
    let expected = eval_manually(expr);
    let actual = eval_expr_result(expr);
    assert_eq!(expected, actual, "expression: {}", expr);
}

fn test_ifs(expr: &TestExpr) {
    test_if(expr);
    test_if_else(expr);
    // If condition expression compilation is different from compilation of the expression,
    // so we explicitly test both cases.
    test_expr_result(expr);
}

fn bool_values() -> [bool; 2] {
    [true, false]
}

fn basic_bool_exprs() -> impl Iterator<Item = TestExpr> {
    bool_values()
        .into_iter()
        .flat_map(|x| [TestExpr::Count(x), TestExpr::Const(x)])
}

#[test]
fn test_basic() {
    test_ifs(&TestExpr::Const(true));
    test_ifs(&TestExpr::Const(false));
    test_ifs(&TestExpr::Count(true));
    test_ifs(&TestExpr::Count(false));
    test_ifs(&TestExpr::Not(Box::new(TestExpr::Const(true))));
    test_ifs(&TestExpr::Not(Box::new(TestExpr::Const(false))));
    test_ifs(&TestExpr::Not(Box::new(TestExpr::Count(true))));
    test_ifs(&TestExpr::Not(Box::new(TestExpr::Count(false))));
}

#[test]
fn test_and() {
    for lhs in basic_bool_exprs() {
        for rhs in basic_bool_exprs() {
            test_ifs(&TestExpr::BinOp(
                TestBinOp::And,
                Box::new((lhs.clone(), rhs)),
            ));
        }
    }
}

#[test]
fn test_or() {
    for lhs in basic_bool_exprs() {
        for rhs in basic_bool_exprs() {
            test_ifs(&TestExpr::BinOp(
                TestBinOp::Or,
                Box::new((lhs.clone(), rhs)),
            ));
        }
    }
}

#[test]
fn test_and_or_not() {
    for lhs in basic_bool_exprs() {
        for rhs in basic_bool_exprs() {
            for negate_lhs in [false, true] {
                for negate_rhs in [false, true] {
                    for bin_op in [TestBinOp::And, TestBinOp::Or] {
                        let lhs = if negate_lhs {
                            TestExpr::Not(Box::new(lhs.clone()))
                        } else {
                            lhs.clone()
                        };
                        let rhs = if negate_rhs {
                            TestExpr::Not(Box::new(rhs.clone()))
                        } else {
                            rhs.clone()
                        };
                        test_ifs(&TestExpr::BinOp(bin_op, Box::new((lhs, rhs))));
                    }
                }
            }
        }
    }
}

const RANDOM_ITERATIONS: usize = 100;

fn max_depth_for_iter(i: usize) -> usize {
    if i < 5 {
        0
    } else if i < RANDOM_ITERATIONS / 50 {
        1
    } else if i < RANDOM_ITERATIONS / 25 {
        2
    } else if i < RANDOM_ITERATIONS / 10 {
        3
    } else if i < RANDOM_ITERATIONS / 3 {
        4
    } else if i < RANDOM_ITERATIONS / 2 {
        5
    } else {
        20
    }
}

fn random_expr(rng: &mut SmallRng, max_depth: usize) -> TestExpr {
    fn random_simple_expr(rng: &mut SmallRng) -> TestExpr {
        match rng.gen_range(0..4) {
            0 => TestExpr::Const(true),
            1 => TestExpr::Const(false),
            2 => TestExpr::Count(true),
            3 => TestExpr::Count(false),
            _ => unreachable!(),
        }
    }

    if max_depth == 0 {
        random_simple_expr(rng)
    } else {
        match rng.gen_range(0..4) {
            0 => random_simple_expr(rng),
            1 => TestExpr::Not(Box::new(random_expr(rng, max_depth - 1))),
            2 => TestExpr::BinOp(
                TestBinOp::And,
                Box::new((
                    random_expr(rng, max_depth - 1),
                    random_expr(rng, max_depth - 1),
                )),
            ),
            3 => TestExpr::BinOp(
                TestBinOp::Or,
                Box::new((
                    random_expr(rng, max_depth - 1),
                    random_expr(rng, max_depth - 1),
                )),
            ),
            _ => unreachable!(),
        }
    }
}

#[test]
fn test_if_random() {
    let mut rng = SmallRng::seed_from_u64(17);
    for i in 0..RANDOM_ITERATIONS {
        let max_depth = max_depth_for_iter(i);
        let expr = random_expr(&mut rng, max_depth);
        test_if(&expr);
    }
}

#[test]
fn test_if_else_random() {
    let mut rng = SmallRng::seed_from_u64(17);
    for i in 0..RANDOM_ITERATIONS {
        let max_depth = max_depth_for_iter(i);
        let expr = random_expr(&mut rng, max_depth);
        test_if_else(&expr);
    }
}

#[test]
fn test_expr_random() {
    let mut rng = SmallRng::seed_from_u64(17);
    for i in 0..RANDOM_ITERATIONS {
        let max_depth = max_depth_for_iter(i);
        let expr = random_expr(&mut rng, max_depth);
        test_expr_result(&expr);
    }
}


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

//! Optimizer tests.

mod constant_folding;
mod def_inline;
mod if_rand;
mod speculative_exec;
mod type_is;
mod types;

use crate::tests::bc::golden::bc_golden_test;

#[test]
fn test_type_is_inlined() {
    bc_golden_test(
        "opt_type_is_inlined",
        r#"
def is_list(x):
    return type(x) == type([])

def test(x):
    return is_list(x)
    "#,
    )
}

#[test]
fn test_private_forward_mutable_module_vars_inlined() {
    bc_golden_test(
        "opt_private_forward_mutable_module_vars_inlined",
        r#"
def test():
    # Reference to module variable should be replaced with constant
    return _private_forward_mutable

_private_forward_mutable = {1: 2}
"#,
    );
}

#[test]
fn test_same_module_struct_getattr_inlined() {
    bc_golden_test(
        "opt_same_module_struct_getattr_inlined",
        r#"
def test():
    return _s.f

_s = struct(f = 1)
"#,
    );
}

#[test]
fn test_list_plus_list() {
    bc_golden_test(
        "opt_list_plus_list",
        r#"
L = [1, 2]

def test():
    return L + [1]
"#,
    );
}

#[test]
fn test_empty_iterable_optimized_away() {
    bc_golden_test(
        "opt_empty_iterable_optimized_away",
        r#"
L = []
def test():
    for x in L:
        print(x)
"#,
    );
}

#[test]
fn test_unreachable_code_optimized_away() {
    bc_golden_test(
        "opt_unreachable_code_optimized_away",
        r#"
def test():
    if True:
        return
    fail("unreachable")
"#,
    );
}

#[test]
fn test_recursion() {
    bc_golden_test(
        "opt_recursion",
        // Test inlining does not fail here.
        "def test(): return test()",
    );
}

#[test]
fn test_mutual_recursion() {
    // Just check we do not enter an infinite recursion in the optimizer here.
    bc_golden_test(
        "opt_mutual_recursion",
        r#"
def test():
    return g()

def g():
    return test()
"#,
    );
}


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

use crate::tests::bc::golden::bc_golden_test;

#[test]
fn test_methods_invoked_speculatively() {
    bc_golden_test(
        "speculative_exec_methods_invoked_speculatively",
        r#"
def test():
    return "foo".startswith("f")
"#,
    )
}

#[test]
fn test_format_speculatively_before_format_instr() {
    bc_golden_test(
        "speculative_exec_format_speculatively_before_format_instr",
        r#"
def test():
    # Test this expression is compiled to constant, not to `FormatOne` instruction.
    return "x{}y".format(1)
"#,
    );
}

#[test]
fn test_speculatively_inline_enum() {
    bc_golden_test(
        "speculative_exec_enum_inline",
        r#"
MyEnum = enum("red", "green", "blue")

def test(x):
    # Test there is no enum evaluation.
    return x == MyEnum("red")
"#,
    );
}


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

use crate::tests::bc::golden::bc_golden_test;

#[test]
fn test_return_type_erased() {
    // Return type is checked in bytecode.
    // In this example, return type is anything.
    // Check it is not present in bytecode.
    bc_golden_test("opt_return_type_erased", "def test() -> '_a': return 1");
}

#[test]
fn test_return_type_present() {
    // Self-check for previous test.
    bc_golden_test(
        "opt_return_type_present",
        "def test() -> 'string': return 'x'",
    );
}


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

//! Test for type-is optimizations.

use starlark_derive::starlark_module;

use crate as starlark;
use crate::assert::Assert;
use crate::environment::GlobalsBuilder;
use crate::eval::compiler::def::Def;
use crate::eval::compiler::def::FrozenDef;
use crate::eval::compiler::def_inline::InlineDefBody;
use crate::values::Value;
use crate::values::ValueLike;

#[starlark_module]
fn globals(builder: &mut GlobalsBuilder) {
    fn returns_type_is<'v>(value: Value<'v>) -> anyhow::Result<bool> {
        Ok(if let Some(def) = value.downcast_ref::<FrozenDef>() {
            matches!(
                def.def_info.inline_def_body,
                Some(InlineDefBody::ReturnTypeIs(..))
            )
        } else if let Some(def) = value.downcast_ref::<Def>() {
            matches!(
                def.def_info.inline_def_body,
                Some(InlineDefBody::ReturnTypeIs(..))
            )
        } else {
            panic!("not def")
        })
    }
}

#[test]
fn returns_type_is() {
    let mut a = Assert::new();
    a.globals_add(globals);

    a.module(
        "types.star",
        "\
def is_list(x):
  return type(x) == type([])
",
    );

    a.pass(
        "\
load('types.star', 'is_list')
assert_true(returns_type_is(is_list))
assert_true(is_list([]))
assert_false(is_list({}))
    ",
    );
}

#[test]
fn does_not_return_type_is() {
    let mut a = Assert::new();
    a.globals_add(globals);
    a.pass(
        "\
def is_not_list(x):
  return type(x) != type([])

def something_else(x, y):
  return type(x) == type([])

assert_false(returns_type_is(is_not_list))
assert_false(returns_type_is(something_else))
    ",
    );
}


/*
 * 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::collections::HashMap;
use std::sync::Arc;

use dupe::Dupe;

use crate::codemap::CodeMap;
use crate::codemap::Span;
use crate::codemap::Spanned;
use crate::eval::compiler::scope::payload::CstAssign;
use crate::eval::compiler::scope::payload::CstExpr;
use crate::eval::compiler::scope::payload::CstPayload;
use crate::eval::compiler::scope::payload::CstStmt;
use crate::eval::compiler::scope::BindingId;
use crate::eval::compiler::scope::ResolvedIdent;
use crate::syntax::ast::AssignOp;
use crate::syntax::ast::AssignP;
use crate::syntax::ast::ClauseP;
use crate::syntax::ast::DefP;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::ForClauseP;
use crate::syntax::ast::IdentP;
use crate::syntax::ast::ParameterP;
use crate::syntax::ast::StmtP;
use crate::syntax::uniplate::Visit;
use crate::typing::error::InternalError;
use crate::typing::function::Param;
use crate::typing::mode::TypecheckMode;
use crate::typing::ty::Approximation;
use crate::typing::ty::Ty;

#[derive(Clone)]
pub(crate) enum BindExpr<'a> {
    Expr(&'a CstExpr),
    /// Get this position from the expression
    GetIndex(usize, Box<BindExpr<'a>>),
    Iter(Box<BindExpr<'a>>),
    AssignOp(&'a CstAssign, AssignOp, &'a CstExpr),
    /// Set this index in the variable
    SetIndex(BindingId, &'a CstExpr, Box<BindExpr<'a>>),
    ListAppend(BindingId, &'a CstExpr),
    ListExtend(BindingId, &'a CstExpr),
}

impl<'a> BindExpr<'a> {
    pub(crate) fn span(&self) -> Span {
        match self {
            BindExpr::Expr(x) => x.span,
            BindExpr::GetIndex(_, x) => x.span(),
            BindExpr::Iter(x) => x.span(),
            BindExpr::AssignOp(x, _, _) => x.span,
            BindExpr::SetIndex(_, x, _) => x.span,
            BindExpr::ListAppend(_, x) => x.span,
            BindExpr::ListExtend(_, x) => x.span,
        }
    }
}

#[derive(Default)]
pub(crate) struct Bindings<'a> {
    pub(crate) expressions: HashMap<BindingId, Vec<BindExpr<'a>>>,
    pub(crate) types: HashMap<BindingId, Ty>,
    pub(crate) check: Vec<&'a CstExpr>,
    pub(crate) check_type: Vec<(Span, Option<&'a CstExpr>, Ty)>,
}

/// Interface representing the types of all bindings in a module.
#[derive(Default, Dupe, Clone, Debug)]
pub struct Interface(Arc<HashMap<String, Ty>>);

impl Interface {
    /// Create an empty interface, with no bindings.
    pub fn empty() -> Self {
        Self::default()
    }

    /// Create a new interface with the given bindings.
    pub fn new(bindings: HashMap<String, Ty>) -> Self {
        Self(Arc::new(bindings))
    }

    /// Get the type for a given binding.
    pub fn get(&self, name: &str) -> Option<&Ty> {
        self.0.get(name)
    }
}

#[derive(Default)]
pub(crate) struct BindingsCollect<'a> {
    pub(crate) bindings: Bindings<'a>,
    pub(crate) approximations: Vec<Approximation>,
}

impl<'a> BindingsCollect<'a> {
    /// Collect all the assignments to variables.
    ///
    /// This function only fails on internal errors.
    pub(crate) fn collect(
        xs: &'a [CstStmt],
        typecheck_mode: TypecheckMode,
        codemap: &CodeMap,
    ) -> Result<Self, InternalError> {
        fn assign<'a>(
            lhs: &'a CstAssign,
            rhs: BindExpr<'a>,
            bindings: &mut BindingsCollect<'a>,
            codemap: &CodeMap,
        ) -> Result<(), InternalError> {
            match &**lhs {
                AssignP::Identifier(x) => {
                    bindings
                        .bindings
                        .expressions
                        .entry(x.resolved_binding_id(codemap)?)
                        .or_default()
                        .push(rhs);
                }
                AssignP::Tuple(xs) => {
                    for (i, x) in xs.iter().enumerate() {
                        assign(
                            x,
                            BindExpr::GetIndex(i, Box::new(rhs.clone())),
                            bindings,
                            codemap,
                        )?;
                    }
                }
                AssignP::Index(array_index) => match &*array_index.0 {
                    ExprP::Identifier(Spanned {
                        span: _,
                        node: IdentP(_name, Some(ResolvedIdent::Slot(_, ident))),
                    }) => {
                        bindings
                            .bindings
                            .expressions
                            .entry(*ident)
                            .or_default()
                            .push(BindExpr::SetIndex(*ident, &array_index.1, Box::new(rhs)));
                    }
                    _ => {
                        bindings.approximations.push(Approximation::new(
                            "Underapproximation",
                            "a.b[x] = .. not handled",
                        ));
                    }
                },
                AssignP::Dot(_, _) => {
                    bindings.approximations.push(Approximation::new(
                        "Underapproximation",
                        "a.b = .. not handled",
                    ));
                }
            }
            Ok(())
        }

        fn visit<'a>(
            x: Visit<'a, CstPayload>,
            return_type: &Ty,
            bindings: &mut BindingsCollect<'a>,
            typecheck_mode: TypecheckMode,
            codemap: &CodeMap,
        ) -> Result<(), InternalError> {
            match x {
                Visit::Stmt(x) => match &**x {
                    StmtP::Assign(lhs, ty_rhs) => {
                        if let Some(ty) = &ty_rhs.0 {
                            let ty2 = Ty::from_type_expr(
                                ty,
                                typecheck_mode,
                                &mut bindings.approximations,
                                codemap,
                            )?;
                            bindings.bindings.check_type.push((
                                ty.span,
                                Some(&ty_rhs.1),
                                ty2.clone(),
                            ));
                            if let AssignP::Identifier(id) = &**lhs {
                                // FIXME: This could be duplicated if you declare the type of a variable twice,
                                // we would only see the second one.
                                bindings
                                    .bindings
                                    .types
                                    .insert(id.resolved_binding_id(codemap)?, ty2);
                            }
                        }
                        assign(lhs, BindExpr::Expr(&ty_rhs.1), bindings, codemap)?
                    }
                    StmtP::AssignModify(lhs, op, rhs) => {
                        assign(lhs, BindExpr::AssignOp(lhs, *op, rhs), bindings, codemap)?
                    }
                    StmtP::For(lhs, iter_body) => assign(
                        lhs,
                        BindExpr::Iter(Box::new(BindExpr::Expr(&iter_body.0))),
                        bindings,
                        codemap,
                    )?,
                    StmtP::Def(DefP {
                        name,
                        params,
                        return_type,
                        ..
                    }) => {
                        let mut params2 = Vec::with_capacity(params.len());
                        let mut seen_no_args = false;
                        for p in params {
                            let name_ty = match &**p {
                                ParameterP::Normal(name, ty)
                                | ParameterP::WithDefaultValue(name, ty, _) => {
                                    let ty = Ty::from_type_expr_opt(
                                        ty,
                                        typecheck_mode,
                                        &mut bindings.approximations,
                                        codemap,
                                    )?;
                                    let mut param = if seen_no_args {
                                        Param::name_only(&name.0, ty.clone())
                                    } else {
                                        Param::pos_or_name(&name.0, ty.clone())
                                    };
                                    if matches!(&**p, ParameterP::WithDefaultValue(..)) {
                                        param = param.optional();
                                    }
                                    params2.push(param);
                                    Some((name, ty))
                                }
                                ParameterP::NoArgs => {
                                    seen_no_args = true;
                                    None
                                }
                                ParameterP::Args(name, ty) => {
                                    // There is the type we require people calling us use (usually any)
                                    // and then separately the type we are when we are running (always tuple)
                                    params2.push(Param::args(Ty::from_type_expr_opt(
                                        ty,
                                        typecheck_mode,
                                        &mut bindings.approximations,
                                        codemap,
                                    )?));
                                    Some((name, Ty::name("tuple")))
                                }
                                ParameterP::KwArgs(name, ty) => {
                                    let ty = Ty::from_type_expr_opt(
                                        ty,
                                        typecheck_mode,
                                        &mut bindings.approximations,
                                        codemap,
                                    )?;
                                    let ty = if ty.is_any() {
                                        Ty::dict(Ty::string(), Ty::Any)
                                    } else {
                                        ty
                                    };
                                    params2.push(Param::kwargs(ty.clone()));
                                    Some((name, ty))
                                }
                            };
                            if let Some((name, ty)) = name_ty {
                                bindings
                                    .bindings
                                    .types
                                    .insert(name.resolved_binding_id(codemap)?, ty);
                            }
                        }
                        let ret_ty = Ty::from_type_expr_opt(
                            return_type,
                            typecheck_mode,
                            &mut bindings.approximations,
                            codemap,
                        )?;
                        bindings.bindings.types.insert(
                            name.resolved_binding_id(codemap)?,
                            Ty::function(params2, ret_ty.clone()),
                        );
                        x.visit_children_err(|x| {
                            visit(x, &ret_ty, bindings, typecheck_mode, codemap)
                        })?;
                        // We do our own visit_children, with a different return type
                        return Ok(());
                    }
                    StmtP::Load(x) => {
                        let mp = &x.payload;
                        for (ident, _load) in &x.args {
                            let ty = mp.get(ident.0.as_str()).cloned().unwrap_or(Ty::Any);
                            bindings
                                .bindings
                                .types
                                .insert(ident.resolved_binding_id(codemap)?, ty);
                        }
                    }
                    StmtP::Return(ret) => bindings.bindings.check_type.push((
                        x.span,
                        ret.as_ref(),
                        return_type.clone(),
                    )),
                    StmtP::Expression(x) => {
                        // We want to find ident.append(), ident.extend(), ident.extend()
                        // to fake up a BindExpr::ListAppend/ListExtend
                        // so that mutating list operations aren't invisible to us
                        if let ExprP::Call(fun, args) = &**x {
                            if let ExprP::Dot(id, attr) = &***fun {
                                if let ExprP::Identifier(id) = &id.node {
                                    let res = match attr.as_str() {
                                        "append" if args.len() == 1 => Some((false, 0)),
                                        "insert" if args.len() == 2 => Some((false, 1)),
                                        "extend" if args.len() == 1 => Some((true, 0)),
                                        _ => None,
                                    };
                                    if let Some((extend, arg)) = res {
                                        if let ResolvedIdent::Slot(_, id) =
                                            id.node.1.as_ref().unwrap()
                                        {
                                            let bind = if extend {
                                                BindExpr::ListExtend(*id, args[arg].expr())
                                            } else {
                                                BindExpr::ListAppend(*id, args[arg].expr())
                                            };
                                            bindings
                                                .bindings
                                                .expressions
                                                .entry(*id)
                                                .or_default()
                                                .push(bind)
                                        }
                                    }
                                }
                            }
                        }

                        bindings.bindings.check.push(x)
                    }
                    StmtP::If(x, _) => bindings.bindings.check.push(x),
                    StmtP::IfElse(x, _) => bindings.bindings.check.push(x),
                    _ => {}
                },
                Visit::Expr(x) => match &**x {
                    ExprP::ListComprehension(_, for1, clauses)
                    | ExprP::DictComprehension(_, for1, clauses) => {
                        fn get_for_clause(
                            x: &ClauseP<CstPayload>,
                        ) -> Option<&ForClauseP<CstPayload>> {
                            match x {
                                ClauseP::For(x) => Some(x),
                                _ => None,
                            }
                        }
                        for x in std::iter::once(&**for1)
                            .chain(clauses.iter().filter_map(get_for_clause))
                        {
                            assign(
                                &x.var,
                                BindExpr::Iter(Box::new(BindExpr::Expr(&x.over))),
                                bindings,
                                codemap,
                            )?
                        }
                    }
                    _ => {}
                },
            }
            x.visit_children_err(|x| visit(x, return_type, bindings, typecheck_mode, codemap))?;
            Ok(())
        }

        let mut res = BindingsCollect::default();
        for x in xs {
            visit(Visit::Stmt(x), &Ty::Any, &mut res, typecheck_mode, codemap)?;
        }
        Ok(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.
 */

use std::cell::RefCell;
use std::collections::HashMap;
use std::fmt::Debug;

use thiserror::Error;

use crate::codemap::Span;
use crate::codemap::Spanned;
use crate::eval::compiler::scope::payload::CstAssign;
use crate::eval::compiler::scope::payload::CstExpr;
use crate::eval::compiler::scope::payload::CstPayload;
use crate::eval::compiler::scope::BindingId;
use crate::eval::compiler::scope::ResolvedIdent;
use crate::slice_vec_ext::SliceExt;
use crate::slice_vec_ext::VecExt;
use crate::syntax::ast::ArgumentP;
use crate::syntax::ast::AssignOp;
use crate::syntax::ast::AssignP;
use crate::syntax::ast::AstLiteral;
use crate::syntax::ast::BinOp;
use crate::syntax::ast::ClauseP;
use crate::syntax::ast::ExprP;
use crate::syntax::ast::ForClauseP;
use crate::typing::bindings::BindExpr;
use crate::typing::error::TypingError;
use crate::typing::function::Arg;
use crate::typing::oracle::ctx::TypingOracleCtx;
use crate::typing::oracle::traits::TypingAttr;
use crate::typing::oracle::traits::TypingBinOp;
use crate::typing::oracle::traits::TypingUnOp;
use crate::typing::ty::Approximation;
use crate::typing::ty::Ty;
use crate::typing::OracleDocs;

#[derive(Error, Debug)]
enum TypingContextError {
    #[error("The attribute `{attr}` is not available on the type `{typ}`")]
    AttributeNotAvailable { typ: String, attr: String },
    #[error("The builtin `{name}` is not known")]
    UnknownBuiltin { name: String },
}

pub(crate) struct TypingContext<'a> {
    pub(crate) oracle: TypingOracleCtx<'a>,
    pub(crate) global_docs: OracleDocs,
    // We'd prefer this to be a &mut self,
    // but that makes writing the code more fiddly, so just RefCell the errors
    pub(crate) errors: RefCell<Vec<TypingError>>,
    pub(crate) approximoations: RefCell<Vec<Approximation>>,
    pub(crate) types: HashMap<BindingId, Ty>,
}

impl TypingContext<'_> {
    fn add_error(&self, span: Span, err: TypingContextError) -> Ty {
        let err = self.oracle.mk_error(span, err);
        self.errors.borrow_mut().push(err);
        Ty::Never
    }

    pub(crate) fn approximation(&self, category: &'static str, message: impl Debug) -> Ty {
        self.approximoations
            .borrow_mut()
            .push(Approximation::new(category, message));
        Ty::Any
    }

    fn validate_call(&self, fun: &Ty, args: &[Spanned<Arg>], span: Span) -> Ty {
        match self.oracle.validate_call(span, fun, args) {
            Ok(ty) => ty,
            Err(e) => {
                self.errors.borrow_mut().push(e);
                Ty::Never
            }
        }
    }

    fn from_iterated(&self, ty: &Ty, span: Span) -> Ty {
        self.expression_attribute(ty, TypingAttr::Iter, span)
    }

    pub(crate) fn validate_type(&self, got: &Ty, require: &Ty, span: Span) {
        if let Err(e) = self.oracle.validate_type(got, require, span) {
            self.errors.borrow_mut().push(e);
        }
    }

    fn builtin(&self, name: &str, span: Span) -> Ty {
        match self.global_docs.builtin(name) {
            Ok(x) => x,
            Err(()) => self.add_error(
                span,
                TypingContextError::UnknownBuiltin {
                    name: name.to_owned(),
                },
            ),
        }
    }

    fn expression_attribute(&self, ty: &Ty, attr: TypingAttr, span: Span) -> Ty {
        match ty.attribute(attr, self) {
            Ok(x) => x,
            Err(()) => self.add_error(
                span,
                TypingContextError::AttributeNotAvailable {
                    typ: ty.to_string(),
                    attr: attr.to_string(),
                },
            ),
        }
    }

    fn expression_primitive_ty(
        &self,
        name: TypingAttr,
        arg0: Ty,
        args: Vec<Spanned<Ty>>,
        span: Span,
    ) -> Ty {
        let fun = self.expression_attribute(&arg0, name, span);
        self.validate_call(&fun, &args.into_map(|x| x.into_map(Arg::Pos)), span)
    }

    fn expression_primitive(&self, name: TypingAttr, args: &[&CstExpr], span: Span) -> Ty {
        let t0 = self.expression_type(args[0]);
        let ts = args[1..].map(|x| Spanned {
            span: x.span,
            node: self.expression_type(x),
        });
        self.expression_primitive_ty(name, t0, ts, span)
    }

    pub(crate) fn expression_bind_type(&self, x: &BindExpr) -> Ty {
        match x {
            BindExpr::Expr(x) => self.expression_type(x),
            BindExpr::GetIndex(i, x) => self.expression_bind_type(x).indexed(*i),
            BindExpr::Iter(x) => self.from_iterated(&self.expression_bind_type(x), x.span()),
            BindExpr::AssignOp(lhs, op, rhs) => {
                let span = lhs.span;
                let rhs_ty = self.expression_type(rhs);
                let lhs = self.expression_assign(lhs);
                let attr = match op {
                    AssignOp::Add => TypingAttr::BinOp(TypingBinOp::Add),
                    AssignOp::Subtract => TypingAttr::BinOp(TypingBinOp::Sub),
                    AssignOp::Multiply => TypingAttr::BinOp(TypingBinOp::Mul),
                    AssignOp::Divide => TypingAttr::BinOp(TypingBinOp::Div),
                    AssignOp::FloorDivide => TypingAttr::BinOp(TypingBinOp::FloorDiv),
                    AssignOp::Percent => TypingAttr::BinOp(TypingBinOp::Percent),
                    AssignOp::BitAnd => TypingAttr::BinOp(TypingBinOp::BitAnd),
                    AssignOp::BitOr => TypingAttr::BinOp(TypingBinOp::BitOr),
                    AssignOp::BitXor => TypingAttr::BinOp(TypingBinOp::BitXor),
                    AssignOp::LeftShift => TypingAttr::BinOp(TypingBinOp::LeftShift),
                    AssignOp::RightShift => TypingAttr::BinOp(TypingBinOp::RightShift),
                };
                self.expression_primitive_ty(
                    attr,
                    lhs,
                    vec![Spanned {
                        span: rhs.span,
                        node: rhs_ty,
                    }],
                    span,
                )
            }
            BindExpr::SetIndex(id, index, e) => {
                let span = index.span;
                let index = self.expression_type(index);
                let e = self.expression_bind_type(e);
                let mut res = Vec::new();
                // We know about list and dict, everything else we just ignore
                if self.types[id].is_list() {
                    // If we know it MUST be a list, then the index must be an int
                    self.validate_type(&index, &Ty::int(), span);
                }
                for ty in self.types[id].iter_union() {
                    match ty {
                        Ty::List(_) => {
                            res.push(Ty::list(e.clone()));
                        }
                        Ty::Dict(_) => {
                            res.push(Ty::dict(index.clone(), e.clone()));
                        }
                        _ => {
                            // Either it's not something we can apply this to, in which case do nothing.
                            // Or it's an Any, in which case we aren't going to change its type or spot errors.
                        }
                    }
                }
                Ty::unions(res)
            }
            BindExpr::ListAppend(id, e) => {
                if Ty::probably_a_list(&self.types[id], self.oracle) {
                    Ty::list(self.expression_type(e))
                } else {
                    // It doesn't seem to be a list, so let's assume the append is non-mutating
                    Ty::Never
                }
            }
            BindExpr::ListExtend(id, e) => {
                if Ty::probably_a_list(&self.types[id], self.oracle) {
                    Ty::list(self.from_iterated(&self.expression_type(e), e.span))
                } else {
                    // It doesn't seem to be a list, so let's assume the extend is non-mutating
                    Ty::Never
                }
            }
        }
    }

    /// Used to get the type of an expression when used as part of a ModifyAssign operation
    fn expression_assign(&self, x: &CstAssign) -> Ty {
        match &**x {
            AssignP::Tuple(_) => self.approximation("expression_assignment", x),
            AssignP::Index(a_b) => {
                self.expression_primitive(TypingAttr::Index, &[&a_b.0, &a_b.1], x.span)
            }
            AssignP::Dot(_, _) => self.approximation("expression_assignment", x),
            AssignP::Identifier(x) => {
                if let Some(i) = x.1 {
                    if let Some(ty) = self.types.get(&i) {
                        return ty.clone();
                    }
                }
                panic!("Unknown identifier")
            }
        }
    }

    /// We don't need the type out of the clauses (it doesn't change the overall type),
    /// but it is important we see through to the nested expressions to raise errors
    fn check_comprehension(&self, for_: &ForClauseP<CstPayload>, clauses: &[ClauseP<CstPayload>]) {
        self.expression_type(&for_.over);
        for x in clauses {
            match x {
                ClauseP::For(x) => self.expression_type(&x.over),
                ClauseP::If(x) => self.expression_type(x),
            };
        }
    }

    pub(crate) fn expression_type_spanned(&self, x: &CstExpr) -> Spanned<Ty> {
        Spanned {
            span: x.span,
            node: self.expression_type(x),
        }
    }

    pub(crate) fn expression_type(&self, x: &CstExpr) -> Ty {
        let span = x.span;
        match &**x {
            ExprP::Tuple(xs) => Ty::Tuple(xs.map(|x| self.expression_type(x))),
            ExprP::Dot(a, b) => {
                self.expression_attribute(&self.expression_type(a), TypingAttr::Regular(b), b.span)
            }
            ExprP::Call(f, args) => {
                let args_ty: Vec<Spanned<Arg>> = args.map(|x| Spanned {
                    span: x.span,
                    node: match &**x {
                        ArgumentP::Positional(x) => Arg::Pos(self.expression_type(x)),
                        ArgumentP::Named(name, x) => {
                            Arg::Name((**name).clone(), self.expression_type(x))
                        }
                        ArgumentP::Args(x) => {
                            let ty = self.expression_type(x);
                            self.from_iterated(&ty, x.span);
                            Arg::Args(ty)
                        }
                        ArgumentP::KwArgs(x) => {
                            let ty = self.expression_type(x);
                            self.validate_type(&ty, &Ty::dict(Ty::string(), Ty::Any), x.span);
                            Arg::Kwargs(ty)
                        }
                    },
                });
                let f_ty = self.expression_type(f);
                // If we can't resolve the types of the arguments, we can't validate the call,
                // but we still know the type of the result since the args don't impact that
                self.validate_call(&f_ty, &args_ty, span)
            }
            ExprP::Index(a_b) => {
                self.expression_primitive(TypingAttr::Index, &[&a_b.0, &a_b.1], span)
            }
            ExprP::Index2(a_i0_i1) => {
                let (a, i0, i1) = &**a_i0_i1;
                self.expression_type(a);
                self.expression_type(i0);
                self.expression_type(i1);
                Ty::Any
            }
            ExprP::Slice(x, start, stop, stride) => {
                for e in [start, stop, stride].iter().copied().flatten() {
                    self.validate_type(&self.expression_type(e), &Ty::int(), e.span);
                }
                self.expression_attribute(&self.expression_type(x), TypingAttr::Slice, span)
            }
            ExprP::Identifier(x) => {
                match &x.node.1 {
                    Some(ResolvedIdent::Slot(_, i)) => {
                        if let Some(ty) = self.types.get(i) {
                            ty.clone()
                        } else {
                            // All types must be resolved to this point,
                            // this code is unreachable.
                            Ty::Any
                        }
                    }
                    Some(ResolvedIdent::Global(g)) => {
                        if let Some(t) = g.to_value().get_ref().typechecker_ty() {
                            t
                        } else {
                            self.builtin(&x.node.0, x.span)
                        }
                    }
                    None => {
                        // All identifiers must be resolved at this point,
                        // but we don't stop after scope resolution error,
                        // so this code is reachable.
                        Ty::Any
                    }
                }
            }
            ExprP::Lambda(_) => {
                self.approximation("We don't type check lambdas", ());
                Ty::name("function")
            }
            ExprP::Literal(x) => match x {
                AstLiteral::Int(_) => Ty::int(),
                AstLiteral::Float(_) => Ty::float(),
                AstLiteral::String(_) => Ty::string(),
            },
            ExprP::Not(x) => {
                if self.expression_type(x).is_never() {
                    Ty::Never
                } else {
                    Ty::bool()
                }
            }
            ExprP::Minus(x) => {
                self.expression_primitive(TypingAttr::UnOp(TypingUnOp::Minus), &[&**x], span)
            }
            ExprP::Plus(x) => {
                self.expression_primitive(TypingAttr::UnOp(TypingUnOp::Plus), &[&**x], span)
            }
            ExprP::BitNot(x) => {
                self.expression_primitive(TypingAttr::UnOp(TypingUnOp::BitNot), &[&**x], span)
            }
            ExprP::Op(lhs, op, rhs) => {
                let lhs = self.expression_type_spanned(lhs);
                let rhs = self.expression_type_spanned(rhs);
                let bool_ret = if lhs.is_never() || rhs.is_never() {
                    Ty::Never
                } else {
                    Ty::bool()
                };
                match op {
                    BinOp::And | BinOp::Or => {
                        if lhs.is_never() {
                            Ty::Never
                        } else {
                            Ty::union2(lhs.node, rhs.node)
                        }
                    }
                    BinOp::Equal | BinOp::NotEqual => {
                        // It's not an error to compare two different types, but it is pointless
                        self.validate_type(&lhs, &rhs, span);
                        bool_ret
                    }
                    BinOp::In | BinOp::NotIn => {
                        // We dispatch `x in y` as y.__in__(x) as that's how we validate
                        self.expression_primitive_ty(
                            TypingAttr::BinOp(TypingBinOp::In),
                            rhs.node,
                            vec![lhs],
                            span,
                        );
                        // Ignore the return type, we know it's always a bool
                        bool_ret
                    }
                    BinOp::Less | BinOp::LessOrEqual | BinOp::Greater | BinOp::GreaterOrEqual => {
                        self.expression_primitive_ty(
                            TypingAttr::BinOp(TypingBinOp::Less),
                            lhs.node,
                            vec![rhs],
                            span,
                        );
                        bool_ret
                    }
                    BinOp::Subtract => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::Sub),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::Add => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::Add),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::Multiply => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::Mul),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::Percent => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::Percent),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::Divide => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::Div),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::FloorDivide => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::FloorDiv),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::BitAnd => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::BitAnd),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::BitOr => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::BitOr),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::BitXor => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::BitXor),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::LeftShift => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::LeftShift),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                    BinOp::RightShift => self.expression_primitive_ty(
                        TypingAttr::BinOp(TypingBinOp::RightShift),
                        lhs.node,
                        vec![rhs],
                        span,
                    ),
                }
            }
            ExprP::If(c_t_f) => {
                let c = self.expression_type(&c_t_f.0);
                let t = self.expression_type(&c_t_f.1);
                let f = self.expression_type(&c_t_f.2);
                if c.is_never() {
                    Ty::Never
                } else {
                    Ty::union2(t, f)
                }
            }
            ExprP::List(xs) => {
                let ts = xs.map(|x| self.expression_type(x));
                Ty::list(Ty::unions(ts))
            }
            ExprP::Dict(xs) => {
                let (ks, vs) = xs
                    .iter()
                    .map(|(k, v)| (self.expression_type(k), self.expression_type(v)))
                    .unzip();
                Ty::dict(Ty::unions(ks), Ty::unions(vs))
            }
            ExprP::ListComprehension(a, b, c) => {
                self.check_comprehension(b, c);
                Ty::list(self.expression_type(a))
            }
            ExprP::DictComprehension(k_v, b, c) => {
                self.check_comprehension(b, c);
                Ty::dict(self.expression_type(&k_v.0), self.expression_type(&k_v.1))
            }
        }
    }
}


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

use std::fmt::Display;

use crate::codemap::CodeMap;
use crate::codemap::Span;
use crate::eval::compiler::EvalException;

/// Internal error, bug in the typechecker.
pub(crate) struct InternalError(EvalException);

impl InternalError {
    #[cold]
    pub(crate) fn msg(message: impl Display, span: Span, codemap: &CodeMap) -> InternalError {
        InternalError(EvalException::new(
            anyhow::anyhow!("{} (internal error)", message),
            span,
            codemap,
        ))
    }

    #[cold]
    pub(crate) fn from_eval_exception(e: EvalException) -> InternalError {
        InternalError(e)
    }

    #[cold]
    pub(crate) fn into_anyhow(self) -> anyhow::Error {
        self.0.into_anyhow()
    }

    #[cold]
    pub(crate) fn into_eval_exception(self) -> EvalException {
        self.0
    }
}

/// Errors used in typechecker API. Error has a span.
pub struct TypingError(EvalException);

impl TypingError {
    // TODO(nga): some errors we create, we ignore later. For example, when typechecking a union,
    //   if either variant is good, we ignore the other variant errors.
    //   So we pay for expensive error creation we ignore. Make this function cheap.
    #[cold]
    pub(crate) fn msg(message: impl Display, span: Span, codemap: &CodeMap) -> TypingError {
        TypingError(EvalException::new(
            anyhow::Error::msg(message.to_string()),
            span,
            codemap,
        ))
    }

    #[cold]
    pub(crate) fn new(error: anyhow::Error, span: Span, codemap: &CodeMap) -> TypingError {
        TypingError(EvalException::new(error, span, codemap))
    }

    #[cold]
    pub(crate) fn into_anyhow(self) -> anyhow::Error {
        self.0.into_anyhow()
    }

    #[cold]
    pub(crate) fn into_eval_exception(self) -> EvalException {
        self.0
    }
}


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

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

use allocative::Allocative;

use crate::codemap::Span;
use crate::codemap::Spanned;
use crate::typing::error::TypingError;
use crate::typing::ty::TyCustomImpl;
use crate::typing::Ty;
use crate::typing::TypingAttr;
use crate::typing::TypingOracleCtx;

/// An argument being passed to a function
#[derive(Debug)]
pub enum Arg {
    /// A positional argument.
    Pos(Ty),
    /// A named argument.
    Name(String, Ty),
    /// A `*args`.
    Args(Ty),
    /// A `**kwargs`.
    Kwargs(Ty),
}

/// The type of a parameter - can be positional, by name, `*args` or `**kwargs`.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Allocative)]
pub enum ParamMode {
    /// Parameter can only be passed by position.
    PosOnly,
    /// Parameter can be passed by position or name.
    PosOrName(String),
    /// Parameter can only be passed by name.
    NameOnly(String),
    /// Parameter is `*args`.
    Args,
    /// Parameter is `**kwargs`.
    Kwargs,
}

/// A parameter argument to a function
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Allocative)]
pub struct Param {
    /// The type of parameter
    pub mode: ParamMode,
    /// Whether the parameter have a default value or is otherwise optional
    pub optional: bool,
    /// The type of the parameter
    pub ty: Ty,
}

impl Param {
    /// Create a [`ParamMode::PosOnly`] parameter.
    pub fn pos_only(ty: Ty) -> Self {
        Self {
            mode: ParamMode::PosOnly,
            optional: false,
            ty,
        }
    }

    /// Create a [`ParamMode::NameOnly`] parameter.
    pub fn name_only(name: &str, ty: Ty) -> Self {
        Self {
            mode: ParamMode::NameOnly(name.to_owned()),
            optional: false,
            ty,
        }
    }

    /// Create a [`ParamMode::PosOrName`] parameter.
    pub fn pos_or_name(name: &str, ty: Ty) -> Self {
        Self {
            mode: ParamMode::PosOrName(name.to_owned()),
            optional: false,
            ty,
        }
    }

    /// Make a parameter optional.
    pub fn optional(self) -> Self {
        Self {
            optional: true,
            ..self
        }
    }

    /// Create a [`ParamMode::Args`] parameter.
    pub fn args(ty: Ty) -> Self {
        Self {
            mode: ParamMode::Args,
            optional: true,
            ty,
        }
    }

    /// Create a [`ParamMode::Kwargs`] parameter.
    pub fn kwargs(ty: Ty) -> Self {
        Self {
            mode: ParamMode::Kwargs,
            optional: true,
            ty,
        }
    }

    pub(crate) fn allows_pos(&self) -> bool {
        match self.mode {
            ParamMode::PosOnly | ParamMode::PosOrName(_) | ParamMode::Args => true,
            ParamMode::NameOnly(_) | ParamMode::Kwargs => false,
        }
    }

    pub(crate) fn allows_many(&self) -> bool {
        match self.mode {
            ParamMode::Args | ParamMode::Kwargs => true,
            _ => false,
        }
    }

    /// Get a display name for this parameter.
    pub fn name(&self) -> &str {
        match &self.mode {
            ParamMode::PosOnly => "_",
            ParamMode::PosOrName(x) => x,
            ParamMode::NameOnly(x) => x,
            ParamMode::Args => "*args",
            ParamMode::Kwargs => "**kwargs",
        }
    }
}

/// Custom function typechecker.
pub trait TyCustomFunctionImpl:
    Clone + Debug + Eq + Ord + Allocative + Send + Sync + 'static
{
    fn has_type_attr(&self) -> bool {
        false
    }

    fn validate_call(
        &self,
        span: Span,
        args: &[Spanned<Arg>],
        oracle: TypingOracleCtx,
    ) -> Result<Ty, TypingError>;
}

#[derive(
    Allocative,
    Eq,
    PartialEq,
    Ord,
    PartialOrd,
    Debug,
    Clone,
    derive_more::Display
)]
#[display(fmt = "\"function\"")]
pub struct TyCustomFunction<F: TyCustomFunctionImpl>(pub F);

impl<F: TyCustomFunctionImpl> TyCustomImpl for TyCustomFunction<F> {
    fn as_name(&self) -> Option<&str> {
        Some("function")
    }

    fn validate_call(
        &self,
        span: Span,
        args: &[Spanned<Arg>],
        oracle: TypingOracleCtx,
    ) -> Result<Ty, TypingError> {
        self.0.validate_call(span, args, oracle)
    }

    fn attribute(&self, attr: TypingAttr) -> Option<Result<Ty, ()>> {
        if attr == TypingAttr::Regular("type") && self.0.has_type_attr() {
            Some(Ok(Ty::string()))
        } else {
            // TODO(nga): some types (like record types) pretend to be functions,
            //   so we don't know which attributes they have.
            None
        }
    }

    fn union2(x: Box<Self>, other: Box<Self>) -> Result<Box<Self>, (Box<Self>, Box<Self>)> {
        if x == other { Ok(x) } else { Err((x, other)) }
    }
}

/// A function.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Allocative)]
pub struct TyFunction {
    /// The `.type` property of the function, often `""`.
    pub type_attr: String,
    /// The parameters to the function.
    pub params: Vec<Param>,
    /// The result type of the function.
    pub result: Box<Ty>,
}

impl Display for TyFunction {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let TyFunction { params, result, .. } = self;
        write!(f, "def(")?;
        let mut first = true;
        for param in params {
            if !first {
                write!(f, ", ")?;
                first = false;
            }
            let opt = if param.optional { "=.." } else { "" };
            match &param.mode {
                ParamMode::PosOnly => write!(f, "#: {}{}", param.ty, opt)?,
                ParamMode::PosOrName(name) => write!(f, "#{}: {}{}", name, param.ty, opt)?,
                ParamMode::NameOnly(name) => write!(f, "{}: {}{}", name, param.ty, opt)?,
                ParamMode::Args => write!(f, "*args: {}", param.ty)?,
                ParamMode::Kwargs => write!(f, "**kwargs: {}", param.ty)?,
            }
        }
        write!(f, ") -> {}", result)
    }
}

impl TyCustomFunctionImpl for TyFunction {
    fn has_type_attr(&self) -> bool {
        !self.type_attr.is_empty()
    }

    fn validate_call(
        &self,
        span: Span,
        args: &[Spanned<Arg>],
        oracle: TypingOracleCtx,
    ) -> Result<Ty, TypingError> {
        oracle.validate_fn_call(span, self, args)
    }
}


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

// Notes:
//   We deal with list.append/list.extend/list.insert, which mutate their list argument
//   We ignore dict.setdefault/dict.update, as these are pretty complex functions
//   We consider "non-sensicle" operations like list.remove and == to have implied types that make them meaningful
//       even if they don't fail when doing something silly

//! Types required to support the [`typecheck`](crate::syntax::AstModule::typecheck) function.

pub(crate) mod bindings;
pub(crate) mod ctx;
pub(crate) mod error;
pub(crate) mod function;
pub(crate) mod mode;
pub(crate) mod oracle;
pub(crate) mod structs;
pub(crate) mod ty;
pub(crate) mod typecheck;

#[cfg(test)]
mod tests;

pub use bindings::Interface;
pub use function::Arg;
pub use function::Param;
pub use function::ParamMode;
pub use function::TyFunction;
pub use oracle::ctx::TypingOracleCtx;
pub use oracle::docs::OracleDocs;
pub use oracle::standard::OracleStandard;
pub use oracle::traits::OracleSeq;
pub use oracle::traits::TypingAttr;
pub use oracle::traits::TypingBinOp;
pub use oracle::traits::TypingOracle;
pub use oracle::traits::TypingUnOp;
pub use structs::TyStruct;
pub use ty::Approximation;
pub use ty::Ty;
pub use ty::TyName;
pub use ty::TyUnion;
pub use typecheck::TypeMap;


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

#[derive(Copy, Clone, Dupe, Debug, Eq, PartialEq)]
pub(crate) enum TypecheckMode {
    /// Lint mode.
    Lint,
    /// Invoked from the compiler.
    Compiler,
}


/*
 * 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::collections::BTreeMap;
use std::fmt;
use std::fmt::Display;
use std::fmt::Formatter;

use allocative::Allocative;

use crate::codemap::Span;
use crate::codemap::Spanned;
use crate::typing::error::TypingError;
use crate::typing::ty::TyCustomImpl;
use crate::typing::Arg;
use crate::typing::Param;
use crate::typing::Ty;
use crate::typing::TypingAttr;
use crate::typing::TypingBinOp;
use crate::typing::TypingOracleCtx;

/// Struct type.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Allocative)]
pub struct TyStruct {
    /// The fields that are definitely present in the struct, with their types.
    pub(crate) fields: BTreeMap<String, Ty>,
    /// [`true`] if there might be additional fields not captured above,
    /// [`false`] if this struct has no extra members.
    pub(crate) extra: bool,
}

impl TyStruct {
    /// Any struct.
    pub fn any() -> TyStruct {
        TyStruct {
            fields: BTreeMap::new(),
            extra: true,
        }
    }
}

impl TyCustomImpl for TyStruct {
    fn as_name(&self) -> Option<&str> {
        Some("struct")
    }

    fn validate_call(
        &self,
        span: Span,
        _args: &[Spanned<Arg>],
        oracle: TypingOracleCtx,
    ) -> Result<Ty, TypingError> {
        Err(oracle.msg_error(span, "struct instances are not callable"))
    }

    fn attribute(&self, attr: TypingAttr) -> Option<Result<Ty, ()>> {
        match attr {
            TypingAttr::Regular(attr) => match self.fields.get(attr) {
                Some(ty) => Some(Ok(ty.clone())),
                None if self.extra => Some(Ok(Ty::Any)),
                _ => Some(Err(())),
            },
            TypingAttr::BinOp(TypingBinOp::Less) => Some(Ok(Ty::function(
                vec![Param::pos_only(Ty::custom(TyStruct::any()))],
                Ty::bool(),
            ))),
            _ => Some(Err(())),
        }
    }

    fn union2(a: Box<Self>, b: Box<Self>) -> Result<Box<Self>, (Box<Self>, Box<Self>)> {
        if a == b {
            // Fast path.
            Ok(a)
        } else if a.extra == b.extra && itertools::equal(a.fields.keys(), b.fields.keys()) {
            let mut fields = BTreeMap::new();
            for ((a_k, a_v), (b_k, b_v)) in a.fields.into_iter().zip(b.fields) {
                assert_eq!(a_k, b_k);
                fields.insert(a_k, Ty::union2(a_v, b_v));
            }
            Ok(Box::new(TyStruct {
                fields,
                extra: a.extra,
            }))
        } else {
            Err((a, b))
        }
    }
}

impl Display for TyStruct {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let TyStruct { fields, extra } = self;
        display_container::fmt_container(
            f,
            "struct(",
            ")",
            display_container::iter_display_chain(
                fields.iter().map(|(k, v)| format!("{} = {}", k, v)),
                extra.then_some(".."),
            ),
        )
    }
}